CL64.3-2-JWT-Replay (#19544)

* make jti required and use it to prevent replay attacks. also, make exp and iat required and enforce 5 minute max token age

* emit metrics and add optional field to verify method in JWT

* revert changes in http_trigger_handler and change metrics name

Author: jinhoonbang

core/services/gateway/handlers/capabilities/v2/http_trigger_handler_test.go

@@ -206,14 +206,14 @@ func TestHttpTriggerHandler_HandleUserTriggerRequest(t *testing.T) {
 			Method:  gateway_common.MethodWorkflowExecute,
 			Params:  &rawParams,
 		}
-		req.Auth = createTestJWTToken(t, req, privateKey)
-
 		// First request should succeed
+		req.Auth = createTestJWTToken(t, req, privateKey)
 		mockDon.EXPECT().SendToNode(mock.Anything, mock.Anything, mock.Anything).Return(nil).Times(3)
 		err = handler.HandleUserTriggerRequest(testutils.Context(t), req, callback1, time.Now())
 		require.NoError(t, err)
 
 		// Second request with same ID should fail
+		req.Auth = createTestJWTToken(t, req, privateKey)
 		err = handler.HandleUserTriggerRequest(testutils.Context(t), req, callback2, time.Now())
 		require.Error(t, err)
 		require.Contains(t, err.Error(), "in-flight request")
@@ -223,6 +223,45 @@ func TestHttpTriggerHandler_HandleUserTriggerRequest(t *testing.T) {
 		requireUserErrorSent(t, r, jsonrpc.ErrConflict)
 	})
 
