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