TUN-10994: Fix dropping the tail of messages larger than the buffer

Integrating the changes introduced here: #1741

This fixes silent data loss in TCP/SOCKS-over-WebSocket connections when a WebSocket message exceeds the caller’s read buffer. Unread bytes are now buffered and delivered by subsequent reads, preserving the complete byte stream in order.

Author: macmarcelino

websocket/connection.go

@@ -31,7 +31,6 @@ type PingPeriodContext string
 // This is still used by access carrier
 type GorillaConn struct {
 	*websocket.Conn
-	log     *zerolog.Logger
 	readBuf bytes.Buffer
 }
 
@@ -42,7 +41,7 @@ func (c *GorillaConn) Read(p []byte) (int, error) {
 		return c.readBuf.Read(p)
 	}
 
-	_, message, err := c.Conn.ReadMessage()
+	_, message, err := c.ReadMessage()
 	if err != nil {
 		return 0, err
 	}
@@ -58,7 +57,7 @@ func (c *GorillaConn) Read(p []byte) (int, error) {
 
 // Write will write messages to the websocket connection
 func (c *GorillaConn) Write(p []byte) (int, error) {
-	if err := c.Conn.WriteMessage(websocket.BinaryMessage, p); err != nil {
+	if err := c.WriteMessage(websocket.BinaryMessage, p); err != nil {
 		return 0, err
 	}
 
@@ -69,10 +68,10 @@ func (c *GorillaConn) Write(p []byte) (int, error) {
 // "It is equivalent to calling both SetReadDeadline and SetWriteDeadline."
 // Note there is no synchronization here, but the gorilla implementation isn't thread safe anyway
 func (c *GorillaConn) SetDeadline(t time.Time) error {
-	if err := c.Conn.SetReadDeadline(t); err != nil {
+	if err := c.SetReadDeadline(t); err != nil {
 		return fmt.Errorf("error setting read deadline: %w", err)
 	}
-	if err := c.Conn.SetWriteDeadline(t); err != nil {
+	if err := c.SetWriteDeadline(t); err != nil {
 		return fmt.Errorf("error setting write deadline: %w", err)
 	}
 	return nil
@@ -81,6 +80,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 []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.
@@ -99,11 +101,26 @@ 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 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)
 	if err != nil {
 		return 0, err
 	}
-	return copy(reader, data), nil
+
+	n := copy(reader, data)
+	if n < len(data) {
+		c.readBuf = data[n:]
+	}
+	return n, nil
 }
 
 // Write will write messages to the websocket connection.

websocket/connection_test.go

@@ -0,0 +1,103 @@
+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) {
+	t.Parallel()
+	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) {
+	t.Parallel()
+	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)
+}