+	t.Run("duplicate JWT token and request ID", func(t *testing.T) {
+		handler, mockDon := createTestTriggerHandler(t)
+		privateKey := createTestPrivateKey(t)
+		registerWorkflow(t, handler, workflowID, privateKey)
+		callback1 := hc.NewCallback()
+		callback2 := hc.NewCallback()
+
+		triggerReq := gateway_common.HTTPTriggerRequest{
+			Workflow: gateway_common.WorkflowSelector{
+				WorkflowID: workflowID,
+			},
+			Input: []byte(`{"key": "value"}`),
+		}
+		reqBytes, err := json.Marshal(triggerReq)
+		require.NoError(t, err)
+
+		rawParams := json.RawMessage(reqBytes)
+		req := &jsonrpc.Request[json.RawMessage]{
+			Version: "2.0",
+			ID:      requestID,
+			Method:  gateway_common.MethodWorkflowExecute,
+			Params:  &rawParams,
+		}
+		// First request should succeed
+		req.Auth = createTestJWTToken(t, req, privateKey)
+		mockDon.EXPECT().SendToNode(mock.Anything, mock.Anything, mock.Anything).Return(nil).Times(3)
+		err = handler.HandleUserTriggerRequest(testutils.Context(t), req, callback1, time.Now())
+		require.NoError(t, err)
+
+		// Second request with same ID should fail
+		err = handler.HandleUserTriggerRequest(testutils.Context(t), req, callback2, time.Now())
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "token has already been used")
+
+		r, err := callback2.Wait(t.Context())
+		require.NoError(t, err)
+		requireUserErrorSent(t, r, jsonrpc.ErrInvalidRequest)
+	})
+
 	t.Run("invalid input JSON", func(t *testing.T) {
 		handler, _ := createTestTriggerHandler(t)
 		callback := hc.NewCallback()
@@ -380,7 +419,6 @@ func TestHttpTriggerHandler_ReapExpiredCallbacks(t *testing.T) {
 		Params:  &rawParams,
 	}
 	privateKey := createTestPrivateKey(t)
-	req.Auth = createTestJWTToken(t, req, privateKey)
 	cfg := ServiceConfig{
 		CleanUpPeriodMs:             100,
 		MaxTriggerRequestDurationMs: 50,
@@ -389,6 +427,7 @@ func TestHttpTriggerHandler_ReapExpiredCallbacks(t *testing.T) {
 	registerWorkflow(t, handler, workflowID, privateKey)
 
 	t.Run("reap expired callbacks", func(t *testing.T) {
+		req.Auth = createTestJWTToken(t, req, privateKey)
 		callback := hc.NewCallback()
 		mockDon.EXPECT().SendToNode(mock.Anything, mock.Anything, mock.Anything).Return(nil).Times(3)
 		err = handler.HandleUserTriggerRequest(testutils.Context(t), req, callback, time.Now())
@@ -413,6 +452,7 @@ func TestHttpTriggerHandler_ReapExpiredCallbacks(t *testing.T) {
 	})
 
 	t.Run("keep non-expired callbacks", func(t *testing.T) {
+		req.Auth = createTestJWTToken(t, req, privateKey)
 		callback := hc.NewCallback()
 
 		mockDon.EXPECT().SendToNode(mock.Anything, mock.Anything, mock.Anything).Return(nil).Times(3)

core/services/gateway/handlers/capabilities/v2/metrics/metrics.go

@@ -56,6 +56,8 @@ type TriggerMetrics struct {
 	metadataRequestCount             metric.Int64Counter
 	metadataObservationsCleanUpCount metric.Int64Counter
 	metadataObservationsCount        metric.Int64Gauge
+	jwtCacheSize                     metric.Int64Gauge
+	jwtCacheCleanUpCount             metric.Int64Counter
 }
 
 // Metrics combines all gateway metrics for dependency injection
@@ -328,6 +330,22 @@ func newTriggerMetrics(meter metric.Meter) (*TriggerMetrics, error) {
 		return nil, fmt.Errorf("failed to create workflow metadata observations count metric: %w", err)
 	}
 
+	m.jwtCacheSize, err = meter.Int64Gauge(
+		"http_trigger_jwt_cache_size",
+		metric.WithDescription("Current number of entries in JWT replay protection cache"),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create HTTP trigger JWT cache size metric: %w", err)
+	}
+
+	m.jwtCacheCleanUpCount, err = meter.Int64Counter(
+		"http_trigger_jwt_cache_cleanup_count",
+		metric.WithDescription("Number of JWT replay protection cache entries cleaned up"),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create HTTP trigger JWT cache cleanup count metric: %w", err)
+	}
+
 	return m, nil
 }
 
@@ -456,3 +474,11 @@ func (m *TriggerMetrics) IncrementMetadataObservationsCleanUpCount(ctx context.C
 func (m *TriggerMetrics) RecordMetadataObservationsCount(ctx context.Context, count int64, lggr logger.Logger) {
 	m.metadataObservationsCount.Record(ctx, count)
 }
+
+func (m *TriggerMetrics) RecordJwtCacheSize(ctx context.Context, size int64, lggr logger.Logger) {
+	m.jwtCacheSize.Record(ctx, size)
+}
+
+func (m *TriggerMetrics) IncrementJwtCacheCleanUpCount(ctx context.Context, count int64, lggr logger.Logger) {
+	m.jwtCacheCleanUpCount.Add(ctx, count)
+}

core/services/gateway/handlers/capabilities/v2/workflow_metadata_handler.go

@@ -28,6 +28,13 @@ type workflowReference struct {
 	workflowTag   string
 }
 
+// jwtReplayCache manages used JWT IDs to prevent replay attacks
+type jwtReplayCache struct {
+	mu            sync.RWMutex
+	cleanupPeriod time.Duration
+	cache         map[string]time.Time // jti -> timestamp
+}
+
 type WorkflowMetadataHandler struct {
 	services.StateMachine
 	lggr            logger.Logger
@@ -41,6 +48,7 @@ type WorkflowMetadataHandler struct {
 	donConfig       *config.DONConfig
 	stopCh          services.StopChan
 	metrics         *metrics.Metrics
+	jwtCache        *jwtReplayCache // JWT replay protection cache
 }
 
 // NewWorkflowMetadataHandler creates a new WorkflowMetadataHandler.
@@ -58,15 +66,22 @@ func NewWorkflowMetadataHandler(lggr logger.Logger, cfg ServiceConfig, don handl
 		config:          cfg,
 		stopCh:          make(services.StopChan),
 		metrics:         metrics,
+		jwtCache:        newJWTReplayCache(time.Duration(cfg.CleanUpPeriodMs) * time.Millisecond),
 	}
 }
 
 func (h *WorkflowMetadataHandler) Authorize(workflowID string, token string, req *jsonrpc.Request[json.RawMessage]) (*gateway.AuthorizedKey, error) {
-	_, signer, err := utils.VerifyRequestJWT(token, *req)
+	claims, signer, err := utils.VerifyRequestJWT(token, *req)
 	if err != nil {
 		h.lggr.Errorw("Failed to verify JWT", "error", err)
 		return nil, err
 	}
+
+	if h.jwtCache.isReplay(claims.ID) {
+		h.lggr.Warnw("JWT token has already been used", "workflowID", workflowID, "signer", signer.Hex(), "jti", claims.ID)
+		return nil, errors.New("JWT token has already been used. Please generate a new one with new id (jti)")
+	}
+
 	keys, exists := h.authorizedKeys[workflowID]
 	if !exists {
 		h.lggr.Errorw("Workflow ID not found in authorized keys", "workflowID", workflowID)
@@ -80,6 +95,8 @@ func (h *WorkflowMetadataHandler) Authorize(workflowID string, token string, req
 		h.lggr.Errorw("Signer not found in authorized keys", "signer", signer.Hex())
 		return nil, errors.New("signer not found in authorized keys")
 	}
+	h.jwtCache.recordUsage(claims.ID)
+
 	return &key, nil
 }
 
@@ -205,6 +222,14 @@ func (h *WorkflowMetadataHandler) Start(ctx context.Context) error {
 			}
 		})
 		h.runTicker(time.Duration(h.config.MetadataAggregationIntervalMs)*time.Millisecond, h.syncMetadata)
+
+		h.runTicker(h.jwtCache.cleanupPeriod, func() {
+			now := time.Now()
+			expiredCount := h.jwtCache.cleanupOldEntries(now.Add(-h.jwtCache.cleanupPeriod))
+			h.metrics.Trigger.IncrementJwtCacheCleanUpCount(context.Background(), int64(expiredCount), h.lggr)
+			h.metrics.Trigger.RecordJwtCacheSize(context.Background(), int64(len(h.jwtCache.cache)), h.lggr)
+			h.lggr.Debugw("Workflow execution cache cleanup completed", "expired_entries", expiredCount, "remaining_entries", len(h.jwtCache.cache))
+		})
 		return nil
 	})
 }
@@ -277,3 +302,39 @@ func (h *WorkflowMetadataHandler) Close() error {
 		return nil
 	})
 }
+
+func newJWTReplayCache(cleanupPeriod time.Duration) *jwtReplayCache {
+	return &jwtReplayCache{
+		cache:         make(map[string]time.Time),
+		cleanupPeriod: cleanupPeriod,
+	}
+}
+
+func (cache *jwtReplayCache) isReplay(jti string) bool {
+	cache.mu.RLock()
+	defer cache.mu.RUnlock()
+
+	_, exists := cache.cache[jti]
+	return exists
+}
+
+func (cache *jwtReplayCache) recordUsage(jti string) {
+	cache.mu.Lock()
+	defer cache.mu.Unlock()
+
+	cache.cache[jti] = time.Now()
+}
+
+// cleanupOldEntries removes expired entries from the cache
+func (cache *jwtReplayCache) cleanupOldEntries(cutoff time.Time) int {
+	cache.mu.Lock()
+	defer cache.mu.Unlock()
+	var expiredCount int
+	for jti, createdAt := range cache.cache {
+		if createdAt.Before(cutoff) {
+			delete(cache.cache, jti)
+			expiredCount++
+		}
+	}
+	return expiredCount
+}

core/services/gateway/handlers/capabilities/v2/workflow_metadata_handler_test.go

@@ -941,6 +941,67 @@ func TestWorkflowMetadataHandler_Authorize(t *testing.T) {
 		require.Contains(t, err.Error(), "JWT digest does not match request digest")
 		require.Nil(t, key)
 	})
