From c9a5a1b4190f6432622067831c26629f641fbdf9 Mon Sep 17 00:00:00 2001 From: Constantine Peresypkin Date: Mon, 14 Sep 2026 12:51:30 -0700 Subject: [PATCH 1/2] TUN: Fix websocket.Conn.Read truncating messages larger than the read buffer Conn.Read returned copy(reader, data) after reading a whole WebSocket message, dropping every byte past len(reader). stream.Pipe feeds it a 16 KiB buffer via cfio.Copy, so client messages larger than 16 KiB reached tcp://, ssh://, rdp://, smb://, bastion and socks ingress origins with only their first 16 KiB. Keep the unread remainder in a bytes.Buffer and drain it on the next Read, as GorillaConn.Read already does. Add tests for a 100 KiB message and for a mix of sizes around the 16 KiB boundary. --- websocket/connection.go | 16 +++++- websocket/connection_test.go | 101 +++++++++++++++++++++++++++++++++++ 2 files changed, 116 insertions(+), 1 deletion(-) create mode 100644 websocket/connection_test.go diff --git a/websocket/connection.go b/websocket/connection.go index 83468b8952d..122f8ceae69 100644 --- a/websocket/connection.go +++ b/websocket/connection.go @@ -81,6 +81,9 @@ func (c *GorillaConn) SetDeadline(t time.Time) error { type Conn struct { rw io.ReadWriter log *zerolog.Logger + // readBuf holds the unread remainder of the last message when it was larger than the buffer passed to Read. + // Without it, everything past len(reader) of a message would be silently dropped. + readBuf bytes.Buffer // writeLock makes sure // 1. Only one write at a time. The pinger and Stream function can both call write. // 2. Close only returns after in progress Write is finished, and no more Write will succeed after calling Close. @@ -99,11 +102,22 @@ func NewConn(ctx context.Context, rw io.ReadWriter, log *zerolog.Logger) *Conn { // Read will read messages from the websocket connection func (c *Conn) Read(reader []byte) (int, error) { + // Intermediate buffer may contain unread bytes from the last read, start there before blocking on a new frame + if c.readBuf.Len() > 0 { + return c.readBuf.Read(reader) + } + data, err := wsutil.ReadClientBinary(c.rw) if err != nil { return 0, err } - return copy(reader, data), nil + + copied := copy(reader, data) + + // Write unread bytes to readBuf; if everything was read this is a no-op + c.readBuf.Write(data[copied:]) + + return copied, nil } // Write will write messages to the websocket connection. diff --git a/websocket/connection_test.go b/websocket/connection_test.go new file mode 100644 index 00000000000..de7a411dcc2 --- /dev/null +++ b/websocket/connection_test.go @@ -0,0 +1,101 @@ +package websocket + +import ( + "bytes" + "context" + "crypto/rand" + "io" + "net" + "testing" + "time" + + "github.com/gobwas/ws/wsutil" + "github.com/rs/zerolog" + "github.com/stretchr/testify/require" +) + +// readThrough reads exactly want bytes from conn using buffers of size bufSize, +// the way cfio.Copy feeds Conn.Read from its fixed-size pool buffer. +func readThrough(t *testing.T, conn *Conn, bufSize, want int) []byte { + t.Helper() + var out bytes.Buffer + buf := make([]byte, bufSize) + for out.Len() < want { + n, err := conn.Read(buf) + require.NoError(t, err) + require.LessOrEqual(t, n, bufSize) + out.Write(buf[:n]) + } + return out.Bytes() +} + +func randomPayload(t *testing.T, size int) []byte { + t.Helper() + payload := make([]byte, size) + _, err := io.ReadFull(rand.Reader, payload) + require.NoError(t, err) + return payload +} + +func newTestConn(t *testing.T) (*Conn, net.Conn) { + t.Helper() + ctx, cancel := context.WithCancel(context.Background()) + server, client := net.Pipe() + log := zerolog.Nop() + // Bound the test: a Read that never sees the rest of a message would otherwise block forever + require.NoError(t, server.SetReadDeadline(time.Now().Add(10*time.Second))) + conn := NewConn(ctx, server, &log) + t.Cleanup(func() { + cancel() + conn.Close() + _ = server.Close() + _ = client.Close() + }) + return conn, client +} + +// TestConnReadMessageLargerThanBuffer verifies that a single client message larger +// than the buffer passed to Read is delivered whole across multiple Read calls, +// rather than truncated to the first len(buffer) bytes. +func TestConnReadMessageLargerThanBuffer(t *testing.T) { + const bufSize = 16 * 1024 + const msgSize = 100 * 1024 + + conn, client := newTestConn(t) + payload := randomPayload(t, msgSize) + + go func() { + _ = wsutil.WriteClientBinary(client, payload) + }() + + got := readThrough(t, conn, bufSize, msgSize) + require.Equal(t, payload, got) +} + +// TestConnReadMixedMessageSizes sends several messages around the buffer boundary +// and verifies all bytes arrive in order. +func TestConnReadMixedMessageSizes(t *testing.T) { + const bufSize = 16 * 1024 + sizes := []int{1, bufSize - 1, bufSize, bufSize + 1, 40000, 3 * bufSize, 100 * 1024, 7, bufSize} + + conn, client := newTestConn(t) + + var expected bytes.Buffer + payloads := make([][]byte, 0, len(sizes)) + for _, size := range sizes { + payload := randomPayload(t, size) + payloads = append(payloads, payload) + expected.Write(payload) + } + + go func() { + for _, payload := range payloads { + if err := wsutil.WriteClientBinary(client, payload); err != nil { + return + } + } + }() + + got := readThrough(t, conn, bufSize, expected.Len()) + require.Equal(t, expected.Bytes(), got) +} From 9675cf7625715257f56a5ff326adbc89ebb06927 Mon Sep 17 00:00:00 2001 From: Constantine Peresypkin Date: Wed, 16 Sep 2026 10:05:44 -0700 Subject: [PATCH 2/2] TUN: Keep the unread message tail as a slice instead of a bytes.Buffer Holding the remainder as a subslice of the message wsutil returned avoids a second copy and releases the memory once the tail is drained, instead of a bytes.Buffer pinning its peak allocation for the life of the connection. The two new tests are parallel-safe, so they call t.Parallel. --- websocket/connection.go | 22 +++++++++++++--------- websocket/connection_test.go | 2 ++ 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/websocket/connection.go b/websocket/connection.go index 122f8ceae69..a97b6c89e98 100644 --- a/websocket/connection.go +++ b/websocket/connection.go @@ -83,7 +83,7 @@ type Conn struct { log *zerolog.Logger // readBuf holds the unread remainder of the last message when it was larger than the buffer passed to Read. // Without it, everything past len(reader) of a message would be silently dropped. - readBuf bytes.Buffer + readBuf []byte // writeLock makes sure // 1. Only one write at a time. The pinger and Stream function can both call write. // 2. Close only returns after in progress Write is finished, and no more Write will succeed after calling Close. @@ -103,8 +103,13 @@ func NewConn(ctx context.Context, rw io.ReadWriter, log *zerolog.Logger) *Conn { // Read will read messages from the websocket connection func (c *Conn) Read(reader []byte) (int, error) { // Intermediate buffer may contain unread bytes from the last read, start there before blocking on a new frame - if c.readBuf.Len() > 0 { - return c.readBuf.Read(reader) + if len(c.readBuf) > 0 { + n := copy(reader, c.readBuf) + c.readBuf = c.readBuf[n:] + if len(c.readBuf) == 0 { + c.readBuf = nil + } + return n, nil } data, err := wsutil.ReadClientBinary(c.rw) @@ -112,12 +117,11 @@ func (c *Conn) Read(reader []byte) (int, error) { return 0, err } - copied := copy(reader, data) - - // Write unread bytes to readBuf; if everything was read this is a no-op - c.readBuf.Write(data[copied:]) - - return copied, nil + n := copy(reader, data) + if n < len(data) { + c.readBuf = data[n:] + } + return n, nil } // Write will write messages to the websocket connection. diff --git a/websocket/connection_test.go b/websocket/connection_test.go index de7a411dcc2..b530c8f6a1b 100644 --- a/websocket/connection_test.go +++ b/websocket/connection_test.go @@ -58,6 +58,7 @@ func newTestConn(t *testing.T) (*Conn, net.Conn) { // than the buffer passed to Read is delivered whole across multiple Read calls, // rather than truncated to the first len(buffer) bytes. func TestConnReadMessageLargerThanBuffer(t *testing.T) { + t.Parallel() const bufSize = 16 * 1024 const msgSize = 100 * 1024 @@ -75,6 +76,7 @@ func TestConnReadMessageLargerThanBuffer(t *testing.T) { // TestConnReadMixedMessageSizes sends several messages around the buffer boundary // and verifies all bytes arrive in order. func TestConnReadMixedMessageSizes(t *testing.T) { + t.Parallel() const bufSize = 16 * 1024 sizes := []int{1, bufSize - 1, bufSize, bufSize + 1, 40000, 3 * bufSize, 100 * 1024, 7, bufSize}