Skip to content

Commit d6fedbe

Browse files
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 <roachdev-claude-bot@cockroachlabs.com>
1 parent c6d9f95 commit d6fedbe

2 files changed

Lines changed: 103 additions & 0 deletions

File tree

drpcmanager/manager.go

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,10 @@ import (
2525

2626
var managerClosed = errs.Class("manager closed")
2727

28+
// maxStreamsExceeded is sent to a peer whose inbound stream is refused for
29+
// exceeding the concurrent-stream cap.
30+
var maxStreamsExceeded = errs.Class("too many concurrent streams")
31+
2832
// Options controls configuration settings for a manager.
2933
type Options struct {
3034
// Reader are passed to any readers the manager creates.
@@ -56,6 +60,11 @@ type Options struct {
5660
// handling. When enabled, the server stream will decode incoming metadata
5761
// into grpc metadata in the context.
5862
GRPCMetadataCompatMode bool
63+
64+
// MaxStreams caps concurrent peer-initiated streams; beyond it, inbound
65+
// streams are refused with a stream-level error rather than admitted. It
66+
// bounds per-stream bookkeeping that byte windows do not. 0 means unlimited.
67+
MaxStreams int
5968
}
6069

6170
// 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 {
294303
return nil
295304
}
296305

306+
// Refuse in the read path so the stream is never admitted; the connection
307+
// stays up and the peer learns via a stream-level error.
308+
if m.opts.MaxStreams > 0 && m.streams.Len() >= m.opts.MaxStreams {
309+
m.rejectStream(pkt.ID.Stream)
310+
delete(m.pendingStreams, fr.ID.Stream)
311+
return nil
312+
}
313+
297314
// Invoke packet completes the sequence. Send to NewServerStream.
298315
select {
299316
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 {
308325
return nil
309326
}
310327

328+
// rejectStream refuses an inbound stream with a stream-level KindError. No
329+
// stream is created. The frame is not marked Control: refusals are unbounded (a
330+
// peer can flood over-cap invokes), so it must respect write-buffer
331+
// backpressure rather than append past MaximumBufferSize. A full buffer parks
332+
// the reader, throttling the flooder; nil cancel lets the error wait for space
333+
// so it still reaches the peer.
334+
func (m *Manager) rejectStream(sid uint64) {
335+
err := maxStreamsExceeded.New("refused: at most %d concurrent streams", m.opts.MaxStreams)
336+
// Message id 1: the peer's receive assembler expects ids to start at 1.
337+
_ = m.wr.WriteFrame(drpcwire.Frame{
338+
ID: drpcwire.ID{Stream: sid, Message: 1},
339+
Kind: drpcwire.KindError,
340+
Data: drpcwire.MarshalError(err),
341+
Done: true,
342+
}, nil)
343+
}
344+
311345
//
312346
// manage streams
313347
//

drpcmanager/max_streams_test.go

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
// Copyright (C) 2026 Cockroach Labs.
2+
// See LICENSE for copying information.
3+
4+
package drpcmanager
5+
6+
import (
7+
"context"
8+
"net"
9+
"strings"
10+
"testing"
11+
12+
"github.com/zeebo/assert"
13+
14+
"storj.io/drpc"
15+
"storj.io/drpc/drpctest"
16+
"storj.io/drpc/drpcwire"
17+
)
18+
19+
// A server with MaxStreams set refuses an inbound stream beyond the cap with a
20+
// stream-level error, without tearing down the connection or admitting the
21+
// stream.
22+
func TestManager_MaxStreamsRefusesExcessInbound(t *testing.T) {
23+
ctx := drpctest.NewTracker(t)
24+
defer ctx.Close()
25+
26+
cconn, sconn := net.Pipe()
27+
defer func() { _ = cconn.Close() }()
28+
defer func() { _ = sconn.Close() }()
29+
30+
cman := New(cconn, Client)
31+
defer func() { _ = cman.Close() }()
32+
sman := NewWithOptions(sconn, Server, Options{MaxStreams: 1})
33+
defer func() { _ = sman.Close() }()
34+
35+
accepted := make(chan struct{})
36+
release := make(chan struct{})
37+
38+
// Server accepts exactly one stream and holds it open (so it keeps counting
39+
// toward the cap) until released.
40+
ctx.Run(func(ctx context.Context) {
41+
s1, _, err := sman.NewServerStream(ctx)
42+
assert.NoError(t, err)
43+
close(accepted)
44+
<-release
45+
_ = s1.Close()
46+
})
47+
48+
ctx.Run(func(ctx context.Context) {
49+
// First stream is admitted.
50+
c1, err := cman.NewClientStream(ctx, "rpc", drpc.CompressionNone)
51+
assert.NoError(t, err)
52+
assert.NoError(t, c1.RawWrite(drpcwire.KindInvoke, []byte("rpc")))
53+
<-accepted // ensure the first stream is active before opening the second
54+
55+
// Second stream: the server is at its cap, so it is refused with a
56+
// stream-level error surfaced on the next receive.
57+
c2, err := cman.NewClientStream(ctx, "rpc", drpc.CompressionNone)
58+
assert.NoError(t, err)
59+
assert.NoError(t, c2.RawWrite(drpcwire.KindInvoke, []byte("rpc")))
60+
_, err = c2.RawRecv()
61+
assert.Error(t, err)
62+
assert.That(t, strings.Contains(err.Error(), "concurrent streams"))
63+
64+
close(release)
65+
_ = c1.Close()
66+
})
67+
68+
ctx.Wait()
69+
}

0 commit comments

Comments
 (0)