From 9a5e229ccc629774e544d10c8d5a0a13b26328de Mon Sep 17 00:00:00 2001 From: Varun Chawla Date: Tue, 10 Feb 2026 00:13:08 -0800 Subject: [PATCH] Fix Upgrader.Upgrade never reusing hijacked write buffer bufio.Writer.AvailableBuffer() returns a slice with len==0 but cap equal to the available buffer space. The existing code checked len(buf) which was always 0, so the buffer reuse path was dead code. This caused an unnecessary allocation on every Upgrade when the hijacked write buffer was large enough to be reused. Changes: - Use cap(buf) instead of len(buf) to check hijacked buffer size - Extend buf to full capacity with buf[:cap(buf)] before passing as writeBuf so downstream code using len(c.writeBuf) works - Use cap(p) instead of len(p) when selecting larger buffer for the response header construction - Fix typos: "effor" -> "effort" (2x in conn.go), "buferred" -> "buffered" (1x in server.go) - Add TestWriteBufferReuse to verify the buffer reuse behavior Fixes #973 --- conn.go | 4 +-- server.go | 8 +++--- server_test.go | 68 ++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 74 insertions(+), 6 deletions(-) diff --git a/conn.go b/conn.go index 9562ffd4..18c53792 100644 --- a/conn.go +++ b/conn.go @@ -991,7 +991,7 @@ func (c *Conn) handleProtocolError(message string) error { if len(data) > maxControlFramePayloadSize { data = data[:maxControlFramePayloadSize] } - // Make a best effor to send a close message describing the problem. + // Make a best effort to send a close message describing the problem. _ = c.WriteControl(CloseMessage, data, time.Now().Add(writeWait)) return errors.New("websocket: " + message) } @@ -1147,7 +1147,7 @@ func (c *Conn) SetCloseHandler(h func(code int, text string) error) { if h == nil { h = func(code int, text string) error { message := FormatCloseMessage(code, "") - // Make a best effor to send the close message. + // Make a best effort to send the close message. _ = c.WriteControl(CloseMessage, message, time.Now().Add(writeWait)) return nil } diff --git a/server.go b/server.go index 02ea01fd..b3f4a422 100644 --- a/server.go +++ b/server.go @@ -205,9 +205,9 @@ func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeade buf := brw.Writer.AvailableBuffer() var writeBuf []byte - if u.WriteBufferPool == nil && u.WriteBufferSize == 0 && len(buf) >= maxFrameHeaderSize+256 { + if u.WriteBufferPool == nil && u.WriteBufferSize == 0 && cap(buf) >= maxFrameHeaderSize+256 { // Reuse hijacked write buffer as connection buffer. - writeBuf = buf + writeBuf = buf[:cap(buf)] } c := newConn(netConn, true, u.ReadBufferSize, u.WriteBufferSize, u.WriteBufferPool, br, writeBuf) @@ -220,7 +220,7 @@ func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeade // Use larger of hijacked buffer and connection write buffer for header. p := buf - if len(c.writeBuf) > len(p) { + if len(c.writeBuf) > cap(p) { p = c.writeBuf } p = p[:0] @@ -353,7 +353,7 @@ type brNetConn struct { func (b *brNetConn) Read(p []byte) (n int, err error) { if b.br != nil { - // Limit read to buferred data. + // Limit read to buffered data. if n := b.br.Buffered(); len(p) > n { p = p[:n] } diff --git a/server_test.go b/server_test.go index bb5f074e..245ed28b 100644 --- a/server_test.go +++ b/server_test.go @@ -169,3 +169,71 @@ func TestHijack_NotSupported(t *testing.T) { t.Fatalf("got err=%T and status_code=%d", err, recorder.Code) } } + +func TestWriteBufferReuse(t *testing.T) { + // Test that Upgrader.Upgrade correctly reuses the hijacked write buffer + // from bufio.Writer.AvailableBuffer(). AvailableBuffer() returns a slice + // with len==0 and cap equal to the available buffer space, so the check + // must use cap(buf) instead of len(buf). + + for _, tt := range []struct { + name string + bufSize int + wantReuse bool + }{ + {"large enough buffer", 4096, true}, + {"too small buffer", 128, false}, + } { + t.Run(tt.name, func(t *testing.T) { + var writeBuf bytes.Buffer + br := bufio.NewReaderSize(strings.NewReader(""), tt.bufSize) + bw := bufio.NewWriterSize(&writeBuf, tt.bufSize) + + // Get the AvailableBuffer to compare addresses later. + availBuf := bw.AvailableBuffer() + if len(availBuf) != 0 { + t.Fatalf("AvailableBuffer len=%d, want 0", len(availBuf)) + } + if cap(availBuf) != tt.bufSize { + t.Fatalf("AvailableBuffer cap=%d, want %d", cap(availBuf), tt.bufSize) + } + + brw := bufio.NewReadWriter(br, bw) + resp := &reuseTestResponseWriter{ + brw: brw, + ResponseWriter: httptest.NewRecorder(), + } + + upgrader := Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, + } + c, err := upgrader.Upgrade(resp, &http.Request{ + Method: http.MethodGet, + Header: http.Header{ + "Upgrade": []string{"websocket"}, + "Connection": []string{"upgrade"}, + "Sec-Websocket-Key": []string{"dGhlIHNhbXBsZSBub25jZQ=="}, + "Sec-Websocket-Version": []string{"13"}, + }, + }, nil) + if err != nil { + t.Fatalf("Upgrade: %v", err) + } + defer c.Close() + + if tt.wantReuse { + // When the buffer is large enough, the connection write buffer + // should be backed by the same underlying array as the hijacked + // writer's AvailableBuffer. + if cap(availBuf) > 0 && len(c.writeBuf) > 0 && &c.writeBuf[0] != &availBuf[:cap(availBuf)][0] { + t.Error("write buffer was not reused from hijacked connection") + } + } else { + // When the buffer is too small, a new buffer should be allocated. + if cap(availBuf) > 0 && len(c.writeBuf) > 0 && &c.writeBuf[0] == &availBuf[:cap(availBuf)][0] { + t.Error("write buffer was unexpectedly reused from small hijacked buffer") + } + } + }) + } +}