Fix websocket.Conn.Read dropping the tail of messages larger than the read buffer

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 []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 +102,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)
+}