ci/image/push.sh

-#!/usr/bin/env bash
-
-set -euo pipefail
-cd "$(dirname "${0}")"
-cd "$(git rev-parse --show-toplevel)"
-
-docker build -f ./ci/image/Dockerfile -t nhooyr/websocket-ci .
-docker push nhooyr/websocket-ci

ci/lint.sh

-#!/usr/bin/env bash
+#!/bin/sh
+set -eu
+cd -- "$(dirname "$0")/.."
 
-set -euo pipefail
-cd "$(dirname "${0}")"
-cd "$(git rev-parse --show-toplevel)"
+STATICCHECK_VERSION=v0.6.1
+GOVULNCHECK_VERSION=v1.1.4
 
-# shellcheck disable=SC2046
-shellcheck -x $(git ls-files "*.sh")
 go vet ./...
-go run golang.org/x/lint/golint -set_exit_status ./...
+GOOS=js GOARCH=wasm go vet ./...
+
+go install honnef.co/go/tools/cmd/staticcheck@${STATICCHECK_VERSION}
+staticcheck ./...
+GOOS=js GOARCH=wasm staticcheck ./...
+
+govulncheck() {
+	tmpf=$(mktemp)
+	if ! command govulncheck "$@" >"$tmpf" 2>&1; then
+		cat "$tmpf"
+	fi
+}
+go install golang.org/x/vuln/cmd/govulncheck@${GOVULNCHECK_VERSION}
+govulncheck ./...
+GOOS=js GOARCH=wasm govulncheck ./...
+
+(
+  cd ./internal/examples
+  go vet ./...
+  staticcheck ./...
+  govulncheck ./...
+)
+(
+  cd ./internal/thirdparty
+  go vet ./...
+  staticcheck ./...
+  govulncheck ./...
+)

ci/out/.gitignore

+*

ci/run.sh

-#!/usr/bin/env bash
-
-# This script is for local testing. See .github/workflows/ci.yml for CI.
-
-set -euo pipefail
-cd "$(dirname "${0}")"
-cd "$(git rev-parse --show-toplevel)"
-
-echo "--- fmt"
-./ci/fmt.sh
-
-echo "--- lint"
-./ci/lint.sh
-
-echo "--- test"
-./ci/test.sh
-
-echo "--- wasm"
-./ci/wasm.sh

ci/test.sh

-#!/usr/bin/env bash
+#!/bin/sh
+set -eu
+cd -- "$(dirname "$0")/.."
 
