Skip to content

Commit 69ff675

Browse files
nhooyrnhooyr
committed
More tests and fixes
1 parent b53f306 commit 69ff675

File tree

9 files changed

+170
-314
lines changed

9 files changed

+170
-314
lines changed

‎close_notjs.go‎

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ func (c *Conn) closeHandshake(code StatusCode, reason string) (err error) {
3535
defer errd.Wrap(&err, "failed to close WebSocket")
3636

3737
err = c.writeClose(code, reason)
38-
if err != nil {
38+
if CloseStatus(err) == -1 {
3939
return err
4040
}
4141

@@ -46,12 +46,6 @@ func (c *Conn) closeHandshake(code StatusCode, reason string) (err error) {
4646
return nil
4747
}
4848

49-
func (c *Conn) writeError(code StatusCode, err error) {
50-
c.setCloseErr(err)
51-
c.writeClose(code, err.Error())
52-
c.close(nil)
53-
}
54-
5549
func (c *Conn) writeClose(code StatusCode, reason string) error {
5650
c.closeMu.Lock()
5751
closing := c.wroteClose
@@ -70,7 +64,12 @@ func (c *Conn) writeClose(code StatusCode, reason string) error {
7064

7165
var p []byte
7266
if ce.Code != StatusNoStatusRcvd {
73-
p = ce.bytes()
67+
var err error
68+
p, err = ce.bytes()
69+
if err != nil {
70+
log.Printf("websocket: %v", err)
71+
return err
72+
}
7473
}
7574

7675
return c.writeControl(context.Background(), opClose, p)
@@ -148,16 +147,16 @@ func validWireCloseCode(code StatusCode) bool {
148147
return false
149148
}
150149

151-
func (ce CloseError) bytes() []byte {
150+
func (ce CloseError) bytes() ([]byte, error) {
152151
p, err := ce.bytesErr()
153152
if err != nil {
154-
log.Printf("websocket: failed to marshal close frame: %v", err)
153+
err = xerrors.Errorf("failed to marshal close frame: %w", err)
155154
ce = CloseError{
156155
Code: StatusInternalError,
157156
}
158157
p, _ = ce.bytesErr()
159158
}
160-
return p
159+
return p, err
161160
}
162161

163162
const maxCloseReason = maxControlPayload - 2

‎compress_notjs.go‎

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,6 @@ import (
1010
)
1111

1212
func (m CompressionMode) opts() *compressionOptions {
13-
if m == CompressionDisabled {
14-
return nil
15-
}
1613
return &compressionOptions{
1714
clientNoContextTakeover: m == CompressionNoContextTakeover,
1815
serverNoContextTakeover: m == CompressionNoContextTakeover,
@@ -42,14 +39,6 @@ func (copts *compressionOptions) setHeader(h http.Header) {
4239
// trying to return more bytes.
4340
const deflateMessageTail = "\x00\x00\xff\xff"
4441

45-
func (c *Conn) writeNoContextTakeOver() bool {
46-
return c.client && c.copts.clientNoContextTakeover || !c.client && c.copts.serverNoContextTakeover
47-
}
48-
49-
func (c *Conn) readNoContextTakeOver() bool {
50-
return !c.client && c.copts.clientNoContextTakeover || c.client && c.copts.serverNoContextTakeover
51-
}
52-
5342
type trimLastFourBytesWriter struct {
5443
w io.Writer
5544
tail []byte

‎conn_notjs.go‎

Lines changed: 7 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -87,12 +87,6 @@ func newConn(cfg connConfig) *Conn {
8787
closed: make(chan struct{}),
8888
activePings: make(map[string]chan<- struct{}),
8989
}
90-
if c.flate() && c.flateThreshold == 0 {
91-
c.flateThreshold = 256
92-
if c.writeNoContextTakeOver() {
93-
c.flateThreshold = 512
94-
}
95-
}
9690

9791
c.readMu = newMu(c)
9892
c.writeFrameMu = newMu(c)
@@ -104,6 +98,13 @@ func newConn(cfg connConfig) *Conn {
10498
c.writeBuf = extractBufioWriterBuf(c.bw, c.rwc)
10599
}
106100

101+
if c.flate() && c.flateThreshold == 0 {
102+
c.flateThreshold = 256
103+
if !c.msgWriter.flateContextTakeover() {
104+
c.flateThreshold = 512
105+
}
106+
}
107+
107108
runtime.SetFinalizer(c, func(c *Conn) {
108109
c.close(xerrors.New("connection garbage collected"))
109110
})
@@ -247,15 +248,6 @@ func (m *mu) Lock(ctx context.Context) error {
247248
}
248249
}
249250

250-
func (m *mu) TryLock() bool {
251-
select {
252-
case m.ch <- struct{}{}:
253-
return true
254-
default:
255-
return false
256-
}
257-
}
258-
259251
func (m *mu) Unlock() {
260252
select {
261253
case <-m.ch:

‎conn_test.go‎

Lines changed: 133 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616
"golang.org/x/xerrors"
1717

1818
"nhooyr.io/websocket"
19+
"nhooyr.io/websocket/internal/test/cmp"
1920
"nhooyr.io/websocket/internal/test/wstest"
2021
"nhooyr.io/websocket/internal/test/xrand"
2122
"nhooyr.io/websocket/internal/xsync"
@@ -31,9 +32,6 @@ func TestConn(t *testing.T) {
3132
t.Run("", func(t *testing.T) {
3233
t.Parallel()
3334

34-
ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
35-
defer cancel()
36-
3735
copts := &websocket.CompressionOptions{
3836
Mode: websocket.CompressionMode(xrand.Int(int(websocket.CompressionDisabled) + 1)),
3937
Threshold: xrand.Int(9999),
@@ -47,11 +45,14 @@ func TestConn(t *testing.T) {
4745
if err != nil {
4846
t.Fatal(err)
4947
}
50-
defer c1.Close(websocket.StatusInternalError, "")
5148
defer c2.Close(websocket.StatusInternalError, "")
49+
defer c1.Close(websocket.StatusInternalError, "")
50+
51+
ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
52+
defer cancel()
5253

5354
echoLoopErr := xsync.Go(func() error {
54-
err := wstest.EchoLoop(ctx, c1)
55+
err := wstest.EchoLoop(ctx, c2)
5556
return assertCloseStatus(websocket.StatusNormalClosure, err)
5657
})
5758
defer func() {
@@ -62,19 +63,143 @@ func TestConn(t *testing.T) {
6263
}()
6364
defer cancel()
6465

65-
c2.SetReadLimit(131072)
66+
c1.SetReadLimit(131072)
6667

6768
for i := 0; i < 5; i++ {
68-
err := wstest.Echo(ctx, c2, 131072)
69+
err := wstest.Echo(ctx, c1, 131072)
6970
if err != nil {
7071
t.Fatal(err)
7172
}
7273
}
7374

74-
c2.Close(websocket.StatusNormalClosure, "")
75+
err = c1.Close(websocket.StatusNormalClosure, "")
76+
if err != nil {
77+
t.Fatalf("unexpected error: %v", err)
78+
}
7579
})
7680
}
7781
})
82+
83+
t.Run("badClose", func(t *testing.T) {
84+
t.Parallel()
85+
86+
c1, c2, err := wstest.Pipe(nil, nil)
87+
if err != nil {
88+
t.Fatal(err)
89+
}
90+
defer c1.Close(websocket.StatusInternalError, "")
91+
defer c2.Close(websocket.StatusInternalError, "")
92+
93+
err = c1.Close(-1, "")
94+
if !cmp.ErrorContains(err, "failed to marshal close frame: status code StatusCode(-1) cannot be set") {
95+
t.Fatalf("unexpected error: %v", err)
96+
}
97+
})
98+
99+
t.Run("ping", func(t *testing.T) {
100+
t.Parallel()
101+
102+
c1, c2, err := wstest.Pipe(nil, nil)
103+
if err != nil {
104+
t.Fatal(err)
105+
}
106+
defer c1.Close(websocket.StatusInternalError, "")
107+
defer c2.Close(websocket.StatusInternalError, "")
108+
109+
ctx, cancel := context.WithTimeout(context.Background(), time.Second*15)
110+
defer cancel()
111+
112+
c2.CloseRead(ctx)
113+
c1.CloseRead(ctx)
114+
115+
for i := 0; i < 10; i++ {
116+
err = c1.Ping(ctx)
117+
if err != nil {
118+
t.Fatal(err)
119+
}
120+
}
121+
122+
err = c1.Close(websocket.StatusNormalClosure, "")
123+
if err != nil {
124+
t.Fatalf("unexpected error: %v", err)
125+
}
126+
})
127+
128+
t.Run("badPing", func(t *testing.T) {
129+
t.Parallel()
130+
131+
c1, c2, err := wstest.Pipe(nil, nil)
132+
if err != nil {
133+
t.Fatal(err)
134+
}
135+
defer c1.Close(websocket.StatusInternalError, "")
136+
defer c2.Close(websocket.StatusInternalError, "")
137+
138+
ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
139+
defer cancel()
140+
141+
c2.CloseRead(ctx)
142+
143+
err = c1.Ping(ctx)
144+
if !cmp.ErrorContains(err, "failed to wait for pong") {
145+
t.Fatalf("unexpected error: %v", err)
146+
}
147+
})
148+
149+
t.Run("concurrentWrite", func(t *testing.T) {
150+
t.Parallel()
151+
152+
c1, c2, err := wstest.Pipe(nil, nil)
153+
if err != nil {
154+
t.Fatal(err)
155+
}
156+
defer c2.Close(websocket.StatusInternalError, "")
157+
defer c1.Close(websocket.StatusInternalError, "")
158+
159+
ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
160+
defer cancel()
161+
162+
discardLoopErr := xsync.Go(func() error {
163+
for {
164+
_, _, err := c2.Read(ctx)
165+
if websocket.CloseStatus(err) == websocket.StatusNormalClosure {
166+
return nil
167+
}
168+
if err != nil {
169+
return err
170+
}
171+
}
172+
})
173+
defer func() {
174+
err := <-discardLoopErr
175+
if err != nil {
176+
t.Errorf("discard loop error: %v", err)
177+
}
178+
}()
179+
defer cancel()
180+
181+
msg := xrand.Bytes(xrand.Int(9999))
182+
const count = 100
183+
errs := make(chan error, count)
184+
185+
for i := 0; i < count; i++ {
186+
go func() {
187+
errs <- c1.Write(ctx, websocket.MessageBinary, msg)
188+
}()
189+
}
190+
191+
for i := 0; i < count; i++ {
192+
err := <-errs
193+
if err != nil {
194+
t.Fatal(err)
195+
}
196+
}
197+
198+
err = c1.Close(websocket.StatusNormalClosure, "")
199+
if err != nil {
200+
t.Fatalf("unexpected error: %v", err)
201+
}
202+
})
78203
}
79204

80205
func TestWasm(t *testing.T) {

‎dial_test.go‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import (
1313
"testing"
1414
"time"
1515

16-
"cdr.dev/slog/sloggers/slogtest/assert"
16+
"nhooyr.io/websocket/internal/test/cmp"
1717
)
1818

1919
func TestBadDials(t *testing.T) {
@@ -70,7 +70,9 @@ func TestBadDials(t *testing.T) {
7070
}
7171

7272
_, _, err := dial(ctx, tc.url, tc.opts, tc.rand)
73-
assert.Error(t, "dial", err)
73+
if err == nil {
74+
t.Fatalf("expected error")
75+
}
7476
})
7577
}
7678
})
@@ -88,7 +90,9 @@ func TestBadDials(t *testing.T) {
8890
}, nil
8991
}),
9092
})
91-
assert.ErrorContains(t, "dial", err, "failed to WebSocket dial: expected handshake response status code 101 but got 0")
93+
if !cmp.ErrorContains(err, "failed to WebSocket dial: expected handshake response status code 101 but got 0") {
94+
t.Fatal(err)
95+
}
9296
})
9397

9498
t.Run("badBody", func(t *testing.T) {
@@ -113,7 +117,9 @@ func TestBadDials(t *testing.T) {
113117
_, _, err := Dial(ctx, "ws://example.com", &DialOptions{
114118
HTTPClient: mockHTTPClient(rt),
115119
})
116-
assert.ErrorContains(t, "dial", err, "response body is not a io.ReadWriteCloser")
120+
if !cmp.ErrorContains(err, "response body is not a io.ReadWriteCloser") {
121+
t.Fatal(err)
122+
}
117123
})
118124
}
119125

@@ -211,7 +217,9 @@ func Test_verifyServerHandshake(t *testing.T) {
211217

212218
r := httptest.NewRequest("GET", "/", nil)
213219
key, err := secWebSocketKey(rand.Reader)
214-
assert.Success(t, "secWebSocketKey", err)
220+
if err != nil {
221+
t.Fatal(err)
222+
}
215223
r.Header.Set("Sec-WebSocket-Key", key)
216224

217225
if resp.Header.Get("Sec-WebSocket-Accept") == "" {

‎go.mod‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,20 +3,12 @@ module nhooyr.io/websocket
33
go 1.12
44

55
require (
6-
cdr.dev/slog v1.3.0
7-
github.com/alecthomas/chroma v0.7.1 // indirect
8-
github.com/fatih/color v1.9.0 // indirect
96
github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee // indirect
107
github.com/gobwas/pool v0.2.0 // indirect
118
github.com/gobwas/ws v1.0.2
12-
github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e // indirect
139
github.com/golang/protobuf v1.3.3
1410
github.com/google/go-cmp v0.4.0
1511
github.com/gorilla/websocket v1.4.1
16-
github.com/mattn/go-isatty v0.0.12 // indirect
17-
go.opencensus.io v0.22.3 // indirect
18-
golang.org/x/crypto v0.0.0-20200208060501-ecb85df21340 // indirect
19-
golang.org/x/sys v0.0.0-20200202164722-d101bd2416d5 // indirect
2012
golang.org/x/time v0.0.0-20191024005414-555d28b269f0
2113
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543
2214
)

0 commit comments

Comments
 (0)