Both modes seem to work :)

Author: nhooyr

accept.go

@@ -111,12 +111,14 @@ func accept(w http.ResponseWriter, r *http.Request, opts *AcceptOptions) (_ *Con
 	brw.Reader.Reset(io.MultiReader(bytes.NewReader(b), netConn))
 
 	return newConn(connConfig{
-		subprotocol: w.Header().Get("Sec-WebSocket-Protocol"),
-		rwc:         netConn,
-		client:      false,
-		copts:       copts,
-		br:          brw.Reader,
-		bw:          brw.Writer,
+		subprotocol:    w.Header().Get("Sec-WebSocket-Protocol"),
+		rwc:            netConn,
+		client:         false,
+		copts:          copts,
+		flateThreshold: opts.CompressionOptions.Threshold,
+
+		br: brw.Reader,
+		bw: brw.Writer,
 	}), nil
 }
 

assert_test.go

@@ -6,6 +6,7 @@ import (
 	"strings"
 	"testing"
 
+	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/slogtest/assert"
 
 	"nhooyr.io/websocket"
@@ -33,7 +34,7 @@ func assertJSONEcho(t *testing.T, ctx context.Context, c *websocket.Conn, n int)
 }
 
 func assertJSONRead(t *testing.T, ctx context.Context, c *websocket.Conn, exp interface{}) {
-	t.Helper()
+	slog.Helper()
 
 	var act interface{}
 	err := wsjson.Read(ctx, c, &act)

compress.go

@@ -148,12 +148,12 @@ func (tw *trimLastFourBytesWriter) Write(p []byte) (int, error) {
 
 var flateReaderPool sync.Pool
 
-func getFlateReader(r io.Reader) io.Reader {
+func getFlateReader(r io.Reader, dict []byte) io.Reader {
 	fr, ok := flateReaderPool.Get().(io.Reader)
 	if !ok {
-		return flate.NewReader(r)
+		return flate.NewReaderDict(r, dict)
 	}
-	fr.(flate.Resetter).Reset(r, nil)
+	fr.(flate.Resetter).Reset(r, dict)
 	return fr
 }
 
@@ -163,10 +163,10 @@ func putFlateReader(fr io.Reader) {
 
 var flateWriterPool sync.Pool
 
-func getFlateWriter(w io.Writer, dict []byte) *flate.Writer {
+func getFlateWriter(w io.Writer) *flate.Writer {
 	fw, ok := flateWriterPool.Get().(*flate.Writer)
 	if !ok {
-		fw, _ = flate.NewWriterDict(w, flate.BestSpeed, dict)
+		fw, _ = flate.NewWriter(w, flate.BestSpeed)
 		return fw
 	}
 	fw.Reset(w)
@@ -177,40 +177,32 @@ func putFlateWriter(w *flate.Writer) {
 	flateWriterPool.Put(w)
 }
 
-type slidingWindowReader struct {
-	window []byte
-
-	r io.Reader
+type slidingWindow struct {
+	r   io.Reader
+	buf []byte
 }
 
-func (r slidingWindowReader) Read(p []byte) (int, error) {
-	n, err := r.r.Read(p)
-	p = p[:n]
-
-	r.append(p)
-
-	return n, err
+func newSlidingWindow(n int) *slidingWindow {
+	return &slidingWindow{
+		buf: make([]byte, 0, n),
+	}
 }
 
-func (r slidingWindowReader) append(p []byte) {
-	if len(r.window) <= cap(r.window) {
-		r.window = append(r.window, p...)
+func (w *slidingWindow) write(p []byte) {
+	if len(p) >= cap(w.buf) {
+		w.buf = w.buf[:cap(w.buf)]
+		p = p[len(p)-cap(w.buf):]
+		copy(w.buf, p)
+		return
 	}
 
-	if len(p) > cap(r.window) {
-		p = p[len(p)-cap(r.window):]
+	left := cap(w.buf) - len(w.buf)
+	if left < len(p) {
+		// We need to shift spaceNeeded bytes from the end to make room for p at the end.
+		spaceNeeded := len(p) - left
+		copy(w.buf, w.buf[spaceNeeded:])
+		w.buf = w.buf[:len(w.buf)-spaceNeeded]
 	}
 
-	// p now contains at max the last window bytes
-	// so we need to be able to append all of it to r.window.
-	// Shift as many bytes from r.window as needed.
-
-	// Maximum window size minus current window minus extra gives
-	// us the number of bytes that need to be shifted.
-	off := len(r.window) + len(p) - cap(r.window)
-
-	r.window = append(r.window[:0], r.window[off:]...)
-	copy(r.window, r.window[off:])
-	copy(r.window[len(r.window)-len(p):], p)
-	return
+	w.buf = append(w.buf, p...)
 }

compress_test.go

@@ -0,0 +1,45 @@
+package websocket
+
+import (
+	"crypto/rand"
+	"encoding/base64"
+	"math/big"
+	"strings"
+	"testing"
+
+	"cdr.dev/slog/sloggers/slogtest/assert"
+)
+
+func Test_slidingWindow(t *testing.T) {
+	t.Parallel()
+
+	const testCount = 99
+	const maxWindow = 99999
+	for i := 0; i < testCount; i++ {
+		input := randStr(t, maxWindow)
+		windowLength := randInt(t, maxWindow)
+		r := newSlidingWindow(windowLength)
+		r.write([]byte(input))
+
+		if cap(r.buf) != windowLength {
+			t.Fatalf("sliding window length changed somehow: %q and windowLength %d", input, windowLength)
+		}
+		assert.True(t, "hasSuffix", strings.HasSuffix(input, string(r.buf)))
+	}
+}
+
+func randStr(t *testing.T, max int) string {
+	n := randInt(t, max)
+
+	b := make([]byte, n)
+	_, err := rand.Read(b)
+	assert.Success(t, "rand.Read", err)
+
+	return base64.StdEncoding.EncodeToString(b)
+}
+
+func randInt(t *testing.T, max int) int {
+	x, err := rand.Int(rand.Reader, big.NewInt(int64(max)))
+	assert.Success(t, "rand.Int", err)
+	return int(x.Int64())
+}

conn.go

@@ -38,12 +38,13 @@ const (
 // On any error from any method, the connection is closed
 // with an appropriate reason.
 type Conn struct {
-	subprotocol string
-	rwc         io.ReadWriteCloser
-	client      bool
-	copts       *compressionOptions
-	br          *bufio.Reader
-	bw          *bufio.Writer
+	subprotocol    string
+	rwc            io.ReadWriteCloser
+	client         bool
+	copts          *compressionOptions
+	flateThreshold int
+	br             *bufio.Reader
+	bw             *bufio.Writer
 
 	readTimeout  chan context.Context
 	writeTimeout chan context.Context
@@ -71,21 +72,23 @@ type Conn struct {
 }
 
 type connConfig struct {
-	subprotocol string
-	rwc         io.ReadWriteCloser
-	client      bool
-	copts       *compressionOptions
+	subprotocol    string
+	rwc            io.ReadWriteCloser
+	client         bool
+	copts          *compressionOptions
+	flateThreshold int
 
 	br *bufio.Reader
 	bw *bufio.Writer
 }
 
 func newConn(cfg connConfig) *Conn {
 	c := &Conn{
-		subprotocol: cfg.subprotocol,
-		rwc:         cfg.rwc,
-		client:      cfg.client,
-		copts:       cfg.copts,
+		subprotocol:    cfg.subprotocol,
+		rwc:            cfg.rwc,
+		client:         cfg.client,
+		copts:          cfg.copts,
+		flateThreshold: cfg.flateThreshold,
 
 		br: cfg.br,
 		bw: cfg.bw,
@@ -96,6 +99,12 @@ func newConn(cfg connConfig) *Conn {
 		closed:      make(chan struct{}),
 		activePings: make(map[string]chan<- struct{}),
 	}
+	if c.flateThreshold == 0 {
+		c.flateThreshold = 256
+		if c.writeNoContextTakeOver() {
+			c.flateThreshold = 512
+		}
+	}
 
 	c.readMu = newMu(c)
 	c.writeFrameMu = newMu(c)
@@ -145,12 +154,10 @@ func (c *Conn) close(err error) {
 		}
 		c.msgWriter.close()
 
+		c.msgReader.close()
 		if c.client {
-			c.readMu.Lock(context.Background())
 			putBufioReader(c.br)
-			c.readMu.Unlock()
 		}
-		c.msgReader.close()
 	}()
 }
 

conn_test.go

@@ -27,13 +27,15 @@ func TestConn(t *testing.T) {
 				Subprotocols:       []string{"echo"},
 				InsecureSkipVerify: true,
 				CompressionOptions: websocket.CompressionOptions{
-					Mode: websocket.CompressionNoContextTakeover,
+					Mode:      websocket.CompressionContextTakeover,
+					Threshold: 1,
 				},
 			})
 			assert.Success(t, "accept", err)
 			defer c.Close(websocket.StatusInternalError, "")
 
 			err = echoLoop(r.Context(), c)
+			t.Logf("server: %v", err)
 			assertCloseStatus(t, websocket.StatusNormalClosure, err)
 		}, false)
 		defer closeFn()
@@ -46,7 +48,8 @@ func TestConn(t *testing.T) {
 		opts := &websocket.DialOptions{
 			Subprotocols: []string{"echo"},
 			CompressionOptions: websocket.CompressionOptions{
-				Mode: websocket.CompressionNoContextTakeover,
+				Mode:      websocket.CompressionContextTakeover,
+				Threshold: 1,
 			},
 		}
 		opts.HTTPClient = s.Client()

dial.go

@@ -99,12 +99,13 @@ func dial(ctx context.Context, urls string, opts *DialOptions, rand io.Reader) (
 	}
 
 	return newConn(connConfig{
-		subprotocol: resp.Header.Get("Sec-WebSocket-Protocol"),
-		rwc:         rwc,
-		client:      true,
-		copts:       copts,
-		br:          getBufioReader(rwc),
-		bw:          getBufioWriter(rwc),
+		subprotocol:    resp.Header.Get("Sec-WebSocket-Protocol"),
+		rwc:            rwc,
+		client:         true,
+		copts:          copts,
+		flateThreshold: opts.CompressionOptions.Threshold,
+		br:             getBufioReader(rwc),
+		bw:             getBufioWriter(rwc),
 	}), resp, nil
 }