-set -euo pipefail
-cd "$(dirname "${0}")"
-cd "$(git rev-parse --show-toplevel)"
-
-argv+=(go test
-  "-parallel=1024"
-  "-coverprofile=ci/out/coverage.prof"
-  "-coverpkg=./..."
+(
+  cd ./internal/examples
+  go test "$@" ./...
+)
+(
+  cd ./internal/thirdparty
+  go test "$@" ./...
 )
 
-if [[ ${CI-} ]]; then
-  argv+=(
-    "-race"
-  )
-fi
+(
+  GOARCH=arm64 go test -c -o ./ci/out/websocket-arm64.test "$@" .
+  if [ "$#" -eq 0 ]; then
+    if [ "${CI-}" ]; then
+      sudo apt-get update
+      sudo apt-get install -y qemu-user-static
+	  ln -s /usr/bin/qemu-aarch64-static /usr/local/bin/qemu-aarch64
+    fi
+    qemu-aarch64 ./ci/out/websocket-arm64.test -test.run=TestMask
+  fi
+)
 
-if [[ $# -gt 0 ]]; then
-  argv+=(
-    "$@"
-  )
-else
-  argv+=(./...)
-fi
 
-mkdir -p ci/out/websocket
-"${argv[@]}"
+go install github.com/agnivade/wasmbrowsertest@8be019f6c6dceae821467b4c589eb195c2b761ce
+go test --race --bench=. --timeout=1h --covermode=atomic --coverprofile=ci/out/coverage.prof --coverpkg=./... "$@" ./...
+sed -i.bak '/stringer\.go/d' ci/out/coverage.prof
+sed -i.bak '/nhooyr.io\/websocket\/internal\/test/d' ci/out/coverage.prof
+sed -i.bak '/examples/d' ci/out/coverage.prof
 
-# Removes coverage of generated files.
-grep -v _string.go < ci/out/coverage.prof > ci/out/coverage2.prof
-mv ci/out/coverage2.prof ci/out/coverage.prof
+# Last line is the total coverage.
+go tool cover -func ci/out/coverage.prof | tail -n1
 
 go tool cover -html=ci/out/coverage.prof -o=ci/out/coverage.html
-if [[ ${CI-} ]]; then
-  bash <(curl -s https://codecov.io/bash) -Z -R . -f ci/out/coverage.prof
-fi

ci/tools.go

-// +build tools
-
-package ci
-
-// See https://github.com/go-modules-by-example/index/blob/master/010_tools/README.md
-import (
-	_ "github.com/agnivade/wasmbrowsertest"
-	_ "go.coder.com/go-tools/cmd/goimports"
-	_ "golang.org/x/lint/golint"
-	_ "golang.org/x/tools/cmd/stringer"
-	_ "gotest.tools/gotestsum"
-	_ "mvdan.cc/sh/cmd/shfmt"
-)

ci/wasm.sh

-#!/usr/bin/env bash
-
-set -euo pipefail
-cd "$(dirname "${0}")"
-cd "$(git rev-parse --show-toplevel)"
-
-GOOS=js GOARCH=wasm go vet ./...
-
-go install golang.org/x/lint/golint
-GOOS=js GOARCH=wasm golint -set_exit_status ./...
-
-wsjstestOut="$(mktemp)"
-go install ./internal/wsjstest
-timeout 30s wsjstest >> "$wsjstestOut" 2>&1 &
-wsjstestPID=$!
-
-# See https://superuser.com/a/900134
-WS_ECHO_SERVER_URL="$( (tail -f -n0 "$wsjstestOut" &) | timeout 10s head -n 1)"
-if [[ -z $WS_ECHO_SERVER_URL ]]; then
-  echo "./internal/wsjstest failed to start in 10s"
-  exit 1
-fi
-
-go install github.com/agnivade/wasmbrowsertest
-export WS_ECHO_SERVER_URL
-GOOS=js GOARCH=wasm go test -exec=wasmbrowsertest ./...
-
-kill "$wsjstestPID"
-if ! wait "$wsjstestPID"; then
-  echo "--- wsjstest exited unsuccessfully"
-  echo "output:"
-  cat "$wsjstestOut"
-  exit 1
-fi

close.go

+//go:build !js
+// +build !js
+
+package websocket
+
+import (
+	"context"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"net"
+	"time"
+
+	"github.com/coder/websocket/internal/errd"
+)
+
+// StatusCode represents a WebSocket status code.
+// https://tools.ietf.org/html/rfc6455#section-7.4
+type StatusCode int
+
+// https://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number
+//
+// These are only the status codes defined by the protocol.
+//
+// You can define custom codes in the 3000-4999 range.
+// The 3000-3999 range is reserved for use by libraries, frameworks and applications.
+// The 4000-4999 range is reserved for private use.
+const (
+	StatusNormalClosure   StatusCode = 1000
+	StatusGoingAway       StatusCode = 1001
+	StatusProtocolError   StatusCode = 1002
+	StatusUnsupportedData StatusCode = 1003
+
+	// 1004 is reserved and so unexported.
+	statusReserved StatusCode = 1004
+
+	// StatusNoStatusRcvd cannot be sent in a close message.
+	// It is reserved for when a close message is received without
+	// a status code.
+	StatusNoStatusRcvd StatusCode = 1005
+
+	// StatusAbnormalClosure is exported for use only with Wasm.
+	// In non Wasm Go, the returned error will indicate whether the
+	// connection was closed abnormally.
+	StatusAbnormalClosure StatusCode = 1006
+
+	StatusInvalidFramePayloadData StatusCode = 1007
+	StatusPolicyViolation         StatusCode = 1008
+	StatusMessageTooBig           StatusCode = 1009
+	StatusMandatoryExtension      StatusCode = 1010
+	StatusInternalError           StatusCode = 1011
+	StatusServiceRestart          StatusCode = 1012
+	StatusTryAgainLater           StatusCode = 1013
+	StatusBadGateway              StatusCode = 1014
+
+	// StatusTLSHandshake is only exported for use with Wasm.
+	// In non Wasm Go, the returned error will indicate whether there was
+	// a TLS handshake failure.
+	StatusTLSHandshake StatusCode = 1015
+)
+
+// CloseError is returned when the connection is closed with a status and reason.
+//
+// Use Go 1.13's errors.As to check for this error.
+// Also see the CloseStatus helper.
+type CloseError struct {
+	Code   StatusCode
+	Reason string
+}
+
+func (ce CloseError) Error() string {
+	return fmt.Sprintf("status = %v and reason = %q", ce.Code, ce.Reason)
+}
+
+// CloseStatus is a convenience wrapper around Go 1.13's errors.As to grab
+// the status code from a CloseError.
+//
+// -1 will be returned if the passed error is nil or not a CloseError.
+func CloseStatus(err error) StatusCode {
+	var ce CloseError
+	if errors.As(err, &ce) {
+		return ce.Code
+	}
+	return -1
+}
+
+// Close performs the WebSocket close handshake with the given status code and reason.
+//
+// It will write a WebSocket close frame with a timeout of 5s and then wait 5s for
+// the peer to send a close frame.
+// All data messages received from the peer during the close handshake will be discarded.
+//
+// The connection can only be closed once. Additional calls to Close
+// are no-ops.
+//
+// The maximum length of reason must be 125 bytes. Avoid sending a dynamic reason.
+//
+// Close will unblock all goroutines interacting with the connection once
+// complete.
+func (c *Conn) Close(code StatusCode, reason string) (err error) {
+	defer errd.Wrap(&err, "failed to close WebSocket")
+
+	if c.casClosing() {
+		err = c.waitGoroutines()
+		if err != nil {
+			return err
+		}
+		return net.ErrClosed
+	}
+	defer func() {
+		if errors.Is(err, net.ErrClosed) {
+			err = nil
+		}
+	}()
+
+	err = c.closeHandshake(code, reason)
+
+	err2 := c.close()
+	if err == nil && err2 != nil {
+		err = err2
+	}
+
+	err2 = c.waitGoroutines()
+	if err == nil && err2 != nil {
+		err = err2
+	}
+
+	return err
+}
+
+// CloseNow closes the WebSocket connection without attempting a close handshake.
+// Use when you do not want the overhead of the close handshake.
+func (c *Conn) CloseNow() (err error) {
+	defer errd.Wrap(&err, "failed to immediately close WebSocket")
+
+	if c.casClosing() {
+		err = c.waitGoroutines()
+		if err != nil {
+			return err
+		}
+		return net.ErrClosed
+	}
+	defer func() {
+		if errors.Is(err, net.ErrClosed) {
+			err = nil
+		}
+	}()
+
+	err = c.close()
+
+	err2 := c.waitGoroutines()
+	if err == nil && err2 != nil {
+		err = err2
+	}
+	return err
+}
+
+func (c *Conn) closeHandshake(code StatusCode, reason string) error {
+	err := c.writeClose(code, reason)
+	if err != nil {
+		return err
+	}
+
+	err = c.waitCloseHandshake()
+	if CloseStatus(err) != code {
+		return err
+	}
+	return nil
+}
+
+func (c *Conn) writeClose(code StatusCode, reason string) error {
+	ce := CloseError{
+		Code:   code,
+		Reason: reason,
+	}
+
+	var p []byte
+	var err error
+	if ce.Code != StatusNoStatusRcvd {
+		p, err = ce.bytes()
+		if err != nil {
+			return err
+		}
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+
+	err = c.writeControl(ctx, opClose, p)
+	// If the connection closed as we're writing we ignore the error as we might
+	// have written the close frame, the peer responded and then someone else read it
+	// and closed the connection.
+	if err != nil && !errors.Is(err, net.ErrClosed) {
+		return err
+	}
+	return nil
+}
+
+func (c *Conn) waitCloseHandshake() error {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+
+	err := c.readMu.lock(ctx)
+	if err != nil {
+		return err
+	}
+	defer c.readMu.unlock()
+
+	for i := int64(0); i < c.msgReader.payloadLength; i++ {
+		_, err := c.br.ReadByte()
+		if err != nil {
+			return err
+		}
+	}
+
+	for {
+		h, err := c.readLoop(ctx)
+		if err != nil {
+			return err
+		}
+
+		for i := int64(0); i < h.payloadLength; i++ {
+			_, err := c.br.ReadByte()
+			if err != nil {
+				return err
+			}
+		}
+	}
+}
+
+func (c *Conn) waitGoroutines() error {
+	t := time.NewTimer(time.Second * 15)
+	defer t.Stop()
+
+	select {
+	case <-c.timeoutLoopDone:
+	case <-t.C:
+		return errors.New("failed to wait for timeoutLoop goroutine to exit")
+	}
+
+	c.closeReadMu.Lock()
+	closeRead := c.closeReadCtx != nil
+	c.closeReadMu.Unlock()
+	if closeRead {
+		select {
+		case <-c.closeReadDone:
+		case <-t.C:
+			return errors.New("failed to wait for close read goroutine to exit")
+		}
+	}
+
+	select {
+	case <-c.closed:
+	case <-t.C:
+		return errors.New("failed to wait for connection to be closed")
+	}
+
+	return nil
+}
+
+func parseClosePayload(p []byte) (CloseError, error) {
+	if len(p) == 0 {
+		return CloseError{
+			Code: StatusNoStatusRcvd,
+		}, nil
+	}
+
+	if len(p) < 2 {
+		return CloseError{}, fmt.Errorf("close payload %q too small, cannot even contain the 2 byte status code", p)
+	}
+
+	ce := CloseError{
+		Code:   StatusCode(binary.BigEndian.Uint16(p)),
+		Reason: string(p[2:]),
+	}
+
+	if !validWireCloseCode(ce.Code) {
+		return CloseError{}, fmt.Errorf("invalid status code %v", ce.Code)
+	}
+
+	return ce, nil
+}
+
+// See http://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number
+// and https://tools.ietf.org/html/rfc6455#section-7.4.1
+func validWireCloseCode(code StatusCode) bool {
+	switch code {
+	case statusReserved, StatusNoStatusRcvd, StatusAbnormalClosure, StatusTLSHandshake:
+		return false
+	}
+
+	if code >= StatusNormalClosure && code <= StatusBadGateway {
+		return true
+	}
+	if code >= 3000 && code <= 4999 {
+		return true
+	}
+
+	return false
+}
+
+func (ce CloseError) bytes() ([]byte, error) {
+	p, err := ce.bytesErr()
+	if err != nil {
+		err = fmt.Errorf("failed to marshal close frame: %w", err)
+		ce = CloseError{
+			Code: StatusInternalError,
+		}
+		p, _ = ce.bytesErr()
+	}
+	return p, err
+}
+
+const maxCloseReason = maxControlPayload - 2
+
+func (ce CloseError) bytesErr() ([]byte, error) {
+	if len(ce.Reason) > maxCloseReason {
+		return nil, fmt.Errorf("reason string max is %v but got %q with length %v", maxCloseReason, ce.Reason, len(ce.Reason))
+	}
+
+	if !validWireCloseCode(ce.Code) {
+		return nil, fmt.Errorf("status code %v cannot be set", ce.Code)
+	}
+
+	buf := make([]byte, 2+len(ce.Reason))
+	binary.BigEndian.PutUint16(buf, uint16(ce.Code))
+	copy(buf[2:], ce.Reason)
+	return buf, nil
+}
+
+func (c *Conn) casClosing() bool {
+	return c.closing.Swap(true)
+}
+
+func (c *Conn) isClosed() bool {
+	select {
+	case <-c.closed:
+		return true
+	default:
+		return false
+	}
+}

close_test.go

+//go:build !js
+// +build !js
+
+package websocket
+
+import (
+	"io"
+	"math"
+	"strings"
+	"testing"
+
+	"github.com/coder/websocket/internal/test/assert"
+)
+
+func TestCloseError(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name    string
+		ce      CloseError
+		success bool
+	}{
+		{
+			name: "normal",
+			ce: CloseError{
+				Code:   StatusNormalClosure,
+				Reason: strings.Repeat("x", maxCloseReason),
+			},
+			success: true,
+		},
+		{
+			name: "bigReason",
+			ce: CloseError{
+				Code:   StatusNormalClosure,
+				Reason: strings.Repeat("x", maxCloseReason+1),
+			},
+			success: false,
+		},
+		{
+			name: "bigCode",
+			ce: CloseError{
+				Code:   math.MaxUint16,
+				Reason: strings.Repeat("x", maxCloseReason),
+			},
+			success: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			_, err := tc.ce.bytesErr()
+			if tc.success {
+				assert.Success(t, err)
+			} else {
+				assert.Error(t, err)
+			}
+		})
+	}
+
+	t.Run("Error", func(t *testing.T) {
+		exp := `status = StatusInternalError and reason = "meow"`
+		act := CloseError{
+			Code:   StatusInternalError,
+			Reason: "meow",
+		}.Error()
+		assert.Equal(t, "CloseError.Error()", exp, act)
+	})
+}
+
+func Test_parseClosePayload(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name    string
+		p       []byte
+		success bool
+		ce      CloseError
+	}{
+		{
+			name:    "normal",
+			p:       append([]byte{0x3, 0xE8}, []byte("hello")...),
+			success: true,
+			ce: CloseError{
+				Code:   StatusNormalClosure,
+				Reason: "hello",
+			},
+		},
+		{
+			name:    "nothing",
+			success: true,
+			ce: CloseError{
+				Code: StatusNoStatusRcvd,
+			},
+		},
+		{
+			name:    "oneByte",
+			p:       []byte{0},
+			success: false,
+		},
+		{
+			name:    "badStatusCode",
+			p:       []byte{0x17, 0x70},
+			success: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			ce, err := parseClosePayload(tc.p)
+			if tc.success {
+				assert.Success(t, err)
+				assert.Equal(t, "close payload", tc.ce, ce)
+			} else {
+				assert.Error(t, err)
+			}
+		})
+	}
+}
+
+func Test_validWireCloseCode(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name  string
+		code  StatusCode
+		valid bool
+	}{
+		{
+			name:  "normal",
+			code:  StatusNormalClosure,
+			valid: true,
+		},
+		{
+			name:  "noStatus",
+			code:  StatusNoStatusRcvd,
+			valid: false,
+		},
+		{
+			name:  "3000",
+			code:  3000,
+			valid: true,
+		},
+		{
+			name:  "4999",
+			code:  4999,
+			valid: true,
+		},
+		{
+			name:  "unknown",
+			code:  5000,
+			valid: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			act := validWireCloseCode(tc.code)
+			assert.Equal(t, "wire close code", tc.valid, act)
+		})
+	}
+}
+
+func TestCloseStatus(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name string
+		in   error
+		exp  StatusCode
+	}{
+		{
+			name: "nil",
+			in:   nil,
+			exp:  -1,
+		},
+		{
+			name: "io.EOF",
+			in:   io.EOF,
+			exp:  -1,
+		},
+		{
+			name: "StatusInternalError",
+			in: CloseError{
+				Code: StatusInternalError,
+			},
+			exp: StatusInternalError,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			act := CloseStatus(tc.in)
+			assert.Equal(t, "close status", tc.exp, act)
+		})
+	}
+}

