drpcstream: bound message size against the flow-control window

drpcmanager/flow_control_test.go

@@ -0,0 +1,200 @@
+// Copyright (C) 2026 Cockroach Labs.
+// See LICENSE for copying information.
+
+package drpcmanager
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/zeebo/assert"
+
+	"storj.io/drpc/drpcstream"
+	"storj.io/drpc/drpctest"
+	"storj.io/drpc/drpcwire"
+)
+
+// Enabling flow control through the manager's Stream options installs windows
+// on every stream it creates (client and server), so a message larger than the
+// send window only completes if credit grants flow across the real connection.
+func TestManager_FlowControlEndToEnd(t *testing.T) {
+	ctx := drpctest.NewTracker(t)
+	defer ctx.Close()
+
+	cconn, sconn := net.Pipe()
+	defer func() { _ = cconn.Close() }()
+	defer func() { _ = sconn.Close() }()
+
+	opts := Options{Stream: drpcstream.Options{
+		SplitSize: 64 << 10,
+		FlowControl: drpcstream.FlowControl{
+			Enabled:        true,
+			StreamWindow:   128 << 10, // 2 frames of initial credit
+			HighWater:      1 << 20,
+			GrantThreshold: 64 << 10,
+		},
+	}}
+
+	cman := NewWithOptions(cconn, Client, opts)
+	defer func() { _ = cman.Close() }()
+	sman := NewWithOptions(sconn, Server, opts)
+	defer func() { _ = sman.Close() }()
+
+	// 256 KiB is four frames: two fit in the initial window, the rest can only
+	// be sent as the server dispatches frames and returns credit.
+	msg := make([]byte, 256<<10)
+
+	ctx.Run(func(ctx context.Context) {
+		stream, err := cman.NewClientStream(ctx, "rpc")
+		assert.NoError(t, err)
+		assert.NoError(t, stream.RawWrite(drpcwire.KindInvoke, []byte("rpc")))
+		assert.NoError(t, stream.RawWrite(drpcwire.KindMessage, msg))
+		assert.NoError(t, stream.Close())
+	})
+
+	ctx.Run(func(ctx context.Context) {
+		stream, _, err := sman.NewServerStream(ctx)
+		assert.NoError(t, err)
+		defer func() { _ = stream.Close() }()
+
+		got, err := stream.RawRecv()
+		assert.NoError(t, err)
+		assert.Equal(t, len(got), len(msg))
+
+		_, err = stream.RawRecv()
+		assert.That(t, errors.Is(err, io.EOF))
+	})
+
+	ctx.Wait()
+}
+
+// Cancelling an RPC's context wakes a sender parked on flow-control credit. The
+// client has flow control with a small window; the server does not, so it never
+// returns credit and the client's oversized send parks. This exercises the
+// production wake path: manageStream sees the cancellation and calls
+// Cancel -> terminate -> sendWindow.close.
+func TestManager_FlowControlContextCancelWakesParkedSend(t *testing.T) {
+	tr := drpctest.NewTracker(t)
+	defer tr.Close()
+
+	cconn, sconn := net.Pipe()
+	defer func() { _ = cconn.Close() }()
+	defer func() { _ = sconn.Close() }()
+
+	cman := NewWithOptions(cconn, Client, Options{Stream: drpcstream.Options{
+		SplitSize: 64 << 10,
+		FlowControl: drpcstream.FlowControl{
+			Enabled:        true,
+			StreamWindow:   128 << 10, // two frames of initial credit
+			HighWater:      1 << 20,
+			GrantThreshold: 64 << 10,
+		},
+	}})
+	defer func() { _ = cman.Close() }()
+	// Server has no flow control, so it never returns credit.
+	sman := New(sconn, Server)
+	defer func() { _ = sman.Close() }()
+
+	// Accept the stream and drain; the message never completes (the sender
+	// parks mid-message), so RawRecv just blocks until teardown.
+	tr.Run(func(ctx context.Context) {
+		stream, _, err := sman.NewServerStream(ctx)
+		if err != nil {
+			return
+		}
+		defer func() { _ = stream.Close() }()
+		for {
+			if _, err := stream.RawRecv(); err != nil {
+				return
+			}
+		}
+	})
+
+	streamCtx, cancel := context.WithCancel(context.Background())
+	stream, err := cman.NewClientStream(streamCtx, "rpc")
+	assert.NoError(t, err)
+	assert.NoError(t, stream.RawWrite(drpcwire.KindInvoke, []byte("rpc")))
+
+	// A 512 KiB message exceeds the 128 KiB window; once the initial credit is
+	// spent the send parks waiting for grants that never come.
+	done := make(chan error, 1)
+	go func() { done <- stream.RawWrite(drpcwire.KindMessage, make([]byte, 512<<10)) }()
+
+	select {
+	case <-done:
+		t.Fatal("send returned before it could park on credit")
+	case <-time.After(50 * time.Millisecond):
+	}
+
+	cancel()
+
+	select {
+	case err := <-done:
+		assert.Error(t, err)
+	case <-time.After(time.Second):
+		t.Fatal("parked send did not wake on context cancellation")
+	}
+}
+
+// A peer that does not honor the receiver's implicit maximum message size is
+// caught on receipt: the server's window is small, and a client with a much
+// larger window sends a message that exceeds the server's maximum. The server
+// fails that stream with a data-overflow error while the connection stays up.
+func TestManager_FlowControlReceiverRejectsOversizedMessage(t *testing.T) {
+	ctx := drpctest.NewTracker(t)
+	defer ctx.Close()
+
+	cconn, sconn := net.Pipe()
+	defer func() { _ = cconn.Close() }()
+	defer func() { _ = sconn.Close() }()
+
+	// The client's window is large enough to send the whole message immediately
+	// without waiting for grants and without tripping its own sender-side limit;
+	// the server's window is small, so the message exceeds the server's implicit
+	// maximum (high_water + window) and must be rejected on receipt.
+	cman := NewWithOptions(cconn, Client, Options{Stream: drpcstream.Options{
+		SplitSize: 64 << 10,
+		FlowControl: drpcstream.FlowControl{
+			Enabled: true, StreamWindow: 1 << 20, HighWater: 8 << 20, GrantThreshold: 256 << 10,
+		},
+	}})
+	defer func() { _ = cman.Close() }()
+	sman := NewWithOptions(sconn, Server, Options{Stream: drpcstream.Options{
+		SplitSize: 64 << 10,
+		FlowControl: drpcstream.FlowControl{
+			Enabled: true, StreamWindow: 64 << 10, HighWater: 64 << 10, GrantThreshold: 32 << 10,
+		},
+	}})
+	defer func() { _ = sman.Close() }()
+
+	// 256 KiB exceeds the server's 128 KiB implicit maximum (64 KiB high-water +
+	// 64 KiB window).
+	msg := make([]byte, 256<<10)
+
+	ctx.Run(func(ctx context.Context) {
+		stream, err := cman.NewClientStream(ctx, "rpc")
+		assert.NoError(t, err)
+		assert.NoError(t, stream.RawWrite(drpcwire.KindInvoke, []byte("rpc")))
+		// Best-effort: the send may complete (credit is ample) or fail once the
+		// server's cancel arrives; either way the server must reject the message.
+		_ = stream.RawWrite(drpcwire.KindMessage, msg)
+		_ = stream.Close()
+	})
+
+	ctx.Run(func(ctx context.Context) {
+		stream, _, err := sman.NewServerStream(ctx)
+		assert.NoError(t, err)
+		defer func() { _ = stream.Close() }()
+
+		_, err = stream.RawRecv()
+		assert.Error(t, err)
+		assert.That(t, strings.Contains(err.Error(), "exceeds"))
+	})
+
+	ctx.Wait()
+}

