websocket.go 10 KB

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