compress.go

+//go:build !js
+// +build !js
+
+package websocket
+
+import (
+	"compress/flate"
+	"io"
+	"sync"
+)
+
+// CompressionMode represents the modes available to the permessage-deflate extension.
+// See https://tools.ietf.org/html/rfc7692
+//
+// Works in all modern browsers except Safari which does not implement the permessage-deflate extension.
+//
+// Compression is only used if the peer supports the mode selected.
+type CompressionMode int
+
+const (
+	// CompressionDisabled disables the negotiation of the permessage-deflate extension.
+	//
+	// This is the default. Do not enable compression without benchmarking for your particular use case first.
+	CompressionDisabled CompressionMode = iota
+
+	// CompressionContextTakeover compresses each message greater than 128 bytes reusing the 32 KB sliding window from
+	// previous messages. i.e compression context across messages is preserved.
+	//
+	// As most WebSocket protocols are text based and repetitive, this compression mode can be very efficient.
+	//
+	// The memory overhead is a fixed 32 KB sliding window, a fixed 1.2 MB flate.Writer and a sync.Pool of 40 KB flate.Reader's
+	// that are used when reading and then returned.
+	//
+	// Thus, it uses more memory than CompressionNoContextTakeover but compresses more efficiently.
+	//
+	// If the peer does not support CompressionContextTakeover then we will fall back to CompressionNoContextTakeover.
+	CompressionContextTakeover
+
+	// CompressionNoContextTakeover compresses each message greater than 512 bytes. Each message is compressed with
+	// a new 1.2 MB flate.Writer pulled from a sync.Pool. Each message is read with a 40 KB flate.Reader pulled from
+	// a sync.Pool.
+	//
+	// This means less efficient compression as the sliding window from previous messages will not be used but the
+	// memory overhead will be lower as there will be no fixed cost for the flate.Writer nor the 32 KB sliding window.
+	// Especially if the connections are long lived and seldom written to.
+	//
+	// Thus, it uses less memory than CompressionContextTakeover but compresses less efficiently.
+	//
+	// If the peer does not support CompressionNoContextTakeover then we will fall back to CompressionDisabled.
+	CompressionNoContextTakeover
+)
+
+func (m CompressionMode) opts() *compressionOptions {
+	return &compressionOptions{
+		clientNoContextTakeover: m == CompressionNoContextTakeover,
+		serverNoContextTakeover: m == CompressionNoContextTakeover,
+	}
+}
+
+type compressionOptions struct {
+	clientNoContextTakeover bool
+	serverNoContextTakeover bool
+}
+
+func (copts *compressionOptions) String() string {
+	s := "permessage-deflate"
+	if copts.clientNoContextTakeover {
+		s += "; client_no_context_takeover"
+	}
+	if copts.serverNoContextTakeover {
+		s += "; server_no_context_takeover"
+	}
+	return s
+}
+
+// These bytes are required to get flate.Reader to return.
+// They are removed when sending to avoid the overhead as
+// WebSocket framing tell's when the message has ended but then
+// we need to add them back otherwise flate.Reader keeps
+// trying to read more bytes.
+const deflateMessageTail = "\x00\x00\xff\xff"
+
+type trimLastFourBytesWriter struct {
+	w    io.Writer
+	tail []byte
+}
+
+func (tw *trimLastFourBytesWriter) reset() {
+	if tw != nil && tw.tail != nil {
+		tw.tail = tw.tail[:0]
+	}
+}
+
+func (tw *trimLastFourBytesWriter) Write(p []byte) (int, error) {
+	if tw.tail == nil {
+		tw.tail = make([]byte, 0, 4)
+	}
+
+	extra := len(tw.tail) + len(p) - 4
+
+	if extra <= 0 {
+		tw.tail = append(tw.tail, p...)
+		return len(p), nil
+	}
+
+	// Now we need to write as many extra bytes as we can from the previous tail.
+	if extra > len(tw.tail) {
+		extra = len(tw.tail)
+	}
+	if extra > 0 {
+		_, err := tw.w.Write(tw.tail[:extra])
+		if err != nil {
+			return 0, err
+		}
+
+		// Shift remaining bytes in tail over.
+		n := copy(tw.tail, tw.tail[extra:])
+		tw.tail = tw.tail[:n]
+	}
+
+	// If p is less than or equal to 4 bytes,
+	// all of it is is part of the tail.
+	if len(p) <= 4 {
+		tw.tail = append(tw.tail, p...)
+		return len(p), nil
+	}
+
+	// Otherwise, only the last 4 bytes are.
+	tw.tail = append(tw.tail, p[len(p)-4:]...)
+
+	p = p[:len(p)-4]
+	n, err := tw.w.Write(p)
+	return n + 4, err
+}
+
+var flateReaderPool sync.Pool
+
+func getFlateReader(r io.Reader, dict []byte) io.Reader {
+	fr, ok := flateReaderPool.Get().(io.Reader)
+	if !ok {
+		return flate.NewReaderDict(r, dict)
+	}
+	fr.(flate.Resetter).Reset(r, dict)
+	return fr
+}
+
+func putFlateReader(fr io.Reader) {
+	flateReaderPool.Put(fr)
+}
+
+var flateWriterPool sync.Pool
+
+func getFlateWriter(w io.Writer) *flate.Writer {
+	fw, ok := flateWriterPool.Get().(*flate.Writer)
+	if !ok {
+		fw, _ = flate.NewWriter(w, flate.BestSpeed)
+		return fw
+	}
+	fw.Reset(w)
+	return fw
+}
+
+func putFlateWriter(w *flate.Writer) {
+	flateWriterPool.Put(w)
+}
+
+type slidingWindow struct {
+	buf []byte
+}
+
+var swPoolMu sync.RWMutex
+var swPool = map[int]*sync.Pool{}
+
+func slidingWindowPool(n int) *sync.Pool {
+	swPoolMu.RLock()
+	p, ok := swPool[n]
+	swPoolMu.RUnlock()
+	if ok {
+		return p
+	}
+
+	p = &sync.Pool{}
+
+	swPoolMu.Lock()
+	swPool[n] = p
+	swPoolMu.Unlock()
+
+	return p
+}
+
+func (sw *slidingWindow) init(n int) {
+	if sw.buf != nil {
+		return
+	}
+
+	if n == 0 {
+		n = 32768
+	}
+
+	p := slidingWindowPool(n)
+	sw2, ok := p.Get().(*slidingWindow)
+	if ok {
+		*sw = *sw2
+	} else {
+		sw.buf = make([]byte, 0, n)
+	}
+}
+
+func (sw *slidingWindow) close() {
+	sw.buf = sw.buf[:0]
+	swPoolMu.Lock()
+	swPool[cap(sw.buf)].Put(sw)
+	swPoolMu.Unlock()
+}
+
+func (sw *slidingWindow) write(p []byte) {
+	if len(p) >= cap(sw.buf) {
+		sw.buf = sw.buf[:cap(sw.buf)]
+		p = p[len(p)-cap(sw.buf):]
+		copy(sw.buf, p)
+		return
+	}
+
+	left := cap(sw.buf) - len(sw.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(sw.buf, sw.buf[spaceNeeded:])
+		sw.buf = sw.buf[:len(sw.buf)-spaceNeeded]
+	}
+
+	sw.buf = append(sw.buf, p...)
+}

