snapshot of https://github.com/centrifugal/centrifugo.git (master, 2026-09-16; blobs > 131072B omitted)

internal/websocket/AUTHORS

@@ -0,0 +1,8 @@
+# This is the official list of Gorilla WebSocket authors for copyright
+# purposes.
+#
+# Please keep the list sorted.
+
+Gary Burd <gary@beagledreams.com>
+Google LLC (https://opensource.google.com/)
+Joachim Bauch <mail@joachim-bauch.de>

internal/websocket/LICENSE

@@ -0,0 +1,22 @@
+Copyright (c) 2013 The Gorilla WebSocket Authors. All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+  Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+  Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

internal/websocket/README.md

@@ -0,0 +1,11 @@
+Fork of Gorilla Websocket adapted for Centrifuge needs. All copyrights were left – huge respect to the original authors.
+
+What was changed:
+
+* no custom Proxy in Client
+* no concurrent use best-effort detection
+* selected subprotocol is not kept internally, returned during Upgrade or Dial
+* lint fixes
+* no possibility to set custom CloseHandler
+* upgrade optimizations (9 -> 3 allocations, see BenchmarkUpgrade)
+* Support WebSocket over HTTP/2 ([RFC 8441](https://www.rfc-editor.org/rfc/rfc8441.html), server side only)

internal/websocket/client.go

@@ -0,0 +1,385 @@
+// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package websocket
+
+import (
+	"bytes"
+	"context"
+	"crypto/tls"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptrace"
+	"net/url"
+	"strings"
+	"time"
+)
+
+// ErrBadHandshake is returned when the server response to opening handshake is
+// invalid.
+var ErrBadHandshake = errors.New("websocket: bad handshake")
+
+var errInvalidCompression = errors.New("websocket: invalid compression negotiation")
+
+// A Dialer contains options for connecting to WebSocket server.
+//
+// It is safe to call Dialer's methods concurrently.
+type Dialer struct {
+	// NetDial specifies the dial function for creating TCP connections. If
+	// NetDial is nil, net.Dial is used.
+	NetDial func(network, addr string) (net.Conn, error)
+
+	// NetDialContext specifies the dial function for creating TCP connections. If
+	// NetDialContext is nil, NetDial is used.
+	NetDialContext func(ctx context.Context, network, addr string) (net.Conn, error)
+
+	// NetDialTLSContext specifies the dial function for creating TLS/TCP connections. If
+	// NetDialTLSContext is nil, NetDialContext is used.
+	// If NetDialTLSContext is set, Dial assumes the TLS handshake is done there and
+	// TLSClientConfig is ignored.
+	NetDialTLSContext func(ctx context.Context, network, addr string) (net.Conn, error)
+
+	// TLSClientConfig specifies the TLS configuration to use with tls.Client.
+	// If nil, the default configuration is used.
+	// If either NetDialTLS or NetDialTLSContext are set, Dial assumes the TLS handshake
+	// is done there and TLSClientConfig is ignored.
+	TLSClientConfig *tls.Config
+
+	// HandshakeTimeout specifies the duration for the handshake to complete.
+	HandshakeTimeout time.Duration
+
+	// ReadBufferSize and WriteBufferSize specify I/O buffer sizes in bytes. If a buffer
+	// size is zero, then a useful default size is used. The I/O buffer sizes
+	// do not limit the size of the messages that can be sent or received.
+	ReadBufferSize, WriteBufferSize int
+
+	// WriteBufferPool is a pool of buffers for write operations. If the value
+	// is not set, then write buffers are allocated to the connection for the
+	// lifetime of the connection.
+	//
+	// A pool is most useful when the application has a modest volume of writes
+	// across a large number of connections.
+	//
+	// Applications should use a single pool for each unique value of
+	// WriteBufferSize.
+	WriteBufferPool BufferPool
+
+	// Subprotocols specifies the client's requested subprotocols.
+	Subprotocols []string
+
+	// EnableCompression specifies if the client should attempt to negotiate
+	// per message compression (RFC 7692). Setting this value to true does not
+	// guarantee that compression will be supported. Currently only "no context
+	// takeover" modes are supported.
+	EnableCompression bool
+
+	// Jar specifies the cookie jar.
+	// If Jar is nil, cookies are not sent in requests and ignored
+	// in responses.
+	Jar http.CookieJar
+}
+
+// Dial creates a new client connection by calling DialContext with a background context.
+func (d *Dialer) Dial(urlStr string, requestHeader http.Header) (*Conn, *http.Response, string, error) {
+	return d.DialContext(context.Background(), urlStr, requestHeader)
+}
+
+var errMalformedURL = errors.New("malformed ws or wss URL")
+
+func hostPortNoPort(u *url.URL) (hostPort, hostNoPort string) {
+	hostPort = u.Host
+	hostNoPort = u.Host
+	if i := strings.LastIndex(u.Host, ":"); i > strings.LastIndex(u.Host, "]") {
+		hostNoPort = hostNoPort[:i]
+	} else {
+		switch u.Scheme {
+		case "wss":
+			hostPort += ":443"
+		case "https":
+			hostPort += ":443"
+		default:
+			hostPort += ":80"
+		}
+	}
+	return hostPort, hostNoPort
+}
+
+// DialContext creates a new client connection. Use requestHeader to specify the
+// origin (Origin), subprotocols (Sec-WebSocket-Protocol) and cookies (Cookie).
+// Use the response.Header to get the selected subprotocol
+// (Sec-WebSocket-Protocol) and cookies (Set-Cookie).
+//
+// The context will be used in the request and in the Dialer.
+//
+// If the WebSocket handshake fails, ErrBadHandshake is returned along with a
+// non-nil *http.Response so that callers can handle redirects, authentication,
+// etcetera. The response body may not contain the entire response and does not
+// need to be closed by the application.
+func (d *Dialer) DialContext(ctx context.Context, urlStr string, requestHeader http.Header) (*Conn, *http.Response, string, error) {
+	if d.HandshakeTimeout == 0 {
+		d.HandshakeTimeout = 15 * time.Second
+	}
+
+	challengeKey, err := generateChallengeKey()
+	if err != nil {
+		return nil, nil, "", err
+	}
+
+	u, err := url.Parse(urlStr)
+	if err != nil {
+		return nil, nil, "", err
+	}
+
+	switch u.Scheme {
+	case "ws":
+		u.Scheme = "http"
+	case "wss":
+		u.Scheme = "https"
+	default:
+		return nil, nil, "", errMalformedURL
+	}
+
+	if u.User != nil {
+		// User name and password are not allowed in websocket URIs.
+		return nil, nil, "", errMalformedURL
+	}
+
+	req := &http.Request{
+		Method:     http.MethodGet,
+		URL:        u,
+		Proto:      "HTTP/1.1",
+		ProtoMajor: 1,
+		ProtoMinor: 1,
+		Header:     make(http.Header),
+		Host:       u.Host,
+	}
+	req = req.WithContext(ctx)
+
+	// Set the cookies present in the cookie jar of the dialer
+	if d.Jar != nil {
+		for _, cookie := range d.Jar.Cookies(u) {
+			req.AddCookie(cookie)
+		}
+	}
+
+	// Set the request headers using the capitalization for names and values in
+	// RFC examples. Although the capitalization shouldn't matter, there are
+	// servers that depend on it. The Header.Set method is not used because the
+	// method canonicalizes the header names.
+	req.Header["Upgrade"] = []string{"websocket"}
+	req.Header["Connection"] = []string{"Upgrade"}
+	req.Header["Sec-WebSocket-Key"] = []string{challengeKey}
+	req.Header["Sec-WebSocket-Version"] = []string{"13"}
+	if len(d.Subprotocols) > 0 {
+		req.Header["Sec-WebSocket-Protocol"] = []string{strings.Join(d.Subprotocols, ", ")}
+	}
+	for k, vs := range requestHeader {
+		switch {
+		case k == "Host":
+			if len(vs) > 0 {
+				req.Host = vs[0]
+			}
+		case k == "Upgrade" ||
+			k == "Connection" ||
+			k == "Sec-Websocket-Key" ||
+			k == "Sec-Websocket-Version" ||
+			k == "Sec-Websocket-Extensions" ||
+			(k == "Sec-Websocket-Protocol" && len(d.Subprotocols) > 0):
+			return nil, nil, "", errors.New("websocket: duplicate header not allowed: " + k)
+		case k == "Sec-Websocket-Protocol":
+			req.Header["Sec-WebSocket-Protocol"] = vs
+		default:
+			req.Header[k] = vs
+		}
+	}
+
+	if d.EnableCompression {
+		req.Header["Sec-WebSocket-Extensions"] = []string{"permessage-deflate; server_no_context_takeover; client_no_context_takeover"}
+	}
+
+	if d.HandshakeTimeout != 0 {
+		var cancel func()
+		ctx, cancel = context.WithTimeout(ctx, d.HandshakeTimeout)
+		defer cancel()
+	}
+
+	// Get network dial function.
+	var netDial func(network, add string) (net.Conn, error)
+
+	switch u.Scheme {
+	case "http":
+		if d.NetDialContext != nil {
+			netDial = func(network, addr string) (net.Conn, error) {
+				return d.NetDialContext(ctx, network, addr)
+			}
+		} else if d.NetDial != nil {
+			netDial = d.NetDial
+		}
+	case "https":
+		if d.NetDialTLSContext != nil {
+			netDial = func(network, addr string) (net.Conn, error) {
+				return d.NetDialTLSContext(ctx, network, addr)
+			}
+		} else if d.NetDialContext != nil {
+			netDial = func(network, addr string) (net.Conn, error) {
+				return d.NetDialContext(ctx, network, addr)
+			}
+		} else if d.NetDial != nil {
+			netDial = d.NetDial
+		}
+	default:
+		return nil, nil, "", errMalformedURL
+	}
+
+	if netDial == nil {
+		netDialer := &net.Dialer{}
+		netDial = func(network, addr string) (net.Conn, error) {
+			return netDialer.DialContext(ctx, network, addr)
+		}
+	}
+
+	// If needed, wrap the dial function to set the connection deadline.
+	if deadline, ok := ctx.Deadline(); ok {
+		forwardDial := netDial
+		netDial = func(network, addr string) (net.Conn, error) {
+			c, err := forwardDial(network, addr)
+			if err != nil {
+				return nil, err
+			}
+			err = c.SetDeadline(deadline)
+			if err != nil {
+				_ = c.Close()
+				return nil, err
+			}
+			return c, nil
+		}
+	}
+
+	hostPort, hostNoPort := hostPortNoPort(u)
+	trace := httptrace.ContextClientTrace(ctx)
+	if trace != nil && trace.GetConn != nil {
+		trace.GetConn(hostPort)
+	}
+
+	netConn, err := netDial("tcp", hostPort)
+	if err != nil {
+		return nil, nil, "", err
+	}
+	if trace != nil && trace.GotConn != nil {
+		trace.GotConn(httptrace.GotConnInfo{
+			Conn: netConn,
+		})
+	}
+
+	defer func() {
+		if netConn != nil {
+			_ = netConn.Close()
+		}
+	}()
+
+	if u.Scheme == "https" && d.NetDialTLSContext == nil {
+		// If NetDialTLSContext is set, assume that the TLS handshake has already been done
+
+		cfg := cloneTLSConfig(d.TLSClientConfig)
+		if cfg.ServerName == "" {
+			cfg.ServerName = hostNoPort
+		}
+		tlsConn := tls.Client(netConn, cfg)
+		netConn = tlsConn
+
+		if trace != nil && trace.TLSHandshakeStart != nil {
+			trace.TLSHandshakeStart()
+		}
+		err := doHandshake(ctx, tlsConn, cfg)
+		if trace != nil && trace.TLSHandshakeDone != nil {
+			trace.TLSHandshakeDone(tlsConn.ConnectionState(), err)
+		}
+
+		if err != nil {
+			return nil, nil, "", err
+		}
+	}
+
+	conn := newConn(netConn, false, d.ReadBufferSize, d.WriteBufferSize, d.WriteBufferPool, nil, nil)
+
+	if err := req.Write(netConn); err != nil {
+		return nil, nil, "", err
+	}
+
+	if trace != nil && trace.GotFirstResponseByte != nil {
+		if peek, err := conn.br.Peek(1); err == nil && len(peek) == 1 {
+			trace.GotFirstResponseByte()
+		}
+	}
+
+	resp, err := http.ReadResponse(conn.br, req)
+	if err != nil {
+		if d.TLSClientConfig != nil {
+			for _, proto := range d.TLSClientConfig.NextProtos {
+				if proto != "http/1.1" {
+					return nil, nil, "", fmt.Errorf(
+						"websocket: protocol %q was given but is not supported;"+
+							"sharing tls.Config with net/http Transport can cause this error: %w",
+						proto, err,
+					)
+				}
+			}
+		}
+		return nil, nil, "", err
+	}
+
+	if d.Jar != nil {
+		if rc := resp.Cookies(); len(rc) > 0 {
+			d.Jar.SetCookies(u, rc)
+		}
+	}
+
+	if resp.StatusCode != 101 ||
+		!tokenListContainsValue(resp.Header, "Upgrade", "websocket") ||
+		!tokenListContainsValue(resp.Header, "Connection", "upgrade") ||
+		resp.Header.Get("Sec-Websocket-Accept") != computeAcceptKey(challengeKey) {
+		// Before closing the network connection on return from this
+		// function, slurp up some of the response to aid application
+		// debugging.
+		buf := make([]byte, 1024)
+		n, _ := io.ReadFull(resp.Body, buf)
+		resp.Body = io.NopCloser(bytes.NewReader(buf[:n]))
+		return nil, resp, "", ErrBadHandshake
+	}
+
+	for _, ext := range parseExtensions(resp.Header) {
+		if ext[""] != "permessage-deflate" {
+			continue
+		}
+		_, snct := ext["server_no_context_takeover"]
+		_, cnct := ext["client_no_context_takeover"]
+		if !snct || !cnct {
+			return nil, resp, "", errInvalidCompression
+		}
+		conn.newCompressionWriter = compressNoContextTakeover
+		conn.newDecompressionReader = decompressNoContextTakeover
+		break
+	}
+
+	resp.Body = io.NopCloser(bytes.NewReader([]byte{}))
+
+	err = netConn.SetDeadline(time.Time{})
+	if err != nil {
+		return conn, resp, "", err
+	}
+	netConn = nil // to avoid close in defer.
+	return conn, resp, resp.Header.Get("Sec-Websocket-Protocol"), nil
+}
+
+func cloneTLSConfig(cfg *tls.Config) *tls.Config {
+	if cfg == nil {
+		return &tls.Config{
+			MinVersion: tls.VersionTLS12,
+		}
+	}
+	return cfg.Clone()
+}

internal/websocket/client_test.go

@@ -0,0 +1,33 @@
+// Copyright 2014 The Gorilla WebSocket Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package websocket
+
+import (
+	"net/url"
+	"testing"
+)
+
+var hostPortNoPortTests = []struct {
+	u          *url.URL
+	hostPort   string
+	hostNoPort string
+}{
+	{&url.URL{Scheme: "ws", Host: "example.com"}, "example.com:80", "example.com"},
+	{&url.URL{Scheme: "wss", Host: "example.com"}, "example.com:443", "example.com"},
+	{&url.URL{Scheme: "ws", Host: "example.com:7777"}, "example.com:7777", "example.com"},
+	{&url.URL{Scheme: "wss", Host: "example.com:7777"}, "example.com:7777", "example.com"},
+}
+
+func TestHostPortNoPort(t *testing.T) {
+	for _, tt := range hostPortNoPortTests {
+		hostPort, hostNoPort := hostPortNoPort(tt.u)
+		if hostPort != tt.hostPort {
+			t.Errorf("hostPortNoPort(%v) returned hostPort %q, want %q", tt.u, hostPort, tt.hostPort)
+		}
+		if hostNoPort != tt.hostNoPort {
+			t.Errorf("hostPortNoPort(%v) returned hostNoPort %q, want %q", tt.u, hostNoPort, tt.hostNoPort)
+		}
+	}
+}

internal/websocket/close_code_test.go

@@ -0,0 +1,72 @@
+package websocket
+
+import (
+	"bytes"
+	"testing"
+	"time"
+)
+
+// TestCloseCode_Unset returns zero values before any close frame is observed.
+func TestCloseCode_Unset(t *testing.T) {
+	c := newTestConn(nil, &bytes.Buffer{}, true)
+	if code, incoming := c.CloseCode(); code != 0 || incoming {
+		t.Fatalf("CloseCode() = (%d, %v), want (0, false)", code, incoming)
+	}
+}
+
+// TestCloseCode_Outgoing records a close code we write as outgoing.
+func TestCloseCode_Outgoing(t *testing.T) {
+	c := newTestConn(nil, &bytes.Buffer{}, true)
+	_ = c.WriteControl(CloseMessage, FormatCloseMessage(CloseMessageTooBig, ""), time.Now().Add(time.Second))
+	if code, incoming := c.CloseCode(); code != CloseMessageTooBig || incoming {
+		t.Fatalf("CloseCode() = (%d, %v), want (%d, false)", code, incoming, CloseMessageTooBig)
+	}
+}
+
+// TestCloseCode_Incoming records a close code received from the peer as incoming
+// (and the subsequent RFC echo must not flip it to outgoing).
+func TestCloseCode_Incoming(t *testing.T) {
+	var buf bytes.Buffer
+	// Client writes a close frame (masked) into buf.
+	wc := newTestConn(nil, &buf, false)
+	_ = wc.WriteControl(CloseMessage, FormatCloseMessage(CloseGoingAway, ""), time.Now().Add(time.Second))
+
+	// Server reads it; advanceFrame records it as incoming and echoes a close.
+	rc := newTestConn(&buf, &bytes.Buffer{}, true)
+	if _, _, err := rc.NextReader(); !IsCloseError(err, CloseGoingAway) {
+		t.Fatalf("NextReader err = %v, want CloseError(%d)", err, CloseGoingAway)
+	}
+	if code, incoming := rc.CloseCode(); code != CloseGoingAway || !incoming {
+		t.Fatalf("CloseCode() = (%d, %v), want (%d, true)", code, incoming, CloseGoingAway)
+	}
+}
+
+// TestCloseCode_FirstWins keeps the initiating close, ignoring later frames.
+func TestCloseCode_FirstWins(t *testing.T) {
+	c := newTestConn(nil, &bytes.Buffer{}, true)
+	deadline := time.Now().Add(time.Second)
+	_ = c.WriteControl(CloseMessage, FormatCloseMessage(CloseMessageTooBig, ""), deadline)
+	_ = c.WriteControl(CloseMessage, FormatCloseMessage(CloseNormalClosure, ""), deadline)
+	if code, incoming := c.CloseCode(); code != CloseMessageTooBig || incoming {
+		t.Fatalf("CloseCode() = (%d, %v), want (%d, false)", code, incoming, CloseMessageTooBig)
+	}
+}
+
+// TestCloseCode_IncomingAppRangeStillIncoming proves an attacker-supplied
+// application-range close code (3000-4999) is recorded as incoming, so the
+// Centrifuge layer (which only records !incoming) will not turn it into a label.
+func TestCloseCode_IncomingAppRangeStillIncoming(t *testing.T) {
+	const appCode = 4321
+	var buf bytes.Buffer
+	wc := newTestConn(nil, &buf, false)
+	_ = wc.WriteControl(CloseMessage, FormatCloseMessage(appCode, ""), time.Now().Add(time.Second))
+
+	rc := newTestConn(&buf, &bytes.Buffer{}, true)
+	if _, _, err := rc.NextReader(); !IsCloseError(err, appCode) {
+		t.Fatalf("NextReader err = %v, want CloseError(%d)", err, appCode)
+	}
+	code, incoming := rc.CloseCode()
+	if code != appCode || !incoming {
+		t.Fatalf("CloseCode() = (%d, %v), want (%d, true)", code, incoming, appCode)
+	}
+}

internal/websocket/compression.go

@@ -0,0 +1,148 @@
+// Copyright 2017 The Gorilla WebSocket Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package websocket
+
+import (
+	"compress/flate"
+	"errors"
+	"io"
+	"strings"
+	"sync"
+)
+
+const (
+	minCompressionLevel     = -2 // flate.HuffmanOnly not defined in Go < 1.6
+	maxCompressionLevel     = flate.BestCompression
+	defaultCompressionLevel = 1
+)
+
+var (
+	flateWriterPools [maxCompressionLevel - minCompressionLevel + 1]sync.Pool
+	flateReaderPool  = sync.Pool{New: func() interface{} {
+		return flate.NewReader(nil)
+	}}
+)
+
+func decompressNoContextTakeover(r io.Reader) io.ReadCloser {
+	const tail =
+	// Add four bytes as specified in RFC
+	"\x00\x00\xff\xff" +
+		// Add final block to squelch unexpected EOF error from flate reader.
+		"\x01\x00\x00\xff\xff"
+
+	fr, _ := flateReaderPool.Get().(io.ReadCloser)
+	_ = fr.(flate.Resetter).Reset(io.MultiReader(r, strings.NewReader(tail)), nil)
+	return &flateReadWrapper{fr}
+}
+
+func isValidCompressionLevel(level int) bool {
+	return minCompressionLevel <= level && level <= maxCompressionLevel
+}
+
+func compressNoContextTakeover(w io.WriteCloser, level int) io.WriteCloser {
+	p := &flateWriterPools[level-minCompressionLevel]
+	tw := &truncWriter{w: w}
+	fw, _ := p.Get().(*flate.Writer)
+	if fw == nil {
+		fw, _ = flate.NewWriter(tw, level)
+	} else {
+		fw.Reset(tw)
+	}
+	return &flateWriteWrapper{fw: fw, tw: tw, p: p}
+}
+
+// truncWriter is an io.Writer that writes all but the last four bytes of the
+// stream to another io.Writer.
+type truncWriter struct {
+	w io.WriteCloser
+	n int
+	p [4]byte
+}
+
+func (w *truncWriter) Write(p []byte) (int, error) {
+	n := 0
+
+	// fill buffer first for simplicity.
+	if w.n < len(w.p) {
+		n = copy(w.p[w.n:], p)
+		p = p[n:]
+		w.n += n
+		if len(p) == 0 {
+			return n, nil
+		}
+	}
+
+	m := len(p)
+	if m > len(w.p) {
+		m = len(w.p)
+	}
+
+	if nn, err := w.w.Write(w.p[:m]); err != nil {
+		return n + nn, err
+	}
+
+	copy(w.p[:], w.p[m:])
+	copy(w.p[len(w.p)-m:], p[len(p)-m:])
+	nn, err := w.w.Write(p[:len(p)-m])
+	return n + nn, err
+}
+
+type flateWriteWrapper struct {
+	fw *flate.Writer
+	tw *truncWriter
+	p  *sync.Pool
+}
+
+func (w *flateWriteWrapper) Write(p []byte) (int, error) {
+	if w.fw == nil {
+		return 0, errWriteClosed
+	}
+	return w.fw.Write(p)
+}
+
+func (w *flateWriteWrapper) Close() error {
+	if w.fw == nil {
+		return errWriteClosed
+	}
+	err1 := w.fw.Flush()
+	w.p.Put(w.fw)
+	w.fw = nil
+	if w.tw.p != [4]byte{0, 0, 0xff, 0xff} {
+		return errors.New("websocket: internal error, unexpected bytes at end of flate stream")
+	}
+	err2 := w.tw.w.Close()
+	if err1 != nil {
+		return err1
+	}
+	return err2
+}
+
+type flateReadWrapper struct {
+	fr io.ReadCloser
+}
+
+func (r *flateReadWrapper) Read(p []byte) (int, error) {
+	if r.fr == nil {
+		return 0, io.ErrClosedPipe
+	}
+	n, err := r.fr.Read(p)
+	if err == io.EOF {
+		// Preemptively place the reader back in the pool. This helps with
+		// scenarios where the application does not call NextReader() soon after
+		// this final read.
+		_ = r.Close()
+	}
+	return n, err
+}
+
+func (r *flateReadWrapper) Close() error {
+	if r.fr == nil {
+		return io.ErrClosedPipe
+	}
+	err := r.fr.Close()
+	flateReaderPool.Put(r.fr)
+	r.fr = nil
+	return err
+}

internal/websocket/compression_limit_test.go

@@ -0,0 +1,391 @@
+package websocket
+
+import (
+	"bytes"
+	"errors"
+	"io"
+	"math/rand"
+	"testing"
+)
+
+// writeCompressedFrame returns a buffer holding a single permessage-deflate
+// compressed binary message carrying payload. The frame is written as a client
+// (masked) so it can be consumed by a server-side reader.
+func writeCompressedFrame(t *testing.T, payload []byte, level int) *bytes.Buffer {
+	t.Helper()
+	var buf bytes.Buffer
+	wc := newTestConn(nil, &buf, false)
+	wc.newCompressionWriter = compressNoContextTakeover
+	if level != 0 {
+		if err := wc.SetCompressionLevel(level); err != nil {
+			t.Fatalf("SetCompressionLevel: %v", err)
+		}
+	}
+	w, err := wc.NextWriter(BinaryMessage)
+	if err != nil {
+		t.Fatalf("NextWriter: %v", err)
+	}
+	if _, err := w.Write(payload); err != nil {
+		t.Fatalf("Write: %v", err)
+	}
+	if err := w.Close(); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+	return &buf
+}
+
+// newCompressedReader builds a server-side reader for frames stored in src,
+// with permessage-deflate decompression enabled and the given decompressed read
+// limit. out captures any control frames (e.g. CloseMessageTooBig) the reader
+// writes back to the peer.
+func newCompressedReader(src *bytes.Buffer, decompressedLimit int64) (rc *Conn, out *bytes.Buffer) {
+	out = &bytes.Buffer{}
+	rc = newTestConn(src, out, true)
+	rc.newDecompressionReader = decompressNoContextTakeover
+	rc.SetDecompressedReadLimit(decompressedLimit)
+	return rc, out
+}
+
+func sentTooBig(out *bytes.Buffer) bool {
+	// FormatCloseMessage encodes the 2-byte close code; the server reader does
+	// not mask control frames, so the encoded code appears verbatim in out.
+	return bytes.Contains(out.Bytes(), FormatCloseMessage(CloseMessageTooBig, ""))
+}
+
+// TestDecompressedReadLimit_Bomb proves that a tiny compressed frame that
+// inflates to a huge payload is rejected with ErrReadLimit, that the close
+// frame is sent, and crucially that memory is bounded: io.ReadAll stops after
+// at most limit+1 bytes instead of materializing the whole bomb.
+func TestDecompressedReadLimit_Bomb(t *testing.T) {
+	const limit = 64 * 1024
+	const decompressedSize = 8 * 1024 * 1024 // 8MB inflates from a few KB.
+
+	src := writeCompressedFrame(t, bytes.Repeat([]byte{'A'}, decompressedSize), 9)
+	if src.Len() >= limit {
+		t.Fatalf("compressed frame %d bytes is not under the compressed read limit %d; test would not isolate the decompressed limit", src.Len(), limit)
+	}
+
+	rc, out := newCompressedReader(src, limit)
+
+	mt, r, err := rc.NextReader()
+	if err != nil {
+		t.Fatalf("NextReader: %v", err)
+	}
+	if mt != BinaryMessage {
+		t.Fatalf("message type = %d, want %d", mt, BinaryMessage)
+	}
+
+	data, err := io.ReadAll(r)
+	if !errors.Is(err, ErrReadLimit) {
+		t.Fatalf("io.ReadAll error = %v, want ErrReadLimit", err)
+	}
+	if int64(len(data)) > limit+1 {
+		t.Fatalf("read %d bytes, want at most %d (memory must stay bounded)", len(data), limit+1)
+	}
+	if !sentTooBig(out) {
+		t.Fatalf("expected CloseMessageTooBig control frame to be sent")
+	}
+	// The server-sent 1009 must be observable as an outgoing close code.
+	if code, incoming := rc.CloseCode(); code != CloseMessageTooBig || incoming {
+		t.Fatalf("CloseCode() = (%d, %v), want (%d, false)", code, incoming, CloseMessageTooBig)
+	}
+}
+
+// TestDecompressedReadLimit_Boundary proves the exact accept/reject boundary:
+// a message of exactly limit bytes is accepted, limit+1 is rejected.
+func TestDecompressedReadLimit_Boundary(t *testing.T) {
+	const limit = 4096
+
+	cases := []struct {
+		name       string
+		size       int
+		wantErr    bool
+		wantTooBig bool
+	}{
+		{"under", limit - 1, false, false},
+		{"exact", limit, false, false},
+		{"over_by_one", limit + 1, true, true},
+		{"over_by_many", limit * 4, true, true},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			payload := bytes.Repeat([]byte{'x'}, tc.size)
+			src := writeCompressedFrame(t, payload, 0)
+			rc, out := newCompressedReader(src, limit)
+
+			_, r, err := rc.NextReader()
+			if err != nil {
+				t.Fatalf("NextReader: %v", err)
+			}
+			data, err := io.ReadAll(r)
+
+			if tc.wantErr {
+				if !errors.Is(err, ErrReadLimit) {
+					t.Fatalf("error = %v, want ErrReadLimit", err)
+				}
+			} else {
+				if err != nil {
+					t.Fatalf("unexpected error: %v", err)
+				}
+				if !bytes.Equal(data, payload) {
+					t.Fatalf("decompressed payload mismatch: got %d bytes, want %d", len(data), len(payload))
+				}
+			}
+			if got := sentTooBig(out); got != tc.wantTooBig {
+				t.Fatalf("sentTooBig = %v, want %v", got, tc.wantTooBig)
+			}
+		})
+	}
+}
+
+// TestDecompressedReadLimit_LegitCompressible proves the fix does not break the
+// legitimate case the report worried about: a message whose compressed size is
+// small but whose decompressed size is well above the compressed bytes is still
+// accepted, as long as it stays within the (higher) decompressed limit.
+func TestDecompressedReadLimit_LegitCompressible(t *testing.T) {
+	const compressedReadLimit = 8 * 1024    // mimics MessageSizeLimit
+	const decompressedReadLimit = 80 * 1024 // mimics MessageSizeLimit * multiplier
+	const payloadSize = 64 * 1024           // > compressed limit, < decompressed limit
+
+	payload := bytes.Repeat([]byte("centrifuge "), payloadSize/11+1)[:payloadSize]
+	src := writeCompressedFrame(t, payload, 9)
+	if src.Len() >= compressedReadLimit {
+		t.Fatalf("compressed size %d not under compressed read limit %d", src.Len(), compressedReadLimit)
+	}
+
+	rc, out := newCompressedReader(src, decompressedReadLimit)
+	rc.SetReadLimit(compressedReadLimit) // compressed-bytes limit also active
+
+	_, r, err := rc.NextReader()
+	if err != nil {
+		t.Fatalf("NextReader: %v", err)
+	}
+	data, err := io.ReadAll(r)
+	if err != nil {
+		t.Fatalf("legit compressible message rejected: %v", err)
+	}
+	if !bytes.Equal(data, payload) {
+		t.Fatalf("payload mismatch: got %d bytes, want %d", len(data), len(payload))
+	}
+	if sentTooBig(out) {
+		t.Fatalf("unexpected CloseMessageTooBig for a legit message")
+	}
+}
+
+// TestDecompressedReadLimit_ReproducesPreFixVulnerability documents the
+// vulnerability the fix addresses: with only the compressed-bytes limit
+// (SetReadLimit, the protection that existed before this change) and no
+// decompressed limit, a small compressed frame that satisfies the compressed
+// limit still inflates without bound. This is exactly the pre-fix behavior
+// (decompressedReadLimit == 0 is equivalent to the code before the fix).
+func TestDecompressedReadLimit_ReproducesPreFixVulnerability(t *testing.T) {
+	const compressedReadLimit = 64 * 1024    // the old message_size_limit default
+	const decompressedSize = 8 * 1024 * 1024 // 8MB from a few KB on the wire
+
+	src := writeCompressedFrame(t, bytes.Repeat([]byte{'A'}, decompressedSize), 9)
+	compressedLen := src.Len() // capture before reading drains the buffer
+	// The compressed frame is comfortably under the compressed read limit, so
+	// the pre-fix defense does not trigger.
+	if compressedLen >= compressedReadLimit {
+		t.Fatalf("compressed frame %d bytes not under compressed limit %d", compressedLen, compressedReadLimit)
+	}
+
+	rc, out := newCompressedReader(src, 0) // 0 == no decompressed limit (pre-fix)
+	rc.SetReadLimit(compressedReadLimit)   // only the old compressed-bytes defense
+
+	_, r, err := rc.NextReader()
+	if err != nil {
+		t.Fatalf("NextReader: %v", err)
+	}
+	data, err := io.ReadAll(r)
+	// Pre-fix: no error and the full 8MB is materialized despite the compressed
+	// read limit, demonstrating the unbounded amplification (the bug).
+	if err != nil {
+		t.Fatalf("pre-fix path unexpectedly errored: %v", err)
+	}
+	if len(data) != decompressedSize {
+		t.Fatalf("inflated to %d bytes, want full %d (unbounded amplification)", len(data), decompressedSize)
+	}
+	if sentTooBig(out) {
+		t.Fatalf("pre-fix path should not send CloseMessageTooBig")
+	}
+	t.Logf("reproduced: %d compressed wire bytes (<= %d limit) inflated to %d bytes unbounded",
+		compressedLen, compressedReadLimit, len(data))
+}
+
+// TestDecompressedReadLimit_Disabled proves that a zero limit means unlimited
+// (the fork primitive's contract): the bomb fully inflates without error.
+func TestDecompressedReadLimit_Disabled(t *testing.T) {
+	const decompressedSize = 1 * 1024 * 1024
+	payload := bytes.Repeat([]byte{'A'}, decompressedSize)
+	src := writeCompressedFrame(t, payload, 9)
+
+	rc, out := newCompressedReader(src, 0) // 0 == unlimited
+
+	_, r, err := rc.NextReader()
+	if err != nil {
+		t.Fatalf("NextReader: %v", err)
+	}
+	data, err := io.ReadAll(r)
+	if err != nil {
+		t.Fatalf("unexpected error with limit disabled: %v", err)
+	}
+	if len(data) != decompressedSize {
+		t.Fatalf("read %d bytes, want %d", len(data), decompressedSize)
+	}
+	if sentTooBig(out) {
+		t.Fatalf("unexpected CloseMessageTooBig with limit disabled")
+	}
+}
+
+// TestDecompressedReadLimit_IgnoredWithoutCompression proves the decompressed
+// limit only applies to compressed frames: an uncompressed message larger than
+// the decompressed limit (but within the compressed read limit) passes through,
+// because the limit is gated on the RSV1/decompress path.
+func TestDecompressedReadLimit_IgnoredWithoutCompression(t *testing.T) {
+	const decompressedReadLimit = 1024
+	const payloadSize = 8 * 1024
+
+	var src bytes.Buffer
+	wc := newTestConn(nil, &src, false) // no compression writer -> plain frame
+	w, err := wc.NextWriter(BinaryMessage)
+	if err != nil {
+		t.Fatalf("NextWriter: %v", err)
+	}
+	payload := bytes.Repeat([]byte{'y'}, payloadSize)
+	if _, err := w.Write(payload); err != nil {
+		t.Fatalf("Write: %v", err)
+	}
+	if err := w.Close(); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+
+	rc, _ := newCompressedReader(&src, decompressedReadLimit)
+	// readLimit (compressed bytes) generous enough to admit the plain frame.
+	rc.SetReadLimit(payloadSize + 1024)
+
+	_, r, err := rc.NextReader()
+	if err != nil {
+		t.Fatalf("NextReader: %v", err)
+	}
+	data, err := io.ReadAll(r)
+	if err != nil {
+		t.Fatalf("uncompressed message wrongly limited: %v", err)
+	}
+	if !bytes.Equal(data, payload) {
+		t.Fatalf("payload mismatch: got %d bytes, want %d", len(data), len(payload))
+	}
+}
+
+// TestDecompressedReadLimit_CompressedLimitStillApplies proves the two limits
+// are independent: with a decompressed limit set, the compressed-bytes limit
+// enforced in advanceFrame still trips for an oversized compressed frame.
+func TestDecompressedReadLimit_CompressedLimitStillApplies(t *testing.T) {
+	const compressedReadLimit = 1024
+	// Incompressible (deterministic pseudo-random) payload so compressed size
+	// stays large and exceeds the compressed read limit before decompression
+	// even matters.
+	payload := make([]byte, 8*1024)
+	rng := rand.New(rand.NewSource(1))
+	_, _ = rng.Read(payload)
+	src := writeCompressedFrame(t, payload, 0)
+	if src.Len() <= compressedReadLimit {
+		t.Fatalf("payload compressed too well (%d bytes) to exceed compressed limit", src.Len())
+	}
+
+	rc, _ := newCompressedReader(src, 10*1024*1024) // generous decompressed limit
+	rc.SetReadLimit(compressedReadLimit)
+
+	_, r, err := rc.NextReader()
+	if err != nil {
+		// Limit may surface from NextReader itself (advanceFrame) ...
+		if !errors.Is(err, ErrReadLimit) {
+			t.Fatalf("NextReader error = %v, want ErrReadLimit", err)
+		}
+		return
+	}
+	// ... or while reading.
+	if _, err := io.ReadAll(r); !errors.Is(err, ErrReadLimit) {
+		t.Fatalf("io.ReadAll error = %v, want ErrReadLimit", err)
+	}
+}
+
+// TestDecompressedReadLimit_ErrorIsPermanent proves that once the decompressed
+// limit trips, the error is sticky: subsequent NextReader calls keep returning
+// ErrReadLimit instead of advancing to the next frame. This mirrors the
+// compressed-bytes limit (enforced in advanceFrame, which sets c.readErr via
+// NextReader) and honors the gorilla contract that read errors are permanent.
+func TestDecompressedReadLimit_ErrorIsPermanent(t *testing.T) {
+	const limit = 4096
+
+	// Two back-to-back compressed frames in one stream. The first is a bomb that
+	// trips the decompressed limit; the second is a small legitimate message that
+	// must NOT be delivered once the connection is in its terminal error state.
+	src := writeCompressedFrame(t, bytes.Repeat([]byte{'A'}, 1*1024*1024), 9)
+	next := writeCompressedFrame(t, []byte("should never be read"), 9)
+	src.Write(next.Bytes())
+
+	rc, out := newCompressedReader(src, limit)
+
+	_, r, err := rc.NextReader()
+	if err != nil {
+		t.Fatalf("NextReader: %v", err)
+	}
+	if _, err := io.ReadAll(r); !errors.Is(err, ErrReadLimit) {
+		t.Fatalf("io.ReadAll error = %v, want ErrReadLimit", err)
+	}
+	if !sentTooBig(out) {
+		t.Fatalf("expected CloseMessageTooBig control frame to be sent")
+	}
+
+	// The error must now be permanent: NextReader keeps returning ErrReadLimit
+	// rather than reading the second frame.
+	for i := 0; i < 3; i++ {
+		mt, r, err := rc.NextReader()
+		if !errors.Is(err, ErrReadLimit) {
+			t.Fatalf("NextReader #%d after limit = (%d, %v), want ErrReadLimit", i, mt, err)
+		}
+		if r != nil {
+			t.Fatalf("NextReader #%d returned non-nil reader after permanent error", i)
+		}
+	}
+}
+
+// TestDecompressedReadLimit_StreamingReads proves the limit is enforced on the
+// incremental NextReader path (used by the bidirectional handler's streaming
+// command decoder), not only via io.ReadAll: reading small chunks still trips
+// the limit and keeps total consumed bytes bounded.
+func TestDecompressedReadLimit_StreamingReads(t *testing.T) {
+	const limit = 4096
+	payload := bytes.Repeat([]byte{'z'}, 1*1024*1024)
+	src := writeCompressedFrame(t, payload, 9)
+
+	rc, out := newCompressedReader(src, limit)
+	_, r, err := rc.NextReader()
+	if err != nil {
+		t.Fatalf("NextReader: %v", err)
+	}
+
+	buf := make([]byte, 256)
+	var total int
+	for {
+		n, err := r.Read(buf)
+		total += n
+		if err != nil {
+			if !errors.Is(err, ErrReadLimit) {
+				t.Fatalf("Read error = %v, want ErrReadLimit", err)
+			}
+			break
+		}
+		if total > limit+1 {
+			t.Fatalf("consumed %d bytes without hitting limit %d", total, limit)
+		}
+	}
+	if total > limit+1 {
+		t.Fatalf("consumed %d bytes, want at most %d", total, limit+1)
+	}
+	if !sentTooBig(out) {
+		t.Fatalf("expected CloseMessageTooBig control frame to be sent")
+	}
+}

internal/websocket/compression_test.go

@@ -0,0 +1,79 @@
+package websocket
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"testing"
+)
+
+type nopCloser struct{ io.Writer }
+
+func (nopCloser) Close() error { return nil }
+
+func TestTruncWriter(t *testing.T) {
+	const data = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijlkmnopqrstuvwxyz987654321"
+	for n := 1; n <= 10; n++ {
+		var b bytes.Buffer
+		w := &truncWriter{w: nopCloser{&b}}
+		p := []byte(data)
+		for len(p) > 0 {
+			m := len(p)
+			if m > n {
+				m = n
+			}
+			_, _ = w.Write(p[:m])
+			p = p[m:]
+		}
+		if b.String() != data[:len(data)-len(w.p)] {
+			t.Errorf("%d: %q", n, b.String())
+		}
+	}
+}
+
+func textMessages(num int) [][]byte {
+	messages := make([][]byte, num)
+	for i := 0; i < num; i++ {
+		msg := fmt.Sprintf("planet: %d, country: %d, city: %d, street: %d", i, i, i, i)
+		messages[i] = []byte(msg)
+	}
+	return messages
+}
+
+func BenchmarkWriteNoCompression(b *testing.B) {
+	w := io.Discard
+	c := newTestConn(nil, w, false)
+	messages := textMessages(100)
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_ = c.WriteMessage(TextMessage, messages[i%len(messages)])
+	}
+	b.ReportAllocs()
+}
+
+func BenchmarkWriteWithCompression(b *testing.B) {
+	w := io.Discard
+	c := newTestConn(nil, w, false)
+	messages := textMessages(100)
+	c.enableWriteCompression = true
+	c.newCompressionWriter = compressNoContextTakeover
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_ = c.WriteMessage(TextMessage, messages[i%len(messages)])
+	}
+	b.ReportAllocs()
+}
+
+func TestValidCompressionLevel(t *testing.T) {
+	c := newTestConn(nil, nil, false)
+	for _, level := range []int{minCompressionLevel - 1, maxCompressionLevel + 1} {
+		if err := c.SetCompressionLevel(level); err == nil {
+			t.Errorf("no error for level %d", level)
+		}
+	}
+	for _, level := range []int{minCompressionLevel, maxCompressionLevel} {
+		if err := c.SetCompressionLevel(level); err != nil {
+			t.Errorf("error for level %d", level)
+		}
+	}
+}

