conn_test.go raw
1 package ws
2
3 import (
4 "bufio"
5 "crypto/rand"
6 "encoding/binary"
7 "fmt"
8 "io"
9 mrand "math/rand"
10 "net"
11 "strings"
12 "testing"
13 "time"
14 )
15
16 // listenTest picks an explicit loopback port. TinyGo's listener reports :0
17 // rather than the assigned port, so a random explicit port with retries is the
18 // portable way to get one for a test.
19 func listenTest(t *testing.T) net.Listener {
20 t.Helper()
21 for i := 0; i < 40; i++ {
22 addr := fmt.Sprintf("127.0.0.1:%d", 20000+mrand.Intn(2000))
23 ln, err := net.Listen("tcp", addr)
24 if err == nil {
25 return ln
26 }
27 }
28 t.Fatal("could not bind a test port")
29 return nil
30 }
31
32 func newFramePair(t *testing.T, max int64) (*Conn, net.Conn) {
33 t.Helper()
34 ln := listenTest(t)
35 defer ln.Close()
36 type accepted struct {
37 c net.Conn
38 }
39 ch := make(chan net.Conn, 1)
40 go func() {
41 c, err := ln.Accept()
42 if err == nil {
43 ch <- c
44 }
45 }()
46 client, err := net.Dial("tcp", ln.Addr().String())
47 if err != nil {
48 t.Fatalf("dial: %v", err)
49 }
50 server := <-ch
51 brw := bufio.NewReadWriter(bufio.NewReader(server), bufio.NewWriter(server))
52 c := &Conn{conn: server, br: brw.Reader, bw: brw.Writer, MaxMessage: max}
53 t.Cleanup(func() {
54 client.Close()
55 server.Close()
56 })
57 return c, client
58 }
59
60 func clientFrame(masked, fin bool, op byte, payload []byte) []byte {
61 b0 := op
62 if fin {
63 b0 |= 0x80
64 }
65 out := []byte{b0}
66 n := len(payload)
67 maskBit := byte(0)
68 if masked {
69 maskBit = 0x80
70 }
71 switch {
72 case n < 126:
73 out = append(out, maskBit|byte(n))
74 case n <= 0xffff:
75 out = append(out, maskBit|126, byte(n>>8), byte(n))
76 default:
77 out = append(out, maskBit|127)
78 var ext [8]byte
79 binary.BigEndian.PutUint64(ext[:], uint64(n))
80 out = append(out, ext[:]...)
81 }
82 if !masked {
83 return append(out, payload...)
84 }
85 var key [4]byte
86 rand.Read(key[:])
87 out = append(out, key[:]...)
88 for i, b := range payload {
89 out = append(out, b^key[i&3])
90 }
91 return out
92 }
93
94 func writeClient(t *testing.T, c net.Conn, masked, fin bool, op byte, payload []byte) {
95 t.Helper()
96 if _, err := c.Write(clientFrame(masked, fin, op, payload)); err != nil {
97 t.Fatalf("client write: %v", err)
98 }
99 }
100
101 func readServerFrame(t *testing.T, c net.Conn) (op byte, payload []byte) {
102 t.Helper()
103 c.SetReadDeadline(time.Now().Add(3 * time.Second))
104 hdr := make([]byte, 2)
105 if _, err := io.ReadFull(c, hdr); err != nil {
106 t.Fatalf("server frame header: %v", err)
107 }
108 if hdr[1]&0x80 != 0 {
109 t.Fatal("server frame is masked")
110 }
111 n := int64(hdr[1] & 0x7f)
112 switch n {
113 case 126:
114 ext := make([]byte, 2)
115 io.ReadFull(c, ext)
116 n = int64(binary.BigEndian.Uint16(ext))
117 case 127:
118 ext := make([]byte, 8)
119 io.ReadFull(c, ext)
120 n = int64(binary.BigEndian.Uint64(ext))
121 }
122 payload = make([]byte, n)
123 if _, err := io.ReadFull(c, payload); err != nil {
124 t.Fatalf("server frame payload: %v", err)
125 }
126 return hdr[0] & 0x0f, payload
127 }
128
129 func TestUpgradeHandshake(t *testing.T) {
130 const key = "dGhlIHNhbXBsZSBub25jZQ=="
131 headers := map[string]string{
132 "upgrade": "websocket",
133 "connection": "keep-alive, Upgrade",
134 "sec-websocket-version": "13",
135 "sec-websocket-key": key,
136 }
137 server, client := net.Pipe()
138 defer server.Close()
139 defer client.Close()
140 go func() {
141 brw := bufio.NewReadWriter(bufio.NewReader(server), bufio.NewWriter(server))
142 Upgrade(server, brw, headers)
143 }()
144
145 br := bufio.NewReader(client)
146 line, err := br.ReadString('\n')
147 if err != nil {
148 t.Fatal(err)
149 }
150 if !strings.Contains(line, "101") {
151 t.Fatalf("status: %q", line)
152 }
153 got := map[string]string{}
154 for {
155 h, err := br.ReadString('\n')
156 if err != nil {
157 t.Fatal(err)
158 }
159 if h == "\r\n" {
160 break
161 }
162 parts := strings.SplitN(strings.TrimSpace(h), ":", 2)
163 got[strings.ToLower(parts[0])] = strings.TrimSpace(parts[1])
164 }
165 // RFC 6455 example accept value for this key.
166 if got["sec-websocket-accept"] != "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=" {
167 t.Fatalf("accept: %q", got["sec-websocket-accept"])
168 }
169 if got["upgrade"] != "websocket" || got["connection"] != "Upgrade" {
170 t.Fatalf("upgrade headers: %+v", got)
171 }
172 }
173
174 func TestUpgradeRejectsBadHeaders(t *testing.T) {
175 base := map[string]string{
176 "upgrade": "websocket",
177 "connection": "Upgrade",
178 "sec-websocket-version": "13",
179 "sec-websocket-key": "dGhlIHNhbXBsZSBub25jZQ==",
180 }
181 clone := func() map[string]string {
182 m := make(map[string]string, len(base))
183 for k, v := range base {
184 m[k] = v
185 }
186 return m
187 }
188 cases := []struct {
189 name string
190 mod func(map[string]string)
191 }{
192 {"no upgrade", func(m map[string]string) { delete(m, "upgrade") }},
193 {"wrong upgrade", func(m map[string]string) { m["upgrade"] = "h2c" }},
194 {"no connection", func(m map[string]string) { delete(m, "connection") }},
195 {"old version", func(m map[string]string) { m["sec-websocket-version"] = "8" }},
196 {"no key", func(m map[string]string) { delete(m, "sec-websocket-key") }},
197 }
198 for _, tc := range cases {
199 t.Run(tc.name, func(t *testing.T) {
200 h := clone()
201 tc.mod(h)
202 server, client := net.Pipe()
203 defer server.Close()
204 defer client.Close()
205 brw := bufio.NewReadWriter(bufio.NewReader(server), bufio.NewWriter(server))
206 if c, err := Upgrade(server, brw, h); err == nil || c != nil {
207 t.Fatalf("expected rejection, got conn=%v err=%v", c, err)
208 }
209 })
210 }
211 }
212
213 func TestReadText(t *testing.T) {
214 c, client := newFramePair(t, 1<<20)
215 writeClient(t, client, true, true, OpText, []byte("hello"))
216 op, data, err := c.ReadMessage()
217 if err != nil || op != OpText || string(data) != "hello" {
218 t.Fatalf("op=%d data=%q err=%v", op, data, err)
219 }
220 }
221
222 func TestReadFragmentedText(t *testing.T) {
223 c, client := newFramePair(t, 1<<20)
224 writeClient(t, client, true, false, OpText, []byte("hel"))
225 writeClient(t, client, true, true, OpContinuation, []byte("lo"))
226 op, data, err := c.ReadMessage()
227 if err != nil || op != OpText || string(data) != "hello" {
228 t.Fatalf("op=%d data=%q err=%v", op, data, err)
229 }
230 }
231
232 func TestPingPongThenText(t *testing.T) {
233 c, client := newFramePair(t, 1<<20)
234 writeClient(t, client, true, true, OpPing, []byte("pingpayload"))
235 writeClient(t, client, true, true, OpText, []byte("after"))
236 // ReadMessage answers the ping before returning the text.
237 op, data, err := c.ReadMessage()
238 if err != nil || op != OpText || string(data) != "after" {
239 t.Fatalf("op=%d data=%q err=%v", op, data, err)
240 }
241 br := bufio.NewReader(client)
242 client.SetReadDeadline(time.Now().Add(2 * time.Second))
243 hdr := make([]byte, 2)
244 if _, err := io.ReadFull(br, hdr); err != nil {
245 t.Fatalf("pong header: %v", err)
246 }
247 if hdr[0]&0x0f != OpPong {
248 t.Fatalf("expected pong, got opcode %d", hdr[0]&0x0f)
249 }
250 payload := make([]byte, hdr[1]&0x7f)
251 io.ReadFull(br, payload)
252 if string(payload) != "pingpayload" {
253 t.Fatalf("pong payload %q", payload)
254 }
255 }
256
257 func TestUnmaskedFrameRejected(t *testing.T) {
258 c, client := newFramePair(t, 1<<20)
259 writeClient(t, client, false, true, OpText, []byte("unmasked"))
260 if _, _, err := c.ReadMessage(); err == nil {
261 t.Fatal("unmasked frame accepted")
262 }
263 op, payload := readServerFrame(t, client)
264 if op != OpClose || len(payload) < 2 || binary.BigEndian.Uint16(payload) != 1002 {
265 t.Fatalf("expected close 1002, op=%d payload=%x", op, payload)
266 }
267 }
268
269 func TestOversizedMessageRejected(t *testing.T) {
270 c, client := newFramePair(t, 8)
271 writeClient(t, client, true, true, OpText, []byte("this is longer than eight"))
272 if _, _, err := c.ReadMessage(); err != ErrTooLarge {
273 t.Fatalf("err=%v, want ErrTooLarge", err)
274 }
275 op, payload := readServerFrame(t, client)
276 if op != OpClose || binary.BigEndian.Uint16(payload) != 1009 {
277 t.Fatalf("expected close 1009, op=%d payload=%x", op, payload)
278 }
279 }
280
281 func TestInvalidUTF8Rejected(t *testing.T) {
282 c, client := newFramePair(t, 1<<20)
283 writeClient(t, client, true, true, OpText, []byte{0xff, 0xfe})
284 if _, _, err := c.ReadMessage(); err == nil {
285 t.Fatal("invalid UTF-8 accepted")
286 }
287 op, payload := readServerFrame(t, client)
288 if op != OpClose || binary.BigEndian.Uint16(payload) != 1007 {
289 t.Fatalf("expected close 1007, op=%d payload=%x", op, payload)
290 }
291 }
292
293 func TestClientClose(t *testing.T) {
294 c, client := newFramePair(t, 1<<20)
295 payload := []byte{0x03, 0xe8} // 1000, no reason
296 writeClient(t, client, true, true, OpClose, payload)
297 if _, _, err := c.ReadMessage(); err != ErrClosed {
298 t.Fatalf("err=%v, want ErrClosed", err)
299 }
300 op, reply := readServerFrame(t, client)
301 if op != OpClose || len(reply) < 2 || binary.BigEndian.Uint16(reply) != 1000 {
302 t.Fatalf("close reply: op=%d payload=%x", op, reply)
303 }
304 }
305
306 func TestWriteFrameLengths(t *testing.T) {
307 sizes := []int{0, 5, 125, 126, 4096, 70000}
308 for _, n := range sizes {
309 c, client := newFramePair(t, 1<<20)
310 payload := make([]byte, n)
311 for i := range payload {
312 payload[i] = byte(i * 7)
313 }
314 if err := c.WriteText(payload); err != nil {
315 t.Fatalf("size %d: %v", n, err)
316 }
317 op, got := readServerFrame(t, client)
318 if op != OpText || len(got) != n {
319 t.Fatalf("size %d: op=%d len=%d", n, op, len(got))
320 }
321 for i := range got {
322 if got[i] != payload[i] {
323 t.Fatalf("size %d: payload mismatch", n)
324 }
325 }
326 client.Close()
327 }
328 }
329
330 func TestCloseIsIdempotent(t *testing.T) {
331 c, _ := newFramePair(t, 1<<20)
332 if err := c.Close(1000, "bye"); err != nil {
333 t.Fatal(err)
334 }
335 if err := c.Close(1000, "again"); err != nil {
336 t.Fatalf("second close: %v", err)
337 }
338 if err := c.WriteText([]byte("x")); err != ErrClosed {
339 t.Fatalf("write after close: %v", err)
340 }
341 }
342
343 func TestRemoteAddr(t *testing.T) {
344 c, _ := newFramePair(t, 1<<20)
345 if c.RemoteAddr() == "" {
346 t.Fatal("empty remote address")
347 }
348 c.SetReadDeadline(time.Now().Add(time.Second))
349 c.SetWriteDeadline(time.Now().Add(time.Second))
350 }
351