compress_test.go

+//go:build !js
+// +build !js
+
+package websocket
+
+import (
+	"bytes"
+	"compress/flate"
+	"io"
+	"strings"
+	"testing"
+
+	"github.com/coder/websocket/internal/test/assert"
+	"github.com/coder/websocket/internal/test/xrand"
+)
+
+func Test_slidingWindow(t *testing.T) {
+	t.Parallel()
+
+	const testCount = 99
+	const maxWindow = 99999
+	for i := 0; i < testCount; i++ {
+		t.Run("", func(t *testing.T) {
+			t.Parallel()
+
+			input := xrand.String(maxWindow)
+			windowLength := xrand.Int(maxWindow)
+			var sw slidingWindow
+			sw.init(windowLength)
+			sw.write([]byte(input))
+
+			assert.Equal(t, "window length", windowLength, cap(sw.buf))
+			if !strings.HasSuffix(input, string(sw.buf)) {
+				t.Fatalf("r.buf is not a suffix of input: %q and %q", input, sw.buf)
+			}
+		})
+	}
+}
+
+func BenchmarkFlateWriter(b *testing.B) {
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		w, _ := flate.NewWriter(io.Discard, flate.BestSpeed)
+		// We have to write a byte to get the writer to allocate to its full extent.
+		w.Write([]byte{'a'})
+		w.Flush()
+	}
+}
+
+func BenchmarkFlateReader(b *testing.B) {
+	b.ReportAllocs()
+
+	var buf bytes.Buffer
+	w, _ := flate.NewWriter(&buf, flate.BestSpeed)
+	w.Write([]byte{'a'})
+	w.Flush()
+
+	for i := 0; i < b.N; i++ {
+		r := flate.NewReader(bytes.NewReader(buf.Bytes()))
+		io.ReadAll(r)
+	}
+}

