conn.go 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053
  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. package websocket
  5. import (
  6. "bufio"
  7. "encoding/binary"
  8. "errors"
  9. "io"
  10. "io/ioutil"
  11. "net"
  12. "strconv"
  13. "sync"
  14. "time"
  15. "unicode/utf8"
  16. )
  17. const (
  18. // Frame header byte 0 bits from Section 5.2 of RFC 6455
  19. finalBit = 1 << 7
  20. rsv1Bit = 1 << 6
  21. rsv2Bit = 1 << 5
  22. rsv3Bit = 1 << 4
  23. // Frame header byte 1 bits from Section 5.2 of RFC 6455
  24. maskBit = 1 << 7
  25. maxFrameHeaderSize = 2 + 8 + 4 // Fixed header + length + mask
  26. maxControlFramePayloadSize = 125
  27. writeWait = time.Second
  28. defaultReadBufferSize = 4096
  29. defaultWriteBufferSize = 4096
  30. continuationFrame = 0
  31. noFrame = -1
  32. )
  33. // Close codes defined in RFC 6455, section 11.7.
  34. const (
  35. CloseNormalClosure = 1000
  36. CloseGoingAway = 1001
  37. CloseProtocolError = 1002
  38. CloseUnsupportedData = 1003
  39. CloseNoStatusReceived = 1005
  40. CloseAbnormalClosure = 1006
  41. CloseInvalidFramePayloadData = 1007
  42. ClosePolicyViolation = 1008
  43. CloseMessageTooBig = 1009
  44. CloseMandatoryExtension = 1010
  45. CloseInternalServerErr = 1011
  46. CloseServiceRestart = 1012
  47. CloseTryAgainLater = 1013
  48. CloseTLSHandshake = 1015
  49. )
  50. // The message types are defined in RFC 6455, section 11.8.
  51. const (
  52. // TextMessage denotes a text data message. The text message payload is
  53. // interpreted as UTF-8 encoded text data.
  54. TextMessage = 1
  55. // BinaryMessage denotes a binary data message.
  56. BinaryMessage = 2
  57. // CloseMessage denotes a close control message. The optional message
  58. // payload contains a numeric code and text. Use the FormatCloseMessage
  59. // function to format a close message payload.
  60. CloseMessage = 8
  61. // PingMessage denotes a ping control message. The optional message payload
  62. // is UTF-8 encoded text.
  63. PingMessage = 9
  64. // PongMessage denotes a ping control message. The optional message payload
  65. // is UTF-8 encoded text.
  66. PongMessage = 10
  67. )
  68. // ErrCloseSent is returned when the application writes a message to the
  69. // connection after sending a close message.
  70. var ErrCloseSent = errors.New("websocket: close sent")
  71. // ErrReadLimit is returned when reading a message that is larger than the
  72. // read limit set for the connection.
  73. var ErrReadLimit = errors.New("websocket: read limit exceeded")
  74. // netError satisfies the net Error interface.
  75. type netError struct {
  76. msg string
  77. temporary bool
  78. timeout bool
  79. }
  80. func (e *netError) Error() string { return e.msg }
  81. func (e *netError) Temporary() bool { return e.temporary }
  82. func (e *netError) Timeout() bool { return e.timeout }
  83. // CloseError represents close frame.
  84. type CloseError struct {
  85. // Code is defined in RFC 6455, section 11.7.
  86. Code int
  87. // Text is the optional text payload.
  88. Text string
  89. }
  90. func (e *CloseError) Error() string {
  91. s := []byte("websocket: close ")
  92. s = strconv.AppendInt(s, int64(e.Code), 10)
  93. switch e.Code {
  94. case CloseNormalClosure:
  95. s = append(s, " (normal)"...)
  96. case CloseGoingAway:
  97. s = append(s, " (going away)"...)
  98. case CloseProtocolError:
  99. s = append(s, " (protocol error)"...)
  100. case CloseUnsupportedData:
  101. s = append(s, " (unsupported data)"...)
  102. case CloseNoStatusReceived:
  103. s = append(s, " (no status)"...)
  104. case CloseAbnormalClosure:
  105. s = append(s, " (abnormal closure)"...)
  106. case CloseInvalidFramePayloadData:
  107. s = append(s, " (invalid payload data)"...)
  108. case ClosePolicyViolation:
  109. s = append(s, " (policy violation)"...)
  110. case CloseMessageTooBig:
  111. s = append(s, " (message too big)"...)
  112. case CloseMandatoryExtension:
  113. s = append(s, " (mandatory extension missing)"...)
  114. case CloseInternalServerErr:
  115. s = append(s, " (internal server error)"...)
  116. case CloseTLSHandshake:
  117. s = append(s, " (TLS handshake error)"...)
  118. }
  119. if e.Text != "" {
  120. s = append(s, ": "...)
  121. s = append(s, e.Text...)
  122. }
  123. return string(s)
  124. }
  125. // IsCloseError returns boolean indicating whether the error is a *CloseError
  126. // with one of the specified codes.
  127. func IsCloseError(err error, codes ...int) bool {
  128. if e, ok := err.(*CloseError); ok {
  129. for _, code := range codes {
  130. if e.Code == code {
  131. return true
  132. }
  133. }
  134. }
  135. return false
  136. }
  137. // IsUnexpectedCloseError returns boolean indicating whether the error is a
  138. // *CloseError with a code not in the list of expected codes.
  139. func IsUnexpectedCloseError(err error, expectedCodes ...int) bool {
  140. if e, ok := err.(*CloseError); ok {
  141. for _, code := range expectedCodes {
  142. if e.Code == code {
  143. return false
  144. }
  145. }
  146. return true
  147. }
  148. return false
  149. }
  150. var (
  151. errWriteTimeout = &netError{msg: "websocket: write timeout", timeout: true, temporary: true}
  152. errUnexpectedEOF = &CloseError{Code: CloseAbnormalClosure, Text: io.ErrUnexpectedEOF.Error()}
  153. errBadWriteOpCode = errors.New("websocket: bad write message type")
  154. errWriteClosed = errors.New("websocket: write closed")
  155. errInvalidControlFrame = errors.New("websocket: invalid control frame")
  156. )
  157. func hideTempErr(err error) error {
  158. if e, ok := err.(net.Error); ok && e.Temporary() {
  159. err = &netError{msg: e.Error(), timeout: e.Timeout()}
  160. }
  161. return err
  162. }
  163. func isControl(frameType int) bool {
  164. return frameType == CloseMessage || frameType == PingMessage || frameType == PongMessage
  165. }
  166. func isData(frameType int) bool {
  167. return frameType == TextMessage || frameType == BinaryMessage
  168. }
  169. var validReceivedCloseCodes = map[int]bool{
  170. // see http://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number
  171. CloseNormalClosure: true,
  172. CloseGoingAway: true,
  173. CloseProtocolError: true,
  174. CloseUnsupportedData: true,
  175. CloseNoStatusReceived: false,
  176. CloseAbnormalClosure: false,
  177. CloseInvalidFramePayloadData: true,
  178. ClosePolicyViolation: true,
  179. CloseMessageTooBig: true,
  180. CloseMandatoryExtension: true,
  181. CloseInternalServerErr: true,
  182. CloseServiceRestart: true,
  183. CloseTryAgainLater: true,
  184. CloseTLSHandshake: false,
  185. }
  186. func isValidReceivedCloseCode(code int) bool {
  187. return validReceivedCloseCodes[code] || (code >= 3000 && code <= 4999)
  188. }
  189. // The Conn type represents a WebSocket connection.
  190. type Conn struct {
  191. conn net.Conn
  192. isServer bool
  193. subprotocol string
  194. // Write fields
  195. mu chan bool // used as mutex to protect write to conn
  196. writeBuf []byte // frame is constructed in this buffer.
  197. writeDeadline time.Time
  198. writer io.WriteCloser // the current writer returned to the application
  199. isWriting bool // for best-effort concurrent write detection
  200. writeErrMu sync.Mutex
  201. writeErr error
  202. enableWriteCompression bool
  203. newCompressionWriter func(io.WriteCloser) (io.WriteCloser, error)
  204. // Read fields
  205. reader io.ReadCloser // the current reader returned to the application
  206. readErr error
  207. br *bufio.Reader
  208. readRemaining int64 // bytes remaining in current frame.
  209. readFinal bool // true the current message has more frames.
  210. readLength int64 // Message size.
  211. readLimit int64 // Maximum message size.
  212. readMaskPos int
  213. readMaskKey [4]byte
  214. handlePong func(string) error
  215. handlePing func(string) error
  216. handleClose func(int, string) error
  217. readErrCount int
  218. messageReader *messageReader // the current low-level reader
  219. readDecompress bool // whether last read frame had RSV1 set
  220. newDecompressionReader func(io.Reader) io.ReadCloser
  221. }
  222. func newConn(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int) *Conn {
  223. mu := make(chan bool, 1)
  224. mu <- true
  225. if readBufferSize == 0 {
  226. readBufferSize = defaultReadBufferSize
  227. }
  228. if readBufferSize < maxControlFramePayloadSize {
  229. readBufferSize = maxControlFramePayloadSize
  230. }
  231. if writeBufferSize == 0 {
  232. writeBufferSize = defaultWriteBufferSize
  233. }
  234. c := &Conn{
  235. isServer: isServer,
  236. br: bufio.NewReaderSize(conn, readBufferSize),
  237. conn: conn,
  238. mu: mu,
  239. readFinal: true,
  240. writeBuf: make([]byte, writeBufferSize+maxFrameHeaderSize),
  241. enableWriteCompression: true,
  242. }
  243. c.SetCloseHandler(nil)
  244. c.SetPingHandler(nil)
  245. c.SetPongHandler(nil)
  246. return c
  247. }
  248. // Subprotocol returns the negotiated protocol for the connection.
  249. func (c *Conn) Subprotocol() string {
  250. return c.subprotocol
  251. }
  252. // Close closes the underlying network connection without sending or waiting for a close frame.
  253. func (c *Conn) Close() error {
  254. return c.conn.Close()
  255. }
  256. // LocalAddr returns the local network address.
  257. func (c *Conn) LocalAddr() net.Addr {
  258. return c.conn.LocalAddr()
  259. }
  260. // RemoteAddr returns the remote network address.
  261. func (c *Conn) RemoteAddr() net.Addr {
  262. return c.conn.RemoteAddr()
  263. }
  264. // Write methods
  265. func (c *Conn) writeFatal(err error) error {
  266. err = hideTempErr(err)
  267. c.writeErrMu.Lock()
  268. if c.writeErr == nil {
  269. c.writeErr = err
  270. }
  271. c.writeErrMu.Unlock()
  272. return err
  273. }
  274. func (c *Conn) write(frameType int, deadline time.Time, bufs ...[]byte) error {
  275. <-c.mu
  276. defer func() { c.mu <- true }()
  277. c.writeErrMu.Lock()
  278. err := c.writeErr
  279. c.writeErrMu.Unlock()
  280. if err != nil {
  281. return err
  282. }
  283. c.conn.SetWriteDeadline(deadline)
  284. for _, buf := range bufs {
  285. if len(buf) > 0 {
  286. _, err := c.conn.Write(buf)
  287. if err != nil {
  288. return c.writeFatal(err)
  289. }
  290. }
  291. }
  292. if frameType == CloseMessage {
  293. c.writeFatal(ErrCloseSent)
  294. }
  295. return nil
  296. }
  297. // WriteControl writes a control message with the given deadline. The allowed
  298. // message types are CloseMessage, PingMessage and PongMessage.
  299. func (c *Conn) WriteControl(messageType int, data []byte, deadline time.Time) error {
  300. if !isControl(messageType) {
  301. return errBadWriteOpCode
  302. }
  303. if len(data) > maxControlFramePayloadSize {
  304. return errInvalidControlFrame
  305. }
  306. b0 := byte(messageType) | finalBit
  307. b1 := byte(len(data))
  308. if !c.isServer {
  309. b1 |= maskBit
  310. }
  311. buf := make([]byte, 0, maxFrameHeaderSize+maxControlFramePayloadSize)
  312. buf = append(buf, b0, b1)
  313. if c.isServer {
  314. buf = append(buf, data...)
  315. } else {
  316. key := newMaskKey()
  317. buf = append(buf, key[:]...)
  318. buf = append(buf, data...)
  319. maskBytes(key, 0, buf[6:])
  320. }
  321. d := time.Hour * 1000
  322. if !deadline.IsZero() {
  323. d = deadline.Sub(time.Now())
  324. if d < 0 {
  325. return errWriteTimeout
  326. }
  327. }
  328. timer := time.NewTimer(d)
  329. select {
  330. case <-c.mu:
  331. timer.Stop()
  332. case <-timer.C:
  333. return errWriteTimeout
  334. }
  335. defer func() { c.mu <- true }()
  336. c.writeErrMu.Lock()
  337. err := c.writeErr
  338. c.writeErrMu.Unlock()
  339. if err != nil {
  340. return err
  341. }
  342. c.conn.SetWriteDeadline(deadline)
  343. _, err = c.conn.Write(buf)
  344. if err != nil {
  345. return c.writeFatal(err)
  346. }
  347. if messageType == CloseMessage {
  348. c.writeFatal(ErrCloseSent)
  349. }
  350. return err
  351. }
  352. func (c *Conn) prepWrite(messageType int) error {
  353. // Close previous writer if not already closed by the application. It's
  354. // probably better to return an error in this situation, but we cannot
  355. // change this without breaking existing applications.
  356. if c.writer != nil {
  357. c.writer.Close()
  358. c.writer = nil
  359. }
  360. if !isControl(messageType) && !isData(messageType) {
  361. return errBadWriteOpCode
  362. }
  363. c.writeErrMu.Lock()
  364. err := c.writeErr
  365. c.writeErrMu.Unlock()
  366. return err
  367. }
  368. // NextWriter returns a writer for the next message to send. The writer's Close
  369. // method flushes the complete message to the network.
  370. //
  371. // There can be at most one open writer on a connection. NextWriter closes the
  372. // previous writer if the application has not already done so.
  373. func (c *Conn) NextWriter(messageType int) (io.WriteCloser, error) {
  374. if err := c.prepWrite(messageType); err != nil {
  375. return nil, err
  376. }
  377. mw := &messageWriter{
  378. c: c,
  379. frameType: messageType,
  380. pos: maxFrameHeaderSize,
  381. }
  382. c.writer = mw
  383. if c.newCompressionWriter != nil && c.enableWriteCompression && isData(messageType) {
  384. w, err := c.newCompressionWriter(c.writer)
  385. if err != nil {
  386. c.writer = nil
  387. return nil, err
  388. }
  389. mw.compress = true
  390. c.writer = w
  391. }
  392. return c.writer, nil
  393. }
  394. type messageWriter struct {
  395. c *Conn
  396. compress bool // whether next call to flushFrame should set RSV1
  397. pos int // end of data in writeBuf.
  398. frameType int // type of the current frame.
  399. err error
  400. }
  401. func (w *messageWriter) fatal(err error) error {
  402. if w.err != nil {
  403. w.err = err
  404. w.c.writer = nil
  405. }
  406. return err
  407. }
  408. // flushFrame writes buffered data and extra as a frame to the network. The
  409. // final argument indicates that this is the last frame in the message.
  410. func (w *messageWriter) flushFrame(final bool, extra []byte) error {
  411. c := w.c
  412. length := w.pos - maxFrameHeaderSize + len(extra)
  413. // Check for invalid control frames.
  414. if isControl(w.frameType) &&
  415. (!final || length > maxControlFramePayloadSize) {
  416. return w.fatal(errInvalidControlFrame)
  417. }
  418. b0 := byte(w.frameType)
  419. if final {
  420. b0 |= finalBit
  421. }
  422. if w.compress {
  423. b0 |= rsv1Bit
  424. }
  425. w.compress = false
  426. b1 := byte(0)
  427. if !c.isServer {
  428. b1 |= maskBit
  429. }
  430. // Assume that the frame starts at beginning of c.writeBuf.
  431. framePos := 0
  432. if c.isServer {
  433. // Adjust up if mask not included in the header.
  434. framePos = 4
  435. }
  436. switch {
  437. case length >= 65536:
  438. c.writeBuf[framePos] = b0
  439. c.writeBuf[framePos+1] = b1 | 127
  440. binary.BigEndian.PutUint64(c.writeBuf[framePos+2:], uint64(length))
  441. case length > 125:
  442. framePos += 6
  443. c.writeBuf[framePos] = b0
  444. c.writeBuf[framePos+1] = b1 | 126
  445. binary.BigEndian.PutUint16(c.writeBuf[framePos+2:], uint16(length))
  446. default:
  447. framePos += 8
  448. c.writeBuf[framePos] = b0
  449. c.writeBuf[framePos+1] = b1 | byte(length)
  450. }
  451. if !c.isServer {
  452. key := newMaskKey()
  453. copy(c.writeBuf[maxFrameHeaderSize-4:], key[:])
  454. maskBytes(key, 0, c.writeBuf[maxFrameHeaderSize:w.pos])
  455. if len(extra) > 0 {
  456. return c.writeFatal(errors.New("websocket: internal error, extra used in client mode"))
  457. }
  458. }
  459. // Write the buffers to the connection with best-effort detection of
  460. // concurrent writes. See the concurrency section in the package
  461. // documentation for more info.
  462. if c.isWriting {
  463. panic("concurrent write to websocket connection")
  464. }
  465. c.isWriting = true
  466. err := c.write(w.frameType, c.writeDeadline, c.writeBuf[framePos:w.pos], extra)
  467. if !c.isWriting {
  468. panic("concurrent write to websocket connection")
  469. }
  470. c.isWriting = false
  471. if err != nil {
  472. return w.fatal(err)
  473. }
  474. if final {
  475. c.writer = nil
  476. return nil
  477. }
  478. // Setup for next frame.
  479. w.pos = maxFrameHeaderSize
  480. w.frameType = continuationFrame
  481. return nil
  482. }
  483. func (w *messageWriter) ncopy(max int) (int, error) {
  484. n := len(w.c.writeBuf) - w.pos
  485. if n <= 0 {
  486. if err := w.flushFrame(false, nil); err != nil {
  487. return 0, err
  488. }
  489. n = len(w.c.writeBuf) - w.pos
  490. }
  491. if n > max {
  492. n = max
  493. }
  494. return n, nil
  495. }
  496. func (w *messageWriter) Write(p []byte) (int, error) {
  497. if w.err != nil {
  498. return 0, w.err
  499. }
  500. if len(p) > 2*len(w.c.writeBuf) && w.c.isServer {
  501. // Don't buffer large messages.
  502. err := w.flushFrame(false, p)
  503. if err != nil {
  504. return 0, err
  505. }
  506. return len(p), nil
  507. }
  508. nn := len(p)
  509. for len(p) > 0 {
  510. n, err := w.ncopy(len(p))
  511. if err != nil {
  512. return 0, err
  513. }
  514. copy(w.c.writeBuf[w.pos:], p[:n])
  515. w.pos += n
  516. p = p[n:]
  517. }
  518. return nn, nil
  519. }
  520. func (w *messageWriter) WriteString(p string) (int, error) {
  521. if w.err != nil {
  522. return 0, w.err
  523. }
  524. nn := len(p)
  525. for len(p) > 0 {
  526. n, err := w.ncopy(len(p))
  527. if err != nil {
  528. return 0, err
  529. }
  530. copy(w.c.writeBuf[w.pos:], p[:n])
  531. w.pos += n
  532. p = p[n:]
  533. }
  534. return nn, nil
  535. }
  536. func (w *messageWriter) ReadFrom(r io.Reader) (nn int64, err error) {
  537. if w.err != nil {
  538. return 0, w.err
  539. }
  540. for {
  541. if w.pos == len(w.c.writeBuf) {
  542. err = w.flushFrame(false, nil)
  543. if err != nil {
  544. break
  545. }
  546. }
  547. var n int
  548. n, err = r.Read(w.c.writeBuf[w.pos:])
  549. w.pos += n
  550. nn += int64(n)
  551. if err != nil {
  552. if err == io.EOF {
  553. err = nil
  554. }
  555. break
  556. }
  557. }
  558. return nn, err
  559. }
  560. func (w *messageWriter) Close() error {
  561. if w.err != nil {
  562. return w.err
  563. }
  564. if err := w.flushFrame(true, nil); err != nil {
  565. return err
  566. }
  567. w.err = errWriteClosed
  568. return nil
  569. }
  570. // WriteMessage is a helper method for getting a writer using NextWriter,
  571. // writing the message and closing the writer.
  572. func (c *Conn) WriteMessage(messageType int, data []byte) error {
  573. if c.isServer && (c.newCompressionWriter == nil || !c.enableWriteCompression) {
  574. // Fast path with no allocations and single frame.
  575. if err := c.prepWrite(messageType); err != nil {
  576. return err
  577. }
  578. mw := messageWriter{c: c, frameType: messageType, pos: maxFrameHeaderSize}
  579. n := copy(c.writeBuf[mw.pos:], data)
  580. mw.pos += n
  581. data = data[n:]
  582. return mw.flushFrame(true, data)
  583. }
  584. w, err := c.NextWriter(messageType)
  585. if err != nil {
  586. return err
  587. }
  588. if _, err = w.Write(data); err != nil {
  589. return err
  590. }
  591. return w.Close()
  592. }
  593. // SetWriteDeadline sets the write deadline on the underlying network
  594. // connection. After a write has timed out, the websocket state is corrupt and
  595. // all future writes will return an error. A zero value for t means writes will
  596. // not time out.
  597. func (c *Conn) SetWriteDeadline(t time.Time) error {
  598. c.writeDeadline = t
  599. return nil
  600. }
  601. // Read methods
  602. func (c *Conn) advanceFrame() (int, error) {
  603. // 1. Skip remainder of previous frame.
  604. if c.readRemaining > 0 {
  605. if _, err := io.CopyN(ioutil.Discard, c.br, c.readRemaining); err != nil {
  606. return noFrame, err
  607. }
  608. }
  609. // 2. Read and parse first two bytes of frame header.
  610. p, err := c.read(2)
  611. if err != nil {
  612. return noFrame, err
  613. }
  614. final := p[0]&finalBit != 0
  615. frameType := int(p[0] & 0xf)
  616. mask := p[1]&maskBit != 0
  617. c.readRemaining = int64(p[1] & 0x7f)
  618. c.readDecompress = false
  619. if c.newDecompressionReader != nil && (p[0]&rsv1Bit) != 0 {
  620. c.readDecompress = true
  621. p[0] &^= rsv1Bit
  622. }
  623. if rsv := p[0] & (rsv1Bit | rsv2Bit | rsv3Bit); rsv != 0 {
  624. return noFrame, c.handleProtocolError("unexpected reserved bits 0x" + strconv.FormatInt(int64(rsv), 16))
  625. }
  626. switch frameType {
  627. case CloseMessage, PingMessage, PongMessage:
  628. if c.readRemaining > maxControlFramePayloadSize {
  629. return noFrame, c.handleProtocolError("control frame length > 125")
  630. }
  631. if !final {
  632. return noFrame, c.handleProtocolError("control frame not final")
  633. }
  634. case TextMessage, BinaryMessage:
  635. if !c.readFinal {
  636. return noFrame, c.handleProtocolError("message start before final message frame")
  637. }
  638. c.readFinal = final
  639. case continuationFrame:
  640. if c.readFinal {
  641. return noFrame, c.handleProtocolError("continuation after final message frame")
  642. }
  643. c.readFinal = final
  644. default:
  645. return noFrame, c.handleProtocolError("unknown opcode " + strconv.Itoa(frameType))
  646. }
  647. // 3. Read and parse frame length.
  648. switch c.readRemaining {
  649. case 126:
  650. p, err := c.read(2)
  651. if err != nil {
  652. return noFrame, err
  653. }
  654. c.readRemaining = int64(binary.BigEndian.Uint16(p))
  655. case 127:
  656. p, err := c.read(8)
  657. if err != nil {
  658. return noFrame, err
  659. }
  660. c.readRemaining = int64(binary.BigEndian.Uint64(p))
  661. }
  662. // 4. Handle frame masking.
  663. if mask != c.isServer {
  664. return noFrame, c.handleProtocolError("incorrect mask flag")
  665. }
  666. if mask {
  667. c.readMaskPos = 0
  668. p, err := c.read(len(c.readMaskKey))
  669. if err != nil {
  670. return noFrame, err
  671. }
  672. copy(c.readMaskKey[:], p)
  673. }
  674. // 5. For text and binary messages, enforce read limit and return.
  675. if frameType == continuationFrame || frameType == TextMessage || frameType == BinaryMessage {
  676. c.readLength += c.readRemaining
  677. if c.readLimit > 0 && c.readLength > c.readLimit {
  678. c.WriteControl(CloseMessage, FormatCloseMessage(CloseMessageTooBig, ""), time.Now().Add(writeWait))
  679. return noFrame, ErrReadLimit
  680. }
  681. return frameType, nil
  682. }
  683. // 6. Read control frame payload.
  684. var payload []byte
  685. if c.readRemaining > 0 {
  686. payload, err = c.read(int(c.readRemaining))
  687. c.readRemaining = 0
  688. if err != nil {
  689. return noFrame, err
  690. }
  691. if c.isServer {
  692. maskBytes(c.readMaskKey, 0, payload)
  693. }
  694. }
  695. // 7. Process control frame payload.
  696. switch frameType {
  697. case PongMessage:
  698. if err := c.handlePong(string(payload)); err != nil {
  699. return noFrame, err
  700. }
  701. case PingMessage:
  702. if err := c.handlePing(string(payload)); err != nil {
  703. return noFrame, err
  704. }
  705. case CloseMessage:
  706. closeCode := CloseNoStatusReceived
  707. closeText := ""
  708. if len(payload) >= 2 {
  709. closeCode = int(binary.BigEndian.Uint16(payload))
  710. if !isValidReceivedCloseCode(closeCode) {
  711. return noFrame, c.handleProtocolError("invalid close code")
  712. }
  713. closeText = string(payload[2:])
  714. if !utf8.ValidString(closeText) {
  715. return noFrame, c.handleProtocolError("invalid utf8 payload in close frame")
  716. }
  717. }
  718. if err := c.handleClose(closeCode, closeText); err != nil {
  719. return noFrame, err
  720. }
  721. return noFrame, &CloseError{Code: closeCode, Text: closeText}
  722. }
  723. return frameType, nil
  724. }
  725. func (c *Conn) handleProtocolError(message string) error {
  726. c.WriteControl(CloseMessage, FormatCloseMessage(CloseProtocolError, message), time.Now().Add(writeWait))
  727. return errors.New("websocket: " + message)
  728. }
  729. // NextReader returns the next data message received from the peer. The
  730. // returned messageType is either TextMessage or BinaryMessage.
  731. //
  732. // There can be at most one open reader on a connection. NextReader discards
  733. // the previous message if the application has not already consumed it.
  734. //
  735. // Applications must break out of the application's read loop when this method
  736. // returns a non-nil error value. Errors returned from this method are
  737. // permanent. Once this method returns a non-nil error, all subsequent calls to
  738. // this method return the same error.
  739. func (c *Conn) NextReader() (messageType int, r io.Reader, err error) {
  740. // Close previous reader, only relevant for decompression.
  741. if c.reader != nil {
  742. c.reader.Close()
  743. c.reader = nil
  744. }
  745. c.messageReader = nil
  746. c.readLength = 0
  747. for c.readErr == nil {
  748. frameType, err := c.advanceFrame()
  749. if err != nil {
  750. c.readErr = hideTempErr(err)
  751. break
  752. }
  753. if frameType == TextMessage || frameType == BinaryMessage {
  754. c.messageReader = &messageReader{c}
  755. c.reader = c.messageReader
  756. if c.readDecompress {
  757. c.reader = c.newDecompressionReader(c.reader)
  758. }
  759. return frameType, c.reader, nil
  760. }
  761. }
  762. // Applications that do handle the error returned from this method spin in
  763. // tight loop on connection failure. To help application developers detect
  764. // this error, panic on repeated reads to the failed connection.
  765. c.readErrCount++
  766. if c.readErrCount >= 1000 {
  767. panic("repeated read on failed websocket connection")
  768. }
  769. return noFrame, nil, c.readErr
  770. }
  771. type messageReader struct{ c *Conn }
  772. func (r *messageReader) Read(b []byte) (int, error) {
  773. c := r.c
  774. if c.messageReader != r {
  775. return 0, io.EOF
  776. }
  777. for c.readErr == nil {
  778. if c.readRemaining > 0 {
  779. if int64(len(b)) > c.readRemaining {
  780. b = b[:c.readRemaining]
  781. }
  782. n, err := c.br.Read(b)
  783. c.readErr = hideTempErr(err)
  784. if c.isServer {
  785. c.readMaskPos = maskBytes(c.readMaskKey, c.readMaskPos, b[:n])
  786. }
  787. c.readRemaining -= int64(n)
  788. if c.readRemaining > 0 && c.readErr == io.EOF {
  789. c.readErr = errUnexpectedEOF
  790. }
  791. return n, c.readErr
  792. }
  793. if c.readFinal {
  794. c.messageReader = nil
  795. return 0, io.EOF
  796. }
  797. frameType, err := c.advanceFrame()
  798. switch {
  799. case err != nil:
  800. c.readErr = hideTempErr(err)
  801. case frameType == TextMessage || frameType == BinaryMessage:
  802. c.readErr = errors.New("websocket: internal error, unexpected text or binary in Reader")
  803. }
  804. }
  805. err := c.readErr
  806. if err == io.EOF && c.messageReader == r {
  807. err = errUnexpectedEOF
  808. }
  809. return 0, err
  810. }
  811. func (r *messageReader) Close() error {
  812. return nil
  813. }
  814. // ReadMessage is a helper method for getting a reader using NextReader and
  815. // reading from that reader to a buffer.
  816. func (c *Conn) ReadMessage() (messageType int, p []byte, err error) {
  817. var r io.Reader
  818. messageType, r, err = c.NextReader()
  819. if err != nil {
  820. return messageType, nil, err
  821. }
  822. p, err = ioutil.ReadAll(r)
  823. return messageType, p, err
  824. }
  825. // SetReadDeadline sets the read deadline on the underlying network connection.
  826. // After a read has timed out, the websocket connection state is corrupt and
  827. // all future reads will return an error. A zero value for t means reads will
  828. // not time out.
  829. func (c *Conn) SetReadDeadline(t time.Time) error {
  830. return c.conn.SetReadDeadline(t)
  831. }
  832. // SetReadLimit sets the maximum size for a message read from the peer. If a
  833. // message exceeds the limit, the connection sends a close frame to the peer
  834. // and returns ErrReadLimit to the application.
  835. func (c *Conn) SetReadLimit(limit int64) {
  836. c.readLimit = limit
  837. }
  838. // CloseHandler returns the current close handler
  839. func (c *Conn) CloseHandler() func(code int, text string) error {
  840. return c.handleClose
  841. }
  842. // SetCloseHandler sets the handler for close messages received from the peer.
  843. // The code argument to h is the received close code or CloseNoStatusReceived
  844. // if the close message is empty. The default close handler sends a close frame
  845. // back to the peer.
  846. func (c *Conn) SetCloseHandler(h func(code int, text string) error) {
  847. if h == nil {
  848. h = func(code int, text string) error {
  849. message := []byte{}
  850. if code != CloseNoStatusReceived {
  851. message = FormatCloseMessage(code, "")
  852. }
  853. c.WriteControl(CloseMessage, message, time.Now().Add(writeWait))
  854. return nil
  855. }
  856. }
  857. c.handleClose = h
  858. }
  859. // PingHandler returns the current ping handler
  860. func (c *Conn) PingHandler() func(appData string) error {
  861. return c.handlePing
  862. }
  863. // SetPingHandler sets the handler for ping messages received from the peer.
  864. // The appData argument to h is the PING frame application data. The default
  865. // ping handler sends a pong to the peer.
  866. func (c *Conn) SetPingHandler(h func(appData string) error) {
  867. if h == nil {
  868. h = func(message string) error {
  869. err := c.WriteControl(PongMessage, []byte(message), time.Now().Add(writeWait))
  870. if err == ErrCloseSent {
  871. return nil
  872. } else if e, ok := err.(net.Error); ok && e.Temporary() {
  873. return nil
  874. }
  875. return err
  876. }
  877. }
  878. c.handlePing = h
  879. }
  880. // PongHandler returns the current pong handler
  881. func (c *Conn) PongHandler() func(appData string) error {
  882. return c.handlePong
  883. }
  884. // SetPongHandler sets the handler for pong messages received from the peer.
  885. // The appData argument to h is the PONG frame application data. The default
  886. // pong handler does nothing.
  887. func (c *Conn) SetPongHandler(h func(appData string) error) {
  888. if h == nil {
  889. h = func(string) error { return nil }
  890. }
  891. c.handlePong = h
  892. }
  893. // UnderlyingConn returns the internal net.Conn. This can be used to further
  894. // modifications to connection specific flags.
  895. func (c *Conn) UnderlyingConn() net.Conn {
  896. return c.conn
  897. }
  898. // EnableWriteCompression enables and disables write compression of
  899. // subsequent text and binary messages. This function is a noop if
  900. // compression was not negotiated with the peer.
  901. func (c *Conn) EnableWriteCompression(enable bool) {
  902. c.enableWriteCompression = enable
  903. }
  904. // FormatCloseMessage formats closeCode and text as a WebSocket close message.
  905. func FormatCloseMessage(closeCode int, text string) []byte {
  906. buf := make([]byte, 2+len(text))
  907. binary.BigEndian.PutUint16(buf, uint16(closeCode))
  908. copy(buf[2:], text)
  909. return buf
  910. }