conn.go raw

   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