conn_export_test.go

-// +build !js
-
-package websocket
-
-import (
-	"bufio"
-	"context"
-	"fmt"
-)
-
-type (
-	Addr   = websocketAddr
-	OpCode int
-)
-
-const (
-	OpClose        = OpCode(opClose)
-	OpBinary       = OpCode(opBinary)
-	OpText         = OpCode(opText)
-	OpPing         = OpCode(opPing)
-	OpPong         = OpCode(opPong)
-	OpContinuation = OpCode(opContinuation)
-)
-
-func (c *Conn) ReadFrame(ctx context.Context) (OpCode, []byte, error) {
-	h, err := c.readFrameHeader(ctx)
-	if err != nil {
-		return 0, nil, err
-	}
-	b := make([]byte, h.payloadLength)
-	_, err = c.readFramePayload(ctx, b)
-	if err != nil {
-		return 0, nil, err
-	}
-	if h.masked {
-		fastXOR(h.maskKey, 0, b)
-	}
-	return OpCode(h.opcode), b, nil
-}
-
-func (c *Conn) WriteFrame(ctx context.Context, fin bool, opc OpCode, p []byte) (int, error) {
-	return c.writeFrame(ctx, fin, opcode(opc), p)
-}
-
-// header represents a WebSocket frame header.
-// See https://tools.ietf.org/html/rfc6455#section-5.2
-type Header struct {
-	Fin    bool
-	Rsv1   bool
-	Rsv2   bool
-	Rsv3   bool
-	OpCode OpCode
-
-	PayloadLength int64
-}
-
-func (c *Conn) WriteHeader(ctx context.Context, h Header) error {
-	headerBytes := writeHeader(c.writeHeaderBuf, header{
-		fin:           h.Fin,
-		rsv1:          h.Rsv1,
-		rsv2:          h.Rsv2,
-		rsv3:          h.Rsv3,
-		opcode:        opcode(h.OpCode),
-		payloadLength: h.PayloadLength,
-		masked:        c.client,
-	})
-	_, err := c.bw.Write(headerBytes)
-	if err != nil {
-		return fmt.Errorf("failed to write header: %w", err)
-	}
-	if h.Fin {
-		err = c.Flush()
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
-func (c *Conn) PingWithPayload(ctx context.Context, p string) error {
-	return c.ping(ctx, p)
-}
-
-func (c *Conn) WriteHalfFrame(ctx context.Context) (int, error) {
-	return c.realWriteFrame(ctx, header{
-		fin:           true,
-		opcode:        opBinary,
-		payloadLength: 10,
-	}, make([]byte, 5))
-}
-
-func (c *Conn) CloseUnderlyingConn() {
-	c.closer.Close()
-}
-
-func (c *Conn) Flush() error {
-	return c.bw.Flush()
-}
-
-func (c CloseError) Bytes() ([]byte, error) {
-	return c.bytes()
-}
-
-func (c *Conn) BW() *bufio.Writer {
-	return c.bw
-}
-
-func (c *Conn) WriteClose(ctx context.Context, code StatusCode, reason string) ([]byte, error) {
-	b, err := CloseError{
-		Code:   code,
-		Reason: reason,
-	}.Bytes()
-	if err != nil {
-		return nil, err
-	}
-	_, err = c.WriteFrame(ctx, true, OpClose, b)
-	if err != nil {
-		return nil, err
-	}
-	return b, nil
-}
-
-func ParseClosePayload(p []byte) (CloseError, error) {
-	return parseClosePayload(p)
-}

dial.go

+//go:build !js
+// +build !js
+
+package websocket
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"crypto/rand"
+	"encoding/base64"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/coder/websocket/internal/errd"
+)
+
+// DialOptions represents Dial's options.
+type DialOptions struct {
+	// HTTPClient is used for the connection.
+	// Its Transport must return writable bodies for WebSocket handshakes.
+	// http.Transport does beginning with Go 1.12.
+	HTTPClient *http.Client
+
+	// HTTPHeader specifies the HTTP headers included in the handshake request.
+	HTTPHeader http.Header
+
+	// Host optionally overrides the Host HTTP header to send. If empty, the value
+	// of URL.Host will be used.
+	Host string
+
+	// Subprotocols lists the WebSocket subprotocols to negotiate with the server.
+	Subprotocols []string
+
+	// CompressionMode controls the compression mode.
+	// Defaults to CompressionDisabled.
+	//
+	// See docs on CompressionMode for details.
+	CompressionMode CompressionMode
+
+	// CompressionThreshold controls the minimum size of a message before compression is applied.
+	//
+	// Defaults to 512 bytes for CompressionNoContextTakeover and 128 bytes
+	// for CompressionContextTakeover.
+	CompressionThreshold int
+
+	// OnPingReceived is an optional callback invoked synchronously when a ping frame is received.
+	//
+	// The payload contains the application data of the ping frame.
+	// If the callback returns false, the subsequent pong frame will not be sent.
+	// To avoid blocking, any expensive processing should be performed asynchronously using a goroutine.
+	OnPingReceived func(ctx context.Context, payload []byte) bool
+
+	// OnPongReceived is an optional callback invoked synchronously when a pong frame is received.
+	//
+	// The payload contains the application data of the pong frame.
+	// To avoid blocking, any expensive processing should be performed asynchronously using a goroutine.
+	//
+	// Unlike OnPingReceived, this callback does not return a value because a pong frame
+	// is a response to a ping and does not trigger any further frame transmission.
+	OnPongReceived func(ctx context.Context, payload []byte)
+}
+
+func (opts *DialOptions) cloneWithDefaults(ctx context.Context) (context.Context, context.CancelFunc, *DialOptions) {
+	var cancel context.CancelFunc
+
+	var o DialOptions
+	if opts != nil {
+		o = *opts
+	}
+	if o.HTTPClient == nil {
+		o.HTTPClient = http.DefaultClient
+	}
+	if o.HTTPClient.Timeout > 0 {
+		ctx, cancel = context.WithTimeout(ctx, o.HTTPClient.Timeout)
+
+		newClient := *o.HTTPClient
+		newClient.Timeout = 0
+		o.HTTPClient = &newClient
+	}
+	if o.HTTPHeader == nil {
+		o.HTTPHeader = http.Header{}
+	}
+	newClient := *o.HTTPClient
+	oldCheckRedirect := o.HTTPClient.CheckRedirect
+	newClient.CheckRedirect = func(req *http.Request, via []*http.Request) error {
+		switch req.URL.Scheme {
+		case "ws":
+			req.URL.Scheme = "http"
+		case "wss":
+			req.URL.Scheme = "https"
+		}
+		if oldCheckRedirect != nil {
+			return oldCheckRedirect(req, via)
+		}
+		return nil
+	}
+	o.HTTPClient = &newClient
+
+	return ctx, cancel, &o
+}
+
+// Dial performs a WebSocket handshake on url.
+//
+// The response is the WebSocket handshake response from the server.
+// You never need to close resp.Body yourself.
+//
+// If an error occurs, the returned response may be non nil.
+// However, you can only read the first 1024 bytes of the body.
+//
+// This function requires at least Go 1.12 as it uses a new feature
+// in net/http to perform WebSocket handshakes.
+// See docs on the HTTPClient option and https://github.com/golang/go/issues/26937#issuecomment-415855861
+//
+// URLs with http/https schemes will work and are interpreted as ws/wss.
+func Dial(ctx context.Context, u string, opts *DialOptions) (*Conn, *http.Response, error) {
+	return dial(ctx, u, opts, nil)
+}
+
+func dial(ctx context.Context, urls string, opts *DialOptions, rand io.Reader) (_ *Conn, _ *http.Response, err error) {
+	defer errd.Wrap(&err, "failed to WebSocket dial")
+
+	var cancel context.CancelFunc
+	ctx, cancel, opts = opts.cloneWithDefaults(ctx)
+	if cancel != nil {
+		defer cancel()
+	}
+
+	secWebSocketKey, err := secWebSocketKey(rand)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to generate Sec-WebSocket-Key: %w", err)
+	}
+
+	var copts *compressionOptions
+	if opts.CompressionMode != CompressionDisabled {
+		copts = opts.CompressionMode.opts()
+	}
+
+	resp, err := handshakeRequest(ctx, urls, opts, copts, secWebSocketKey)
+	if err != nil {
+		return nil, resp, err
+	}
+	respBody := resp.Body
+	resp.Body = nil
+	defer func() {
+		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, _ := io.ReadAll(r)
+			respBody.Close()
+			resp.Body = io.NopCloser(bytes.NewReader(b))
+		}
+	}()
+
+	copts, err = verifyServerResponse(opts, copts, secWebSocketKey, resp)
+	if err != nil {
+		return nil, resp, err
+	}
+
+	rwc, ok := respBody.(io.ReadWriteCloser)
+	if !ok {
+		return nil, resp, fmt.Errorf("response body is not a io.ReadWriteCloser: %T", respBody)
+	}
+
+	return newConn(connConfig{
+		subprotocol:    resp.Header.Get("Sec-WebSocket-Protocol"),
+		rwc:            rwc,
+		client:         true,
+		copts:          copts,
+		flateThreshold: opts.CompressionThreshold,
+		onPingReceived: opts.OnPingReceived,
+		onPongReceived: opts.OnPongReceived,
+		br:             getBufioReader(rwc),
+		bw:             getBufioWriter(rwc),
+	}), resp, nil
+}
+
+func handshakeRequest(ctx context.Context, urls string, opts *DialOptions, copts *compressionOptions, secWebSocketKey string) (*http.Response, error) {
+	u, err := url.Parse(urls)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse url: %w", err)
+	}
+
+	switch u.Scheme {
+	case "ws":
+		u.Scheme = "http"
+	case "wss":
+		u.Scheme = "https"
+	case "http", "https":
+	default:
+		return nil, fmt.Errorf("unexpected url scheme: %q", u.Scheme)
+	}
+
+	req, err := http.NewRequestWithContext(ctx, "GET", u.String(), nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create new http request: %w", err)
+	}
+	if len(opts.Host) > 0 {
+		req.Host = opts.Host
+	}
+	req.Header = opts.HTTPHeader.Clone()
+	req.Header.Set("Connection", "Upgrade")
+	req.Header.Set("Upgrade", "websocket")
+	req.Header.Set("Sec-WebSocket-Version", "13")
+	req.Header.Set("Sec-WebSocket-Key", secWebSocketKey)
+	if len(opts.Subprotocols) > 0 {
+		req.Header.Set("Sec-WebSocket-Protocol", strings.Join(opts.Subprotocols, ","))
+	}
+	if copts != nil {
+		req.Header.Set("Sec-WebSocket-Extensions", copts.String())
+	}
+
+	resp, err := opts.HTTPClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to send handshake request: %w", err)
+	}
+	return resp, nil
+}
+
+func secWebSocketKey(rr io.Reader) (string, error) {
+	if rr == nil {
+		rr = rand.Reader
+	}
+	b := make([]byte, 16)
+	_, err := io.ReadFull(rr, b)
+	if err != nil {
+		return "", fmt.Errorf("failed to read random data from rand.Reader: %w", err)
+	}
+	return base64.StdEncoding.EncodeToString(b), nil
+}
+
+func verifyServerResponse(opts *DialOptions, copts *compressionOptions, secWebSocketKey string, resp *http.Response) (*compressionOptions, error) {
+	if resp.StatusCode != http.StatusSwitchingProtocols {
+		return nil, fmt.Errorf("expected handshake response status code %v but got %v", http.StatusSwitchingProtocols, resp.StatusCode)
+	}
+
+	if !headerContainsTokenIgnoreCase(resp.Header, "Connection", "Upgrade") {
+		return nil, fmt.Errorf("WebSocket protocol violation: Connection header %q does not contain Upgrade", resp.Header.Get("Connection"))
+	}
+
+	if !headerContainsTokenIgnoreCase(resp.Header, "Upgrade", "WebSocket") {
+		return nil, fmt.Errorf("WebSocket protocol violation: Upgrade header %q does not contain websocket", resp.Header.Get("Upgrade"))
+	}
+
+	if resp.Header.Get("Sec-WebSocket-Accept") != secWebSocketAccept(secWebSocketKey) {
+		return nil, fmt.Errorf("WebSocket protocol violation: invalid Sec-WebSocket-Accept %q, key %q",
+			resp.Header.Get("Sec-WebSocket-Accept"),
+			secWebSocketKey,
+		)
+	}
+
+	err := verifySubprotocol(opts.Subprotocols, resp)
+	if err != nil {
+		return nil, err
+	}
+
+	return verifyServerExtensions(copts, resp.Header)
+}
+
+func verifySubprotocol(subprotos []string, resp *http.Response) error {
+	proto := resp.Header.Get("Sec-WebSocket-Protocol")
+	if proto == "" {
+		return nil
+	}
+
+	for _, sp2 := range subprotos {
+		if strings.EqualFold(sp2, proto) {
+			return nil
+		}
+	}
+
+	return fmt.Errorf("WebSocket protocol violation: unexpected Sec-WebSocket-Protocol from server: %q", proto)
+}
+
+func verifyServerExtensions(copts *compressionOptions, h http.Header) (*compressionOptions, error) {
+	exts := websocketExtensions(h)
+	if len(exts) == 0 {
+		return nil, nil
+	}
+
+	ext := exts[0]
+	if ext.name != "permessage-deflate" || len(exts) > 1 || copts == nil {
+		return nil, fmt.Errorf("WebSocket protcol violation: unsupported extensions from server: %+v", exts[1:])
+	}
+
+	_copts := *copts
+	copts = &_copts
+
+	for _, p := range ext.params {
+		switch p {
+		case "client_no_context_takeover":
+			copts.clientNoContextTakeover = true
+			continue
+		case "server_no_context_takeover":
+			copts.serverNoContextTakeover = true
+			continue
+		}
+		if strings.HasPrefix(p, "server_max_window_bits=") {
+			// We can't adjust the deflate window, but decoding with a larger window is acceptable.
+			continue
+		}
+
+		return nil, fmt.Errorf("unsupported permessage-deflate parameter: %q", p)
+	}
+
+	return copts, nil
+}
+
+var bufioReaderPool sync.Pool
+
+func getBufioReader(r io.Reader) *bufio.Reader {
+	br, ok := bufioReaderPool.Get().(*bufio.Reader)
+	if !ok {
+		return bufio.NewReader(r)
+	}
+	br.Reset(r)
+	return br
+}
+
+func putBufioReader(br *bufio.Reader) {
+	bufioReaderPool.Put(br)
+}
+
+var bufioWriterPool sync.Pool
+
+func getBufioWriter(w io.Writer) *bufio.Writer {
+	bw, ok := bufioWriterPool.Get().(*bufio.Writer)
+	if !ok {
+		return bufio.NewWriter(w)
+	}
+	bw.Reset(w)
+	return bw
+}
+
+func putBufioWriter(bw *bufio.Writer) {
+	bufioWriterPool.Put(bw)
+}