internal/websocket/conn_broadcast_test.go

@@ -0,0 +1,130 @@
+// Copyright 2017 The Gorilla WebSocket Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package websocket
+
+import (
+	"io"
+	"sync/atomic"
+	"testing"
+)
+
+// broadcastBench allows to run broadcast benchmarks.
+// In every broadcast benchmark we create many connections, then send the same
+// message into every connection and wait for all writes complete. This emulates
+// an application where many connections listen to the same data - i.e. PUB/SUB
+// scenarios with many subscribers in one channel.
+type broadcastBench struct {
+	w           io.Writer
+	closeCh     chan struct{}
+	doneCh      chan struct{}
+	count       int32
+	conns       []*broadcastConn
+	compression bool
+	usePrepared bool
+}
+
+type broadcastMessage struct {
+	payload  []byte
+	prepared *PreparedMessage
+}
+
+type broadcastConn struct {
+	conn  *Conn
+	msgCh chan *broadcastMessage
+}
+
+func newBroadcastConn(c *Conn) *broadcastConn {
+	return &broadcastConn{
+		conn:  c,
+		msgCh: make(chan *broadcastMessage, 1),
+	}
+}
+
+func newBroadcastBench(usePrepared, compression bool) *broadcastBench {
+	bench := &broadcastBench{
+		w:           io.Discard,
+		doneCh:      make(chan struct{}),
+		closeCh:     make(chan struct{}),
+		usePrepared: usePrepared,
+		compression: compression,
+	}
+	bench.makeConns(10000)
+	return bench
+}
+
+func (b *broadcastBench) makeConns(numConns int) {
+	conns := make([]*broadcastConn, numConns)
+
+	for i := 0; i < numConns; i++ {
+		c := newTestConn(nil, b.w, true)
+		if b.compression {
+			c.enableWriteCompression = true
+			c.newCompressionWriter = compressNoContextTakeover
+		}
+		conns[i] = newBroadcastConn(c)
+		go func(c *broadcastConn) {
+			for {
+				select {
+				case msg := <-c.msgCh:
+					if msg.prepared != nil {
+						_ = c.conn.WritePreparedMessage(msg.prepared)
+					} else {
+						_ = c.conn.WriteMessage(TextMessage, msg.payload)
+					}
+					val := atomic.AddInt32(&b.count, 1)
+					if val%int32(numConns) == 0 {
+						b.doneCh <- struct{}{}
+					}
+				case <-b.closeCh:
+					return
+				}
+			}
+		}(conns[i])
+	}
+	b.conns = conns
+}
+
+func (b *broadcastBench) close() {
+	close(b.closeCh)
+}
+
+func (b *broadcastBench) broadcastOnce(msg *broadcastMessage) {
+	for _, c := range b.conns {
+		c.msgCh <- msg
+	}
+	<-b.doneCh
+}
+
+func BenchmarkBroadcast(b *testing.B) {
+	benchmarks := []struct {
+		name        string
+		usePrepared bool
+		compression bool
+	}{
+		{"NoCompression", false, false},
+		{"Compression", false, true},
+		{"NoCompressionPrepared", true, false},
+		{"CompressionPrepared", true, true},
+	}
+	payload := textMessages(1)[0]
+	for _, bm := range benchmarks {
+		b.Run(bm.name, func(b *testing.B) {
+			bench := newBroadcastBench(bm.usePrepared, bm.compression)
+			defer bench.close()
+			message := &broadcastMessage{
+				payload: payload,
+			}
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				if bench.usePrepared {
+					pm, _ := NewPreparedMessage(TextMessage, message.payload)
+					message.prepared = pm
+				}
+				bench.broadcastOnce(message)
+			}
+			b.ReportAllocs()
+		})
+	}
+}

