Source file src/github.com/gorilla/websocket/conn.go

Documentation: github.com/gorilla/websocket

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

View as plain text