doc.go

+//go:build !js
 // +build !js
 
-// Package websocket is a minimal and idiomatic implementation of the WebSocket protocol.
+// Package websocket implements the RFC 6455 WebSocket protocol.
 //
 // https://tools.ietf.org/html/rfc6455
 //
-// Conn, Dial, and Accept are the main entrypoints into this package. Use Dial to dial
-// a WebSocket server, Accept to accept a WebSocket client dial and then Conn to interact
-// with the resulting WebSocket connections.
+// Use Dial to dial a WebSocket server.
 //
-// The examples are the best way to understand how to correctly use the library.
+// Use Accept to accept a WebSocket client.
+//
+// Conn represents the resulting WebSocket connection.
 //
-// The wsjson and wspb subpackages contain helpers for JSON and ProtoBuf messages.
+// The examples are the best way to understand how to correctly use the library.
 //
-// See https://nhooyr.io/websocket for more overview docs and a
-// comparison with existing implementations.
+// The wsjson subpackage contain helpers for JSON and protobuf messages.
 //
-// Use the errors.As function new in Go 1.13 to check for websocket.CloseError.
-// See the CloseError example.
+// More documentation at https://github.com/coder/websocket.
 //
-// Wasm
+// # Wasm
 //
-// The client side fully supports compiling to Wasm.
+// The client side supports compiling to Wasm.
 // It wraps the WebSocket browser API.
 //
 // See https://developer.mozilla.org/en-US/docs/Web/API/WebSocket
 //
