From 19810798a7496e095988f646a13f314d7215f1a7 Mon Sep 17 00:00:00 2001 From: Jay Date: Mon, 10 Aug 2026 12:16:23 -0400 Subject: [PATCH] config: add a read limit option to the connection. By default, the read limit is never set by NewConnection. When using the default pool dialer, a read limit of 512KB is set. When the read limit option is set, all sockets passed to NewConnection have that read limit set. --- CONFIG.md | 3 +++ honeybeetest/mocks.go | 6 ++++++ pool.go | 7 +++++++ transport/connection.go | 22 +++++++++++++++++++++- transport/connection_test.go | 33 +++++++++++++++++++++++++++++++++ types/types.go | 1 + 6 files changed, 71 insertions(+), 1 deletion(-) diff --git a/CONFIG.md b/CONFIG.md index 7872f1c..6365f45 100644 --- a/CONFIG.md +++ b/CONFIG.md @@ -16,6 +16,7 @@ All configuration uses option functions applied at construction time. Validation | Connection | `PingInterval` | 20s | `0` | ±10% jitter applied per interval | | Connection | `IncomingBufferSize` | 100 | — | Must be positive | | Connection | `ErrorsBufferSize` | 10 | — | Must be positive | +| Connection | `ReadLimit` | nil | 0 | Does nothing when unset, 0 = unlimited | | Retry | `MaxRetries` | 0 | — | 0 means infinite | | Retry | `InitialDelay` | 1s | — | Must be positive | | Retry | `MaxDelay` | 60s | — | Must be ≥ `InitialDelay` | @@ -61,6 +62,8 @@ These are passed to `transport.NewConnectionConfig`. Import `git.wisehodl.dev/ja **`WithErrorsBufferSize(int)`** — sets the capacity of the channel carrying connection-level errors. Must be at least 1. +**`WithReadLimit(int64)`** — sets the limit on incoming socket messages. Does nothing when unset. Incoming message size is unlimited when set to zero. The default pool dialer sets a read limit of 512KB. Setting this option overrides that default. + ## Retry Options These are passed to `transport.NewRetryConfig`. Import `git.wisehodl.dev/jay/go-honeybee/transport`. The resulting `RetryConfig` is then passed to `NewSessionConfig` via `WithRetryConfig`. diff --git a/honeybeetest/mocks.go b/honeybeetest/mocks.go index 17c0ce6..5bb0210 100644 --- a/honeybeetest/mocks.go +++ b/honeybeetest/mocks.go @@ -45,6 +45,7 @@ type MockSocket struct { CloseFunc func() error SetCloseHandlerFunc func(func(int, string) error) SetPongHandlerFunc func(func(string) error) + SetReadLimitFunc func(limit int64) Closed chan struct{} Once sync.Once Mu sync.Mutex @@ -61,6 +62,7 @@ func NewMockSocket() *MockSocket { SetWriteDeadlineFunc: func(time.Time) error { return nil }, SetCloseHandlerFunc: func(func(int, string) error) {}, SetPongHandlerFunc: func(func(string) error) {}, + SetReadLimitFunc: func(limit int64) {}, Closed: make(chan struct{}), } @@ -99,6 +101,10 @@ func (m *MockSocket) SetPongHandler(h func(s string) error) { m.SetPongHandlerFunc(h) } +func (m *MockSocket) SetReadLimit(limit int64) { + m.SetReadLimitFunc(limit) +} + // ---------------------------------------------------------------------------- // Logging mocks // ---------------------------------------------------------------------------- diff --git a/pool.go b/pool.go index 20798c0..b030f0e 100644 --- a/pool.go +++ b/pool.go @@ -349,6 +349,13 @@ func (p *Pool) Connect(url string, opts ...ConnectOption) error { } dialFn := func(ctx context.Context) (types.Socket, error) { socket, _, err := p.dialer.DialContext(ctx, url, hdr) + + // Set a sane read limit when using the default dialer + // Can be overridden with ConnectionConfig.ReadLimit + if socket != nil { + socket.SetReadLimit(512 * 1024) + } + return socket, err } diff --git a/transport/connection.go b/transport/connection.go index d6ac6d9..4f81da6 100644 --- a/transport/connection.go +++ b/transport/connection.go @@ -51,6 +51,7 @@ type ConnectionConfig struct { PingInterval time.Duration IncomingBufferSize int ErrorsBufferSize int + ReadLimit *int64 } func NewConnectionConfig(opts ...ConnectionOption) (*ConnectionConfig, error) { @@ -59,6 +60,7 @@ func NewConnectionConfig(opts ...ConnectionOption) (*ConnectionConfig, error) { PingInterval: 20 * time.Second, IncomingBufferSize: 100, ErrorsBufferSize: 10, + ReadLimit: nil, } for _, o := range opts { o(cfg) @@ -97,6 +99,14 @@ func WithErrorsBufferSize(value int) ConnectionOption { } } +// When ReadLimit is set to zero, reads are unlimited. +func WithReadLimit(value int64) ConnectionOption { + return func(c *ConnectionConfig) { + ptr := &value + c.ReadLimit = ptr + } +} + func ValidateConnectionConfig(c ConnectionConfig) error { if c.WriteTimeout < 0 { return fmt.Errorf("invalid write timeout: %v", c.WriteTimeout) @@ -112,6 +122,9 @@ func ValidateConnectionConfig(c ConnectionConfig) error { return fmt.Errorf("invalid errors buffer size: %d", c.ErrorsBufferSize) } + if c.ReadLimit != nil && *c.ReadLimit < 0 { + return fmt.Errorf("invalid read limit: %d", *c.ReadLimit) + } return nil } @@ -146,7 +159,10 @@ type Connection struct { } func NewConnection( - ctx context.Context, socket types.Socket, config *ConnectionConfig, handler slog.Handler, + ctx context.Context, + socket types.Socket, + config *ConnectionConfig, + handler slog.Handler, ) (*Connection, error) { if socket == nil { return nil, ErrNilSocket @@ -179,6 +195,10 @@ func NewConnection( done: make(chan struct{}), } + if config.ReadLimit != nil { + conn.socket.SetReadLimit(*config.ReadLimit) + } + if handler != nil { comp := component.FromContext(ctx) conn.logger = slog.New(handler).With(slog.Any("component", comp)) diff --git a/transport/connection_test.go b/transport/connection_test.go index 12f40d2..404bf0e 100644 --- a/transport/connection_test.go +++ b/transport/connection_test.go @@ -108,6 +108,39 @@ func TestNewConnection(t *testing.T) { } } +func TestNewConnection_ReadLimit(t *testing.T) { + socket, _, _ := honeybeetest.SetupTestSocket(t) + + var limitWasSet bool + var setLimit int64 + socket.SetReadLimitFunc = func(limit int64) { + limitWasSet = true + setLimit = limit + } + + // when ReadLimit is nil, SetReadLimit should not be called + config, _ := NewConnectionConfig() + NewConnection(context.Background(), socket, config, nil) + assert.False(t, limitWasSet) + + // when ReadLimit is set to zero, SetReadLimit should be called + config, _ = NewConnectionConfig( + WithReadLimit(0), + ) + NewConnection(context.Background(), socket, config, nil) + assert.True(t, limitWasSet) + assert.Equal(t, int64(0), setLimit) + + // when ReadLimit is set to a positive value, SetReadLimit should be called + limitWasSet = false + config, _ = NewConnectionConfig( + WithReadLimit(100), + ) + NewConnection(context.Background(), socket, config, nil) + assert.True(t, limitWasSet) + assert.Equal(t, int64(100), setLimit) +} + // ---------------------------------------------------------------------------- // Accessors // ---------------------------------------------------------------------------- diff --git a/types/types.go b/types/types.go index 8618b37..fa5d5ba 100644 --- a/types/types.go +++ b/types/types.go @@ -23,6 +23,7 @@ type Socket interface { SetWriteDeadline(t time.Time) error SetCloseHandler(h func(code int, text string) error) SetPongHandler(h func(appData string) error) + SetReadLimit(limit int64) } type ReceivedMessage struct {