drpcmanager/manager.go

@@ -28,6 +28,10 @@ import (
 
 var managerClosed = errs.Class("manager closed")
 
+// maxStreamsExceeded is returned to a peer whose inbound stream is refused
+// because the manager is already at its concurrent-stream cap.
+var maxStreamsExceeded = errs.Class("too many concurrent streams")
+
 // Options controls configuration settings for a manager.
 type Options struct {
 	// Reader are passed to any readers the manager creates.
@@ -46,6 +50,13 @@ type Options struct {
 	// handling. When enabled, the server stream will decode incoming metadata
 	// into grpc metadata in the context.
 	GRPCMetadataCompatMode bool
+
+	// MaxStreams caps the number of concurrent streams the peer may open against
+	// this manager. When the cap is reached, further inbound streams are refused
+	// with a stream-level error rather than admitted, bounding the per-stream
+	// bookkeeping (goroutines, stream-table entries) that byte windows do not.
+	// 0 means unlimited.
+	MaxStreams int
 }
 
 // Manager handles the logic of managing a transport for a drpc client or
@@ -246,6 +257,16 @@ func (m *Manager) handleInvokeFrame(fr drpcwire.Frame) error {
 		return nil
 	}
 
+	// Enforce the concurrent-stream cap before admitting the stream. Refuse it
+	// here, in the read path, so the per-stream state the cap protects is never
+	// allocated; the peer learns via a stream-level error and the connection
+	// stays up.
+	if m.opts.MaxStreams > 0 && m.streams.Len() >= m.opts.MaxStreams {
+		m.rejectStream(pkt.ID.Stream)
+		delete(m.pendingStreams, fr.ID.Stream)
+		return nil
+	}
+
 	// Invoke packet completes the sequence. Send to NewServerStream.
 	select {
 	case m.invokes <- invokeInfo{sid: pkt.ID.Stream, data: pkt.Data, metadata: ps.metadata}:
@@ -260,6 +281,23 @@ func (m *Manager) handleInvokeFrame(fr drpcwire.Frame) error {
 	return nil
 }
 
+// rejectStream refuses an inbound stream that would exceed MaxStreams by
+// sending a stream-level KindError to the peer. The control bit is set so the
+// frame is appended immediately without blocking the reader goroutine on write
+// backpressure (as with an abortive cancel). No stream is created.
+func (m *Manager) rejectStream(sid uint64) {
+	err := maxStreamsExceeded.New("refused: at most %d concurrent streams", m.opts.MaxStreams)
+	// Message id 1: this is the first (and only) frame the peer receives on the
+	// refused stream, and its receive assembler expects ids to start at 1.
+	_ = m.wr.WriteFrame(drpcwire.Frame{
+		ID:      drpcwire.ID{Stream: sid, Message: 1},
+		Kind:    drpcwire.KindError,
+		Data:    drpcwire.MarshalError(err),
+		Done:    true,
+		Control: true,
+	}, nil)
+}
+
 //
 // manage streams
 //

drpcmanager/max_streams_test.go

@@ -0,0 +1,68 @@
+// Copyright (C) 2026 Cockroach Labs.
+// See LICENSE for copying information.
+
+package drpcmanager
+
+import (
+	"context"
+	"net"
+	"strings"
+	"testing"
+
+	"github.com/zeebo/assert"
+
+	"storj.io/drpc/drpctest"
+	"storj.io/drpc/drpcwire"
+)
+
+// A server with MaxStreams set refuses an inbound stream beyond the cap with a
+// stream-level error, without tearing down the connection or admitting the
+// stream.
+func TestManager_MaxStreamsRefusesExcessInbound(t *testing.T) {
+	ctx := drpctest.NewTracker(t)
+	defer ctx.Close()
+
+	cconn, sconn := net.Pipe()
+	defer func() { _ = cconn.Close() }()
+	defer func() { _ = sconn.Close() }()
+
+	cman := New(cconn, Client)
+	defer func() { _ = cman.Close() }()
+	sman := NewWithOptions(sconn, Server, Options{MaxStreams: 1})
+	defer func() { _ = sman.Close() }()
+
+	accepted := make(chan struct{})
+	release := make(chan struct{})
+
+	// Server accepts exactly one stream and holds it open (so it keeps counting
+	// toward the cap) until released.
+	ctx.Run(func(ctx context.Context) {
+		s1, _, err := sman.NewServerStream(ctx)
+		assert.NoError(t, err)
+		close(accepted)
+		<-release
+		_ = s1.Close()
+	})
+
+	ctx.Run(func(ctx context.Context) {
+		// First stream is admitted.
+		c1, err := cman.NewClientStream(ctx, "rpc")
+		assert.NoError(t, err)
+		assert.NoError(t, c1.RawWrite(drpcwire.KindInvoke, []byte("rpc")))
+		<-accepted // ensure the first stream is active before opening the second
+
+		// Second stream: the server is at its cap, so it is refused with a
+		// stream-level error surfaced on the next receive.
+		c2, err := cman.NewClientStream(ctx, "rpc")
+		assert.NoError(t, err)
+		assert.NoError(t, c2.RawWrite(drpcwire.KindInvoke, []byte("rpc")))
+		_, err = c2.RawRecv()
+		assert.Error(t, err)
+		assert.That(t, strings.Contains(err.Error(), "concurrent streams"))
+
+		close(release)
+		_ = c1.Close()
+	})
+
+	ctx.Wait()
+}