-// Thus the unsupported features (not compiled in) for Wasm are:
-//
-//  - Accept and AcceptOptions
-//  - Conn.Ping
-//  - HTTPClient and HTTPHeader fields in DialOptions
-//
-// The *http.Response returned by Dial will always either be nil or &http.Response{} as
-// we do not have access to the handshake response in the browser.
-//
-// The Writer method on the Conn buffers everything in memory and then sends it as a message
-// when the writer is closed.
-//
-// The Reader method also reads the entire response and then returns a reader that
-// reads from the byte slice.
-//
-// SetReadLimit cannot actually limit the number of bytes read from the connection so instead
-// when a message beyond the limit is fully read, it throws an error.
-//
-// Writes are also always async so the passed context is no-op.
-//
-// Everything else is fully supported. This includes the wsjson and wspb helper packages.
+// Some important caveats to be aware of:
 //
-// Once https://github.com/gopherjs/gopherjs/issues/929 is closed, GopherJS should be supported
-// as well.
-package websocket // import "nhooyr.io/websocket"
+//   - Accept always errors out
+//   - Conn.Ping is no-op
+//   - Conn.CloseNow is Close(StatusGoingAway, "")
+//   - HTTPClient, HTTPHeader and CompressionMode in DialOptions are no-op
+//   - *http.Response from Dial is &http.Response{} with a 101 status code on success
+package websocket // import "github.com/coder/websocket"

