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 	"sync"

    17 	"time"

    18 	"unicode/utf8"

    19 )

    20 

    21 const (

    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.

    44 const (

    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.

    62 const (

    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 

    79 	// PongMessage denotes a pong control message. The optional message payload

    80 	// 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.

    86 var 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.

    90 var ErrReadLimit = errors.New("websocket: read limit exceeded")

    91 

    92 // netError satisfies the net Error interface.

    93 type netError struct {

    94 	msg       string

    95 	temporary bool

    96 	timeout   bool

    97 }

    98 

    99 func (e *netError) Error() string   { return e.msg }

   100 func (e *netError) Temporary() bool { return e.temporary }

   101 func (e *netError) Timeout() bool   { return e.timeout }

   102 

   103 // CloseError represents a close message.

   104 type CloseError struct {

   105 	// 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 

   112 func (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.

   150 func 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.

   163 func 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 

   175 var (

   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 

   183 func newMaskKey() [4]byte {

   184 	n := rand.Uint32()

   185 	return [4]byte{byte(n), byte(n >> 8), byte(n >> 16), byte(n >> 24)}

   186 }

   187 

   188 func 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 

   195 func isControl(frameType int) bool {

   196 	return frameType == CloseMessage || frameType == PingMessage || frameType == PongMessage

   197 }

   198 

   199 func isData(frameType int) bool {

   200 	return frameType == TextMessage || frameType == BinaryMessage

   201 }

   202 

   203 var 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 

   222 func isValidReceivedCloseCode(code int) bool {

   223 	return validReceivedCloseCodes[code] || (code >= 3000 && code <= 4999)

   224 }

   225 

   226 // 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.

   228 type 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.

   238 type writePoolData struct{ buf []byte }

   239 

   240 // The Conn type represents a WebSocket connection.

   241 type Conn struct {

   242 	conn        net.Conn

   243 	isServer    bool

   244 	subprotocol string

   245 

   246 	// Write fields

   247 	mu            chan bool // used as mutex to protect write to conn

   248 	writeBuf      []byte    // frame is constructed in this buffer.

   249 	writePool     BufferPool

   250 	writeBufSize  int

   251 	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

   263 	reader        io.ReadCloser // the current reader returned to the application

   264 	readErr       error

   265 	br            *bufio.Reader

   266 	readRemaining int64 // bytes remaining in current frame.

   267 	readFinal     bool  // true the current message has more frames.

   268 	readLength    int64 // Message size.

   269 	readLimit     int64 // Maximum message size.

   270 	readMaskPos   int

   271 	readMaskKey   [4]byte

   272 	handlePong    func(string) error

   273 	handlePing    func(string) error

   274 	handleClose   func(int, string) error

   275 	readErrCount  int

   276 	messageReader *messageReader // the current low-level reader

   277 

   278 	readDecompress         bool // whether last read frame had RSV1 set

   279 	newDecompressionReader func(io.Reader) io.ReadCloser

   280 }

   281 

   282 func newConn(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int, writeBufferPool BufferPool, br *bufio.Reader, writeBuf []byte) *Conn {

   283 

   284 	if br == nil {

   285 		if readBufferSize == 0 {

   286 			readBufferSize = defaultReadBufferSize

   287 		} else if readBufferSize < maxControlFramePayloadSize {

   288 			// must be large enough for control frame

   289 			readBufferSize = maxControlFramePayloadSize

   290 		}

   291 		br = bufio.NewReaderSize(conn, readBufferSize)

   292 	}

   293 

   294 	if writeBufferSize <= 0 {

   295 		writeBufferSize = defaultWriteBufferSize

   296 	}

   297 	writeBufferSize += maxFrameHeaderSize

   298 

   299 	if writeBuf == nil && writeBufferPool == nil {

   300 		writeBuf = make([]byte, writeBufferSize)

   301 	}

   302 

   303 	mu := make(chan bool, 1)

   304 	mu <- true

   305 	c := &Conn{

   306 		isServer:               isServer,

   307 		br:                     br,

   308 		conn:                   conn,

   309 		mu:                     mu,

   310 		readFinal:              true,

   311 		writeBuf:               writeBuf,

   312 		writePool:              writeBufferPool,

   313 		writeBufSize:           writeBufferSize,

   314 		enableWriteCompression: true,

   315 		compressionLevel:       defaultCompressionLevel,

   316 	}

   317 	c.SetCloseHandler(nil)

   318 	c.SetPingHandler(nil)

   319 	c.SetPongHandler(nil)

   320 	return c

   321 }

   322 

   323 // Subprotocol returns the negotiated protocol for the connection.

   324 func (c *Conn) Subprotocol() string {

   325 	return c.subprotocol

   326 }

   327 

   328 // Close closes the underlying network connection without sending or waiting

   329 // for a close message.

   330 func (c *Conn) Close() error {

   331 	return c.conn.Close()

   332 }

   333 

   334 // LocalAddr returns the local network address.

   335 func (c *Conn) LocalAddr() net.Addr {

   336 	return c.conn.LocalAddr()

   337 }

   338 

   339 // RemoteAddr returns the remote network address.

   340 func (c *Conn) RemoteAddr() net.Addr {

   341 	return c.conn.RemoteAddr()

   342 }

   343 

   344 // Write methods

   345 

   346 func (c *Conn) writeFatal(err error) error {

   347 	err = hideTempErr(err)

   348 	c.writeErrMu.Lock()

   349 	if c.writeErr == nil {

   350 		c.writeErr = err

   351 	}

   352 	c.writeErrMu.Unlock()

   353 	return err

   354 }

   355 

   356 func (c *Conn) read(n int) ([]byte, error) {

   357 	p, err := c.br.Peek(n)

   358 	if err == io.EOF {

   359 		err = errUnexpectedEOF

   360 	}

   361 	c.br.Discard(len(p))

   362 	return p, err

   363 }

   364 

   365 func (c *Conn) write(frameType int, deadline time.Time, buf0, buf1 []byte) error {

   366 	<-c.mu

   367 	defer func() { c.mu <- true }()

   368 

   369 	c.writeErrMu.Lock()

   370 	err := c.writeErr

   371 	c.writeErrMu.Unlock()

   372 	if err != nil {

   373 		return err

   374 	}

   375 

   376 	c.conn.SetWriteDeadline(deadline)

   377 	if len(buf1) == 0 {

   378 		_, err = c.conn.Write(buf0)

   379 	} else {

   380 		err = c.writeBufs(buf0, buf1)

   381 	}

   382 	if err != nil {

   383 		return c.writeFatal(err)

   384 	}

   385 	if frameType == CloseMessage {

   386 		c.writeFatal(ErrCloseSent)

   387 	}

   388 	return nil

   389 }

   390 

   391 // WriteControl writes a control message with the given deadline. The allowed

   392 // message types are CloseMessage, PingMessage and PongMessage.

   393 func (c *Conn) WriteControl(messageType int, data []byte, deadline time.Time) error {

   394 	if !isControl(messageType) {

   395 		return errBadWriteOpCode

   396 	}

   397 	if len(data) > maxControlFramePayloadSize {

   398 		return errInvalidControlFrame

   399 	}

   400 

   401 	b0 := byte(messageType) | finalBit

   402 	b1 := byte(len(data))

   403 	if !c.isServer {

   404 		b1 |= maskBit

   405 	}

   406 

   407 	buf := make([]byte, 0, maxFrameHeaderSize+maxControlFramePayloadSize)

   408 	buf = append(buf, b0, b1)

   409 

   410 	if c.isServer {

   411 		buf = append(buf, data...)

   412 	} else {

   413 		key := newMaskKey()

   414 		buf = append(buf, key[:]...)

   415 		buf = append(buf, data...)

   416 		maskBytes(key, 0, buf[6:])

   417 	}

   418 

   419 	d := time.Hour * 1000

   420 	if !deadline.IsZero() {

   421 		d = deadline.Sub(time.Now())

   422 		if d < 0 {

   423 			return errWriteTimeout

   424 		}

   425 	}

   426 

   427 	timer := time.NewTimer(d)

   428 	select {

   429 	case <-c.mu:

   430 		timer.Stop()

   431 	case <-timer.C:

   432 		return errWriteTimeout

   433 	}

   434 	defer func() { c.mu <- true }()

   435 

   436 	c.writeErrMu.Lock()

   437 	err := c.writeErr

   438 	c.writeErrMu.Unlock()

   439 	if err != nil {

   440 		return err

   441 	}

   442 

   443 	c.conn.SetWriteDeadline(deadline)

   444 	_, err = c.conn.Write(buf)

   445 	if err != nil {

   446 		return c.writeFatal(err)

   447 	}

   448 	if messageType == CloseMessage {

   449 		c.writeFatal(ErrCloseSent)

   450 	}

   451 	return err

   452 }

   453 

   454 func (c *Conn) prepWrite(messageType int) error {

   455 	// Close previous writer if not already closed by the application. It's

   456 	// probably better to return an error in this situation, but we cannot

   457 	// change this without breaking existing applications.

   458 	if c.writer != nil {

   459 		c.writer.Close()

   460 		c.writer = nil

   461 	}

   462 

   463 	if !isControl(messageType) && !isData(messageType) {

   464 		return errBadWriteOpCode

   465 	}

   466 

   467 	c.writeErrMu.Lock()

   468 	err := c.writeErr

   469 	c.writeErrMu.Unlock()

   470 	if err != nil {

   471 		return err

   472 	}

   473 

   474 	if c.writeBuf == nil {

   475 		wpd, ok := c.writePool.Get().(writePoolData)

   476 		if ok {

   477 			c.writeBuf = wpd.buf

   478 		} else {

   479 			c.writeBuf = make([]byte, c.writeBufSize)

   480 		}

   481 	}

   482 	return nil

   483 }

   484 

   485 // NextWriter returns a writer for the next message to send. The writer's Close

   486 // method flushes the complete message to the network.

   487 //

   488 // There can be at most one open writer on a connection. NextWriter closes the

   489 // previous writer if the application has not already done so.

   490 //

   491 // All message types (TextMessage, BinaryMessage, CloseMessage, PingMessage and

   492 // PongMessage) are supported.

   493 func (c *Conn) NextWriter(messageType int) (io.WriteCloser, error) {

   494 	if err := c.prepWrite(messageType); err != nil {

   495 		return nil, err

   496 	}

   497 

   498 	mw := &messageWriter{

   499 		c:         c,

   500 		frameType: messageType,

   501 		pos:       maxFrameHeaderSize,

   502 	}

   503 	c.writer = mw

   504 	if c.newCompressionWriter != nil && c.enableWriteCompression && isData(messageType) {

   505 		w := c.newCompressionWriter(c.writer, c.compressionLevel)

   506 		mw.compress = true

   507 		c.writer = w

   508 	}

   509 	return c.writer, nil

   510 }

   511 

   512 type messageWriter struct {

   513 	c         *Conn

   514 	compress  bool // whether next call to flushFrame should set RSV1

   515 	pos       int  // end of data in writeBuf.

   516 	frameType int  // type of the current frame.

   517 	err       error

   518 }

   519 

   520 func (w *messageWriter) fatal(err error) error {

   521 	if w.err != nil {

   522 		w.err = err

   523 		w.c.writer = nil

   524 	}

   525 	return err

   526 }

   527 

   528 // flushFrame writes buffered data and extra as a frame to the network. The

   529 // final argument indicates that this is the last frame in the message.

   530 func (w *messageWriter) flushFrame(final bool, extra []byte) error {

   531 	c := w.c

   532 	length := w.pos - maxFrameHeaderSize + len(extra)

   533 

   534 	// Check for invalid control frames.

   535 	if isControl(w.frameType) &&

   536 		(!final || length > maxControlFramePayloadSize) {

   537 		return w.fatal(errInvalidControlFrame)

   538 	}

   539 

   540 	b0 := byte(w.frameType)

   541 	if final {

   542 		b0 |= finalBit

   543 	}

   544 	if w.compress {

   545 		b0 |= rsv1Bit

   546 	}

   547 	w.compress = false

   548 

   549 	b1 := byte(0)

   550 	if !c.isServer {

   551 		b1 |= maskBit

   552 	}

   553 

   554 	// Assume that the frame starts at beginning of c.writeBuf.

   555 	framePos := 0

   556 	if c.isServer {

   557 		// Adjust up if mask not included in the header.

   558 		framePos = 4

   559 	}

   560 

   561 	switch {

   562 	case length >= 65536:

   563 		c.writeBuf[framePos] = b0

   564 		c.writeBuf[framePos+1] = b1 | 127

   565 		binary.BigEndian.PutUint64(c.writeBuf[framePos+2:], uint64(length))

   566 	case length > 125:

   567 		framePos += 6

   568 		c.writeBuf[framePos] = b0

   569 		c.writeBuf[framePos+1] = b1 | 126

   570 		binary.BigEndian.PutUint16(c.writeBuf[framePos+2:], uint16(length))

   571 	default:

   572 		framePos += 8

   573 		c.writeBuf[framePos] = b0

   574 		c.writeBuf[framePos+1] = b1 | byte(length)

   575 	}

   576 

   577 	if !c.isServer {

   578 		key := newMaskKey()

   579 		copy(c.writeBuf[maxFrameHeaderSize-4:], key[:])

   580 		maskBytes(key, 0, c.writeBuf[maxFrameHeaderSize:w.pos])

   581 		if len(extra) > 0 {

   582 			return c.writeFatal(errors.New("websocket: internal error, extra used in client mode"))

   583 		}

   584 	}

   585 

   586 	// Write the buffers to the connection with best-effort detection of

   587 	// concurrent writes. See the concurrency section in the package

   588 	// documentation for more info.

   589 

   590 	if c.isWriting {

   591 		panic("concurrent write to websocket connection")

   592 	}

   593 	c.isWriting = true

   594 

   595 	err := c.write(w.frameType, c.writeDeadline, c.writeBuf[framePos:w.pos], extra)

   596 

   597 	if !c.isWriting {

   598 		panic("concurrent write to websocket connection")

   599 	}

   600 	c.isWriting = false

   601 

   602 	if err != nil {

   603 		return w.fatal(err)

   604 	}

   605 

   606 	if final {

   607 		c.writer = nil

   608 		if c.writePool != nil {

   609 			c.writePool.Put(writePoolData{buf: c.writeBuf})

   610 			c.writeBuf = nil

   611 		}

   612 		return nil

   613 	}

   614 

   615 	// Setup for next frame.

   616 	w.pos = maxFrameHeaderSize

   617 	w.frameType = continuationFrame

   618 	return nil

   619 }

   620 

   621 func (w *messageWriter) ncopy(max int) (int, error) {

   622 	n := len(w.c.writeBuf) - w.pos

   623 	if n <= 0 {

   624 		if err := w.flushFrame(false, nil); err != nil {

   625 			return 0, err

   626 		}

   627 		n = len(w.c.writeBuf) - w.pos

   628 	}

   629 	if n > max {

   630 		n = max

   631 	}

   632 	return n, nil

   633 }

   634 

   635 func (w *messageWriter) Write(p []byte) (int, error) {

   636 	if w.err != nil {

   637 		return 0, w.err

   638 	}

   639 

   640 	if len(p) > 2*len(w.c.writeBuf) && w.c.isServer {

   641 		// Don't buffer large messages.

   642 		err := w.flushFrame(false, p)

   643 		if err != nil {

   644 			return 0, err

   645 		}

   646 		return len(p), nil

   647 	}

   648 

   649 	nn := len(p)

   650 	for len(p) > 0 {

   651 		n, err := w.ncopy(len(p))

   652 		if err != nil {

   653 			return 0, err

   654 		}

   655 		copy(w.c.writeBuf[w.pos:], p[:n])

   656 		w.pos += n

   657 		p = p[n:]

   658 	}

   659 	return nn, nil

   660 }

   661 

   662 func (w *messageWriter) WriteString(p string) (int, error) {

   663 	if w.err != nil {

   664 		return 0, w.err

   665 	}

   666 

   667 	nn := len(p)

   668 	for len(p) > 0 {

   669 		n, err := w.ncopy(len(p))

   670 		if err != nil {

   671 			return 0, err

   672 		}

   673 		copy(w.c.writeBuf[w.pos:], p[:n])

   674 		w.pos += n

   675 		p = p[n:]

   676 	}

   677 	return nn, nil

   678 }

   679 

   680 func (w *messageWriter) ReadFrom(r io.Reader) (nn int64, err error) {

   681 	if w.err != nil {

   682 		return 0, w.err

   683 	}

   684 	for {

   685 		if w.pos == len(w.c.writeBuf) {

   686 			err = w.flushFrame(false, nil)

   687 			if err != nil {

   688 				break

   689 			}

   690 		}

   691 		var n int

   692 		n, err = r.Read(w.c.writeBuf[w.pos:])

   693 		w.pos += n

   694 		nn += int64(n)

   695 		if err != nil {

   696 			if err == io.EOF {

   697 				err = nil

   698 			}

   699 			break

   700 		}

   701 	}

   702 	return nn, err

   703 }

   704 

   705 func (w *messageWriter) Close() error {

   706 	if w.err != nil {

   707 		return w.err

   708 	}

   709 	if err := w.flushFrame(true, nil); err != nil {

   710 		return err

   711 	}

   712 	w.err = errWriteClosed

   713 	return nil

   714 }

   715 

   716 // WritePreparedMessage writes prepared message into connection.

   717 func (c *Conn) WritePreparedMessage(pm *PreparedMessage) error {

   718 	frameType, frameData, err := pm.frame(prepareKey{

   719 		isServer:         c.isServer,

   720 		compress:         c.newCompressionWriter != nil && c.enableWriteCompression && isData(pm.messageType),

   721 		compressionLevel: c.compressionLevel,

   722 	})

   723 	if err != nil {

   724 		return err

   725 	}

   726 	if c.isWriting {

   727 		panic("concurrent write to websocket connection")

   728 	}

   729 	c.isWriting = true

   730 	err = c.write(frameType, c.writeDeadline, frameData, nil)

   731 	if !c.isWriting {

   732 		panic("concurrent write to websocket connection")

   733 	}

   734 	c.isWriting = false

   735 	return err

   736 }

   737 

   738 // WriteMessage is a helper method for getting a writer using NextWriter,

   739 // writing the message and closing the writer.

   740 func (c *Conn) WriteMessage(messageType int, data []byte) error {

   741 

   742 	if c.isServer && (c.newCompressionWriter == nil || !c.enableWriteCompression) {

   743 		// Fast path with no allocations and single frame.

   744 

   745 		if err := c.prepWrite(messageType); err != nil {

   746 			return err

   747 		}

   748 		mw := messageWriter{c: c, frameType: messageType, pos: maxFrameHeaderSize}

   749 		n := copy(c.writeBuf[mw.pos:], data)

   750 		mw.pos += n

   751 		data = data[n:]

   752 		return mw.flushFrame(true, data)

   753 	}

   754 

   755 	w, err := c.NextWriter(messageType)

   756 	if err != nil {

   757 		return err

   758 	}

   759 	if _, err = w.Write(data); err != nil {

   760 		return err

   761 	}

   762 	return w.Close()

   763 }

   764 

   765 // SetWriteDeadline sets the write deadline on the underlying network

   766 // connection. After a write has timed out, the websocket state is corrupt and

   767 // all future writes will return an error. A zero value for t means writes will

   768 // not time out.

   769 func (c *Conn) SetWriteDeadline(t time.Time) error {

   770 	c.writeDeadline = t

   771 	return nil

   772 }

   773 

   774 // Read methods

   775 

   776 func (c *Conn) advanceFrame() (int, error) {

   777 	// 1. Skip remainder of previous frame.

   778 

   779 	if c.readRemaining > 0 {

   780 		if _, err := io.CopyN(ioutil.Discard, c.br, c.readRemaining); err != nil {

   781 			return noFrame, err

   782 		}

   783 	}

   784 

   785 	// 2. Read and parse first two bytes of frame header.

   786 

   787 	p, err := c.read(2)

   788 	if err != nil {

   789 		return noFrame, err

   790 	}

   791 

   792 	final := p[0]&finalBit != 0

   793 	frameType := int(p[0] & 0xf)

   794 	mask := p[1]&maskBit != 0

   795 	c.readRemaining = int64(p[1] & 0x7f)

   796 

   797 	c.readDecompress = false

   798 	if c.newDecompressionReader != nil && (p[0]&rsv1Bit) != 0 {

   799 		c.readDecompress = true

   800 		p[0] &^= rsv1Bit

   801 	}

   802 

   803 	if rsv := p[0] & (rsv1Bit | rsv2Bit | rsv3Bit); rsv != 0 {

   804 		return noFrame, c.handleProtocolError("unexpected reserved bits 0x" + strconv.FormatInt(int64(rsv), 16))

   805 	}

   806 

   807 	switch frameType {

   808 	case CloseMessage, PingMessage, PongMessage:

   809 		if c.readRemaining > maxControlFramePayloadSize {

   810 			return noFrame, c.handleProtocolError("control frame length > 125")

   811 		}

   812 		if !final {

   813 			return noFrame, c.handleProtocolError("control frame not final")

   814 		}

   815 	case TextMessage, BinaryMessage:

   816 		if !c.readFinal {

   817 			return noFrame, c.handleProtocolError("message start before final message frame")

   818 		}

   819 		c.readFinal = final

   820 	case continuationFrame:

   821 		if c.readFinal {

   822 			return noFrame, c.handleProtocolError("continuation after final message frame")

   823 		}

   824 		c.readFinal = final

   825 	default:

   826 		return noFrame, c.handleProtocolError("unknown opcode " + strconv.Itoa(frameType))

   827 	}

   828 

   829 	// 3. Read and parse frame length.

   830 

   831 	switch c.readRemaining {

   832 	case 126:

   833 		p, err := c.read(2)

   834 		if err != nil {

   835 			return noFrame, err

   836 		}

   837 		c.readRemaining = int64(binary.BigEndian.Uint16(p))

   838 	case 127:

   839 		p, err := c.read(8)

   840 		if err != nil {

   841 			return noFrame, err

   842 		}

   843 		c.readRemaining = int64(binary.BigEndian.Uint64(p))

   844 	}

   845 

   846 	// 4. Handle frame masking.

   847 

   848 	if mask != c.isServer {

   849 		return noFrame, c.handleProtocolError("incorrect mask flag")

   850 	}

   851 

   852 	if mask {

   853 		c.readMaskPos = 0

   854 		p, err := c.read(len(c.readMaskKey))

   855 		if err != nil {

   856 			return noFrame, err

   857 		}

   858 		copy(c.readMaskKey[:], p)

   859 	}

   860 

   861 	// 5. For text and binary messages, enforce read limit and return.

   862 

   863 	if frameType == continuationFrame || frameType == TextMessage || frameType == BinaryMessage {

   864 

   865 		c.readLength += c.readRemaining

   866 		if c.readLimit > 0 && c.readLength > c.readLimit {

   867 			c.WriteControl(CloseMessage, FormatCloseMessage(CloseMessageTooBig, ""), time.Now().Add(writeWait))

   868 			return noFrame, ErrReadLimit

   869 		}

   870 

   871 		return frameType, nil

   872 	}

   873 

   874 	// 6. Read control frame payload.

   875 

   876 	var payload []byte

   877 	if c.readRemaining > 0 {

   878 		payload, err = c.read(int(c.readRemaining))

   879 		c.readRemaining = 0

   880 		if err != nil {

   881 			return noFrame, err

   882 		}

   883 		if c.isServer {

   884 			maskBytes(c.readMaskKey, 0, payload)

   885 		}

   886 	}

   887 

   888 	// 7. Process control frame payload.

   889 

   890 	switch frameType {

   891 	case PongMessage:

   892 		if err := c.handlePong(string(payload)); err != nil {

   893 			return noFrame, err

   894 		}

   895 	case PingMessage:

   896 		if err := c.handlePing(string(payload)); err != nil {

   897 			return noFrame, err

   898 		}

   899 	case CloseMessage:

   900 		closeCode := CloseNoStatusReceived

   901 		closeText := ""

   902 		if len(payload) >= 2 {

   903 			closeCode = int(binary.BigEndian.Uint16(payload))

   904 			if !isValidReceivedCloseCode(closeCode) {

   905 				return noFrame, c.handleProtocolError("invalid close code")

   906 			}

   907 			closeText = string(payload[2:])

   908 			if !utf8.ValidString(closeText) {

   909 				return noFrame, c.handleProtocolError("invalid utf8 payload in close frame")

   910 			}

   911 		}

   912 		if err := c.handleClose(closeCode, closeText); err != nil {

   913 			return noFrame, err

   914 		}

   915 		return noFrame, &CloseError{Code: closeCode, Text: closeText}

   916 	}

   917 

   918 	return frameType, nil

   919 }

   920 

   921 func (c *Conn) handleProtocolError(message string) error {

   922 	c.WriteControl(CloseMessage, FormatCloseMessage(CloseProtocolError, message), time.Now().Add(writeWait))

   923 	return errors.New("websocket: " + message)

   924 }

   925 

   926 // NextReader returns the next data message received from the peer. The

   927 // returned messageType is either TextMessage or BinaryMessage.

   928 //

   929 // There can be at most one open reader on a connection. NextReader discards

   930 // the previous message if the application has not already consumed it.

   931 //

   932 // Applications must break out of the application's read loop when this method

   933 // returns a non-nil error value. Errors returned from this method are

   934 // permanent. Once this method returns a non-nil error, all subsequent calls to

   935 // this method return the same error.

   936 func (c *Conn) NextReader() (messageType int, r io.Reader, err error) {

   937 	// Close previous reader, only relevant for decompression.

   938 	if c.reader != nil {

   939 		c.reader.Close()

   940 		c.reader = nil

   941 	}

   942 

   943 	c.messageReader = nil

   944 	c.readLength = 0

   945 

   946 	for c.readErr == nil {

   947 		frameType, err := c.advanceFrame()

   948 		if err != nil {

   949 			c.readErr = hideTempErr(err)

   950 			break

   951 		}

   952 		if frameType == TextMessage || frameType == BinaryMessage {

   953 			c.messageReader = &messageReader{c}

   954 			c.reader = c.messageReader

   955 			if c.readDecompress {

   956 				c.reader = c.newDecompressionReader(c.reader)

   957 			}

   958 			return frameType, c.reader, nil

   959 		}

   960 	}

   961 

   962 	// Applications that do handle the error returned from this method spin in

   963 	// tight loop on connection failure. To help application developers detect

   964 	// this error, panic on repeated reads to the failed connection.

   965 	c.readErrCount++

   966 	if c.readErrCount >= 1000 {

   967 		panic("repeated read on failed websocket connection")

   968 	}

   969 

   970 	return noFrame, nil, c.readErr

   971 }

   972 

   973 type messageReader struct{ c *Conn }

   974 

   975 func (r *messageReader) Read(b []byte) (int, error) {

   976 	c := r.c

   977 	if c.messageReader != r {

   978 		return 0, io.EOF

   979 	}

   980 

   981 	for c.readErr == nil {

   982 

   983 		if c.readRemaining > 0 {

   984 			if int64(len(b)) > c.readRemaining {

   985 				b = b[:c.readRemaining]

   986 			}

   987 			n, err := c.br.Read(b)

   988 			c.readErr = hideTempErr(err)

   989 			if c.isServer {

   990 				c.readMaskPos = maskBytes(c.readMaskKey, c.readMaskPos, b[:n])

   991 			}

   992 			c.readRemaining -= int64(n)

   993 			if c.readRemaining > 0 && c.readErr == io.EOF {

   994 				c.readErr = errUnexpectedEOF

   995 			}

   996 			return n, c.readErr

   997 		}

   998 

   999 		if c.readFinal {

  1000 			c.messageReader = nil

  1001 			return 0, io.EOF

  1002 		}

  1003 

  1004 		frameType, err := c.advanceFrame()

  1005 		switch {

  1006 		case err != nil:

  1007 			c.readErr = hideTempErr(err)

  1008 		case frameType == TextMessage || frameType == BinaryMessage:

  1009 			c.readErr = errors.New("websocket: internal error, unexpected text or binary in Reader")

  1010 		}

  1011 	}

  1012 

  1013 	err := c.readErr

  1014 	if err == io.EOF && c.messageReader == r {

  1015 		err = errUnexpectedEOF

  1016 	}

  1017 	return 0, err

  1018 }

  1019 

  1020 func (r *messageReader) Close() error {

  1021 	return nil

  1022 }

  1023 

  1024 // ReadMessage is a helper method for getting a reader using NextReader and

  1025 // reading from that reader to a buffer.

  1026 func (c *Conn) ReadMessage() (messageType int, p []byte, err error) {

  1027 	var r io.Reader

  1028 	messageType, r, err = c.NextReader()

  1029 	if err != nil {

  1030 		return messageType, nil, err

  1031 	}

  1032 	p, err = ioutil.ReadAll(r)

  1033 	return messageType, p, err

  1034 }

  1035 

  1036 // SetReadDeadline sets the read deadline on the underlying network connection.

  1037 // After a read has timed out, the websocket connection state is corrupt and

  1038 // all future reads will return an error. A zero value for t means reads will

  1039 // not time out.

  1040 func (c *Conn) SetReadDeadline(t time.Time) error {

  1041 	return c.conn.SetReadDeadline(t)

  1042 }

  1043 

  1044 // SetReadLimit sets the maximum size for a message read from the peer. If a

  1045 // message exceeds the limit, the connection sends a close message to the peer

  1046 // and returns ErrReadLimit to the application.

  1047 func (c *Conn) SetReadLimit(limit int64) {

  1048 	c.readLimit = limit

  1049 }

  1050 

  1051 // CloseHandler returns the current close handler

  1052 func (c *Conn) CloseHandler() func(code int, text string) error {

  1053 	return c.handleClose

  1054 }

  1055 

  1056 // SetCloseHandler sets the handler for close messages received from the peer.

  1057 // The code argument to h is the received close code or CloseNoStatusReceived

  1058 // if the close message is empty. The default close handler sends a close

  1059 // message back to the peer.

  1060 //

  1061 // The handler function is called from the NextReader, ReadMessage and message

  1062 // reader Read methods. The application must read the connection to process

  1063 // close messages as described in the section on Control Messages above.

  1064 //

  1065 // The connection read methods return a CloseError when a close message is

  1066 // received. Most applications should handle close messages as part of their

  1067 // normal error handling. Applications should only set a close handler when the

  1068 // application must perform some action before sending a close message back to

  1069 // the peer.

  1070 func (c *Conn) SetCloseHandler(h func(code int, text string) error) {

  1071 	if h == nil {

  1072 		h = func(code int, text string) error {

  1073 			message := FormatCloseMessage(code, "")

  1074 			c.WriteControl(CloseMessage, message, time.Now().Add(writeWait))

  1075 			return nil

  1076 		}

  1077 	}

  1078 	c.handleClose = h

  1079 }

  1080 

  1081 // PingHandler returns the current ping handler

  1082 func (c *Conn) PingHandler() func(appData string) error {

  1083 	return c.handlePing

  1084 }

  1085 

  1086 // SetPingHandler sets the handler for ping messages received from the peer.

  1087 // The appData argument to h is the PING message application data. The default

  1088 // ping handler sends a pong to the peer.

  1089 //

  1090 // The handler function is called from the NextReader, ReadMessage and message

  1091 // reader Read methods. The application must read the connection to process

  1092 // ping messages as described in the section on Control Messages above.

  1093 func (c *Conn) SetPingHandler(h func(appData string) error) {

  1094 	if h == nil {

  1095 		h = func(message string) error {

  1096 			err := c.WriteControl(PongMessage, []byte(message), time.Now().Add(writeWait))

  1097 			if err == ErrCloseSent {

  1098 				return nil

  1099 			} else if e, ok := err.(net.Error); ok && e.Temporary() {

  1100 				return nil

  1101 			}

  1102 			return err

  1103 		}

  1104 	}

  1105 	c.handlePing = h

  1106 }

  1107 

  1108 // PongHandler returns the current pong handler

  1109 func (c *Conn) PongHandler() func(appData string) error {

  1110 	return c.handlePong

  1111 }

  1112 

  1113 // SetPongHandler sets the handler for pong messages received from the peer.

  1114 // The appData argument to h is the PONG message application data. The default

  1115 // pong handler does nothing.

  1116 //

  1117 // The handler function is called from the NextReader, ReadMessage and message

  1118 // reader Read methods. The application must read the connection to process

  1119 // pong messages as described in the section on Control Messages above.

  1120 func (c *Conn) SetPongHandler(h func(appData string) error) {

  1121 	if h == nil {

  1122 		h = func(string) error { return nil }

  1123 	}

  1124 	c.handlePong = h

  1125 }

  1126 

  1127 // UnderlyingConn returns the internal net.Conn. This can be used to further

  1128 // modifications to connection specific flags.

  1129 func (c *Conn) UnderlyingConn() net.Conn {

  1130 	return c.conn

  1131 }

  1132 

  1133 // EnableWriteCompression enables and disables write compression of

  1134 // subsequent text and binary messages. This function is a noop if

  1135 // compression was not negotiated with the peer.