client.go raw
1 // Package ws implements the client side of RFC 6455 WebSockets.
2 //
3 // It is deliberately small: a relay connection needs the opening handshake,
4 // masked client frames, fragmentation, ping/pong, close, and one bounded
5 // message reader. Server frames are never masked. One goroutine reads and one
6 // goroutine writes; WriteText serialises concurrent writers.
7 package ws
8
9 import (
10 "bufio"
11 "crypto/rand"
12 "crypto/sha1"
13 "crypto/tls"
14 "encoding/base64"
15 "encoding/binary"
16 "errors"
17 "fmt"
18 "io"
19 "net"
20 "net/url"
21 "strings"
22 "sync"
23 "time"
24 )
25
26 // Frame opcodes.
27 const (
28 OpContinuation = 0x0
29 OpText = 0x1
30 OpBinary = 0x2
31 OpClose = 0x8
32 OpPing = 0x9
33 OpPong = 0xa
34 )
35
36 const acceptGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
37
38 var (
39 // ErrProtocol is returned for a framing or handshake violation.
40 ErrProtocol = errors.New("ws: protocol error")
41 // ErrTooLarge is returned when a message exceeds MaxMessage.
42 ErrTooLarge = errors.New("ws: message too large")
43 // ErrClosed is returned after the peer closed the connection.
44 ErrClosed = errors.New("ws: closed")
45 )
46
47 // Dialer opens WebSocket connections.
48 type Dialer struct {
49 // Timeout bounds the TCP connect and TLS handshake, not the WebSocket
50 // lifetime. Zero means no bound.
51 Timeout time.Duration
52 // MaxMessage bounds one (possibly fragmented) message. Zero means 8 MiB.
53 MaxMessage int64
54 // TLSConfig overrides the default TLS configuration for wss URLs. The
55 // ServerName and an empty RootCAs are filled from the URL when unset.
56 TLSConfig *tls.Config
57 // Header holds extra handshake headers, for example Origin or
58 // Authorization. Sec-WebSocket-* headers cannot be overridden.
59 Header map[string]string
60 // PingInterval sends a WebSocket ping when the connection has been idle
61 // for this long. Zero means 45s; a negative value disables keepalive.
62 // Relays commonly close idle sockets; a ping is enough to hold the
63 // subscription open.
64 PingInterval time.Duration
65 }
66
67 // Conn is a client-side WebSocket connection.
68 type Conn struct {
69 conn net.Conn
70 br *bufio.Reader
71
72 wmu sync.Mutex
73 dead bool
74
75 maxMessage int64
76 closeOnce sync.Once
77 closeErr error
78
79 pingStop chan struct{}
80 }
81
82 // Dial opens a WebSocket connection to a ws:// or wss:// URL.
83 func Dial(rawURL string) (*Conn, error) {
84 return (&Dialer{}).Dial(rawURL)
85 }
86
87 // Dial opens a WebSocket connection with d's settings.
88 func (d *Dialer) Dial(rawURL string) (*Conn, error) {
89 u, err := url.Parse(rawURL)
90 if err != nil {
91 return nil, fmt.Errorf("ws: parse url: %w", err)
92 }
93 var tlsInUse bool
94 switch strings.ToLower(u.Scheme) {
95 case "ws":
96 case "wss":
97 tlsInUse = true
98 default:
99 return nil, fmt.Errorf("ws: unsupported scheme %q", u.Scheme)
100 }
101 host := u.Host
102 if host == "" {
103 return nil, errors.New("ws: url has no host")
104 }
105 hostname := u.Hostname()
106 port := u.Port()
107 if port == "" {
108 if tlsInUse {
109 port = "443"
110 } else {
111 port = "80"
112 }
113 }
114 addr := net.JoinHostPort(hostname, port)
115
116 var conn net.Conn
117 if d.Timeout > 0 {
118 conn, err = net.DialTimeout("tcp", addr, d.Timeout)
119 } else {
120 conn, err = net.Dial("tcp", addr)
121 }
122 if err != nil {
123 return nil, fmt.Errorf("ws: dial %s: %w", addr, err)
124 }
125 deadline := time.Time{}
126 if d.Timeout > 0 {
127 deadline = time.Now().Add(d.Timeout)
128 conn.SetDeadline(deadline)
129 }
130 if tlsInUse {
131 cfg := &tls.Config{}
132 if d.TLSConfig != nil {
133 cfg = d.TLSConfig.Clone()
134 }
135 if cfg.ServerName == "" {
136 cfg.ServerName = hostname
137 }
138 tconn := tls.Client(conn, cfg)
139 if err := tconn.Handshake(); err != nil {
140 conn.Close()
141 return nil, fmt.Errorf("ws: tls handshake: %w", err)
142 }
143 conn = tconn
144 }
145
146 path := u.RequestURI()
147 if path == "" {
148 path = "/"
149 }
150 keyBytes := make([]byte, 16)
151 if _, err := rand.Read(keyBytes); err != nil {
152 conn.Close()
153 return nil, fmt.Errorf("ws: random key: %w", err)
154 }
155 key := base64.StdEncoding.EncodeToString(keyBytes)
156
157 var sb strings.Builder
158 sb.WriteString("GET " + path + " HTTP/1.1\r\n")
159 sb.WriteString("Host: " + host + "\r\n")
160 sb.WriteString("Upgrade: websocket\r\n")
161 sb.WriteString("Connection: Upgrade\r\n")
162 sb.WriteString("Sec-WebSocket-Key: " + key + "\r\n")
163 sb.WriteString("Sec-WebSocket-Version: 13\r\n")
164 for name, value := range d.Header {
165 lower := strings.ToLower(name)
166 if strings.HasPrefix(lower, "sec-websocket-") ||
167 lower == "upgrade" || lower == "connection" || lower == "host" {
168 continue
169 }
170 sb.WriteString(name + ": " + value + "\r\n")
171 }
172 sb.WriteString("\r\n")
173 if _, err := conn.Write([]byte(sb.String())); err != nil {
174 conn.Close()
175 return nil, fmt.Errorf("ws: handshake write: %w", err)
176 }
177
178 br := bufio.NewReaderSize(conn, 16<<10)
179 status, err := br.ReadString('\n')
180 if err != nil {
181 conn.Close()
182 return nil, fmt.Errorf("ws: handshake status: %w", err)
183 }
184 if !strings.Contains(status, " 101 ") {
185 conn.Close()
186 return nil, fmt.Errorf("%w: handshake status %q", ErrProtocol, strings.TrimSpace(status))
187 }
188 headers := make(map[string]string, 16)
189 for {
190 line, err := br.ReadString('\n')
191 if err != nil {
192 conn.Close()
193 return nil, fmt.Errorf("ws: handshake headers: %w", err)
194 }
195 line = strings.TrimRight(line, "\r\n")
196 if line == "" {
197 break
198 }
199 name, value, ok := strings.Cut(line, ":")
200 if !ok {
201 conn.Close()
202 return nil, fmt.Errorf("%w: malformed header %q", ErrProtocol, line)
203 }
204 headers[strings.ToLower(strings.TrimSpace(name))] = strings.TrimSpace(value)
205 }
206 if !headerContains(headers["upgrade"], "websocket") ||
207 !headerContains(headers["connection"], "upgrade") {
208 conn.Close()
209 return nil, fmt.Errorf("%w: missing upgrade headers", ErrProtocol)
210 }
211 expected := acceptKey(key)
212 if headers["sec-websocket-accept"] != expected {
213 conn.Close()
214 return nil, fmt.Errorf("%w: bad accept key", ErrProtocol)
215 }
216 if d.Timeout > 0 {
217 conn.SetDeadline(time.Time{})
218 }
219 c := &Conn{conn: conn, br: br, maxMessage: d.MaxMessage}
220 interval := d.PingInterval
221 if interval == 0 {
222 interval = 45 * time.Second
223 }
224 if interval > 0 {
225 c.pingStop = make(chan struct{})
226 go c.keepalive(interval)
227 }
228 return c, nil
229 }
230
231 func (c *Conn) keepalive(interval time.Duration) {
232 ticker := time.NewTicker(interval)
233 defer ticker.Stop()
234 for {
235 select {
236 case <-c.pingStop:
237 return
238 case <-ticker.C:
239 if err := c.writeFrame(OpPing, []byte("nostrino")); err != nil {
240 c.closeConn()
241 return
242 }
243 }
244 }
245 }
246
247 func headerContains(value, token string) bool {
248 for _, part := range strings.Split(value, ",") {
249 if strings.EqualFold(strings.TrimSpace(part), token) {
250 return true
251 }
252 }
253 return false
254 }
255
256 func acceptKey(key string) string {
257 h := sha1.Sum([]byte(key + acceptGUID))
258 return base64.StdEncoding.EncodeToString(h[:])
259 }
260
261 // WriteText writes one unfragmented text message.
262 func (c *Conn) WriteText(payload []byte) error {
263 return c.writeFrame(OpText, payload)
264 }
265
266 // WriteBinary writes one unfragmented binary message.
267 func (c *Conn) WriteBinary(payload []byte) error {
268 return c.writeFrame(OpBinary, payload)
269 }
270
271 func (c *Conn) writeFrame(opcode byte, payload []byte) error {
272 c.wmu.Lock()
273 defer c.wmu.Unlock()
274 if c.dead {
275 return ErrClosed
276 }
277 var hdr [14]byte
278 hdr[0] = 0x80 | opcode
279 n := len(payload)
280 var hlen int
281 switch {
282 case n < 126:
283 hdr[1] = 0x80 | byte(n)
284 hlen = 2
285 case n <= 0xffff:
286 hdr[1] = 0x80 | 126
287 binary.BigEndian.PutUint16(hdr[2:4], uint16(n))
288 hlen = 4
289 default:
290 hdr[1] = 0x80 | 127
291 binary.BigEndian.PutUint64(hdr[2:10], uint64(n))
292 hlen = 10
293 }
294 var mask [4]byte
295 if _, err := rand.Read(mask[:]); err != nil {
296 return fmt.Errorf("ws: mask: %w", err)
297 }
298 copy(hdr[hlen:hlen+4], mask[:])
299 hlen += 4
300 buf := make([]byte, hlen+n)
301 copy(buf, hdr[:hlen])
302 for i := 0; i < n; i++ {
303 buf[hlen+i] = payload[i] ^ mask[i&3]
304 }
305 if _, err := c.conn.Write(buf); err != nil {
306 return err
307 }
308 return nil
309 }
310
311 // ReadMessage reads the next complete text or binary message. Control frames
312 // are handled internally: ping is answered with pong and close ends the
313 // connection. A close frame is reported as io.EOF after ErrClosed is stored.
314 func (c *Conn) ReadMessage() ([]byte, error) {
315 var message []byte
316 var messageOp byte
317 inMessage := false
318 for {
319 op, fin, payload, err := c.readFrame()
320 if err != nil {
321 return nil, err
322 }
323 switch op {
324 case OpPing:
325 if err := c.writeFrame(OpPong, payload); err != nil {
326 return nil, err
327 }
328 continue
329 case OpPong:
330 continue
331 case OpClose:
332 c.writeFrame(OpClose, payload)
333 c.closeConn()
334 return nil, io.EOF
335 case OpText, OpBinary:
336 if inMessage {
337 return nil, fmt.Errorf("%w: new data frame inside a message", ErrProtocol)
338 }
339 messageOp = op
340 message = append(message[:0], payload...)
341 inMessage = true
342 case OpContinuation:
343 if !inMessage {
344 return nil, fmt.Errorf("%w: continuation without a message", ErrProtocol)
345 }
346 message = append(message, payload...)
347 default:
348 return nil, fmt.Errorf("%w: unknown opcode %#x", ErrProtocol, op)
349 }
350 if int64(len(message)) > c.max() {
351 return nil, ErrTooLarge
352 }
353 if fin {
354 if messageOp == 0 {
355 return nil, fmt.Errorf("%w: empty message", ErrProtocol)
356 }
357 return message, nil
358 }
359 }
360 }
361
362 func (c *Conn) max() int64 {
363 if c.maxMessage > 0 {
364 return c.maxMessage
365 }
366 return 8 << 20
367 }
368
369 func (c *Conn) readFrame() (op byte, fin bool, payload []byte, err error) {
370 var hdr [2]byte
371 if _, err = io.ReadFull(c.br, hdr[:]); err != nil {
372 return 0, false, nil, err
373 }
374 fin = hdr[0]&0x80 != 0
375 if hdr[0]&0x70 != 0 {
376 return 0, false, nil, fmt.Errorf("%w: reserved bits set", ErrProtocol)
377 }
378 op = hdr[0] & 0x0f
379 masked := hdr[1]&0x80 != 0
380 length := uint64(hdr[1] & 0x7f)
381 switch length {
382 case 126:
383 var ext [2]byte
384 if _, err = io.ReadFull(c.br, ext[:]); err != nil {
385 return 0, false, nil, err
386 }
387 length = uint64(binary.BigEndian.Uint16(ext[:]))
388 case 127:
389 var ext [8]byte
390 if _, err = io.ReadFull(c.br, ext[:]); err != nil {
391 return 0, false, nil, err
392 }
393 length = binary.BigEndian.Uint64(ext[:])
394 if length>>63 != 0 {
395 return 0, false, nil, fmt.Errorf("%w: negative length", ErrProtocol)
396 }
397 }
398 if length > uint64(c.max()) {
399 return 0, false, nil, ErrTooLarge
400 }
401 if op >= OpClose && (!fin || length > 125) {
402 return 0, false, nil, fmt.Errorf("%w: bad control frame", ErrProtocol)
403 }
404 if masked {
405 // A server must not mask frames. Tolerate and unmask rather than fail,
406 // because some reverse proxies do it; the RFC violation is harmless here.
407 var mask [4]byte
408 if _, err = io.ReadFull(c.br, mask[:]); err != nil {
409 return 0, false, nil, err
410 }
411 payload = make([]byte, length)
412 if _, err = io.ReadFull(c.br, payload); err != nil {
413 return 0, false, nil, err
414 }
415 for i := range payload {
416 payload[i] ^= mask[i&3]
417 }
418 return op, fin, payload, nil
419 }
420 payload = make([]byte, length)
421 if _, err = io.ReadFull(c.br, payload); err != nil {
422 return 0, false, nil, err
423 }
424 return op, fin, payload, nil
425 }
426
427 // Close sends a close frame and closes the underlying socket.
428 func (c *Conn) Close() error {
429 c.closeOnce.Do(func() {
430 if c.pingStop != nil {
431 close(c.pingStop)
432 }
433 c.writeFrame(OpClose, nil)
434 c.closeErr = c.closeConn()
435 })
436 return c.closeErr
437 }
438
439 func (c *Conn) closeConn() error {
440 c.wmu.Lock()
441 c.dead = true
442 c.wmu.Unlock()
443 return c.conn.Close()
444 }
445
446 // SetReadDeadline sets the deadline for future reads.
447 func (c *Conn) SetReadDeadline(t time.Time) error { return c.conn.SetReadDeadline(t) }
448
449 // SetWriteDeadline sets the deadline for future writes.
450 func (c *Conn) SetWriteDeadline(t time.Time) error { return c.conn.SetWriteDeadline(t) }
451
452 // SetDeadline sets both deadlines.
453 func (c *Conn) SetDeadline(t time.Time) error { return c.conn.SetDeadline(t) }
454