Add inbound message size validation

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
 	}
 

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)

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

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)