websocket.go 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451
  1. // Copyright 2009 The Go 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 implements a client and server for the WebSocket protocol
  5. // as specified in RFC 6455.
  6. //
  7. // This package currently lacks some features found in alternative
  8. // and more actively maintained WebSocket packages:
  9. //
  10. // https://godoc.org/github.com/gorilla/websocket
  11. // https://godoc.org/nhooyr.io/websocket
  12. package websocket // import "golang.org/x/net/websocket"
  13. import (
  14. "bufio"
  15. "crypto/tls"
  16. "encoding/json"
  17. "errors"
  18. "io"
  19. "io/ioutil"
  20. "net"
  21. "net/http"
  22. "net/url"
  23. "sync"
  24. "time"
  25. )
  26. const (
  27. ProtocolVersionHybi13 = 13
  28. ProtocolVersionHybi = ProtocolVersionHybi13
  29. SupportedProtocolVersion = "13"
  30. ContinuationFrame = 0
  31. TextFrame = 1
  32. BinaryFrame = 2
  33. CloseFrame = 8
  34. PingFrame = 9
  35. PongFrame = 10
  36. UnknownFrame = 255
  37. DefaultMaxPayloadBytes = 32 << 20 // 32MB
  38. )
  39. // ProtocolError represents WebSocket protocol errors.
  40. type ProtocolError struct {
  41. ErrorString string
  42. }
  43. func (err *ProtocolError) Error() string { return err.ErrorString }
  44. var (
  45. ErrBadProtocolVersion = &ProtocolError{"bad protocol version"}
  46. ErrBadScheme = &ProtocolError{"bad scheme"}
  47. ErrBadStatus = &ProtocolError{"bad status"}
  48. ErrBadUpgrade = &ProtocolError{"missing or bad upgrade"}
  49. ErrBadWebSocketOrigin = &ProtocolError{"missing or bad WebSocket-Origin"}
  50. ErrBadWebSocketLocation = &ProtocolError{"missing or bad WebSocket-Location"}
  51. ErrBadWebSocketProtocol = &ProtocolError{"missing or bad WebSocket-Protocol"}
  52. ErrBadWebSocketVersion = &ProtocolError{"missing or bad WebSocket Version"}
  53. ErrChallengeResponse = &ProtocolError{"mismatch challenge/response"}
  54. ErrBadFrame = &ProtocolError{"bad frame"}
  55. ErrBadFrameBoundary = &ProtocolError{"not on frame boundary"}
  56. ErrNotWebSocket = &ProtocolError{"not websocket protocol"}
  57. ErrBadRequestMethod = &ProtocolError{"bad method"}
  58. ErrNotSupported = &ProtocolError{"not supported"}
  59. )
  60. // ErrFrameTooLarge is returned by Codec's Receive method if payload size
  61. // exceeds limit set by Conn.MaxPayloadBytes
  62. var ErrFrameTooLarge = errors.New("websocket: frame payload size exceeds limit")
  63. // Addr is an implementation of net.Addr for WebSocket.
  64. type Addr struct {
  65. *url.URL
  66. }
  67. // Network returns the network type for a WebSocket, "websocket".
  68. func (addr *Addr) Network() string { return "websocket" }
  69. // Config is a WebSocket configuration
  70. type Config struct {
  71. // A WebSocket server address.
  72. Location *url.URL
  73. // A Websocket client origin.
  74. Origin *url.URL
  75. // WebSocket subprotocols.
  76. Protocol []string
  77. // WebSocket protocol version.
  78. Version int
  79. // TLS config for secure WebSocket (wss).
  80. TlsConfig *tls.Config
  81. // Additional header fields to be sent in WebSocket opening handshake.
  82. Header http.Header
  83. // Dialer used when opening websocket connections.
  84. Dialer *net.Dialer
  85. handshakeData map[string]string
  86. }
  87. // serverHandshaker is an interface to handle WebSocket server side handshake.
  88. type serverHandshaker interface {
  89. // ReadHandshake reads handshake request message from client.
  90. // Returns http response code and error if any.
  91. ReadHandshake(buf *bufio.Reader, req *http.Request) (code int, err error)
  92. // AcceptHandshake accepts the client handshake request and sends
  93. // handshake response back to client.
  94. AcceptHandshake(buf *bufio.Writer) (err error)
  95. // NewServerConn creates a new WebSocket connection.
  96. NewServerConn(buf *bufio.ReadWriter, rwc io.ReadWriteCloser, request *http.Request) (conn *Conn)
  97. }
  98. // frameReader is an interface to read a WebSocket frame.
  99. type frameReader interface {
  100. // Reader is to read payload of the frame.
  101. io.Reader
  102. // PayloadType returns payload type.
  103. PayloadType() byte
  104. // HeaderReader returns a reader to read header of the frame.
  105. HeaderReader() io.Reader
  106. // TrailerReader returns a reader to read trailer of the frame.
  107. // If it returns nil, there is no trailer in the frame.
  108. TrailerReader() io.Reader
  109. // Len returns total length of the frame, including header and trailer.
  110. Len() int
  111. }
  112. // frameReaderFactory is an interface to creates new frame reader.
  113. type frameReaderFactory interface {
  114. NewFrameReader() (r frameReader, err error)
  115. }
  116. // frameWriter is an interface to write a WebSocket frame.
  117. type frameWriter interface {
  118. // Writer is to write payload of the frame.
  119. io.WriteCloser
  120. }
  121. // frameWriterFactory is an interface to create new frame writer.
  122. type frameWriterFactory interface {
  123. NewFrameWriter(payloadType byte) (w frameWriter, err error)
  124. }
  125. type frameHandler interface {
  126. HandleFrame(frame frameReader) (r frameReader, err error)
  127. WriteClose(status int) (err error)
  128. }
  129. // Conn represents a WebSocket connection.
  130. //
  131. // Multiple goroutines may invoke methods on a Conn simultaneously.
  132. type Conn struct {
  133. config *Config
  134. request *http.Request
  135. buf *bufio.ReadWriter
  136. rwc io.ReadWriteCloser
  137. rio sync.Mutex
  138. frameReaderFactory
  139. frameReader
  140. wio sync.Mutex
  141. frameWriterFactory
  142. frameHandler
  143. PayloadType byte
  144. defaultCloseStatus int
  145. // MaxPayloadBytes limits the size of frame payload received over Conn
  146. // by Codec's Receive method. If zero, DefaultMaxPayloadBytes is used.
  147. MaxPayloadBytes int
  148. }
  149. // Read implements the io.Reader interface:
  150. // it reads data of a frame from the WebSocket connection.
  151. // if msg is not large enough for the frame data, it fills the msg and next Read
  152. // will read the rest of the frame data.
  153. // it reads Text frame or Binary frame.
  154. func (ws *Conn) Read(msg []byte) (n int, err error) {
  155. ws.rio.Lock()
  156. defer ws.rio.Unlock()
  157. again:
  158. if ws.frameReader == nil {
  159. frame, err := ws.frameReaderFactory.NewFrameReader()
  160. if err != nil {
  161. return 0, err
  162. }
  163. ws.frameReader, err = ws.frameHandler.HandleFrame(frame)
  164. if err != nil {
  165. return 0, err
  166. }
  167. if ws.frameReader == nil {
  168. goto again
  169. }
  170. }
  171. n, err = ws.frameReader.Read(msg)
  172. if err == io.EOF {
  173. if trailer := ws.frameReader.TrailerReader(); trailer != nil {
  174. io.Copy(ioutil.Discard, trailer)
  175. }
  176. ws.frameReader = nil
  177. goto again
  178. }
  179. return n, err
  180. }
  181. // Write implements the io.Writer interface:
  182. // it writes data as a frame to the WebSocket connection.
  183. func (ws *Conn) Write(msg []byte) (n int, err error) {
  184. ws.wio.Lock()
  185. defer ws.wio.Unlock()
  186. w, err := ws.frameWriterFactory.NewFrameWriter(ws.PayloadType)
  187. if err != nil {
  188. return 0, err
  189. }
  190. n, err = w.Write(msg)
  191. w.Close()
  192. return n, err
  193. }
  194. // Close implements the io.Closer interface.
  195. func (ws *Conn) Close() error {
  196. err := ws.frameHandler.WriteClose(ws.defaultCloseStatus)
  197. err1 := ws.rwc.Close()
  198. if err != nil {
  199. return err
  200. }
  201. return err1
  202. }
  203. // IsClientConn reports whether ws is a client-side connection.
  204. func (ws *Conn) IsClientConn() bool { return ws.request == nil }
  205. // IsServerConn reports whether ws is a server-side connection.
  206. func (ws *Conn) IsServerConn() bool { return ws.request != nil }
  207. // LocalAddr returns the WebSocket Origin for the connection for client, or
  208. // the WebSocket location for server.
  209. func (ws *Conn) LocalAddr() net.Addr {
  210. if ws.IsClientConn() {
  211. return &Addr{ws.config.Origin}
  212. }
  213. return &Addr{ws.config.Location}
  214. }
  215. // RemoteAddr returns the WebSocket location for the connection for client, or
  216. // the Websocket Origin for server.
  217. func (ws *Conn) RemoteAddr() net.Addr {
  218. if ws.IsClientConn() {
  219. return &Addr{ws.config.Location}
  220. }
  221. return &Addr{ws.config.Origin}
  222. }
  223. var errSetDeadline = errors.New("websocket: cannot set deadline: not using a net.Conn")
  224. // SetDeadline sets the connection's network read & write deadlines.
  225. func (ws *Conn) SetDeadline(t time.Time) error {
  226. if conn, ok := ws.rwc.(net.Conn); ok {
  227. return conn.SetDeadline(t)
  228. }
  229. return errSetDeadline
  230. }
  231. // SetReadDeadline sets the connection's network read deadline.
  232. func (ws *Conn) SetReadDeadline(t time.Time) error {
  233. if conn, ok := ws.rwc.(net.Conn); ok {
  234. return conn.SetReadDeadline(t)
  235. }
  236. return errSetDeadline
  237. }
  238. // SetWriteDeadline sets the connection's network write deadline.
  239. func (ws *Conn) SetWriteDeadline(t time.Time) error {
  240. if conn, ok := ws.rwc.(net.Conn); ok {
  241. return conn.SetWriteDeadline(t)
  242. }
  243. return errSetDeadline
  244. }
  245. // Config returns the WebSocket config.
  246. func (ws *Conn) Config() *Config { return ws.config }
  247. // Request returns the http request upgraded to the WebSocket.
  248. // It is nil for client side.
  249. func (ws *Conn) Request() *http.Request { return ws.request }
  250. // Codec represents a symmetric pair of functions that implement a codec.
  251. type Codec struct {
  252. Marshal func(v interface{}) (data []byte, payloadType byte, err error)
  253. Unmarshal func(data []byte, payloadType byte, v interface{}) (err error)
  254. }
  255. // Send sends v marshaled by cd.Marshal as single frame to ws.
  256. func (cd Codec) Send(ws *Conn, v interface{}) (err error) {
  257. data, payloadType, err := cd.Marshal(v)
  258. if err != nil {
  259. return err
  260. }
  261. ws.wio.Lock()
  262. defer ws.wio.Unlock()
  263. w, err := ws.frameWriterFactory.NewFrameWriter(payloadType)
  264. if err != nil {
  265. return err
  266. }
  267. _, err = w.Write(data)
  268. w.Close()
  269. return err
  270. }
  271. // Receive receives single frame from ws, unmarshaled by cd.Unmarshal and stores
  272. // in v. The whole frame payload is read to an in-memory buffer; max size of
  273. // payload is defined by ws.MaxPayloadBytes. If frame payload size exceeds
  274. // limit, ErrFrameTooLarge is returned; in this case frame is not read off wire
  275. // completely. The next call to Receive would read and discard leftover data of
  276. // previous oversized frame before processing next frame.
  277. func (cd Codec) Receive(ws *Conn, v interface{}) (err error) {
  278. ws.rio.Lock()
  279. defer ws.rio.Unlock()
  280. if ws.frameReader != nil {
  281. _, err = io.Copy(ioutil.Discard, ws.frameReader)
  282. if err != nil {
  283. return err
  284. }
  285. ws.frameReader = nil
  286. }
  287. again:
  288. frame, err := ws.frameReaderFactory.NewFrameReader()
  289. if err != nil {
  290. return err
  291. }
  292. frame, err = ws.frameHandler.HandleFrame(frame)
  293. if err != nil {
  294. return err
  295. }
  296. if frame == nil {
  297. goto again
  298. }
  299. maxPayloadBytes := ws.MaxPayloadBytes
  300. if maxPayloadBytes == 0 {
  301. maxPayloadBytes = DefaultMaxPayloadBytes
  302. }
  303. if hf, ok := frame.(*hybiFrameReader); ok && hf.header.Length > int64(maxPayloadBytes) {
  304. // payload size exceeds limit, no need to call Unmarshal
  305. //
  306. // set frameReader to current oversized frame so that
  307. // the next call to this function can drain leftover
  308. // data before processing the next frame
  309. ws.frameReader = frame
  310. return ErrFrameTooLarge
  311. }
  312. payloadType := frame.PayloadType()
  313. data, err := ioutil.ReadAll(frame)
  314. if err != nil {
  315. return err
  316. }
  317. return cd.Unmarshal(data, payloadType, v)
  318. }
  319. func marshal(v interface{}) (msg []byte, payloadType byte, err error) {
  320. switch data := v.(type) {
  321. case string:
  322. return []byte(data), TextFrame, nil
  323. case []byte:
  324. return data, BinaryFrame, nil
  325. }
  326. return nil, UnknownFrame, ErrNotSupported
  327. }
  328. func unmarshal(msg []byte, payloadType byte, v interface{}) (err error) {
  329. switch data := v.(type) {
  330. case *string:
  331. *data = string(msg)
  332. return nil
  333. case *[]byte:
  334. *data = msg
  335. return nil
  336. }
  337. return ErrNotSupported
  338. }
  339. /*
  340. Message is a codec to send/receive text/binary data in a frame on WebSocket connection.
  341. To send/receive text frame, use string type.
  342. To send/receive binary frame, use []byte type.
  343. Trivial usage:
  344. import "websocket"
  345. // receive text frame
  346. var message string
  347. websocket.Message.Receive(ws, &message)
  348. // send text frame
  349. message = "hello"
  350. websocket.Message.Send(ws, message)
  351. // receive binary frame
  352. var data []byte
  353. websocket.Message.Receive(ws, &data)
  354. // send binary frame
  355. data = []byte{0, 1, 2}
  356. websocket.Message.Send(ws, data)
  357. */
  358. var Message = Codec{marshal, unmarshal}
  359. func jsonMarshal(v interface{}) (msg []byte, payloadType byte, err error) {
  360. msg, err = json.Marshal(v)
  361. return msg, TextFrame, err
  362. }
  363. func jsonUnmarshal(msg []byte, payloadType byte, v interface{}) (err error) {
  364. return json.Unmarshal(msg, v)
  365. }
  366. /*
  367. JSON is a codec to send/receive JSON data in a frame from a WebSocket connection.
  368. Trivial usage:
  369. import "websocket"
  370. type T struct {
  371. Msg string
  372. Count int
  373. }
  374. // receive JSON type T
  375. var data T
  376. websocket.JSON.Receive(ws, &data)
  377. // send JSON type T
  378. websocket.JSON.Send(ws, data)
  379. */
  380. var JSON = Codec{jsonMarshal, jsonUnmarshal}