conn.go 26 KB

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