conn.go 28 KB

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