internal/websocket/conn_test.go

@@ -0,0 +1,654 @@
+// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package websocket
+
+import (
+	"bufio"
+	"bytes"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"reflect"
+	"sync"
+	"testing"
+	"testing/iotest"
+	"time"
+)
+
+var _ net.Error = errWriteTimeout
+
+type fakeNetConn struct {
+	io.Reader
+	io.Writer
+}
+
+func (c fakeNetConn) Close() error                       { return nil }
+func (c fakeNetConn) LocalAddr() net.Addr                { return localAddr }
+func (c fakeNetConn) RemoteAddr() net.Addr               { return remoteAddr }
+func (c fakeNetConn) SetDeadline(_ time.Time) error      { return nil }
+func (c fakeNetConn) SetReadDeadline(_ time.Time) error  { return nil }
+func (c fakeNetConn) SetWriteDeadline(_ time.Time) error { return nil }
+
+type fakeAddr int
+
+var (
+	localAddr  = fakeAddr(1)
+	remoteAddr = fakeAddr(2)
+)
+
+func (a fakeAddr) Network() string {
+	return "net"
+}
+
+func (a fakeAddr) String() string {
+	return "str"
+}
+
+// newTestConn creates a connection backed by a fake network connection using
+// default values for buffering.
+func newTestConn(r io.Reader, w io.Writer, isServer bool) *Conn {
+	return newConn(fakeNetConn{Reader: r, Writer: w}, isServer, 1024, 1024, nil, nil, nil)
+}
+
+func TestFraming(t *testing.T) {
+	frameSizes := []int{
+		0, 1, 2, 124, 125, 126, 127, 128, 129, 65534, 65535,
+		// 65536, 65537
+	}
+	var readChunkers = []struct {
+		name string
+		f    func(io.Reader) io.Reader
+	}{
+		{"half", iotest.HalfReader},
+		{"one", iotest.OneByteReader},
+		{"asis", func(r io.Reader) io.Reader { return r }},
+	}
+	writeBuf := make([]byte, 65537)
+	for i := range writeBuf {
+		writeBuf[i] = byte(i)
+	}
+	var writers = []struct {
+		name string
+		f    func(w io.Writer, n int) (int, error)
+	}{
+		{"iocopy", func(w io.Writer, n int) (int, error) {
+			nn, err := io.Copy(w, bytes.NewReader(writeBuf[:n]))
+			return int(nn), err
+		}},
+		{"write", func(w io.Writer, n int) (int, error) {
+			return w.Write(writeBuf[:n])
+		}},
+	}
+
+	for _, compress := range []bool{false, true} {
+		for _, isServer := range []bool{true, false} {
+			for _, chunker := range readChunkers {
+
+				var connBuf bytes.Buffer
+				wc := newTestConn(nil, &connBuf, isServer)
+				rc := newTestConn(chunker.f(&connBuf), nil, !isServer)
+				if compress {
+					wc.newCompressionWriter = compressNoContextTakeover
+					rc.newDecompressionReader = decompressNoContextTakeover
+				}
+				for _, n := range frameSizes {
+					for _, writer := range writers {
+						name := fmt.Sprintf("z:%v, s:%v, r:%s, n:%d w:%s", compress, isServer, chunker.name, n, writer.name)
+
+						w, err := wc.NextWriter(TextMessage)
+						if err != nil {
+							t.Errorf("%s: wc.NextWriter() returned %v", name, err)
+							continue
+						}
+						nn, err := writer.f(w, n)
+						if err != nil || nn != n {
+							t.Errorf("%s: w.Write(writeBuf[:n]) returned %d, %v", name, nn, err)
+							continue
+						}
+						err = w.Close()
+						if err != nil {
+							t.Errorf("%s: w.Close() returned %v", name, err)
+							continue
+						}
+
+						opCode, r, err := rc.NextReader()
+						if err != nil || opCode != TextMessage {
+							t.Errorf("%s: NextReader() returned %d, r, %v", name, opCode, err)
+							continue
+						}
+
+						t.Logf("frame size: %d", n)
+						rbuf, err := io.ReadAll(r)
+						if err != nil {
+							t.Errorf("%s: ReadFull() returned rbuf, %v", name, err)
+							continue
+						}
+
+						if len(rbuf) != n {
+							t.Errorf("%s: len(rbuf) is %d, want %d", name, len(rbuf), n)
+							continue
+						}
+
+						for i, b := range rbuf {
+							if byte(i) != b {
+								t.Errorf("%s: bad byte at offset %d", name, i)
+								break
+							}
+						}
+					}
+				}
+			}
+		}
+	}
+}
+
+func TestWriteControlDeadline(t *testing.T) {
+	t.Parallel()
+	message := []byte("hello")
+	var connBuf bytes.Buffer
+	c := newTestConn(nil, &connBuf, true)
+	if err := c.WriteControl(PongMessage, message, time.Time{}); err != nil {
+		t.Errorf("WriteControl(..., zero deadline) = %v, want nil", err)
+	}
+	if err := c.WriteControl(PongMessage, message, time.Now().Add(time.Second)); err != nil {
+		t.Errorf("WriteControl(..., future deadline) = %v, want nil", err)
+	}
+	if err := c.WriteControl(PongMessage, message, time.Now().Add(-time.Second)); err == nil {
+		t.Errorf("WriteControl(..., past deadline) = nil, want timeout error")
+	}
+}
+
+func TestControl(t *testing.T) {
+	const message = "this is a ping/pong message"
+	for _, isServer := range []bool{true, false} {
+		for _, isWriteControl := range []bool{true, false} {
+			name := fmt.Sprintf("s:%v, wc:%v", isServer, isWriteControl)
+			var connBuf bytes.Buffer
+			wc := newTestConn(nil, &connBuf, isServer)
+			rc := newTestConn(&connBuf, nil, !isServer)
+			if isWriteControl {
+				_ = wc.WriteControl(PongMessage, []byte(message), time.Now().Add(time.Second))
+			} else {
+				w, err := wc.NextWriter(PongMessage)
+				if err != nil {
+					t.Errorf("%s: wc.NextWriter() returned %v", name, err)
+					continue
+				}
+				if _, err := w.Write([]byte(message)); err != nil {
+					t.Errorf("%s: w.Write() returned %v", name, err)
+					continue
+				}
+				if err := w.Close(); err != nil {
+					t.Errorf("%s: w.Close() returned %v", name, err)
+					continue
+				}
+				var actualMessage string
+				rc.SetPongHandler(func(s []byte) error { actualMessage = string(s); return nil })
+				_, _, _ = rc.NextReader()
+				if actualMessage != message {
+					t.Errorf("%s: pong=%q, want %q", name, actualMessage, message)
+					continue
+				}
+			}
+		}
+	}
+}
+
+// simpleBufferPool is an implementation of BufferPool for TestWriteBufferPool.
+type simpleBufferPool struct {
+	v interface{}
+}
+
+func (p *simpleBufferPool) Get() interface{} {
+	v := p.v
+	p.v = nil
+	return v
+}
+
+func (p *simpleBufferPool) Put(v interface{}) {
+	p.v = v
+}
+
+func TestWriteBufferPool(t *testing.T) {
+	const message = "Now is the time for all good people to come to the aid of the party."
+
+	var buf bytes.Buffer
+	var pool simpleBufferPool
+	rc := newTestConn(&buf, nil, false)
+
+	// Specify writeBufferSize smaller than message size to ensure that pooling
+	// works with fragmented messages.
+	wc := newConn(fakeNetConn{Writer: &buf}, true, 1024, len(message)-1, &pool, nil, nil)
+
+	if wc.writeBuf != nil {
+		t.Fatal("writeBuf not nil after create")
+	}
+
+	// Part 1: test NextWriter/Write/Close
+
+	w, err := wc.NextWriter(TextMessage)
+	if err != nil {
+		t.Fatalf("wc.NextWriter() returned %v", err)
+	}
+
+	if wc.writeBuf == nil {
+		t.Fatal("writeBuf is nil after NextWriter")
+	}
+
+	writeBufAddr := &wc.writeBuf[0]
+
+	if _, err := io.WriteString(w, message); err != nil {
+		t.Fatalf("io.WriteString(w, message) returned %v", err)
+	}
+
+	if err := w.Close(); err != nil {
+		t.Fatalf("w.Close() returned %v", err)
+	}
+
+	if wc.writeBuf != nil {
+		t.Fatal("writeBuf not nil after w.Close()")
+	}
+
+	if wpd, ok := pool.v.(writePoolData); !ok || len(wpd.buf) == 0 || &wpd.buf[0] != writeBufAddr {
+		t.Fatal("writeBuf not returned to pool")
+	}
+
+	opCode, p, err := rc.ReadMessage()
+	if opCode != TextMessage || err != nil {
+		t.Fatalf("ReadMessage() returned %d, p, %v", opCode, err)
+	}
+
+	if s := string(p); s != message {
+		t.Fatalf("message is %s, want %s", s, message)
+	}
+
+	// Part 2: Test WriteMessage.
+
+	if err := wc.WriteMessage(TextMessage, []byte(message)); err != nil {
+		t.Fatalf("wc.WriteMessage() returned %v", err)
+	}
+
+	if wc.writeBuf != nil {
+		t.Fatal("writeBuf not nil after wc.WriteMessage()")
+	}
+
+	if wpd, ok := pool.v.(writePoolData); !ok || len(wpd.buf) == 0 || &wpd.buf[0] != writeBufAddr {
+		t.Fatal("writeBuf not returned to pool after WriteMessage")
+	}
+
+	opCode, p, err = rc.ReadMessage()
+	if opCode != TextMessage || err != nil {
+		t.Fatalf("ReadMessage() returned %d, p, %v", opCode, err)
+	}
+
+	if s := string(p); s != message {
+		t.Fatalf("message is %s, want %s", s, message)
+	}
+}
+
+// TestWriteBufferPoolSync ensures that *sync.Pool works as a buffer pool.
+func TestWriteBufferPoolSync(t *testing.T) {
+	var buf bytes.Buffer
+	var pool sync.Pool
+	wc := newConn(fakeNetConn{Writer: &buf}, true, 1024, 1024, &pool, nil, nil)
+	rc := newTestConn(&buf, nil, false)
+
+	const message = "Hello World!"
+	for i := 0; i < 3; i++ {
+		if err := wc.WriteMessage(TextMessage, []byte(message)); err != nil {
+			t.Fatalf("wc.WriteMessage() returned %v", err)
+		}
+		opCode, p, err := rc.ReadMessage()
+		if opCode != TextMessage || err != nil {
+			t.Fatalf("ReadMessage() returned %d, p, %v", opCode, err)
+		}
+		if s := string(p); s != message {
+			t.Fatalf("message is %s, want %s", s, message)
+		}
+	}
+}
+
+// errorWriter is an io.Writer than returns an error on all writes.
+type errorWriter struct{}
+
+func (ew errorWriter) Write(_ []byte) (int, error) { return 0, errors.New("error") }
+
+// TestWriteBufferPoolError ensures that buffer is returned to pool after error
+// on write.
+func TestWriteBufferPoolError(t *testing.T) {
+
+	// Part 1: Test NextWriter/Write/Close
+
+	var pool simpleBufferPool
+	wc := newConn(fakeNetConn{Writer: errorWriter{}}, true, 1024, 1024, &pool, nil, nil)
+
+	w, err := wc.NextWriter(TextMessage)
+	if err != nil {
+		t.Fatalf("wc.NextWriter() returned %v", err)
+	}
+
+	if wc.writeBuf == nil {
+		t.Fatal("writeBuf is nil after NextWriter")
+	}
+
+	writeBufAddr := &wc.writeBuf[0]
+
+	if _, err := io.WriteString(w, "Hello"); err != nil {
+		t.Fatalf("io.WriteString(w, message) returned %v", err)
+	}
+
+	if err := w.Close(); err == nil {
+		t.Fatalf("w.Close() did not return error")
+	}
+
+	if wpd, ok := pool.v.(writePoolData); !ok || len(wpd.buf) == 0 || &wpd.buf[0] != writeBufAddr {
+		t.Fatal("writeBuf not returned to pool")
+	}
+
+	// Part 2: Test WriteMessage
+
+	wc = newConn(fakeNetConn{Writer: errorWriter{}}, true, 1024, 1024, &pool, nil, nil)
+
+	if err := wc.WriteMessage(TextMessage, []byte("Hello")); err == nil {
+		t.Fatalf("wc.WriteMessage did not return error")
+	}
+
+	if wpd, ok := pool.v.(writePoolData); !ok || len(wpd.buf) == 0 || &wpd.buf[0] != writeBufAddr {
+		t.Fatal("writeBuf not returned to pool")
+	}
+}
+
+func TestCloseFrameBeforeFinalMessageFrame(t *testing.T) {
+	const bufSize = 512
+
+	expectedErr := &CloseError{Code: CloseNormalClosure, Text: "hello"}
+
+	var b1, b2 bytes.Buffer
+	wc := newConn(&fakeNetConn{Reader: nil, Writer: &b1}, false, 1024, bufSize, nil, nil, nil)
+	rc := newTestConn(&b1, &b2, true)
+
+	w, _ := wc.NextWriter(BinaryMessage)
+	_, _ = w.Write(make([]byte, bufSize+bufSize/2))
+	_ = wc.WriteControl(CloseMessage, FormatCloseMessage(expectedErr.Code, expectedErr.Text), time.Now().Add(10*time.Second))
+	_ = w.Close()
+
+	op, r, err := rc.NextReader()
+	if op != BinaryMessage || err != nil {
+		t.Fatalf("NextReader() returned %d, %v", op, err)
+	}
+	_, err = io.Copy(io.Discard, r)
+	if !reflect.DeepEqual(err, expectedErr) {
+		t.Fatalf("io.Copy() returned %v, want %v", err, expectedErr)
+	}
+	_, _, err = rc.NextReader()
+	if !reflect.DeepEqual(err, expectedErr) {
+		t.Fatalf("NextReader() returned %v, want %v", err, expectedErr)
+	}
+}
+
+func TestEOFWithinFrame(t *testing.T) {
+	const bufSize = 64
+
+	for n := 0; ; n++ {
+		var b bytes.Buffer
+		wc := newTestConn(nil, &b, false)
+		rc := newTestConn(&b, nil, true)
+
+		w, _ := wc.NextWriter(BinaryMessage)
+		_, _ = w.Write(make([]byte, bufSize))
+		_ = w.Close()
+
+		if n >= b.Len() {
+			break
+		}
+		b.Truncate(n)
+
+		op, r, err := rc.NextReader()
+		if err == errUnexpectedEOF {
+			continue
+		}
+		if op != BinaryMessage || err != nil {
+			t.Fatalf("%d: NextReader() returned %d, %v", n, op, err)
+		}
+		_, err = io.Copy(io.Discard, r)
+		if err != errUnexpectedEOF {
+			t.Fatalf("%d: io.Copy() returned %v, want %v", n, err, errUnexpectedEOF)
+		}
+		_, _, err = rc.NextReader()
+		if err != errUnexpectedEOF {
+			t.Fatalf("%d: NextReader() returned %v, want %v", n, err, errUnexpectedEOF)
+		}
+	}
+}
+
+func TestEOFBeforeFinalFrame(t *testing.T) {
+	const bufSize = 512
+
+	var b1, b2 bytes.Buffer
+	wc := newConn(&fakeNetConn{Writer: &b1}, false, 1024, bufSize, nil, nil, nil)
+	rc := newTestConn(&b1, &b2, true)
+
+	w, _ := wc.NextWriter(BinaryMessage)
+	_, _ = w.Write(make([]byte, bufSize+bufSize/2))
+
+	op, r, err := rc.NextReader()
+	if op != BinaryMessage || err != nil {
+		t.Fatalf("NextReader() returned %d, %v", op, err)
+	}
+	_, err = io.Copy(io.Discard, r)
+	if err != errUnexpectedEOF {
+		t.Fatalf("io.Copy() returned %v, want %v", err, errUnexpectedEOF)
+	}
+	_, _, err = rc.NextReader()
+	if err != errUnexpectedEOF {
+		t.Fatalf("NextReader() returned %v, want %v", err, errUnexpectedEOF)
+	}
+}
+
+func TestWriteAfterMessageWriterClose(t *testing.T) {
+	wc := newTestConn(nil, &bytes.Buffer{}, false)
+	w, _ := wc.NextWriter(BinaryMessage)
+	_, _ = io.WriteString(w, "hello")
+	if err := w.Close(); err != nil {
+		t.Fatalf("unexpected error closing message writer, %v", err)
+	}
+
+	if _, err := io.WriteString(w, "world"); err == nil {
+		t.Fatalf("no error writing after close")
+	}
+
+	w, _ = wc.NextWriter(BinaryMessage)
+	_, _ = io.WriteString(w, "hello")
+
+	// close w by getting next writer
+	_, err := wc.NextWriter(BinaryMessage)
+	if err != nil {
+		t.Fatalf("unexpected error getting next writer, %v", err)
+	}
+
+	if _, err := io.WriteString(w, "world"); err == nil {
+		t.Fatalf("no error writing after close")
+	}
+}
+
+func TestReadLimit(t *testing.T) {
+	t.Run("Test ReadLimit is enforced", func(t *testing.T) {
+		const readLimit = 512
+		message := make([]byte, readLimit+1)
+
+		var b1, b2 bytes.Buffer
+		wc := newConn(&fakeNetConn{Writer: &b1}, false, 1024, readLimit-2, nil, nil, nil)
+		rc := newTestConn(&b1, &b2, true)
+		rc.SetReadLimit(readLimit)
+
+		// Send message at the limit with interleaved pong.
+		w, _ := wc.NextWriter(BinaryMessage)
+		_, _ = w.Write(message[:readLimit-1])
+		_ = wc.WriteControl(PongMessage, []byte("this is a pong"), time.Now().Add(10*time.Second))
+		_, _ = w.Write(message[:1])
+		_ = w.Close()
+
+		// Send message larger than the limit.
+		_ = wc.WriteMessage(BinaryMessage, message[:readLimit+1])
+
+		op, _, err := rc.NextReader()
+		if op != BinaryMessage || err != nil {
+			t.Fatalf("1: NextReader() returned %d, %v", op, err)
+		}
+		op, r, err := rc.NextReader()
+		if op != BinaryMessage || err != nil {
+			t.Fatalf("2: NextReader() returned %d, %v", op, err)
+		}
+		_, err = io.Copy(io.Discard, r)
+		if err != ErrReadLimit {
+			t.Fatalf("io.Copy() returned %v", err)
+		}
+	})
+
+	t.Run("Test that ReadLimit cannot be overflowed", func(t *testing.T) {
+		const readLimit = 1
+
+		var b1, b2 bytes.Buffer
+		rc := newTestConn(&b1, &b2, true)
+		rc.SetReadLimit(readLimit)
+
+		// First, send a non-final binary message
+		b1.Write([]byte("\x02\x81"))
+
+		// Mask key
+		b1.Write([]byte("\x00\x00\x00\x00"))
+
+		// First payload
+		b1.Write([]byte("A"))
+
+		// Next, send a negative-length, non-final continuation frame
+		b1.Write([]byte("\x00\xFF\x80\x00\x00\x00\x00\x00\x00\x00"))
+
+		// Mask key
+		b1.Write([]byte("\x00\x00\x00\x00"))
+
+		// Next, send a too long, final continuation frame
+		b1.Write([]byte("\x80\xFF\x00\x00\x00\x00\x00\x00\x00\x05"))
+
+		// Mask key
+		b1.Write([]byte("\x00\x00\x00\x00"))
+
+		// Too-long payload
+		b1.Write([]byte("BCDEF"))
+
+		op, r, err := rc.NextReader()
+		if op != BinaryMessage || err != nil {
+			t.Fatalf("1: NextReader() returned %d, %v", op, err)
+		}
+
+		var buf [10]byte
+		var read int
+		n, err := r.Read(buf[:])
+		if err != nil && err != ErrReadLimit {
+			t.Fatalf("unexpected error testing read limit: %v", err)
+		}
+		read += n
+
+		n, err = r.Read(buf[:])
+		if err != nil && err != ErrReadLimit {
+			t.Fatalf("unexpected error testing read limit: %v", err)
+		}
+		read += n
+
+		if err == nil && read > readLimit {
+			t.Fatalf("read limit exceeded: limit %d, read %d", readLimit, read)
+		}
+	})
+}
+
+func TestAddrs(t *testing.T) {
+	c := newTestConn(nil, nil, true)
+	if c.LocalAddr() != localAddr {
+		t.Errorf("LocalAddr = %v, want %v", c.LocalAddr(), localAddr)
+	}
+	if c.RemoteAddr() != remoteAddr {
+		t.Errorf("RemoteAddr = %v, want %v", c.RemoteAddr(), remoteAddr)
+	}
+}
+
+func TestNetConn(t *testing.T) {
+	var b1, b2 bytes.Buffer
+	fc := fakeNetConn{Reader: &b1, Writer: &b2}
+	c := newConn(fc, true, 1024, 1024, nil, nil, nil)
+	ul := c.NetConn()
+	if ul != fc {
+		t.Fatalf("Underlying conn is not what it should be.")
+	}
+}
+
+func TestBufioReadBytes(t *testing.T) {
+	// Test calling bufio.ReadBytes for value longer than read buffer size.
+
+	m := make([]byte, 512)
+	m[len(m)-1] = '\n'
+
+	var b1, b2 bytes.Buffer
+	wc := newConn(fakeNetConn{Writer: &b1}, false, len(m)+64, len(m)+64, nil, nil, nil)
+	rc := newConn(fakeNetConn{Reader: &b1, Writer: &b2}, true, len(m)-64, len(m)-64, nil, nil, nil)
+
+	w, _ := wc.NextWriter(BinaryMessage)
+	_, _ = w.Write(m)
+	_ = w.Close()
+
+	op, r, err := rc.NextReader()
+	if op != BinaryMessage || err != nil {
+		t.Fatalf("NextReader() returned %d, %v", op, err)
+	}
+
+	br := bufio.NewReader(r)
+	p, err := br.ReadBytes('\n')
+	if err != nil {
+		t.Fatalf("ReadBytes() returned %v", err)
+	}
+	if len(p) != len(m) {
+		t.Fatalf("read returned %d bytes, want %d bytes", len(p), len(m))
+	}
+}
+
+var closeErrorTests = []struct {
+	err   error
+	codes []int
+	ok    bool
+}{
+	{&CloseError{Code: CloseNormalClosure}, []int{CloseNormalClosure}, true},
+	{&CloseError{Code: CloseNormalClosure}, []int{CloseNoStatusReceived}, false},
+	{&CloseError{Code: CloseNormalClosure}, []int{CloseNoStatusReceived, CloseNormalClosure}, true},
+	{errors.New("hello"), []int{CloseNormalClosure}, false},
+}
+
+func TestCloseError(t *testing.T) {
+	for _, tt := range closeErrorTests {
+		ok := IsCloseError(tt.err, tt.codes...)
+		if ok != tt.ok {
+			t.Errorf("IsCloseError(%#v, %#v) returned %v, want %v", tt.err, tt.codes, ok, tt.ok)
+		}
+	}
+}
+
+var unexpectedCloseErrorTests = []struct {
+	err   error
+	codes []int
+	ok    bool
+}{
+	{&CloseError{Code: CloseNormalClosure}, []int{CloseNormalClosure}, false},
+	{&CloseError{Code: CloseNormalClosure}, []int{CloseNoStatusReceived}, true},
+	{&CloseError{Code: CloseNormalClosure}, []int{CloseNoStatusReceived, CloseNormalClosure}, false},
+	{errors.New("hello"), []int{CloseNormalClosure}, false},
+}
+
+func TestUnexpectedCloseErrors(t *testing.T) {
+	for _, tt := range unexpectedCloseErrorTests {
+		ok := IsUnexpectedCloseError(tt.err, tt.codes...)
+		if ok != tt.ok {
+			t.Errorf("IsUnexpectedCloseError(%#v, %#v) returned %v, want %v", tt.err, tt.codes, ok, tt.ok)
+		}
+	}
+}

