diff --git a/token/services/utils/json/session/envelope.go b/token/services/utils/json/session/envelope.go index 79aaed6242..fb8477e551 100644 --- a/token/services/utils/json/session/envelope.go +++ b/token/services/utils/json/session/envelope.go @@ -169,6 +169,11 @@ func ReceiveTypedWithTimeout(s *session.S, expectedType string, dst any, d time. func ReceiveTypedWithTimeoutAndMetrics(s *session.S, expectedType string, dst any, d time.Duration, m *EnvelopeMetrics) error { raw, err := s.ReceiveRawWithTimeout(d) if err != nil { + // Count oversized-payload rejections; other transport errors are left as-is. + if errors.Is(err, session.ErrMessageTooLarge) { + m.observeError("too_large") + } + return err } diff --git a/token/services/utils/json/session/envelope_test.go b/token/services/utils/json/session/envelope_test.go index 77df4bf5d9..2bdb6577aa 100644 --- a/token/services/utils/json/session/envelope_test.go +++ b/token/services/utils/json/session/envelope_test.go @@ -292,6 +292,32 @@ func TestReceiveTypedWithTimeout_TypeMismatch(t *testing.T) { require.ErrorIs(t, err, jsession.ErrTypeMismatch) } +func TestReceiveTypedWithTimeout_TooLarge(t *testing.T) { + body, _ := json.Marshal(map[string]string{"name": "bob"}) + raw, _ := json.Marshal(jsession.Envelope{ + Version: jsession.CurrentVersion, + Type: testTypeB, + Body: body, + }) + + mockSession := &mock.Session{} + ch := make(chan *view.Message, 1) + ch <- &view.Message{Payload: raw, Status: int32(view.OK)} + mockSession.ReceiveReturns(ch) + mockSession.InfoReturns(view.SessionInfo{ID: "test"}) + + // Configure a limit below the envelope size so it is rejected before unwrap. + s := utilsession.New(mockSession, t.Context(), jsession.JSONMarshaller{}, utilsession.WithMaxRecvMessageSize(len(raw)-1)) + + var dst struct { + Name string `json:"name"` + } + err := jsession.ReceiveTypedWithTimeout(s, testTypeB, &dst, 1*time.Second) + require.Error(t, err) + require.ErrorIs(t, err, utilsession.ErrMessageTooLarge) + assert.Empty(t, dst.Name, "body must not be deserialized when rejected for size") +} + func TestReceiveTypedWithTimeout_Timeout(t *testing.T) { mockSession := &mock.Session{} ch := make(chan *view.Message) diff --git a/token/services/utils/session/session.go b/token/services/utils/session/session.go index 5feb3cf7c2..53a0d98570 100644 --- a/token/services/utils/session/session.go +++ b/token/services/utils/session/session.go @@ -19,6 +19,14 @@ import ( const DefaultReceiveTimeout = 10 * time.Second +// DefaultMaxRecvMessageSize is the default upper bound, in bytes, on the size +// of a single inbound message accepted by a session. It is generous enough for +// legitimate token transactions (matching common gRPC/Fabric maximums) while +// still bounding the memory a remote party can force the receiver to allocate +// before any deserialization happens. Override per session with +// WithMaxRecvMessageSize. +const DefaultMaxRecvMessageSize = 100 * 1024 * 1024 // 100 MiB + // ErrTimeout is returned when a session receive operation times out. var ErrTimeout = errors.New("session timeout") @@ -28,6 +36,11 @@ var ErrNilMessage = errors.New("received message is nil") // ErrContextDone is returned when the context is done. var ErrContextDone = errors.New("context done") +// ErrMessageTooLarge is returned when an inbound message exceeds the session's +// configured maximum size. The oversized payload is rejected before it is +// deserialized or otherwise processed. +var ErrMessageTooLarge = errors.New("received message exceeds maximum allowed size") + var logger = logging.MustGetLogger() type Session = view.Session @@ -39,15 +52,31 @@ type Marshaller interface { } type S struct { - s Session - ctx context.Context - marshaller Marshaller + s Session + ctx context.Context + marshaller Marshaller + maxRecvSize int +} + +// Option customizes a session created with New. +type Option func(*S) + +// WithMaxRecvMessageSize sets the maximum size, in bytes, of a single inbound +// message. Messages larger than this are rejected with ErrMessageTooLarge +// before any deserialization. A value <= 0 disables the limit. +func WithMaxRecvMessageSize(maxBytes int) Option { + return func(s *S) { s.maxRecvSize = maxBytes } } -func New(s Session, ctx context.Context, m Marshaller) *S { +func New(s Session, ctx context.Context, m Marshaller, opts ...Option) *S { logger.DebugfContext(ctx, "Open session to [%s]", logging.Eval(s.Info)) - return &S{s: s, ctx: ctx, marshaller: m} + sess := &S{s: s, ctx: ctx, marshaller: m, maxRecvSize: DefaultMaxRecvMessageSize} + for _, opt := range opts { + opt(sess) + } + + return sess } func (j *S) Receive(state any) error { @@ -90,6 +119,13 @@ func (j *S) ReceiveRawWithTimeout(d time.Duration) ([]byte, error) { return nil, errors.Errorf("received error from remote [%s]", string(msg.Payload)) } + // Reject oversized payloads before any deserialization so a remote party + // cannot force the receiver to allocate unbounded memory. + if j.maxRecvSize > 0 && len(msg.Payload) > j.maxRecvSize { + logger.ErrorfContext(j.ctx, "received message of size [%d] exceeds maximum [%d] on session [%s]", len(msg.Payload), j.maxRecvSize, j.Info().ID) + + return nil, errors.Join(errors.Errorf("message size [%d] exceeds maximum [%d]", len(msg.Payload), j.maxRecvSize), ErrMessageTooLarge) + } logger.DebugfContext(j.ctx, "session, received message [%s]", logging.SHA256Base64(msg.Payload)) return msg.Payload, nil diff --git a/token/services/utils/session/session_test.go b/token/services/utils/session/session_test.go index cee15256a3..fa6ac585ae 100644 --- a/token/services/utils/session/session_test.go +++ b/token/services/utils/session/session_test.go @@ -96,6 +96,68 @@ func TestReceiveRawWithTimeout_ErrorMessage(t *testing.T) { assert.Contains(t, err.Error(), "remote error") } +func TestReceiveRawWithTimeout_TooLarge(t *testing.T) { + mockSession := &mock.Session{} + ch := make(chan *view.Message, 1) + ch <- &view.Message{Payload: make([]byte, 100), Status: int32(view.OK)} + mockSession.ReceiveReturns(ch) + mockSession.InfoReturns(view.SessionInfo{ID: "test-session"}) + + m := &mock.Marshaller{} + sess := utilsession.New(mockSession, t.Context(), m, utilsession.WithMaxRecvMessageSize(10)) + + _, err := sess.ReceiveRawWithTimeout(1 * time.Second) + require.Error(t, err) + assert.ErrorIs(t, err, utilsession.ErrMessageTooLarge) +} + +func TestReceiveRawWithTimeout_AtLimit(t *testing.T) { + mockSession := &mock.Session{} + ch := make(chan *view.Message, 1) + payload := make([]byte, 10) + ch <- &view.Message{Payload: payload, Status: int32(view.OK)} + mockSession.ReceiveReturns(ch) + + m := &mock.Marshaller{} + sess := utilsession.New(mockSession, t.Context(), m, utilsession.WithMaxRecvMessageSize(10)) + + raw, err := sess.ReceiveRawWithTimeout(1 * time.Second) + require.NoError(t, err) + assert.Len(t, raw, 10) +} + +func TestReceiveRawWithTimeout_UnlimitedWhenZero(t *testing.T) { + mockSession := &mock.Session{} + ch := make(chan *view.Message, 1) + ch <- &view.Message{Payload: make([]byte, 1024), Status: int32(view.OK)} + mockSession.ReceiveReturns(ch) + + m := &mock.Marshaller{} + sess := utilsession.New(mockSession, t.Context(), m, utilsession.WithMaxRecvMessageSize(0)) + + raw, err := sess.ReceiveRawWithTimeout(1 * time.Second) + require.NoError(t, err) + assert.Len(t, raw, 1024) +} + +// TestReceiveWithTimeout_TooLargeRejectedBeforeUnmarshal asserts the oversized +// payload is dropped before any deserialization is attempted. +func TestReceiveWithTimeout_TooLargeRejectedBeforeUnmarshal(t *testing.T) { + mockSession := &mock.Session{} + ch := make(chan *view.Message, 1) + ch <- &view.Message{Payload: make([]byte, 100), Status: int32(view.OK)} + mockSession.ReceiveReturns(ch) + mockSession.InfoReturns(view.SessionInfo{ID: "test-session"}) + + m := &mock.Marshaller{} + sess := utilsession.New(mockSession, t.Context(), m, utilsession.WithMaxRecvMessageSize(10)) + + var result string + err := sess.ReceiveWithTimeout(&result, 1*time.Second) + require.ErrorIs(t, err, utilsession.ErrMessageTooLarge) + assert.Equal(t, 0, m.UnmarshalCallCount(), "payload must be rejected before deserialization") +} + func TestSend_Success(t *testing.T) { mockSession := &mock.Session{} mockSession.SendWithContextReturns(nil)