From 43358a2ef6817101683d88b3c356802280b66d10 Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Wed, 1 Jul 2026 21:36:59 +0530 Subject: [PATCH] drpcmanager: cap concurrent streams with MaxStreams Add Options.MaxStreams, a per-connection cap on how many concurrent streams a peer may open against this manager. When the cap is reached an inbound stream is refused in the read path -- before any per-stream state is allocated -- with a stream-level KindError, leaving the connection up. 0 means unlimited. This bounds the per-stream bookkeeping (goroutines, stream-table entries, window counters) that the byte-based flow-control windows do not, giving DoS resistance against a peer that opens many zero-byte streams. It is the third flow-control bound alongside the per-stream and connection windows, and is independent of the credit path. Co-Authored-By: roachdev-claude --- drpcmanager/manager.go | 34 ++++++++++++++++ drpcmanager/max_streams_test.go | 69 +++++++++++++++++++++++++++++++++ 2 files changed, 103 insertions(+) create mode 100644 drpcmanager/max_streams_test.go diff --git a/drpcmanager/manager.go b/drpcmanager/manager.go index c602e0e..232830b 100644 --- a/drpcmanager/manager.go +++ b/drpcmanager/manager.go @@ -25,6 +25,10 @@ import ( var managerClosed = errs.Class("manager closed") +// maxStreamsExceeded is sent to a peer whose inbound stream is refused for +// exceeding the 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. @@ -56,6 +60,11 @@ type Options struct { // handling. When enabled, the server stream will decode incoming metadata // into grpc metadata in the context. GRPCMetadataCompatMode bool + + // MaxStreams caps concurrent peer-initiated streams; beyond it, inbound + // streams are refused with a stream-level error rather than admitted. It + // bounds per-stream bookkeeping that byte windows do not. 0 means unlimited. + MaxStreams int } // Manager handles the logic of managing a transport for a drpc client or @@ -294,6 +303,14 @@ func (m *Manager) handleInvokeFrame(fr drpcwire.Frame) error { return nil } + // Refuse in the read path so the stream is never admitted; the connection + // stays up and the peer learns via a stream-level error. + 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}: @@ -308,6 +325,23 @@ func (m *Manager) handleInvokeFrame(fr drpcwire.Frame) error { return nil } +// rejectStream refuses an inbound stream with a stream-level KindError. No +// stream is created. The frame is not marked Control: refusals are unbounded (a +// peer can flood over-cap invokes), so it must respect write-buffer +// backpressure rather than append past MaximumBufferSize. A full buffer parks +// the reader, throttling the flooder; nil cancel lets the error wait for space +// so it still reaches the peer. +func (m *Manager) rejectStream(sid uint64) { + err := maxStreamsExceeded.New("refused: at most %d concurrent streams", m.opts.MaxStreams) + // Message id 1: the peer's 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, + }, nil) +} + // // manage streams // diff --git a/drpcmanager/max_streams_test.go b/drpcmanager/max_streams_test.go new file mode 100644 index 0000000..8ac8dd4 --- /dev/null +++ b/drpcmanager/max_streams_test.go @@ -0,0 +1,69 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcmanager + +import ( + "context" + "net" + "strings" + "testing" + + "github.com/zeebo/assert" + + "storj.io/drpc" + "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", drpc.CompressionNone) + 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", drpc.CompressionNone) + 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() +}