server.go 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113
  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
  5. import (
  6. "bufio"
  7. "fmt"
  8. "io"
  9. "net/http"
  10. )
  11. func newServerConn(rwc io.ReadWriteCloser, buf *bufio.ReadWriter, req *http.Request, config *Config, handshake func(*Config, *http.Request) error) (conn *Conn, err error) {
  12. var hs serverHandshaker = &hybiServerHandshaker{Config: config}
  13. code, err := hs.ReadHandshake(buf.Reader, req)
  14. if err == ErrBadWebSocketVersion {
  15. fmt.Fprintf(buf, "HTTP/1.1 %03d %s\r\n", code, http.StatusText(code))
  16. fmt.Fprintf(buf, "Sec-WebSocket-Version: %s\r\n", SupportedProtocolVersion)
  17. buf.WriteString("\r\n")
  18. buf.WriteString(err.Error())
  19. buf.Flush()
  20. return
  21. }
  22. if err != nil {
  23. fmt.Fprintf(buf, "HTTP/1.1 %03d %s\r\n", code, http.StatusText(code))
  24. buf.WriteString("\r\n")
  25. buf.WriteString(err.Error())
  26. buf.Flush()
  27. return
  28. }
  29. if handshake != nil {
  30. err = handshake(config, req)
  31. if err != nil {
  32. code = http.StatusForbidden
  33. fmt.Fprintf(buf, "HTTP/1.1 %03d %s\r\n", code, http.StatusText(code))
  34. buf.WriteString("\r\n")
  35. buf.Flush()
  36. return
  37. }
  38. }
  39. err = hs.AcceptHandshake(buf.Writer)
  40. if err != nil {
  41. code = http.StatusBadRequest
  42. fmt.Fprintf(buf, "HTTP/1.1 %03d %s\r\n", code, http.StatusText(code))
  43. buf.WriteString("\r\n")
  44. buf.Flush()
  45. return
  46. }
  47. conn = hs.NewServerConn(buf, rwc, req)
  48. return
  49. }
  50. // Server represents a server of a WebSocket.
  51. type Server struct {
  52. // Config is a WebSocket configuration for new WebSocket connection.
  53. Config
  54. // Handshake is an optional function in WebSocket handshake.
  55. // For example, you can check, or don't check Origin header.
  56. // Another example, you can select config.Protocol.
  57. Handshake func(*Config, *http.Request) error
  58. // Handler handles a WebSocket connection.
  59. Handler
  60. }
  61. // ServeHTTP implements the http.Handler interface for a WebSocket
  62. func (s Server) ServeHTTP(w http.ResponseWriter, req *http.Request) {
  63. s.serveWebSocket(w, req)
  64. }
  65. func (s Server) serveWebSocket(w http.ResponseWriter, req *http.Request) {
  66. rwc, buf, err := w.(http.Hijacker).Hijack()
  67. if err != nil {
  68. panic("Hijack failed: " + err.Error())
  69. }
  70. // The server should abort the WebSocket connection if it finds
  71. // the client did not send a handshake that matches with protocol
  72. // specification.
  73. defer rwc.Close()
  74. conn, err := newServerConn(rwc, buf, req, &s.Config, s.Handshake)
  75. if err != nil {
  76. return
  77. }
  78. if conn == nil {
  79. panic("unexpected nil conn")
  80. }
  81. s.Handler(conn)
  82. }
  83. // Handler is a simple interface to a WebSocket browser client.
  84. // It checks if Origin header is valid URL by default.
  85. // You might want to verify websocket.Conn.Config().Origin in the func.
  86. // If you use Server instead of Handler, you could call websocket.Origin and
  87. // check the origin in your Handshake func. So, if you want to accept
  88. // non-browser clients, which do not send an Origin header, set a
  89. // Server.Handshake that does not check the origin.
  90. type Handler func(*Conn)
  91. func checkOrigin(config *Config, req *http.Request) (err error) {
  92. config.Origin, err = Origin(config, req)
  93. if err == nil && config.Origin == nil {
  94. return fmt.Errorf("null origin")
  95. }
  96. return err
  97. }
  98. // ServeHTTP implements the http.Handler interface for a WebSocket
  99. func (h Handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
  100. s := Server{Handler: h, Handshake: checkOrigin}
  101. s.serveWebSocket(w, req)
  102. }