Minor fixes

Author: nhooyr

accept.go

@@ -12,6 +12,7 @@ import (
 	"net/http"
 	"net/textproto"
 	"net/url"
+	"strconv"
 	"strings"
 
 	"nhooyr.io/websocket/internal/errd"
@@ -208,6 +209,7 @@ func acceptCompression(r *http.Request, w http.ResponseWriter, mode CompressionM
 
 func acceptDeflate(w http.ResponseWriter, ext websocketExtension, mode CompressionMode) (*compressionOptions, error) {
 	copts := mode.opts()
+	copts.serverMaxWindowBits = 8
 
 	for _, p := range ext.params {
 		switch p {
@@ -219,7 +221,27 @@ func acceptDeflate(w http.ResponseWriter, ext websocketExtension, mode Compressi
 			continue
 		}
 
-		if strings.HasPrefix(p, "client_max_window_bits") || strings.HasPrefix(p, "server_max_window_bits") {
+		if strings.HasPrefix(p, "client_max_window_bits") {
+			continue
+
+			// bits, ok := parseExtensionParameter(p, 15)
+			// if !ok || bits < 8 || bits > 16 {
+			// 	err := fmt.Errorf("invalid client_max_window_bits: %q", p)
+			// 	http.Error(w, err.Error(), http.StatusBadRequest)
+			// 	return nil, err
+			// }
+			// copts.clientMaxWindowBits = bits
+			// continue
+		}
+
+		if false && strings.HasPrefix(p, "server_max_window_bits") {
+			// We always send back 8 but make sure to validate.
+			bits, ok := parseExtensionParameter(p, 0)
+			if !ok || bits < 8 || bits > 16 {
+				err := fmt.Errorf("invalid server_max_window_bits: %q", p)
+				http.Error(w, err.Error(), http.StatusBadRequest)
+				return nil, err
+			}
 			continue
 		}
 
@@ -233,6 +255,21 @@ func acceptDeflate(w http.ResponseWriter, ext websocketExtension, mode Compressi
 	return copts, nil
 }
 
+// parseExtensionParameter parses the value in the extension parameter p.
+// It falls back to defaultVal if there is no value.
+// If defaultVal == 0, then ok == false if there is no value.
+func parseExtensionParameter(p string, defaultVal int) (int, bool) {
+	ps := strings.Split(p, "=")
+	if len(ps) == 1 {
+		if defaultVal > 0 {
+			return defaultVal, true
+		}
+		return 0, false
+	}
+	i, e := strconv.Atoi(strings.Trim(ps[1], `"`))
+	return i, e == nil
+}
+
 func acceptWebkitDeflate(w http.ResponseWriter, ext websocketExtension, mode CompressionMode) (*compressionOptions, error) {
 	copts := mode.opts()
 	// The peer must explicitly request it.

accept_test.go

@@ -327,6 +327,7 @@ func Test_acceptCompression(t *testing.T) {
 			expCopts: &compressionOptions{
 				clientNoContextTakeover: true,
 				serverNoContextTakeover: true,
+				serverMaxWindowBits:     8,
 			},
 		},
 		{

compress_notjs.go

@@ -3,6 +3,7 @@
 package websocket
 
 import (
+	"fmt"
 	"io"
 	"net/http"
 	"sync"
@@ -19,7 +20,10 @@ func (m CompressionMode) opts() *compressionOptions {
 
 type compressionOptions struct {
 	clientNoContextTakeover bool
+	clientMaxWindowBits     int
+
 	serverNoContextTakeover bool
+	serverMaxWindowBits     int
 }
 
 func (copts *compressionOptions) setHeader(h http.Header) {
@@ -30,6 +34,12 @@ func (copts *compressionOptions) setHeader(h http.Header) {
 	if copts.serverNoContextTakeover {
 		s += "; server_no_context_takeover"
 	}
+	if false && copts.serverMaxWindowBits > 0 {
+		s += fmt.Sprintf("; server_max_window_bits=%v", copts.serverMaxWindowBits)
+	}
+	if false && copts.clientMaxWindowBits > 0 {
+		s += fmt.Sprintf("; client_max_window_bits=%v", copts.clientMaxWindowBits)
+	}
 	h.Set("Sec-WebSocket-Extensions", s)
 }
 
@@ -152,9 +162,8 @@ func (sw *slidingWindow) close() {
 	}
 
 	swPoolMu.Lock()
-	defer swPoolMu.Unlock()
-
 	swPool[cap(sw.buf)].Put(sw.buf)
+	swPoolMu.Unlock()
 	sw.buf = nil
 }
 

conn_test.go

@@ -114,13 +114,21 @@ func TestConn(t *testing.T) {
 
 		for i := 0; i < count; i++ {
 			go func() {
-				errs <- c1.Write(tt.ctx, websocket.MessageBinary, msg)
+				select {
+				case errs <- c1.Write(tt.ctx, websocket.MessageBinary, msg):
+				case <-tt.ctx.Done():
+					return
+				}
 			}()
 		}
 
 		for i := 0; i < count; i++ {
-			err := <-errs
-			assert.Success(t, err)
+			select {
+			case err := <-errs:
+				assert.Success(t, err)
+			case <-tt.ctx.Done():
+				t.Fatal(tt.ctx.Err())
+			}
 		}
 
 		err := c1.Close(websocket.StatusNormalClosure, "")
@@ -171,8 +179,12 @@ func TestConn(t *testing.T) {
 		_, err = n1.Read(nil)
 		assert.Equal(t, "read error", err, io.EOF)
 
-		err = <-errs
-		assert.Success(t, err)
+		select {
+		case err := <-errs:
+			assert.Success(t, err)
+		case <-tt.ctx.Done():
+			t.Fatal(tt.ctx.Err())
+		}
 
 		assert.Equal(t, "read msg", []byte("hello"), b)
 	})
@@ -195,8 +207,12 @@ func TestConn(t *testing.T) {
 		_, err := ioutil.ReadAll(n1)
 		assert.Contains(t, err, `unexpected frame type read (expected MessageBinary): MessageText`)
 
-		err = <-errs
-		assert.Success(t, err)
+		select {
+		case err := <-errs:
+			assert.Success(t, err)
+		case <-tt.ctx.Done():
+			t.Fatal(tt.ctx.Err())
+		}
 	})
 
 	t.Run("wsjson", func(t *testing.T) {
@@ -218,8 +234,12 @@ func TestConn(t *testing.T) {
 		assert.Success(t, err)
 		assert.Equal(t, "read msg", exp, act)
 
-		err = <-werr
-		assert.Success(t, err)
+		select {
+		case err := <-werr:
+			assert.Success(t, err)
+		case <-tt.ctx.Done():
+			t.Fatal(tt.ctx.Err())
+		}
 
 		err = c1.Close(websocket.StatusNormalClosure, "")
 		assert.Success(t, err)
@@ -411,14 +431,22 @@ func BenchmarkConn(b *testing.B) {
 
 			go func() {
 				for range writes {
-					werrs <- c1.Write(bb.ctx, websocket.MessageText, msg)
+					select {
+					case werrs <- c1.Write(bb.ctx, websocket.MessageText, msg):
+					case <-bb.ctx.Done():
+						return
+					}
 				}
 			}()
 			b.SetBytes(int64(len(msg)))
 			b.ReportAllocs()
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
-				writes <- struct{}{}
+				select {
+				case writes <- struct{}{}:
+				case <-bb.ctx.Done():
+					b.Fatal(bb.ctx.Err())
+				}
 
 				typ, r, err := c1.Reader(bb.ctx)
 				if err != nil {
@@ -445,7 +473,11 @@ func BenchmarkConn(b *testing.B) {
 					assert.Equal(b, "msg", msg, readBuf)
 				}
 
-				err = <-werrs
+				select {
+				case err = <-werrs:
+				case <-bb.ctx.Done():
+					b.Fatal(bb.ctx.Err())
+				}
 				if err != nil {
 					b.Fatal(err)
 				}

dial.go

@@ -16,6 +16,7 @@ import (
 	"net/url"
 	"strings"
 	"sync"
+	"time"
 
 	"nhooyr.io/websocket/internal/errd"
 )
@@ -91,6 +92,12 @@ func dial(ctx context.Context, urls string, opts *DialOptions, rand io.Reader) (
 		if err != nil {
 			// We read a bit of the body for easier debugging.
 			r := io.LimitReader(respBody, 1024)
+
+			timer := time.AfterFunc(time.Second*3, func() {
+				respBody.Close()
+			})
+			defer timer.Stop()
+
 			b, _ := ioutil.ReadAll(r)
 			respBody.Close()
 			resp.Body = ioutil.NopCloser(bytes.NewReader(b))
@@ -148,6 +155,7 @@ func handshakeRequest(ctx context.Context, urls string, opts *DialOptions, secWe
 	}
 	if opts.CompressionMode != CompressionDisabled {
 		copts := opts.CompressionMode.opts()
+		copts.clientMaxWindowBits = 8
 		copts.setHeader(req.Header)
 	}
 
@@ -225,15 +233,36 @@ func verifyServerExtensions(h http.Header) (*compressionOptions, error) {
 	}
 
 	copts := &compressionOptions{}
+	copts.clientMaxWindowBits = 8
 	for _, p := range ext.params {
 		switch p {
 		case "client_no_context_takeover":
 			copts.clientNoContextTakeover = true
+			continue
 		case "server_no_context_takeover":
 			copts.serverNoContextTakeover = true
-		default:
-			return nil, fmt.Errorf("unsupported permessage-deflate parameter: %q", p)
+			continue
+		}
+
+		if false && strings.HasPrefix(p, "server_max_window_bits") {
+			bits, ok := parseExtensionParameter(p, 0)
+			if !ok || bits < 8 || bits > 16 {
+				return nil, fmt.Errorf("invalid server_max_window_bits: %q", p)
+			}
+			copts.serverMaxWindowBits = bits
+			continue
 		}
+
+		if false && strings.HasPrefix(p, "client_max_window_bits") {
+			bits, ok := parseExtensionParameter(p, 0)
+			if !ok || bits < 8 || bits > 16 {
+				return nil, fmt.Errorf("invalid client_max_window_bits: %q", p)
+			}
+			copts.clientMaxWindowBits = 8
+			continue
+		}
+
+		return nil, fmt.Errorf("unsupported permessage-deflate parameter: %q", p)
 	}
 
 	return copts, nil

internal/test/assert/assert.go

@@ -2,17 +2,27 @@ package assert
 
 import (
 	"fmt"
+	"reflect"
 	"strings"
 	"testing"
 
-	"nhooyr.io/websocket/internal/test/cmp"
+	"github.com/golang/protobuf/proto"
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 )
 
+// Diff returns a human readable diff between v1 and v2
+func Diff(v1, v2 interface{}) string {
+	return cmp.Diff(v1, v2, cmpopts.EquateErrors(), cmp.Exporter(func(r reflect.Type) bool {
+		return true
+	}), cmp.Comparer(proto.Equal))
+}
+
 // Equal asserts exp == act.
 func Equal(t testing.TB, name string, exp, act interface{}) {
 	t.Helper()
 
-	if diff := cmp.Diff(exp, act); diff != "" {
+	if diff := Diff(exp, act); diff != "" {
 		t.Fatalf("unexpected %v: %v", name, diff)
 	}
 }

internal/test/cmp/cmp.go

@@ -8,7 +8,7 @@ import (
 	"time"
 
 	"nhooyr.io/websocket"
-	"nhooyr.io/websocket/internal/test/cmp"
+	"nhooyr.io/websocket/internal/test/assert"
 	"nhooyr.io/websocket/internal/test/xrand"
 	"nhooyr.io/websocket/internal/xsync"
 )
@@ -76,7 +76,7 @@ func Echo(ctx context.Context, c *websocket.Conn, max int) error {
 	}
 
 	if !bytes.Equal(msg, act) {
-		return fmt.Errorf("unexpected msg read: %v", cmp.Diff(msg, act))
+		return fmt.Errorf("unexpected msg read: %v", assert.Diff(msg, act))
 	}
 
 	return nil