example_echo_test.go

-// +build !js
-
-package websocket_test
-
-import (
-	"context"
-	"fmt"
-	"io"
-	"log"
-	"net"
-	"net/http"
-	"time"
-
-	"golang.org/x/time/rate"
-
-	"nhooyr.io/websocket"
-	"nhooyr.io/websocket/wsjson"
-)
-
-// This example starts a WebSocket echo server,
-// dials the server and then sends 5 different messages
-// and prints out the server's responses.
-func Example_echo() {
-	// First we listen on port 0 which means the OS will
-	// assign us a random free port. This is the listener
-	// the server will serve on and the client will connect to.
-	l, err := net.Listen("tcp", "localhost:0")
-	if err != nil {
-		log.Fatalf("failed to listen: %v", err)
-	}
-	defer l.Close()
-
-	s := &http.Server{
-		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			err := echoServer(w, r)
-			if err != nil {
-				log.Printf("echo server: %v", err)
-			}
-		}),
-		ReadTimeout:  time.Second * 15,
-		WriteTimeout: time.Second * 15,
-	}
-	defer s.Close()
-
-	// This starts the echo server on the listener.
-	go func() {
-		err := s.Serve(l)
-		if err != http.ErrServerClosed {
-			log.Fatalf("failed to listen and serve: %v", err)
-		}
-	}()
-
-	// Now we dial the server, send the messages and echo the responses.
-	err = client("ws://" + l.Addr().String())
-	if err != nil {
-		log.Fatalf("client failed: %v", err)
-	}
-	// Output:
-	// received: map[i:0]
-	// received: map[i:1]
-	// received: map[i:2]
-	// received: map[i:3]
-	// received: map[i:4]
-}
-
-// echoServer is the WebSocket echo server implementation.
-// It ensures the client speaks the echo subprotocol and
-// only allows one message every 100ms with a 10 message burst.
-func echoServer(w http.ResponseWriter, r *http.Request) error {
-	log.Printf("serving %v", r.RemoteAddr)
-
-	c, err := websocket.Accept(w, r, &websocket.AcceptOptions{
-		Subprotocols: []string{"echo"},
-	})
-	if err != nil {
-		return err
-	}
-	defer c.Close(websocket.StatusInternalError, "the sky is falling")
-
-	if c.Subprotocol() != "echo" {
-		c.Close(websocket.StatusPolicyViolation, "client must speak the echo subprotocol")
-		return fmt.Errorf("client does not speak echo sub protocol")
-	}
-
-	l := rate.NewLimiter(rate.Every(time.Millisecond*100), 10)
-	for {
-		err = echo(r.Context(), c, l)
-		if err != nil {
-			return fmt.Errorf("failed to echo with %v: %w", r.RemoteAddr, err)
-		}
-	}
-}
-
-// echo reads from the websocket connection and then writes
-// the received message back to it.
-// The entire function has 10s to complete.
-func echo(ctx context.Context, c *websocket.Conn, l *rate.Limiter) error {
-	ctx, cancel := context.WithTimeout(ctx, time.Second*10)
-	defer cancel()
-
-	err := l.Wait(ctx)
-	if err != nil {
-		return err
-	}
-
-	typ, r, err := c.Reader(ctx)
-	if err != nil {
-		return err
-	}
-
-	w, err := c.Writer(ctx, typ)
-	if err != nil {
-		return err
-	}
-
-	_, err = io.Copy(w, r)
-	if err != nil {
-		return fmt.Errorf("failed to io.Copy: %w", err)
-	}
-
-	err = w.Close()
-	return err
-}
-
-// client dials the WebSocket echo server at the given url.
-// It then sends it 5 different messages and echo's the server's
-// response to each.
-func client(url string) error {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
-	defer cancel()
-
-	c, _, err := websocket.Dial(ctx, url, &websocket.DialOptions{
-		Subprotocols: []string{"echo"},
-	})
-	if err != nil {
-		return err
-	}
-	defer c.Close(websocket.StatusInternalError, "the sky is falling")
-
-	for i := 0; i < 5; i++ {
-		err = wsjson.Write(ctx, c, map[string]int{
-			"i": i,
-		})
-		if err != nil {
-			return err
-		}
-
-		v := map[string]int{}
-		err = wsjson.Read(ctx, c, &v)
-		if err != nil {
-			return err
-		}
-
-		fmt.Printf("received: %v\n", v)
-	}
-
-	c.Close(websocket.StatusNormalClosure, "")
-	return nil
-}

export_test.go

+//go:build !js
+// +build !js
+
+package websocket
+
+import (
+	"net"
+
+	"github.com/coder/websocket/internal/util"
+)
+
+func (c *Conn) RecordBytesWritten() *int {
+	var bytesWritten int
+	c.bw.Reset(util.WriterFunc(func(p []byte) (int, error) {
+		bytesWritten += len(p)
+		return c.rwc.Write(p)
+	}))
+	return &bytesWritten
+}
+
+func (c *Conn) RecordBytesRead() *int {
+	var bytesRead int
+	c.br.Reset(util.ReaderFunc(func(p []byte) (int, error) {
+		n, err := c.rwc.Read(p)
+		bytesRead += n
+		return n, err
+	}))
+	return &bytesRead
+}
+
+var ErrClosed = net.ErrClosed
+
+var ExportedDial = dial
+var SecWebSocketAccept = secWebSocketAccept
+var SecWebSocketKey = secWebSocketKey
+var VerifyServerResponse = verifyServerResponse
+
+var CompressionModeOpts = CompressionMode.opts