blob: ca46d2f793c24eb7593d6ee6ba834d4f437c53da [file] [log] [blame]
khenaidoo59ce9dd2019-11-11 13:05:32 -05001// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package websocket
6
7import (
8 "bufio"
9 "encoding/binary"
10 "errors"
11 "io"
12 "io/ioutil"
13 "math/rand"
14 "net"
15 "strconv"
16 "sync"
17 "time"
18 "unicode/utf8"
19)
20
21const (
22 // Frame header byte 0 bits from Section 5.2 of RFC 6455
23 finalBit = 1 << 7
24 rsv1Bit = 1 << 6
25 rsv2Bit = 1 << 5
26 rsv3Bit = 1 << 4
27
28 // Frame header byte 1 bits from Section 5.2 of RFC 6455
29 maskBit = 1 << 7
30
31 maxFrameHeaderSize = 2 + 8 + 4 // Fixed header + length + mask
32 maxControlFramePayloadSize = 125
33
34 writeWait = time.Second
35
36 defaultReadBufferSize = 4096
37 defaultWriteBufferSize = 4096
38
39 continuationFrame = 0
40 noFrame = -1
41)
42
43// Close codes defined in RFC 6455, section 11.7.
44const (
45 CloseNormalClosure = 1000
46 CloseGoingAway = 1001
47 CloseProtocolError = 1002
48 CloseUnsupportedData = 1003
49 CloseNoStatusReceived = 1005
50 CloseAbnormalClosure = 1006
51 CloseInvalidFramePayloadData = 1007
52 ClosePolicyViolation = 1008
53 CloseMessageTooBig = 1009
54 CloseMandatoryExtension = 1010
55 CloseInternalServerErr = 1011
56 CloseServiceRestart = 1012
57 CloseTryAgainLater = 1013
58 CloseTLSHandshake = 1015
59)
60
61// The message types are defined in RFC 6455, section 11.8.
62const (
63 // TextMessage denotes a text data message. The text message payload is
64 // interpreted as UTF-8 encoded text data.
65 TextMessage = 1
66
67 // BinaryMessage denotes a binary data message.
68 BinaryMessage = 2
69
70 // CloseMessage denotes a close control message. The optional message
71 // payload contains a numeric code and text. Use the FormatCloseMessage
72 // function to format a close message payload.
73 CloseMessage = 8
74
75 // PingMessage denotes a ping control message. The optional message payload
76 // is UTF-8 encoded text.
77 PingMessage = 9
78
khenaidoo26721882021-08-11 17:42:52 -040079 // PongMessage denotes a pong control message. The optional message payload
khenaidoo59ce9dd2019-11-11 13:05:32 -050080 // is UTF-8 encoded text.
81 PongMessage = 10
82)
83
84// ErrCloseSent is returned when the application writes a message to the
85// connection after sending a close message.
86var ErrCloseSent = errors.New("websocket: close sent")
87
88// ErrReadLimit is returned when reading a message that is larger than the
89// read limit set for the connection.
90var ErrReadLimit = errors.New("websocket: read limit exceeded")
91
92// netError satisfies the net Error interface.
93type netError struct {
94 msg string
95 temporary bool
96 timeout bool
97}
98
99func (e *netError) Error() string { return e.msg }
100func (e *netError) Temporary() bool { return e.temporary }
101func (e *netError) Timeout() bool { return e.timeout }
102
khenaidoo26721882021-08-11 17:42:52 -0400103// CloseError represents a close message.
khenaidoo59ce9dd2019-11-11 13:05:32 -0500104type CloseError struct {
khenaidoo59ce9dd2019-11-11 13:05:32 -0500105 // Code is defined in RFC 6455, section 11.7.
106 Code int
107
108 // Text is the optional text payload.
109 Text string
110}
111
112func (e *CloseError) Error() string {
113 s := []byte("websocket: close ")
114 s = strconv.AppendInt(s, int64(e.Code), 10)
115 switch e.Code {
116 case CloseNormalClosure:
117 s = append(s, " (normal)"...)
118 case CloseGoingAway:
119 s = append(s, " (going away)"...)
120 case CloseProtocolError:
121 s = append(s, " (protocol error)"...)
122 case CloseUnsupportedData:
123 s = append(s, " (unsupported data)"...)
124 case CloseNoStatusReceived:
125 s = append(s, " (no status)"...)
126 case CloseAbnormalClosure:
127 s = append(s, " (abnormal closure)"...)
128 case CloseInvalidFramePayloadData:
129 s = append(s, " (invalid payload data)"...)
130 case ClosePolicyViolation:
131 s = append(s, " (policy violation)"...)
132 case CloseMessageTooBig:
133 s = append(s, " (message too big)"...)
134 case CloseMandatoryExtension:
135 s = append(s, " (mandatory extension missing)"...)
136 case CloseInternalServerErr:
137 s = append(s, " (internal server error)"...)
138 case CloseTLSHandshake:
139 s = append(s, " (TLS handshake error)"...)
140 }
141 if e.Text != "" {
142 s = append(s, ": "...)
143 s = append(s, e.Text...)
144 }
145 return string(s)
146}
147
148// IsCloseError returns boolean indicating whether the error is a *CloseError
149// with one of the specified codes.
150func IsCloseError(err error, codes ...int) bool {
151 if e, ok := err.(*CloseError); ok {
152 for _, code := range codes {
153 if e.Code == code {
154 return true
155 }
156 }
157 }
158 return false
159}
160
161// IsUnexpectedCloseError returns boolean indicating whether the error is a
162// *CloseError with a code not in the list of expected codes.
163func IsUnexpectedCloseError(err error, expectedCodes ...int) bool {
164 if e, ok := err.(*CloseError); ok {
165 for _, code := range expectedCodes {
166 if e.Code == code {
167 return false
168 }
169 }
170 return true
171 }
172 return false
173}
174
175var (
176 errWriteTimeout = &netError{msg: "websocket: write timeout", timeout: true, temporary: true}
177 errUnexpectedEOF = &CloseError{Code: CloseAbnormalClosure, Text: io.ErrUnexpectedEOF.Error()}
178 errBadWriteOpCode = errors.New("websocket: bad write message type")
179 errWriteClosed = errors.New("websocket: write closed")
180 errInvalidControlFrame = errors.New("websocket: invalid control frame")
181)
182
183func newMaskKey() [4]byte {
184 n := rand.Uint32()
185 return [4]byte{byte(n), byte(n >> 8), byte(n >> 16), byte(n >> 24)}
186}
187
188func hideTempErr(err error) error {
189 if e, ok := err.(net.Error); ok && e.Temporary() {
190 err = &netError{msg: e.Error(), timeout: e.Timeout()}
191 }
192 return err
193}
194
195func isControl(frameType int) bool {
196 return frameType == CloseMessage || frameType == PingMessage || frameType == PongMessage
197}
198
199func isData(frameType int) bool {
200 return frameType == TextMessage || frameType == BinaryMessage
201}
202
203var validReceivedCloseCodes = map[int]bool{
204 // see http://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number
205
206 CloseNormalClosure: true,
207 CloseGoingAway: true,
208 CloseProtocolError: true,
209 CloseUnsupportedData: true,
210 CloseNoStatusReceived: false,
211 CloseAbnormalClosure: false,
212 CloseInvalidFramePayloadData: true,
213 ClosePolicyViolation: true,
214 CloseMessageTooBig: true,
215 CloseMandatoryExtension: true,
216 CloseInternalServerErr: true,
217 CloseServiceRestart: true,
218 CloseTryAgainLater: true,
219 CloseTLSHandshake: false,
220}
221
222func isValidReceivedCloseCode(code int) bool {
223 return validReceivedCloseCodes[code] || (code >= 3000 && code <= 4999)
224}
225
khenaidoo26721882021-08-11 17:42:52 -0400226// BufferPool represents a pool of buffers. The *sync.Pool type satisfies this
227// interface. The type of the value stored in a pool is not specified.
228type BufferPool interface {
229 // Get gets a value from the pool or returns nil if the pool is empty.
230 Get() interface{}
231 // Put adds a value to the pool.
232 Put(interface{})
233}
234
235// writePoolData is the type added to the write buffer pool. This wrapper is
236// used to prevent applications from peeking at and depending on the values
237// added to the pool.
238type writePoolData struct{ buf []byte }
239
khenaidoo59ce9dd2019-11-11 13:05:32 -0500240// The Conn type represents a WebSocket connection.
241type Conn struct {
242 conn net.Conn
243 isServer bool
244 subprotocol string
245
246 // Write fields
khenaidoo26721882021-08-11 17:42:52 -0400247 mu chan struct{} // used as mutex to protect write to conn
248 writeBuf []byte // frame is constructed in this buffer.
249 writePool BufferPool
250 writeBufSize int
khenaidoo59ce9dd2019-11-11 13:05:32 -0500251 writeDeadline time.Time
252 writer io.WriteCloser // the current writer returned to the application
253 isWriting bool // for best-effort concurrent write detection
254
255 writeErrMu sync.Mutex
256 writeErr error
257
258 enableWriteCompression bool
259 compressionLevel int
260 newCompressionWriter func(io.WriteCloser, int) io.WriteCloser
261
262 // Read fields
khenaidoo26721882021-08-11 17:42:52 -0400263 reader io.ReadCloser // the current reader returned to the application
264 readErr error
265 br *bufio.Reader
266 // bytes remaining in current frame.
267 // set setReadRemaining to safely update this value and prevent overflow
268 readRemaining int64
khenaidoo59ce9dd2019-11-11 13:05:32 -0500269 readFinal bool // true the current message has more frames.
270 readLength int64 // Message size.
271 readLimit int64 // Maximum message size.
272 readMaskPos int
273 readMaskKey [4]byte
274 handlePong func(string) error
275 handlePing func(string) error
276 handleClose func(int, string) error
277 readErrCount int
278 messageReader *messageReader // the current low-level reader
279
280 readDecompress bool // whether last read frame had RSV1 set
281 newDecompressionReader func(io.Reader) io.ReadCloser
282}
283
khenaidoo26721882021-08-11 17:42:52 -0400284func newConn(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int, writeBufferPool BufferPool, br *bufio.Reader, writeBuf []byte) *Conn {
khenaidoo59ce9dd2019-11-11 13:05:32 -0500285
khenaidoo59ce9dd2019-11-11 13:05:32 -0500286 if br == nil {
287 if readBufferSize == 0 {
288 readBufferSize = defaultReadBufferSize
khenaidoo26721882021-08-11 17:42:52 -0400289 } else if readBufferSize < maxControlFramePayloadSize {
290 // must be large enough for control frame
khenaidoo59ce9dd2019-11-11 13:05:32 -0500291 readBufferSize = maxControlFramePayloadSize
292 }
293 br = bufio.NewReaderSize(conn, readBufferSize)
294 }
295
khenaidoo26721882021-08-11 17:42:52 -0400296 if writeBufferSize <= 0 {
297 writeBufferSize = defaultWriteBufferSize
298 }
299 writeBufferSize += maxFrameHeaderSize
300
301 if writeBuf == nil && writeBufferPool == nil {
302 writeBuf = make([]byte, writeBufferSize)
khenaidoo59ce9dd2019-11-11 13:05:32 -0500303 }
304
khenaidoo26721882021-08-11 17:42:52 -0400305 mu := make(chan struct{}, 1)
306 mu <- struct{}{}
khenaidoo59ce9dd2019-11-11 13:05:32 -0500307 c := &Conn{
308 isServer: isServer,
309 br: br,
310 conn: conn,
311 mu: mu,
312 readFinal: true,
313 writeBuf: writeBuf,
khenaidoo26721882021-08-11 17:42:52 -0400314 writePool: writeBufferPool,
315 writeBufSize: writeBufferSize,
khenaidoo59ce9dd2019-11-11 13:05:32 -0500316 enableWriteCompression: true,
317 compressionLevel: defaultCompressionLevel,
318 }
319 c.SetCloseHandler(nil)
320 c.SetPingHandler(nil)
321 c.SetPongHandler(nil)
322 return c
323}
324
khenaidoo26721882021-08-11 17:42:52 -0400325// setReadRemaining tracks the number of bytes remaining on the connection. If n
326// overflows, an ErrReadLimit is returned.
327func (c *Conn) setReadRemaining(n int64) error {
328 if n < 0 {
329 return ErrReadLimit
330 }
331
332 c.readRemaining = n
333 return nil
334}
335
khenaidoo59ce9dd2019-11-11 13:05:32 -0500336// Subprotocol returns the negotiated protocol for the connection.
337func (c *Conn) Subprotocol() string {
338 return c.subprotocol
339}
340
khenaidoo26721882021-08-11 17:42:52 -0400341// Close closes the underlying network connection without sending or waiting
342// for a close message.
khenaidoo59ce9dd2019-11-11 13:05:32 -0500343func (c *Conn) Close() error {
344 return c.conn.Close()
345}
346
347// LocalAddr returns the local network address.
348func (c *Conn) LocalAddr() net.Addr {
349 return c.conn.LocalAddr()
350}
351
352// RemoteAddr returns the remote network address.
353func (c *Conn) RemoteAddr() net.Addr {
354 return c.conn.RemoteAddr()
355}
356
357// Write methods
358
359func (c *Conn) writeFatal(err error) error {
360 err = hideTempErr(err)
361 c.writeErrMu.Lock()
362 if c.writeErr == nil {
363 c.writeErr = err
364 }
365 c.writeErrMu.Unlock()
366 return err
367}
368
khenaidoo26721882021-08-11 17:42:52 -0400369func (c *Conn) read(n int) ([]byte, error) {
370 p, err := c.br.Peek(n)
371 if err == io.EOF {
372 err = errUnexpectedEOF
373 }
374 c.br.Discard(len(p))
375 return p, err
376}
377
378func (c *Conn) write(frameType int, deadline time.Time, buf0, buf1 []byte) error {
khenaidoo59ce9dd2019-11-11 13:05:32 -0500379 <-c.mu
khenaidoo26721882021-08-11 17:42:52 -0400380 defer func() { c.mu <- struct{}{} }()
khenaidoo59ce9dd2019-11-11 13:05:32 -0500381
382 c.writeErrMu.Lock()
383 err := c.writeErr
384 c.writeErrMu.Unlock()
385 if err != nil {
386 return err
387 }
388
389 c.conn.SetWriteDeadline(deadline)
khenaidoo26721882021-08-11 17:42:52 -0400390 if len(buf1) == 0 {
391 _, err = c.conn.Write(buf0)
392 } else {
393 err = c.writeBufs(buf0, buf1)
khenaidoo59ce9dd2019-11-11 13:05:32 -0500394 }
khenaidoo26721882021-08-11 17:42:52 -0400395 if err != nil {
396 return c.writeFatal(err)
397 }
khenaidoo59ce9dd2019-11-11 13:05:32 -0500398 if frameType == CloseMessage {
399 c.writeFatal(ErrCloseSent)
400 }
401 return nil
402}
403
404// WriteControl writes a control message with the given deadline. The allowed
405// message types are CloseMessage, PingMessage and PongMessage.
406func (c *Conn) WriteControl(messageType int, data []byte, deadline time.Time) error {
407 if !isControl(messageType) {
408 return errBadWriteOpCode
409 }
410 if len(data) > maxControlFramePayloadSize {
411 return errInvalidControlFrame
412 }
413
414 b0 := byte(messageType) | finalBit
415 b1 := byte(len(data))
416 if !c.isServer {
417 b1 |= maskBit
418 }
419
420 buf := make([]byte, 0, maxFrameHeaderSize+maxControlFramePayloadSize)
421 buf = append(buf, b0, b1)
422
423 if c.isServer {
424 buf = append(buf, data...)
425 } else {
426 key := newMaskKey()
427 buf = append(buf, key[:]...)
428 buf = append(buf, data...)
429 maskBytes(key, 0, buf[6:])
430 }
431
khenaidoo26721882021-08-11 17:42:52 -0400432 d := 1000 * time.Hour
khenaidoo59ce9dd2019-11-11 13:05:32 -0500433 if !deadline.IsZero() {
434 d = deadline.Sub(time.Now())
435 if d < 0 {
436 return errWriteTimeout
437 }
438 }
439
440 timer := time.NewTimer(d)
441 select {
442 case <-c.mu:
443 timer.Stop()
444 case <-timer.C:
445 return errWriteTimeout
446 }
khenaidoo26721882021-08-11 17:42:52 -0400447 defer func() { c.mu <- struct{}{} }()
khenaidoo59ce9dd2019-11-11 13:05:32 -0500448
449 c.writeErrMu.Lock()
450 err := c.writeErr
451 c.writeErrMu.Unlock()
452 if err != nil {
453 return err
454 }
455
456 c.conn.SetWriteDeadline(deadline)
457 _, err = c.conn.Write(buf)
458 if err != nil {
459 return c.writeFatal(err)
460 }
461 if messageType == CloseMessage {
462 c.writeFatal(ErrCloseSent)
463 }
464 return err
465}
466
khenaidoo26721882021-08-11 17:42:52 -0400467// beginMessage prepares a connection and message writer for a new message.
468func (c *Conn) beginMessage(mw *messageWriter, messageType int) error {
khenaidoo59ce9dd2019-11-11 13:05:32 -0500469 // Close previous writer if not already closed by the application. It's
470 // probably better to return an error in this situation, but we cannot
471 // change this without breaking existing applications.
472 if c.writer != nil {
473 c.writer.Close()
474 c.writer = nil
475 }
476
477 if !isControl(messageType) && !isData(messageType) {
478 return errBadWriteOpCode
479 }
480
481 c.writeErrMu.Lock()
482 err := c.writeErr
483 c.writeErrMu.Unlock()
khenaidoo26721882021-08-11 17:42:52 -0400484 if err != nil {
485 return err
486 }
487
488 mw.c = c
489 mw.frameType = messageType
490 mw.pos = maxFrameHeaderSize
491
492 if c.writeBuf == nil {
493 wpd, ok := c.writePool.Get().(writePoolData)
494 if ok {
495 c.writeBuf = wpd.buf
496 } else {
497 c.writeBuf = make([]byte, c.writeBufSize)
498 }
499 }
500 return nil
khenaidoo59ce9dd2019-11-11 13:05:32 -0500501}
502
503// NextWriter returns a writer for the next message to send. The writer's Close
504// method flushes the complete message to the network.
505//
506// There can be at most one open writer on a connection. NextWriter closes the
507// previous writer if the application has not already done so.
khenaidoo26721882021-08-11 17:42:52 -0400508//
509// All message types (TextMessage, BinaryMessage, CloseMessage, PingMessage and
510// PongMessage) are supported.
khenaidoo59ce9dd2019-11-11 13:05:32 -0500511func (c *Conn) NextWriter(messageType int) (io.WriteCloser, error) {
khenaidoo26721882021-08-11 17:42:52 -0400512 var mw messageWriter
513 if err := c.beginMessage(&mw, messageType); err != nil {
khenaidoo59ce9dd2019-11-11 13:05:32 -0500514 return nil, err
515 }
khenaidoo26721882021-08-11 17:42:52 -0400516 c.writer = &mw
khenaidoo59ce9dd2019-11-11 13:05:32 -0500517 if c.newCompressionWriter != nil && c.enableWriteCompression && isData(messageType) {
518 w := c.newCompressionWriter(c.writer, c.compressionLevel)
519 mw.compress = true
520 c.writer = w
521 }
522 return c.writer, nil
523}
524
525type messageWriter struct {
526 c *Conn
527 compress bool // whether next call to flushFrame should set RSV1
528 pos int // end of data in writeBuf.
529 frameType int // type of the current frame.
530 err error
531}
532
khenaidoo26721882021-08-11 17:42:52 -0400533func (w *messageWriter) endMessage(err error) error {
khenaidoo59ce9dd2019-11-11 13:05:32 -0500534 if w.err != nil {
khenaidoo26721882021-08-11 17:42:52 -0400535 return err
536 }
537 c := w.c
538 w.err = err
539 c.writer = nil
540 if c.writePool != nil {
541 c.writePool.Put(writePoolData{buf: c.writeBuf})
542 c.writeBuf = nil
khenaidoo59ce9dd2019-11-11 13:05:32 -0500543 }
544 return err
545}
546
547// flushFrame writes buffered data and extra as a frame to the network. The
548// final argument indicates that this is the last frame in the message.
549func (w *messageWriter) flushFrame(final bool, extra []byte) error {
550 c := w.c
551 length := w.pos - maxFrameHeaderSize + len(extra)
552
553 // Check for invalid control frames.
554 if isControl(w.frameType) &&
555 (!final || length > maxControlFramePayloadSize) {
khenaidoo26721882021-08-11 17:42:52 -0400556 return w.endMessage(errInvalidControlFrame)
khenaidoo59ce9dd2019-11-11 13:05:32 -0500557 }
558
559 b0 := byte(w.frameType)
560 if final {
561 b0 |= finalBit
562 }
563 if w.compress {
564 b0 |= rsv1Bit
565 }
566 w.compress = false
567
568 b1 := byte(0)
569 if !c.isServer {
570 b1 |= maskBit
571 }
572
573 // Assume that the frame starts at beginning of c.writeBuf.
574 framePos := 0
575 if c.isServer {
576 // Adjust up if mask not included in the header.
577 framePos = 4
578 }
579
580 switch {
581 case length >= 65536:
582 c.writeBuf[framePos] = b0
583 c.writeBuf[framePos+1] = b1 | 127
584 binary.BigEndian.PutUint64(c.writeBuf[framePos+2:], uint64(length))
585 case length > 125:
586 framePos += 6
587 c.writeBuf[framePos] = b0
588 c.writeBuf[framePos+1] = b1 | 126
589 binary.BigEndian.PutUint16(c.writeBuf[framePos+2:], uint16(length))
590 default:
591 framePos += 8
592 c.writeBuf[framePos] = b0
593 c.writeBuf[framePos+1] = b1 | byte(length)
594 }
595
596 if !c.isServer {
597 key := newMaskKey()
598 copy(c.writeBuf[maxFrameHeaderSize-4:], key[:])
599 maskBytes(key, 0, c.writeBuf[maxFrameHeaderSize:w.pos])
600 if len(extra) > 0 {
khenaidoo26721882021-08-11 17:42:52 -0400601 return w.endMessage(c.writeFatal(errors.New("websocket: internal error, extra used in client mode")))
khenaidoo59ce9dd2019-11-11 13:05:32 -0500602 }
603 }
604
605 // Write the buffers to the connection with best-effort detection of
606 // concurrent writes. See the concurrency section in the package
607 // documentation for more info.
608
609 if c.isWriting {
610 panic("concurrent write to websocket connection")
611 }
612 c.isWriting = true
613
614 err := c.write(w.frameType, c.writeDeadline, c.writeBuf[framePos:w.pos], extra)
615
616 if !c.isWriting {
617 panic("concurrent write to websocket connection")
618 }
619 c.isWriting = false
620
621 if err != nil {
khenaidoo26721882021-08-11 17:42:52 -0400622 return w.endMessage(err)
khenaidoo59ce9dd2019-11-11 13:05:32 -0500623 }
624
625 if final {
khenaidoo26721882021-08-11 17:42:52 -0400626 w.endMessage(errWriteClosed)
khenaidoo59ce9dd2019-11-11 13:05:32 -0500627 return nil
628 }
629
630 // Setup for next frame.
631 w.pos = maxFrameHeaderSize
632 w.frameType = continuationFrame
633 return nil
634}
635
636func (w *messageWriter) ncopy(max int) (int, error) {
637 n := len(w.c.writeBuf) - w.pos
638 if n <= 0 {
639 if err := w.flushFrame(false, nil); err != nil {
640 return 0, err
641 }
642 n = len(w.c.writeBuf) - w.pos
643 }
644 if n > max {
645 n = max
646 }
647 return n, nil
648}
649
650func (w *messageWriter) Write(p []byte) (int, error) {
651 if w.err != nil {
652 return 0, w.err
653 }
654
655 if len(p) > 2*len(w.c.writeBuf) && w.c.isServer {
656 // Don't buffer large messages.
657 err := w.flushFrame(false, p)
658 if err != nil {
659 return 0, err
660 }
661 return len(p), nil
662 }
663
664 nn := len(p)
665 for len(p) > 0 {
666 n, err := w.ncopy(len(p))
667 if err != nil {
668 return 0, err
669 }
670 copy(w.c.writeBuf[w.pos:], p[:n])
671 w.pos += n
672 p = p[n:]
673 }
674 return nn, nil
675}
676
677func (w *messageWriter) WriteString(p string) (int, error) {
678 if w.err != nil {
679 return 0, w.err
680 }
681
682 nn := len(p)
683 for len(p) > 0 {
684 n, err := w.ncopy(len(p))
685 if err != nil {
686 return 0, err
687 }
688 copy(w.c.writeBuf[w.pos:], p[:n])
689 w.pos += n
690 p = p[n:]
691 }
692 return nn, nil
693}
694
695func (w *messageWriter) ReadFrom(r io.Reader) (nn int64, err error) {
696 if w.err != nil {
697 return 0, w.err
698 }
699 for {
700 if w.pos == len(w.c.writeBuf) {
701 err = w.flushFrame(false, nil)
702 if err != nil {
703 break
704 }
705 }
706 var n int
707 n, err = r.Read(w.c.writeBuf[w.pos:])
708 w.pos += n
709 nn += int64(n)
710 if err != nil {
711 if err == io.EOF {
712 err = nil
713 }
714 break
715 }
716 }
717 return nn, err
718}
719
720func (w *messageWriter) Close() error {
721 if w.err != nil {
722 return w.err
723 }
khenaidoo26721882021-08-11 17:42:52 -0400724 return w.flushFrame(true, nil)
khenaidoo59ce9dd2019-11-11 13:05:32 -0500725}
726
727// WritePreparedMessage writes prepared message into connection.
728func (c *Conn) WritePreparedMessage(pm *PreparedMessage) error {
729 frameType, frameData, err := pm.frame(prepareKey{
730 isServer: c.isServer,
731 compress: c.newCompressionWriter != nil && c.enableWriteCompression && isData(pm.messageType),
732 compressionLevel: c.compressionLevel,
733 })
734 if err != nil {
735 return err
736 }
737 if c.isWriting {
738 panic("concurrent write to websocket connection")
739 }
740 c.isWriting = true
741 err = c.write(frameType, c.writeDeadline, frameData, nil)
742 if !c.isWriting {
743 panic("concurrent write to websocket connection")
744 }
745 c.isWriting = false
746 return err
747}
748
749// WriteMessage is a helper method for getting a writer using NextWriter,
750// writing the message and closing the writer.
751func (c *Conn) WriteMessage(messageType int, data []byte) error {
752
753 if c.isServer && (c.newCompressionWriter == nil || !c.enableWriteCompression) {
754 // Fast path with no allocations and single frame.
755
khenaidoo26721882021-08-11 17:42:52 -0400756 var mw messageWriter
757 if err := c.beginMessage(&mw, messageType); err != nil {
khenaidoo59ce9dd2019-11-11 13:05:32 -0500758 return err
759 }
khenaidoo59ce9dd2019-11-11 13:05:32 -0500760 n := copy(c.writeBuf[mw.pos:], data)
761 mw.pos += n
762 data = data[n:]
763 return mw.flushFrame(true, data)
764 }
765
766 w, err := c.NextWriter(messageType)
767 if err != nil {
768 return err
769 }
770 if _, err = w.Write(data); err != nil {
771 return err
772 }
773 return w.Close()
774}
775
776// SetWriteDeadline sets the write deadline on the underlying network
777// connection. After a write has timed out, the websocket state is corrupt and
778// all future writes will return an error. A zero value for t means writes will
779// not time out.
780func (c *Conn) SetWriteDeadline(t time.Time) error {
781 c.writeDeadline = t
782 return nil
783}
784
785// Read methods
786
787func (c *Conn) advanceFrame() (int, error) {
khenaidoo59ce9dd2019-11-11 13:05:32 -0500788 // 1. Skip remainder of previous frame.
789
790 if c.readRemaining > 0 {
791 if _, err := io.CopyN(ioutil.Discard, c.br, c.readRemaining); err != nil {
792 return noFrame, err
793 }
794 }
795
796 // 2. Read and parse first two bytes of frame header.
797
798 p, err := c.read(2)
799 if err != nil {
800 return noFrame, err
801 }
802
803 final := p[0]&finalBit != 0
804 frameType := int(p[0] & 0xf)
805 mask := p[1]&maskBit != 0
khenaidoo26721882021-08-11 17:42:52 -0400806 c.setReadRemaining(int64(p[1] & 0x7f))
khenaidoo59ce9dd2019-11-11 13:05:32 -0500807
808 c.readDecompress = false
809 if c.newDecompressionReader != nil && (p[0]&rsv1Bit) != 0 {
810 c.readDecompress = true
811 p[0] &^= rsv1Bit
812 }
813
814 if rsv := p[0] & (rsv1Bit | rsv2Bit | rsv3Bit); rsv != 0 {
815 return noFrame, c.handleProtocolError("unexpected reserved bits 0x" + strconv.FormatInt(int64(rsv), 16))
816 }
817
818 switch frameType {
819 case CloseMessage, PingMessage, PongMessage:
820 if c.readRemaining > maxControlFramePayloadSize {
821 return noFrame, c.handleProtocolError("control frame length > 125")
822 }
823 if !final {
824 return noFrame, c.handleProtocolError("control frame not final")
825 }
826 case TextMessage, BinaryMessage:
827 if !c.readFinal {
828 return noFrame, c.handleProtocolError("message start before final message frame")
829 }
830 c.readFinal = final
831 case continuationFrame:
832 if c.readFinal {
833 return noFrame, c.handleProtocolError("continuation after final message frame")
834 }
835 c.readFinal = final
836 default:
837 return noFrame, c.handleProtocolError("unknown opcode " + strconv.Itoa(frameType))
838 }
839
khenaidoo26721882021-08-11 17:42:52 -0400840 // 3. Read and parse frame length as per
841 // https://tools.ietf.org/html/rfc6455#section-5.2
842 //
843 // The length of the "Payload data", in bytes: if 0-125, that is the payload
844 // length.
845 // - If 126, the following 2 bytes interpreted as a 16-bit unsigned
846 // integer are the payload length.
847 // - If 127, the following 8 bytes interpreted as
848 // a 64-bit unsigned integer (the most significant bit MUST be 0) are the
849 // payload length. Multibyte length quantities are expressed in network byte
850 // order.
khenaidoo59ce9dd2019-11-11 13:05:32 -0500851
852 switch c.readRemaining {
853 case 126:
854 p, err := c.read(2)
855 if err != nil {
856 return noFrame, err
857 }
khenaidoo26721882021-08-11 17:42:52 -0400858
859 if err := c.setReadRemaining(int64(binary.BigEndian.Uint16(p))); err != nil {
860 return noFrame, err
861 }
khenaidoo59ce9dd2019-11-11 13:05:32 -0500862 case 127:
863 p, err := c.read(8)
864 if err != nil {
865 return noFrame, err
866 }
khenaidoo26721882021-08-11 17:42:52 -0400867
868 if err := c.setReadRemaining(int64(binary.BigEndian.Uint64(p))); err != nil {
869 return noFrame, err
870 }
khenaidoo59ce9dd2019-11-11 13:05:32 -0500871 }
872
873 // 4. Handle frame masking.
874
875 if mask != c.isServer {
876 return noFrame, c.handleProtocolError("incorrect mask flag")
877 }
878
879 if mask {
880 c.readMaskPos = 0
881 p, err := c.read(len(c.readMaskKey))
882 if err != nil {
883 return noFrame, err
884 }
885 copy(c.readMaskKey[:], p)
886 }
887
888 // 5. For text and binary messages, enforce read limit and return.
889
890 if frameType == continuationFrame || frameType == TextMessage || frameType == BinaryMessage {
891
892 c.readLength += c.readRemaining
khenaidoo26721882021-08-11 17:42:52 -0400893 // Don't allow readLength to overflow in the presence of a large readRemaining
894 // counter.
895 if c.readLength < 0 {
896 return noFrame, ErrReadLimit
897 }
898
khenaidoo59ce9dd2019-11-11 13:05:32 -0500899 if c.readLimit > 0 && c.readLength > c.readLimit {
900 c.WriteControl(CloseMessage, FormatCloseMessage(CloseMessageTooBig, ""), time.Now().Add(writeWait))
901 return noFrame, ErrReadLimit
902 }
903
904 return frameType, nil
905 }
906
907 // 6. Read control frame payload.
908
909 var payload []byte
910 if c.readRemaining > 0 {
911 payload, err = c.read(int(c.readRemaining))
khenaidoo26721882021-08-11 17:42:52 -0400912 c.setReadRemaining(0)
khenaidoo59ce9dd2019-11-11 13:05:32 -0500913 if err != nil {
914 return noFrame, err
915 }
916 if c.isServer {
917 maskBytes(c.readMaskKey, 0, payload)
918 }
919 }
920
921 // 7. Process control frame payload.
922
923 switch frameType {
924 case PongMessage:
925 if err := c.handlePong(string(payload)); err != nil {
926 return noFrame, err
927 }
928 case PingMessage:
929 if err := c.handlePing(string(payload)); err != nil {
930 return noFrame, err
931 }
932 case CloseMessage:
933 closeCode := CloseNoStatusReceived
934 closeText := ""
935 if len(payload) >= 2 {
936 closeCode = int(binary.BigEndian.Uint16(payload))
937 if !isValidReceivedCloseCode(closeCode) {
938 return noFrame, c.handleProtocolError("invalid close code")
939 }
940 closeText = string(payload[2:])
941 if !utf8.ValidString(closeText) {
942 return noFrame, c.handleProtocolError("invalid utf8 payload in close frame")
943 }
944 }
945 if err := c.handleClose(closeCode, closeText); err != nil {
946 return noFrame, err
947 }
948 return noFrame, &CloseError{Code: closeCode, Text: closeText}
949 }
950
951 return frameType, nil
952}
953
954func (c *Conn) handleProtocolError(message string) error {
955 c.WriteControl(CloseMessage, FormatCloseMessage(CloseProtocolError, message), time.Now().Add(writeWait))
956 return errors.New("websocket: " + message)
957}
958
959// NextReader returns the next data message received from the peer. The
960// returned messageType is either TextMessage or BinaryMessage.
961//
962// There can be at most one open reader on a connection. NextReader discards
963// the previous message if the application has not already consumed it.
964//
965// Applications must break out of the application's read loop when this method
966// returns a non-nil error value. Errors returned from this method are
967// permanent. Once this method returns a non-nil error, all subsequent calls to
968// this method return the same error.
969func (c *Conn) NextReader() (messageType int, r io.Reader, err error) {
970 // Close previous reader, only relevant for decompression.
971 if c.reader != nil {
972 c.reader.Close()
973 c.reader = nil
974 }
975
976 c.messageReader = nil
977 c.readLength = 0
978
979 for c.readErr == nil {
980 frameType, err := c.advanceFrame()
981 if err != nil {
982 c.readErr = hideTempErr(err)
983 break
984 }
khenaidoo26721882021-08-11 17:42:52 -0400985
khenaidoo59ce9dd2019-11-11 13:05:32 -0500986 if frameType == TextMessage || frameType == BinaryMessage {
987 c.messageReader = &messageReader{c}
988 c.reader = c.messageReader
989 if c.readDecompress {
990 c.reader = c.newDecompressionReader(c.reader)
991 }
992 return frameType, c.reader, nil
993 }
994 }
995
996 // Applications that do handle the error returned from this method spin in
997 // tight loop on connection failure. To help application developers detect
998 // this error, panic on repeated reads to the failed connection.
999 c.readErrCount++
1000 if c.readErrCount >= 1000 {
1001 panic("repeated read on failed websocket connection")
1002 }
1003
1004 return noFrame, nil, c.readErr
1005}
1006
1007type messageReader struct{ c *Conn }
1008
1009func (r *messageReader) Read(b []byte) (int, error) {
1010 c := r.c
1011 if c.messageReader != r {
1012 return 0, io.EOF
1013 }
1014
1015 for c.readErr == nil {
1016
1017 if c.readRemaining > 0 {
1018 if int64(len(b)) > c.readRemaining {
1019 b = b[:c.readRemaining]
1020 }
1021 n, err := c.br.Read(b)
1022 c.readErr = hideTempErr(err)
1023 if c.isServer {
1024 c.readMaskPos = maskBytes(c.readMaskKey, c.readMaskPos, b[:n])
1025 }
khenaidoo26721882021-08-11 17:42:52 -04001026 rem := c.readRemaining
1027 rem -= int64(n)
1028 c.setReadRemaining(rem)
khenaidoo59ce9dd2019-11-11 13:05:32 -05001029 if c.readRemaining > 0 && c.readErr == io.EOF {
1030 c.readErr = errUnexpectedEOF
1031 }
1032 return n, c.readErr
1033 }
1034
1035 if c.readFinal {
1036 c.messageReader = nil
1037 return 0, io.EOF
1038 }
1039
1040 frameType, err := c.advanceFrame()
1041 switch {
1042 case err != nil:
1043 c.readErr = hideTempErr(err)
1044 case frameType == TextMessage || frameType == BinaryMessage:
1045 c.readErr = errors.New("websocket: internal error, unexpected text or binary in Reader")
1046 }
1047 }
1048
1049 err := c.readErr
1050 if err == io.EOF && c.messageReader == r {
1051 err = errUnexpectedEOF
1052 }
1053 return 0, err
1054}
1055
1056func (r *messageReader) Close() error {
1057 return nil
1058}
1059
1060// ReadMessage is a helper method for getting a reader using NextReader and
1061// reading from that reader to a buffer.
1062func (c *Conn) ReadMessage() (messageType int, p []byte, err error) {
1063 var r io.Reader
1064 messageType, r, err = c.NextReader()
1065 if err != nil {
1066 return messageType, nil, err
1067 }
1068 p, err = ioutil.ReadAll(r)
1069 return messageType, p, err
1070}
1071
1072// SetReadDeadline sets the read deadline on the underlying network connection.
1073// After a read has timed out, the websocket connection state is corrupt and
1074// all future reads will return an error. A zero value for t means reads will
1075// not time out.
1076func (c *Conn) SetReadDeadline(t time.Time) error {
1077 return c.conn.SetReadDeadline(t)
1078}
1079
khenaidoo26721882021-08-11 17:42:52 -04001080// SetReadLimit sets the maximum size in bytes for a message read from the peer. If a
1081// message exceeds the limit, the connection sends a close message to the peer
khenaidoo59ce9dd2019-11-11 13:05:32 -05001082// and returns ErrReadLimit to the application.
1083func (c *Conn) SetReadLimit(limit int64) {
1084 c.readLimit = limit
1085}
1086
1087// CloseHandler returns the current close handler
1088func (c *Conn) CloseHandler() func(code int, text string) error {
1089 return c.handleClose
1090}
1091
1092// SetCloseHandler sets the handler for close messages received from the peer.
1093// The code argument to h is the received close code or CloseNoStatusReceived
khenaidoo26721882021-08-11 17:42:52 -04001094// if the close message is empty. The default close handler sends a close
1095// message back to the peer.
khenaidoo59ce9dd2019-11-11 13:05:32 -05001096//
khenaidoo26721882021-08-11 17:42:52 -04001097// The handler function is called from the NextReader, ReadMessage and message
1098// reader Read methods. The application must read the connection to process
1099// close messages as described in the section on Control Messages above.
khenaidoo59ce9dd2019-11-11 13:05:32 -05001100//
khenaidoo26721882021-08-11 17:42:52 -04001101// The connection read methods return a CloseError when a close message is
khenaidoo59ce9dd2019-11-11 13:05:32 -05001102// received. Most applications should handle close messages as part of their
1103// normal error handling. Applications should only set a close handler when the
khenaidoo26721882021-08-11 17:42:52 -04001104// application must perform some action before sending a close message back to
khenaidoo59ce9dd2019-11-11 13:05:32 -05001105// the peer.
1106func (c *Conn) SetCloseHandler(h func(code int, text string) error) {
1107 if h == nil {
1108 h = func(code int, text string) error {
khenaidoo26721882021-08-11 17:42:52 -04001109 message := FormatCloseMessage(code, "")
khenaidoo59ce9dd2019-11-11 13:05:32 -05001110 c.WriteControl(CloseMessage, message, time.Now().Add(writeWait))
1111 return nil
1112 }
1113 }
1114 c.handleClose = h
1115}
1116
1117// PingHandler returns the current ping handler
1118func (c *Conn) PingHandler() func(appData string) error {
1119 return c.handlePing
1120}
1121
1122// SetPingHandler sets the handler for ping messages received from the peer.
khenaidoo26721882021-08-11 17:42:52 -04001123// The appData argument to h is the PING message application data. The default
khenaidoo59ce9dd2019-11-11 13:05:32 -05001124// ping handler sends a pong to the peer.
1125//
khenaidoo26721882021-08-11 17:42:52 -04001126// The handler function is called from the NextReader, ReadMessage and message
1127// reader Read methods. The application must read the connection to process
1128// ping messages as described in the section on Control Messages above.
khenaidoo59ce9dd2019-11-11 13:05:32 -05001129func (c *Conn) SetPingHandler(h func(appData string) error) {
1130 if h == nil {
1131 h = func(message string) error {
1132 err := c.WriteControl(PongMessage, []byte(message), time.Now().Add(writeWait))
1133 if err == ErrCloseSent {
1134 return nil
1135 } else if e, ok := err.(net.Error); ok && e.Temporary() {
1136 return nil
1137 }
1138 return err
1139 }
1140 }
1141 c.handlePing = h
1142}
1143
1144// PongHandler returns the current pong handler
1145func (c *Conn) PongHandler() func(appData string) error {
1146 return c.handlePong
1147}
1148
1149// SetPongHandler sets the handler for pong messages received from the peer.
khenaidoo26721882021-08-11 17:42:52 -04001150// The appData argument to h is the PONG message application data. The default
khenaidoo59ce9dd2019-11-11 13:05:32 -05001151// pong handler does nothing.
1152//
khenaidoo26721882021-08-11 17:42:52 -04001153// The handler function is called from the NextReader, ReadMessage and message
1154// reader Read methods. The application must read the connection to process
1155// pong messages as described in the section on Control Messages above.
khenaidoo59ce9dd2019-11-11 13:05:32 -05001156func (c *Conn) SetPongHandler(h func(appData string) error) {
1157 if h == nil {
1158 h = func(string) error { return nil }
1159 }
1160 c.handlePong = h
1161}
1162
1163// UnderlyingConn returns the internal net.Conn. This can be used to further
1164// modifications to connection specific flags.
1165func (c *Conn) UnderlyingConn() net.Conn {
1166 return c.conn
1167}
1168
1169// EnableWriteCompression enables and disables write compression of
1170// subsequent text and binary messages. This function is a noop if
1171// compression was not negotiated with the peer.
1172func (c *Conn) EnableWriteCompression(enable bool) {
1173 c.enableWriteCompression = enable
1174}
1175
1176// SetCompressionLevel sets the flate compression level for subsequent text and
1177// binary messages. This function is a noop if compression was not negotiated
1178// with the peer. See the compress/flate package for a description of
1179// compression levels.
1180func (c *Conn) SetCompressionLevel(level int) error {
1181 if !isValidCompressionLevel(level) {
1182 return errors.New("websocket: invalid compression level")
1183 }
1184 c.compressionLevel = level
1185 return nil
1186}
1187
1188// FormatCloseMessage formats closeCode and text as a WebSocket close message.
khenaidoo26721882021-08-11 17:42:52 -04001189// An empty message is returned for code CloseNoStatusReceived.
khenaidoo59ce9dd2019-11-11 13:05:32 -05001190func FormatCloseMessage(closeCode int, text string) []byte {
khenaidoo26721882021-08-11 17:42:52 -04001191 if closeCode == CloseNoStatusReceived {
1192 // Return empty message because it's illegal to send
1193 // CloseNoStatusReceived. Return non-nil value in case application
1194 // checks for nil.
1195 return []byte{}
1196 }
khenaidoo59ce9dd2019-11-11 13:05:32 -05001197 buf := make([]byte, 2+len(text))
1198 binary.BigEndian.PutUint16(buf, uint16(closeCode))
1199 copy(buf[2:], text)
1200 return buf
1201}