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.

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