1
2
3
4
5 package websocket
6
7 import (
8 "bufio"
9 "bytes"
10 "net"
11 "net/http"
12 "reflect"
13 "strings"
14 "testing"
15 )
16
17 var subprotocolTests = []struct {
18 h string
19 protocols []string
20 }{
21 {"", nil},
22 {"foo", []string{"foo"}},
23 {"foo,bar", []string{"foo", "bar"}},
24 {"foo, bar", []string{"foo", "bar"}},
25 {" foo, bar", []string{"foo", "bar"}},
26 {" foo, bar ", []string{"foo", "bar"}},
27 }
28
29 func TestSubprotocols(t *testing.T) {
30 for _, st := range subprotocolTests {
31 r := http.Request{Header: http.Header{"Sec-Websocket-Protocol": {st.h}}}
32 protocols := Subprotocols(&r)
33 if !reflect.DeepEqual(st.protocols, protocols) {
34 t.Errorf("SubProtocols(%q) returned %#v, want %#v", st.h, protocols, st.protocols)
35 }
36 }
37 }
38
39 var isWebSocketUpgradeTests = []struct {
40 ok bool
41 h http.Header
42 }{
43 {false, http.Header{"Upgrade": {"websocket"}}},
44 {false, http.Header{"Connection": {"upgrade"}}},
45 {true, http.Header{"Connection": {"upgRade"}, "Upgrade": {"WebSocket"}}},
46 }
47
48 func TestIsWebSocketUpgrade(t *testing.T) {
49 for _, tt := range isWebSocketUpgradeTests {
50 ok := IsWebSocketUpgrade(&http.Request{Header: tt.h})
51 if tt.ok != ok {
52 t.Errorf("IsWebSocketUpgrade(%v) returned %v, want %v", tt.h, ok, tt.ok)
53 }
54 }
55 }
56
57 var checkSameOriginTests = []struct {
58 ok bool
59 r *http.Request
60 }{
61 {false, &http.Request{Host: "example.org", Header: map[string][]string{"Origin": {"https://other.org"}}}},
62 {true, &http.Request{Host: "example.org", Header: map[string][]string{"Origin": {"https://example.org"}}}},
63 {true, &http.Request{Host: "Example.org", Header: map[string][]string{"Origin": {"https://example.org"}}}},
64 }
65
66 func TestCheckSameOrigin(t *testing.T) {
67 for _, tt := range checkSameOriginTests {
68 ok := checkSameOrigin(tt.r)
69 if tt.ok != ok {
70 t.Errorf("checkSameOrigin(%+v) returned %v, want %v", tt.r, ok, tt.ok)
71 }
72 }
73 }
74
75 type reuseTestResponseWriter struct {
76 brw *bufio.ReadWriter
77 http.ResponseWriter
78 }
79
80 func (resp *reuseTestResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
81 return fakeNetConn{strings.NewReader(""), &bytes.Buffer{}}, resp.brw, nil
82 }
83
84 var bufioReuseTests = []struct {
85 n int
86 reuse bool
87 }{
88 {4096, true},
89 {128, false},
90 }
91
92 func TestBufioReuse(t *testing.T) {
93 for i, tt := range bufioReuseTests {
94 br := bufio.NewReaderSize(strings.NewReader(""), tt.n)
95 bw := bufio.NewWriterSize(&bytes.Buffer{}, tt.n)
96 resp := &reuseTestResponseWriter{
97 brw: bufio.NewReadWriter(br, bw),
98 }
99 upgrader := Upgrader{}
100 c, err := upgrader.Upgrade(resp, &http.Request{
101 Method: http.MethodGet,
102 Header: http.Header{
103 "Upgrade": []string{"websocket"},
104 "Connection": []string{"upgrade"},
105 "Sec-Websocket-Key": []string{"dGhlIHNhbXBsZSBub25jZQ=="},
106 "Sec-Websocket-Version": []string{"13"},
107 }}, nil)
108 if err != nil {
109 t.Fatal(err)
110 }
111 if reuse := c.br == br; reuse != tt.reuse {
112 t.Errorf("%d: buffered reader reuse=%v, want %v", i, reuse, tt.reuse)
113 }
114 writeBuf := bufioWriterBuffer(c.NetConn(), bw)
115 if reuse := &c.writeBuf[0] == &writeBuf[0]; reuse != tt.reuse {
116 t.Errorf("%d: write buffer reuse=%v, want %v", i, reuse, tt.reuse)
117 }
118 }
119 }
120
View as plain text