internal/websocket/convert_safe.go

@@ -0,0 +1,7 @@
+//go:build appengine
+
+package websocket
+
+func stringToBytes(s string) []byte {
+	return []byte(s)
+}

internal/websocket/convert_unsafe.go

@@ -0,0 +1,10 @@
+//go:build !appengine
+
+package websocket
+
+import "unsafe"
+
+// StringToBytes converts string to byte slice.
+func stringToBytes(s string) []byte {
+	return unsafe.Slice(unsafe.StringData(s), len(s)) //nolint:gosec // Audited.
+}

internal/websocket/convert_unsafe_test.go

@@ -0,0 +1,17 @@
+package websocket
+
+import "testing"
+
+var tests = []string{
+	"abc",
+	"hello world",
+	"tests",
+}
+
+func TestStringToBytes(t *testing.T) {
+	for _, want := range tests {
+		if got := stringToBytes(want); string(got) != want {
+			t.Errorf("StringToBytes() = %s, want %s", got, want)
+		}
+	}
+}

internal/websocket/doc.go

@@ -0,0 +1,6 @@
+// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package websocket implements the WebSocket protocol defined in RFC 6455.
+package websocket

internal/websocket/http2.go

@@ -0,0 +1,70 @@
+package websocket
+
+import (
+	"io"
+	"net"
+	"net/http"
+	"time"
+)
+
+var (
+	_ net.Conn = (*http2Stream)(nil)
+)
+
+// http2Stream is a wrapper for HTTP/2 extended CONNECT tunnel.
+type http2Stream struct {
+	io.ReadCloser
+	io.Writer
+	rc *http.ResponseController
+}
+
+func (s *http2Stream) Read(p []byte) (int, error) {
+	return s.ReadCloser.Read(p)
+}
+
+func (s *http2Stream) Write(p []byte) (int, error) {
+	n, err := s.Writer.Write(p)
+	if err != nil {
+		return n, err
+	}
+	err = s.rc.Flush()
+	return n, err
+}
+
+func (s *http2Stream) Flush() error {
+	return s.rc.Flush()
+}
+
+func (s *http2Stream) Close() error {
+	return s.ReadCloser.Close()
+}
+
+// LocalAddr is not implemented for HTTP/2 streams.
+// May be taken from request if needed.
+func (s *http2Stream) LocalAddr() net.Addr {
+	return &net.TCPAddr{}
+}
+
+// RemoteAddr is not implemented for HTTP/2 streams.
+// May be taken from request if needed.
+func (s *http2Stream) RemoteAddr() net.Addr {
+	return &net.TCPAddr{}
+}
+
+// SetDeadline ...
+func (s *http2Stream) SetDeadline(t time.Time) error {
+	if err := s.rc.SetWriteDeadline(t); err != nil {
+		return err
+	}
+	return s.rc.SetReadDeadline(t)
+}
+
+// SetReadDeadline ...
+func (s *http2Stream) SetReadDeadline(t time.Time) error {
+	return s.rc.SetReadDeadline(t)
+}
+
+// SetWriteDeadline ...
+func (s *http2Stream) SetWriteDeadline(t time.Time) error {
+	return s.rc.SetWriteDeadline(t)
+}

