1
0

conn.go 30 KB

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