server.go 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365
  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. "errors"
  8. "io"
  9. "net/http"
  10. "net/url"
  11. "strings"
  12. "time"
  13. )
  14. // HandshakeError describes an error with the handshake from the peer.
  15. type HandshakeError struct {
  16. message string
  17. }
  18. func (e HandshakeError) Error() string { return e.message }
  19. // Upgrader specifies parameters for upgrading an HTTP connection to a
  20. // WebSocket connection.
  21. //
  22. // It is safe to call Upgrader's methods concurrently.
  23. type Upgrader struct {
  24. // HandshakeTimeout specifies the duration for the handshake to complete.
  25. HandshakeTimeout time.Duration
  26. // ReadBufferSize and WriteBufferSize specify I/O buffer sizes in bytes. If a buffer
  27. // size is zero, then buffers allocated by the HTTP server are used. The
  28. // I/O buffer sizes do not limit the size of the messages that can be sent
  29. // or received.
  30. ReadBufferSize, WriteBufferSize int
  31. // WriteBufferPool is a pool of buffers for write operations. If the value
  32. // is not set, then write buffers are allocated to the connection for the
  33. // lifetime of the connection.
  34. //
  35. // A pool is most useful when the application has a modest volume of writes
  36. // across a large number of connections.
  37. //
  38. // Applications should use a single pool for each unique value of
  39. // WriteBufferSize.
  40. WriteBufferPool BufferPool
  41. // Subprotocols specifies the server's supported protocols in order of
  42. // preference. If this field is not nil, then the Upgrade method negotiates a
  43. // subprotocol by selecting the first match in this list with a protocol
  44. // requested by the client. If there's no match, then no protocol is
  45. // negotiated (the Sec-Websocket-Protocol header is not included in the
  46. // handshake response).
  47. Subprotocols []string
  48. // Error specifies the function for generating HTTP error responses. If Error
  49. // is nil, then http.Error is used to generate the HTTP response.
  50. Error func(w http.ResponseWriter, r *http.Request, status int, reason error)
  51. // CheckOrigin returns true if the request Origin header is acceptable. If
  52. // CheckOrigin is nil, then a safe default is used: return false if the
  53. // Origin request header is present and the origin host is not equal to
  54. // request Host header.
  55. //
  56. // A CheckOrigin function should carefully validate the request origin to
  57. // prevent cross-site request forgery.
  58. CheckOrigin func(r *http.Request) bool
  59. // EnableCompression specify if the server should attempt to negotiate per
  60. // message compression (RFC 7692). Setting this value to true does not
  61. // guarantee that compression will be supported. Currently only "no context
  62. // takeover" modes are supported.
  63. EnableCompression bool
  64. }
  65. func (u *Upgrader) returnError(w http.ResponseWriter, r *http.Request, status int, reason string) (*Conn, error) {
  66. err := HandshakeError{reason}
  67. if u.Error != nil {
  68. u.Error(w, r, status, err)
  69. } else {
  70. w.Header().Set("Sec-Websocket-Version", "13")
  71. http.Error(w, http.StatusText(status), status)
  72. }
  73. return nil, err
  74. }
  75. // checkSameOrigin returns true if the origin is not set or is equal to the request host.
  76. func checkSameOrigin(r *http.Request) bool {
  77. origin := r.Header["Origin"]
  78. if len(origin) == 0 {
  79. return true
  80. }
  81. u, err := url.Parse(origin[0])
  82. if err != nil {
  83. return false
  84. }
  85. return equalASCIIFold(u.Host, r.Host)
  86. }
  87. func (u *Upgrader) selectSubprotocol(r *http.Request, responseHeader http.Header) string {
  88. if u.Subprotocols != nil {
  89. clientProtocols := Subprotocols(r)
  90. for _, serverProtocol := range u.Subprotocols {
  91. for _, clientProtocol := range clientProtocols {
  92. if clientProtocol == serverProtocol {
  93. return clientProtocol
  94. }
  95. }
  96. }
  97. } else if responseHeader != nil {
  98. return responseHeader.Get("Sec-Websocket-Protocol")
  99. }
  100. return ""
  101. }
  102. // Upgrade upgrades the HTTP server connection to the WebSocket protocol.
  103. //
  104. // The responseHeader is included in the response to the client's upgrade
  105. // request. Use the responseHeader to specify cookies (Set-Cookie). To specify
  106. // subprotocols supported by the server, set Upgrader.Subprotocols directly.
  107. //
  108. // If the upgrade fails, then Upgrade replies to the client with an HTTP error
  109. // response.
  110. func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header) (*Conn, error) {
  111. const badHandshake = "websocket: the client is not using the websocket protocol: "
  112. if !tokenListContainsValue(r.Header, "Connection", "upgrade") {
  113. return u.returnError(w, r, http.StatusBadRequest, badHandshake+"'upgrade' token not found in 'Connection' header")
  114. }
  115. if !tokenListContainsValue(r.Header, "Upgrade", "websocket") {
  116. return u.returnError(w, r, http.StatusBadRequest, badHandshake+"'websocket' token not found in 'Upgrade' header")
  117. }
  118. if r.Method != http.MethodGet {
  119. return u.returnError(w, r, http.StatusMethodNotAllowed, badHandshake+"request method is not GET")
  120. }
  121. if !tokenListContainsValue(r.Header, "Sec-Websocket-Version", "13") {
  122. return u.returnError(w, r, http.StatusBadRequest, "websocket: unsupported version: 13 not found in 'Sec-Websocket-Version' header")
  123. }
  124. if _, ok := responseHeader["Sec-Websocket-Extensions"]; ok {
  125. return u.returnError(w, r, http.StatusInternalServerError, "websocket: application specific 'Sec-WebSocket-Extensions' headers are unsupported")
  126. }
  127. checkOrigin := u.CheckOrigin
  128. if checkOrigin == nil {
  129. checkOrigin = checkSameOrigin
  130. }
  131. if !checkOrigin(r) {
  132. return u.returnError(w, r, http.StatusForbidden, "websocket: request origin not allowed by Upgrader.CheckOrigin")
  133. }
  134. challengeKey := r.Header.Get("Sec-Websocket-Key")
  135. if challengeKey == "" {
  136. return u.returnError(w, r, http.StatusBadRequest, "websocket: not a websocket handshake: 'Sec-WebSocket-Key' header is missing or blank")
  137. }
  138. subprotocol := u.selectSubprotocol(r, responseHeader)
  139. // Negotiate PMCE
  140. var compress bool
  141. if u.EnableCompression {
  142. for _, ext := range parseExtensions(r.Header) {
  143. if ext[""] != "permessage-deflate" {
  144. continue
  145. }
  146. compress = true
  147. break
  148. }
  149. }
  150. h, ok := w.(http.Hijacker)
  151. if !ok {
  152. return u.returnError(w, r, http.StatusInternalServerError, "websocket: response does not implement http.Hijacker")
  153. }
  154. var brw *bufio.ReadWriter
  155. netConn, brw, err := h.Hijack()
  156. if err != nil {
  157. return u.returnError(w, r, http.StatusInternalServerError, err.Error())
  158. }
  159. if brw.Reader.Buffered() > 0 {
  160. netConn.Close()
  161. return nil, errors.New("websocket: client sent data before handshake is complete")
  162. }
  163. var br *bufio.Reader
  164. if u.ReadBufferSize == 0 && bufioReaderSize(netConn, brw.Reader) > 256 {
  165. // Reuse hijacked buffered reader as connection reader.
  166. br = brw.Reader
  167. }
  168. buf := bufioWriterBuffer(netConn, brw.Writer)
  169. var writeBuf []byte
  170. if u.WriteBufferPool == nil && u.WriteBufferSize == 0 && len(buf) >= maxFrameHeaderSize+256 {
  171. // Reuse hijacked write buffer as connection buffer.
  172. writeBuf = buf
  173. }
  174. c := newConn(netConn, true, u.ReadBufferSize, u.WriteBufferSize, u.WriteBufferPool, br, writeBuf)
  175. c.subprotocol = subprotocol
  176. if compress {
  177. c.newCompressionWriter = compressNoContextTakeover
  178. c.newDecompressionReader = decompressNoContextTakeover
  179. }
  180. // Use larger of hijacked buffer and connection write buffer for header.
  181. p := buf
  182. if len(c.writeBuf) > len(p) {
  183. p = c.writeBuf
  184. }
  185. p = p[:0]
  186. p = append(p, "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: "...)
  187. p = append(p, computeAcceptKey(challengeKey)...)
  188. p = append(p, "\r\n"...)
  189. if c.subprotocol != "" {
  190. p = append(p, "Sec-WebSocket-Protocol: "...)
  191. p = append(p, c.subprotocol...)
  192. p = append(p, "\r\n"...)
  193. }
  194. if compress {
  195. p = append(p, "Sec-WebSocket-Extensions: permessage-deflate; server_no_context_takeover; client_no_context_takeover\r\n"...)
  196. }
  197. for k, vs := range responseHeader {
  198. if k == "Sec-Websocket-Protocol" {
  199. continue
  200. }
  201. for _, v := range vs {
  202. p = append(p, k...)
  203. p = append(p, ": "...)
  204. for i := 0; i < len(v); i++ {
  205. b := v[i]
  206. if b <= 31 {
  207. // prevent response splitting.
  208. b = ' '
  209. }
  210. p = append(p, b)
  211. }
  212. p = append(p, "\r\n"...)
  213. }
  214. }
  215. p = append(p, "\r\n"...)
  216. // Clear deadlines set by HTTP server.
  217. netConn.SetDeadline(time.Time{})
  218. if u.HandshakeTimeout > 0 {
  219. netConn.SetWriteDeadline(time.Now().Add(u.HandshakeTimeout))
  220. }
  221. if _, err = netConn.Write(p); err != nil {
  222. netConn.Close()
  223. return nil, err
  224. }
  225. if u.HandshakeTimeout > 0 {
  226. netConn.SetWriteDeadline(time.Time{})
  227. }
  228. return c, nil
  229. }
  230. // Upgrade upgrades the HTTP server connection to the WebSocket protocol.
  231. //
  232. // Deprecated: Use websocket.Upgrader instead.
  233. //
  234. // Upgrade does not perform origin checking. The application is responsible for
  235. // checking the Origin header before calling Upgrade. An example implementation
  236. // of the same origin policy check is:
  237. //
  238. // if req.Header.Get("Origin") != "http://"+req.Host {
  239. // http.Error(w, "Origin not allowed", http.StatusForbidden)
  240. // return
  241. // }
  242. //
  243. // If the endpoint supports subprotocols, then the application is responsible
  244. // for negotiating the protocol used on the connection. Use the Subprotocols()
  245. // function to get the subprotocols requested by the client. Use the
  246. // Sec-Websocket-Protocol response header to specify the subprotocol selected
  247. // by the application.
  248. //
  249. // The responseHeader is included in the response to the client's upgrade
  250. // request. Use the responseHeader to specify cookies (Set-Cookie) and the
  251. // negotiated subprotocol (Sec-Websocket-Protocol).
  252. //
  253. // The connection buffers IO to the underlying network connection. The
  254. // readBufSize and writeBufSize parameters specify the size of the buffers to
  255. // use. Messages can be larger than the buffers.
  256. //
  257. // If the request is not a valid WebSocket handshake, then Upgrade returns an
  258. // error of type HandshakeError. Applications should handle this error by
  259. // replying to the client with an HTTP error response.
  260. func Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header, readBufSize, writeBufSize int) (*Conn, error) {
  261. u := Upgrader{ReadBufferSize: readBufSize, WriteBufferSize: writeBufSize}
  262. u.Error = func(w http.ResponseWriter, r *http.Request, status int, reason error) {
  263. // don't return errors to maintain backwards compatibility
  264. }
  265. u.CheckOrigin = func(r *http.Request) bool {
  266. // allow all connections by default
  267. return true
  268. }
  269. return u.Upgrade(w, r, responseHeader)
  270. }
  271. // Subprotocols returns the subprotocols requested by the client in the
  272. // Sec-Websocket-Protocol header.
  273. func Subprotocols(r *http.Request) []string {
  274. h := strings.TrimSpace(r.Header.Get("Sec-Websocket-Protocol"))
  275. if h == "" {
  276. return nil
  277. }
  278. protocols := strings.Split(h, ",")
  279. for i := range protocols {
  280. protocols[i] = strings.TrimSpace(protocols[i])
  281. }
  282. return protocols
  283. }
  284. // IsWebSocketUpgrade returns true if the client requested upgrade to the
  285. // WebSocket protocol.
  286. func IsWebSocketUpgrade(r *http.Request) bool {
  287. return tokenListContainsValue(r.Header, "Connection", "upgrade") &&
  288. tokenListContainsValue(r.Header, "Upgrade", "websocket")
  289. }
  290. // bufioReaderSize size returns the size of a bufio.Reader.
  291. func bufioReaderSize(originalReader io.Reader, br *bufio.Reader) int {
  292. // This code assumes that peek on a reset reader returns
  293. // bufio.Reader.buf[:0].
  294. // TODO: Use bufio.Reader.Size() after Go 1.10
  295. br.Reset(originalReader)
  296. if p, err := br.Peek(0); err == nil {
  297. return cap(p)
  298. }
  299. return 0
  300. }
  301. // writeHook is an io.Writer that records the last slice passed to it vio
  302. // io.Writer.Write.
  303. type writeHook struct {
  304. p []byte
  305. }
  306. func (wh *writeHook) Write(p []byte) (int, error) {
  307. wh.p = p
  308. return len(p), nil
  309. }
  310. // bufioWriterBuffer grabs the buffer from a bufio.Writer.
  311. func bufioWriterBuffer(originalWriter io.Writer, bw *bufio.Writer) []byte {
  312. // This code assumes that bufio.Writer.buf[:1] is passed to the
  313. // bufio.Writer's underlying writer.
  314. var wh writeHook
  315. bw.Reset(&wh)
  316. bw.WriteByte(0)
  317. bw.Flush()
  318. bw.Reset(originalWriter)
  319. return wh.p[:cap(wh.p)]
  320. }