diff --git a/.codecov.yml b/.codecov.yml index a7a1a9e0..87940329 100644 --- a/.codecov.yml +++ b/.codecov.yml @@ -7,3 +7,4 @@ coverage: ignore: - "internal/controlpb/control*" # generated code - "client_experimental.go" # experimental code + - "internal/websocket/*" # embedded Gorilla WebSocket fork diff --git a/README.md b/README.md index 7789b2d6..76150dc5 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ Library highlights: * Fast and optimized for low-latency communication with millions of client connections. See [test stand with 1 million connections in Kubernetes](https://centrifugal.dev/blog/2020/02/10/million-connections-with-centrifugo) * WebSocket bidirectional transport using JSON or binary Protobuf formats, both based on a [strict Protobuf schema](https://github.com/centrifugal/protocol/blob/master/definitions/client.proto). Code generation is used to push both JSON and Protobuf serialization performance to the limits -* Our [own WebSocket emulation layer](https://centrifugal.dev/blog/2022/07/19/centrifugo-v4-released#modern-websocket-emulation-in-javascript) over HTTP-streaming (JSON + Protobuf) and Eventsource (JSON) without sticky sessions requirement for distributed setup, or SockJS (JSON only) +* Our [own WebSocket emulation layer](https://centrifugal.dev/blog/2022/07/19/centrifugo-v4-released#modern-websocket-emulation-in-javascript) over HTTP-streaming (JSON + Protobuf) and Eventsource (JSON) without sticky sessions requirement for distributed setup * Possibility to use unidirectional transports without using custom Centrifuge SDK library: see examples for [GRPC](https://github.com/centrifugal/centrifuge/tree/master/_examples/unidirectional_grpc), [EventSource(SSE)](https://github.com/centrifugal/centrifuge/tree/master/_examples/unidirectional_sse), [HTTP-streaming](https://github.com/centrifugal/centrifuge/tree/master/_examples/unidirectional_http_stream), [Unidirectional WebSocket](https://github.com/centrifugal/centrifuge/tree/master/_examples/unidirectional_ws) * Built-in horizontal scalability with Redis PUB/SUB, consistent Redis sharding, Redis Sentinel and Redis Cluster support, [super-optimized Redis communication layer](https://centrifugal.dev/blog/2022/12/20/improving-redis-engine-performance) * Effective non-blocking broadcasts towards client connections using individual queues @@ -244,7 +244,7 @@ Some useful advices about library here. Let's describe some aspects related to connection life cycle and event handling in Centrifuge: -* If you set middleware for transport handlers (`WebsocketHandler`, `SockjsHandler`) – then it will be called first before a client sent any command to a server and handler had a chance to start working. Just like a regular HTTP middleware. You can put `Credentials` to `Context` to authenticate connection. +* If you set middleware for transport handlers (like `WebsocketHandler`) – then it will be called first before a client sent any command to a server and handler had a chance to start working. Just like a regular HTTP middleware. You can put `Credentials` to `Context` to authenticate connection. * `node.OnConnecting` called as soon as client sent `Connect` command to server. At this point no `Client` instance exists. You have incoming `Context` and `Transport` information. You still can authenticate Client at this point (based on string token sent from client side or any other way). Also, you can add extra data to context and return modified context to Centrifuge. Context cancelled as soon as client connection closes. This handler is synchronous and connection read loop can't proceed until you return `ConnectReply`. * `node.OnConnect` then called (after a reply to `Connect` command already written to connection). Inside `OnConnect` closure you have a possibility to define per-connection event handlers. If particular handler not set then client will get `ErrorNotAvailable` errors requesting it. Remember that none of event handlers available in Centrifuge should block forever – do minimal work, start separate goroutines if you need blocking code. * Client initiated request handlers called one by one from connection reading goroutine. This includes `OnSubscribe`, `OnPublish`, `OnPresence`, `OnPresenceStats`, `OnHistory`, client-side `OnRefresh`, client-side `OnSubRefresh`. diff --git a/_examples/chat_json/main.go b/_examples/chat_json/main.go index e3347d5b..cfcca570 100644 --- a/_examples/chat_json/main.go +++ b/_examples/chat_json/main.go @@ -245,14 +245,6 @@ func main() { }) mux.Handle("/connection/websocket", authMiddleware(websocketHandler)) - sockjsHandler := centrifuge.NewSockjsHandler(node, centrifuge.SockjsConfig{ - URL: "https://cdn.jsdelivr.net/npm/sockjs-client@1/dist/sockjs.min.js", - HandlerPrefix: "/connection/sockjs", - WebsocketReadBufferSize: 1024, - WebsocketWriteBufferSize: 1024, - }) - mux.Handle("/connection/sockjs/", authMiddleware(sockjsHandler)) - mux.Handle("/metrics", promhttp.Handler()) mux.Handle("/", http.FileServer(http.Dir("./"))) diff --git a/_examples/concurrency/main.go b/_examples/concurrency/main.go index c6ffdcd1..0ca3177c 100644 --- a/_examples/concurrency/main.go +++ b/_examples/concurrency/main.go @@ -111,14 +111,6 @@ func main() { }) http.Handle("/connection/websocket", authMiddleware(websocketHandler)) - sockjsHandler := centrifuge.NewSockjsHandler(node, centrifuge.SockjsConfig{ - URL: "https://cdn.jsdelivr.net/npm/sockjs-client@1/dist/sockjs.min.js", - HandlerPrefix: "/connection/sockjs", - WebsocketReadBufferSize: 1024, - WebsocketWriteBufferSize: 1024, - }) - http.Handle("/connection/sockjs/", authMiddleware(sockjsHandler)) - http.Handle("/", http.FileServer(http.Dir("./"))) go func() { diff --git a/_examples/gin_auth/README.md b/_examples/gin_auth/README.md index 5ec5cf54..e5480121 100644 --- a/_examples/gin_auth/README.md +++ b/_examples/gin_auth/README.md @@ -1,7 +1,5 @@ Example demonstrates a simple chat with JSON protocol sharing session auth with [gin-gonic](https://github.com/gin-gonic/gin). -Client uses Websocket by default, but you can simply uncomment one line in `chat.html` to use SockJS instead. - To start example run the following command from example directory: ``` @@ -15,4 +13,4 @@ There is only one email/pass combination : `email@email.com:password` [gin-gonic]: https://github.com/gin-gonic/gin -_Credits to the example by FZambia from whom I took most of the centrifuge code_ \ No newline at end of file +_Credits to the example by FZambia from whom I took most of the centrifuge code_ diff --git a/_examples/go.mod b/_examples/go.mod index 6542e101..d143d269 100644 --- a/_examples/go.mod +++ b/_examples/go.mod @@ -51,7 +51,6 @@ require ( github.com/google/pprof v0.0.0-20230207041349-798e818bf904 // indirect github.com/gorilla/context v1.1.1 // indirect github.com/gorilla/securecookie v1.1.1 // indirect - github.com/igm/sockjs-go/v3 v3.0.2 // indirect github.com/josharian/intern v1.0.0 // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/klauspost/compress v1.17.0 // indirect diff --git a/_examples/go.sum b/_examples/go.sum index 08d946ce..7a825b9f 100644 --- a/_examples/go.sum +++ b/_examples/go.sum @@ -81,11 +81,8 @@ github.com/gorilla/sessions v1.1.1/go.mod h1:8KCfur6+4Mqcc6S0FEfKuN15Vl5MgXW92AE github.com/gorilla/sessions v1.1.3/go.mod h1:8KCfur6+4Mqcc6S0FEfKuN15Vl5MgXW92AE8ovaJD0w= github.com/gorilla/sessions v1.2.1 h1:DHd3rPN5lE3Ts3D8rKkQ8x/0kqfeNmBAaiSi+o7FsgI= github.com/gorilla/sessions v1.2.1/go.mod h1:dk2InVEVJ0sfLlnXv9EAgkf6ecYs/i80K/zI+bUmuGM= -github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc= github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= -github.com/igm/sockjs-go/v3 v3.0.2 h1:2m0k53w0DBiGozeQUIEPR6snZFmpFpYvVsGnfLPNXbE= -github.com/igm/sockjs-go/v3 v3.0.2/go.mod h1:UqchsOjeagIBFHvd+RZpLaVRbCwGilEC08EDHsD1jYE= github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= github.com/json-iterator/go v1.1.7/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4= @@ -158,7 +155,6 @@ github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSS github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4= -github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA= github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= diff --git a/_examples/single_connection/main.go b/_examples/single_connection/main.go index d18d75bf..129739a3 100644 --- a/_examples/single_connection/main.go +++ b/_examples/single_connection/main.go @@ -93,13 +93,6 @@ func main() { }) http.Handle("/connection/websocket", authMiddleware(websocketHandler)) - sockjsHandler := centrifuge.NewSockjsHandler(node, centrifuge.SockjsConfig{ - URL: "https://cdn.jsdelivr.net/npm/sockjs-client@1/dist/sockjs.min.js", - HandlerPrefix: "/connection/sockjs", - WebsocketReadBufferSize: 1024, - WebsocketWriteBufferSize: 1024, - }) - http.Handle("/connection/sockjs/", authMiddleware(sockjsHandler)) http.Handle("/", http.FileServer(http.Dir("./"))) go func() { diff --git a/doc.go b/doc.go index 04906336..e8726880 100644 --- a/doc.go +++ b/doc.go @@ -1,8 +1,8 @@ // Package centrifuge is a real-time messaging library that abstracts // several bidirectional transports (Websocket and its emulation over -// HTTP-Streaming, SSE/EventSource, SockJS) and provides primitives -// to build scalable real-time applications with Go. It's also used as -// a core of Centrifugo server (https://github.com/centrifugal/centrifugo). +// HTTP-Streaming, SSE/EventSource) and provides primitives to build +// scalable real-time applications with Go. It's also used as a core +// of Centrifugo server (https://github.com/centrifugal/centrifugo). // // Centrifuge library provides several features on top of plain Websocket // implementation and comes with several client SDKs – see more details in diff --git a/go.mod b/go.mod index 02a67522..d8fc9671 100644 --- a/go.mod +++ b/go.mod @@ -6,8 +6,6 @@ require ( github.com/FZambia/eagle v0.1.0 github.com/centrifugal/protocol v0.11.0 github.com/google/uuid v1.5.0 - github.com/gorilla/websocket v1.5.0 - github.com/igm/sockjs-go/v3 v3.0.2 github.com/prometheus/client_golang v1.18.0 github.com/redis/rueidis v1.0.27 github.com/segmentio/encoding v0.4.0 diff --git a/go.sum b/go.sum index da118729..345ba410 100644 --- a/go.sum +++ b/go.sum @@ -6,7 +6,6 @@ github.com/centrifugal/protocol v0.11.0 h1:pQKfVT4c3/uiRNszaOenE4NqJqL4VBlzBku8s github.com/centrifugal/protocol v0.11.0/go.mod h1:33nZhrA2iRoR6jT+oVzu1ARx+iWIgxOgYuZMhWMWVM4= github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= -github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= @@ -14,11 +13,6 @@ github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/ github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/uuid v1.5.0 h1:1p67kYwdtXjb0gL0BPiP1Av9wiZPo5A8z2cWkTZ+eyU= github.com/google/uuid v1.5.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= -github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= -github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc= -github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= -github.com/igm/sockjs-go/v3 v3.0.2 h1:2m0k53w0DBiGozeQUIEPR6snZFmpFpYvVsGnfLPNXbE= -github.com/igm/sockjs-go/v3 v3.0.2/go.mod h1:UqchsOjeagIBFHvd+RZpLaVRbCwGilEC08EDHsD1jYE= github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= @@ -45,8 +39,6 @@ github.com/segmentio/asm v1.2.0 h1:9BQrFxC+YOHJlTlHGkTrFWf59nbL3XnCoFLTwDCI7ys= github.com/segmentio/asm v1.2.0/go.mod h1:BqMnlJP91P8d+4ibuonYZw9mfnzI9HfxselHZr5aAcs= github.com/segmentio/encoding v0.4.0 h1:MEBYvRqiUB2nfR2criEXWqwdY6HJOUrCn5hboVOVmy8= github.com/segmentio/encoding v0.4.0/go.mod h1:/d03Cd8PoaDeceuhUUUQWjU0KhWjrmYrWPgtJHYZSnI= -github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= -github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA= github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk= github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw= @@ -63,6 +55,5 @@ google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= -gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/handler_sockjs.go b/handler_sockjs.go deleted file mode 100644 index 634f0fdc..00000000 --- a/handler_sockjs.go +++ /dev/null @@ -1,287 +0,0 @@ -package centrifuge - -import ( - "net/http" - "sync" - "time" - - "github.com/centrifugal/centrifuge/internal/cancelctx" - "github.com/centrifugal/centrifuge/internal/readerpool" - - "github.com/centrifugal/protocol" - "github.com/gorilla/websocket" - "github.com/igm/sockjs-go/v3/sockjs" -) - -// SockjsConfig represents config for SockJS handler. -type SockjsConfig struct { - // HandlerPrefix sets prefix for SockJS handler endpoint path. - HandlerPrefix string - - // URL is an address to SockJS client javascript library. Required for iframe-based - // transports to work. This URL should lead to the same SockJS client version as used - // for connecting on the client side. - URL string - - // CheckOrigin allows deciding whether to use CORS or not in XHR case. - // When false returned then CORS headers won't be set. - CheckOrigin func(*http.Request) bool - - // WebsocketCheckOrigin allows setting custom CheckOrigin func for underlying - // Gorilla Websocket based websocket.Upgrader. - WebsocketCheckOrigin func(*http.Request) bool - - // WebsocketReadBufferSize is a parameter that is used for raw websocket.Upgrader. - // If set to zero reasonable default value will be used. - WebsocketReadBufferSize int - - // WebsocketWriteBufferSize is a parameter that is used for raw websocket.Upgrader. - // If set to zero reasonable default value will be used. - WebsocketWriteBufferSize int - - // WebsocketUseWriteBufferPool enables using buffer pool for writes in Websocket transport. - WebsocketUseWriteBufferPool bool - - // WebsocketWriteTimeout is maximum time of write message operation. - // Slow client will be disconnected. - // By default, 1 * time.Second will be used. - WebsocketWriteTimeout time.Duration - - PingPongConfig -} - -// SockjsHandler accepts SockJS connections. SockJS has a bunch of fallback -// transports when WebSocket connection is not supported. It comes with additional -// costs though: small protocol framing overhead, lack of binary support, more -// goroutines per connection, and you need to use sticky session mechanism on -// your load balancer in case you are using HTTP-based SockJS fallbacks and have -// more than one Centrifuge Node on a backend (so SockJS to be able to emulate -// bidirectional protocol). So if you can afford it - use WebsocketHandler only. -type SockjsHandler struct { - node *Node - config SockjsConfig - handler http.Handler -} - -// NewSockjsHandler creates new SockjsHandler. -func NewSockjsHandler(node *Node, config SockjsConfig) *SockjsHandler { - options := sockjs.DefaultOptions - - wsUpgrader := &websocket.Upgrader{ - ReadBufferSize: config.WebsocketReadBufferSize, - WriteBufferSize: config.WebsocketWriteBufferSize, - Error: func(w http.ResponseWriter, r *http.Request, status int, reason error) {}, - } - if config.WebsocketCheckOrigin != nil { - wsUpgrader.CheckOrigin = config.WebsocketCheckOrigin - } else { - wsUpgrader.CheckOrigin = sameHostOriginCheck(node) - } - if config.WebsocketUseWriteBufferPool { - wsUpgrader.WriteBufferPool = writeBufferPool - } else { - wsUpgrader.WriteBufferSize = config.WebsocketWriteBufferSize - } - options.WebsocketUpgrader = wsUpgrader - - // Override sockjs url. It's important to use the same SockJS - // library version on client and server sides when using iframe - // based SockJS transports, otherwise SockJS will raise error - // about version mismatch. - options.SockJSURL = config.URL - if config.CheckOrigin != nil { - options.CheckOrigin = config.CheckOrigin - } else { - options.CheckOrigin = sameHostOriginCheck(node) - } - - wsWriteTimeout := config.WebsocketWriteTimeout - if wsWriteTimeout == 0 { - wsWriteTimeout = 1 * time.Second - } - options.WebsocketWriteTimeout = wsWriteTimeout - - s := &SockjsHandler{ - node: node, - config: config, - } - - options.HeartbeatDelay = 0 - s.handler = sockjs.NewHandler(config.HandlerPrefix, options, s.sockJSHandler) - return s -} - -func (s *SockjsHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) { - s.handler.ServeHTTP(rw, r) -} - -// sockJSHandler called when new client connection comes to SockJS endpoint. -func (s *SockjsHandler) sockJSHandler(sess sockjs.Session) { - s.handleSession(sess) -} - -// sockJSHandler called when new client connection comes to SockJS endpoint. -func (s *SockjsHandler) handleSession(sess sockjs.Session) { - s.node.metrics.incTransportConnect(transportSockJS) - - // Separate goroutine for better GC of caller's data. - go func() { - transport := newSockjsTransport(sess, sockjsTransportOptions{ - pingPong: s.config.PingPongConfig, - }) - - select { - case <-s.node.NotifyShutdown(): - _ = transport.Close(DisconnectShutdown) - return - default: - } - - ctxCh := make(chan struct{}) - defer close(ctxCh) - c, closeFn, err := NewClient(cancelctx.New(sess.Request().Context(), ctxCh), s.node, transport) - if err != nil { - s.node.logger.log(newLogEntry(LogLevelError, "error creating client", map[string]any{"transport": transportSockJS})) - return - } - defer func() { _ = closeFn() }() - - if s.node.LogEnabled(LogLevelDebug) { - s.node.logger.log(newLogEntry(LogLevelDebug, "client connection established", map[string]any{"client": c.ID(), "transport": transportSockJS})) - defer func(started time.Time) { - s.node.logger.log(newLogEntry(LogLevelDebug, "client connection completed", map[string]any{"client": c.ID(), "transport": transportSockJS, "duration": time.Since(started)})) - }(time.Now()) - } - - var needWaitLoop bool - - for { - if msg, err := sess.Recv(); err == nil { - reader := readerpool.GetStringReader(msg) - if ok := HandleReadFrame(c, reader); !ok { - readerpool.PutStringReader(reader) - needWaitLoop = true - break - } - readerpool.PutStringReader(reader) - continue - } - break - } - - if needWaitLoop { - // One extra loop till we get an error from session, - // this is required to wait until close frame will be sent - // into connection inside Client implementation and transport - // closed with proper disconnect reason. - for { - if _, err := sess.Recv(); err != nil { - break - } - } - } - }() -} - -const ( - transportSockJS = "sockjs" -) - -type sockjsTransportOptions struct { - pingPong PingPongConfig -} - -type sockjsTransport struct { - mu sync.RWMutex - closeCh chan struct{} - session sockjs.Session - opts sockjsTransportOptions - closed bool -} - -func newSockjsTransport(s sockjs.Session, opts sockjsTransportOptions) *sockjsTransport { - t := &sockjsTransport{ - session: s, - closeCh: make(chan struct{}), - opts: opts, - } - return t -} - -// Name returns name of transport. -func (t *sockjsTransport) Name() string { - return transportSockJS -} - -// Protocol returns transport protocol. -func (t *sockjsTransport) Protocol() ProtocolType { - return ProtocolTypeJSON -} - -// ProtocolVersion returns transport ProtocolVersion. -func (t *sockjsTransport) ProtocolVersion() ProtocolVersion { - return ProtocolVersion2 -} - -// Unidirectional returns whether transport is unidirectional. -func (t *sockjsTransport) Unidirectional() bool { - return false -} - -// Emulation ... -func (t *sockjsTransport) Emulation() bool { - return false -} - -// DisabledPushFlags ... -func (t *sockjsTransport) DisabledPushFlags() uint64 { - // SockJS has its own close frames to mimic WebSocket Close frames, - // so we don't need to send Disconnect pushes. - return PushFlagDisconnect -} - -// PingPongConfig ... -func (t *sockjsTransport) PingPongConfig() PingPongConfig { - return t.opts.pingPong -} - -// Write data to transport. -func (t *sockjsTransport) Write(message []byte) error { - select { - case <-t.closeCh: - return nil - default: - // No need to use protocol encoders here since - // SockJS only supports JSON. - return t.session.Send(string(message)) - } -} - -// WriteMany messages to transport. -func (t *sockjsTransport) WriteMany(messages ...[]byte) error { - select { - case <-t.closeCh: - return nil - default: - encoder := protocol.GetDataEncoder(ProtocolTypeJSON.toProto()) - defer protocol.PutDataEncoder(ProtocolTypeJSON.toProto(), encoder) - for i := range messages { - _ = encoder.Encode(messages[i]) - } - return t.session.Send(string(encoder.Finish())) - } -} - -// Close closes transport. -func (t *sockjsTransport) Close(disconnect Disconnect) error { - t.mu.Lock() - if t.closed { - // Already closed, noop. - t.mu.Unlock() - return nil - } - t.closed = true - close(t.closeCh) - t.mu.Unlock() - return t.session.Close(disconnect.Code, disconnect.Reason) -} diff --git a/handler_sockjs_test.go b/handler_sockjs_test.go deleted file mode 100644 index 85cd3c12..00000000 --- a/handler_sockjs_test.go +++ /dev/null @@ -1,253 +0,0 @@ -package centrifuge - -import ( - "context" - "encoding/json" - "fmt" - "net/http" - "net/http/httptest" - "strings" - "testing" - - "github.com/centrifugal/protocol" - "github.com/gorilla/websocket" - "github.com/stretchr/testify/require" -) - -func sockjsData(data []byte) []byte { - quoted, _ := json.Marshal(string(data)) - return []byte(fmt.Sprintf("[%s]", string(quoted))) -} - -func TestSockjsHandler(t *testing.T) { - n, _ := New(Config{}) - require.NoError(t, n.Run()) - defer func() { _ = n.Shutdown(context.Background()) }() - mux := http.NewServeMux() - - n.OnConnecting(func(ctx context.Context, event ConnectEvent) (ConnectReply, error) { - require.Equal(t, transportSockJS, event.Transport.Name()) - require.Equal(t, ProtocolTypeJSON, event.Transport.Protocol()) - return ConnectReply{ - Credentials: &Credentials{UserID: "user"}, - Data: []byte(`{"SockJS connect response": 1}`), - }, nil - }) - - doneCh := make(chan struct{}) - - n.OnConnect(func(client *Client) { - err := client.Send([]byte(`{"SockJS write": 1}`)) - require.NoError(t, err) - client.Disconnect(DisconnectForceReconnect) - }) - - mux.Handle("/connection/sockjs/", NewSockjsHandler(n, SockjsConfig{ - HandlerPrefix: "/connection/sockjs", - })) - server := httptest.NewServer(mux) - defer server.Close() - - url := "ws" + server.URL[4:] - - conn, resp, err := websocket.DefaultDialer.Dial(url+"/connection/sockjs/220/fi0988475/websocket", nil) - require.NoError(t, err) - defer func() { _ = resp.Body.Close() }() - require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) - require.NotNil(t, conn) - defer func() { _ = conn.Close() }() - _, p, err := conn.ReadMessage() - require.NoError(t, err) - require.Equal(t, "o", string(p)) // open frame of SockJS protocol. - - connectRequest := &protocol.ConnectRequest{ - Token: "boom", - } - cmd := &protocol.Command{ - Id: 1, - Connect: connectRequest, - } - cmdBytes, _ := json.Marshal(cmd) - err = conn.WriteMessage(websocket.TextMessage, sockjsData(cmdBytes)) - require.NoError(t, err) - - go func() { - pos := 0 - contentExpected := []string{ - "SockJS connect response", - "SockJS write", - "force reconnect", - } - - loop: - for { - _, p, err = conn.ReadMessage() - if err != nil { - break loop - } - - for { - if strings.Contains(string(p), contentExpected[pos]) { - pos++ - if pos >= len(contentExpected) { - close(doneCh) - break loop - } - } else { - break - } - } - } - }() - - waitWithTimeout(t, doneCh) -} - -func TestSockjsTransportWrite(t *testing.T) { - node := defaultNodeNoHandlers() - defer func() { _ = node.Shutdown(context.Background()) }() - - node.OnConnecting(func(ctx context.Context, event ConnectEvent) (ConnectReply, error) { - require.Equal(t, event.Transport.Protocol(), ProtocolTypeJSON) - transport := event.Transport.(Transport) - // Write to transport directly - this is only valid for tests, in normal situation - // we write over client methods. - require.NoError(t, transport.Write([]byte("hello"))) - return ConnectReply{}, DisconnectForceNoReconnect - }) - - mux := http.NewServeMux() - mux.Handle("/connection/sockjs/", NewSockjsHandler(node, SockjsConfig{ - HandlerPrefix: "/connection/sockjs", - })) - server := httptest.NewServer(mux) - defer server.Close() - - url := "ws" + server.URL[4:] - - conn, resp, err := websocket.DefaultDialer.Dial(url+"/connection/sockjs/220/fi0988475/websocket", nil) - require.NoError(t, err) - defer func() { _ = resp.Body.Close() }() - require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) - require.NotNil(t, conn) - defer func() { _ = conn.Close() }() - _, p, err := conn.ReadMessage() - require.NoError(t, err) - require.Equal(t, "o", string(p)) // open frame of SockJS protocol. - - connectRequest := &protocol.ConnectRequest{ - Token: "boom", - } - cmd := &protocol.Command{ - Id: 1, - Connect: connectRequest, - } - cmdBytes, _ := json.Marshal(cmd) - err = conn.WriteMessage(websocket.TextMessage, sockjsData(cmdBytes)) - require.NoError(t, err) - - _, p, err = conn.ReadMessage() - require.NoError(t, err) - require.Equal(t, "a[\"hello\"]", string(p)) -} - -func TestSockjsTransportWriteMany(t *testing.T) { - node := defaultNodeNoHandlers() - defer func() { _ = node.Shutdown(context.Background()) }() - - node.OnConnecting(func(ctx context.Context, event ConnectEvent) (ConnectReply, error) { - require.Equal(t, event.Transport.Protocol(), ProtocolTypeJSON) - transport := event.Transport.(Transport) - // Write to transport directly - this is only valid for tests, in normal situation - // we write over client methods. - require.NoError(t, transport.WriteMany([]byte("1"), []byte("22"))) - return ConnectReply{}, DisconnectForceNoReconnect - }) - - mux := http.NewServeMux() - mux.Handle("/connection/sockjs/", NewSockjsHandler(node, SockjsConfig{ - HandlerPrefix: "/connection/sockjs", - })) - server := httptest.NewServer(mux) - defer server.Close() - - url := "ws" + server.URL[4:] - - conn, resp, err := websocket.DefaultDialer.Dial(url+"/connection/sockjs/220/fi0988475/websocket", nil) - require.NoError(t, err) - defer func() { _ = resp.Body.Close() }() - require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) - require.NotNil(t, conn) - defer func() { _ = conn.Close() }() - _, p, err := conn.ReadMessage() - require.NoError(t, err) - require.Equal(t, "o", string(p)) // open frame of SockJS protocol. - - connectRequest := &protocol.ConnectRequest{ - Token: "boom", - } - cmd := &protocol.Command{ - Id: 1, - Connect: connectRequest, - } - cmdBytes, _ := json.Marshal(cmd) - err = conn.WriteMessage(websocket.TextMessage, sockjsData(cmdBytes)) - require.NoError(t, err) - - _, p, err = conn.ReadMessage() - require.NoError(t, err) - require.Equal(t, "a[\"1\\n22\"]", string(p)) -} - -func TestSockjsHandlerURLParams(t *testing.T) { - n, _ := New(Config{}) - require.NoError(t, n.Run()) - defer func() { _ = n.Shutdown(context.Background()) }() - mux := http.NewServeMux() - - n.OnConnecting(func(ctx context.Context, event ConnectEvent) (ConnectReply, error) { - return ConnectReply{ - Credentials: &Credentials{UserID: "user"}, - }, nil - }) - - doneCh := make(chan struct{}) - - n.OnConnect(func(client *Client) { - require.Equal(t, transportSockJS, client.Transport().Name()) - require.Equal(t, ProtocolTypeJSON, client.Transport().Protocol()) - require.Equal(t, ProtocolVersion2, client.Transport().ProtocolVersion()) - close(doneCh) - }) - - mux.Handle("/connection/sockjs/", NewSockjsHandler(n, SockjsConfig{ - HandlerPrefix: "/connection/sockjs", - })) - server := httptest.NewServer(mux) - defer server.Close() - - url := "ws" + server.URL[4:] - - conn, resp, err := websocket.DefaultDialer.Dial(url+"/connection/sockjs/220/fi0988475/websocket?cf_protocol_version=v2", nil) - require.NoError(t, err) - defer func() { _ = resp.Body.Close() }() - require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) - require.NotNil(t, conn) - defer func() { _ = conn.Close() }() - _, p, err := conn.ReadMessage() - require.NoError(t, err) - require.Equal(t, "o", string(p)) // open frame of SockJS protocol. - - connectRequest := &protocol.ConnectRequest{ - Token: "boom", - } - cmd := &protocol.Command{ - Id: 1, - Connect: connectRequest, - } - cmdBytes, _ := json.Marshal(cmd) - err = conn.WriteMessage(websocket.TextMessage, sockjsData(cmdBytes)) - require.NoError(t, err) - - waitWithTimeout(t, doneCh) -} diff --git a/handler_websocket.go b/handler_websocket.go index b94d6f37..53efc354 100644 --- a/handler_websocket.go +++ b/handler_websocket.go @@ -11,9 +11,9 @@ import ( "github.com/centrifugal/centrifuge/internal/cancelctx" "github.com/centrifugal/centrifuge/internal/timers" + "github.com/centrifugal/centrifuge/internal/websocket" "github.com/centrifugal/protocol" - "github.com/gorilla/websocket" ) // WebsocketConfig represents config for WebsocketHandler. @@ -122,7 +122,7 @@ func (s *WebsocketHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) { compressionLevel := s.config.CompressionLevel compressionMinSize := s.config.CompressionMinSize - conn, err := s.upgrade.Upgrade(rw, r, nil) + conn, subProtocol, err := s.upgrade.Upgrade(rw, r, nil) if err != nil { s.node.logger.log(newLogEntry(LogLevelDebug, "websocket upgrade error", map[string]any{"error": err.Error()})) return @@ -147,7 +147,6 @@ func (s *WebsocketHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) { conn.SetReadLimit(int64(messageSizeLimit)) } - subProtocol := conn.Subprotocol() if subProtocol == "centrifuge-protobuf" { protoType = ProtocolTypeProtobuf } diff --git a/handler_websocket_test.go b/handler_websocket_test.go index 5469e91e..3ca53341 100644 --- a/handler_websocket_test.go +++ b/handler_websocket_test.go @@ -14,8 +14,9 @@ import ( "testing" "time" + "github.com/centrifugal/centrifuge/internal/websocket" + "github.com/centrifugal/protocol" - "github.com/gorilla/websocket" "github.com/stretchr/testify/require" ) @@ -33,8 +34,10 @@ func TestWebsocketHandler(t *testing.T) { server := httptest.NewServer(mux) defer server.Close() + dialer := &websocket.Dialer{} + url := "ws" + server.URL[4:] - conn, resp, err := websocket.DefaultDialer.Dial(url+"/connection/websocket", nil) + conn, resp, _, err := dialer.Dial(url+"/connection/websocket", nil) require.NoError(t, err) defer func() { _ = resp.Body.Close() }() require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) @@ -57,8 +60,9 @@ func TestWebsocketHandlerProtocolV2(t *testing.T) { server := httptest.NewServer(mux) defer server.Close() + dialer := &websocket.Dialer{} url := "ws" + server.URL[4:] - conn, resp, err := websocket.DefaultDialer.Dial(url+"/connection/websocket?cf_protocol_version=v2", nil) + conn, resp, _, err := dialer.Dial(url+"/connection/websocket?cf_protocol_version=v2", nil) require.NoError(t, err) defer func() { _ = resp.Body.Close() }() require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) @@ -83,13 +87,10 @@ func TestWebsocketHandlerSubprotocol(t *testing.T) { server := httptest.NewServer(mux) defer server.Close() + dialer := &websocket.Dialer{} url := "ws" + server.URL[4:] - dialer := &websocket.Dialer{ - Proxy: http.ProxyFromEnvironment, - HandshakeTimeout: 45 * time.Second, - } dialer.Subprotocols = []string{"centrifuge-protobuf"} - conn, resp, err := dialer.Dial(url+"/connection/websocket", nil) + conn, resp, _, err := dialer.Dial(url+"/connection/websocket", nil) require.NoError(t, err) defer func() { _ = resp.Body.Close() }() require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) @@ -119,12 +120,9 @@ func TestWebsocketHandlerURLParams(t *testing.T) { defer server.Close() url := "ws" + server.URL[4:] - dialer := &websocket.Dialer{ - Proxy: http.ProxyFromEnvironment, - HandshakeTimeout: 45 * time.Second, - } + dialer := &websocket.Dialer{} - conn, resp, err := dialer.Dial(url+"/connection/websocket?cf_protocol=protobuf&cf_protocol_version=v1", nil) + conn, resp, _, err := dialer.Dial(url+"/connection/websocket?cf_protocol=protobuf&cf_protocol_version=v1", nil) require.NoError(t, err) defer func() { _ = resp.Body.Close() }() require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) @@ -155,11 +153,9 @@ func TestWebsocketTransportWrite(t *testing.T) { url := "ws" + server.URL[4:] dialer := &websocket.Dialer{ - Proxy: http.ProxyFromEnvironment, - HandshakeTimeout: 45 * time.Second, + Subprotocols: []string{"centrifuge-protobuf"}, } - dialer.Subprotocols = []string{"centrifuge-protobuf"} - conn, resp, err := dialer.Dial(url+"/connection/websocket", nil) + conn, resp, subprotocol, err := dialer.Dial(url+"/connection/websocket", nil) require.NoError(t, err) defer func() { _ = resp.Body.Close() }() require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) @@ -167,6 +163,7 @@ func TestWebsocketTransportWrite(t *testing.T) { defer func() { _ = conn.Close() }() err = conn.WriteMessage(websocket.BinaryMessage, getConnectCommandProtobuf(t)) require.NoError(t, err) + require.Equal(t, "centrifuge-protobuf", subprotocol) msgType, msg, err := conn.ReadMessage() require.NoError(t, err) @@ -195,11 +192,9 @@ func TestWebsocketTransportWriteMany(t *testing.T) { url := "ws" + server.URL[4:] dialer := &websocket.Dialer{ - Proxy: http.ProxyFromEnvironment, - HandshakeTimeout: 45 * time.Second, + Subprotocols: []string{"centrifuge-protobuf"}, } - dialer.Subprotocols = []string{"centrifuge-protobuf"} - conn, resp, err := dialer.Dial(url+"/connection/websocket", nil) + conn, resp, _, err := dialer.Dial(url+"/connection/websocket", nil) require.NoError(t, err) defer func() { _ = resp.Body.Close() }() require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) @@ -251,8 +246,9 @@ func TestWebsocketHandlerProtobuf(t *testing.T) { server := httptest.NewServer(mux) defer server.Close() + dialer := &websocket.Dialer{} url := "ws" + server.URL[4:] - conn, resp, err := websocket.DefaultDialer.Dial(url+"/connection/websocket?format=protobuf", nil) + conn, resp, _, err := dialer.Dial(url+"/connection/websocket?format=protobuf", nil) require.NoError(t, err) defer func() { _ = resp.Body.Close() }() require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) @@ -283,7 +279,8 @@ func TestWebsocketHandlerPing(t *testing.T) { url := "ws" + server.URL[4:] - conn, resp, err := websocket.DefaultDialer.Dial(url+"/connection/websocket", nil) + dialer := &websocket.Dialer{} + conn, resp, _, err := dialer.Dial(url+"/connection/websocket", nil) require.NoError(t, err) defer func() { _ = resp.Body.Close() }() require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) @@ -336,7 +333,8 @@ func TestWebsocketHandler_FramePingPong(t *testing.T) { url := "ws" + server.URL[4:] - conn, resp, err := websocket.DefaultDialer.Dial(url+"/connection/websocket?cf_ws_frame_ping_pong=true", nil) + dialer := &websocket.Dialer{} + conn, resp, _, err := dialer.Dial(url+"/connection/websocket?cf_ws_frame_ping_pong=true", nil) require.NoError(t, err) defer func() { _ = resp.Body.Close() }() require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) @@ -382,8 +380,8 @@ func TestWebsocketHandlerCustomDisconnect(t *testing.T) { defer server.Close() url := "ws" + server.URL[4:] - - conn, resp, err := websocket.DefaultDialer.Dial(url+"/connection/websocket", nil) + dialer := &websocket.Dialer{} + conn, resp, _, err := dialer.Dial(url+"/connection/websocket", nil) require.NoError(t, err) defer func() { _ = resp.Body.Close() }() @@ -672,7 +670,8 @@ func BenchmarkWsConnectV2(b *testing.B) { } func newRealConnJSONConnectV2(b testing.TB, url string) *websocket.Conn { - conn, resp, err := websocket.DefaultDialer.Dial(url+"/connection/websocket", nil) + dialer := &websocket.Dialer{} + conn, resp, _, err := dialer.Dial(url+"/connection/websocket", nil) require.NoError(b, err) defer func() { _ = resp.Body.Close() }() @@ -689,7 +688,8 @@ func newRealConnJSONConnectV2(b testing.TB, url string) *websocket.Conn { } func newRealConnProtobufConnectV2(b testing.TB, url string) *websocket.Conn { - conn, resp, err := websocket.DefaultDialer.Dial(url+"/connection/websocket?format=protobuf", nil) + dialer := &websocket.Dialer{} + conn, resp, _, err := dialer.Dial(url+"/connection/websocket?format=protobuf", nil) require.NoError(b, err) defer func() { _ = resp.Body.Close() }() diff --git a/internal/websocket/AUTHORS b/internal/websocket/AUTHORS new file mode 100644 index 00000000..124965e6 --- /dev/null +++ b/internal/websocket/AUTHORS @@ -0,0 +1,8 @@ +# This is the official list of Gorilla WebSocket authors for copyright +# purposes. +# +# Please keep the list sorted. + +Gary Burd +Google LLC (https://opensource.google.com/) +Joachim Bauch diff --git a/internal/websocket/LICENSE b/internal/websocket/LICENSE new file mode 100644 index 00000000..9171c972 --- /dev/null +++ b/internal/websocket/LICENSE @@ -0,0 +1,22 @@ +Copyright (c) 2013 The Gorilla WebSocket Authors. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + + Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/internal/websocket/README.md b/internal/websocket/README.md new file mode 100644 index 00000000..bfe16c53 --- /dev/null +++ b/internal/websocket/README.md @@ -0,0 +1,8 @@ +Fork of Gorilla Websocket adapted for Centrifuge needs. All copyrights were left – huge respect to the original authors. + +What was changed: + +* no custom Proxy in Client +* no concurrent use best-effort detection +* selected subprotocol is not kept internally, returned during Upgrade or Dial +* lint fixes diff --git a/internal/websocket/client.go b/internal/websocket/client.go new file mode 100644 index 00000000..23ba6c1a --- /dev/null +++ b/internal/websocket/client.go @@ -0,0 +1,383 @@ +// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package websocket + +import ( + "bytes" + "context" + "crypto/tls" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/http/httptrace" + "net/url" + "strings" + "time" +) + +// ErrBadHandshake is returned when the server response to opening handshake is +// invalid. +var ErrBadHandshake = errors.New("websocket: bad handshake") + +var errInvalidCompression = errors.New("websocket: invalid compression negotiation") + +// A Dialer contains options for connecting to WebSocket server. +// +// It is safe to call Dialer's methods concurrently. +type Dialer struct { + // NetDial specifies the dial function for creating TCP connections. If + // NetDial is nil, net.Dial is used. + NetDial func(network, addr string) (net.Conn, error) + + // NetDialContext specifies the dial function for creating TCP connections. If + // NetDialContext is nil, NetDial is used. + NetDialContext func(ctx context.Context, network, addr string) (net.Conn, error) + + // NetDialTLSContext specifies the dial function for creating TLS/TCP connections. If + // NetDialTLSContext is nil, NetDialContext is used. + // If NetDialTLSContext is set, Dial assumes the TLS handshake is done there and + // TLSClientConfig is ignored. + NetDialTLSContext func(ctx context.Context, network, addr string) (net.Conn, error) + + // TLSClientConfig specifies the TLS configuration to use with tls.Client. + // If nil, the default configuration is used. + // If either NetDialTLS or NetDialTLSContext are set, Dial assumes the TLS handshake + // is done there and TLSClientConfig is ignored. + TLSClientConfig *tls.Config + + // HandshakeTimeout specifies the duration for the handshake to complete. + HandshakeTimeout time.Duration + + // ReadBufferSize and WriteBufferSize specify I/O buffer sizes in bytes. If a buffer + // size is zero, then a useful default size is used. The I/O buffer sizes + // do not limit the size of the messages that can be sent or received. + ReadBufferSize, WriteBufferSize int + + // WriteBufferPool is a pool of buffers for write operations. If the value + // is not set, then write buffers are allocated to the connection for the + // lifetime of the connection. + // + // A pool is most useful when the application has a modest volume of writes + // across a large number of connections. + // + // Applications should use a single pool for each unique value of + // WriteBufferSize. + WriteBufferPool BufferPool + + // Subprotocols specifies the client's requested subprotocols. + Subprotocols []string + + // EnableCompression specifies if the client should attempt to negotiate + // per message compression (RFC 7692). Setting this value to true does not + // guarantee that compression will be supported. Currently only "no context + // takeover" modes are supported. + EnableCompression bool + + // Jar specifies the cookie jar. + // If Jar is nil, cookies are not sent in requests and ignored + // in responses. + Jar http.CookieJar +} + +// Dial creates a new client connection by calling DialContext with a background context. +func (d *Dialer) Dial(urlStr string, requestHeader http.Header) (*Conn, *http.Response, string, error) { + return d.DialContext(context.Background(), urlStr, requestHeader) +} + +var errMalformedURL = errors.New("malformed ws or wss URL") + +func hostPortNoPort(u *url.URL) (hostPort, hostNoPort string) { + hostPort = u.Host + hostNoPort = u.Host + if i := strings.LastIndex(u.Host, ":"); i > strings.LastIndex(u.Host, "]") { + hostNoPort = hostNoPort[:i] + } else { + switch u.Scheme { + case "wss": + hostPort += ":443" + case "https": + hostPort += ":443" + default: + hostPort += ":80" + } + } + return hostPort, hostNoPort +} + +// DialContext creates a new client connection. Use requestHeader to specify the +// origin (Origin), subprotocols (Sec-WebSocket-Protocol) and cookies (Cookie). +// Use the response.Header to get the selected subprotocol +// (Sec-WebSocket-Protocol) and cookies (Set-Cookie). +// +// The context will be used in the request and in the Dialer. +// +// If the WebSocket handshake fails, ErrBadHandshake is returned along with a +// non-nil *http.Response so that callers can handle redirects, authentication, +// etcetera. The response body may not contain the entire response and does not +// need to be closed by the application. +func (d *Dialer) DialContext(ctx context.Context, urlStr string, requestHeader http.Header) (*Conn, *http.Response, string, error) { + if d.HandshakeTimeout == 0 { + d.HandshakeTimeout = 15 * time.Second + } + + challengeKey, err := generateChallengeKey() + if err != nil { + return nil, nil, "", err + } + + u, err := url.Parse(urlStr) + if err != nil { + return nil, nil, "", err + } + + switch u.Scheme { + case "ws": + u.Scheme = "http" + case "wss": + u.Scheme = "https" + default: + return nil, nil, "", errMalformedURL + } + + if u.User != nil { + // User name and password are not allowed in websocket URIs. + return nil, nil, "", errMalformedURL + } + + req := &http.Request{ + Method: http.MethodGet, + URL: u, + Proto: "HTTP/1.1", + ProtoMajor: 1, + ProtoMinor: 1, + Header: make(http.Header), + Host: u.Host, + } + req = req.WithContext(ctx) + + // Set the cookies present in the cookie jar of the dialer + if d.Jar != nil { + for _, cookie := range d.Jar.Cookies(u) { + req.AddCookie(cookie) + } + } + + // Set the request headers using the capitalization for names and values in + // RFC examples. Although the capitalization shouldn't matter, there are + // servers that depend on it. The Header.Set method is not used because the + // method canonicalizes the header names. + req.Header["Upgrade"] = []string{"websocket"} + req.Header["Connection"] = []string{"Upgrade"} + req.Header["Sec-WebSocket-Key"] = []string{challengeKey} + req.Header["Sec-WebSocket-Version"] = []string{"13"} + if len(d.Subprotocols) > 0 { + req.Header["Sec-WebSocket-Protocol"] = []string{strings.Join(d.Subprotocols, ", ")} + } + for k, vs := range requestHeader { + switch { + case k == "Host": + if len(vs) > 0 { + req.Host = vs[0] + } + case k == "Upgrade" || + k == "Connection" || + k == "Sec-Websocket-Key" || + k == "Sec-Websocket-Version" || + k == "Sec-Websocket-Extensions" || + (k == "Sec-Websocket-Protocol" && len(d.Subprotocols) > 0): + return nil, nil, "", errors.New("websocket: duplicate header not allowed: " + k) + case k == "Sec-Websocket-Protocol": + req.Header["Sec-WebSocket-Protocol"] = vs + default: + req.Header[k] = vs + } + } + + if d.EnableCompression { + req.Header["Sec-WebSocket-Extensions"] = []string{"permessage-deflate; server_no_context_takeover; client_no_context_takeover"} + } + + if d.HandshakeTimeout != 0 { + var cancel func() + ctx, cancel = context.WithTimeout(ctx, d.HandshakeTimeout) + defer cancel() + } + + // Get network dial function. + var netDial func(network, add string) (net.Conn, error) + + switch u.Scheme { + case "http": + if d.NetDialContext != nil { + netDial = func(network, addr string) (net.Conn, error) { + return d.NetDialContext(ctx, network, addr) + } + } else if d.NetDial != nil { + netDial = d.NetDial + } + case "https": + if d.NetDialTLSContext != nil { + netDial = func(network, addr string) (net.Conn, error) { + return d.NetDialTLSContext(ctx, network, addr) + } + } else if d.NetDialContext != nil { + netDial = func(network, addr string) (net.Conn, error) { + return d.NetDialContext(ctx, network, addr) + } + } else if d.NetDial != nil { + netDial = d.NetDial + } + default: + return nil, nil, "", errMalformedURL + } + + if netDial == nil { + netDialer := &net.Dialer{} + netDial = func(network, addr string) (net.Conn, error) { + return netDialer.DialContext(ctx, network, addr) + } + } + + // If needed, wrap the dial function to set the connection deadline. + if deadline, ok := ctx.Deadline(); ok { + forwardDial := netDial + netDial = func(network, addr string) (net.Conn, error) { + c, err := forwardDial(network, addr) + if err != nil { + return nil, err + } + err = c.SetDeadline(deadline) + if err != nil { + _ = c.Close() + return nil, err + } + return c, nil + } + } + + hostPort, hostNoPort := hostPortNoPort(u) + trace := httptrace.ContextClientTrace(ctx) + if trace != nil && trace.GetConn != nil { + trace.GetConn(hostPort) + } + + netConn, err := netDial("tcp", hostPort) + if err != nil { + return nil, nil, "", err + } + if trace != nil && trace.GotConn != nil { + trace.GotConn(httptrace.GotConnInfo{ + Conn: netConn, + }) + } + + defer func() { + if netConn != nil { + _ = netConn.Close() + } + }() + + if u.Scheme == "https" && d.NetDialTLSContext == nil { + // If NetDialTLSContext is set, assume that the TLS handshake has already been done + + cfg := cloneTLSConfig(d.TLSClientConfig) + if cfg.ServerName == "" { + cfg.ServerName = hostNoPort + } + tlsConn := tls.Client(netConn, cfg) + netConn = tlsConn + + if trace != nil && trace.TLSHandshakeStart != nil { + trace.TLSHandshakeStart() + } + err := doHandshake(ctx, tlsConn, cfg) + if trace != nil && trace.TLSHandshakeDone != nil { + trace.TLSHandshakeDone(tlsConn.ConnectionState(), err) + } + + if err != nil { + return nil, nil, "", err + } + } + + conn := newConn(netConn, false, d.ReadBufferSize, d.WriteBufferSize, d.WriteBufferPool, nil, nil) + + if err := req.Write(netConn); err != nil { + return nil, nil, "", err + } + + if trace != nil && trace.GotFirstResponseByte != nil { + if peek, err := conn.br.Peek(1); err == nil && len(peek) == 1 { + trace.GotFirstResponseByte() + } + } + + resp, err := http.ReadResponse(conn.br, req) + if err != nil { + if d.TLSClientConfig != nil { + for _, proto := range d.TLSClientConfig.NextProtos { + if proto != "http/1.1" { + return nil, nil, "", fmt.Errorf( + "websocket: protocol %q was given but is not supported;"+ + "sharing tls.Config with net/http Transport can cause this error: %w", + proto, err, + ) + } + } + } + return nil, nil, "", err + } + + if d.Jar != nil { + if rc := resp.Cookies(); len(rc) > 0 { + d.Jar.SetCookies(u, rc) + } + } + + if resp.StatusCode != 101 || + !tokenListContainsValue(resp.Header, "Upgrade", "websocket") || + !tokenListContainsValue(resp.Header, "Connection", "upgrade") || + resp.Header.Get("Sec-Websocket-Accept") != computeAcceptKey(challengeKey) { + // Before closing the network connection on return from this + // function, slurp up some of the response to aid application + // debugging. + buf := make([]byte, 1024) + n, _ := io.ReadFull(resp.Body, buf) + resp.Body = io.NopCloser(bytes.NewReader(buf[:n])) + return nil, resp, "", ErrBadHandshake + } + + for _, ext := range parseExtensions(resp.Header) { + if ext[""] != "permessage-deflate" { + continue + } + _, snct := ext["server_no_context_takeover"] + _, cnct := ext["client_no_context_takeover"] + if !snct || !cnct { + return nil, resp, "", errInvalidCompression + } + conn.newCompressionWriter = compressNoContextTakeover + conn.newDecompressionReader = decompressNoContextTakeover + break + } + + resp.Body = io.NopCloser(bytes.NewReader([]byte{})) + + err = netConn.SetDeadline(time.Time{}) + if err != nil { + return conn, resp, "", err + } + netConn = nil // to avoid close in defer. + return conn, resp, resp.Header.Get("Sec-Websocket-Protocol"), nil +} + +func cloneTLSConfig(cfg *tls.Config) *tls.Config { + if cfg == nil { + return &tls.Config{} + } + return cfg.Clone() +} diff --git a/internal/websocket/client_server_test.go b/internal/websocket/client_server_test.go new file mode 100644 index 00000000..4588da1e --- /dev/null +++ b/internal/websocket/client_server_test.go @@ -0,0 +1,991 @@ +// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package websocket + +import ( + "context" + "crypto/tls" + "crypto/x509" + "errors" + "fmt" + "io" + "log" + "net" + "net/http" + "net/http/cookiejar" + "net/http/httptest" + "net/http/httptrace" + "net/url" + "reflect" + "strings" + "testing" + "time" +) + +var cstUpgrader = Upgrader{ + Subprotocols: []string{"p0", "p1"}, + ReadBufferSize: 1024, + WriteBufferSize: 1024, + EnableCompression: true, + Error: func(w http.ResponseWriter, r *http.Request, status int, reason error) { + http.Error(w, reason.Error(), status) + }, +} + +var cstDialer = Dialer{ + Subprotocols: []string{"p1", "p2"}, + ReadBufferSize: 1024, + WriteBufferSize: 1024, + HandshakeTimeout: 30 * time.Second, +} + +type cstHandler struct{ *testing.T } + +type cstServer struct { + *httptest.Server + URL string +} + +const ( + cstPath = "/a/b" + cstRawQuery = "x=y" + cstRequestURI = cstPath + "?" + cstRawQuery +) + +func newServer(t *testing.T) *cstServer { + var s cstServer + s.Server = httptest.NewServer(cstHandler{t}) + s.Server.URL += cstRequestURI + s.URL = makeWsProto(s.Server.URL) + return &s +} + +func newTLSServer(t *testing.T) *cstServer { + var s cstServer + s.Server = httptest.NewTLSServer(cstHandler{t}) + s.Server.URL += cstRequestURI + s.URL = makeWsProto(s.Server.URL) + return &s +} + +func (t cstHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != cstPath { + t.Logf("path=%v, want %v", r.URL.Path, cstPath) + http.Error(w, "bad path", http.StatusBadRequest) + return + } + if r.URL.RawQuery != cstRawQuery { + t.Logf("query=%v, want %v", r.URL.RawQuery, cstRawQuery) + http.Error(w, "bad path", http.StatusBadRequest) + return + } + subprotos := Subprotocols(r) + if !reflect.DeepEqual(subprotos, cstDialer.Subprotocols) { + t.Logf("subprotols=%v, want %v", subprotos, cstDialer.Subprotocols) + http.Error(w, "bad protocol", http.StatusBadRequest) + return + } + ws, subprotocol, err := cstUpgrader.Upgrade(w, r, http.Header{"Set-Cookie": {"sessionID=1234"}}) + if err != nil { + t.Logf("Upgrade: %v", err) + return + } + defer func() { + _ = ws.Close() + }() + + if subprotocol != "p1" { + t.Logf("Subprotocol() = %s, want p1", subprotocol) + _ = ws.Close() + return + } + op, rd, err := ws.NextReader() + if err != nil { + t.Logf("NextReader: %v", err) + return + } + wr, err := ws.NextWriter(op) + if err != nil { + t.Logf("NextWriter: %v", err) + return + } + if _, err = io.Copy(wr, rd); err != nil { + t.Logf("NextWriter: %v", err) + return + } + if err := wr.Close(); err != nil { + t.Logf("Close: %v", err) + return + } +} + +func makeWsProto(s string) string { + return "ws" + strings.TrimPrefix(s, "http") +} + +func sendRecv(t *testing.T, ws *Conn) { + const message = "Hello World!" + if err := ws.SetWriteDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatalf("SetWriteDeadline: %v", err) + } + if err := ws.WriteMessage(TextMessage, []byte(message)); err != nil { + t.Fatalf("WriteMessage: %v", err) + } + if err := ws.SetReadDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatalf("SetReadDeadline: %v", err) + } + _, p, err := ws.ReadMessage() + if err != nil { + t.Fatalf("ReadMessage: %v", err) + } + if string(p) != message { + t.Fatalf("message=%s, want %s", p, message) + } +} + +func TestDial(t *testing.T) { + s := newServer(t) + defer s.Close() + + ws, _, _, err := cstDialer.Dial(s.URL, nil) + if err != nil { + t.Fatalf("Dial: %v", err) + } + defer func() { + _ = ws.Close() + }() + sendRecv(t, ws) +} + +func TestDialCookieJar(t *testing.T) { + s := newServer(t) + defer s.Close() + + jar, _ := cookiejar.New(nil) + d := cstDialer + d.Jar = jar + + u, _ := url.Parse(s.URL) + + switch u.Scheme { + case "ws": + u.Scheme = "http" + case "wss": + u.Scheme = "https" + } + + cookies := []*http.Cookie{{Name: "gorilla", Value: "ws", Path: "/"}} + d.Jar.SetCookies(u, cookies) + + ws, _, _, err := d.Dial(s.URL, nil) + if err != nil { + t.Fatalf("Dial: %v", err) + } + defer func() { + _ = ws.Close() + }() + + var gorilla string + var sessionID string + for _, c := range d.Jar.Cookies(u) { + if c.Name == "gorilla" { + gorilla = c.Value + } + + if c.Name == "sessionID" { + sessionID = c.Value + } + } + if gorilla != "ws" { + t.Error("Cookie not present in jar.") + } + + if sessionID != "1234" { + t.Error("Set-Cookie not received from the server.") + } + + sendRecv(t, ws) +} + +func rootCAs(t *testing.T, s *httptest.Server) *x509.CertPool { + certs := x509.NewCertPool() + for _, c := range s.TLS.Certificates { + roots, err := x509.ParseCertificates(c.Certificate[len(c.Certificate)-1]) + if err != nil { + t.Fatalf("error parsing server's root cert: %v", err) + } + for _, root := range roots { + certs.AddCert(root) + } + } + return certs +} + +func TestDialTLS(t *testing.T) { + s := newTLSServer(t) + defer s.Close() + + d := cstDialer + d.TLSClientConfig = &tls.Config{RootCAs: rootCAs(t, s.Server)} + ws, _, _, err := d.Dial(s.URL, nil) + if err != nil { + t.Fatalf("Dial: %v", err) + } + defer func() { + _ = ws.Close() + }() + sendRecv(t, ws) +} + +func TestDialTimeout(t *testing.T) { + s := newServer(t) + defer s.Close() + + d := cstDialer + d.HandshakeTimeout = -1 + ws, _, _, err := d.Dial(s.URL, nil) + if err == nil { + _ = ws.Close() + t.Fatalf("Dial: nil") + } +} + +// requireDeadlineNetConn fails the current test when Read or Write are called +// with no deadline. +type requireDeadlineNetConn struct { + t *testing.T + c net.Conn + readDeadlineIsSet bool + writeDeadlineIsSet bool +} + +func (c *requireDeadlineNetConn) SetDeadline(t time.Time) error { + c.writeDeadlineIsSet = !t.Equal(time.Time{}) + c.readDeadlineIsSet = c.writeDeadlineIsSet + return c.c.SetDeadline(t) +} + +func (c *requireDeadlineNetConn) SetReadDeadline(t time.Time) error { + c.readDeadlineIsSet = !t.Equal(time.Time{}) + return c.c.SetDeadline(t) +} + +func (c *requireDeadlineNetConn) SetWriteDeadline(t time.Time) error { + c.writeDeadlineIsSet = !t.Equal(time.Time{}) + return c.c.SetDeadline(t) +} + +func (c *requireDeadlineNetConn) Write(p []byte) (int, error) { + if !c.writeDeadlineIsSet { + c.t.Fatalf("write with no deadline") + } + return c.c.Write(p) +} + +func (c *requireDeadlineNetConn) Read(p []byte) (int, error) { + if !c.readDeadlineIsSet { + c.t.Fatalf("read with no deadline") + } + return c.c.Read(p) +} + +func (c *requireDeadlineNetConn) Close() error { return c.c.Close() } +func (c *requireDeadlineNetConn) LocalAddr() net.Addr { return c.c.LocalAddr() } +func (c *requireDeadlineNetConn) RemoteAddr() net.Addr { return c.c.RemoteAddr() } + +func TestHandshakeTimeout(t *testing.T) { + s := newServer(t) + defer s.Close() + + d := cstDialer + d.NetDial = func(n, a string) (net.Conn, error) { + c, err := net.Dial(n, a) + return &requireDeadlineNetConn{c: c, t: t}, err + } + ws, _, _, err := d.Dial(s.URL, nil) + if err != nil { + t.Fatal("Dial:", err) + } + _ = ws.Close() +} + +func TestHandshakeTimeoutInContext(t *testing.T) { + s := newServer(t) + defer s.Close() + + d := cstDialer + d.HandshakeTimeout = 0 + d.NetDialContext = func(ctx context.Context, n, a string) (net.Conn, error) { + netDialer := &net.Dialer{} + c, err := netDialer.DialContext(ctx, n, a) + return &requireDeadlineNetConn{c: c, t: t}, err + } + + ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(30*time.Second)) + defer cancel() + ws, _, _, err := d.DialContext(ctx, s.URL, nil) + if err != nil { + t.Fatal("Dial:", err) + } + _ = ws.Close() +} + +func TestDialBadScheme(t *testing.T) { + s := newServer(t) + defer s.Close() + + ws, _, _, err := cstDialer.Dial(s.Server.URL, nil) + if err == nil { + _ = ws.Close() + t.Fatalf("Dial: nil") + } +} + +func TestDialBadOrigin(t *testing.T) { + s := newServer(t) + defer s.Close() + + ws, resp, _, err := cstDialer.Dial(s.URL, http.Header{"Origin": {"bad"}}) + if err == nil { + _ = ws.Close() + t.Fatalf("Dial: nil") + } + if resp == nil { + t.Fatalf("resp=nil, err=%v", err) + } + if resp.StatusCode != http.StatusForbidden { + t.Fatalf("status=%d, want %d", resp.StatusCode, http.StatusForbidden) + } +} + +func TestDialBadHeader(t *testing.T) { + s := newServer(t) + defer s.Close() + + for _, k := range []string{"Upgrade", + "Connection", + "Sec-Websocket-Key", + "Sec-Websocket-Version", + "Sec-Websocket-Protocol"} { + h := http.Header{} + h.Set(k, "bad") + ws, _, _, err := cstDialer.Dial(s.URL, http.Header{"Origin": {"bad"}}) + if err == nil { + _ = ws.Close() + t.Errorf("Dial with header %s returned nil", k) + } + } +} + +func TestBadMethod(t *testing.T) { + s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ws, _, err := cstUpgrader.Upgrade(w, r, nil) + if err == nil { + t.Errorf("handshake succeeded, expect fail") + _ = ws.Close() + } + })) + defer s.Close() + + req, err := http.NewRequest(http.MethodPost, s.URL, strings.NewReader("")) + if err != nil { + t.Fatalf("NewRequest returned error %v", err) + } + req.Header.Set("Connection", "upgrade") + req.Header.Set("Upgrade", "websocket") + req.Header.Set("Sec-Websocket-Version", "13") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("Do returned error %v", err) + } + _ = resp.Body.Close() + if resp.StatusCode != http.StatusMethodNotAllowed { + t.Errorf("Status = %d, want %d", resp.StatusCode, http.StatusMethodNotAllowed) + } +} + +func TestDialExtraTokensInRespHeaders(t *testing.T) { + s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + challengeKey := r.Header.Get("Sec-Websocket-Key") + w.Header().Set("Upgrade", "foo, websocket") + w.Header().Set("Connection", "upgrade, keep-alive") + w.Header().Set("Sec-Websocket-Accept", computeAcceptKey(challengeKey)) + w.WriteHeader(101) + })) + defer s.Close() + + ws, _, _, err := cstDialer.Dial(makeWsProto(s.URL), nil) + if err != nil { + t.Fatalf("Dial: %v", err) + } + defer func() { + _ = ws.Close() + }() +} + +func TestHandshake(t *testing.T) { + s := newServer(t) + defer s.Close() + + ws, resp, subprotocol, err := cstDialer.Dial(s.URL, http.Header{"Origin": {s.URL}}) + if err != nil { + t.Fatalf("Dial: %v", err) + } + defer func() { + _ = ws.Close() + }() + + var sessionID string + for _, c := range resp.Cookies() { + if c.Name == "sessionID" { + sessionID = c.Value + } + } + if sessionID != "1234" { + t.Error("Set-Cookie not received from the server.") + } + + if subprotocol != "p1" { + t.Errorf("ws.Subprotocol() = %s, want p1", subprotocol) + } + sendRecv(t, ws) +} + +func TestRespOnBadHandshake(t *testing.T) { + const expectedStatus = http.StatusGone + const expectedBody = "This is the response body." + + s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(expectedStatus) + _, _ = io.WriteString(w, expectedBody) + })) + defer s.Close() + + ws, resp, _, err := cstDialer.Dial(makeWsProto(s.URL), nil) + if err == nil { + _ = ws.Close() + t.Fatalf("Dial: nil") + } + + if resp == nil { + t.Fatalf("resp=nil, err=%v", err) + } + + if resp.StatusCode != expectedStatus { + t.Errorf("resp.StatusCode=%d, want %d", resp.StatusCode, expectedStatus) + } + + p, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("ReadFull(resp.Body) returned error %v", err) + } + + if string(p) != expectedBody { + t.Errorf("resp.Body=%s, want %s", p, expectedBody) + } +} + +type testLogWriter struct { + t *testing.T +} + +func (w testLogWriter) Write(p []byte) (int, error) { + w.t.Logf("%s", p) + return len(p), nil +} + +// TestHost tests handling of host names and confirms that it matches net/http. +func TestHost(t *testing.T) { + + upgrader := Upgrader{} + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if IsWebSocketUpgrade(r) { + c, _, err := upgrader.Upgrade(w, r, http.Header{"X-Test-Host": {r.Host}}) + if err != nil { + t.Fatal(err) + } + _ = c.Close() + } else { + w.Header().Set("X-Test-Host", r.Host) + } + }) + + server := httptest.NewServer(handler) + defer server.Close() + + tlsServer := httptest.NewTLSServer(handler) + defer tlsServer.Close() + + addrs := map[*httptest.Server]string{server: server.Listener.Addr().String(), tlsServer: tlsServer.Listener.Addr().String()} + wsProtos := map[*httptest.Server]string{server: "ws://", tlsServer: "wss://"} + httpProtos := map[*httptest.Server]string{server: "http://", tlsServer: "https://"} + + // Avoid log noise from net/http server by logging to testing.T + server.Config.ErrorLog = log.New(testLogWriter{t}, "", 0) + tlsServer.Config.ErrorLog = server.Config.ErrorLog + + cas := rootCAs(t, tlsServer) + + tests := []struct { + fail bool // true if dial / get should fail + server *httptest.Server // server to use + url string // host for request URI + header string // optional request host header + tls string // optional host for tls ServerName + wantAddr string // expected host for dial + wantHeader string // expected request header on server + insecureSkipVerify bool + }{ + { + server: server, + url: addrs[server], + wantAddr: addrs[server], + wantHeader: addrs[server], + }, + { + server: tlsServer, + url: addrs[tlsServer], + wantAddr: addrs[tlsServer], + wantHeader: addrs[tlsServer], + }, + + { + server: server, + url: addrs[server], + header: "badhost.com", + wantAddr: addrs[server], + wantHeader: "badhost.com", + }, + { + server: tlsServer, + url: addrs[tlsServer], + header: "badhost.com", + wantAddr: addrs[tlsServer], + wantHeader: "badhost.com", + }, + + { + server: server, + url: "example.com", + header: "badhost.com", + wantAddr: "example.com:80", + wantHeader: "badhost.com", + }, + { + server: tlsServer, + url: "example.com", + header: "badhost.com", + wantAddr: "example.com:443", + wantHeader: "badhost.com", + }, + + { + server: server, + url: "badhost.com", + header: "example.com", + wantAddr: "badhost.com:80", + wantHeader: "example.com", + }, + { + fail: true, + server: tlsServer, + url: "badhost.com", + header: "example.com", + wantAddr: "badhost.com:443", + }, + { + server: tlsServer, + url: "badhost.com", + insecureSkipVerify: true, + wantAddr: "badhost.com:443", + wantHeader: "badhost.com", + }, + { + server: tlsServer, + url: "badhost.com", + tls: "example.com", + wantAddr: "badhost.com:443", + wantHeader: "badhost.com", + }, + } + + for i, tt := range tests { + + tlsConfig := &tls.Config{ + RootCAs: cas, + ServerName: tt.tls, + InsecureSkipVerify: tt.insecureSkipVerify, + } + + var gotAddr string + dialer := Dialer{ + NetDial: func(network, addr string) (net.Conn, error) { + gotAddr = addr + return net.Dial(network, addrs[tt.server]) + }, + TLSClientConfig: tlsConfig, + } + + // Test websocket dial + + h := http.Header{} + if tt.header != "" { + h.Set("Host", tt.header) + } + c, resp, _, err := dialer.Dial(wsProtos[tt.server]+tt.url+"/", h) + if err == nil { + _ = c.Close() + } + + check := func(protos map[*httptest.Server]string) { + name := fmt.Sprintf("%d: %s%s/ header[Host]=%q, tls.ServerName=%q", i+1, protos[tt.server], tt.url, tt.header, tt.tls) + if gotAddr != tt.wantAddr { + t.Errorf("%s: got addr %s, want %s", name, gotAddr, tt.wantAddr) + } + switch { + case tt.fail && err == nil: + t.Errorf("%s: unexpected success", name) + case !tt.fail && err != nil: + t.Errorf("%s: unexpected error %v", name, err) + case !tt.fail && err == nil: + if gotHost := resp.Header.Get("X-Test-Host"); gotHost != tt.wantHeader { + t.Errorf("%s: got host %s, want %s", name, gotHost, tt.wantHeader) + } + } + } + + check(wsProtos) + + // Confirm that net/http has same result + + transport := &http.Transport{ + Dial: dialer.NetDial, + TLSClientConfig: dialer.TLSClientConfig, + } + req, _ := http.NewRequest(http.MethodGet, httpProtos[tt.server]+tt.url+"/", nil) + if tt.header != "" { + req.Host = tt.header + } + client := &http.Client{Transport: transport} + resp, err = client.Do(req) + if err == nil { + _ = resp.Body.Close() + } + transport.CloseIdleConnections() + check(httpProtos) + } +} + +func TestDialCompression(t *testing.T) { + s := newServer(t) + defer s.Close() + + dialer := cstDialer + dialer.EnableCompression = true + ws, _, _, err := dialer.Dial(s.URL, nil) + if err != nil { + t.Fatalf("Dial: %v", err) + } + defer func() { + _ = ws.Close() + }() + sendRecv(t, ws) +} + +func TestTracingDialWithContext(t *testing.T) { + + var headersWrote, requestWrote, getConn, gotConn, connectDone, gotFirstResponseByte bool + trace := &httptrace.ClientTrace{ + WroteHeaders: func() { + headersWrote = true + }, + WroteRequest: func(httptrace.WroteRequestInfo) { + requestWrote = true + }, + GetConn: func(hostPort string) { + getConn = true + }, + GotConn: func(info httptrace.GotConnInfo) { + gotConn = true + }, + ConnectDone: func(network, addr string, err error) { + connectDone = true + }, + GotFirstResponseByte: func() { + gotFirstResponseByte = true + }, + } + ctx := httptrace.WithClientTrace(context.Background(), trace) + + s := newTLSServer(t) + defer s.Close() + + d := cstDialer + d.TLSClientConfig = &tls.Config{RootCAs: rootCAs(t, s.Server)} + + ws, _, _, err := d.DialContext(ctx, s.URL, nil) + if err != nil { + t.Fatalf("Dial: %v", err) + } + + if !headersWrote { + t.Fatal("Headers was not written") + } + if !requestWrote { + t.Fatal("Request was not written") + } + if !getConn { + t.Fatal("getConn was not called") + } + if !gotConn { + t.Fatal("gotConn was not called") + } + if !connectDone { + t.Fatal("connectDone was not called") + } + if !gotFirstResponseByte { + t.Fatal("GotFirstResponseByte was not called") + } + + defer func() { + _ = ws.Close() + }() + sendRecv(t, ws) +} + +func TestEmptyTracingDialWithContext(t *testing.T) { + + trace := &httptrace.ClientTrace{} + ctx := httptrace.WithClientTrace(context.Background(), trace) + + s := newTLSServer(t) + defer s.Close() + + d := cstDialer + d.TLSClientConfig = &tls.Config{RootCAs: rootCAs(t, s.Server)} + + ws, _, _, err := d.DialContext(ctx, s.URL, nil) + if err != nil { + t.Fatalf("Dial: %v", err) + } + + defer func() { + _ = ws.Close() + }() + sendRecv(t, ws) +} + +// TestNetDialConnect tests selection of dial method between NetDial, NetDialContext, NetDialTLS or NetDialTLSContext +func TestNetDialConnect(t *testing.T) { + + upgrader := Upgrader{} + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if IsWebSocketUpgrade(r) { + c, _, err := upgrader.Upgrade(w, r, http.Header{"X-Test-Host": {r.Host}}) + if err != nil { + t.Fatal(err) + } + _ = c.Close() + } else { + w.Header().Set("X-Test-Host", r.Host) + } + }) + + server := httptest.NewServer(handler) + defer server.Close() + + tlsServer := httptest.NewTLSServer(handler) + defer tlsServer.Close() + + testUrls := map[*httptest.Server]string{ + server: "ws://" + server.Listener.Addr().String() + "/", + tlsServer: "wss://" + tlsServer.Listener.Addr().String() + "/", + } + + cas := rootCAs(t, tlsServer) + tlsConfig := &tls.Config{ + RootCAs: cas, + ServerName: "example.com", + InsecureSkipVerify: false, + } + + tests := []struct { + name string + server *httptest.Server // server to use + netDial func(network, addr string) (net.Conn, error) + netDialContext func(ctx context.Context, network, addr string) (net.Conn, error) + netDialTLSContext func(ctx context.Context, network, addr string) (net.Conn, error) + tlsClientConfig *tls.Config + }{ + + { + name: "HTTP server, all NetDial* defined, shall use NetDialContext", + server: server, + netDial: func(network, addr string) (net.Conn, error) { + return nil, errors.New("NetDial should not be called") + }, + netDialContext: func(_ context.Context, network, addr string) (net.Conn, error) { + return net.Dial(network, addr) + }, + netDialTLSContext: func(_ context.Context, network, addr string) (net.Conn, error) { + return nil, errors.New("NetDialTLSContext should not be called") + }, + tlsClientConfig: nil, + }, + { + name: "HTTP server, all NetDial* undefined", + server: server, + netDial: nil, + netDialContext: nil, + netDialTLSContext: nil, + tlsClientConfig: nil, + }, + { + name: "HTTP server, NetDialContext undefined, shall fallback to NetDial", + server: server, + netDial: func(network, addr string) (net.Conn, error) { + return net.Dial(network, addr) + }, + netDialContext: nil, + netDialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + return nil, errors.New("NetDialTLSContext should not be called") + }, + tlsClientConfig: nil, + }, + { + name: "HTTPS server, all NetDial* defined, shall use NetDialTLSContext", + server: tlsServer, + netDial: func(network, addr string) (net.Conn, error) { + return nil, errors.New("NetDial should not be called") + }, + netDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + return nil, errors.New("NetDialContext should not be called") + }, + netDialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + netConn, err := net.Dial(network, addr) + if err != nil { + return nil, err + } + tlsConn := tls.Client(netConn, tlsConfig) + err = tlsConn.Handshake() + if err != nil { + return nil, err + } + return tlsConn, nil + }, + tlsClientConfig: nil, + }, + { + name: "HTTPS server, NetDialTLSContext undefined, shall fallback to NetDialContext and do handshake", + server: tlsServer, + netDial: func(network, addr string) (net.Conn, error) { + return nil, errors.New("NetDial should not be called") + }, + netDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + return net.Dial(network, addr) + }, + netDialTLSContext: nil, + tlsClientConfig: tlsConfig, + }, + { + name: "HTTPS server, NetDialTLSContext and NetDialContext undefined, shall fallback to NetDial and do handshake", + server: tlsServer, + netDial: func(network, addr string) (net.Conn, error) { + return net.Dial(network, addr) + }, + netDialContext: nil, + netDialTLSContext: nil, + tlsClientConfig: tlsConfig, + }, + { + name: "HTTPS server, all NetDial* undefined", + server: tlsServer, + netDial: nil, + netDialContext: nil, + netDialTLSContext: nil, + tlsClientConfig: tlsConfig, + }, + { + name: "HTTPS server, all NetDialTLSContext defined, dummy TlsClientConfig defined, shall not do handshake", + server: tlsServer, + netDial: func(network, addr string) (net.Conn, error) { + return nil, errors.New("NetDial should not be called") + }, + netDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + return nil, errors.New("NetDialContext should not be called") + }, + netDialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + netConn, err := net.Dial(network, addr) + if err != nil { + return nil, err + } + tlsConn := tls.Client(netConn, tlsConfig) + err = tlsConn.Handshake() + if err != nil { + return nil, err + } + return tlsConn, nil + }, + tlsClientConfig: &tls.Config{ + RootCAs: nil, + ServerName: "badserver.com", + InsecureSkipVerify: false, + }, + }, + } + + for _, tc := range tests { + dialer := Dialer{ + NetDial: tc.netDial, + NetDialContext: tc.netDialContext, + NetDialTLSContext: tc.netDialTLSContext, + TLSClientConfig: tc.tlsClientConfig, + } + + // Test websocket dial + c, _, _, err := dialer.Dial(testUrls[tc.server], nil) + if err != nil { + t.Errorf("FAILED %s, err: %s", tc.name, err.Error()) + } else { + _ = c.Close() + } + } +} +func TestNextProtos(t *testing.T) { + ts := httptest.NewUnstartedServer( + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}), + ) + ts.EnableHTTP2 = true + ts.StartTLS() + defer ts.Close() + + d := Dialer{ + TLSClientConfig: ts.Client().Transport.(*http.Transport).TLSClientConfig, + } + + r, err := ts.Client().Get(ts.URL) + if err != nil { + t.Fatalf("Get: %v", err) + } + _ = r.Body.Close() + + // Asserts that Dialer.TLSClientConfig.NextProtos contains "h2" + // after the Client.Get call from net/http above. + containsHTTP2 := false + for _, proto := range d.TLSClientConfig.NextProtos { + if proto == "h2" { + containsHTTP2 = true + } + } + if !containsHTTP2 { + t.Fatalf("Dialer.TLSClientConfig.NextProtos does not contain \"h2\"") + } + + _, _, _, err = d.Dial(makeWsProto(ts.URL), nil) + if err == nil { + t.Fatalf("Dial succeeded, expect fail ") + } +} diff --git a/internal/websocket/client_test.go b/internal/websocket/client_test.go new file mode 100644 index 00000000..5aa27b37 --- /dev/null +++ b/internal/websocket/client_test.go @@ -0,0 +1,32 @@ +// Copyright 2014 The Gorilla WebSocket Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package websocket + +import ( + "net/url" + "testing" +) + +var hostPortNoPortTests = []struct { + u *url.URL + hostPort, hostNoPort string +}{ + {&url.URL{Scheme: "ws", Host: "example.com"}, "example.com:80", "example.com"}, + {&url.URL{Scheme: "wss", Host: "example.com"}, "example.com:443", "example.com"}, + {&url.URL{Scheme: "ws", Host: "example.com:7777"}, "example.com:7777", "example.com"}, + {&url.URL{Scheme: "wss", Host: "example.com:7777"}, "example.com:7777", "example.com"}, +} + +func TestHostPortNoPort(t *testing.T) { + for _, tt := range hostPortNoPortTests { + hostPort, hostNoPort := hostPortNoPort(tt.u) + if hostPort != tt.hostPort { + t.Errorf("hostPortNoPort(%v) returned hostPort %q, want %q", tt.u, hostPort, tt.hostPort) + } + if hostNoPort != tt.hostNoPort { + t.Errorf("hostPortNoPort(%v) returned hostNoPort %q, want %q", tt.u, hostNoPort, tt.hostNoPort) + } + } +} diff --git a/internal/websocket/compression.go b/internal/websocket/compression.go new file mode 100644 index 00000000..9a13cd28 --- /dev/null +++ b/internal/websocket/compression.go @@ -0,0 +1,148 @@ +// Copyright 2017 The Gorilla WebSocket Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package websocket + +import ( + "compress/flate" + "errors" + "io" + "strings" + "sync" +) + +const ( + minCompressionLevel = -2 // flate.HuffmanOnly not defined in Go < 1.6 + maxCompressionLevel = flate.BestCompression + defaultCompressionLevel = 1 +) + +var ( + flateWriterPools [maxCompressionLevel - minCompressionLevel + 1]sync.Pool + flateReaderPool = sync.Pool{New: func() interface{} { + return flate.NewReader(nil) + }} +) + +func decompressNoContextTakeover(r io.Reader) io.ReadCloser { + const tail = + // Add four bytes as specified in RFC + "\x00\x00\xff\xff" + + // Add final block to squelch unexpected EOF error from flate reader. + "\x01\x00\x00\xff\xff" + + fr, _ := flateReaderPool.Get().(io.ReadCloser) + _ = fr.(flate.Resetter).Reset(io.MultiReader(r, strings.NewReader(tail)), nil) + return &flateReadWrapper{fr} +} + +func isValidCompressionLevel(level int) bool { + return minCompressionLevel <= level && level <= maxCompressionLevel +} + +func compressNoContextTakeover(w io.WriteCloser, level int) io.WriteCloser { + p := &flateWriterPools[level-minCompressionLevel] + tw := &truncWriter{w: w} + fw, _ := p.Get().(*flate.Writer) + if fw == nil { + fw, _ = flate.NewWriter(tw, level) + } else { + fw.Reset(tw) + } + return &flateWriteWrapper{fw: fw, tw: tw, p: p} +} + +// truncWriter is an io.Writer that writes all but the last four bytes of the +// stream to another io.Writer. +type truncWriter struct { + w io.WriteCloser + n int + p [4]byte +} + +func (w *truncWriter) Write(p []byte) (int, error) { + n := 0 + + // fill buffer first for simplicity. + if w.n < len(w.p) { + n = copy(w.p[w.n:], p) + p = p[n:] + w.n += n + if len(p) == 0 { + return n, nil + } + } + + m := len(p) + if m > len(w.p) { + m = len(w.p) + } + + if nn, err := w.w.Write(w.p[:m]); err != nil { + return n + nn, err + } + + copy(w.p[:], w.p[m:]) + copy(w.p[len(w.p)-m:], p[len(p)-m:]) + nn, err := w.w.Write(p[:len(p)-m]) + return n + nn, err +} + +type flateWriteWrapper struct { + fw *flate.Writer + tw *truncWriter + p *sync.Pool +} + +func (w *flateWriteWrapper) Write(p []byte) (int, error) { + if w.fw == nil { + return 0, errWriteClosed + } + return w.fw.Write(p) +} + +func (w *flateWriteWrapper) Close() error { + if w.fw == nil { + return errWriteClosed + } + err1 := w.fw.Flush() + w.p.Put(w.fw) + w.fw = nil + if w.tw.p != [4]byte{0, 0, 0xff, 0xff} { + return errors.New("websocket: internal error, unexpected bytes at end of flate stream") + } + err2 := w.tw.w.Close() + if err1 != nil { + return err1 + } + return err2 +} + +type flateReadWrapper struct { + fr io.ReadCloser +} + +func (r *flateReadWrapper) Read(p []byte) (int, error) { + if r.fr == nil { + return 0, io.ErrClosedPipe + } + n, err := r.fr.Read(p) + if err == io.EOF { + // Preemptively place the reader back in the pool. This helps with + // scenarios where the application does not call NextReader() soon after + // this final read. + _ = r.Close() + } + return n, err +} + +func (r *flateReadWrapper) Close() error { + if r.fr == nil { + return io.ErrClosedPipe + } + err := r.fr.Close() + flateReaderPool.Put(r.fr) + r.fr = nil + return err +} diff --git a/internal/websocket/compression_test.go b/internal/websocket/compression_test.go new file mode 100644 index 00000000..00ae42f1 --- /dev/null +++ b/internal/websocket/compression_test.go @@ -0,0 +1,79 @@ +package websocket + +import ( + "bytes" + "fmt" + "io" + "testing" +) + +type nopCloser struct{ io.Writer } + +func (nopCloser) Close() error { return nil } + +func TestTruncWriter(t *testing.T) { + const data = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijlkmnopqrstuvwxyz987654321" + for n := 1; n <= 10; n++ { + var b bytes.Buffer + w := &truncWriter{w: nopCloser{&b}} + p := []byte(data) + for len(p) > 0 { + m := len(p) + if m > n { + m = n + } + _, _ = w.Write(p[:m]) + p = p[m:] + } + if b.String() != data[:len(data)-len(w.p)] { + t.Errorf("%d: %q", n, b.String()) + } + } +} + +func textMessages(num int) [][]byte { + messages := make([][]byte, num) + for i := 0; i < num; i++ { + msg := fmt.Sprintf("planet: %d, country: %d, city: %d, street: %d", i, i, i, i) + messages[i] = []byte(msg) + } + return messages +} + +func BenchmarkWriteNoCompression(b *testing.B) { + w := io.Discard + c := newTestConn(nil, w, false) + messages := textMessages(100) + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = c.WriteMessage(TextMessage, messages[i%len(messages)]) + } + b.ReportAllocs() +} + +func BenchmarkWriteWithCompression(b *testing.B) { + w := io.Discard + c := newTestConn(nil, w, false) + messages := textMessages(100) + c.enableWriteCompression = true + c.newCompressionWriter = compressNoContextTakeover + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = c.WriteMessage(TextMessage, messages[i%len(messages)]) + } + b.ReportAllocs() +} + +func TestValidCompressionLevel(t *testing.T) { + c := newTestConn(nil, nil, false) + for _, level := range []int{minCompressionLevel - 1, maxCompressionLevel + 1} { + if err := c.SetCompressionLevel(level); err == nil { + t.Errorf("no error for level %d", level) + } + } + for _, level := range []int{minCompressionLevel, maxCompressionLevel} { + if err := c.SetCompressionLevel(level); err != nil { + t.Errorf("error for level %d", level) + } + } +} diff --git a/internal/websocket/conn.go b/internal/websocket/conn.go new file mode 100644 index 00000000..5664aec6 --- /dev/null +++ b/internal/websocket/conn.go @@ -0,0 +1,1187 @@ +// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package websocket + +import ( + "bufio" + "encoding/binary" + "errors" + "io" + "math/rand" + "net" + "strconv" + "strings" + "sync" + "time" + "unicode/utf8" +) + +const ( + // Frame header byte 0 bits from Section 5.2 of RFC 6455 + finalBit = 1 << 7 + rsv1Bit = 1 << 6 + rsv2Bit = 1 << 5 + rsv3Bit = 1 << 4 + + // Frame header byte 1 bits from Section 5.2 of RFC 6455 + maskBit = 1 << 7 + + maxFrameHeaderSize = 2 + 8 + 4 // Fixed header + length + mask + maxControlFramePayloadSize = 125 + + writeWait = time.Second + + defaultReadBufferSize = 4096 + defaultWriteBufferSize = 4096 + + continuationFrame = 0 + noFrame = -1 +) + +// Close codes defined in RFC 6455, section 11.7. +const ( + CloseNormalClosure = 1000 + CloseGoingAway = 1001 + CloseProtocolError = 1002 + CloseUnsupportedData = 1003 + CloseNoStatusReceived = 1005 + CloseAbnormalClosure = 1006 + CloseInvalidFramePayloadData = 1007 + ClosePolicyViolation = 1008 + CloseMessageTooBig = 1009 + CloseMandatoryExtension = 1010 + CloseInternalServerErr = 1011 + CloseServiceRestart = 1012 + CloseTryAgainLater = 1013 + CloseTLSHandshake = 1015 +) + +// The message types are defined in RFC 6455, section 11.8. +const ( + // TextMessage denotes a text data message. The text message payload is + // interpreted as UTF-8 encoded text data. + TextMessage = 1 + + // BinaryMessage denotes a binary data message. + BinaryMessage = 2 + + // CloseMessage denotes a close control message. The optional message + // payload contains a numeric code and text. Use the FormatCloseMessage + // function to format a close message payload. + CloseMessage = 8 + + // PingMessage denotes a ping control message. The optional message payload + // is UTF-8 encoded text. + PingMessage = 9 + + // PongMessage denotes a pong control message. The optional message payload + // is UTF-8 encoded text. + PongMessage = 10 +) + +// ErrCloseSent is returned when the application writes a message to the +// connection after sending a close message. +var ErrCloseSent = errors.New("websocket: close sent") + +// ErrReadLimit is returned when reading a message that is larger than the +// read limit set for the connection. +var ErrReadLimit = errors.New("websocket: read limit exceeded") + +// netError satisfies the net Error interface. +type netError struct { + msg string + temporary bool + timeout bool +} + +func (e *netError) Error() string { return e.msg } +func (e *netError) Temporary() bool { return e.temporary } +func (e *netError) Timeout() bool { return e.timeout } + +// CloseError represents a close message. +type CloseError struct { + // Text is the optional text payload. + Text string + // Code is defined in RFC 6455, section 11.7. + Code int +} + +func (e *CloseError) Error() string { + s := []byte("websocket: close ") + s = strconv.AppendInt(s, int64(e.Code), 10) + switch e.Code { + case CloseNormalClosure: + s = append(s, " (normal)"...) + case CloseGoingAway: + s = append(s, " (going away)"...) + case CloseProtocolError: + s = append(s, " (protocol error)"...) + case CloseUnsupportedData: + s = append(s, " (unsupported data)"...) + case CloseNoStatusReceived: + s = append(s, " (no status)"...) + case CloseAbnormalClosure: + s = append(s, " (abnormal closure)"...) + case CloseInvalidFramePayloadData: + s = append(s, " (invalid payload data)"...) + case ClosePolicyViolation: + s = append(s, " (policy violation)"...) + case CloseMessageTooBig: + s = append(s, " (message too big)"...) + case CloseMandatoryExtension: + s = append(s, " (mandatory extension missing)"...) + case CloseInternalServerErr: + s = append(s, " (internal server error)"...) + case CloseTLSHandshake: + s = append(s, " (TLS handshake error)"...) + } + if e.Text != "" { + s = append(s, ": "...) + s = append(s, e.Text...) + } + return string(s) +} + +// IsCloseError returns boolean indicating whether the error is a *CloseError +// with one of the specified codes. +func IsCloseError(err error, codes ...int) bool { + if e, ok := err.(*CloseError); ok { + for _, code := range codes { + if e.Code == code { + return true + } + } + } + return false +} + +// IsUnexpectedCloseError returns boolean indicating whether the error is a +// *CloseError with a code not in the list of expected codes. +func IsUnexpectedCloseError(err error, expectedCodes ...int) bool { + if e, ok := err.(*CloseError); ok { + for _, code := range expectedCodes { + if e.Code == code { + return false + } + } + return true + } + return false +} + +var ( + errWriteTimeout = &netError{msg: "websocket: write timeout", timeout: true, temporary: true} + errUnexpectedEOF = &CloseError{Code: CloseAbnormalClosure, Text: io.ErrUnexpectedEOF.Error()} + errBadWriteOpCode = errors.New("websocket: bad write message type") + errWriteClosed = errors.New("websocket: write closed") + errInvalidControlFrame = errors.New("websocket: invalid control frame") +) + +func newMaskKey() [4]byte { + n := rand.Uint32() + return [4]byte{byte(n), byte(n >> 8), byte(n >> 16), byte(n >> 24)} +} + +func hideTempErr(err error) error { + if e, ok := err.(net.Error); ok && e.Temporary() { //nolint:staticcheck + err = &netError{msg: e.Error(), timeout: e.Timeout()} + } + return err +} + +func isControl(frameType int) bool { + return frameType == CloseMessage || frameType == PingMessage || frameType == PongMessage +} + +func isData(frameType int) bool { + return frameType == TextMessage || frameType == BinaryMessage +} + +var validReceivedCloseCodes = map[int]bool{ + // see http://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number + + CloseNormalClosure: true, + CloseGoingAway: true, + CloseProtocolError: true, + CloseUnsupportedData: true, + CloseNoStatusReceived: false, + CloseAbnormalClosure: false, + CloseInvalidFramePayloadData: true, + ClosePolicyViolation: true, + CloseMessageTooBig: true, + CloseMandatoryExtension: true, + CloseInternalServerErr: true, + CloseServiceRestart: true, + CloseTryAgainLater: true, + CloseTLSHandshake: false, +} + +func isValidReceivedCloseCode(code int) bool { + return validReceivedCloseCodes[code] || (code >= 3000 && code <= 4999) +} + +// BufferPool represents a pool of buffers. The *sync.Pool type satisfies this +// interface. The type of the value stored in a pool is not specified. +type BufferPool interface { + // Get gets a value from the pool or returns nil if the pool is empty. + Get() interface{} + // Put adds a value to the pool. + Put(interface{}) +} + +// writePoolData is the type added to the write buffer pool. This wrapper is +// used to prevent applications from peeking at and depending on the values +// added to the pool. +type writePoolData struct{ buf []byte } + +// The Conn type represents a WebSocket connection. +type Conn struct { + conn net.Conn + + reader io.ReadCloser // the current reader returned to the application + readErr error + + mu chan struct{} // used as mutex to protect write to conn + writeDeadline time.Time + writer io.WriteCloser // the current writer returned to the application + writePool BufferPool + writeErr error + + newDecompressionReader func(io.Reader) io.ReadCloser + messageReader *messageReader // the current low-level reader + handleClose func(int, string) error + handlePing func(string) error + handlePong func(string) error + newCompressionWriter func(io.WriteCloser, int) io.WriteCloser + br *bufio.Reader + writeBuf []byte // frame is constructed in this buffer. + readLength int64 // Message size. + readLimit int64 // Maximum message size. + compressionLevel int + readMaskPos int + writeBufSize int + writeErrMu sync.Mutex + // bytes remaining in current frame. + // set setReadRemaining to safely update this value and prevent overflow + readRemaining int64 + readMaskKey [4]byte + readDecompress bool // whether last read frame had RSV1 set + readFinal bool // true the current message has more frames. + enableWriteCompression bool + isServer bool +} + +func newConn(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int, writeBufferPool BufferPool, br *bufio.Reader, writeBuf []byte) *Conn { + if br == nil { + if readBufferSize == 0 { + readBufferSize = defaultReadBufferSize + } else if readBufferSize < maxControlFramePayloadSize { + // must be large enough for control frame + readBufferSize = maxControlFramePayloadSize + } + br = bufio.NewReaderSize(conn, readBufferSize) + } + + if writeBufferSize <= 0 { + writeBufferSize = defaultWriteBufferSize + } + writeBufferSize += maxFrameHeaderSize + + if writeBuf == nil && writeBufferPool == nil { + writeBuf = make([]byte, writeBufferSize) + } + + mu := make(chan struct{}, 1) + mu <- struct{}{} + c := &Conn{ + isServer: isServer, + br: br, + conn: conn, + mu: mu, + readFinal: true, + writeBuf: writeBuf, + writePool: writeBufferPool, + writeBufSize: writeBufferSize, + enableWriteCompression: true, + compressionLevel: defaultCompressionLevel, + } + c.SetCloseHandler(nil) + c.SetPingHandler(nil) + c.SetPongHandler(nil) + return c +} + +// setReadRemaining tracks the number of bytes remaining on the connection. If n +// overflows, an ErrReadLimit is returned. +func (c *Conn) setReadRemaining(n int64) error { + if n < 0 { + return ErrReadLimit + } + + c.readRemaining = n + return nil +} + +// Close closes the underlying network connection without sending or waiting +// for a close message. +func (c *Conn) Close() error { + return c.conn.Close() +} + +// LocalAddr returns the local network address. +func (c *Conn) LocalAddr() net.Addr { + return c.conn.LocalAddr() +} + +// RemoteAddr returns the remote network address. +func (c *Conn) RemoteAddr() net.Addr { + return c.conn.RemoteAddr() +} + +// Write methods + +func (c *Conn) writeFatal(err error) error { + err = hideTempErr(err) + c.writeErrMu.Lock() + if c.writeErr == nil { + c.writeErr = err + } + c.writeErrMu.Unlock() + return err +} + +func (c *Conn) read(n int) ([]byte, error) { + p, err := c.br.Peek(n) + if err == io.EOF { + err = errUnexpectedEOF + } + _, _ = c.br.Discard(len(p)) + return p, err +} + +func (c *Conn) write(frameType int, deadline time.Time, buf0, buf1 []byte) error { + <-c.mu + defer func() { c.mu <- struct{}{} }() + + c.writeErrMu.Lock() + err := c.writeErr + c.writeErrMu.Unlock() + if err != nil { + return err + } + + _ = c.conn.SetWriteDeadline(deadline) + if len(buf1) == 0 { + _, err = c.conn.Write(buf0) + } else { + err = c.writeBufs(buf0, buf1) + } + if err != nil { + return c.writeFatal(err) + } + if frameType == CloseMessage { + _ = c.writeFatal(ErrCloseSent) + } + return nil +} + +func (c *Conn) writeBufs(bufs ...[]byte) error { + b := net.Buffers(bufs) + _, err := b.WriteTo(c.conn) + return err +} + +// WriteControl writes a control message with the given deadline. The allowed +// message types are CloseMessage, PingMessage and PongMessage. +func (c *Conn) WriteControl(messageType int, data []byte, deadline time.Time) error { + if !isControl(messageType) { + return errBadWriteOpCode + } + if len(data) > maxControlFramePayloadSize { + return errInvalidControlFrame + } + + b0 := byte(messageType) | finalBit + b1 := byte(len(data)) + if !c.isServer { + b1 |= maskBit + } + + buf := make([]byte, 0, maxFrameHeaderSize+maxControlFramePayloadSize) + buf = append(buf, b0, b1) + + if c.isServer { + buf = append(buf, data...) + } else { + key := newMaskKey() + buf = append(buf, key[:]...) + buf = append(buf, data...) + maskBytes(key, 0, buf[6:]) + } + + d := 1000 * time.Hour + if !deadline.IsZero() { + d = time.Until(deadline) + if d < 0 { + return errWriteTimeout + } + } + + select { + case <-c.mu: + default: + timer := time.NewTimer(d) + select { + case <-c.mu: + timer.Stop() + case <-timer.C: + return errWriteTimeout + } + } + + defer func() { c.mu <- struct{}{} }() + + c.writeErrMu.Lock() + err := c.writeErr + c.writeErrMu.Unlock() + if err != nil { + return err + } + + _ = c.conn.SetWriteDeadline(deadline) + _, err = c.conn.Write(buf) + if err != nil { + return c.writeFatal(err) + } + if messageType == CloseMessage { + _ = c.writeFatal(ErrCloseSent) + } + return err +} + +// beginMessage prepares a connection and message writer for a new message. +func (c *Conn) beginMessage(mw *messageWriter, messageType int) error { + // Close previous writer if not already closed by the application. It's + // probably better to return an error in this situation, but we cannot + // change this without breaking existing applications. + if c.writer != nil { + _ = c.writer.Close() + c.writer = nil + } + + if !isControl(messageType) && !isData(messageType) { + return errBadWriteOpCode + } + + c.writeErrMu.Lock() + err := c.writeErr + c.writeErrMu.Unlock() + if err != nil { + return err + } + + mw.c = c + mw.frameType = messageType + mw.pos = maxFrameHeaderSize + + if c.writeBuf == nil { + wpd, ok := c.writePool.Get().(writePoolData) + if ok { + c.writeBuf = wpd.buf + } else { + c.writeBuf = make([]byte, c.writeBufSize) + } + } + return nil +} + +// NextWriter returns a writer for the next message to send. The writer's Close +// method flushes the complete message to the network. +// +// There can be at most one open writer on a connection. NextWriter closes the +// previous writer if the application has not already done so. +// +// All message types (TextMessage, BinaryMessage, CloseMessage, PingMessage and +// PongMessage) are supported. +func (c *Conn) NextWriter(messageType int) (io.WriteCloser, error) { + var mw messageWriter + if err := c.beginMessage(&mw, messageType); err != nil { + return nil, err + } + c.writer = &mw + if c.newCompressionWriter != nil && c.enableWriteCompression && isData(messageType) { + w := c.newCompressionWriter(c.writer, c.compressionLevel) + mw.compress = true + c.writer = w + } + return c.writer, nil +} + +type messageWriter struct { + c *Conn + err error + pos int // end of data in writeBuf. + frameType int // type of the current frame. + compress bool // whether next call to flushFrame should set RSV1 +} + +func (w *messageWriter) endMessage(err error) error { + if w.err != nil { + return err + } + c := w.c + w.err = err + c.writer = nil + if c.writePool != nil { + c.writePool.Put(writePoolData{buf: c.writeBuf}) + c.writeBuf = nil + } + return err +} + +// flushFrame writes buffered data and extra as a frame to the network. The +// final argument indicates that this is the last frame in the message. +func (w *messageWriter) flushFrame(final bool, extra []byte) error { + c := w.c + length := w.pos - maxFrameHeaderSize + len(extra) + + // Check for invalid control frames. + if isControl(w.frameType) && + (!final || length > maxControlFramePayloadSize) { + return w.endMessage(errInvalidControlFrame) + } + + b0 := byte(w.frameType) + if final { + b0 |= finalBit + } + if w.compress { + b0 |= rsv1Bit + } + w.compress = false + + b1 := byte(0) + if !c.isServer { + b1 |= maskBit + } + + // Assume that the frame starts at beginning of c.writeBuf. + framePos := 0 + if c.isServer { + // Adjust up if mask not included in the header. + framePos = 4 + } + + switch { + case length >= 65536: + c.writeBuf[framePos] = b0 + c.writeBuf[framePos+1] = b1 | 127 + binary.BigEndian.PutUint64(c.writeBuf[framePos+2:], uint64(length)) + case length > 125: + framePos += 6 + c.writeBuf[framePos] = b0 + c.writeBuf[framePos+1] = b1 | 126 + binary.BigEndian.PutUint16(c.writeBuf[framePos+2:], uint16(length)) + default: + framePos += 8 + c.writeBuf[framePos] = b0 + c.writeBuf[framePos+1] = b1 | byte(length) + } + + if !c.isServer { + key := newMaskKey() + copy(c.writeBuf[maxFrameHeaderSize-4:], key[:]) + maskBytes(key, 0, c.writeBuf[maxFrameHeaderSize:w.pos]) + if len(extra) > 0 { + return w.endMessage(c.writeFatal(errors.New("websocket: internal error, extra used in client mode"))) + } + } + + err := c.write(w.frameType, c.writeDeadline, c.writeBuf[framePos:w.pos], extra) + if err != nil { + return w.endMessage(err) + } + + if final { + _ = w.endMessage(errWriteClosed) + return nil + } + + // Setup for next frame. + w.pos = maxFrameHeaderSize + w.frameType = continuationFrame + return nil +} + +func (w *messageWriter) ncopy(max int) (int, error) { + n := len(w.c.writeBuf) - w.pos + if n <= 0 { + if err := w.flushFrame(false, nil); err != nil { + return 0, err + } + n = len(w.c.writeBuf) - w.pos + } + if n > max { + n = max + } + return n, nil +} + +func (w *messageWriter) Write(p []byte) (int, error) { + if w.err != nil { + return 0, w.err + } + + if len(p) > 2*len(w.c.writeBuf) && w.c.isServer { + // Don't buffer large messages. + err := w.flushFrame(false, p) + if err != nil { + return 0, err + } + return len(p), nil + } + + nn := len(p) + for len(p) > 0 { + n, err := w.ncopy(len(p)) + if err != nil { + return 0, err + } + copy(w.c.writeBuf[w.pos:], p[:n]) + w.pos += n + p = p[n:] + } + return nn, nil +} + +func (w *messageWriter) WriteString(p string) (int, error) { + if w.err != nil { + return 0, w.err + } + + nn := len(p) + for len(p) > 0 { + n, err := w.ncopy(len(p)) + if err != nil { + return 0, err + } + copy(w.c.writeBuf[w.pos:], p[:n]) + w.pos += n + p = p[n:] + } + return nn, nil +} + +func (w *messageWriter) ReadFrom(r io.Reader) (nn int64, err error) { + if w.err != nil { + return 0, w.err + } + for { + if w.pos == len(w.c.writeBuf) { + err = w.flushFrame(false, nil) + if err != nil { + break + } + } + var n int + n, err = r.Read(w.c.writeBuf[w.pos:]) + w.pos += n + nn += int64(n) + if err != nil { + if err == io.EOF { + err = nil + } + break + } + } + return nn, err +} + +func (w *messageWriter) Close() error { + if w.err != nil { + return w.err + } + return w.flushFrame(true, nil) +} + +// WritePreparedMessage writes prepared message into connection. +func (c *Conn) WritePreparedMessage(pm *PreparedMessage) error { + frameType, frameData, err := pm.frame(prepareKey{ + isServer: c.isServer, + compress: c.newCompressionWriter != nil && c.enableWriteCompression && isData(pm.messageType), + compressionLevel: c.compressionLevel, + }) + if err != nil { + return err + } + return c.write(frameType, c.writeDeadline, frameData, nil) +} + +// WriteMessage is a helper method for getting a writer using NextWriter, +// writing the message and closing the writer. +func (c *Conn) WriteMessage(messageType int, data []byte) error { + if c.isServer && (c.newCompressionWriter == nil || !c.enableWriteCompression) { + // Fast path with no allocations and single frame. + var mw messageWriter + if err := c.beginMessage(&mw, messageType); err != nil { + return err + } + n := copy(c.writeBuf[mw.pos:], data) + mw.pos += n + data = data[n:] + return mw.flushFrame(true, data) + } + + w, err := c.NextWriter(messageType) + if err != nil { + return err + } + if _, err = w.Write(data); err != nil { + return err + } + return w.Close() +} + +// SetWriteDeadline sets the write deadline on the underlying network +// connection. After a write has timed out, the websocket state is corrupt and +// all future writes will return an error. A zero value for t means writes will +// not time out. +func (c *Conn) SetWriteDeadline(t time.Time) error { + c.writeDeadline = t + return nil +} + +// Read methods + +func (c *Conn) advanceFrame() (int, error) { + // 1. Skip remainder of previous frame. + + if c.readRemaining > 0 { + if _, err := io.CopyN(io.Discard, c.br, c.readRemaining); err != nil { + return noFrame, err + } + } + + // 2. Read and parse first two bytes of frame header. + // To aid debugging, collect and report all errors in the first two bytes + // of the header. + + var errs []string + + p, err := c.read(2) + if err != nil { + return noFrame, err + } + + frameType := int(p[0] & 0xf) + final := p[0]&finalBit != 0 + rsv1 := p[0]&rsv1Bit != 0 + rsv2 := p[0]&rsv2Bit != 0 + rsv3 := p[0]&rsv3Bit != 0 + mask := p[1]&maskBit != 0 + _ = c.setReadRemaining(int64(p[1] & 0x7f)) + + c.readDecompress = false + if rsv1 { + if c.newDecompressionReader != nil { + c.readDecompress = true + } else { + errs = append(errs, "RSV1 set") + } + } + + if rsv2 { + errs = append(errs, "RSV2 set") + } + + if rsv3 { + errs = append(errs, "RSV3 set") + } + + switch frameType { + case CloseMessage, PingMessage, PongMessage: + if c.readRemaining > maxControlFramePayloadSize { + errs = append(errs, "len > 125 for control") + } + if !final { + errs = append(errs, "FIN not set on control") + } + case TextMessage, BinaryMessage: + if !c.readFinal { + errs = append(errs, "data before FIN") + } + c.readFinal = final + case continuationFrame: + if c.readFinal { + errs = append(errs, "continuation after FIN") + } + c.readFinal = final + default: + errs = append(errs, "bad opcode "+strconv.Itoa(frameType)) + } + + if mask != c.isServer { + errs = append(errs, "bad MASK") + } + + if len(errs) > 0 { + return noFrame, c.handleProtocolError(strings.Join(errs, ", ")) + } + + // 3. Read and parse frame length as per + // https://tools.ietf.org/html/rfc6455#section-5.2 + // + // The length of the "Payload data", in bytes: if 0-125, that is the payload + // length. + // - If 126, the following 2 bytes interpreted as a 16-bit unsigned + // integer are the payload length. + // - If 127, the following 8 bytes interpreted as + // a 64-bit unsigned integer (the most significant bit MUST be 0) are the + // payload length. Multibyte length quantities are expressed in network byte + // order. + + switch c.readRemaining { + case 126: + p, err := c.read(2) + if err != nil { + return noFrame, err + } + + if err := c.setReadRemaining(int64(binary.BigEndian.Uint16(p))); err != nil { + return noFrame, err + } + case 127: + p, err := c.read(8) + if err != nil { + return noFrame, err + } + + if err := c.setReadRemaining(int64(binary.BigEndian.Uint64(p))); err != nil { + return noFrame, err + } + } + + // 4. Handle frame masking. + + if mask { + c.readMaskPos = 0 + p, err := c.read(len(c.readMaskKey)) + if err != nil { + return noFrame, err + } + copy(c.readMaskKey[:], p) + } + + // 5. For text and binary messages, enforce read limit and return. + + if frameType == continuationFrame || frameType == TextMessage || frameType == BinaryMessage { + + c.readLength += c.readRemaining + // Don't allow readLength to overflow in the presence of a large readRemaining + // counter. + if c.readLength < 0 { + return noFrame, ErrReadLimit + } + + if c.readLimit > 0 && c.readLength > c.readLimit { + _ = c.WriteControl(CloseMessage, FormatCloseMessage(CloseMessageTooBig, ""), time.Now().Add(writeWait)) + return noFrame, ErrReadLimit + } + + return frameType, nil + } + + // 6. Read control frame payload. + + var payload []byte + if c.readRemaining > 0 { + payload, err = c.read(int(c.readRemaining)) + _ = c.setReadRemaining(0) + if err != nil { + return noFrame, err + } + if c.isServer { + maskBytes(c.readMaskKey, 0, payload) + } + } + + // 7. Process control frame payload. + + switch frameType { + case PongMessage: + if err := c.handlePong(string(payload)); err != nil { + return noFrame, err + } + case PingMessage: + if err := c.handlePing(string(payload)); err != nil { + return noFrame, err + } + case CloseMessage: + closeCode := CloseNoStatusReceived + closeText := "" + if len(payload) >= 2 { + closeCode = int(binary.BigEndian.Uint16(payload)) + if !isValidReceivedCloseCode(closeCode) { + return noFrame, c.handleProtocolError("bad close code " + strconv.Itoa(closeCode)) + } + closeText = string(payload[2:]) + if !utf8.ValidString(closeText) { + return noFrame, c.handleProtocolError("invalid utf8 payload in close frame") + } + } + if err := c.handleClose(closeCode, closeText); err != nil { + return noFrame, err + } + return noFrame, &CloseError{Code: closeCode, Text: closeText} + } + + return frameType, nil +} + +func (c *Conn) handleProtocolError(message string) error { + data := FormatCloseMessage(CloseProtocolError, message) + if len(data) > maxControlFramePayloadSize { + data = data[:maxControlFramePayloadSize] + } + _ = c.WriteControl(CloseMessage, data, time.Now().Add(writeWait)) + return errors.New("websocket: " + message) +} + +// NextReader returns the next data message received from the peer. The +// returned messageType is either TextMessage or BinaryMessage. +// +// There can be at most one open reader on a connection. NextReader discards +// the previous message if the application has not already consumed it. +// +// Applications must break out of the application's read loop when this method +// returns a non-nil error value. Errors returned from this method are +// permanent. Once this method returns a non-nil error, all subsequent calls to +// this method return the same error. +func (c *Conn) NextReader() (messageType int, r io.Reader, err error) { + // Close previous reader, only relevant for decompression. + if c.reader != nil { + _ = c.reader.Close() + c.reader = nil + } + + c.messageReader = nil + c.readLength = 0 + + for c.readErr == nil { + frameType, err := c.advanceFrame() + if err != nil { + c.readErr = hideTempErr(err) + break + } + + if frameType == TextMessage || frameType == BinaryMessage { + c.messageReader = &messageReader{c} + c.reader = c.messageReader + if c.readDecompress { + c.reader = c.newDecompressionReader(c.reader) + } + return frameType, c.reader, nil + } + } + + return noFrame, nil, c.readErr +} + +type messageReader struct{ c *Conn } + +func (r *messageReader) Read(b []byte) (int, error) { + c := r.c + if c.messageReader != r { + return 0, io.EOF + } + + for c.readErr == nil { + + if c.readRemaining > 0 { + if int64(len(b)) > c.readRemaining { + b = b[:c.readRemaining] + } + n, err := c.br.Read(b) + c.readErr = hideTempErr(err) + if c.isServer { + c.readMaskPos = maskBytes(c.readMaskKey, c.readMaskPos, b[:n]) + } + rem := c.readRemaining + rem -= int64(n) + _ = c.setReadRemaining(rem) + if c.readRemaining > 0 && c.readErr == io.EOF { + c.readErr = errUnexpectedEOF + } + return n, c.readErr + } + + if c.readFinal { + c.messageReader = nil + return 0, io.EOF + } + + frameType, err := c.advanceFrame() + switch { + case err != nil: + c.readErr = hideTempErr(err) + case frameType == TextMessage || frameType == BinaryMessage: + c.readErr = errors.New("websocket: internal error, unexpected text or binary in Reader") + } + } + + err := c.readErr + if err == io.EOF && c.messageReader == r { + err = errUnexpectedEOF + } + return 0, err +} + +func (r *messageReader) Close() error { + return nil +} + +// ReadMessage is a helper method for getting a reader using NextReader and +// reading from that reader to a buffer. +func (c *Conn) ReadMessage() (messageType int, p []byte, err error) { + var r io.Reader + messageType, r, err = c.NextReader() + if err != nil { + return messageType, nil, err + } + p, err = io.ReadAll(r) + return messageType, p, err +} + +// SetReadDeadline sets the read deadline on the underlying network connection. +// After a read has timed out, the websocket connection state is corrupt and +// all future reads will return an error. A zero value for t means reads will +// not time out. +func (c *Conn) SetReadDeadline(t time.Time) error { + return c.conn.SetReadDeadline(t) +} + +// SetReadLimit sets the maximum size in bytes for a message read from the peer. If a +// message exceeds the limit, the connection sends a close message to the peer +// and returns ErrReadLimit to the application. +func (c *Conn) SetReadLimit(limit int64) { + c.readLimit = limit +} + +// CloseHandler returns the current close handler +func (c *Conn) CloseHandler() func(code int, text string) error { + return c.handleClose +} + +// SetCloseHandler sets the handler for close messages received from the peer. +// The code argument to h is the received close code or CloseNoStatusReceived +// if the close message is empty. The default close handler sends a close +// message back to the peer. +// +// The handler function is called from the NextReader, ReadMessage and message +// reader Read methods. The application must read the connection to process +// close messages as described in the section on Control Messages above. +// +// The connection read methods return a CloseError when a close message is +// received. Most applications should handle close messages as part of their +// normal error handling. Applications should only set a close handler when the +// application must perform some action before sending a close message back to +// the peer. +func (c *Conn) SetCloseHandler(h func(code int, text string) error) { + if h == nil { + h = func(code int, text string) error { + message := FormatCloseMessage(code, "") + _ = c.WriteControl(CloseMessage, message, time.Now().Add(writeWait)) + return nil + } + } + c.handleClose = h +} + +// PingHandler returns the current ping handler +func (c *Conn) PingHandler() func(appData string) error { + return c.handlePing +} + +// SetPingHandler sets the handler for ping messages received from the peer. +// The appData argument to h is the PING message application data. The default +// ping handler sends a pong to the peer. +// +// The handler function is called from the NextReader, ReadMessage and message +// reader Read methods. The application must read the connection to process +// ping messages as described in the section on Control Messages above. +func (c *Conn) SetPingHandler(h func(appData string) error) { + if h == nil { + h = func(message string) error { + err := c.WriteControl(PongMessage, []byte(message), time.Now().Add(writeWait)) + if err == ErrCloseSent { + return nil + } else if e, ok := err.(net.Error); ok && e.Temporary() { //nolint:staticcheck + return nil + } + return err + } + } + c.handlePing = h +} + +// PongHandler returns the current pong handler +func (c *Conn) PongHandler() func(appData string) error { + return c.handlePong +} + +// SetPongHandler sets the handler for pong messages received from the peer. +// The appData argument to h is the PONG message application data. The default +// pong handler does nothing. +// +// The handler function is called from the NextReader, ReadMessage and message +// reader Read methods. The application must read the connection to process +// pong messages as described in the section on Control Messages above. +func (c *Conn) SetPongHandler(h func(appData string) error) { + if h == nil { + h = func(string) error { return nil } + } + c.handlePong = h +} + +// NetConn returns the underlying connection that is wrapped by c. +// Note that writing to or reading from this connection directly will corrupt the +// WebSocket connection. +func (c *Conn) NetConn() net.Conn { + return c.conn +} + +// EnableWriteCompression enables and disables write compression of +// subsequent text and binary messages. This function is a noop if +// compression was not negotiated with the peer. +func (c *Conn) EnableWriteCompression(enable bool) { + c.enableWriteCompression = enable +} + +// SetCompressionLevel sets the flate compression level for subsequent text and +// binary messages. This function is a noop if compression was not negotiated +// with the peer. See the compress/flate package for a description of +// compression levels. +func (c *Conn) SetCompressionLevel(level int) error { + if !isValidCompressionLevel(level) { + return errors.New("websocket: invalid compression level") + } + c.compressionLevel = level + return nil +} + +// FormatCloseMessage formats closeCode and text as a WebSocket close message. +// An empty message is returned for code CloseNoStatusReceived. +func FormatCloseMessage(closeCode int, text string) []byte { + if closeCode == CloseNoStatusReceived { + // Return empty message because it's illegal to send + // CloseNoStatusReceived. Return non-nil value in case application + // checks for nil. + return []byte{} + } + buf := make([]byte, 2+len(text)) + binary.BigEndian.PutUint16(buf, uint16(closeCode)) + copy(buf[2:], text) + return buf +} diff --git a/internal/websocket/conn_broadcast_test.go b/internal/websocket/conn_broadcast_test.go new file mode 100644 index 00000000..b954aad4 --- /dev/null +++ b/internal/websocket/conn_broadcast_test.go @@ -0,0 +1,130 @@ +// Copyright 2017 The Gorilla WebSocket Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package websocket + +import ( + "io" + "sync/atomic" + "testing" +) + +// broadcastBench allows to run broadcast benchmarks. +// In every broadcast benchmark we create many connections, then send the same +// message into every connection and wait for all writes complete. This emulates +// an application where many connections listen to the same data - i.e. PUB/SUB +// scenarios with many subscribers in one channel. +type broadcastBench struct { + w io.Writer + closeCh chan struct{} + doneCh chan struct{} + count int32 + conns []*broadcastConn + compression bool + usePrepared bool +} + +type broadcastMessage struct { + payload []byte + prepared *PreparedMessage +} + +type broadcastConn struct { + conn *Conn + msgCh chan *broadcastMessage +} + +func newBroadcastConn(c *Conn) *broadcastConn { + return &broadcastConn{ + conn: c, + msgCh: make(chan *broadcastMessage, 1), + } +} + +func newBroadcastBench(usePrepared, compression bool) *broadcastBench { + bench := &broadcastBench{ + w: io.Discard, + doneCh: make(chan struct{}), + closeCh: make(chan struct{}), + usePrepared: usePrepared, + compression: compression, + } + bench.makeConns(10000) + return bench +} + +func (b *broadcastBench) makeConns(numConns int) { + conns := make([]*broadcastConn, numConns) + + for i := 0; i < numConns; i++ { + c := newTestConn(nil, b.w, true) + if b.compression { + c.enableWriteCompression = true + c.newCompressionWriter = compressNoContextTakeover + } + conns[i] = newBroadcastConn(c) + go func(c *broadcastConn) { + for { + select { + case msg := <-c.msgCh: + if msg.prepared != nil { + _ = c.conn.WritePreparedMessage(msg.prepared) + } else { + _ = c.conn.WriteMessage(TextMessage, msg.payload) + } + val := atomic.AddInt32(&b.count, 1) + if val%int32(numConns) == 0 { + b.doneCh <- struct{}{} + } + case <-b.closeCh: + return + } + } + }(conns[i]) + } + b.conns = conns +} + +func (b *broadcastBench) close() { + close(b.closeCh) +} + +func (b *broadcastBench) broadcastOnce(msg *broadcastMessage) { + for _, c := range b.conns { + c.msgCh <- msg + } + <-b.doneCh +} + +func BenchmarkBroadcast(b *testing.B) { + benchmarks := []struct { + name string + usePrepared bool + compression bool + }{ + {"NoCompression", false, false}, + {"Compression", false, true}, + {"NoCompressionPrepared", true, false}, + {"CompressionPrepared", true, true}, + } + payload := textMessages(1)[0] + for _, bm := range benchmarks { + b.Run(bm.name, func(b *testing.B) { + bench := newBroadcastBench(bm.usePrepared, bm.compression) + defer bench.close() + message := &broadcastMessage{ + payload: payload, + } + b.ResetTimer() + for i := 0; i < b.N; i++ { + if bench.usePrepared { + pm, _ := NewPreparedMessage(TextMessage, message.payload) + message.prepared = pm + } + bench.broadcastOnce(message) + } + b.ReportAllocs() + }) + } +} diff --git a/internal/websocket/conn_test.go b/internal/websocket/conn_test.go new file mode 100644 index 00000000..c107d39b --- /dev/null +++ b/internal/websocket/conn_test.go @@ -0,0 +1,641 @@ +// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package websocket + +import ( + "bufio" + "bytes" + "errors" + "fmt" + "io" + "net" + "reflect" + "sync" + "testing" + "testing/iotest" + "time" +) + +var _ net.Error = errWriteTimeout + +type fakeNetConn struct { + io.Reader + io.Writer +} + +func (c fakeNetConn) Close() error { return nil } +func (c fakeNetConn) LocalAddr() net.Addr { return localAddr } +func (c fakeNetConn) RemoteAddr() net.Addr { return remoteAddr } +func (c fakeNetConn) SetDeadline(_ time.Time) error { return nil } +func (c fakeNetConn) SetReadDeadline(_ time.Time) error { return nil } +func (c fakeNetConn) SetWriteDeadline(_ time.Time) error { return nil } + +type fakeAddr int + +var ( + localAddr = fakeAddr(1) + remoteAddr = fakeAddr(2) +) + +func (a fakeAddr) Network() string { + return "net" +} + +func (a fakeAddr) String() string { + return "str" +} + +// newTestConn creates a connnection backed by a fake network connection using +// default values for buffering. +func newTestConn(r io.Reader, w io.Writer, isServer bool) *Conn { + return newConn(fakeNetConn{Reader: r, Writer: w}, isServer, 1024, 1024, nil, nil, nil) +} + +func TestFraming(t *testing.T) { + frameSizes := []int{ + 0, 1, 2, 124, 125, 126, 127, 128, 129, 65534, 65535, + // 65536, 65537 + } + var readChunkers = []struct { + name string + f func(io.Reader) io.Reader + }{ + {"half", iotest.HalfReader}, + {"one", iotest.OneByteReader}, + {"asis", func(r io.Reader) io.Reader { return r }}, + } + writeBuf := make([]byte, 65537) + for i := range writeBuf { + writeBuf[i] = byte(i) + } + var writers = []struct { + name string + f func(w io.Writer, n int) (int, error) + }{ + {"iocopy", func(w io.Writer, n int) (int, error) { + nn, err := io.Copy(w, bytes.NewReader(writeBuf[:n])) + return int(nn), err + }}, + {"write", func(w io.Writer, n int) (int, error) { + return w.Write(writeBuf[:n]) + }}, + {"string", func(w io.Writer, n int) (int, error) { + return io.WriteString(w, string(writeBuf[:n])) + }}, + } + + for _, compress := range []bool{false, true} { + for _, isServer := range []bool{true, false} { + for _, chunker := range readChunkers { + + var connBuf bytes.Buffer + wc := newTestConn(nil, &connBuf, isServer) + rc := newTestConn(chunker.f(&connBuf), nil, !isServer) + if compress { + wc.newCompressionWriter = compressNoContextTakeover + rc.newDecompressionReader = decompressNoContextTakeover + } + for _, n := range frameSizes { + for _, writer := range writers { + name := fmt.Sprintf("z:%v, s:%v, r:%s, n:%d w:%s", compress, isServer, chunker.name, n, writer.name) + + w, err := wc.NextWriter(TextMessage) + if err != nil { + t.Errorf("%s: wc.NextWriter() returned %v", name, err) + continue + } + nn, err := writer.f(w, n) + if err != nil || nn != n { + t.Errorf("%s: w.Write(writeBuf[:n]) returned %d, %v", name, nn, err) + continue + } + err = w.Close() + if err != nil { + t.Errorf("%s: w.Close() returned %v", name, err) + continue + } + + opCode, r, err := rc.NextReader() + if err != nil || opCode != TextMessage { + t.Errorf("%s: NextReader() returned %d, r, %v", name, opCode, err) + continue + } + + t.Logf("frame size: %d", n) + rbuf, err := io.ReadAll(r) + if err != nil { + t.Errorf("%s: ReadFull() returned rbuf, %v", name, err) + continue + } + + if len(rbuf) != n { + t.Errorf("%s: len(rbuf) is %d, want %d", name, len(rbuf), n) + continue + } + + for i, b := range rbuf { + if byte(i) != b { + t.Errorf("%s: bad byte at offset %d", name, i) + break + } + } + } + } + } + } + } +} + +func TestControl(t *testing.T) { + const message = "this is a ping/pong messsage" + for _, isServer := range []bool{true, false} { + for _, isWriteControl := range []bool{true, false} { + name := fmt.Sprintf("s:%v, wc:%v", isServer, isWriteControl) + var connBuf bytes.Buffer + wc := newTestConn(nil, &connBuf, isServer) + rc := newTestConn(&connBuf, nil, !isServer) + if isWriteControl { + _ = wc.WriteControl(PongMessage, []byte(message), time.Now().Add(time.Second)) + } else { + w, err := wc.NextWriter(PongMessage) + if err != nil { + t.Errorf("%s: wc.NextWriter() returned %v", name, err) + continue + } + if _, err := w.Write([]byte(message)); err != nil { + t.Errorf("%s: w.Write() returned %v", name, err) + continue + } + if err := w.Close(); err != nil { + t.Errorf("%s: w.Close() returned %v", name, err) + continue + } + var actualMessage string + rc.SetPongHandler(func(s string) error { actualMessage = s; return nil }) + _, _, _ = rc.NextReader() + if actualMessage != message { + t.Errorf("%s: pong=%q, want %q", name, actualMessage, message) + continue + } + } + } + } +} + +// simpleBufferPool is an implementation of BufferPool for TestWriteBufferPool. +type simpleBufferPool struct { + v interface{} +} + +func (p *simpleBufferPool) Get() interface{} { + v := p.v + p.v = nil + return v +} + +func (p *simpleBufferPool) Put(v interface{}) { + p.v = v +} + +func TestWriteBufferPool(t *testing.T) { + const message = "Now is the time for all good people to come to the aid of the party." + + var buf bytes.Buffer + var pool simpleBufferPool + rc := newTestConn(&buf, nil, false) + + // Specify writeBufferSize smaller than message size to ensure that pooling + // works with fragmented messages. + wc := newConn(fakeNetConn{Writer: &buf}, true, 1024, len(message)-1, &pool, nil, nil) + + if wc.writeBuf != nil { + t.Fatal("writeBuf not nil after create") + } + + // Part 1: test NextWriter/Write/Close + + w, err := wc.NextWriter(TextMessage) + if err != nil { + t.Fatalf("wc.NextWriter() returned %v", err) + } + + if wc.writeBuf == nil { + t.Fatal("writeBuf is nil after NextWriter") + } + + writeBufAddr := &wc.writeBuf[0] + + if _, err := io.WriteString(w, message); err != nil { + t.Fatalf("io.WriteString(w, message) returned %v", err) + } + + if err := w.Close(); err != nil { + t.Fatalf("w.Close() returned %v", err) + } + + if wc.writeBuf != nil { + t.Fatal("writeBuf not nil after w.Close()") + } + + if wpd, ok := pool.v.(writePoolData); !ok || len(wpd.buf) == 0 || &wpd.buf[0] != writeBufAddr { + t.Fatal("writeBuf not returned to pool") + } + + opCode, p, err := rc.ReadMessage() + if opCode != TextMessage || err != nil { + t.Fatalf("ReadMessage() returned %d, p, %v", opCode, err) + } + + if s := string(p); s != message { + t.Fatalf("message is %s, want %s", s, message) + } + + // Part 2: Test WriteMessage. + + if err := wc.WriteMessage(TextMessage, []byte(message)); err != nil { + t.Fatalf("wc.WriteMessage() returned %v", err) + } + + if wc.writeBuf != nil { + t.Fatal("writeBuf not nil after wc.WriteMessage()") + } + + if wpd, ok := pool.v.(writePoolData); !ok || len(wpd.buf) == 0 || &wpd.buf[0] != writeBufAddr { + t.Fatal("writeBuf not returned to pool after WriteMessage") + } + + opCode, p, err = rc.ReadMessage() + if opCode != TextMessage || err != nil { + t.Fatalf("ReadMessage() returned %d, p, %v", opCode, err) + } + + if s := string(p); s != message { + t.Fatalf("message is %s, want %s", s, message) + } +} + +// TestWriteBufferPoolSync ensures that *sync.Pool works as a buffer pool. +func TestWriteBufferPoolSync(t *testing.T) { + var buf bytes.Buffer + var pool sync.Pool + wc := newConn(fakeNetConn{Writer: &buf}, true, 1024, 1024, &pool, nil, nil) + rc := newTestConn(&buf, nil, false) + + const message = "Hello World!" + for i := 0; i < 3; i++ { + if err := wc.WriteMessage(TextMessage, []byte(message)); err != nil { + t.Fatalf("wc.WriteMessage() returned %v", err) + } + opCode, p, err := rc.ReadMessage() + if opCode != TextMessage || err != nil { + t.Fatalf("ReadMessage() returned %d, p, %v", opCode, err) + } + if s := string(p); s != message { + t.Fatalf("message is %s, want %s", s, message) + } + } +} + +// errorWriter is an io.Writer than returns an error on all writes. +type errorWriter struct{} + +func (ew errorWriter) Write(_ []byte) (int, error) { return 0, errors.New("error") } + +// TestWriteBufferPoolError ensures that buffer is returned to pool after error +// on write. +func TestWriteBufferPoolError(t *testing.T) { + + // Part 1: Test NextWriter/Write/Close + + var pool simpleBufferPool + wc := newConn(fakeNetConn{Writer: errorWriter{}}, true, 1024, 1024, &pool, nil, nil) + + w, err := wc.NextWriter(TextMessage) + if err != nil { + t.Fatalf("wc.NextWriter() returned %v", err) + } + + if wc.writeBuf == nil { + t.Fatal("writeBuf is nil after NextWriter") + } + + writeBufAddr := &wc.writeBuf[0] + + if _, err := io.WriteString(w, "Hello"); err != nil { + t.Fatalf("io.WriteString(w, message) returned %v", err) + } + + if err := w.Close(); err == nil { + t.Fatalf("w.Close() did not return error") + } + + if wpd, ok := pool.v.(writePoolData); !ok || len(wpd.buf) == 0 || &wpd.buf[0] != writeBufAddr { + t.Fatal("writeBuf not returned to pool") + } + + // Part 2: Test WriteMessage + + wc = newConn(fakeNetConn{Writer: errorWriter{}}, true, 1024, 1024, &pool, nil, nil) + + if err := wc.WriteMessage(TextMessage, []byte("Hello")); err == nil { + t.Fatalf("wc.WriteMessage did not return error") + } + + if wpd, ok := pool.v.(writePoolData); !ok || len(wpd.buf) == 0 || &wpd.buf[0] != writeBufAddr { + t.Fatal("writeBuf not returned to pool") + } +} + +func TestCloseFrameBeforeFinalMessageFrame(t *testing.T) { + const bufSize = 512 + + expectedErr := &CloseError{Code: CloseNormalClosure, Text: "hello"} + + var b1, b2 bytes.Buffer + wc := newConn(&fakeNetConn{Reader: nil, Writer: &b1}, false, 1024, bufSize, nil, nil, nil) + rc := newTestConn(&b1, &b2, true) + + w, _ := wc.NextWriter(BinaryMessage) + _, _ = w.Write(make([]byte, bufSize+bufSize/2)) + _ = wc.WriteControl(CloseMessage, FormatCloseMessage(expectedErr.Code, expectedErr.Text), time.Now().Add(10*time.Second)) + _ = w.Close() + + op, r, err := rc.NextReader() + if op != BinaryMessage || err != nil { + t.Fatalf("NextReader() returned %d, %v", op, err) + } + _, err = io.Copy(io.Discard, r) + if !reflect.DeepEqual(err, expectedErr) { + t.Fatalf("io.Copy() returned %v, want %v", err, expectedErr) + } + _, _, err = rc.NextReader() + if !reflect.DeepEqual(err, expectedErr) { + t.Fatalf("NextReader() returned %v, want %v", err, expectedErr) + } +} + +func TestEOFWithinFrame(t *testing.T) { + const bufSize = 64 + + for n := 0; ; n++ { + var b bytes.Buffer + wc := newTestConn(nil, &b, false) + rc := newTestConn(&b, nil, true) + + w, _ := wc.NextWriter(BinaryMessage) + _, _ = w.Write(make([]byte, bufSize)) + _ = w.Close() + + if n >= b.Len() { + break + } + b.Truncate(n) + + op, r, err := rc.NextReader() + if err == errUnexpectedEOF { + continue + } + if op != BinaryMessage || err != nil { + t.Fatalf("%d: NextReader() returned %d, %v", n, op, err) + } + _, err = io.Copy(io.Discard, r) + if err != errUnexpectedEOF { + t.Fatalf("%d: io.Copy() returned %v, want %v", n, err, errUnexpectedEOF) + } + _, _, err = rc.NextReader() + if err != errUnexpectedEOF { + t.Fatalf("%d: NextReader() returned %v, want %v", n, err, errUnexpectedEOF) + } + } +} + +func TestEOFBeforeFinalFrame(t *testing.T) { + const bufSize = 512 + + var b1, b2 bytes.Buffer + wc := newConn(&fakeNetConn{Writer: &b1}, false, 1024, bufSize, nil, nil, nil) + rc := newTestConn(&b1, &b2, true) + + w, _ := wc.NextWriter(BinaryMessage) + _, _ = w.Write(make([]byte, bufSize+bufSize/2)) + + op, r, err := rc.NextReader() + if op != BinaryMessage || err != nil { + t.Fatalf("NextReader() returned %d, %v", op, err) + } + _, err = io.Copy(io.Discard, r) + if err != errUnexpectedEOF { + t.Fatalf("io.Copy() returned %v, want %v", err, errUnexpectedEOF) + } + _, _, err = rc.NextReader() + if err != errUnexpectedEOF { + t.Fatalf("NextReader() returned %v, want %v", err, errUnexpectedEOF) + } +} + +func TestWriteAfterMessageWriterClose(t *testing.T) { + wc := newTestConn(nil, &bytes.Buffer{}, false) + w, _ := wc.NextWriter(BinaryMessage) + _, _ = io.WriteString(w, "hello") + if err := w.Close(); err != nil { + t.Fatalf("unxpected error closing message writer, %v", err) + } + + if _, err := io.WriteString(w, "world"); err == nil { + t.Fatalf("no error writing after close") + } + + w, _ = wc.NextWriter(BinaryMessage) + _, _ = io.WriteString(w, "hello") + + // close w by getting next writer + _, err := wc.NextWriter(BinaryMessage) + if err != nil { + t.Fatalf("unexpected error getting next writer, %v", err) + } + + if _, err := io.WriteString(w, "world"); err == nil { + t.Fatalf("no error writing after close") + } +} + +func TestReadLimit(t *testing.T) { + t.Run("Test ReadLimit is enforced", func(t *testing.T) { + const readLimit = 512 + message := make([]byte, readLimit+1) + + var b1, b2 bytes.Buffer + wc := newConn(&fakeNetConn{Writer: &b1}, false, 1024, readLimit-2, nil, nil, nil) + rc := newTestConn(&b1, &b2, true) + rc.SetReadLimit(readLimit) + + // Send message at the limit with interleaved pong. + w, _ := wc.NextWriter(BinaryMessage) + _, _ = w.Write(message[:readLimit-1]) + _ = wc.WriteControl(PongMessage, []byte("this is a pong"), time.Now().Add(10*time.Second)) + _, _ = w.Write(message[:1]) + _ = w.Close() + + // Send message larger than the limit. + _ = wc.WriteMessage(BinaryMessage, message[:readLimit+1]) + + op, _, err := rc.NextReader() + if op != BinaryMessage || err != nil { + t.Fatalf("1: NextReader() returned %d, %v", op, err) + } + op, r, err := rc.NextReader() + if op != BinaryMessage || err != nil { + t.Fatalf("2: NextReader() returned %d, %v", op, err) + } + _, err = io.Copy(io.Discard, r) + if err != ErrReadLimit { + t.Fatalf("io.Copy() returned %v", err) + } + }) + + t.Run("Test that ReadLimit cannot be overflowed", func(t *testing.T) { + const readLimit = 1 + + var b1, b2 bytes.Buffer + rc := newTestConn(&b1, &b2, true) + rc.SetReadLimit(readLimit) + + // First, send a non-final binary message + b1.Write([]byte("\x02\x81")) + + // Mask key + b1.Write([]byte("\x00\x00\x00\x00")) + + // First payload + b1.Write([]byte("A")) + + // Next, send a negative-length, non-final continuation frame + b1.Write([]byte("\x00\xFF\x80\x00\x00\x00\x00\x00\x00\x00")) + + // Mask key + b1.Write([]byte("\x00\x00\x00\x00")) + + // Next, send a too long, final continuation frame + b1.Write([]byte("\x80\xFF\x00\x00\x00\x00\x00\x00\x00\x05")) + + // Mask key + b1.Write([]byte("\x00\x00\x00\x00")) + + // Too-long payload + b1.Write([]byte("BCDEF")) + + op, r, err := rc.NextReader() + if op != BinaryMessage || err != nil { + t.Fatalf("1: NextReader() returned %d, %v", op, err) + } + + var buf [10]byte + var read int + n, err := r.Read(buf[:]) + if err != nil && err != ErrReadLimit { + t.Fatalf("unexpected error testing read limit: %v", err) + } + read += n + + n, err = r.Read(buf[:]) + if err != nil && err != ErrReadLimit { + t.Fatalf("unexpected error testing read limit: %v", err) + } + read += n + + if err == nil && read > readLimit { + t.Fatalf("read limit exceeded: limit %d, read %d", readLimit, read) + } + }) +} + +func TestAddrs(t *testing.T) { + c := newTestConn(nil, nil, true) + if c.LocalAddr() != localAddr { + t.Errorf("LocalAddr = %v, want %v", c.LocalAddr(), localAddr) + } + if c.RemoteAddr() != remoteAddr { + t.Errorf("RemoteAddr = %v, want %v", c.RemoteAddr(), remoteAddr) + } +} + +func TestNetConn(t *testing.T) { + var b1, b2 bytes.Buffer + fc := fakeNetConn{Reader: &b1, Writer: &b2} + c := newConn(fc, true, 1024, 1024, nil, nil, nil) + ul := c.NetConn() + if ul != fc { + t.Fatalf("Underlying conn is not what it should be.") + } +} + +func TestBufioReadBytes(t *testing.T) { + // Test calling bufio.ReadBytes for value longer than read buffer size. + + m := make([]byte, 512) + m[len(m)-1] = '\n' + + var b1, b2 bytes.Buffer + wc := newConn(fakeNetConn{Writer: &b1}, false, len(m)+64, len(m)+64, nil, nil, nil) + rc := newConn(fakeNetConn{Reader: &b1, Writer: &b2}, true, len(m)-64, len(m)-64, nil, nil, nil) + + w, _ := wc.NextWriter(BinaryMessage) + _, _ = w.Write(m) + _ = w.Close() + + op, r, err := rc.NextReader() + if op != BinaryMessage || err != nil { + t.Fatalf("NextReader() returned %d, %v", op, err) + } + + br := bufio.NewReader(r) + p, err := br.ReadBytes('\n') + if err != nil { + t.Fatalf("ReadBytes() returned %v", err) + } + if len(p) != len(m) { + t.Fatalf("read returned %d bytes, want %d bytes", len(p), len(m)) + } +} + +var closeErrorTests = []struct { + err error + codes []int + ok bool +}{ + {&CloseError{Code: CloseNormalClosure}, []int{CloseNormalClosure}, true}, + {&CloseError{Code: CloseNormalClosure}, []int{CloseNoStatusReceived}, false}, + {&CloseError{Code: CloseNormalClosure}, []int{CloseNoStatusReceived, CloseNormalClosure}, true}, + {errors.New("hello"), []int{CloseNormalClosure}, false}, +} + +func TestCloseError(t *testing.T) { + for _, tt := range closeErrorTests { + ok := IsCloseError(tt.err, tt.codes...) + if ok != tt.ok { + t.Errorf("IsCloseError(%#v, %#v) returned %v, want %v", tt.err, tt.codes, ok, tt.ok) + } + } +} + +var unexpectedCloseErrorTests = []struct { + err error + codes []int + ok bool +}{ + {&CloseError{Code: CloseNormalClosure}, []int{CloseNormalClosure}, false}, + {&CloseError{Code: CloseNormalClosure}, []int{CloseNoStatusReceived}, true}, + {&CloseError{Code: CloseNormalClosure}, []int{CloseNoStatusReceived, CloseNormalClosure}, false}, + {errors.New("hello"), []int{CloseNormalClosure}, false}, +} + +func TestUnexpectedCloseErrors(t *testing.T) { + for _, tt := range unexpectedCloseErrorTests { + ok := IsUnexpectedCloseError(tt.err, tt.codes...) + if ok != tt.ok { + t.Errorf("IsUnexpectedCloseError(%#v, %#v) returned %v, want %v", tt.err, tt.codes, ok, tt.ok) + } + } +} diff --git a/internal/websocket/doc.go b/internal/websocket/doc.go new file mode 100644 index 00000000..4c39802f --- /dev/null +++ b/internal/websocket/doc.go @@ -0,0 +1,6 @@ +// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// Package websocket implements the WebSocket protocol defined in RFC 6455. +package websocket diff --git a/internal/websocket/examples/autobahn/README.md b/internal/websocket/examples/autobahn/README.md new file mode 100644 index 00000000..cc954fea --- /dev/null +++ b/internal/websocket/examples/autobahn/README.md @@ -0,0 +1,18 @@ +# Test Server + +This package contains a server for the [Autobahn WebSockets Test Suite](https://github.com/crossbario/autobahn-testsuite). + +To test the server, run + + go run server.go + +and start the client test driver + + mkdir -p reports + docker run -it --rm \ + -v ${PWD}/config:/config \ + -v ${PWD}/reports:/reports \ + crossbario/autobahn-testsuite \ + wstest -m fuzzingclient -s /config/fuzzingclient.json + +When the client completes, it writes a report to reports/index.html. diff --git a/internal/websocket/examples/autobahn/config/fuzzingclient.json b/internal/websocket/examples/autobahn/config/fuzzingclient.json new file mode 100644 index 00000000..eda4e663 --- /dev/null +++ b/internal/websocket/examples/autobahn/config/fuzzingclient.json @@ -0,0 +1,29 @@ +{ + "cases": ["*"], + "exclude-cases": [], + "exclude-agent-cases": {}, + "outdir": "/reports", + "options": {"failByDrop": false}, + "servers": [ + { + "agent": "ReadAllWriteMessage", + "url": "ws://host.docker.internal:9000/m" + }, + { + "agent": "ReadAllWritePreparedMessage", + "url": "ws://host.docker.internal:9000/p" + }, + { + "agent": "CopyFull", + "url": "ws://host.docker.internal:9000/f" + }, + { + "agent": "ReadAllWrite", + "url": "ws://host.docker.internal:9000/r" + }, + { + "agent": "CopyWriterOnly", + "url": "ws://host.docker.internal:9000/c" + } + ] +} diff --git a/internal/websocket/examples/autobahn/server.go b/internal/websocket/examples/autobahn/server.go new file mode 100644 index 00000000..5361902e --- /dev/null +++ b/internal/websocket/examples/autobahn/server.go @@ -0,0 +1,269 @@ +// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// Command server is a test server for the Autobahn WebSockets Test Suite. +package main + +import ( + "errors" + "flag" + "io" + "log" + "net/http" + "time" + "unicode/utf8" + + "github.com/centrifugal/centrifuge/internal/websocket" +) + +var upgrader = websocket.Upgrader{ + ReadBufferSize: 4096, + WriteBufferSize: 4096, + EnableCompression: true, + CheckOrigin: func(r *http.Request) bool { + return true + }, +} + +// echoCopy echoes messages from the client using io.Copy. +func echoCopy(w http.ResponseWriter, r *http.Request, writerOnly bool) { + conn, _, err := upgrader.Upgrade(w, r, nil) + if err != nil { + log.Println("Upgrade:", err) + return + } + defer func() { + _ = conn.Close() + }() + for { + mt, r, err := conn.NextReader() + if err != nil { + if err != io.EOF { + log.Println("NextReader:", err) + } + return + } + if mt == websocket.TextMessage { + r = &validator{r: r} + } + w, err := conn.NextWriter(mt) + if err != nil { + log.Println("NextWriter:", err) + return + } + if mt == websocket.TextMessage { + r = &validator{r: r} + } + if writerOnly { + _, err = io.Copy(struct{ io.Writer }{w}, r) + } else { + _, err = io.Copy(w, r) + } + if err != nil { + if err == errInvalidUTF8 { + _ = conn.WriteControl(websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseInvalidFramePayloadData, ""), + time.Time{}) + } + log.Println("Copy:", err) + return + } + err = w.Close() + if err != nil { + log.Println("Close:", err) + return + } + } +} + +func echoCopyWriterOnly(w http.ResponseWriter, r *http.Request) { + echoCopy(w, r, true) +} + +func echoCopyFull(w http.ResponseWriter, r *http.Request) { + echoCopy(w, r, false) +} + +// echoReadAll echoes messages from the client by reading the entire message +// with ioutil.ReadAll. +func echoReadAll(w http.ResponseWriter, r *http.Request, writeMessage, writePrepared bool) { + conn, _, err := upgrader.Upgrade(w, r, nil) + if err != nil { + log.Println("Upgrade:", err) + return + } + defer func() { + _ = conn.Close() + }() + for { + mt, b, err := conn.ReadMessage() + if err != nil { + if err != io.EOF { + log.Println("NextReader:", err) + } + return + } + if mt == websocket.TextMessage { + if !utf8.Valid(b) { + _ = conn.WriteControl(websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseInvalidFramePayloadData, ""), + time.Time{}) + log.Println("ReadAll: invalid utf8") + } + } + if writeMessage { + if !writePrepared { + err = conn.WriteMessage(mt, b) + if err != nil { + log.Println("WriteMessage:", err) + } + } else { + pm, err := websocket.NewPreparedMessage(mt, b) + if err != nil { + log.Println("NewPreparedMessage:", err) + return + } + err = conn.WritePreparedMessage(pm) + if err != nil { + log.Println("WritePreparedMessage:", err) + } + } + } else { + w, err := conn.NextWriter(mt) + if err != nil { + log.Println("NextWriter:", err) + return + } + if _, err := w.Write(b); err != nil { + log.Println("Writer:", err) + return + } + if err := w.Close(); err != nil { + log.Println("Close:", err) + return + } + } + } +} + +func echoReadAllWriter(w http.ResponseWriter, r *http.Request) { + echoReadAll(w, r, false, false) +} + +func echoReadAllWriteMessage(w http.ResponseWriter, r *http.Request) { + echoReadAll(w, r, true, false) +} + +func echoReadAllWritePreparedMessage(w http.ResponseWriter, r *http.Request) { + echoReadAll(w, r, true, true) +} + +func serveHome(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/" { + http.Error(w, "Not found.", http.StatusNotFound) + return + } + if r.Method != http.MethodGet { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + w.Header().Set("Content-Type", "text/html; charset=utf-8") + _, _ = io.WriteString(w, "Echo Server") +} + +var addr = flag.String("addr", ":9000", "http service address") + +func main() { + flag.Parse() + http.HandleFunc("/", serveHome) + http.HandleFunc("/c", echoCopyWriterOnly) + http.HandleFunc("/f", echoCopyFull) + http.HandleFunc("/r", echoReadAllWriter) + http.HandleFunc("/m", echoReadAllWriteMessage) + http.HandleFunc("/p", echoReadAllWritePreparedMessage) + err := http.ListenAndServe(*addr, nil) + if err != nil { + log.Fatal("ListenAndServe: ", err) + } +} + +type validator struct { + state int + x rune + r io.Reader +} + +var errInvalidUTF8 = errors.New("invalid utf8") + +func (r *validator) Read(p []byte) (int, error) { + n, err := r.r.Read(p) + state := r.state + x := r.x + for _, b := range p[:n] { + state, x = decode(state, x, b) + if state == utf8Reject { + break + } + } + r.state = state + r.x = x + if state == utf8Reject || (err == io.EOF && state != utf8Accept) { + return n, errInvalidUTF8 + } + return n, err +} + +// UTF-8 decoder from http://bjoern.hoehrmann.de/utf-8/decoder/dfa/ +// +// Copyright (c) 2008-2009 Bjoern Hoehrmann +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to +// deal in the Software without restriction, including without limitation the +// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or +// sell copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in +// all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS +// IN THE SOFTWARE. +var utf8d = [...]byte{ + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 00..1f + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 20..3f + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 40..5f + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 60..7f + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, // 80..9f + 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, // a0..bf + 8, 8, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, // c0..df + 0xa, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x4, 0x3, 0x3, // e0..ef + 0xb, 0x6, 0x6, 0x6, 0x5, 0x8, 0x8, 0x8, 0x8, 0x8, 0x8, 0x8, 0x8, 0x8, 0x8, 0x8, // f0..ff + 0x0, 0x1, 0x2, 0x3, 0x5, 0x8, 0x7, 0x1, 0x1, 0x1, 0x4, 0x6, 0x1, 0x1, 0x1, 0x1, // s0..s0 + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, // s1..s2 + 1, 2, 1, 1, 1, 1, 1, 2, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, // s3..s4 + 1, 2, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3, 1, 3, 1, 1, 1, 1, 1, 1, // s5..s6 + 1, 3, 1, 1, 1, 1, 1, 3, 1, 3, 1, 1, 1, 1, 1, 1, 1, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, // s7..s8 +} + +const ( + utf8Accept = 0 + utf8Reject = 1 +) + +func decode(state int, x rune, b byte) (int, rune) { + t := utf8d[b] + if state != utf8Accept { + x = rune(b&0x3f) | (x << 6) + } else { + x = rune((0xff >> t) & b) + } + state = int(utf8d[256+state*16+int(t)]) + return state, x +} diff --git a/internal/websocket/mask.go b/internal/websocket/mask.go new file mode 100644 index 00000000..6e5004fb --- /dev/null +++ b/internal/websocket/mask.go @@ -0,0 +1,54 @@ +// Copyright 2016 The Gorilla WebSocket Authors. All rights reserved. Use of +// this source code is governed by a BSD-style license that can be found in the +// LICENSE file. + +//go:build !appengine + +package websocket + +import "unsafe" + +const wordSize = int(unsafe.Sizeof(uintptr(0))) + +func maskBytes(key [4]byte, pos int, b []byte) int { + // Mask one byte at a time for small buffers. + if len(b) < 2*wordSize { + for i := range b { + b[i] ^= key[pos&3] + pos++ + } + return pos & 3 + } + + // Mask one byte at a time to word boundary. + if n := int(uintptr(unsafe.Pointer(&b[0]))) % wordSize; n != 0 { + n = wordSize - n + for i := range b[:n] { + b[i] ^= key[pos&3] + pos++ + } + b = b[n:] + } + + // Create aligned word size key. + var k [wordSize]byte + for i := range k { + k[i] = key[(pos+i)&3] + } + kw := *(*uintptr)(unsafe.Pointer(&k)) + + // Mask one word at a time. + n := (len(b) / wordSize) * wordSize + for i := 0; i < n; i += wordSize { + *(*uintptr)(unsafe.Pointer(uintptr(unsafe.Pointer(&b[0])) + uintptr(i))) ^= kw + } + + // Mask one byte at a time for remaining bytes. + b = b[n:] + for i := range b { + b[i] ^= key[pos&3] + pos++ + } + + return pos & 3 +} diff --git a/internal/websocket/mask_safe.go b/internal/websocket/mask_safe.go new file mode 100644 index 00000000..66706daf --- /dev/null +++ b/internal/websocket/mask_safe.go @@ -0,0 +1,15 @@ +// Copyright 2016 The Gorilla WebSocket Authors. All rights reserved. Use of +// this source code is governed by a BSD-style license that can be found in the +// LICENSE file. + +//go:build appengine + +package websocket + +func maskBytes(key [4]byte, pos int, b []byte) int { + for i := range b { + b[i] ^= key[pos&3] + pos++ + } + return pos & 3 +} diff --git a/internal/websocket/mask_test.go b/internal/websocket/mask_test.go new file mode 100644 index 00000000..6389f436 --- /dev/null +++ b/internal/websocket/mask_test.go @@ -0,0 +1,72 @@ +// Copyright 2016 The Gorilla WebSocket Authors. All rights reserved. Use of +// this source code is governed by a BSD-style license that can be found in the +// LICENSE file. + +// !appengine + +package websocket + +import ( + "fmt" + "testing" +) + +func maskBytesByByte(key [4]byte, pos int, b []byte) int { + for i := range b { + b[i] ^= key[pos&3] + pos++ + } + return pos & 3 +} + +func notzero(b []byte) int { + for i := range b { + if b[i] != 0 { + return i + } + } + return -1 +} + +func TestMaskBytes(t *testing.T) { + key := [4]byte{1, 2, 3, 4} + for size := 1; size <= 1024; size++ { + for align := 0; align < wordSize; align++ { + for pos := 0; pos < 4; pos++ { + b := make([]byte, size+align)[align:] + maskBytes(key, pos, b) + maskBytesByByte(key, pos, b) + if i := notzero(b); i >= 0 { + t.Errorf("size:%d, align:%d, pos:%d, offset:%d", size, align, pos, i) + } + } + } + } +} + +func BenchmarkMaskBytes(b *testing.B) { + for _, size := range []int{2, 4, 8, 16, 32, 512, 1024} { + b.Run(fmt.Sprintf("size-%d", size), func(b *testing.B) { + for _, align := range []int{wordSize / 2} { + b.Run(fmt.Sprintf("align-%d", align), func(b *testing.B) { + for _, fn := range []struct { + name string + fn func(key [4]byte, pos int, b []byte) int + }{ + {"byte", maskBytesByByte}, + {"word", maskBytes}, + } { + b.Run(fn.name, func(b *testing.B) { + key := newMaskKey() + data := make([]byte, size+align)[align:] + for i := 0; i < b.N; i++ { + fn.fn(key, 0, data) + } + b.SetBytes(int64(len(data))) + }) + } + }) + } + }) + } +} diff --git a/internal/websocket/prepared.go b/internal/websocket/prepared.go new file mode 100644 index 00000000..eb40158a --- /dev/null +++ b/internal/websocket/prepared.go @@ -0,0 +1,102 @@ +// Copyright 2017 The Gorilla WebSocket Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package websocket + +import ( + "bytes" + "net" + "sync" + "time" +) + +// PreparedMessage caches on the wire representations of a message payload. +// Use PreparedMessage to efficiently send a message payload to multiple +// connections. PreparedMessage is especially useful when compression is used +// because the CPU and memory expensive compression operation can be executed +// once for a given set of compression options. +type PreparedMessage struct { + messageType int + data []byte + mu sync.Mutex + frames map[prepareKey]*preparedFrame +} + +// prepareKey defines a unique set of options to cache prepared frames in PreparedMessage. +type prepareKey struct { + isServer bool + compress bool + compressionLevel int +} + +// preparedFrame contains data in wire representation. +type preparedFrame struct { + once sync.Once + data []byte +} + +// NewPreparedMessage returns an initialized PreparedMessage. You can then send +// it to connection using WritePreparedMessage method. Valid wire +// representation will be calculated lazily only once for a set of current +// connection options. +func NewPreparedMessage(messageType int, data []byte) (*PreparedMessage, error) { + pm := &PreparedMessage{ + messageType: messageType, + frames: make(map[prepareKey]*preparedFrame), + data: data, + } + + // Prepare a plain server frame. + _, frameData, err := pm.frame(prepareKey{isServer: true, compress: false}) + if err != nil { + return nil, err + } + + // To protect against caller modifying the data argument, remember the data + // copied to the plain server frame. + pm.data = frameData[len(frameData)-len(data):] + return pm, nil +} + +func (pm *PreparedMessage) frame(key prepareKey) (int, []byte, error) { + pm.mu.Lock() + frame, ok := pm.frames[key] + if !ok { + frame = &preparedFrame{} + pm.frames[key] = frame + } + pm.mu.Unlock() + + var err error + frame.once.Do(func() { + // Prepare a frame using a 'fake' connection. + // TODO: Refactor code in conn.go to allow more direct construction of + // the frame. + mu := make(chan struct{}, 1) + mu <- struct{}{} + var nc prepareConn + c := &Conn{ + conn: &nc, + mu: mu, + isServer: key.isServer, + compressionLevel: key.compressionLevel, + enableWriteCompression: true, + writeBuf: make([]byte, defaultWriteBufferSize+maxFrameHeaderSize), + } + if key.compress { + c.newCompressionWriter = compressNoContextTakeover + } + err = c.WriteMessage(pm.messageType, pm.data) + frame.data = nc.buf.Bytes() + }) + return pm.messageType, frame.data, err +} + +type prepareConn struct { + buf bytes.Buffer + net.Conn +} + +func (pc *prepareConn) Write(p []byte) (int, error) { return pc.buf.Write(p) } +func (pc *prepareConn) SetWriteDeadline(_ time.Time) error { return nil } diff --git a/internal/websocket/prepared_test.go b/internal/websocket/prepared_test.go new file mode 100644 index 00000000..c15d9d5a --- /dev/null +++ b/internal/websocket/prepared_test.go @@ -0,0 +1,74 @@ +// Copyright 2017 The Gorilla WebSocket Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package websocket + +import ( + "bytes" + "compress/flate" + "math/rand" + "testing" +) + +var preparedMessageTests = []struct { + messageType int + isServer bool + enableWriteCompression bool + compressionLevel int +}{ + // Server + {TextMessage, true, false, flate.BestSpeed}, + {TextMessage, true, true, flate.BestSpeed}, + {TextMessage, true, true, flate.BestCompression}, + {PingMessage, true, false, flate.BestSpeed}, + {PingMessage, true, true, flate.BestSpeed}, + + // Client + {TextMessage, false, false, flate.BestSpeed}, + {TextMessage, false, true, flate.BestSpeed}, + {TextMessage, false, true, flate.BestCompression}, + {PingMessage, false, false, flate.BestSpeed}, + {PingMessage, false, true, flate.BestSpeed}, +} + +func TestPreparedMessage(t *testing.T) { + for _, tt := range preparedMessageTests { + var data = []byte("this is a test") + var buf bytes.Buffer + c := newTestConn(nil, &buf, tt.isServer) + if tt.enableWriteCompression { + c.newCompressionWriter = compressNoContextTakeover + } + _ = c.SetCompressionLevel(tt.compressionLevel) + + // Seed random number generator for consistent frame mask. + rand.Seed(1234) //nolint:staticcheck + + if err := c.WriteMessage(tt.messageType, data); err != nil { + t.Fatal(err) + } + want := buf.String() + + pm, err := NewPreparedMessage(tt.messageType, data) + if err != nil { + t.Fatal(err) + } + + // Scribble on data to ensure that NewPreparedMessage takes a snapshot. + copy(data, "hello world") + + // Seed random number generator for consistent frame mask. + rand.Seed(1234) //nolint:staticcheck + + buf.Reset() + if err := c.WritePreparedMessage(pm); err != nil { + t.Fatal(err) + } + got := buf.String() + + if got != want { + t.Errorf("write message != prepared message for %+v", tt) + } + } +} diff --git a/internal/websocket/server.go b/internal/websocket/server.go new file mode 100644 index 00000000..d6165e15 --- /dev/null +++ b/internal/websocket/server.go @@ -0,0 +1,311 @@ +// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package websocket + +import ( + "bufio" + "errors" + "io" + "net/http" + "net/url" + "strings" + "time" +) + +// HandshakeError describes an error with the handshake from the peer. +type HandshakeError struct { + message string +} + +func (e HandshakeError) Error() string { return e.message } + +// Upgrader specifies parameters for upgrading an HTTP connection to a +// WebSocket connection. +// +// It is safe to call Upgrader's methods concurrently. +type Upgrader struct { + // HandshakeTimeout specifies the duration for the handshake to complete. + HandshakeTimeout time.Duration + + // ReadBufferSize and WriteBufferSize specify I/O buffer sizes in bytes. If a buffer + // size is zero, then buffers allocated by the HTTP server are used. The + // I/O buffer sizes do not limit the size of the messages that can be sent + // or received. + // The default value is 4096 bytes, 4kb. + ReadBufferSize, WriteBufferSize int + + // WriteBufferPool is a pool of buffers for write operations. If the value + // is not set, then write buffers are allocated to the connection for the + // lifetime of the connection. + // + // A pool is most useful when the application has a modest volume of writes + // across a large number of connections. + // + // Applications should use a single pool for each unique value of + // WriteBufferSize. + WriteBufferPool BufferPool + + // Subprotocols specifies the server's supported protocols in order of + // preference. If this field is not nil, then the Upgrade method negotiates a + // subprotocol by selecting the first match in this list with a protocol + // requested by the client. If there's no match, then no protocol is + // negotiated (the Sec-Websocket-Protocol header is not included in the + // handshake response). + Subprotocols []string + + // Error specifies the function for generating HTTP error responses. If Error + // is nil, then http.Error is used to generate the HTTP response. + Error func(w http.ResponseWriter, r *http.Request, status int, reason error) + + // CheckOrigin returns true if the request Origin header is acceptable. If + // CheckOrigin is nil, then a safe default is used: return false if the + // Origin request header is present and the origin host is not equal to + // request Host header. + // + // A CheckOrigin function should carefully validate the request origin to + // prevent cross-site request forgery. + CheckOrigin func(r *http.Request) bool + + // EnableCompression specify if the server should attempt to negotiate per + // message compression (RFC 7692). Setting this value to true does not + // guarantee that compression will be supported. Currently only "no context + // takeover" modes are supported. + EnableCompression bool +} + +func (u *Upgrader) returnError(w http.ResponseWriter, r *http.Request, status int, reason string) (*Conn, string, error) { + err := HandshakeError{reason} + if u.Error != nil { + u.Error(w, r, status, err) + } else { + w.Header().Set("Sec-Websocket-Version", "13") + http.Error(w, http.StatusText(status), status) + } + return nil, "", err +} + +// checkSameOrigin returns true if the origin is not set or is equal to the request host. +func checkSameOrigin(r *http.Request) bool { + origin := r.Header["Origin"] + if len(origin) == 0 { + return true + } + u, err := url.Parse(origin[0]) + if err != nil { + return false + } + return equalASCIIFold(u.Host, r.Host) +} + +func (u *Upgrader) selectSubprotocol(r *http.Request, responseHeader http.Header) string { + if u.Subprotocols != nil { + clientProtocols := Subprotocols(r) + for _, clientProtocol := range clientProtocols { + for _, serverProtocol := range u.Subprotocols { + if clientProtocol == serverProtocol { + return clientProtocol + } + } + } + } else if responseHeader != nil { + return responseHeader.Get("Sec-Websocket-Protocol") + } + return "" +} + +// Upgrade upgrades the HTTP server connection to the WebSocket protocol. +// +// The responseHeader is included in the response to the client's upgrade +// request. Use the responseHeader to specify cookies (Set-Cookie). To specify +// subprotocols supported by the server, set Upgrader.Subprotocols directly. +// +// If the upgrade fails, then Upgrade replies to the client with an HTTP error +// response. +func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header) (*Conn, string, error) { + const badHandshake = "websocket: the client is not using the websocket protocol: " + + if !tokenListContainsValue(r.Header, "Connection", "upgrade") { + return u.returnError(w, r, http.StatusBadRequest, badHandshake+"'upgrade' token not found in 'Connection' header") + } + + if !tokenListContainsValue(r.Header, "Upgrade", "websocket") { + return u.returnError(w, r, http.StatusBadRequest, badHandshake+"'websocket' token not found in 'Upgrade' header") + } + + if r.Method != http.MethodGet { + return u.returnError(w, r, http.StatusMethodNotAllowed, badHandshake+"request method is not GET") + } + + if !tokenListContainsValue(r.Header, "Sec-Websocket-Version", "13") { + return u.returnError(w, r, http.StatusBadRequest, "websocket: unsupported version: 13 not found in 'Sec-Websocket-Version' header") + } + + if _, ok := responseHeader["Sec-Websocket-Extensions"]; ok { + return u.returnError(w, r, http.StatusInternalServerError, "websocket: application specific 'Sec-WebSocket-Extensions' headers are unsupported") + } + + checkOrigin := u.CheckOrigin + if checkOrigin == nil { + checkOrigin = checkSameOrigin + } + if !checkOrigin(r) { + return u.returnError(w, r, http.StatusForbidden, "websocket: request origin not allowed by Upgrader.CheckOrigin") + } + + challengeKey := r.Header.Get("Sec-Websocket-Key") + if !isValidChallengeKey(challengeKey) { + return u.returnError(w, r, http.StatusBadRequest, "websocket: not a websocket handshake: 'Sec-WebSocket-Key' header must be Base64 encoded value of 16-byte in length") + } + + subprotocol := u.selectSubprotocol(r, responseHeader) + + // Negotiate PMCE + var compress bool + if u.EnableCompression { + for _, ext := range parseExtensions(r.Header) { + if ext[""] != "permessage-deflate" { + continue + } + compress = true + break + } + } + + h, ok := w.(http.Hijacker) + if !ok { + return u.returnError(w, r, http.StatusInternalServerError, "websocket: response does not implement http.Hijacker") + } + + netConn, brw, err := h.Hijack() + if err != nil { + return u.returnError(w, r, http.StatusInternalServerError, err.Error()) + } + + if brw.Reader.Buffered() > 0 { + _ = netConn.Close() + return nil, "", errors.New("websocket: client sent data before handshake is complete") + } + + var br *bufio.Reader + if u.ReadBufferSize == 0 && brw.Reader.Size() > 256 { + // Reuse hijacked buffered reader as connection reader. + br = brw.Reader + } + + buf := bufioWriterBuffer(netConn, brw.Writer) + + var writeBuf []byte + if u.WriteBufferPool == nil && u.WriteBufferSize == 0 && len(buf) >= maxFrameHeaderSize+256 { + // Reuse hijacked write buffer as connection buffer. + writeBuf = buf + } + + c := newConn(netConn, true, u.ReadBufferSize, u.WriteBufferSize, u.WriteBufferPool, br, writeBuf) + + if compress { + c.newCompressionWriter = compressNoContextTakeover + c.newDecompressionReader = decompressNoContextTakeover + } + + // Use larger of hijacked buffer and connection write buffer for header. + p := buf + if len(c.writeBuf) > len(p) { + p = c.writeBuf + } + p = p[:0] + + p = append(p, "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: "...) + p = append(p, computeAcceptKey(challengeKey)...) + p = append(p, "\r\n"...) + if subprotocol != "" { + p = append(p, "Sec-WebSocket-Protocol: "...) + p = append(p, subprotocol...) + p = append(p, "\r\n"...) + } + if compress { + p = append(p, "Sec-WebSocket-Extensions: permessage-deflate; server_no_context_takeover; client_no_context_takeover\r\n"...) + } + for k, vs := range responseHeader { + if k == "Sec-Websocket-Protocol" { + continue + } + for _, v := range vs { + p = append(p, k...) + p = append(p, ": "...) + for i := 0; i < len(v); i++ { + b := v[i] + if b <= 31 { + // prevent response splitting. + b = ' ' + } + p = append(p, b) + } + p = append(p, "\r\n"...) + } + } + p = append(p, "\r\n"...) + + // Clear deadlines set by HTTP server. + _ = netConn.SetDeadline(time.Time{}) + + if u.HandshakeTimeout > 0 { + _ = netConn.SetWriteDeadline(time.Now().Add(u.HandshakeTimeout)) + } + if _, err = netConn.Write(p); err != nil { + _ = netConn.Close() + return nil, "", err + } + if u.HandshakeTimeout > 0 { + _ = netConn.SetWriteDeadline(time.Time{}) + } + + return c, subprotocol, nil +} + +// Subprotocols returns the subprotocols requested by the client in the +// Sec-Websocket-Protocol header. +func Subprotocols(r *http.Request) []string { + h := strings.TrimSpace(r.Header.Get("Sec-Websocket-Protocol")) + if h == "" { + return nil + } + protocols := strings.Split(h, ",") + for i := range protocols { + protocols[i] = strings.TrimSpace(protocols[i]) + } + return protocols +} + +// IsWebSocketUpgrade returns true if the client requested upgrade to the +// WebSocket protocol. +func IsWebSocketUpgrade(r *http.Request) bool { + return tokenListContainsValue(r.Header, "Connection", "upgrade") && + tokenListContainsValue(r.Header, "Upgrade", "websocket") +} + +// writeHook is an io.Writer that records the last slice passed to it vio +// io.Writer.Write. +type writeHook struct { + p []byte +} + +func (wh *writeHook) Write(p []byte) (int, error) { + wh.p = p + return len(p), nil +} + +// bufioWriterBuffer grabs the buffer from a bufio.Writer. +func bufioWriterBuffer(originalWriter io.Writer, bw *bufio.Writer) []byte { + // This code assumes that bufio.Writer.buf[:1] is passed to the + // bufio.Writer's underlying writer. + var wh writeHook + bw.Reset(&wh) + _ = bw.WriteByte(0) + _ = bw.Flush() + + bw.Reset(originalWriter) + + return wh.p[:cap(wh.p)] +} diff --git a/internal/websocket/server_test.go b/internal/websocket/server_test.go new file mode 100644 index 00000000..6562e003 --- /dev/null +++ b/internal/websocket/server_test.go @@ -0,0 +1,149 @@ +// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package websocket + +import ( + "bufio" + "bytes" + "net" + "net/http" + "reflect" + "strings" + "testing" +) + +var subprotocolTests = []struct { + h string + protocols []string +}{ + {"", nil}, + {"foo", []string{"foo"}}, + {"foo,bar", []string{"foo", "bar"}}, + {"foo, bar", []string{"foo", "bar"}}, + {" foo, bar", []string{"foo", "bar"}}, + {" foo, bar ", []string{"foo", "bar"}}, +} + +func TestSubprotocols(t *testing.T) { + for _, st := range subprotocolTests { + r := http.Request{Header: http.Header{"Sec-Websocket-Protocol": {st.h}}} + protocols := Subprotocols(&r) + if !reflect.DeepEqual(st.protocols, protocols) { + t.Errorf("SubProtocols(%q) returned %#v, want %#v", st.h, protocols, st.protocols) + } + } +} + +var isWebSocketUpgradeTests = []struct { + ok bool + h http.Header +}{ + {false, http.Header{"Upgrade": {"websocket"}}}, + {false, http.Header{"Connection": {"upgrade"}}}, + {true, http.Header{"Connection": {"upgRade"}, "Upgrade": {"WebSocket"}}}, +} + +func TestIsWebSocketUpgrade(t *testing.T) { + for _, tt := range isWebSocketUpgradeTests { + ok := IsWebSocketUpgrade(&http.Request{Header: tt.h}) + if tt.ok != ok { + t.Errorf("IsWebSocketUpgrade(%v) returned %v, want %v", tt.h, ok, tt.ok) + } + } +} + +func TestSubProtocolSelection(t *testing.T) { + upgrader := Upgrader{ + Subprotocols: []string{"foo", "bar", "baz"}, + } + + r := http.Request{Header: http.Header{"Sec-Websocket-Protocol": {"foo", "bar"}}} + s := upgrader.selectSubprotocol(&r, nil) + if s != "foo" { + t.Errorf("Upgrader.selectSubprotocol returned %v, want %v", s, "foo") + } + + r = http.Request{Header: http.Header{"Sec-Websocket-Protocol": {"bar", "foo"}}} + s = upgrader.selectSubprotocol(&r, nil) + if s != "bar" { + t.Errorf("Upgrader.selectSubprotocol returned %v, want %v", s, "bar") + } + + r = http.Request{Header: http.Header{"Sec-Websocket-Protocol": {"baz"}}} + s = upgrader.selectSubprotocol(&r, nil) + if s != "baz" { + t.Errorf("Upgrader.selectSubprotocol returned %v, want %v", s, "baz") + } + + r = http.Request{Header: http.Header{"Sec-Websocket-Protocol": {"quux"}}} + s = upgrader.selectSubprotocol(&r, nil) + if s != "" { + t.Errorf("Upgrader.selectSubprotocol returned %v, want %v", s, "empty string") + } +} + +var checkSameOriginTests = []struct { + ok bool + r *http.Request +}{ + {false, &http.Request{Host: "example.org", Header: map[string][]string{"Origin": {"https://other.org"}}}}, + {true, &http.Request{Host: "example.org", Header: map[string][]string{"Origin": {"https://example.org"}}}}, + {true, &http.Request{Host: "Example.org", Header: map[string][]string{"Origin": {"https://example.org"}}}}, +} + +func TestCheckSameOrigin(t *testing.T) { + for _, tt := range checkSameOriginTests { + ok := checkSameOrigin(tt.r) + if tt.ok != ok { + t.Errorf("checkSameOrigin(%+v) returned %v, want %v", tt.r, ok, tt.ok) + } + } +} + +type reuseTestResponseWriter struct { + brw *bufio.ReadWriter + http.ResponseWriter +} + +func (resp *reuseTestResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + return fakeNetConn{strings.NewReader(""), &bytes.Buffer{}}, resp.brw, nil +} + +var bufioReuseTests = []struct { + n int + reuse bool +}{ + {4096, true}, + {128, false}, +} + +func TestBufioReuse(t *testing.T) { + for i, tt := range bufioReuseTests { + br := bufio.NewReaderSize(strings.NewReader(""), tt.n) + bw := bufio.NewWriterSize(&bytes.Buffer{}, tt.n) + resp := &reuseTestResponseWriter{ + brw: bufio.NewReadWriter(br, bw), + } + upgrader := Upgrader{} + c, _, err := upgrader.Upgrade(resp, &http.Request{ + Method: http.MethodGet, + Header: http.Header{ + "Upgrade": []string{"websocket"}, + "Connection": []string{"upgrade"}, + "Sec-Websocket-Key": []string{"dGhlIHNhbXBsZSBub25jZQ=="}, + "Sec-Websocket-Version": []string{"13"}, + }}, nil) + if err != nil { + t.Fatal(err) + } + if reuse := c.br == br; reuse != tt.reuse { + t.Errorf("%d: buffered reader reuse=%v, want %v", i, reuse, tt.reuse) + } + writeBuf := bufioWriterBuffer(c.NetConn(), bw) + if reuse := &c.writeBuf[0] == &writeBuf[0]; reuse != tt.reuse { + t.Errorf("%d: write buffer reuse=%v, want %v", i, reuse, tt.reuse) + } + } +} diff --git a/internal/websocket/tls_handshake.go b/internal/websocket/tls_handshake.go new file mode 100644 index 00000000..7f386453 --- /dev/null +++ b/internal/websocket/tls_handshake.go @@ -0,0 +1,18 @@ +package websocket + +import ( + "context" + "crypto/tls" +) + +func doHandshake(ctx context.Context, tlsConn *tls.Conn, cfg *tls.Config) error { + if err := tlsConn.HandshakeContext(ctx); err != nil { + return err + } + if !cfg.InsecureSkipVerify { + if err := tlsConn.VerifyHostname(cfg.ServerName); err != nil { + return err + } + } + return nil +} diff --git a/internal/websocket/util.go b/internal/websocket/util.go new file mode 100644 index 00000000..31a5dee6 --- /dev/null +++ b/internal/websocket/util.go @@ -0,0 +1,298 @@ +// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package websocket + +import ( + "crypto/rand" + "crypto/sha1" + "encoding/base64" + "io" + "net/http" + "strings" + "unicode/utf8" +) + +var keyGUID = []byte("258EAFA5-E914-47DA-95CA-C5AB0DC85B11") + +func computeAcceptKey(challengeKey string) string { + h := sha1.New() + h.Write([]byte(challengeKey)) + h.Write(keyGUID) + return base64.StdEncoding.EncodeToString(h.Sum(nil)) +} + +func generateChallengeKey() (string, error) { + p := make([]byte, 16) + if _, err := io.ReadFull(rand.Reader, p); err != nil { + return "", err + } + return base64.StdEncoding.EncodeToString(p), nil +} + +// Token octets per RFC 2616. +var isTokenOctet = [256]bool{ + '!': true, + '#': true, + '$': true, + '%': true, + '&': true, + '\'': true, + '*': true, + '+': true, + '-': true, + '.': true, + '0': true, + '1': true, + '2': true, + '3': true, + '4': true, + '5': true, + '6': true, + '7': true, + '8': true, + '9': true, + 'A': true, + 'B': true, + 'C': true, + 'D': true, + 'E': true, + 'F': true, + 'G': true, + 'H': true, + 'I': true, + 'J': true, + 'K': true, + 'L': true, + 'M': true, + 'N': true, + 'O': true, + 'P': true, + 'Q': true, + 'R': true, + 'S': true, + 'T': true, + 'U': true, + 'W': true, + 'V': true, + 'X': true, + 'Y': true, + 'Z': true, + '^': true, + '_': true, + '`': true, + 'a': true, + 'b': true, + 'c': true, + 'd': true, + 'e': true, + 'f': true, + 'g': true, + 'h': true, + 'i': true, + 'j': true, + 'k': true, + 'l': true, + 'm': true, + 'n': true, + 'o': true, + 'p': true, + 'q': true, + 'r': true, + 's': true, + 't': true, + 'u': true, + 'v': true, + 'w': true, + 'x': true, + 'y': true, + 'z': true, + '|': true, + '~': true, +} + +// skipSpace returns a slice of the string s with all leading RFC 2616 linear +// whitespace removed. +func skipSpace(s string) (rest string) { + i := 0 + for ; i < len(s); i++ { + if b := s[i]; b != ' ' && b != '\t' { + break + } + } + return s[i:] +} + +// nextToken returns the leading RFC 2616 token of s and the string following +// the token. +func nextToken(s string) (token, rest string) { + i := 0 + for ; i < len(s); i++ { + if !isTokenOctet[s[i]] { + break + } + } + return s[:i], s[i:] +} + +// nextTokenOrQuoted returns the leading token or quoted string per RFC 2616 +// and the string following the token or quoted string. +func nextTokenOrQuoted(s string) (value string, rest string) { + if !strings.HasPrefix(s, "\"") { + return nextToken(s) + } + s = s[1:] + for i := 0; i < len(s); i++ { + switch s[i] { + case '"': + return s[:i], s[i+1:] + case '\\': + p := make([]byte, len(s)-1) + j := copy(p, s[:i]) + escape := true + for i = i + 1; i < len(s); i++ { + b := s[i] + switch { + case escape: + escape = false + p[j] = b + j++ + case b == '\\': + escape = true + case b == '"': + return string(p[:j]), s[i+1:] + default: + p[j] = b + j++ + } + } + return "", "" + } + } + return "", "" +} + +// equalASCIIFold returns true if s is equal to t with ASCII case folding as +// defined in RFC 4790. +func equalASCIIFold(s, t string) bool { + for s != "" && t != "" { + sr, size := utf8.DecodeRuneInString(s) + s = s[size:] + tr, size := utf8.DecodeRuneInString(t) + t = t[size:] + if sr == tr { + continue + } + if 'A' <= sr && sr <= 'Z' { + sr = sr + 'a' - 'A' + } + if 'A' <= tr && tr <= 'Z' { + tr = tr + 'a' - 'A' + } + if sr != tr { + return false + } + } + return s == t +} + +// tokenListContainsValue returns true if the 1#token header with the given +// name contains a token equal to value with ASCII case folding. +func tokenListContainsValue(header http.Header, name string, value string) bool { +headers: + for _, s := range header[name] { + for { + var t string + t, s = nextToken(skipSpace(s)) + if t == "" { + continue headers + } + s = skipSpace(s) + if s != "" && s[0] != ',' { + continue headers + } + if equalASCIIFold(t, value) { + return true + } + if s == "" { + continue headers + } + s = s[1:] + } + } + return false +} + +// parseExtensions parses WebSocket extensions from a header. +func parseExtensions(header http.Header) []map[string]string { + // From RFC 6455: + // + // Sec-WebSocket-Extensions = extension-list + // extension-list = 1#extension + // extension = extension-token *( ";" extension-param ) + // extension-token = registered-token + // registered-token = token + // extension-param = token [ "=" (token | quoted-string) ] + // ;When using the quoted-string syntax variant, the value + // ;after quoted-string unescaping MUST conform to the + // ;'token' ABNF. + + var result []map[string]string +headers: + for _, s := range header["Sec-Websocket-Extensions"] { + for { + var t string + t, s = nextToken(skipSpace(s)) + if t == "" { + continue headers + } + ext := map[string]string{"": t} + for { + s = skipSpace(s) + if !strings.HasPrefix(s, ";") { + break + } + var k string + k, s = nextToken(skipSpace(s[1:])) + if k == "" { + continue headers + } + s = skipSpace(s) + var v string + if strings.HasPrefix(s, "=") { + v, s = nextTokenOrQuoted(skipSpace(s[1:])) + s = skipSpace(s) + } + if s != "" && s[0] != ',' && s[0] != ';' { + continue headers + } + ext[k] = v + } + if s != "" && s[0] != ',' { + continue headers + } + result = append(result, ext) + if s == "" { + continue headers + } + s = s[1:] + } + } + return result +} + +// isValidChallengeKey checks if the argument meets RFC6455 specification. +func isValidChallengeKey(s string) bool { + // From RFC6455: + // + // A |Sec-WebSocket-Key| header field with a base64-encoded (see + // Section 4 of [RFC4648]) value that, when decoded, is 16 bytes in + // length. + + if s == "" { + return false + } + decoded, err := base64.StdEncoding.DecodeString(s) + return err == nil && len(decoded) == 16 +} diff --git a/internal/websocket/util_test.go b/internal/websocket/util_test.go new file mode 100644 index 00000000..f14d69a1 --- /dev/null +++ b/internal/websocket/util_test.go @@ -0,0 +1,115 @@ +// Copyright 2014 The Gorilla WebSocket Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package websocket + +import ( + "net/http" + "reflect" + "testing" +) + +var equalASCIIFoldTests = []struct { + t, s string + eq bool +}{ + {"WebSocket", "websocket", true}, + {"websocket", "WebSocket", true}, + {"Öyster", "öyster", false}, + {"WebSocket", "WetSocket", false}, +} + +func TestEqualASCIIFold(t *testing.T) { + for _, tt := range equalASCIIFoldTests { + eq := equalASCIIFold(tt.s, tt.t) + if eq != tt.eq { + t.Errorf("equalASCIIFold(%q, %q) = %v, want %v", tt.s, tt.t, eq, tt.eq) + } + } +} + +var tokenListContainsValueTests = []struct { + value string + ok bool +}{ + {"WebSocket", true}, + {"WEBSOCKET", true}, + {"websocket", true}, + {"websockets", false}, + {"x websocket", false}, + {"websocket x", false}, + {"other,websocket,more", true}, + {"other, websocket, more", true}, +} + +func TestTokenListContainsValue(t *testing.T) { + for _, tt := range tokenListContainsValueTests { + h := http.Header{"Upgrade": {tt.value}} + ok := tokenListContainsValue(h, "Upgrade", "websocket") + if ok != tt.ok { + t.Errorf("tokenListContainsValue(h, n, %q) = %v, want %v", tt.value, ok, tt.ok) + } + } +} + +var isValidChallengeKeyTests = []struct { + key string + ok bool +}{ + {"dGhlIHNhbXBsZSBub25jZQ==", true}, + {"", false}, + {"InvalidKey", false}, + {"WHQ4eXhscUtKYjBvOGN3WEdtOEQ=", false}, +} + +func TestIsValidChallengeKey(t *testing.T) { + for _, tt := range isValidChallengeKeyTests { + ok := isValidChallengeKey(tt.key) + if ok != tt.ok { + t.Errorf("isValidChallengeKey returns %v, want %v", ok, tt.ok) + } + } +} + +var parseExtensionTests = []struct { + value string + extensions []map[string]string +}{ + {`foo`, []map[string]string{{"": "foo"}}}, + {`foo, bar; baz=2`, []map[string]string{ + {"": "foo"}, + {"": "bar", "baz": "2"}}}, + {`foo; bar="b,a;z"`, []map[string]string{ + {"": "foo", "bar": "b,a;z"}}}, + {`foo , bar; baz = 2`, []map[string]string{ + {"": "foo"}, + {"": "bar", "baz": "2"}}}, + {`foo, bar; baz=2 junk`, []map[string]string{ + {"": "foo"}}}, + {`foo junk, bar; baz=2 junk`, nil}, + {`mux; max-channels=4; flow-control, deflate-stream`, []map[string]string{ + {"": "mux", "max-channels": "4", "flow-control": ""}, + {"": "deflate-stream"}}}, + {`permessage-foo; x="10"`, []map[string]string{ + {"": "permessage-foo", "x": "10"}}}, + {`permessage-foo; use_y, permessage-foo`, []map[string]string{ + {"": "permessage-foo", "use_y": ""}, + {"": "permessage-foo"}}}, + {`permessage-deflate; client_max_window_bits; server_max_window_bits=10 , permessage-deflate; client_max_window_bits`, []map[string]string{ + {"": "permessage-deflate", "client_max_window_bits": "", "server_max_window_bits": "10"}, + {"": "permessage-deflate", "client_max_window_bits": ""}}}, + {"permessage-deflate; server_no_context_takeover; client_max_window_bits=15", []map[string]string{ + {"": "permessage-deflate", "server_no_context_takeover": "", "client_max_window_bits": "15"}, + }}, +} + +func TestParseExtensions(t *testing.T) { + for _, tt := range parseExtensionTests { + h := http.Header{http.CanonicalHeaderKey("Sec-WebSocket-Extensions"): {tt.value}} + extensions := parseExtensions(h) + if !reflect.DeepEqual(extensions, tt.extensions) { + t.Errorf("parseExtensions(%q)\n = %v,\nwant %v", tt.value, extensions, tt.extensions) + } + } +} diff --git a/metrics.go b/metrics.go index 3af6b0b8..4359a953 100644 --- a/metrics.go +++ b/metrics.go @@ -66,7 +66,6 @@ type metrics struct { recoverCountNo prometheus.Counter transportConnectCountWebsocket prometheus.Counter - transportConnectCountSockJS prometheus.Counter transportConnectCountSSE prometheus.Counter transportConnectCountHTTPStream prometheus.Counter @@ -156,8 +155,6 @@ func (m *metrics) incTransportConnect(transport string) { switch transport { case transportWebsocket: m.transportConnectCountWebsocket.Inc() - case transportSockJS: - m.transportConnectCountSockJS.Inc() case transportSSE: m.transportConnectCountSSE.Inc() case transportHTTPStream: @@ -476,7 +473,6 @@ func initMetricsRegistry(registry prometheus.Registerer, metricsNamespace string m.recoverCountNo = m.recoverCount.WithLabelValues("no") m.transportConnectCountWebsocket = m.transportConnectCount.WithLabelValues(transportWebsocket) - m.transportConnectCountSockJS = m.transportConnectCount.WithLabelValues(transportSockJS) m.transportConnectCountHTTPStream = m.transportConnectCount.WithLabelValues(transportHTTPStream) m.transportConnectCountSSE = m.transportConnectCount.WithLabelValues(transportSSE)