drpc: fix concurrent large message corruption in multiplexed streams

With stream multiplexing, multiple streams write concurrently to a shared
buffer. The stream's rawWriteLocked used to split large messages into
multiple frames (via SplitData) and call WriteFrame for each chunk. Each
WriteFrame acquires the shared mutex independently, so frames from
different streams can interleave in the buffer. The reader on the other
side doesn't handle interleaved frames from different streams mid-packet.
When it sees a frame from a different stream, it resets the partial
packet, silently corrupting data for messages larger than SplitSize.

The fix changes the StreamWriter interface from WriteFrame(Frame) to
WritePacket(Packet). The stream hands off the full message data in a
single call, and the writer serializes it atomically under one mutex
hold. rawWriteLocked no longer splits messages into frames, so there is
nothing to interleave.

Splitting may have been useful before multiplexing. The manageWriter
goroutine already batches all pending data from the shared buffer into a
single transport write, so splitting at the stream level adds no value.
If we ever need to limit per-write size, that belongs in the writer
implementation, not in the stream's rawWrite path.

Author: shubhamdhama

drpcmanager/frame_queue_test.go

@@ -7,21 +7,19 @@ import (
 	"testing"
 
 	"github.com/zeebo/assert"
-
 	"storj.io/drpc/drpcwire"
 )
 
 func TestSharedWriteBuf_AppendDrain(t *testing.T) {
 	sw := newSharedWriteBuf()
 
-	fr := drpcwire.Frame{
+	pkt := drpcwire.Packet{
 		Data: []byte("hello"),
 		ID:   drpcwire.ID{Stream: 1, Message: 2},
 		Kind: drpcwire.KindMessage,
-		Done: true,
 	}
 
-	assert.NoError(t, sw.Append(fr))
+	assert.NoError(t, sw.Append(pkt))
 
 	// Drain should return serialized bytes.
 	data := sw.Drain(nil)
@@ -31,11 +29,11 @@ func TestSharedWriteBuf_AppendDrain(t *testing.T) {
 	_, got, ok, err := drpcwire.ParseFrame(data)
 	assert.NoError(t, err)
 	assert.That(t, ok)
-	assert.DeepEqual(t, got.Data, fr.Data)
-	assert.Equal(t, got.ID.Stream, fr.ID.Stream)
-	assert.Equal(t, got.ID.Message, fr.ID.Message)
-	assert.Equal(t, got.Kind, fr.Kind)
-	assert.Equal(t, got.Done, fr.Done)
+	assert.DeepEqual(t, got.Data, pkt.Data)
+	assert.Equal(t, got.ID.Stream, pkt.ID.Stream)
+	assert.Equal(t, got.ID.Message, pkt.ID.Message)
+	assert.Equal(t, got.Kind, pkt.Kind)
+	assert.Equal(t, got.Done, true)
 }
 
 func TestSharedWriteBuf_CloseIdempotent(t *testing.T) {
@@ -48,7 +46,7 @@ func TestSharedWriteBuf_AppendAfterClose(t *testing.T) {
 	sw := newSharedWriteBuf()
 	sw.Close()
 
-	err := sw.Append(drpcwire.Frame{})
+	err := sw.Append(drpcwire.Packet{})
 	assert.Error(t, err)
 }
 
@@ -64,7 +62,7 @@ func TestSharedWriteBuf_WaitAndDrainBlocks(t *testing.T) {
 	}()
 
 	// Append should wake the blocked WaitAndDrain.
-	assert.NoError(t, sw.Append(drpcwire.Frame{Data: []byte("a")}))
+	assert.NoError(t, sw.Append(drpcwire.Packet{Data: []byte("a")}))
 	<-done
 }
 

drpcmanager/manager_test.go

@@ -15,7 +15,7 @@ import (
 	"github.com/zeebo/assert"
 	grpcmetadata "google.golang.org/grpc/metadata"
 	"storj.io/drpc/drpcmetadata"
-
+	"storj.io/drpc/drpcstream"
 	"storj.io/drpc/drpctest"
 	"storj.io/drpc/drpcwire"
 )
@@ -245,6 +245,86 @@ func TestNewServerStreamUnreadMessageDoesNotBlockOtherStreams(t *testing.T) {
 	defer func() { _ = srvStream2.Close() }()
 }
 
+// TestConcurrentLargeMessages verifies that two streams writing messages larger
+// than SplitSize concurrently do not corrupt each other's data. With the current
+// implementation, rawWriteLocked splits messages into multiple frames and each
+// frame is appended to the shared write buffer independently. Frames from
+// different streams can interleave in the buffer, and the reader resets partial
+// packets when it sees a frame from a different stream, silently corrupting data.
+func TestConcurrentLargeMessages(t *testing.T) {
+	ctx := drpctest.NewTracker(t)
+	defer ctx.Close()
+
+	cconn, sconn := net.Pipe()
+	defer func() { _ = cconn.Close() }()
+	defer func() { _ = sconn.Close() }()
+
+	// Use a small SplitSize to force even small messages to be split into
+	// multiple frames, making interleaving likely.
+	streamOpts := drpcstream.Options{SplitSize: 5}
+
+	cman := NewWithOptions(cconn, Options{Stream: streamOpts})
+	defer func() { _ = cman.Close() }()
+
+	sman := NewWithOptions(sconn, Options{Stream: streamOpts})
+	defer func() { _ = sman.Close() }()
+
+	// Create two client streams and send invoke + message concurrently.
+	stream1, err := cman.NewClientStream(ctx, "rpc-1")
+	assert.NoError(t, err)
+	defer func() { _ = stream1.Close() }()
+
+	stream2, err := cman.NewClientStream(ctx, "rpc-2")
+	assert.NoError(t, err)
+	defer func() { _ = stream2.Close() }()
+
+	msg1 := []byte("AAAAAAAAAAAAAAAAAAAA") // 20 bytes, split into 4 frames of 5 bytes
+	msg2 := []byte("BBBBBBBBBBBBBBBBBBBB") // 20 bytes, split into 4 frames of 5 bytes
+
+	// Send invokes first (these are small, no splitting).
+	assert.NoError(t, stream1.RawWrite(drpcwire.KindInvoke, []byte("rpc-1")))
+	assert.NoError(t, stream2.RawWrite(drpcwire.KindInvoke, []byte("rpc-2")))
+
+	// Accept both server streams before sending messages, so the streams are
+	// registered and the reader can route packets.
+	srvStream1, rpc1, err := sman.NewServerStream(ctx)
+	assert.NoError(t, err)
+	assert.Equal(t, "rpc-1", rpc1)
+	defer func() { _ = srvStream1.Close() }()
+
+	srvStream2, rpc2, err := sman.NewServerStream(ctx)
+	assert.NoError(t, err)
+	assert.Equal(t, "rpc-2", rpc2)
+	defer func() { _ = srvStream2.Close() }()
+
+	// Write messages concurrently from both streams. With SplitSize=5, each
+	// 20-byte message becomes 4 frames. The frames should not interleave.
+	ready := make(chan struct{})
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		<-ready
+		assert.NoError(t, stream1.RawWrite(drpcwire.KindMessage, msg1))
+	}()
+	go func() {
+		defer wg.Done()
+		<-ready
+		assert.NoError(t, stream2.RawWrite(drpcwire.KindMessage, msg2))
+	}()
+	close(ready)
+	wg.Wait()
+
+	// Read from both server streams and verify correctness.
+	got1, err := srvStream1.RawRecv()
+	assert.NoError(t, err)
+	assert.DeepEqual(t, got1, msg1)
+
+	got2, err := srvStream2.RawRecv()
+	assert.NoError(t, err)
+	assert.DeepEqual(t, got2, msg2)
+}
+
 type blockingTransport chan struct{}
 
 func (b blockingTransport) Read(p []byte) (n int, err error)  { <-b; return 0, io.EOF }

drpcmanager/mux_writer.go

@@ -9,23 +9,20 @@ import (
 	"storj.io/drpc/drpcwire"
 )
 
-// muxWriter implements drpcwire.StreamWriter by serializing frame bytes into a
+// muxWriter implements drpcwire.StreamWriter by serializing packet bytes into a
 // shared write buffer. The manageWriter goroutine drains the buffer and writes
 // directly to the transport.
 //
-// Compared to the previous frameQueue approach, this avoids:
-//   - copying frame payload into an intermediate queue slot,
-//   - drpcwire.Writer mutex overhead in the writer goroutine.
-//
-// Frames are serialized (via AppendFrame) into the shared buffer under a
-// short-held mutex. The frame's Data slice is consumed before WriteFrame
-// returns, so callers may safely reuse their buffers afterward.
+// The entire packet is serialized as a single frame (via AppendFrame) under one
+// mutex hold, so frames from concurrent streams never interleave on the wire.
+// The packet's Data slice is consumed (copied) before WritePacket returns, so
+// callers may safely reuse their buffers afterward.
 type muxWriter struct {
 	sw *sharedWriteBuf
 }
 
-func (w *muxWriter) WriteFrame(fr drpcwire.Frame) error {
-	return w.sw.Append(fr)
+func (w *muxWriter) WritePacket(pkt drpcwire.Packet) error {
+	return w.sw.Append(pkt)
 }
 
 // Flush is a no-op because the manageWriter goroutine flushes to the
@@ -50,15 +47,21 @@ func newSharedWriteBuf() *sharedWriteBuf {
 	return sw
 }
 
-// Append serializes fr into the shared buffer. The frame's Data slice is
-// consumed (copied by AppendFrame) before Append returns.
-func (sw *sharedWriteBuf) Append(fr drpcwire.Frame) error {
+// Append serializes pkt as a single frame into the shared buffer. The packet's
+// Data slice is consumed (copied by AppendFrame) before Append returns.
+func (sw *sharedWriteBuf) Append(pkt drpcwire.Packet) error {
 	sw.mu.Lock()
 	if sw.closed {
 		sw.mu.Unlock()
 		return managerClosed.New("enqueue")
 	}
-	sw.buf = drpcwire.AppendFrame(sw.buf, fr)
+	sw.buf = drpcwire.AppendFrame(sw.buf, drpcwire.Frame{
+		Data:    pkt.Data,
+		ID:      pkt.ID,
+		Kind:    pkt.Kind,
+		Control: pkt.Control,
+		Done:    true,
+	})
 	sw.mu.Unlock()
 
 	sw.cond.Signal()

drpcmanager/mux_writer_test.go

@@ -7,49 +7,46 @@ import (
 	"testing"
 
 	"github.com/zeebo/assert"
-
 	"storj.io/drpc/drpcwire"
 )
 
-func TestMuxWriter_WriteFrame(t *testing.T) {
+func TestMuxWriter_WritePacket(t *testing.T) {
 	sw := newSharedWriteBuf()
 	w := &muxWriter{sw: sw}
 
-	fr := drpcwire.Frame{
+	pkt := drpcwire.Packet{
 		Data: []byte("hello"),
 		ID:   drpcwire.ID{Stream: 1, Message: 2},
 		Kind: drpcwire.KindMessage,
-		Done: true,
 	}
 
-	assert.NoError(t, w.WriteFrame(fr))
+	assert.NoError(t, w.WritePacket(pkt))
 
 	data := sw.Drain(nil)
 	_, got, ok, err := drpcwire.ParseFrame(data)
 	assert.NoError(t, err)
 	assert.That(t, ok)
-	assert.DeepEqual(t, got.Data, fr.Data)
-	assert.Equal(t, got.ID.Stream, fr.ID.Stream)
-	assert.Equal(t, got.ID.Message, fr.ID.Message)
-	assert.Equal(t, got.Kind, fr.Kind)
-	assert.Equal(t, got.Done, fr.Done)
+	assert.DeepEqual(t, got.Data, pkt.Data)
+	assert.Equal(t, got.ID.Stream, pkt.ID.Stream)
+	assert.Equal(t, got.ID.Message, pkt.ID.Message)
+	assert.Equal(t, got.Kind, pkt.Kind)
+	assert.Equal(t, got.Done, true)
 }
 
-func TestMuxWriter_WriteFrameIsolatesData(t *testing.T) {
+func TestMuxWriter_WritePacketIsolatesData(t *testing.T) {
 	sw := newSharedWriteBuf()
 	w := &muxWriter{sw: sw}
 
 	data := []byte("hello")
-	fr := drpcwire.Frame{
+	pkt := drpcwire.Packet{
 		Data: data,
 		ID:   drpcwire.ID{Stream: 1, Message: 2},
 		Kind: drpcwire.KindMessage,
-		Done: true,
 	}
 
-	assert.NoError(t, w.WriteFrame(fr))
+	assert.NoError(t, w.WritePacket(pkt))
 
-	// Mutate the original source buffer after WriteFrame.
+	// Mutate the original source buffer after WritePacket.
 	data[0] = 'j'
 
 	// The serialized data in the shared buffer should be unaffected because
@@ -73,11 +70,11 @@ func TestMuxWriter_Empty(t *testing.T) {
 	assert.That(t, w.Empty())
 }
 
-func TestMuxWriter_WriteFrameAfterClose(t *testing.T) {
+func TestMuxWriter_WritePacketAfterClose(t *testing.T) {
 	sw := newSharedWriteBuf()
 	w := &muxWriter{sw: sw}
 	sw.Close()
 
-	err := w.WriteFrame(drpcwire.Frame{})
+	err := w.WritePacket(drpcwire.Packet{})
 	assert.Error(t, err)
 }

drpcstream/stream.go

@@ -11,7 +11,6 @@ import (
 	"sync"
 
 	"github.com/zeebo/errs"
-
 	"storj.io/drpc"
 	"storj.io/drpc/drpcctx"
 	"storj.io/drpc/drpcdebug"
@@ -80,7 +79,9 @@ func New(ctx context.Context, sid uint64, wr drpcwire.StreamWriter) *Stream {
 // stream id and will use the writer to write messages on. It is important use
 // monotonically increasing stream ids within a single transport. The options
 // are used to control details of how the Stream operates.
-func NewWithOptions(ctx context.Context, sid uint64, wr drpcwire.StreamWriter, opts Options) *Stream {
+func NewWithOptions(
+	ctx context.Context, sid uint64, wr drpcwire.StreamWriter, opts Options,
+) *Stream {
 	var task *trace.Task
 	if trace.IsEnabled() {
 		kind, rpc := drpcopts.GetStreamKind(&opts.Internal), drpcopts.GetStreamRPC(&opts.Internal)
@@ -312,26 +313,28 @@ func (s *Stream) checkCancelError(err error) error {
 	return err
 }
 
-// newFrameLocked bumps the internal message id and returns a frame. It must be
+// nextID bumps the internal message id and returns the new ID. It must be
 // called under a mutex.
-func (s *Stream) newFrameLocked(kind drpcwire.Kind) drpcwire.Frame {
+func (s *Stream) nextID() drpcwire.ID {
 	s.id.Message++
-	return drpcwire.Frame{ID: s.id, Kind: kind}
+	return s.id
 }
 
 // sendPacketLocked sends the packet in a single write and flushes. It does not
 // check for any conditions to stop it from writing and is meant for internal
 // stream use to do things like signal errors or closes to the remote side.
 func (s *Stream) sendPacketLocked(kind drpcwire.Kind, control bool, data []byte) (err error) {
-	fr := s.newFrameLocked(kind)
-	fr.Data = data
-	fr.Control = control
-	fr.Done = true
+	pkt := drpcwire.Packet{
+		ID:      s.nextID(),
+		Kind:    kind,
+		Data:    data,
+		Control: control,
+	}
 
 	drpcopts.GetStreamStats(&s.opts.Internal).AddWritten(uint64(len(data)))
-	s.log("SEND", fr.String)
+	s.log("SEND", pkt.String)
 
-	if err := s.wr.WriteFrame(fr); err != nil {
+	if err := s.wr.WritePacket(pkt); err != nil {
 		return errs.Wrap(err)
 	}
 	if err := s.wr.Flush(); err != nil {
@@ -374,29 +377,26 @@ func (s *Stream) RawWrite(kind drpcwire.Kind, data []byte) (err error) {
 // rawWriteLocked does the body of RawWrite assuming the caller is holding the
 // appropriate locks.
 func (s *Stream) rawWriteLocked(kind drpcwire.Kind, data []byte) (err error) {
-	fr := s.newFrameLocked(kind)
-	n := s.opts.SplitSize
-
-	for {
-		switch {
-		case s.sigs.send.IsSet():
-			return s.sigs.send.Err()
-		case s.sigs.term.IsSet():
-			return s.sigs.term.Err()
-		}
+	switch {
+	case s.sigs.send.IsSet():
+		return s.sigs.send.Err()
+	case s.sigs.term.IsSet():
+		return s.sigs.term.Err()
+	}
 
-		fr.Data, data = drpcwire.SplitData(data, n)
-		fr.Done = len(data) == 0
+	pkt := drpcwire.Packet{
+		ID:   s.nextID(),
+		Kind: kind,
+		Data: data,
+	}
 
-		drpcopts.GetStreamStats(&s.opts.Internal).AddWritten(uint64(len(fr.Data)))
-		s.log("SEND", fr.String)
+	drpcopts.GetStreamStats(&s.opts.Internal).AddWritten(uint64(len(data)))
+	s.log("SEND", pkt.String)
 
-		if err := s.wr.WriteFrame(fr); err != nil {
-			return s.checkCancelError(errs.Wrap(err))
-		} else if fr.Done {
-			return nil
-		}
+	if err := s.wr.WritePacket(pkt); err != nil {
+		return s.checkCancelError(errs.Wrap(err))
 	}
+	return nil
 }
 
 // RawFlush flushes any buffers of data.