internal/websocket/mask.go

@@ -0,0 +1,54 @@
+// Copyright 2016 The Gorilla WebSocket Authors. All rights reserved.  Use of
+// this source code is governed by a BSD-style license that can be found in the
+// LICENSE file.
+
+//go:build !appengine
+
+package websocket
+
+import "unsafe"
+
+const wordSize = int(unsafe.Sizeof(uintptr(0))) //nolint:gosec // Audited.
+
+func maskBytes(key [4]byte, pos int, b []byte) int {
+	// Mask one byte at a time for small buffers.
+	if len(b) < 2*wordSize {
+		for i := range b {
+			b[i] ^= key[pos&3]
+			pos++
+		}
+		return pos & 3
+	}
+
+	// Mask one byte at a time to word boundary.
+	if n := int(uintptr(unsafe.Pointer(&b[0]))) % wordSize; n != 0 { //nolint:gosec // Audited.
+		n = wordSize - n
+		for i := range b[:n] {
+			b[i] ^= key[pos&3]
+			pos++
+		}
+		b = b[n:]
+	}
+
+	// Create aligned word size key.
+	var k [wordSize]byte
+	for i := range k {
+		k[i] = key[(pos+i)&3]
+	}
+	kw := *(*uintptr)(unsafe.Pointer(&k)) //nolint:gosec // Audited.
+
+	// Mask one word at a time.
+	n := (len(b) / wordSize) * wordSize
+	for i := 0; i < n; i += wordSize {
+		*(*uintptr)(unsafe.Pointer(uintptr(unsafe.Pointer(&b[0])) + uintptr(i))) ^= kw //nolint:gosec // Audited.
+	}
+
+	// Mask one byte at a time for remaining bytes.
+	b = b[n:]
+	for i := range b {
+		b[i] ^= key[pos&3]
+		pos++
+	}
+
+	return pos & 3
+}

internal/websocket/mask_safe.go

@@ -0,0 +1,15 @@
+// Copyright 2016 The Gorilla WebSocket Authors. All rights reserved.  Use of
+// this source code is governed by a BSD-style license that can be found in the
+// LICENSE file.
+
+//go:build appengine
+
+package websocket
+
+func maskBytes(key [4]byte, pos int, b []byte) int {
+	for i := range b {
+		b[i] ^= key[pos&3]
+		pos++
+	}
+	return pos & 3
+}

internal/websocket/mask_test.go

@@ -0,0 +1,72 @@
+// Copyright 2016 The Gorilla WebSocket Authors. All rights reserved.  Use of
+// this source code is governed by a BSD-style license that can be found in the
+// LICENSE file.
+
+// !appengine
+
+package websocket
+
+import (
+	"fmt"
+	"testing"
+)
+
+func maskBytesByByte(key [4]byte, pos int, b []byte) int {
+	for i := range b {
+		b[i] ^= key[pos&3]
+		pos++
+	}
+	return pos & 3
+}
+
+func notzero(b []byte) int {
+	for i := range b {
+		if b[i] != 0 {
+			return i
+		}
+	}
+	return -1
+}
+
+func TestMaskBytes(t *testing.T) {
+	key := [4]byte{1, 2, 3, 4}
+	for size := 1; size <= 1024; size++ {
+		for align := 0; align < wordSize; align++ {
+			for pos := 0; pos < 4; pos++ {
+				b := make([]byte, size+align)[align:]
+				maskBytes(key, pos, b)
+				maskBytesByByte(key, pos, b)
+				if i := notzero(b); i >= 0 {
+					t.Errorf("size:%d, align:%d, pos:%d, offset:%d", size, align, pos, i)
+				}
+			}
+		}
+	}
+}
+
+func BenchmarkMaskBytes(b *testing.B) {
+	for _, size := range []int{2, 4, 8, 16, 32, 512, 1024} {
+		b.Run(fmt.Sprintf("size-%d", size), func(b *testing.B) {
+			for _, align := range []int{wordSize / 2} {
+				b.Run(fmt.Sprintf("align-%d", align), func(b *testing.B) {
+					for _, fn := range []struct {
+						name string
+						fn   func(key [4]byte, pos int, b []byte) int
+					}{
+						{"byte", maskBytesByByte},
+						{"word", maskBytes},
+					} {
+						b.Run(fn.name, func(b *testing.B) {
+							key := newMaskKey()
+							data := make([]byte, size+align)[align:]
+							for i := 0; i < b.N; i++ {
+								fn.fn(key, 0, data)
+							}
+							b.SetBytes(int64(len(data)))
+						})
+					}
+				})
+			}
+		})
+	}
+}

internal/websocket/prepared.go

