conn.go 30 KB

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