conn.go 25 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006
  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. type Conn struct {
  190. conn net.Conn
  191. isServer bool
  192. subprotocol string
  193. // Write fields
  194. mu chan bool // used as mutex to protect write to conn
  195. writeBuf []byte // frame is constructed in this buffer.
  196. writeDeadline time.Time
  197. writer io.WriteCloser // the current writer returned to the application
  198. isWriting bool // for best-effort concurrent write detection
  199. writeErrMu sync.Mutex
  200. writeErr error
  201. enableWriteCompression bool
  202. newCompressionWriter func(io.WriteCloser) (io.WriteCloser, error)
  203. // Read fields
  204. readErr error
  205. br *bufio.Reader
  206. readRemaining int64 // bytes remaining in current frame.
  207. readFinal bool // true the current message has more frames.
  208. readLength int64 // Message size.
  209. readLimit int64 // Maximum message size.
  210. readMaskPos int
  211. readMaskKey [4]byte
  212. handlePong func(string) error
  213. handlePing func(string) error
  214. readErrCount int
  215. messageReader *messageReader // the current low-level reader
  216. readDecompress bool // whether last read frame had RSV1 set
  217. newDecompressionReader func(io.Reader) io.Reader
  218. }
  219. func newConn(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int) *Conn {
  220. mu := make(chan bool, 1)
  221. mu <- true
  222. if readBufferSize == 0 {
  223. readBufferSize = defaultReadBufferSize
  224. }
  225. if readBufferSize < maxControlFramePayloadSize {
  226. readBufferSize = maxControlFramePayloadSize
  227. }
  228. if writeBufferSize == 0 {
  229. writeBufferSize = defaultWriteBufferSize
  230. }
  231. c := &Conn{
  232. isServer: isServer,
  233. br: bufio.NewReaderSize(conn, readBufferSize),
  234. conn: conn,
  235. mu: mu,
  236. readFinal: true,
  237. writeBuf: make([]byte, writeBufferSize+maxFrameHeaderSize),
  238. enableWriteCompression: true,
  239. }
  240. c.SetPingHandler(nil)
  241. c.SetPongHandler(nil)
  242. return c
  243. }
  244. // Subprotocol returns the negotiated protocol for the connection.
  245. func (c *Conn) Subprotocol() string {
  246. return c.subprotocol
  247. }
  248. // Close closes the underlying network connection without sending or waiting for a close frame.
  249. func (c *Conn) Close() error {
  250. return c.conn.Close()
  251. }
  252. // LocalAddr returns the local network address.
  253. func (c *Conn) LocalAddr() net.Addr {
  254. return c.conn.LocalAddr()
  255. }
  256. // RemoteAddr returns the remote network address.
  257. func (c *Conn) RemoteAddr() net.Addr {
  258. return c.conn.RemoteAddr()
  259. }
  260. // Write methods
  261. func (c *Conn) writeFatal(err error) error {
  262. err = hideTempErr(err)
  263. c.writeErrMu.Lock()
  264. if c.writeErr == nil {
  265. c.writeErr = err
  266. }
  267. c.writeErrMu.Unlock()
  268. return err
  269. }
  270. func (c *Conn) write(frameType int, deadline time.Time, bufs ...[]byte) error {
  271. <-c.mu
  272. defer func() { c.mu <- true }()
  273. c.writeErrMu.Lock()
  274. err := c.writeErr
  275. c.writeErrMu.Unlock()
  276. if err != nil {
  277. return err
  278. }
  279. c.conn.SetWriteDeadline(deadline)
  280. for _, buf := range bufs {
  281. if len(buf) > 0 {
  282. _, err := c.conn.Write(buf)
  283. if err != nil {
  284. return c.writeFatal(err)
  285. }
  286. }
  287. }
  288. if frameType == CloseMessage {
  289. c.writeFatal(ErrCloseSent)
  290. }
  291. return nil
  292. }
  293. // WriteControl writes a control message with the given deadline. The allowed
  294. // message types are CloseMessage, PingMessage and PongMessage.
  295. func (c *Conn) WriteControl(messageType int, data []byte, deadline time.Time) error {
  296. if !isControl(messageType) {
  297. return errBadWriteOpCode
  298. }
  299. if len(data) > maxControlFramePayloadSize {
  300. return errInvalidControlFrame
  301. }
  302. b0 := byte(messageType) | finalBit
  303. b1 := byte(len(data))
  304. if !c.isServer {
  305. b1 |= maskBit
  306. }
  307. buf := make([]byte, 0, maxFrameHeaderSize+maxControlFramePayloadSize)
  308. buf = append(buf, b0, b1)
  309. if c.isServer {
  310. buf = append(buf, data...)
  311. } else {
  312. key := newMaskKey()
  313. buf = append(buf, key[:]...)
  314. buf = append(buf, data...)
  315. maskBytes(key, 0, buf[6:])
  316. }
  317. d := time.Hour * 1000
  318. if !deadline.IsZero() {
  319. d = deadline.Sub(time.Now())
  320. if d < 0 {
  321. return errWriteTimeout
  322. }
  323. }
  324. timer := time.NewTimer(d)
  325. select {
  326. case <-c.mu:
  327. timer.Stop()
  328. case <-timer.C:
  329. return errWriteTimeout
  330. }
  331. defer func() { c.mu <- true }()
  332. c.writeErrMu.Lock()
  333. err := c.writeErr
  334. c.writeErrMu.Unlock()
  335. if err != nil {
  336. return err
  337. }
  338. c.conn.SetWriteDeadline(deadline)
  339. _, err = c.conn.Write(buf)
  340. if err != nil {
  341. return c.writeFatal(err)
  342. }
  343. if messageType == CloseMessage {
  344. c.writeFatal(ErrCloseSent)
  345. }
  346. return err
  347. }
  348. // NextWriter returns a writer for the next message to send. The writer's Close
  349. // method flushes the complete message to the network.
  350. //
  351. // There can be at most one open writer on a connection. NextWriter closes the
  352. // previous writer if the application has not already done so.
  353. func (c *Conn) NextWriter(messageType int) (io.WriteCloser, error) {
  354. // Close previous writer if not already closed by the application. It's
  355. // probably better to return an error in this situation, but we cannot
  356. // change this without breaking existing applications.
  357. if c.writer != nil {
  358. c.writer.Close()
  359. c.writer = nil
  360. }
  361. if !isControl(messageType) && !isData(messageType) {
  362. return nil, errBadWriteOpCode
  363. }
  364. c.writeErrMu.Lock()
  365. err := c.writeErr
  366. c.writeErrMu.Unlock()
  367. if err != nil {
  368. return nil, err
  369. }
  370. mw := &messageWriter{
  371. c: c,
  372. frameType: messageType,
  373. pos: maxFrameHeaderSize,
  374. }
  375. c.writer = mw
  376. if c.newCompressionWriter != nil && c.enableWriteCompression && isData(messageType) {
  377. w, err := c.newCompressionWriter(c.writer)
  378. if err != nil {
  379. c.writer = nil
  380. return nil, err
  381. }
  382. mw.compress = true
  383. c.writer = w
  384. }
  385. return c.writer, nil
  386. }
  387. type messageWriter struct {
  388. c *Conn
  389. compress bool // whether next call to flushFrame should set RSV1
  390. pos int // end of data in writeBuf.
  391. frameType int // type of the current frame.
  392. err error
  393. }
  394. func (w *messageWriter) fatal(err error) error {
  395. if w.err != nil {
  396. w.err = err
  397. w.c.writer = nil
  398. }
  399. return err
  400. }
  401. // flushFrame writes buffered data and extra as a frame to the network. The
  402. // final argument indicates that this is the last frame in the message.
  403. func (w *messageWriter) flushFrame(final bool, extra []byte) error {
  404. c := w.c
  405. length := w.pos - maxFrameHeaderSize + len(extra)
  406. // Check for invalid control frames.
  407. if isControl(w.frameType) &&
  408. (!final || length > maxControlFramePayloadSize) {
  409. return w.fatal(errInvalidControlFrame)
  410. }
  411. b0 := byte(w.frameType)
  412. if final {
  413. b0 |= finalBit
  414. }
  415. if w.compress {
  416. b0 |= rsv1Bit
  417. }
  418. w.compress = false
  419. b1 := byte(0)
  420. if !c.isServer {
  421. b1 |= maskBit
  422. }
  423. // Assume that the frame starts at beginning of c.writeBuf.
  424. framePos := 0
  425. if c.isServer {
  426. // Adjust up if mask not included in the header.
  427. framePos = 4
  428. }
  429. switch {
  430. case length >= 65536:
  431. c.writeBuf[framePos] = b0
  432. c.writeBuf[framePos+1] = b1 | 127
  433. binary.BigEndian.PutUint64(c.writeBuf[framePos+2:], uint64(length))
  434. case length > 125:
  435. framePos += 6
  436. c.writeBuf[framePos] = b0
  437. c.writeBuf[framePos+1] = b1 | 126
  438. binary.BigEndian.PutUint16(c.writeBuf[framePos+2:], uint16(length))
  439. default:
  440. framePos += 8
  441. c.writeBuf[framePos] = b0
  442. c.writeBuf[framePos+1] = b1 | byte(length)
  443. }
  444. if !c.isServer {
  445. key := newMaskKey()
  446. copy(c.writeBuf[maxFrameHeaderSize-4:], key[:])
  447. maskBytes(key, 0, c.writeBuf[maxFrameHeaderSize:w.pos])
  448. if len(extra) > 0 {
  449. return c.writeFatal(errors.New("websocket: internal error, extra used in client mode"))
  450. }
  451. }
  452. // Write the buffers to the connection with best-effort detection of
  453. // concurrent writes. See the concurrency section in the package
  454. // documentation for more info.
  455. if c.isWriting {
  456. panic("concurrent write to websocket connection")
  457. }
  458. c.isWriting = true
  459. err := c.write(w.frameType, c.writeDeadline, c.writeBuf[framePos:w.pos], extra)
  460. if !c.isWriting {
  461. panic("concurrent write to websocket connection")
  462. }
  463. c.isWriting = false
  464. if err != nil {
  465. return w.fatal(err)
  466. }
  467. if final {
  468. c.writer = nil
  469. return nil
  470. }
  471. // Setup for next frame.
  472. w.pos = maxFrameHeaderSize
  473. w.frameType = continuationFrame
  474. return nil
  475. }
  476. func (w *messageWriter) ncopy(max int) (int, error) {
  477. n := len(w.c.writeBuf) - w.pos
  478. if n <= 0 {
  479. if err := w.flushFrame(false, nil); err != nil {
  480. return 0, err
  481. }
  482. n = len(w.c.writeBuf) - w.pos
  483. }
  484. if n > max {
  485. n = max
  486. }
  487. return n, nil
  488. }
  489. func (w *messageWriter) Write(p []byte) (int, error) {
  490. if w.err != nil {
  491. return 0, w.err
  492. }
  493. if len(p) > 2*len(w.c.writeBuf) && w.c.isServer {
  494. // Don't buffer large messages.
  495. err := w.flushFrame(false, p)
  496. if err != nil {
  497. return 0, err
  498. }
  499. return len(p), nil
  500. }
  501. nn := len(p)
  502. for len(p) > 0 {
  503. n, err := w.ncopy(len(p))
  504. if err != nil {
  505. return 0, err
  506. }
  507. copy(w.c.writeBuf[w.pos:], p[:n])
  508. w.pos += n
  509. p = p[n:]
  510. }
  511. return nn, nil
  512. }
  513. func (w *messageWriter) WriteString(p string) (int, error) {
  514. if w.err != nil {
  515. return 0, w.err
  516. }
  517. nn := len(p)
  518. for len(p) > 0 {
  519. n, err := w.ncopy(len(p))
  520. if err != nil {
  521. return 0, err
  522. }
  523. copy(w.c.writeBuf[w.pos:], p[:n])
  524. w.pos += n
  525. p = p[n:]
  526. }
  527. return nn, nil
  528. }
  529. func (w *messageWriter) ReadFrom(r io.Reader) (nn int64, err error) {
  530. if w.err != nil {
  531. return 0, w.err
  532. }
  533. for {
  534. if w.pos == len(w.c.writeBuf) {
  535. err = w.flushFrame(false, nil)
  536. if err != nil {
  537. break
  538. }
  539. }
  540. var n int
  541. n, err = r.Read(w.c.writeBuf[w.pos:])
  542. w.pos += n
  543. nn += int64(n)
  544. if err != nil {
  545. if err == io.EOF {
  546. err = nil
  547. }
  548. break
  549. }
  550. }
  551. return nn, err
  552. }
  553. func (w *messageWriter) Close() error {
  554. if w.err != nil {
  555. return w.err
  556. }
  557. if err := w.flushFrame(true, nil); err != nil {
  558. return err
  559. }
  560. w.err = errWriteClosed
  561. return nil
  562. }
  563. // WriteMessage is a helper method for getting a writer using NextWriter,
  564. // writing the message and closing the writer.
  565. func (c *Conn) WriteMessage(messageType int, data []byte) error {
  566. w, err := c.NextWriter(messageType)
  567. if err != nil {
  568. return err
  569. }
  570. if mw, ok := w.(*messageWriter); ok && c.isServer {
  571. // Optimize write as a single frame.
  572. n := copy(c.writeBuf[mw.pos:], data)
  573. mw.pos += n
  574. data = data[n:]
  575. err = mw.flushFrame(true, data)
  576. return err
  577. }
  578. if _, err = w.Write(data); err != nil {
  579. return err
  580. }
  581. return w.Close()
  582. }
  583. // SetWriteDeadline sets the write deadline on the underlying network
  584. // connection. After a write has timed out, the websocket state is corrupt and
  585. // all future writes will return an error. A zero value for t means writes will
  586. // not time out.
  587. func (c *Conn) SetWriteDeadline(t time.Time) error {
  588. c.writeDeadline = t
  589. return nil
  590. }
  591. // Read methods
  592. func (c *Conn) advanceFrame() (int, error) {
  593. // 1. Skip remainder of previous frame.
  594. if c.readRemaining > 0 {
  595. if _, err := io.CopyN(ioutil.Discard, c.br, c.readRemaining); err != nil {
  596. return noFrame, err
  597. }
  598. }
  599. // 2. Read and parse first two bytes of frame header.
  600. p, err := c.read(2)
  601. if err != nil {
  602. return noFrame, err
  603. }
  604. final := p[0]&finalBit != 0
  605. frameType := int(p[0] & 0xf)
  606. mask := p[1]&maskBit != 0
  607. c.readRemaining = int64(p[1] & 0x7f)
  608. c.readDecompress = false
  609. if c.newDecompressionReader != nil && (p[0]&rsv1Bit) != 0 {
  610. c.readDecompress = true
  611. p[0] &^= rsv1Bit
  612. }
  613. if rsv := p[0] & (rsv1Bit | rsv2Bit | rsv3Bit); rsv != 0 {
  614. return noFrame, c.handleProtocolError("unexpected reserved bits 0x" + strconv.FormatInt(int64(rsv), 16))
  615. }
  616. switch frameType {
  617. case CloseMessage, PingMessage, PongMessage:
  618. if c.readRemaining > maxControlFramePayloadSize {
  619. return noFrame, c.handleProtocolError("control frame length > 125")
  620. }
  621. if !final {
  622. return noFrame, c.handleProtocolError("control frame not final")
  623. }
  624. case TextMessage, BinaryMessage:
  625. if !c.readFinal {
  626. return noFrame, c.handleProtocolError("message start before final message frame")
  627. }
  628. c.readFinal = final
  629. case continuationFrame:
  630. if c.readFinal {
  631. return noFrame, c.handleProtocolError("continuation after final message frame")
  632. }
  633. c.readFinal = final
  634. default:
  635. return noFrame, c.handleProtocolError("unknown opcode " + strconv.Itoa(frameType))
  636. }
  637. // 3. Read and parse frame length.
  638. switch c.readRemaining {
  639. case 126:
  640. p, err := c.read(2)
  641. if err != nil {
  642. return noFrame, err
  643. }
  644. c.readRemaining = int64(binary.BigEndian.Uint16(p))
  645. case 127:
  646. p, err := c.read(8)
  647. if err != nil {
  648. return noFrame, err
  649. }
  650. c.readRemaining = int64(binary.BigEndian.Uint64(p))
  651. }
  652. // 4. Handle frame masking.
  653. if mask != c.isServer {
  654. return noFrame, c.handleProtocolError("incorrect mask flag")
  655. }
  656. if mask {
  657. c.readMaskPos = 0
  658. p, err := c.read(len(c.readMaskKey))
  659. if err != nil {
  660. return noFrame, err
  661. }
  662. copy(c.readMaskKey[:], p)
  663. }
  664. // 5. For text and binary messages, enforce read limit and return.
  665. if frameType == continuationFrame || frameType == TextMessage || frameType == BinaryMessage {
  666. c.readLength += c.readRemaining
  667. if c.readLimit > 0 && c.readLength > c.readLimit {
  668. c.WriteControl(CloseMessage, FormatCloseMessage(CloseMessageTooBig, ""), time.Now().Add(writeWait))
  669. return noFrame, ErrReadLimit
  670. }
  671. return frameType, nil
  672. }
  673. // 6. Read control frame payload.
  674. var payload []byte
  675. if c.readRemaining > 0 {
  676. payload, err = c.read(int(c.readRemaining))
  677. c.readRemaining = 0
  678. if err != nil {
  679. return noFrame, err
  680. }
  681. if c.isServer {
  682. maskBytes(c.readMaskKey, 0, payload)
  683. }
  684. }
  685. // 7. Process control frame payload.
  686. switch frameType {
  687. case PongMessage:
  688. if err := c.handlePong(string(payload)); err != nil {
  689. return noFrame, err
  690. }
  691. case PingMessage:
  692. if err := c.handlePing(string(payload)); err != nil {
  693. return noFrame, err
  694. }
  695. case CloseMessage:
  696. echoMessage := []byte{}
  697. closeCode := CloseNoStatusReceived
  698. closeText := ""
  699. if len(payload) >= 2 {
  700. echoMessage = payload[:2]
  701. closeCode = int(binary.BigEndian.Uint16(payload))
  702. if !isValidReceivedCloseCode(closeCode) {
  703. return noFrame, c.handleProtocolError("invalid close code")
  704. }
  705. closeText = string(payload[2:])
  706. if !utf8.ValidString(closeText) {
  707. return noFrame, c.handleProtocolError("invalid utf8 payload in close frame")
  708. }
  709. }
  710. c.WriteControl(CloseMessage, echoMessage, time.Now().Add(writeWait))
  711. return noFrame, &CloseError{Code: closeCode, Text: closeText}
  712. }
  713. return frameType, nil
  714. }
  715. func (c *Conn) handleProtocolError(message string) error {
  716. c.WriteControl(CloseMessage, FormatCloseMessage(CloseProtocolError, message), time.Now().Add(writeWait))
  717. return errors.New("websocket: " + message)
  718. }
  719. // NextReader returns the next data message received from the peer. The
  720. // returned messageType is either TextMessage or BinaryMessage.
  721. //
  722. // There can be at most one open reader on a connection. NextReader discards
  723. // the previous message if the application has not already consumed it.
  724. //
  725. // Applications must break out of the application's read loop when this method
  726. // returns a non-nil error value. Errors returned from this method are
  727. // permanent. Once this method returns a non-nil error, all subsequent calls to
  728. // this method return the same error.
  729. func (c *Conn) NextReader() (messageType int, r io.Reader, err error) {
  730. c.messageReader = nil
  731. c.readLength = 0
  732. for c.readErr == nil {
  733. frameType, err := c.advanceFrame()
  734. if err != nil {
  735. c.readErr = hideTempErr(err)
  736. break
  737. }
  738. if frameType == TextMessage || frameType == BinaryMessage {
  739. c.messageReader = &messageReader{c}
  740. var r io.Reader = c.messageReader
  741. if c.readDecompress {
  742. r = c.newDecompressionReader(r)
  743. }
  744. return frameType, r, nil
  745. }
  746. }
  747. // Applications that do handle the error returned from this method spin in
  748. // tight loop on connection failure. To help application developers detect
  749. // this error, panic on repeated reads to the failed connection.
  750. c.readErrCount++
  751. if c.readErrCount >= 1000 {
  752. panic("repeated read on failed websocket connection")
  753. }
  754. return noFrame, nil, c.readErr
  755. }
  756. type messageReader struct{ c *Conn }
  757. func (r *messageReader) Read(b []byte) (int, error) {
  758. c := r.c
  759. if c.messageReader != r {
  760. return 0, io.EOF
  761. }
  762. for c.readErr == nil {
  763. if c.readRemaining > 0 {
  764. if int64(len(b)) > c.readRemaining {
  765. b = b[:c.readRemaining]
  766. }
  767. n, err := c.br.Read(b)
  768. c.readErr = hideTempErr(err)
  769. if c.isServer {
  770. c.readMaskPos = maskBytes(c.readMaskKey, c.readMaskPos, b[:n])
  771. }
  772. c.readRemaining -= int64(n)
  773. if c.readRemaining > 0 && c.readErr == io.EOF {
  774. c.readErr = errUnexpectedEOF
  775. }
  776. return n, c.readErr
  777. }
  778. if c.readFinal {
  779. c.messageReader = nil
  780. return 0, io.EOF
  781. }
  782. frameType, err := c.advanceFrame()
  783. switch {
  784. case err != nil:
  785. c.readErr = hideTempErr(err)
  786. case frameType == TextMessage || frameType == BinaryMessage:
  787. c.readErr = errors.New("websocket: internal error, unexpected text or binary in Reader")
  788. }
  789. }
  790. err := c.readErr
  791. if err == io.EOF && c.messageReader == r {
  792. err = errUnexpectedEOF
  793. }
  794. return 0, err
  795. }
  796. // ReadMessage is a helper method for getting a reader using NextReader and
  797. // reading from that reader to a buffer.
  798. func (c *Conn) ReadMessage() (messageType int, p []byte, err error) {
  799. var r io.Reader
  800. messageType, r, err = c.NextReader()
  801. if err != nil {
  802. return messageType, nil, err
  803. }
  804. p, err = ioutil.ReadAll(r)
  805. return messageType, p, err
  806. }
  807. // SetReadDeadline sets the read deadline on the underlying network connection.
  808. // After a read has timed out, the websocket connection state is corrupt and
  809. // all future reads will return an error. A zero value for t means reads will
  810. // not time out.
  811. func (c *Conn) SetReadDeadline(t time.Time) error {
  812. return c.conn.SetReadDeadline(t)
  813. }
  814. // SetReadLimit sets the maximum size for a message read from the peer. If a
  815. // message exceeds the limit, the connection sends a close frame to the peer
  816. // and returns ErrReadLimit to the application.
  817. func (c *Conn) SetReadLimit(limit int64) {
  818. c.readLimit = limit
  819. }
  820. // PingHandler returns the current ping handler
  821. func (c *Conn) PingHandler() func(appData string) error {
  822. return c.handlePing
  823. }
  824. // SetPingHandler sets the handler for ping messages received from the peer.
  825. // The appData argument to h is the PING frame application data. The default
  826. // ping handler sends a pong to the peer.
  827. func (c *Conn) SetPingHandler(h func(appData string) error) {
  828. if h == nil {
  829. h = func(message string) error {
  830. err := c.WriteControl(PongMessage, []byte(message), time.Now().Add(writeWait))
  831. if err == ErrCloseSent {
  832. return nil
  833. } else if e, ok := err.(net.Error); ok && e.Temporary() {
  834. return nil
  835. }
  836. return err
  837. }
  838. }
  839. c.handlePing = h
  840. }
  841. // PongHandler returns the current pong handler
  842. func (c *Conn) PongHandler() func(appData string) error {
  843. return c.handlePong
  844. }
  845. // SetPongHandler sets the handler for pong messages received from the peer.
  846. // The appData argument to h is the PONG frame application data. The default
  847. // pong handler does nothing.
  848. func (c *Conn) SetPongHandler(h func(appData string) error) {
  849. if h == nil {
  850. h = func(string) error { return nil }
  851. }
  852. c.handlePong = h
  853. }
  854. // UnderlyingConn returns the internal net.Conn. This can be used to further
  855. // modifications to connection specific flags.
  856. func (c *Conn) UnderlyingConn() net.Conn {
  857. return c.conn
  858. }
  859. // EnableWriteCompression enables and disables write compression of
  860. // subsequent text and binary messages. This function is a noop if
  861. // compression was not negotiated with the peer.
  862. func (c *Conn) EnableWriteCompression(enable bool) {
  863. c.enableWriteCompression = enable
  864. }
  865. // FormatCloseMessage formats closeCode and text as a WebSocket close message.
  866. func FormatCloseMessage(closeCode int, text string) []byte {
  867. buf := make([]byte, 2+len(text))
  868. binary.BigEndian.PutUint16(buf, uint16(closeCode))
  869. copy(buf[2:], text)
  870. return buf
  871. }