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

@@ -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`.

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
 // ----------------------------------------------------------------------------

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
 	}
 

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))

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
 // ----------------------------------------------------------------------------

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 {