+
+	t.Run("JWT replay protection", func(t *testing.T) {
+		params := json.RawMessage(`{"test": "data"}`)
+		req := &jsonrpc.Request[json.RawMessage]{
+			Version: "2.0",
+			ID:      "test-request-id-replay",
+			Method:  gateway_common.MethodWorkflowExecute,
+			Params:  &params,
+		}
+
+		token, err := utils.CreateRequestJWT(*req)
+		require.NoError(t, err)
+
+		tokenString, err := token.SignedString(privateKey)
+		require.NoError(t, err)
+
+		key, err := handler.Authorize(workflowID, tokenString, req)
+		require.NoError(t, err)
+		require.NotNil(t, key)
+
+		// Second authorization with same JWT should fail (replay attack)
+		key, err = handler.Authorize(workflowID, tokenString, req)
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "JWT token has already been used. Please generate a new one with new id (jti)")
+		require.Nil(t, key)
+	})
+
+	t.Run("different JWT IDs should work", func(t *testing.T) {
+		params := json.RawMessage(`{"test": "data"}`)
+		req1 := &jsonrpc.Request[json.RawMessage]{
+			Version: "2.0",
+			ID:      "test-request-id-1",
+			Method:  gateway_common.MethodWorkflowExecute,
+			Params:  &params,
+		}
+
+		req2 := &jsonrpc.Request[json.RawMessage]{
+			Version: "2.0",
+			ID:      "test-request-id-2",
+			Method:  gateway_common.MethodWorkflowExecute,
+			Params:  &params,
+		}
+
+		token1, err := utils.CreateRequestJWT(*req1)
+		require.NoError(t, err)
+		tokenString1, err := token1.SignedString(privateKey)
+		require.NoError(t, err)
+
+		key1, err := handler.Authorize(workflowID, tokenString1, req1)
+		require.NoError(t, err)
+		require.NotNil(t, key1)
+
+		token2, err := utils.CreateRequestJWT(*req2)
+		require.NoError(t, err)
+		tokenString2, err := token2.SignedString(privateKey)
+		require.NoError(t, err)
+
+		key2, err := handler.Authorize(workflowID, tokenString2, req2)
+		require.NoError(t, err)
+		require.NotNil(t, key2)
+	})
 }
 
 func TestWorkflowMetadataHandler_GetWorkflowID(t *testing.T) {