├── .gitignore ├── .goreleaser.yml ├── go.mod ├── Makefile ├── doc.go ├── .github └── workflows │ └── build-go.yaml ├── config.go ├── go.sum ├── mock_test.go ├── interface.go ├── mock.go ├── README.md ├── wsc.go ├── LICENSE └── wsc_test.go /.gitignore: -------------------------------------------------------------------------------- 1 | unit_coverage.out 2 | cov.report 3 | dist 4 | -------------------------------------------------------------------------------- /.goreleaser.yml: -------------------------------------------------------------------------------- 1 | snapshot: 2 | name_template: "{{ .Tag }}-next" 3 | changelog: 4 | sort: asc 5 | filters: 6 | exclude: 7 | - '^docs:' 8 | - '^test:' 9 | - '^examples:' 10 | builds: 11 | - skip: true 12 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module go.aporeto.io/wsc 2 | 3 | go 1.18 4 | 5 | require ( 6 | github.com/gorilla/websocket v1.5.0 7 | github.com/smartystreets/goconvey v1.7.2 8 | ) 9 | 10 | require ( 11 | github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1 // indirect 12 | github.com/jtolds/gls v4.20.0+incompatible // indirect 13 | github.com/smartystreets/assertions v1.2.0 // indirect 14 | ) 15 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | MAKEFLAGS += --warn-undefined-variables 2 | SHELL := /bin/bash -o pipefail 3 | 4 | export GO111MODULE = on 5 | 6 | default: lint test 7 | 8 | lint: 9 | golangci-lint run \ 10 | --disable-all \ 11 | --exclude-use-default=false \ 12 | --exclude=package-comments \ 13 | --enable=errcheck \ 14 | --enable=goimports \ 15 | --enable=ineffassign \ 16 | --enable=revive \ 17 | --enable=unused \ 18 | --enable=staticcheck \ 19 | --enable=unconvert \ 20 | --enable=misspell \ 21 | --enable=prealloc \ 22 | --enable=nakedret \ 23 | --enable=unparam \ 24 | ./... 25 | 26 | test: 27 | go test ./... -race -cover -covermode=atomic -coverprofile=unit_coverage.out 28 | 29 | sec: 30 | gosec -quiet ./... 31 | -------------------------------------------------------------------------------- /doc.go: -------------------------------------------------------------------------------- 1 | // Copyright 2019 Aporeto Inc. 2 | // Licensed under the Apache License, Version 2.0 (the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at 5 | // http://www.apache.org/licenses/LICENSE-2.0 6 | // Unless required by applicable law or agreed to in writing, software 7 | // distributed under the License is distributed on an "AS IS" BASIS, 8 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 9 | // See the License for the specific language governing permissions and 10 | // limitations under the License. 11 | 12 | // Package wsc provides an idiomatic way to deal with 13 | // websockets using channels. 14 | package wsc // import "go.aporeto.io/wsc" 15 | -------------------------------------------------------------------------------- /.github/workflows/build-go.yaml: -------------------------------------------------------------------------------- 1 | name: build-go 2 | on: 3 | push: 4 | branches: 5 | - master 6 | pull_request: 7 | 8 | defaults: 9 | run: 10 | shell: bash 11 | 12 | env: 13 | GO111MODULE: on 14 | 15 | jobs: 16 | build: 17 | runs-on: ubuntu-latest 18 | strategy: 19 | fail-fast: false 20 | matrix: 21 | go: 22 | - "1.20" 23 | - "1.21" 24 | steps: 25 | - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 26 | 27 | - uses: actions/setup-go@93397bea11091df50f3d7e59dc26a7711a8bcfbe # v4 28 | with: 29 | go-version: ${{ matrix.go }} 30 | cache: true 31 | 32 | - name: setup 33 | run: | 34 | go install github.com/golangci/golangci-lint/cmd/golangci-lint@v1.54.2 35 | 36 | - name: build 37 | run: | 38 | make 39 | -------------------------------------------------------------------------------- /config.go: -------------------------------------------------------------------------------- 1 | // Copyright 2019 Aporeto Inc. 2 | // Licensed under the Apache License, Version 2.0 (the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at 5 | // http://www.apache.org/licenses/LICENSE-2.0 6 | // Unless required by applicable law or agreed to in writing, software 7 | // distributed under the License is distributed on an "AS IS" BASIS, 8 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 9 | // See the License for the specific language governing permissions and 10 | // limitations under the License. 11 | 12 | package wsc 13 | 14 | import ( 15 | "context" 16 | "crypto/tls" 17 | "net" 18 | "net/http" 19 | "time" 20 | ) 21 | 22 | // Config contains configuration for the webbsocket. 23 | type Config struct { 24 | TLSConfig *tls.Config 25 | Headers http.Header 26 | NetDialContextFunc func(ctx context.Context, network, addr string) (net.Conn, error) 27 | WriteWait time.Duration 28 | PongWait time.Duration 29 | PingPeriod time.Duration 30 | ReadBufferSize int 31 | ReadChanSize int 32 | WriteBufferSize int 33 | WriteChanSize int 34 | EnableCompression bool 35 | } 36 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1 h1:EGx4pi6eqNxGaHF6qqu48+N2wcFQ5qg5FXgOdqsJ5d8= 2 | github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY= 3 | github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc= 4 | github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= 5 | github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo= 6 | github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU= 7 | github.com/smartystreets/assertions v1.2.0 h1:42S6lae5dvLc7BrLu/0ugRtcFVjoJNMC/N3yZFZkDFs= 8 | github.com/smartystreets/assertions v1.2.0/go.mod h1:tcbTF8ujkAEcZ8TElKY+i30BzYlVhC/LOxJk7iOWnoo= 9 | github.com/smartystreets/goconvey v1.7.2 h1:9RBaZCeXEQ3UselpuwUQHltGVXvdwm6cv1hgR6gDIPg= 10 | github.com/smartystreets/goconvey v1.7.2/go.mod h1:Vw0tHAZW6lzCRk3xgdin6fKYcG+G3Pg9vgXWeJpQFMM= 11 | golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= 12 | golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= 13 | golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= 14 | golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= 15 | golang.org/x/tools v0.0.0-20190328211700-ab21143f2384/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs= 16 | -------------------------------------------------------------------------------- /mock_test.go: -------------------------------------------------------------------------------- 1 | // Copyright 2019 Aporeto Inc. 2 | // Licensed under the Apache License, Version 2.0 (the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at 5 | // http://www.apache.org/licenses/LICENSE-2.0 6 | // Unless required by applicable law or agreed to in writing, software 7 | // distributed under the License is distributed on an "AS IS" BASIS, 8 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 9 | // See the License for the specific language governing permissions and 10 | // limitations under the License. 11 | 12 | package wsc 13 | 14 | import ( 15 | "context" 16 | "errors" 17 | "testing" 18 | 19 | . "github.com/smartystreets/goconvey/convey" 20 | ) 21 | 22 | func TestWSC_Mock(t *testing.T) { 23 | 24 | Convey("Given I have a MockWebsocket", t, func() { 25 | 26 | m := NewMockWebsocket(context.Background()) 27 | 28 | Convey("When I set NextRead then read", func() { 29 | 30 | m.NextRead([]byte("hello")) 31 | 32 | out := <-m.Read() 33 | 34 | Convey("Then out should be correct", func() { 35 | So(string(out), ShouldEqual, "hello") 36 | }) 37 | }) 38 | 39 | Convey("When I write someting then get LastWrite", func() { 40 | 41 | m.Write([]byte("hello")) 42 | 43 | out := <-m.LastWrite() 44 | 45 | Convey("Then out should be correct", func() { 46 | So(string(out), ShouldEqual, "hello") 47 | }) 48 | }) 49 | 50 | Convey("When I send done then get LastDone", func() { 51 | 52 | m.NextDone(errors.New("boom")) 53 | 54 | out := <-m.Done() 55 | 56 | Convey("Then out should be correct", func() { 57 | So(out.Error(), ShouldEqual, "boom") 58 | }) 59 | }) 60 | 61 | Convey("When I call close", func() { 62 | 63 | m.Close(1) 64 | 65 | out := <-m.Done() 66 | 67 | Convey("Then out should be correct", func() { 68 | So(out.Error(), ShouldEqual, "1") 69 | }) 70 | }) 71 | }) 72 | 73 | } 74 | -------------------------------------------------------------------------------- /interface.go: -------------------------------------------------------------------------------- 1 | // Copyright 2019 Aporeto Inc. 2 | // Licensed under the Apache License, Version 2.0 (the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at 5 | // http://www.apache.org/licenses/LICENSE-2.0 6 | // Unless required by applicable law or agreed to in writing, software 7 | // distributed under the License is distributed on an "AS IS" BASIS, 8 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 9 | // See the License for the specific language governing permissions and 10 | // limitations under the License. 11 | 12 | package wsc 13 | 14 | // Websocket is the interface of channel based websocket. 15 | type Websocket interface { 16 | 17 | // Reads returns a channel where the incoming messages are published. 18 | // If nothing pumps the Read() while it is full, new messages will be 19 | // discarded. 20 | // 21 | // You can configure the size of the read chan in Config. 22 | // The default is 64 messages. 23 | Read() chan []byte 24 | 25 | // Write writes the given []byte in to the websocket. 26 | // If the other side of the websocket cannot get all messages 27 | // while the internal write channel is full, new messages will 28 | // be discarded. 29 | // 30 | // You can configure the size of the write chan in Config. 31 | // The default is 64 messages. 32 | Write([]byte) 33 | 34 | // Done returns a channel that will return when the connection 35 | // is closed. 36 | // 37 | // The content will be nil for clean disconnection or 38 | // the error that caused the disconnection. If nothing pumps the 39 | // Done() channel, the error will be discarded. 40 | Done() chan error 41 | 42 | // Close closes the websocket. 43 | // 44 | // Closing the websocket a second time has no effect. 45 | // A closed Websocket cannot be reused. 46 | Close(code int) 47 | 48 | // Error returns a channel that will return errors like 49 | // read or write discards and other errors that are not 50 | // terminating the connection. 51 | // 52 | // If nothing pumps the Error() channel, the error will be discarded. 53 | Error() chan error 54 | } 55 | -------------------------------------------------------------------------------- /mock.go: -------------------------------------------------------------------------------- 1 | // Copyright 2019 Aporeto Inc. 2 | // Licensed under the Apache License, Version 2.0 (the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at 5 | // http://www.apache.org/licenses/LICENSE-2.0 6 | // Unless required by applicable law or agreed to in writing, software 7 | // distributed under the License is distributed on an "AS IS" BASIS, 8 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 9 | // See the License for the specific language governing permissions and 10 | // limitations under the License. 11 | 12 | package wsc 13 | 14 | import ( 15 | "context" 16 | "fmt" 17 | ) 18 | 19 | // A MockWebsocket is a utility to write unit 20 | // tests on websockets. 21 | type MockWebsocket interface { 22 | NextRead(data []byte) 23 | LastWrite() chan []byte 24 | NextDone(err error) 25 | 26 | Websocket 27 | } 28 | 29 | type mockWebsocket struct { 30 | readChan chan []byte 31 | writeChan chan []byte 32 | doneChan chan error 33 | errChan chan error 34 | cancel context.CancelFunc 35 | } 36 | 37 | // NewMockWebsocket returns a mocked Websocket that can be used 38 | // in unit tests. 39 | func NewMockWebsocket(ctx context.Context) MockWebsocket { 40 | 41 | _, cancel := context.WithCancel(ctx) 42 | 43 | return &mockWebsocket{ 44 | readChan: make(chan []byte, 64), 45 | writeChan: make(chan []byte, 64), 46 | doneChan: make(chan error, 64), 47 | errChan: make(chan error, 1), 48 | cancel: cancel, 49 | } 50 | } 51 | 52 | func (s *mockWebsocket) Write(data []byte) { s.writeChan <- data } 53 | func (s *mockWebsocket) Read() chan []byte { return s.readChan } 54 | func (s *mockWebsocket) Done() chan error { return s.doneChan } 55 | func (s *mockWebsocket) Close(code int) { s.doneChan <- fmt.Errorf("%d", code) } 56 | func (s *mockWebsocket) NextRead(data []byte) { s.readChan <- data } 57 | func (s *mockWebsocket) LastWrite() chan []byte { return s.writeChan } 58 | func (s *mockWebsocket) NextDone(err error) { s.doneChan <- err } 59 | func (s *mockWebsocket) Error() chan error { return s.errChan } 60 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # WSC 2 | 3 | [![Codacy Badge](https://app.codacy.com/project/badge/Grade/74723c4b258640808518c6ec17e357c0)](https://www.codacy.com/gh/PaloAltoNetworks/wsc/dashboard?utm_source=github.com&utm_medium=referral&utm_content=PaloAltoNetworks/wsc&utm_campaign=Badge_Grade) [![Codacy Badge](https://app.codacy.com/project/badge/Coverage/74723c4b258640808518c6ec17e357c0)](https://www.codacy.com/gh/PaloAltoNetworks/wsc/dashboard?utm_source=github.com&utm_medium=referral&utm_content=PaloAltoNetworks/wsc&utm_campaign=Badge_Coverage) 4 | 5 | WSC (WebSocket Channel) is a library that can be used to manage github.com/gorilla/websocket using channels. 6 | 7 | It provides 2 main functions: 8 | 9 | - `func Connect(ctx context.Context, url string, config Config) (Websocket, *http.Response, error)` 10 | - `func Accept(ctx context.Context, conn WSConnection, config Config) (Websocket, error)` 11 | 12 | The interface `WSConnection` is used to interract with the websocket: 13 | 14 | ```go 15 | // Websocket is the interface of channel based websocket. 16 | type Websocket interface { 17 | 18 | // Reads returns a channel where the incoming messages are published. 19 | // If nothing pumps the Read() while it is full, new messages will be 20 | // discarded. 21 | // 22 | // You can configure the size of the read chan in Config. 23 | // The default is 64 messages. 24 | Read() chan []byte 25 | 26 | // Write writes the given []byte in to the websocket. 27 | // If the other side of the websocket cannot get all messages 28 | // while the internal write channel is full, new messages will 29 | // be discarded. 30 | // 31 | // You can configure the size of the write chan in Config. 32 | // The default is 64 messages. 33 | Write([]byte) 34 | 35 | // Done returns a channel that will return when the connection 36 | // is closed. 37 | // 38 | // The content will be nil for clean disconnection or 39 | // the error that caused the disconnection. If nothing pumps the 40 | // Done() channel, the error will be discarded. 41 | Done() chan error 42 | 43 | // Close closes the websocket. 44 | // 45 | // Closing the websocket a second time has no effect. 46 | // A closed Websocket cannot be reused. 47 | Close(code int) 48 | 49 | // Error returns a channel that will return errors like 50 | // read or write discards and other errors that are not 51 | // terminating the connection. 52 | // 53 | // If nothing pumps the Error() channel, the error will be discarded. 54 | Error() chan error 55 | } 56 | ``` 57 | -------------------------------------------------------------------------------- /wsc.go: -------------------------------------------------------------------------------- 1 | // Copyright 2019 Aporeto Inc. 2 | // Licensed under the Apache License, Version 2.0 (the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at 5 | // http://www.apache.org/licenses/LICENSE-2.0 6 | // Unless required by applicable law or agreed to in writing, software 7 | // distributed under the License is distributed on an "AS IS" BASIS, 8 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 9 | // See the License for the specific language governing permissions and 10 | // limitations under the License. 11 | 12 | package wsc 13 | 14 | import ( 15 | "context" 16 | "encoding/binary" 17 | "fmt" 18 | "log" 19 | "net/http" 20 | "time" 21 | 22 | "github.com/gorilla/websocket" 23 | ) 24 | 25 | // various error messages. 26 | var ( 27 | ErrWriteMessageDiscarded = fmt.Errorf("write chan full: one or more messages has not been sent") 28 | ErrReadMessageDiscarded = fmt.Errorf("read chan full: one or more messages has not been received") 29 | ) 30 | 31 | // WSConnection is the interface that must be implemented 32 | // as a websocket. github.com/gorilla/websocket implements 33 | // this interface. 34 | type WSConnection interface { 35 | SetReadDeadline(time.Time) error 36 | SetWriteDeadline(time.Time) error 37 | SetCloseHandler(func(code int, text string) error) 38 | SetPongHandler(func(string) error) 39 | ReadMessage() (int, []byte, error) 40 | WriteMessage(int, []byte) error 41 | WriteControl(int, []byte, time.Time) error 42 | Close() error 43 | } 44 | 45 | type ws struct { 46 | conn WSConnection 47 | readChan chan []byte 48 | writeChan chan []byte 49 | doneChan chan error 50 | errChan chan error 51 | cancel context.CancelFunc 52 | closeCodeCh chan int 53 | config Config 54 | } 55 | 56 | // Connect connects to the url and returns a Websocket. 57 | func Connect(ctx context.Context, url string, config Config) (Websocket, *http.Response, error) { 58 | 59 | dialer := &websocket.Dialer{ 60 | Proxy: http.ProxyFromEnvironment, 61 | TLSClientConfig: config.TLSConfig, 62 | ReadBufferSize: config.ReadBufferSize, 63 | WriteBufferSize: config.WriteBufferSize, 64 | EnableCompression: config.EnableCompression, 65 | NetDialContext: config.NetDialContextFunc, 66 | } 67 | 68 | conn, resp, err := dialer.DialContext(ctx, url, config.Headers) 69 | if err != nil { 70 | return nil, resp, err 71 | } 72 | 73 | s, err := Accept(ctx, conn, config) 74 | 75 | return s, resp, err 76 | } 77 | 78 | // Accept handles an already connect *websocket.Conn and returns a Websocket. 79 | func Accept(ctx context.Context, conn WSConnection, config Config) (Websocket, error) { 80 | 81 | if config.PongWait == 0 { 82 | config.PongWait = 30 * time.Second 83 | } 84 | if config.WriteWait == 0 { 85 | config.WriteWait = 10 * time.Second 86 | } 87 | if config.PingPeriod == 0 { 88 | config.PingPeriod = 15 * time.Second 89 | } 90 | if config.WriteChanSize == 0 { 91 | config.WriteChanSize = 64 92 | } 93 | if config.ReadChanSize == 0 { 94 | config.ReadChanSize = 64 95 | } 96 | 97 | if err := conn.SetReadDeadline(time.Now().Add(config.PongWait)); err != nil { 98 | return nil, err 99 | } 100 | 101 | subCtx, cancel := context.WithCancel(ctx) 102 | 103 | s := &ws{ 104 | conn: conn, 105 | readChan: make(chan []byte, config.ReadChanSize), 106 | writeChan: make(chan []byte, config.WriteChanSize), 107 | doneChan: make(chan error, 2), 108 | errChan: make(chan error, 10), 109 | closeCodeCh: make(chan int, 1), 110 | cancel: cancel, 111 | config: config, 112 | } 113 | 114 | s.conn.SetCloseHandler(func(code int, text string) error { 115 | s.cancel() 116 | return nil 117 | }) 118 | 119 | s.conn.SetPongHandler(func(string) error { 120 | return s.conn.SetReadDeadline(time.Now().Add(s.config.PongWait)) 121 | }) 122 | 123 | go s.readPump() 124 | go s.writePump(subCtx) 125 | 126 | return s, nil 127 | } 128 | 129 | // Write is part of the the Websocket interface implementation. 130 | func (s *ws) Write(data []byte) { 131 | 132 | select { 133 | case s.writeChan <- data: 134 | default: 135 | s.error(ErrWriteMessageDiscarded) 136 | } 137 | } 138 | 139 | // Read is part of the the Websocket interface implementation. 140 | func (s *ws) Read() chan []byte { 141 | 142 | return s.readChan 143 | } 144 | 145 | // Error is part of the the Websocket interface implementation. 146 | func (s *ws) Error() chan error { 147 | 148 | return s.errChan 149 | } 150 | 151 | // Done is part of the the Websocket interface implementation. 152 | func (s *ws) Done() chan error { 153 | 154 | return s.doneChan 155 | } 156 | 157 | // Close is part of the the Websocket interface implementation. 158 | func (s *ws) Close(code int) { 159 | 160 | if code != 0 { 161 | select { 162 | case s.closeCodeCh <- code: 163 | default: 164 | } 165 | } 166 | 167 | s.cancel() 168 | } 169 | 170 | func (s *ws) readPump() { 171 | 172 | var err error 173 | var msg []byte 174 | var msgType int 175 | 176 | for { 177 | if msgType, msg, err = s.conn.ReadMessage(); err != nil { 178 | s.done(err) 179 | return 180 | } 181 | 182 | switch msgType { 183 | 184 | case websocket.TextMessage, websocket.BinaryMessage: 185 | select { 186 | case s.readChan <- msg: 187 | default: 188 | s.error(ErrReadMessageDiscarded) 189 | } 190 | 191 | case websocket.CloseMessage: 192 | return 193 | } 194 | } 195 | } 196 | 197 | func (s *ws) writePump(ctx context.Context) { 198 | 199 | var err error 200 | 201 | ticker := time.NewTicker(s.config.PingPeriod) 202 | defer ticker.Stop() 203 | 204 | for { 205 | select { 206 | 207 | case message := <-s.writeChan: 208 | 209 | _ = s.conn.SetWriteDeadline(time.Now().Add(s.config.WriteWait)) // nolint: errcheck 210 | if err = s.conn.WriteMessage(websocket.TextMessage, message); err != nil { 211 | s.done(err) 212 | return 213 | } 214 | 215 | case <-ticker.C: 216 | 217 | _ = s.conn.SetWriteDeadline(time.Now().Add(s.config.WriteWait)) // nolint: errcheck 218 | if err = s.conn.WriteMessage(websocket.PingMessage, nil); err != nil { 219 | s.done(err) 220 | return 221 | } 222 | 223 | case <-ctx.Done(): 224 | 225 | code := websocket.CloseGoingAway 226 | select { 227 | case code = <-s.closeCodeCh: 228 | default: 229 | } 230 | 231 | enc := make([]byte, 2) 232 | binary.BigEndian.PutUint16(enc, uint16(code)) 233 | 234 | s.done( 235 | s.conn.WriteControl( 236 | websocket.CloseMessage, 237 | enc, 238 | time.Now().Add(1*time.Second), 239 | ), 240 | ) 241 | 242 | _ = s.conn.Close() 243 | 244 | return 245 | } 246 | } 247 | } 248 | 249 | func (s *ws) done(err error) { 250 | 251 | select { 252 | case s.doneChan <- err: 253 | default: 254 | if err != nil { 255 | log.Printf("wsc: error: unable to send done message: error: %s\n", err) 256 | } 257 | } 258 | } 259 | 260 | func (s *ws) error(err error) { 261 | 262 | select { 263 | case s.errChan <- err: 264 | default: 265 | if err != nil { 266 | log.Printf("wsc: error: unable to send error: %s\n", err) 267 | } 268 | } 269 | } 270 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | Apache License 3 | Version 2.0, January 2004 4 | http://www.apache.org/licenses/ 5 | 6 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 7 | 8 | 1. Definitions. 9 | 10 | "License" shall mean the terms and conditions for use, reproduction, 11 | and distribution as defined by Sections 1 through 9 of this document. 12 | 13 | "Licensor" shall mean the copyright owner or entity authorized by 14 | the copyright owner that is granting the License. 15 | 16 | "Legal Entity" shall mean the union of the acting entity and all 17 | other entities that control, are controlled by, or are under common 18 | control with that entity. For the purposes of this definition, 19 | "control" means (i) the power, direct or indirect, to cause the 20 | direction or management of such entity, whether by contract or 21 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 22 | outstanding shares, or (iii) beneficial ownership of such entity. 23 | 24 | "You" (or "Your") shall mean an individual or Legal Entity 25 | exercising permissions granted by this License. 26 | 27 | "Source" form shall mean the preferred form for making modifications, 28 | including but not limited to software source code, documentation 29 | source, and configuration files. 30 | 31 | "Object" form shall mean any form resulting from mechanical 32 | transformation or translation of a Source form, including but 33 | not limited to compiled object code, generated documentation, 34 | and conversions to other media types. 35 | 36 | "Work" shall mean the work of authorship, whether in Source or 37 | Object form, made available under the License, as indicated by a 38 | copyright notice that is included in or attached to the work 39 | (an example is provided in the Appendix below). 40 | 41 | "Derivative Works" shall mean any work, whether in Source or Object 42 | form, that is based on (or derived from) the Work and for which the 43 | editorial revisions, annotations, elaborations, or other modifications 44 | represent, as a whole, an original work of authorship. For the purposes 45 | of this License, Derivative Works shall not include works that remain 46 | separable from, or merely link (or bind by name) to the interfaces of, 47 | the Work and Derivative Works thereof. 48 | 49 | "Contribution" shall mean any work of authorship, including 50 | the original version of the Work and any modifications or additions 51 | to that Work or Derivative Works thereof, that is intentionally 52 | submitted to Licensor for inclusion in the Work by the copyright owner 53 | or by an individual or Legal Entity authorized to submit on behalf of 54 | the copyright owner. For the purposes of this definition, "submitted" 55 | means any form of electronic, verbal, or written communication sent 56 | to the Licensor or its representatives, including but not limited to 57 | communication on electronic mailing lists, source code control systems, 58 | and issue tracking systems that are managed by, or on behalf of, the 59 | Licensor for the purpose of discussing and improving the Work, but 60 | excluding communication that is conspicuously marked or otherwise 61 | designated in writing by the copyright owner as "Not a Contribution." 62 | 63 | "Contributor" shall mean Licensor and any individual or Legal Entity 64 | on behalf of whom a Contribution has been received by Licensor and 65 | subsequently incorporated within the Work. 66 | 67 | 2. Grant of Copyright License. Subject to the terms and conditions of 68 | this License, each Contributor hereby grants to You a perpetual, 69 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 70 | copyright license to reproduce, prepare Derivative Works of, 71 | publicly display, publicly perform, sublicense, and distribute the 72 | Work and such Derivative Works in Source or Object form. 73 | 74 | 3. Grant of Patent License. Subject to the terms and conditions of 75 | this License, each Contributor hereby grants to You a perpetual, 76 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 77 | (except as stated in this section) patent license to make, have made, 78 | use, offer to sell, sell, import, and otherwise transfer the Work, 79 | where such license applies only to those patent claims licensable 80 | by such Contributor that are necessarily infringed by their 81 | Contribution(s) alone or by combination of their Contribution(s) 82 | with the Work to which such Contribution(s) was submitted. If You 83 | institute patent litigation against any entity (including a 84 | cross-claim or counterclaim in a lawsuit) alleging that the Work 85 | or a Contribution incorporated within the Work constitutes direct 86 | or contributory patent infringement, then any patent licenses 87 | granted to You under this License for that Work shall terminate 88 | as of the date such litigation is filed. 89 | 90 | 4. Redistribution. You may reproduce and distribute copies of the 91 | Work or Derivative Works thereof in any medium, with or without 92 | modifications, and in Source or Object form, provided that You 93 | meet the following conditions: 94 | 95 | (a) You must give any other recipients of the Work or 96 | Derivative Works a copy of this License; and 97 | 98 | (b) You must cause any modified files to carry prominent notices 99 | stating that You changed the files; and 100 | 101 | (c) You must retain, in the Source form of any Derivative Works 102 | that You distribute, all copyright, patent, trademark, and 103 | attribution notices from the Source form of the Work, 104 | excluding those notices that do not pertain to any part of 105 | the Derivative Works; and 106 | 107 | (d) If the Work includes a "NOTICE" text file as part of its 108 | distribution, then any Derivative Works that You distribute must 109 | include a readable copy of the attribution notices contained 110 | within such NOTICE file, excluding those notices that do not 111 | pertain to any part of the Derivative Works, in at least one 112 | of the following places: within a NOTICE text file distributed 113 | as part of the Derivative Works; within the Source form or 114 | documentation, if provided along with the Derivative Works; or, 115 | within a display generated by the Derivative Works, if and 116 | wherever such third-party notices normally appear. The contents 117 | of the NOTICE file are for informational purposes only and 118 | do not modify the License. You may add Your own attribution 119 | notices within Derivative Works that You distribute, alongside 120 | or as an addendum to the NOTICE text from the Work, provided 121 | that such additional attribution notices cannot be construed 122 | as modifying the License. 123 | 124 | You may add Your own copyright statement to Your modifications and 125 | may provide additional or different license terms and conditions 126 | for use, reproduction, or distribution of Your modifications, or 127 | for any such Derivative Works as a whole, provided Your use, 128 | reproduction, and distribution of the Work otherwise complies with 129 | the conditions stated in this License. 130 | 131 | 5. Submission of Contributions. Unless You explicitly state otherwise, 132 | any Contribution intentionally submitted for inclusion in the Work 133 | by You to the Licensor shall be under the terms and conditions of 134 | this License, without any additional terms or conditions. 135 | Notwithstanding the above, nothing herein shall supersede or modify 136 | the terms of any separate license agreement you may have executed 137 | with Licensor regarding such Contributions. 138 | 139 | 6. Trademarks. This License does not grant permission to use the trade 140 | names, trademarks, service marks, or product names of the Licensor, 141 | except as required for reasonable and customary use in describing the 142 | origin of the Work and reproducing the content of the NOTICE file. 143 | 144 | 7. Disclaimer of Warranty. Unless required by applicable law or 145 | agreed to in writing, Licensor provides the Work (and each 146 | Contributor provides its Contributions) on an "AS IS" BASIS, 147 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 148 | implied, including, without limitation, any warranties or conditions 149 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 150 | PARTICULAR PURPOSE. You are solely responsible for determining the 151 | appropriateness of using or redistributing the Work and assume any 152 | risks associated with Your exercise of permissions under this License. 153 | 154 | 8. Limitation of Liability. In no event and under no legal theory, 155 | whether in tort (including negligence), contract, or otherwise, 156 | unless required by applicable law (such as deliberate and grossly 157 | negligent acts) or agreed to in writing, shall any Contributor be 158 | liable to You for damages, including any direct, indirect, special, 159 | incidental, or consequential damages of any character arising as a 160 | result of this License or out of the use or inability to use the 161 | Work (including but not limited to damages for loss of goodwill, 162 | work stoppage, computer failure or malfunction, or any and all 163 | other commercial damages or losses), even if such Contributor 164 | has been advised of the possibility of such damages. 165 | 166 | 9. Accepting Warranty or Additional Liability. While redistributing 167 | the Work or Derivative Works thereof, You may choose to offer, 168 | and charge a fee for, acceptance of support, warranty, indemnity, 169 | or other liability obligations and/or rights consistent with this 170 | License. However, in accepting such obligations, You may act only 171 | on Your own behalf and on Your sole responsibility, not on behalf 172 | of any other Contributor, and only if You agree to indemnify, 173 | defend, and hold each Contributor harmless for any liability 174 | incurred by, or claims asserted against, such Contributor by reason 175 | of your accepting any such warranty or additional liability. 176 | 177 | END OF TERMS AND CONDITIONS 178 | 179 | APPENDIX: How to apply the Apache License to your work. 180 | 181 | To apply the Apache License to your work, attach the following 182 | boilerplate notice, with the fields enclosed by brackets "[]" 183 | replaced with your own identifying information. (Don't include 184 | the brackets!) The text should be enclosed in the appropriate 185 | comment syntax for the file format. We also recommend that a 186 | file or class name and description of purpose be included on the 187 | same "printed page" as the copyright notice for easier 188 | identification within third-party archives. 189 | 190 | Copyright [yyyy] [name of copyright owner] 191 | 192 | Licensed under the Apache License, Version 2.0 (the "License"); 193 | you may not use this file except in compliance with the License. 194 | You may obtain a copy of the License at 195 | 196 | http://www.apache.org/licenses/LICENSE-2.0 197 | 198 | Unless required by applicable law or agreed to in writing, software 199 | distributed under the License is distributed on an "AS IS" BASIS, 200 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 201 | See the License for the specific language governing permissions and 202 | limitations under the License. 203 | -------------------------------------------------------------------------------- /wsc_test.go: -------------------------------------------------------------------------------- 1 | // Copyright 2019 Aporeto Inc. 2 | // Licensed under the Apache License, Version 2.0 (the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at 5 | // http://www.apache.org/licenses/LICENSE-2.0 6 | // Unless required by applicable law or agreed to in writing, software 7 | // distributed under the License is distributed on an "AS IS" BASIS, 8 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 9 | // See the License for the specific language governing permissions and 10 | // limitations under the License. 11 | 12 | package wsc 13 | 14 | import ( 15 | "context" 16 | "errors" 17 | "fmt" 18 | "net/http" 19 | "net/http/httptest" 20 | "strings" 21 | "testing" 22 | "time" 23 | 24 | "github.com/gorilla/websocket" 25 | . "github.com/smartystreets/goconvey/convey" 26 | ) 27 | 28 | type fakeWSConnection struct { 29 | readDeadlineError error 30 | writeDeadlineError error 31 | readMessageError error 32 | writeMessageError error 33 | writeControlError error 34 | closeError error 35 | 36 | pongHandler func(string) error 37 | } 38 | 39 | func (c *fakeWSConnection) SetReadDeadline(time.Time) error { return c.readDeadlineError } 40 | func (c *fakeWSConnection) SetWriteDeadline(time.Time) error { return c.writeDeadlineError } 41 | func (c *fakeWSConnection) SetCloseHandler(func(int, string) error) {} 42 | func (c *fakeWSConnection) SetPongHandler(h func(string) error) { c.pongHandler = h } 43 | func (c *fakeWSConnection) ReadMessage() (int, []byte, error) { return 0, nil, c.readMessageError } 44 | func (c *fakeWSConnection) WriteMessage(int, []byte) error { return c.writeMessageError } 45 | func (c *fakeWSConnection) WriteControl(int, []byte, time.Time) error { return c.writeControlError } 46 | func (c *fakeWSConnection) Close() error { return c.closeError } 47 | 48 | func TestWSC_ReadWrite(t *testing.T) { 49 | 50 | Convey("Given I have a webserver that works", t, func() { 51 | 52 | ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) 53 | defer cancel() 54 | 55 | var upgrader = websocket.Upgrader{ 56 | CheckOrigin: func(r *http.Request) bool { return true }, 57 | } 58 | 59 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 60 | 61 | s, err := upgrader.Upgrade(w, r, nil) 62 | if err != nil { 63 | panic(err) 64 | } 65 | 66 | h, err := Accept(ctx, s, Config{}) 67 | if err != nil { 68 | panic(err) 69 | } 70 | 71 | h.Write(<-h.Read()) 72 | 73 | <-ctx.Done() 74 | })) 75 | defer ts.Close() 76 | 77 | Convey("When I connect to the webserver", func() { 78 | 79 | s, resp, err := Connect( 80 | ctx, 81 | strings.Replace(ts.URL, "http://", "ws://", 1), 82 | Config{}, 83 | ) 84 | 85 | Convey("Then err should be nil", func() { 86 | So(err, ShouldBeNil) 87 | }) 88 | 89 | Convey("Then resp should be correct", func() { 90 | So(resp, ShouldNotBeNil) 91 | So(resp.Status, ShouldEqual, "101 Switching Protocols") 92 | }) 93 | 94 | Convey("When I listen for a message", func() { 95 | 96 | s.Write([]byte("hello")) 97 | msg := <-s.Read() 98 | 99 | Convey("Then msg should be correct", func() { 100 | So(string(msg), ShouldEqual, "hello") 101 | }) 102 | 103 | Convey("When I close the connection", func() { 104 | 105 | doneErr := make(chan error) 106 | go func() { 107 | select { 108 | case e := <-s.Done(): 109 | doneErr <- e 110 | case <-ctx.Done(): 111 | doneErr <- errors.New("test: no response in time") 112 | } 113 | }() 114 | 115 | s.Close(0) 116 | 117 | Convey("Then doneErr should be nil", func() { 118 | So(<-doneErr, ShouldBeNil) 119 | }) 120 | }) 121 | }) 122 | }) 123 | }) 124 | } 125 | 126 | func TestWSC_ReadFull(t *testing.T) { 127 | 128 | Convey("Given I have a webserver that works", t, func() { 129 | 130 | ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) 131 | defer cancel() 132 | 133 | var upgrader = websocket.Upgrader{ 134 | CheckOrigin: func(r *http.Request) bool { return true }, 135 | } 136 | 137 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 138 | 139 | s, err := upgrader.Upgrade(w, r, nil) 140 | if err != nil { 141 | panic(err) 142 | } 143 | 144 | h, err := Accept(ctx, s, Config{}) 145 | if err != nil { 146 | panic(err) 147 | } 148 | 149 | h.Write([]byte{}) 150 | h.Write([]byte{}) 151 | 152 | <-ctx.Done() 153 | })) 154 | defer ts.Close() 155 | 156 | Convey("When I connect to the webserver", func() { 157 | 158 | s, _, _ := Connect( 159 | ctx, 160 | strings.Replace(ts.URL, "http://", "ws://", 1), 161 | Config{ 162 | ReadChanSize: 1, 163 | }, 164 | ) 165 | 166 | Convey("When I send for a message", func() { 167 | 168 | s.Write([]byte("hello")) 169 | <-time.After(300 * time.Millisecond) 170 | 171 | var err error 172 | select { 173 | case err = <-s.Error(): 174 | case <-time.After(time.Second): 175 | panic("did not receive error in time") 176 | } 177 | 178 | Convey("Then err should be correct", func() { 179 | So(err, ShouldNotBeNil) 180 | So(err, ShouldEqual, ErrReadMessageDiscarded) 181 | }) 182 | }) 183 | }) 184 | }) 185 | } 186 | 187 | func TestWSC_WriteFull(t *testing.T) { 188 | 189 | Convey("Given I have a webserver that works", t, func() { 190 | 191 | ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) 192 | defer cancel() 193 | 194 | var upgrader = websocket.Upgrader{ 195 | CheckOrigin: func(r *http.Request) bool { return true }, 196 | } 197 | 198 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 199 | 200 | s, err := upgrader.Upgrade(w, r, nil) 201 | if err != nil { 202 | panic(err) 203 | } 204 | 205 | _, err = Accept(ctx, s, Config{}) 206 | if err != nil { 207 | panic(err) 208 | } 209 | 210 | <-ctx.Done() 211 | })) 212 | defer ts.Close() 213 | 214 | Convey("When I connect to the webserver", func() { 215 | 216 | s, _, _ := Connect( 217 | ctx, 218 | strings.Replace(ts.URL, "http://", "ws://", 1), 219 | Config{ 220 | WriteChanSize: 1, 221 | }, 222 | ) 223 | 224 | Convey("When I send for a message", func() { 225 | 226 | s.Write([]byte{}) 227 | s.Write([]byte{}) 228 | s.Write([]byte{}) 229 | s.Write([]byte{}) 230 | s.Write([]byte{}) 231 | 232 | var err error 233 | select { 234 | case err = <-s.Error(): 235 | case <-time.After(time.Second): 236 | panic("did not receive error in time") 237 | } 238 | 239 | Convey("Then err should be correct", func() { 240 | So(err, ShouldNotBeNil) 241 | So(err, ShouldEqual, ErrWriteMessageDiscarded) 242 | }) 243 | }) 244 | }) 245 | }) 246 | } 247 | 248 | func TestWSC_ConnectToServerWithHTTPError(t *testing.T) { 249 | 250 | Convey("Given I have a webserver that returns an http error", t, func() { 251 | 252 | ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) 253 | defer cancel() 254 | 255 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 256 | http.Error(w, "nope man", http.StatusForbidden) 257 | })) 258 | defer ts.Close() 259 | 260 | Convey("When I connect to the webserver", func() { 261 | 262 | ws, resp, err := Connect(ctx, strings.Replace(ts.URL, "http://", "ws://", 1), Config{}) 263 | 264 | Convey("Then ws should be nil", func() { 265 | So(ws, ShouldBeNil) 266 | }) 267 | 268 | Convey("Then err should not be nil", func() { 269 | So(err, ShouldNotBeNil) 270 | So(err.Error(), ShouldEqual, "websocket: bad handshake") 271 | }) 272 | 273 | Convey("Then resp should be correct", func() { 274 | So(resp, ShouldNotBeNil) 275 | So(resp.Status, ShouldEqual, "403 Forbidden") 276 | }) 277 | }) 278 | }) 279 | } 280 | 281 | func TestWSC_CannotConnect(t *testing.T) { 282 | 283 | Convey("Given I have a no webserver", t, func() { 284 | 285 | ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) 286 | defer cancel() 287 | 288 | Convey("When I connect to the non existing server", func() { 289 | 290 | ws, resp, err := Connect(ctx, "ws://127.0.0.1:7745", Config{}) 291 | 292 | Convey("Then ws should be nil", func() { 293 | So(ws, ShouldBeNil) 294 | }) 295 | 296 | Convey("Then err should not be nil", func() { 297 | So(err, ShouldNotBeNil) 298 | So(err.Error(), ShouldEndWith, "connection refused") 299 | }) 300 | 301 | Convey("Then resp should be nil", func() { 302 | So(resp, ShouldBeNil) 303 | }) 304 | }) 305 | }) 306 | } 307 | 308 | func TestWSC_GentleServerDisconnection(t *testing.T) { 309 | 310 | Convey("Given I have a webserver", t, func() { 311 | 312 | ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) 313 | defer cancel() 314 | 315 | var upgrader = websocket.Upgrader{ 316 | CheckOrigin: func(r *http.Request) bool { return true }, 317 | } 318 | 319 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 320 | 321 | ws, err := upgrader.Upgrade(w, r, nil) 322 | if err != nil { 323 | panic(err) 324 | } 325 | 326 | h, err := Accept(ctx, ws, Config{}) 327 | if err != nil { 328 | panic(err) 329 | } 330 | 331 | h.Close(0) 332 | })) 333 | defer ts.Close() 334 | 335 | Convey("When I connect to the webserver", func() { 336 | 337 | ws, _, _ := Connect(ctx, strings.Replace(ts.URL, "http://", "ws://", 1), Config{}) 338 | 339 | Convey("When I wait for a message", func() { 340 | 341 | var err error 342 | select { 343 | case err = <-ws.Done(): 344 | case <-ws.Read(): 345 | panic("test: should not have received message") 346 | case <-ctx.Done(): 347 | panic("test: no response in time") 348 | } 349 | 350 | Convey("Then err should be nil", func() { 351 | So(err, ShouldNotBeNil) 352 | So(err.Error(), ShouldEqual, "websocket: close 1001 (going away)") 353 | }) 354 | }) 355 | }) 356 | }) 357 | } 358 | 359 | func TestWSC_BrutalServerDisconnection(t *testing.T) { 360 | 361 | Convey("Given I have a webserver", t, func() { 362 | 363 | ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) 364 | defer cancel() 365 | 366 | var upgrader = websocket.Upgrader{ 367 | CheckOrigin: func(r *http.Request) bool { return true }, 368 | } 369 | 370 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 371 | 372 | ws, err := upgrader.Upgrade(w, r, nil) 373 | if err != nil { 374 | panic(err) 375 | } 376 | ws.Close() // nolint: errcheck 377 | })) 378 | defer ts.Close() 379 | 380 | Convey("When I connect to the webserver", func() { 381 | 382 | ws, _, _ := Connect(ctx, strings.Replace(ts.URL, "http://", "ws://", 1), Config{}) 383 | 384 | Convey("When I wait for a message", func() { 385 | 386 | var err error 387 | select { 388 | case err = <-ws.Done(): 389 | case <-ws.Read(): 390 | panic("test: should not have received message") 391 | case <-ctx.Done(): 392 | panic("test: no response in time") 393 | } 394 | 395 | Convey("Then err should be nil", func() { 396 | So(err, ShouldNotBeNil) 397 | So(err.Error(), ShouldEqual, "websocket: close 1006 (abnormal closure): unexpected EOF") 398 | }) 399 | }) 400 | }) 401 | }) 402 | } 403 | 404 | func TestWSC_GentleClientDisconnection(t *testing.T) { 405 | 406 | Convey("Given I have a webserver", t, func() { 407 | 408 | ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) 409 | defer cancel() 410 | 411 | var upgrader = websocket.Upgrader{ 412 | CheckOrigin: func(r *http.Request) bool { return true }, 413 | } 414 | 415 | rcvmsg := make(chan []byte) 416 | rcvdone := make(chan error) 417 | 418 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 419 | 420 | ws, err := upgrader.Upgrade(w, r, nil) 421 | if err != nil { 422 | panic(err) 423 | } 424 | 425 | h, err := Accept(ctx, ws, Config{}) 426 | if err != nil { 427 | panic(err) 428 | } 429 | 430 | select { 431 | case err = <-h.Done(): 432 | rcvdone <- err 433 | case msg := <-h.Read(): 434 | rcvmsg <- msg 435 | case <-ctx.Done(): 436 | panic("test: no response in time") 437 | } 438 | 439 | })) 440 | defer ts.Close() 441 | 442 | Convey("When I connect to the webserver", func() { 443 | 444 | ws, _, _ := Connect(ctx, strings.Replace(ts.URL, "http://", "ws://", 1), Config{}) 445 | 446 | Convey("When I gracefully stop the connection", func() { 447 | 448 | ws.Close(websocket.CloseInvalidFramePayloadData) 449 | 450 | var err error 451 | var msg []byte 452 | select { 453 | case err = <-rcvdone: 454 | case msg = <-rcvmsg: 455 | case <-time.After(1 * time.Second): 456 | panic("test: no response in time") 457 | } 458 | 459 | Convey("Then the err received by the client not be nil", func() { 460 | So(err, ShouldNotBeNil) 461 | So(err.Error(), ShouldEqual, "websocket: close 1007 (invalid payload data)") 462 | }) 463 | 464 | Convey("Then no msg should be received by the client", func() { 465 | So(msg, ShouldBeNil) 466 | }) 467 | }) 468 | }) 469 | }) 470 | } 471 | 472 | func TestWSC_BrutalClientDisconnection(t *testing.T) { 473 | 474 | Convey("Given I have a webserver", t, func() { 475 | 476 | ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) 477 | defer cancel() 478 | 479 | var upgrader = websocket.Upgrader{ 480 | CheckOrigin: func(r *http.Request) bool { return true }, 481 | } 482 | 483 | rcvmsg := make(chan []byte) 484 | rcvdone := make(chan error) 485 | 486 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 487 | 488 | ws, err := upgrader.Upgrade(w, r, nil) 489 | if err != nil { 490 | panic(err) 491 | } 492 | 493 | h, err := Accept(ctx, ws, Config{}) 494 | if err != nil { 495 | panic(err) 496 | } 497 | 498 | select { 499 | case err = <-h.Done(): 500 | rcvdone <- err 501 | case msg := <-h.Read(): 502 | rcvmsg <- msg 503 | case <-ctx.Done(): 504 | panic("test: no response in time") 505 | } 506 | })) 507 | defer ts.Close() 508 | 509 | Convey("When I connect to the webserver", func() { 510 | 511 | w, _, _ := Connect(ctx, strings.Replace(ts.URL, "http://", "ws://", 1), Config{}) 512 | 513 | Convey("When I gracefully stop the connection", func() { 514 | 515 | w.(*ws).conn.Close() // nolint: errcheck 516 | 517 | var err error 518 | var msg []byte 519 | select { 520 | case err = <-rcvdone: 521 | case msg = <-rcvmsg: 522 | case <-ctx.Done(): 523 | panic("test: no response in time") 524 | } 525 | 526 | Convey("Then the err received by the server not be nil", func() { 527 | So(err, ShouldNotBeNil) 528 | So(err.Error(), ShouldEqual, "websocket: close 1006 (abnormal closure): unexpected EOF") 529 | }) 530 | 531 | Convey("Then no msg should be received by the server", func() { 532 | So(msg, ShouldBeNil) 533 | }) 534 | }) 535 | }) 536 | }) 537 | } 538 | 539 | func TestWSC_ServerMissingPong(t *testing.T) { 540 | 541 | Convey("Given I have a webserver", t, func() { 542 | 543 | ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) 544 | defer cancel() 545 | 546 | var upgrader = websocket.Upgrader{ 547 | CheckOrigin: func(r *http.Request) bool { return true }, 548 | } 549 | 550 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 551 | 552 | ws, err := upgrader.Upgrade(w, r, nil) 553 | if err != nil { 554 | panic(err) 555 | } 556 | 557 | _, err = Accept(ctx, ws, Config{}) 558 | if err != nil { 559 | panic(err) 560 | } 561 | 562 | <-ctx.Done() 563 | })) 564 | defer ts.Close() 565 | 566 | Convey("When I connect to the webserver", func() { 567 | 568 | ws, _, _ := Connect( 569 | ctx, strings.Replace(ts.URL, "http://", "ws://", 1), Config{ 570 | PongWait: 1 * time.Nanosecond, // we wait for nothing 571 | PingPeriod: 50 * time.Millisecond, 572 | }) 573 | 574 | Convey("When I wait for a message", func() { 575 | 576 | <-time.After(300 * time.Millisecond) 577 | 578 | var err error 579 | var msg []byte 580 | select { 581 | case err = <-ws.Done(): 582 | case msg = <-ws.Read(): 583 | case <-ctx.Done(): 584 | panic("test: no response in time") 585 | } 586 | 587 | Convey("Then the err received by the client not be nil", func() { 588 | So(err, ShouldNotBeNil) 589 | So(err.Error(), ShouldEndWith, "i/o timeout") 590 | }) 591 | 592 | Convey("Then no msg should be received by the client", func() { 593 | So(msg, ShouldBeNil) 594 | }) 595 | }) 596 | }) 597 | }) 598 | } 599 | 600 | func TestWSC_ClientMissingPong(t *testing.T) { 601 | 602 | Convey("Given I have a webserver", t, func() { 603 | 604 | ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) 605 | defer cancel() 606 | 607 | var upgrader = websocket.Upgrader{ 608 | CheckOrigin: func(r *http.Request) bool { return true }, 609 | } 610 | 611 | rcvmsg := make(chan []byte) 612 | rcvdone := make(chan error) 613 | 614 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 615 | 616 | ws, err := upgrader.Upgrade(w, r, nil) 617 | if err != nil { 618 | panic(err) 619 | } 620 | 621 | h, err := Accept(ctx, ws, Config{ 622 | PongWait: 1 * time.Millisecond, 623 | PingPeriod: 50 * time.Millisecond, 624 | }) 625 | if err != nil { 626 | panic(err) 627 | } 628 | 629 | select { 630 | case err = <-h.Done(): 631 | rcvdone <- err 632 | case msg := <-h.Read(): 633 | rcvmsg <- msg 634 | case <-ctx.Done(): 635 | panic("test: no response in time") 636 | } 637 | 638 | })) 639 | defer ts.Close() 640 | 641 | Convey("When I connect to the webserver", func() { 642 | 643 | _, _, _ = Connect(ctx, strings.Replace(ts.URL, "http://", "ws://", 1), Config{}) 644 | 645 | Convey("When I wait for a message", func() { 646 | 647 | <-time.After(300 * time.Millisecond) 648 | 649 | var err error 650 | var msg []byte 651 | select { 652 | case err = <-rcvdone: 653 | case msg = <-rcvmsg: 654 | case <-ctx.Done(): 655 | panic("test: no response in time") 656 | } 657 | 658 | Convey("Then the err received by the server not be nil", func() { 659 | So(err, ShouldNotBeNil) 660 | So(err.Error(), ShouldEndWith, "i/o timeout") 661 | }) 662 | 663 | Convey("Then no msg should be received by the server", func() { 664 | So(msg, ShouldBeNil) 665 | }) 666 | }) 667 | }) 668 | }) 669 | } 670 | 671 | func TestWWS_AcceptWithFailedReadDeadline(t *testing.T) { 672 | 673 | Convey("Given I have a wsconn", t, func() { 674 | 675 | conn := &fakeWSConnection{ 676 | readDeadlineError: fmt.Errorf("failed"), 677 | } 678 | 679 | Convey("When I call Accept", func() { 680 | 681 | ws, err := Accept(context.Background(), conn, Config{}) 682 | 683 | Convey("Then err should be correct", func() { 684 | So(err, ShouldEqual, conn.readDeadlineError) 685 | }) 686 | 687 | Convey("Then ws should be nil", func() { 688 | So(ws, ShouldBeNil) 689 | }) 690 | }) 691 | }) 692 | } 693 | 694 | func TestWSC_writePumpWithWriteErrorForPing(t *testing.T) { 695 | 696 | Convey("Given i have wsconn and ws with a running write pump", t, func() { 697 | 698 | conn := &fakeWSConnection{ 699 | writeMessageError: fmt.Errorf("failed"), 700 | } 701 | 702 | s := &ws{ 703 | conn: conn, 704 | doneChan: make(chan error, 1), 705 | config: Config{ 706 | PingPeriod: 1 * time.Millisecond, 707 | }, 708 | } 709 | 710 | ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) 711 | defer cancel() 712 | 713 | errCh := make(chan error) 714 | go func() { 715 | select { 716 | case e := <-s.Done(): 717 | errCh <- e 718 | case <-ctx.Done(): 719 | panic("did not receive the expected error in time") 720 | } 721 | }() 722 | 723 | go s.writePump(ctx) 724 | 725 | Convey("When I read the errors", func() { 726 | 727 | Convey("Then the error should be correct", func() { 728 | So(<-errCh, ShouldEqual, conn.writeMessageError) 729 | }) 730 | }) 731 | }) 732 | } 733 | 734 | func TestWSC_writePumpWithWriteErrorForWrite(t *testing.T) { 735 | 736 | Convey("Given i have wsconn and ws with a running write pump", t, func() { 737 | 738 | conn := &fakeWSConnection{ 739 | writeMessageError: fmt.Errorf("failed"), 740 | } 741 | 742 | s := &ws{ 743 | conn: conn, 744 | doneChan: make(chan error, 1), 745 | writeChan: make(chan []byte, 2), 746 | config: Config{ 747 | PingPeriod: 10 * time.Millisecond, 748 | }, 749 | } 750 | 751 | ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) 752 | defer cancel() 753 | 754 | errCh := make(chan error) 755 | go func() { 756 | select { 757 | case e := <-s.Done(): 758 | errCh <- e 759 | case <-ctx.Done(): 760 | panic("did not receive the expected error in time") 761 | } 762 | }() 763 | 764 | go s.writePump(ctx) 765 | 766 | s.writeChan <- []byte{} 767 | Convey("When I read the errors", func() { 768 | 769 | Convey("Then the error should be correct", func() { 770 | So(<-errCh, ShouldEqual, conn.writeMessageError) 771 | }) 772 | }) 773 | }) 774 | } 775 | 776 | func TestWSC_PongHandlerWithError(t *testing.T) { 777 | 778 | Convey("Given I have a wsconn", t, func() { 779 | 780 | conn := &fakeWSConnection{} 781 | 782 | Convey("When I call Accept", func() { 783 | 784 | _, _ = Accept(context.Background(), conn, Config{}) 785 | 786 | err := conn.pongHandler("hello") 787 | 788 | Convey("Then err should be correct", func() { 789 | So(err, ShouldBeNil) 790 | }) 791 | }) 792 | }) 793 | } 794 | --------------------------------------------------------------------------------