@@ -0,0 +1,102 @@
+// Copyright 2017 The Gorilla WebSocket Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package websocket
+
+import (
+	"bytes"
+	"net"
+	"sync"
+	"time"
+)
+
+// PreparedMessage caches on the wire representations of a message payload.
+// Use PreparedMessage to efficiently send a message payload to multiple
+// connections. PreparedMessage is especially useful when compression is used
+// because the CPU and memory expensive compression operation can be executed
+// once for a given set of compression options.
+type PreparedMessage struct {
+	messageType int
+	data        []byte
+	mu          sync.Mutex
+	frames      map[prepareKey]*preparedFrame
+}
+
+// prepareKey defines a unique set of options to cache prepared frames in PreparedMessage.
+type prepareKey struct {
+	isServer         bool
+	compress         bool
+	compressionLevel int
+}
+
+// preparedFrame contains data in wire representation.
+type preparedFrame struct {
+	once sync.Once
+	data []byte
+}
+
+// NewPreparedMessage returns an initialized PreparedMessage. You can then send
+// it to connection using WritePreparedMessage method. Valid wire
+// representation will be calculated lazily only once for a set of current
+// connection options.
+func NewPreparedMessage(messageType int, data []byte) (*PreparedMessage, error) {
+	pm := &PreparedMessage{
+		messageType: messageType,
+		frames:      make(map[prepareKey]*preparedFrame),
+		data:        data,
+	}
+
+	// Prepare a plain server frame.
+	_, frameData, err := pm.frame(prepareKey{isServer: true, compress: false})
+	if err != nil {
+		return nil, err
+	}
+
+	// To protect against caller modifying the data argument, remember the data
+	// copied to the plain server frame.
+	pm.data = frameData[len(frameData)-len(data):]
+	return pm, nil
+}
+
+func (pm *PreparedMessage) frame(key prepareKey) (int, []byte, error) {
+	pm.mu.Lock()
+	frame, ok := pm.frames[key]
+	if !ok {
+		frame = &preparedFrame{}
+		pm.frames[key] = frame
+	}
+	pm.mu.Unlock()
+
+	var err error
+	frame.once.Do(func() {
+		// Prepare a frame using a 'fake' connection.
+		// TODO: Refactor code in conn.go to allow more direct construction of
+		// the frame.
+		mu := make(chan struct{}, 1)
+		mu <- struct{}{}
+		var nc prepareConn
+		c := &Conn{
+			conn:                   &nc,
+			mu:                     mu,
+			isServer:               key.isServer,
+			compressionLevel:       key.compressionLevel,
+			enableWriteCompression: true,
+			writeBuf:               make([]byte, defaultWriteBufferSize+maxFrameHeaderSize),
+		}
+		if key.compress {
+			c.newCompressionWriter = compressNoContextTakeover
+		}
+		err = c.WriteMessage(pm.messageType, pm.data)
+		frame.data = nc.buf.Bytes()
+	})
+	return pm.messageType, frame.data, err
+}
+
+type prepareConn struct {
+	buf bytes.Buffer
+	net.Conn
+}
+
+func (pc *prepareConn) Write(p []byte) (int, error)        { return pc.buf.Write(p) }
+func (pc *prepareConn) SetWriteDeadline(_ time.Time) error { return nil }

internal/websocket/prepared_test.go

@@ -0,0 +1,62 @@
+// Copyright 2017 The Gorilla WebSocket Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package websocket
+
+import (
+	"bytes"
+	"compress/flate"
+	"testing"
+)
+
+var preparedMessageTests = []struct {
+	messageType            int
+	isServer               bool
+	enableWriteCompression bool
+	compressionLevel       int
+}{
+	// Server
+	{TextMessage, true, false, flate.BestSpeed},
+	{TextMessage, true, true, flate.BestSpeed},
+	{TextMessage, true, true, flate.BestCompression},
+	{PingMessage, true, false, flate.BestSpeed},
+	{PingMessage, true, true, flate.BestSpeed},
+
+	// Client tests will fail due to random masking, skipping them.
+}
+
+func TestPreparedMessage(t *testing.T) {
+	for _, tt := range preparedMessageTests {
+		var data = []byte("this is a test")
+		var buf bytes.Buffer
+		c := newTestConn(nil, &buf, tt.isServer)
+		if tt.enableWriteCompression {
+			c.newCompressionWriter = compressNoContextTakeover
+		}
+		_ = c.SetCompressionLevel(tt.compressionLevel)
+
+		if err := c.WriteMessage(tt.messageType, data); err != nil {
+			t.Fatal(err)
+		}
+		want := buf.String()
+
+		pm, err := NewPreparedMessage(tt.messageType, data)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		// Scribble on data to ensure that NewPreparedMessage takes a snapshot.
+		copy(data, "hello world")
+
+		buf.Reset()
+		if err := c.WritePreparedMessage(pm); err != nil {
+			t.Fatal(err)
+		}
+		got := buf.String()
+
+		if got != want {
+			t.Errorf("write message != prepared message for %+v", tt)
+		}
+	}
+}

internal/websocket/server.go

@@ -0,0 +1,431 @@
+// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package websocket
+
+import (
+	"bufio"
+	"errors"
+	"io"
+	"net/http"
+	"net/url"
+	"strings"
+	"time"
+)
+
+// HandshakeError describes an error with the handshake from the peer.
+type HandshakeError struct {
+	message string
+}
+
+func (e HandshakeError) Error() string { return e.message }
+
+// Upgrader specifies parameters for upgrading an HTTP connection to a
+// WebSocket connection.
+//
+// It is safe to call Upgrader's methods concurrently.
+type Upgrader struct {
+	// HandshakeTimeout specifies the duration for the handshake to complete.
+	HandshakeTimeout time.Duration
+
+	// ReadBufferSize and WriteBufferSize specify I/O buffer sizes in bytes. If a buffer
+	// size is zero, then buffers allocated by the HTTP server are used. The
+	// I/O buffer sizes do not limit the size of the messages that can be sent
+	// or received.
+	// The default value is 4096 bytes, 4kb.
+	// For HTTP/2 connections via Extended Connect ReadBufferSize is ignored.
+	ReadBufferSize, WriteBufferSize int
+
+	// WriteBufferPool is a pool of buffers for write operations. If the value
+	// is not set, then write buffers are allocated to the connection for the
+	// lifetime of the connection.
+	//
+	// A pool is most useful when the application has a modest volume of writes
+	// across a large number of connections.
+	//
+	// Applications should use a single pool for each unique value of
+	// WriteBufferSize.
+	WriteBufferPool BufferPool
+
+	// Subprotocols specifies the server's supported protocols in order of
+	// preference. If this field is not nil, then the Upgrade method negotiates a
+	// subprotocol by selecting the first match in this list with a protocol
+	// requested by the client. If there's no match, then no protocol is
+	// negotiated (the Sec-Websocket-Protocol header is not included in the
+	// handshake response).
+	Subprotocols []string
+
+	// Error specifies the function for generating HTTP error responses. If Error
+	// is nil, then http.Error is used to generate the HTTP response.
+	Error func(w http.ResponseWriter, r *http.Request, status int, reason error)
+
+	// CheckOrigin returns true if the request Origin header is acceptable. If
+	// CheckOrigin is nil, then a safe default is used: return false if the
+	// Origin request header is present and the origin host is not equal to
+	// request Host header.
+	//
+	// A CheckOrigin function should carefully validate the request origin to
+	// prevent cross-site request forgery.
+	CheckOrigin func(r *http.Request) bool
+
+	// EnableCompression specify if the server should attempt to negotiate per
+	// message compression (RFC 7692). Setting this value to true does not
+	// guarantee that compression will be supported. Currently only "no context
+	// takeover" modes are supported.
+	EnableCompression bool
+
+	// DisableHTTP1Upgrade disables support for HTTP/1.1 Upgrade WebSocket handshakes.
+	// When true, server only accepts WebSocket connections over HTTP/2 Extended Connect
+	// (for now requires GODEBUG=http2xconnect=1).
+	// Experimental: This feature is experimental.
+	DisableHTTP1Upgrade bool
+}
+
+func (u *Upgrader) returnError(w http.ResponseWriter, r *http.Request, status int, reason string) (*Conn, string, error) {
+	err := HandshakeError{reason}
+	if u.Error != nil {
+		u.Error(w, r, status, err)
+	} else {
+		w.Header().Set("Sec-Websocket-Version", "13")
+		http.Error(w, http.StatusText(status), status)
+	}
+	return nil, "", err
+}
+
+// checkSameOrigin returns true if the origin is not set or is equal to the request host.
+func checkSameOrigin(r *http.Request) bool {
+	origin := r.Header["Origin"]
+	if len(origin) == 0 {
+		return true
+	}
+	u, err := url.Parse(origin[0])
+	if err != nil {
+		return false
+	}
+	return equalASCIIFold(u.Host, r.Host)
+}
+
+func (u *Upgrader) selectSubprotocol(r *http.Request, responseHeader http.Header) string {
+	if u.Subprotocols != nil {
+		header := r.Header.Get("Sec-Websocket-Protocol")
+		if header == "" {
+			return ""
+		}
+
+		start := 0
+		for i, c := range header {
+			if c == ',' {
+				proto := strings.TrimSpace(header[start:i])
+				for _, serverProtocol := range u.Subprotocols {
+					if proto == serverProtocol {
+						return proto
+					}
+				}
+				start = i + 1
+			}
+		}
+		// Last protocol (after last comma).
+		if start < len(header) {
+			proto := strings.TrimSpace(header[start:])
+			for _, serverProtocol := range u.Subprotocols {
+				if proto == serverProtocol {
+					return proto
+				}
+			}
+		}
+	} else if responseHeader != nil {
+		return responseHeader.Get("Sec-Websocket-Protocol")
+	}
+	return ""
+}
+
+// Subprotocols returns the subprotocols requested by the client in the
+// Sec-Websocket-Protocol header.
+func Subprotocols(r *http.Request) []string {
+	h := strings.TrimSpace(r.Header.Get("Sec-Websocket-Protocol"))
+	if h == "" {
+		return nil
+	}
+	protocols := strings.Split(h, ",")
+	for i := range protocols {
+		protocols[i] = strings.TrimSpace(protocols[i])
+	}
+	return protocols
+}
+
+// Upgrade upgrades the HTTP server connection to the WebSocket protocol.
+//
+// The responseHeader is included in the response to the client's upgrade
+// request. Use the responseHeader to specify cookies (Set-Cookie). To specify
+// subprotocols supported by the server, set Upgrader.Subprotocols directly.
+//
+// If the upgrade fails, then Upgrade replies to the client with an HTTP error
+// response.
+func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header) (*Conn, string, error) {
+	var challengeKey string
+
+	switch r.ProtoMajor {
+	case 1:
+		if u.DisableHTTP1Upgrade {
+			return u.returnError(w, r, http.StatusBadRequest, "websocket: HTTP/1.1 Upgrade not enabled")
+		}
+
+		const badHandshake = "websocket: the client is not using the websocket protocol: "
+
+		if !tokenListContainsValue(r.Header, "Connection", "upgrade") {
+			return u.returnError(w, r, http.StatusBadRequest, badHandshake+"'upgrade' token not found in 'Connection' header")
+		}
+
+		if !tokenListContainsValue(r.Header, "Upgrade", "websocket") {
+			return u.returnError(w, r, http.StatusBadRequest, badHandshake+"'websocket' token not found in 'Upgrade' header")
+		}
+
+		if r.Method != http.MethodGet {
+			return u.returnError(w, r, http.StatusMethodNotAllowed, badHandshake+"request method is not GET")
+		}
+
+		if !tokenListContainsValue(r.Header, "Sec-Websocket-Version", "13") {
+			return u.returnError(w, r, http.StatusBadRequest, "websocket: unsupported version: 13 not found in 'Sec-Websocket-Version' header")
+		}
+
+		challengeKey = r.Header.Get("Sec-Websocket-Key")
+		if !isValidChallengeKey(challengeKey) {
+			return u.returnError(w, r, http.StatusBadRequest, "websocket: not a websocket handshake: 'Sec-WebSocket-Key' header must be Base64 encoded value of 16-byte in length")
+		}
+	case 2:
+		if r.Header.Get(":protocol") != "websocket" {
+			return u.returnError(w, r, http.StatusBadRequest, "websocket: HTTP/2 Extended CONNECT requires :protocol header with 'websocket' value")
+		}
+		if r.Method != http.MethodConnect {
+			return u.returnError(w, r, http.StatusMethodNotAllowed, "websocket: HTTP/2 handshake request method must be CONNECT")
+		}
+		if !tokenListContainsValue(r.Header, "Sec-Websocket-Version", "13") {
+			return u.returnError(w, r, http.StatusBadRequest, "websocket: unsupported version: 13 not found in 'Sec-Websocket-Version' header")
+		}
+	default:
+		return u.returnError(w, r, http.StatusBadRequest, "websocket: unsupported HTTP protocol version")
+	}
+
+	if _, ok := responseHeader["Sec-Websocket-Extensions"]; ok {
+		return u.returnError(w, r, http.StatusInternalServerError, "websocket: application specific 'Sec-WebSocket-Extensions' headers are unsupported")
+	}
+
+	checkOrigin := u.CheckOrigin
+	if checkOrigin == nil {
+		checkOrigin = checkSameOrigin
+	}
+	if !checkOrigin(r) {
+		return u.returnError(w, r, http.StatusForbidden, "websocket: request origin not allowed by Upgrader.CheckOrigin")
+	}
+
+	subprotocol := u.selectSubprotocol(r, responseHeader)
+
+	// Negotiate PMCE.
+	var compress bool
+	if u.EnableCompression {
+		for _, ext := range parseExtensions(r.Header) {
+			if ext[""] != "permessage-deflate" {
+				continue
+			}
+			compress = true
+			break
+		}
+	}
+
+	if r.ProtoMajor == 2 {
+		// HTTP/2 extended CONNECT (RFC 8441).
+		return u.upgradeH2(w, r, responseHeader, subprotocol, compress)
+	}
+	// HTTP/1.1 Upgrade (RFC 6455).
+	return u.upgradeH1(w, r, responseHeader, challengeKey, subprotocol, compress)
+}
+
+// upgradeH1 handles the HTTP/1.1 Upgrade handshake.
+func (u *Upgrader) upgradeH1(w http.ResponseWriter, r *http.Request, responseHeader http.Header, challengeKey, subprotocol string, compress bool) (*Conn, string, error) {
+	h, ok := w.(http.Hijacker)
+	if !ok {
+		return u.returnError(w, r, http.StatusInternalServerError, "websocket: response does not implement http.Hijacker")
+	}
+
+	netConn, brw, err := h.Hijack()
+	if err != nil {
+		return u.returnError(w, r, http.StatusInternalServerError, err.Error())
+	}
+
+	if brw.Reader.Buffered() > 0 {
+		_ = netConn.Close()
+		return nil, "", errors.New("websocket: client sent data before handshake is complete")
+	}
+
+	var br *bufio.Reader
+	if u.ReadBufferSize == 0 && brw.Reader.Size() > 256 {
+		// Reuse hijacked buffered reader as connection reader.
+		br = brw.Reader
+	}
+
+	buf := bufioWriterBuffer(netConn, brw.Writer)
+
+	var writeBuf []byte
+	if u.WriteBufferPool == nil && u.WriteBufferSize == 0 && len(buf) >= maxFrameHeaderSize+256 {
+		// Reuse hijacked write buffer as connection buffer.
+		writeBuf = buf
+	}
+
+	c := newConn(netConn, true, u.ReadBufferSize, u.WriteBufferSize, u.WriteBufferPool, br, writeBuf)
+
+	if compress {
+		c.newCompressionWriter = compressNoContextTakeover
+		c.newDecompressionReader = decompressNoContextTakeover
+	}
+
+	// Use larger of hijacked buffer and connection write buffer for header.
+	p := buf
+	if len(c.writeBuf) > len(p) {
+		p = c.writeBuf
+	}
+	p = p[:0]
+
+	p = append(p, "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: "...)
+	p = encodeAcceptKey(challengeKey, p)
+
+	//p = append(p, computeAcceptKey(challengeKey)...)
+	p = append(p, "\r\n"...)
+	if subprotocol != "" {
+		p = append(p, "Sec-WebSocket-Protocol: "...)
+		p = append(p, subprotocol...)
+		p = append(p, "\r\n"...)
+	}
+	if compress {
+		p = append(p, "Sec-WebSocket-Extensions: permessage-deflate; server_no_context_takeover; client_no_context_takeover\r\n"...)
+	}
+	for k, vs := range responseHeader {
+		if k == "Sec-Websocket-Protocol" {
+			continue
+		}
+		for _, v := range vs {
+			p = append(p, k...)
+			p = append(p, ": "...)
+			for i := 0; i < len(v); i++ {
+				b := v[i]
+				if b <= 31 {
+					// prevent response splitting.
+					b = ' '
+				}
+				p = append(p, b)
+			}
+			p = append(p, "\r\n"...)
+		}
+	}
+	p = append(p, "\r\n"...)
+
+	// Clear deadlines set by HTTP server.
+	_ = netConn.SetDeadline(time.Time{})
+
+	if u.HandshakeTimeout > 0 {
+		_ = netConn.SetWriteDeadline(time.Now().Add(u.HandshakeTimeout))
+	}
+	if _, err = netConn.Write(p); err != nil {
+		_ = netConn.Close()
+		return nil, "", err
+	}
+	if u.HandshakeTimeout > 0 {
+		_ = netConn.SetWriteDeadline(time.Time{})
+	}
+
+	return c, subprotocol, nil
+}
+
+// upgradeH2 handles the HTTP/2 extended CONNECT handshake.
+func (u *Upgrader) upgradeH2(w http.ResponseWriter, r *http.Request, responseHeader http.Header, subprotocol string, compress bool) (*Conn, string, error) {
+	// https://www.rfc-editor.org/rfc/rfc8441.html:
+	// Implementations using this extended CONNECT to bootstrap WebSockets do not do the processing of
+	// the Sec-WebSocket-Key and Sec-WebSocket-Accept header fields of [RFC6455] as that functionality
+	// has been superseded by the :protocol pseudo-header field.
+
+	if subprotocol != "" {
+		w.Header().Set("Sec-WebSocket-Protocol", subprotocol)
+	}
+
+	if compress {
+		w.Header().Set("Sec-WebSocket-Extensions", "permessage-deflate; server_no_context_takeover; client_no_context_takeover")
+	}
+
+	// Copy additional response headers.
+	for k, vs := range responseHeader {
+		if k == "Sec-Websocket-Protocol" {
+			continue
+		}
+		for _, v := range vs {
+			w.Header().Add(k, v)
+		}
+	}
+
+	// RFC 8441 requires a 2xx response for extended CONNECT.
+	w.WriteHeader(http.StatusOK)
+
+	// Flush the response immediately to complete the extended CONNECT
+	// handshake before we start streaming on the tunnel.
+	rc := http.NewResponseController(w)
+	if err := rc.Flush(); err != nil {
+		return nil, "", err
+	}
+	err := rc.SetReadDeadline(time.Time{})
+	if err != nil {
+		return nil, "", err
+	}
+
+	stream := &http2Stream{
+		ReadCloser: r.Body,
+		Writer:     w,
+		rc:         rc,
+	}
+
+	// HTTP/2 stream already has internal buffering. Make small br to avoid allocating a new
+	// large intermediary buffer inside newConn.
+	// Small reads will be sufficient with 16 bytes buffer, for large reads our intermediary
+	// buffer will be bypassed avoiding any overhead.
+	// This means that for HTTP/2 it's not possible to control the read buffer size
+	// via Upgrader.ReadBufferSize. For write buffers it's better to always use a pool
+	// by setting Upgrader.WriteBufferPool.
+	br := bufio.NewReaderSize(stream, 16)
+	c := newConn(stream, true, u.ReadBufferSize, u.WriteBufferSize, u.WriteBufferPool, br, nil)
+	if compress {
+		c.newCompressionWriter = compressNoContextTakeover
+		c.newDecompressionReader = decompressNoContextTakeover
+	}
+
+	return c, subprotocol, nil
+}
+
+// IsWebSocketUpgrade returns true if the client requested upgrade to the
+// WebSocket protocol.
+func IsWebSocketUpgrade(r *http.Request) bool {
+	return tokenListContainsValue(r.Header, "Connection", "upgrade") &&
+		tokenListContainsValue(r.Header, "Upgrade", "websocket")
+}
+
+// writeHook is an io.Writer that records the last slice passed to it vio
+// io.Writer.Write.
+type writeHook struct {
+	p []byte
+}
+
+func (wh *writeHook) Write(p []byte) (int, error) {
+	wh.p = p
+	return len(p), nil
+}
+
+// bufioWriterBuffer grabs the buffer from a bufio.Writer.
+func bufioWriterBuffer(originalWriter io.Writer, bw *bufio.Writer) []byte {
+	// This code assumes that bufio.Writer.buf[:1] is passed to the
+	// bufio.Writer's underlying writer.
+	var wh writeHook
+	bw.Reset(&wh)
+	_ = bw.WriteByte(0)
+	_ = bw.Flush()
+
+	bw.Reset(originalWriter)
+
+	return wh.p[:cap(wh.p)]
+}

internal/websocket/server_test.go

@@ -0,0 +1,311 @@
+// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package websocket
+
+import (
+	"bufio"
+	"bytes"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"reflect"
+	"strings"
+	"testing"
+	"time"
+)
+
+var subprotocolTests = []struct {
+	h         string
+	protocols []string
+}{
+	{"", nil},
+	{"foo", []string{"foo"}},
+	{"foo,bar", []string{"foo", "bar"}},
+	{"foo, bar", []string{"foo", "bar"}},
+	{" foo, bar", []string{"foo", "bar"}},
+	{" foo, bar ", []string{"foo", "bar"}},
+}
+
+func TestSubprotocols(t *testing.T) {
+	for _, st := range subprotocolTests {
+		r := http.Request{Header: http.Header{"Sec-Websocket-Protocol": {st.h}}}
+		protocols := Subprotocols(&r)
+		if !reflect.DeepEqual(st.protocols, protocols) {
+			t.Errorf("SubProtocols(%q) returned %#v, want %#v", st.h, protocols, st.protocols)
+		}
+	}
+}
+
+var isWebSocketUpgradeTests = []struct {
+	ok bool
+	h  http.Header
+}{
+	{false, http.Header{"Upgrade": {"websocket"}}},
+	{false, http.Header{"Connection": {"upgrade"}}},
+	{true, http.Header{"Connection": {"upgRade"}, "Upgrade": {"WebSocket"}}},
+}
+
+func TestIsWebSocketUpgrade(t *testing.T) {
+	for _, tt := range isWebSocketUpgradeTests {
+		ok := IsWebSocketUpgrade(&http.Request{Header: tt.h})
+		if tt.ok != ok {
+			t.Errorf("IsWebSocketUpgrade(%v) returned %v, want %v", tt.h, ok, tt.ok)
+		}
+	}
+}
+
+func TestSubProtocolSelection(t *testing.T) {
+	upgrader := Upgrader{
+		Subprotocols: []string{"foo", "bar", "baz"},
+	}
+
+	r := http.Request{Header: http.Header{"Sec-Websocket-Protocol": {"foo", "bar"}}}
+	s := upgrader.selectSubprotocol(&r, nil)
+	if s != "foo" {
+		t.Errorf("Upgrader.selectSubprotocol returned %v, want %v", s, "foo")
+	}
+
+	r = http.Request{Header: http.Header{"Sec-Websocket-Protocol": {"bar", "foo"}}}
+	s = upgrader.selectSubprotocol(&r, nil)
+	if s != "bar" {
+		t.Errorf("Upgrader.selectSubprotocol returned %v, want %v", s, "bar")
+	}
+
+	r = http.Request{Header: http.Header{"Sec-Websocket-Protocol": {"baz"}}}
+	s = upgrader.selectSubprotocol(&r, nil)
+	if s != "baz" {
+		t.Errorf("Upgrader.selectSubprotocol returned %v, want %v", s, "baz")
+	}
+
+	r = http.Request{Header: http.Header{"Sec-Websocket-Protocol": {"quux"}}}
+	s = upgrader.selectSubprotocol(&r, nil)
+	if s != "" {
+		t.Errorf("Upgrader.selectSubprotocol returned %v, want %v", s, "empty string")
+	}
+
+	upgrader = Upgrader{
+		Subprotocols: nil,
+	}
+	r = http.Request{Header: http.Header{"Sec-Websocket-Protocol": {"foo"}}}
+	s = upgrader.selectSubprotocol(&r, nil)
+	if s != "" {
+		t.Errorf("Upgrader.selectSubprotocol returned %v, want %v", s, "empty string")
+	}
+}
+
+var checkSameOriginTests = []struct {
+	ok bool
+	r  *http.Request
+}{
+	{false, &http.Request{Host: "example.org", Header: map[string][]string{"Origin": {"https://other.org"}}}},
+	{true, &http.Request{Host: "example.org", Header: map[string][]string{"Origin": {"https://example.org"}}}},
+	{true, &http.Request{Host: "Example.org", Header: map[string][]string{"Origin": {"https://example.org"}}}},
+}
+
+func TestCheckSameOrigin(t *testing.T) {
+	for _, tt := range checkSameOriginTests {
+		ok := checkSameOrigin(tt.r)
+		if tt.ok != ok {
+			t.Errorf("checkSameOrigin(%+v) returned %v, want %v", tt.r, ok, tt.ok)
+		}
+	}
+}
+
+type reuseTestResponseWriter struct {
+	brw *bufio.ReadWriter
+	http.ResponseWriter
+}
+
+func (resp *reuseTestResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return fakeNetConn{strings.NewReader(""), &bytes.Buffer{}}, resp.brw, nil
+}
+
+var bufioReuseTests = []struct {
+	n     int
+	reuse bool
+}{
+	{4096, true},
+	{128, false},
+}
+
+func TestBufioReuse(t *testing.T) {
+	for i, tt := range bufioReuseTests {
+		br := bufio.NewReaderSize(strings.NewReader(""), tt.n)
+		bw := bufio.NewWriterSize(&bytes.Buffer{}, tt.n)
+		resp := &reuseTestResponseWriter{
+			brw: bufio.NewReadWriter(br, bw),
+		}
+
+		request, err := http.NewRequest(http.MethodGet, "http://example.com", nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		request.Header.Set("Upgrade", "websocket")
+		request.Header.Set("Connection", "upgrade")
+		request.Header.Set("Sec-Websocket-Key", "dGhlIHNhbXBsZSBub25jZQ==")
+		request.Header.Set("Sec-Websocket-Version", "13")
+
+		upgrader := Upgrader{}
+		c, _, err := upgrader.Upgrade(resp, request, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if reuse := c.br == br; reuse != tt.reuse {
+			t.Errorf("%d: buffered reader reuse=%v, want %v", i, reuse, tt.reuse)
+		}
+		writeBuf := bufioWriterBuffer(c.NetConn(), bw)
+		if reuse := &c.writeBuf[0] == &writeBuf[0]; reuse != tt.reuse {
+			t.Errorf("%d: write buffer reuse=%v, want %v", i, reuse, tt.reuse)
+		}
+	}
+}
+
+func BenchmarkUpgrade(b *testing.B) {
+	// Create a sample request with necessary WebSocket headers
+	req, err := http.NewRequest("GET", "http://example.com/socket", nil)
+	if err != nil {
+		b.Fatal("Error creating request:", err)
+	}
+	req.Header.Add("Connection", "Upgrade")
+	req.Header.Add("Upgrade", "websocket")
+	req.Header.Add("Sec-Websocket-Key", "dGhlIHNhbXBsZSBub25jZQ==")
+	req.Header.Add("Sec-Websocket-Version", "13")
+
+	// Mock connection and buffer to fulfill the Hijack requirement
+	conn := newMockConn()
+	buf := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
+
+	// Create a response recorder with Hijack capability
+	w := &HijackableResponseRecorder{
+		ResponseRecorder: httptest.NewRecorder(),
+		HijackFunc: func() (net.Conn, *bufio.ReadWriter, error) {
+			return conn, buf, nil
+		},
+	}
+
+	//// Create a response recorder
+	//w := httptest.NewRecorder()
+
+	// Initialize the Upgrader
+	upgrader := &Upgrader{
+		// Configure your Upgrader here
+	}
+
+	// Pre-create the responseHeader map
+	responseHeader := http.Header{}
+
+	b.ResetTimer() // Start the benchmark timer
+	for i := 0; i < b.N; i++ {
+		// Run the Upgrade function
+		_, _, err := upgrader.Upgrade(w, req, responseHeader)
+		if err != nil {
+			b.Error("Upgrade failed:", err)
+		}
+	}
+}
+
+func BenchmarkUpgradeSubprotocol(b *testing.B) {
+	// Create a sample request with necessary WebSocket headers
+	req, err := http.NewRequest("GET", "http://example.com/socket", nil)
+	if err != nil {
+		b.Fatal("Error creating request:", err)
+	}
+	req.Header.Add("Connection", "Upgrade")
+	req.Header.Add("Upgrade", "websocket")
+	req.Header.Add("Sec-Websocket-Key", "dGhlIHNhbXBsZSBub25jZQ==")
+	req.Header.Add("Sec-Websocket-Version", "13")
+	req.Header.Add("Sec-WebSocket-Protocol", "centrifuge-protobuf")
+
+	// Mock connection and buffer to fulfill the Hijack requirement.
+	conn := newMockConn()
+	buf := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
+
+	// Create a response recorder with Hijack capability
+	w := &HijackableResponseRecorder{
+		ResponseRecorder: httptest.NewRecorder(),
+		HijackFunc: func() (net.Conn, *bufio.ReadWriter, error) {
+			return conn, buf, nil
+		},
+	}
+
+	// Initialize the Upgrader
+	upgrader := &Upgrader{
+		// Configure your Upgrader here
+		Subprotocols: []string{"centrifuge-json", "centrifuge-protobuf"},
+	}
+
+	// Pre-create the responseHeader map
+	responseHeader := http.Header{}
+
+	b.ResetTimer() // Start the benchmark timer
+	for i := 0; i < b.N; i++ {
+		// Run the Upgrade function
+		_, sub, err := upgrader.Upgrade(w, req, responseHeader)
+		if err != nil {
+			b.Error("Upgrade failed:", err)
+		}
+		if sub != "centrifuge-protobuf" {
+			b.Errorf("Expected subprotocol 'centrifuge-protobuf', got '%s'", sub)
+		}
+	}
+}
+
+// Mock connection
+func newMockConn() net.Conn {
+	return &mockConn{writeBuf: io.Discard}
+}
+
+type mockConn struct {
+	readBuf  bytes.Buffer
+	writeBuf io.Writer
+}
+
+func (m *mockConn) LocalAddr() net.Addr {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (m *mockConn) RemoteAddr() net.Addr {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (m *mockConn) SetDeadline(t time.Time) error {
+	return nil
+}
+
+func (m *mockConn) SetReadDeadline(t time.Time) error {
+	return nil
+}
+
+func (m *mockConn) SetWriteDeadline(t time.Time) error {
+	return nil
+}
+
+func (m *mockConn) Read(b []byte) (int, error) {
+	return m.readBuf.Read(b)
+}
+
+func (m *mockConn) Write(b []byte) (int, error) {
+	return m.writeBuf.Write(b)
+}
+
+func (m *mockConn) Close() error {
+	return nil
+}
+
+type HijackableResponseRecorder struct {
+	*httptest.ResponseRecorder
+	HijackFunc func() (net.Conn, *bufio.ReadWriter, error)
+}
+
+func (h *HijackableResponseRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if h.HijackFunc != nil {
+		return h.HijackFunc()
+	}
+	return nil, nil, fmt.Errorf("Hijack not implemented")
+}

internal/websocket/tls_handshake.go

@@ -0,0 +1,18 @@
+package websocket
+
+import (
+	"context"
+	"crypto/tls"
+)
+
+func doHandshake(ctx context.Context, tlsConn *tls.Conn, cfg *tls.Config) error {
+	if err := tlsConn.HandshakeContext(ctx); err != nil {
+		return err
+	}
+	if !cfg.InsecureSkipVerify {
+		if err := tlsConn.VerifyHostname(cfg.ServerName); err != nil {
+			return err
+		}
+	}
+	return nil
+}

internal/websocket/util.go

@@ -0,0 +1,319 @@
+// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package websocket
+
+import (
+	"crypto/rand"
+	"crypto/sha1" //nolint:gosec // Used for WebSocket accept key hashing as per RFC 6455. Not security-sensitive.
+	"encoding/base64"
+	"io"
+	"net/http"
+	"strings"
+	"sync"
+	"unicode/utf8"
+)
+
+var keyGUID = []byte("258EAFA5-E914-47DA-95CA-C5AB0DC85B11")
+
+func computeAcceptKey(challengeKey string) string {
+	h := sha1.New() //nolint:gosec // Used for WebSocket accept key hashing as per RFC 6455. Not security-sensitive.
+	h.Write([]byte(challengeKey))
+	h.Write(keyGUID)
+	return base64.StdEncoding.EncodeToString(h.Sum(nil))
+}
+
+var acceptKeyBufferPool = sync.Pool{
+	New: func() interface{} {
+		b := make([]byte, 0, sha1.Size)
+		return &b
+	},
+}
+
+func encodeAcceptKey(challengeKey string, p []byte) []byte {
+	h := sha1.New() //nolint:gosec // Used for WebSocket accept key hashing as per RFC 6455. Not security-sensitive.
+	h.Write(stringToBytes(challengeKey))
+	h.Write(keyGUID)
+
+	bufPtr := acceptKeyBufferPool.Get().(*[]byte)
+	defer acceptKeyBufferPool.Put(bufPtr)
+	*bufPtr = (*bufPtr)[:0]
+	sum := h.Sum(*bufPtr)
+
+	return base64.StdEncoding.AppendEncode(p, sum)
+}
+
+func generateChallengeKey() (string, error) {
+	p := make([]byte, 16)
+	if _, err := io.ReadFull(rand.Reader, p); err != nil {
+		return "", err
+	}
+	return base64.StdEncoding.EncodeToString(p), nil
+}
+
+// Token octets per RFC 2616.
+var isTokenOctet = [256]bool{
+	'!':  true,
+	'#':  true,
+	'$':  true,
+	'%':  true,
+	'&':  true,
+	'\'': true,
+	'*':  true,
+	'+':  true,
+	'-':  true,
+	'.':  true,
+	'0':  true,
+	'1':  true,
+	'2':  true,
+	'3':  true,
+	'4':  true,
+	'5':  true,
+	'6':  true,
+	'7':  true,
+	'8':  true,
+	'9':  true,
+	'A':  true,
+	'B':  true,
+	'C':  true,
+	'D':  true,
+	'E':  true,
+	'F':  true,
+	'G':  true,
+	'H':  true,
+	'I':  true,
+	'J':  true,
+	'K':  true,
+	'L':  true,
+	'M':  true,
+	'N':  true,
+	'O':  true,
+	'P':  true,
+	'Q':  true,
+	'R':  true,
+	'S':  true,
+	'T':  true,
+	'U':  true,
+	'W':  true,
+	'V':  true,
+	'X':  true,
+	'Y':  true,
+	'Z':  true,
+	'^':  true,
+	'_':  true,
+	'`':  true,
+	'a':  true,
+	'b':  true,
+	'c':  true,
+	'd':  true,
+	'e':  true,
+	'f':  true,
+	'g':  true,
+	'h':  true,
+	'i':  true,
+	'j':  true,
+	'k':  true,
+	'l':  true,
+	'm':  true,
+	'n':  true,
+	'o':  true,
+	'p':  true,
+	'q':  true,
+	'r':  true,
+	's':  true,
+	't':  true,
+	'u':  true,
+	'v':  true,
+	'w':  true,
+	'x':  true,
+	'y':  true,
+	'z':  true,
+	'|':  true,
+	'~':  true,
+}
+
+// skipSpace returns a slice of the string s with all leading RFC 2616 linear
+// whitespace removed.
+func skipSpace(s string) (rest string) {
+	i := 0
+	for ; i < len(s); i++ {
+		if b := s[i]; b != ' ' && b != '\t' {
+			break
+		}
+	}
+	return s[i:]
+}
+
+// nextToken returns the leading RFC 2616 token of s and the string following
+// the token.
+func nextToken(s string) (token, rest string) {
+	i := 0
+	for ; i < len(s); i++ {
+		if !isTokenOctet[s[i]] {
+			break
+		}
+	}
+	return s[:i], s[i:]
+}
+
+// nextTokenOrQuoted returns the leading token or quoted string per RFC 2616
+// and the string following the token or quoted string.
+func nextTokenOrQuoted(s string) (value string, rest string) {
+	if !strings.HasPrefix(s, "\"") {
+		return nextToken(s)
+	}
+	s = s[1:]
+	for i := 0; i < len(s); i++ {
+		switch s[i] {
+		case '"':
+			return s[:i], s[i+1:]
+		case '\\':
+			p := make([]byte, len(s)-1)
+			j := copy(p, s[:i])
+			escape := true
+			for i = i + 1; i < len(s); i++ {
+				b := s[i]
+				switch {
+				case escape:
+					escape = false
+					p[j] = b
+					j++
+				case b == '\\':
+					escape = true
+				case b == '"':
+					return string(p[:j]), s[i+1:]
+				default:
+					p[j] = b
+					j++
+				}
+			}
+			return "", ""
+		}
+	}
+	return "", ""
+}
+
+// equalASCIIFold returns true if s is equal to t with ASCII case folding as
+// defined in RFC 4790.
+func equalASCIIFold(s, t string) bool {
+	for s != "" && t != "" {
+		sr, size := utf8.DecodeRuneInString(s)
+		s = s[size:]
+		tr, size := utf8.DecodeRuneInString(t)
+		t = t[size:]
+		if sr == tr {
+			continue
+		}
+		if 'A' <= sr && sr <= 'Z' {
+			sr = sr + 'a' - 'A'
+		}
+		if 'A' <= tr && tr <= 'Z' {
+			tr = tr + 'a' - 'A'
+		}
+		if sr != tr {
+			return false
+		}
+	}
+	return s == t
+}
+
+// tokenListContainsValue returns true if the 1#token header with the given
+// name contains a token equal to value with ASCII case folding.
+func tokenListContainsValue(header http.Header, name string, value string) bool {
+headers:
+	for _, s := range header[name] {
+		for {
+			var t string
+			t, s = nextToken(skipSpace(s))
+			if t == "" {
+				continue headers
+			}
+			s = skipSpace(s)
+			if s != "" && s[0] != ',' {
+				continue headers
+			}
+			if equalASCIIFold(t, value) {
+				return true
+			}
+			if s == "" {
+				continue headers
+			}
+			s = s[1:]
+		}
+	}
+	return false
+}
+
+// parseExtensions parses WebSocket extensions from a header.
+func parseExtensions(header http.Header) []map[string]string {
+	// From RFC 6455:
+	//
+	//  Sec-WebSocket-Extensions = extension-list
+	//  extension-list = 1#extension
+	//  extension = extension-token *( ";" extension-param )
+	//  extension-token = registered-token
+	//  registered-token = token
+	//  extension-param = token [ "=" (token | quoted-string) ]
+	//     ;When using the quoted-string syntax variant, the value
+	//     ;after quoted-string unescaping MUST conform to the
+	//     ;'token' ABNF.
+
+	var result []map[string]string
+headers:
+	for _, s := range header["Sec-Websocket-Extensions"] {
+		for {
+			var t string
+			t, s = nextToken(skipSpace(s))
+			if t == "" {
+				continue headers
+			}
+			ext := map[string]string{"": t}
+			for {
+				s = skipSpace(s)
+				if !strings.HasPrefix(s, ";") {
+					break
+				}
+				var k string
+				k, s = nextToken(skipSpace(s[1:]))
+				if k == "" {
+					continue headers
+				}
+				s = skipSpace(s)
+				var v string
+				if strings.HasPrefix(s, "=") {
+					v, s = nextTokenOrQuoted(skipSpace(s[1:]))
+					s = skipSpace(s)
+				}
+				if s != "" && s[0] != ',' && s[0] != ';' {
+					continue headers
+				}
+				ext[k] = v
+			}
+			if s != "" && s[0] != ',' {
+				continue headers
+			}
+			result = append(result, ext)
+			if s == "" {
+				continue headers
+			}
+			s = s[1:]
+		}
+	}
+	return result
+}
+
+// isValidChallengeKey checks if the argument meets RFC6455 specification.
+func isValidChallengeKey(s string) bool {
+	// From RFC6455:
+	//
+	// A |Sec-WebSocket-Key| header field with a base64-encoded (see
+	// Section 4 of [RFC4648]) value that, when decoded, is 16 bytes in
+	// length.
+	if len(s) != 24 { // 16 bytes should always be 24 characters long when base64 encoded.
+		return false
+	}
+	buf := make([]byte, 16) // Pre-allocated buffer for 16 bytes.
+	n, err := base64.StdEncoding.Decode(buf, []byte(s))
+	return err == nil && n == 16
+}

internal/websocket/util_test.go

@@ -0,0 +1,153 @@
+// Copyright 2014 The Gorilla WebSocket Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package websocket
+
+import (
+	"crypto/sha1"
+	"net/http"
+	"reflect"
+	"strings"
+	"testing"
+)
+
+var equalASCIIFoldTests = []struct {
+	t, s string
+	eq   bool
+}{
+	{"WebSocket", "websocket", true},
+	{"websocket", "WebSocket", true},
+	{"Öyster", "öyster", false},
+	{"WebSocket", "WetSocket", false},
+}
+
+func TestEqualASCIIFold(t *testing.T) {
+	for _, tt := range equalASCIIFoldTests {
+		eq := equalASCIIFold(tt.s, tt.t)
+		if eq != tt.eq {
+			t.Errorf("equalASCIIFold(%q, %q) = %v, want %v", tt.s, tt.t, eq, tt.eq)
+		}
+	}
+}
+
+var tokenListContainsValueTests = []struct {
+	value string
+	ok    bool
+}{
+	{"WebSocket", true},
+	{"WEBSOCKET", true},
+	{"websocket", true},
+	{"websockets", false},
+	{"x websocket", false},
+	{"websocket x", false},
+	{"other,websocket,more", true},
+	{"other, websocket, more", true},
+}
+
+func TestTokenListContainsValue(t *testing.T) {
+	for _, tt := range tokenListContainsValueTests {
+		h := http.Header{"Upgrade": {tt.value}}
+		ok := tokenListContainsValue(h, "Upgrade", "websocket")
+		if ok != tt.ok {
+			t.Errorf("tokenListContainsValue(h, n, %q) = %v, want %v", tt.value, ok, tt.ok)
+		}
+	}
+}
+
+var isValidChallengeKeyTests = []struct {
+	key string
+	ok  bool
+}{
+	{"dGhlIHNhbXBsZSBub25jZQ==", true},
+	{"", false},
+	{"InvalidKey", false},
+	{"WHQ4eXhscUtKYjBvOGN3WEdtOEQ=", false},
+}
+
+func TestIsValidChallengeKey(t *testing.T) {
+	for _, tt := range isValidChallengeKeyTests {
+		ok := isValidChallengeKey(tt.key)
+		if ok != tt.ok {
+			t.Errorf("isValidChallengeKey returns %v, want %v", ok, tt.ok)
+		}
+	}
+}
+
+var parseExtensionTests = []struct {
+	value      string
+	extensions []map[string]string
+}{
+	{`foo`, []map[string]string{{"": "foo"}}},
+	{`foo, bar; baz=2`, []map[string]string{
+		{"": "foo"},
+		{"": "bar", "baz": "2"}}},
+	{`foo; bar="b,a;z"`, []map[string]string{
+		{"": "foo", "bar": "b,a;z"}}},
+	{`foo , bar; baz = 2`, []map[string]string{
+		{"": "foo"},
+		{"": "bar", "baz": "2"}}},
+	{`foo, bar; baz=2 junk`, []map[string]string{
+		{"": "foo"}}},
+	{`foo junk, bar; baz=2 junk`, nil},
+	{`mux; max-channels=4; flow-control, deflate-stream`, []map[string]string{
+		{"": "mux", "max-channels": "4", "flow-control": ""},
+		{"": "deflate-stream"}}},
+	{`permessage-foo; x="10"`, []map[string]string{
+		{"": "permessage-foo", "x": "10"}}},
+	{`permessage-foo; use_y, permessage-foo`, []map[string]string{
+		{"": "permessage-foo", "use_y": ""},
+		{"": "permessage-foo"}}},
+	{`permessage-deflate; client_max_window_bits; server_max_window_bits=10 , permessage-deflate; client_max_window_bits`, []map[string]string{
+		{"": "permessage-deflate", "client_max_window_bits": "", "server_max_window_bits": "10"},
+		{"": "permessage-deflate", "client_max_window_bits": ""}}},
+	{"permessage-deflate; server_no_context_takeover; client_max_window_bits=15", []map[string]string{
+		{"": "permessage-deflate", "server_no_context_takeover": "", "client_max_window_bits": "15"},
+	}},
+}
+
+func TestParseExtensions(t *testing.T) {
+	for _, tt := range parseExtensionTests {
+		h := http.Header{http.CanonicalHeaderKey("Sec-WebSocket-Extensions"): {tt.value}}
+		extensions := parseExtensions(h)
+		if !reflect.DeepEqual(extensions, tt.extensions) {
+			t.Errorf("parseExtensions(%q)\n    = %v,\nwant %v", tt.value, extensions, tt.extensions)
+		}
+	}
+}
+
+func TestEncodeAcceptKey(t *testing.T) {
+	challengeKey := "dGhlIHNhbXBsZSBub25jZQ=="
+	expectedKey := computeAcceptKey(challengeKey)
+
+	// Initial encoding.
+	result := encodeAcceptKey(challengeKey, []byte{})
+	if !strings.EqualFold(expectedKey, string(result)) {
+		t.Errorf("Expected %s, got %s", expectedKey, string(result))
+	}
+
+	// Test reuse of pooled buffer to see if it correctly resets and does not grow unexpectedly.
+	result2 := encodeAcceptKey(challengeKey, []byte{})
+	if len(result2) != len(result) {
+		t.Errorf("Buffer reused improperly, expected length %d, got length %d", len(result), len(result2))
+	}
+
+	// Test that we really append to the buffer.
+	result3 := encodeAcceptKey(challengeKey, []byte{0})
+	if result3[0] != 0 {
+		t.Errorf("appended improperly, expected 0, got %d", result3[0])
+	}
+	if string(result3[1:]) != expectedKey {
+		t.Errorf("Expected %s, got %s", expectedKey, string(result3[1:]))
+	}
+
+	// Check if buffer returns the same size after multiple uses
+	for i := 0; i < 10; i++ {
+		_ = encodeAcceptKey(challengeKey, []byte{})
+	}
+	bufPtr := acceptKeyBufferPool.Get().(*[]byte)
+	if cap(*bufPtr) != sha1.Size {
+		t.Errorf("Expected buffer capacity to be %d, got %d", sha1.Size, cap(*bufPtr))
+	}
+	acceptKeyBufferPool.Put(bufPtr)
+}