diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index a02f5efa4c..d736c87186 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -2021,6 +2021,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { MaxReasoningEffort: maxReasoningEffort, ReasoningEffortMappings: reasoningEffortMappings, BeforeRequest: func(turn int, payload []byte, originalModel string) error { + c.Set(securityAuditWSTurnContextKey, turn) if turn == 1 { return nil } diff --git a/backend/internal/handler/security_audit_helper.go b/backend/internal/handler/security_audit_helper.go index eab6ab340f..210c84976e 100644 --- a/backend/internal/handler/security_audit_helper.go +++ b/backend/internal/handler/security_audit_helper.go @@ -1,6 +1,7 @@ package handler import ( + "crypto/sha256" "net/http" "strings" @@ -12,6 +13,15 @@ import ( ) const securityAuditCompletedContextKey = "sub2api.security_audit.completed" +const securityAuditWSTurnContextKey = "sub2api.security_audit.ws_turn" +const securityAuditWSDedupeContextKey = "sub2api.security_audit.ws_dedupe" + +type securityAuditWSDedupeEntry struct { + stage string + turn int + bodyHash [sha256.Size]byte + decision securityaudit.Decision +} // cachesSecurityAuditCompletion reports whether a successful audit may be // reused for the rest of the gin request. WebSocket turns share one Context @@ -25,6 +35,15 @@ func cachesSecurityAuditCompletion(stage string) bool { } } +func isSecurityAuditWebSocketStage(stage string) bool { + switch strings.TrimSpace(stage) { + case "first_turn", "subsequent_turn": + return true + default: + return false + } +} + func (h *GatewayHandler) checkSecurityAudit(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol, model string, body []byte) *securityaudit.Decision { if h == nil { return nil @@ -76,6 +95,25 @@ func runSecurityAudit(c *gin.Context, reqLog *zap.Logger, coordinator *securitya return &decision } request := buildSecurityAuditRequest(c, apiKey, subject, protocol, model, body, stage) + if isSecurityAuditWebSocketStage(request.Stage) { + if turnNo, ok := securityAuditWSTurn(c); ok { + bodyHash := sha256.Sum256(body) + if cached, exists := c.Get(securityAuditWSDedupeContextKey); exists { + if entry, ok := cached.(securityAuditWSDedupeEntry); ok && + entry.stage == request.Stage && entry.turn == turnNo && entry.bodyHash == bodyHash { + decision := entry.decision + return &decision + } + } + decision := coordinator.Check(c.Request.Context(), request) + if decision.Kind == securityaudit.DecisionAllow { + c.Set(securityAuditWSDedupeContextKey, securityAuditWSDedupeEntry{ + stage: request.Stage, turn: turnNo, bodyHash: bodyHash, decision: decision, + }) + } + return &decision + } + } if reqLog != nil { reqLog.Info("security_audit.gateway_check_start", zap.String("request_id", request.RequestID), zap.Int64("user_id", request.UserID), @@ -97,6 +135,15 @@ func runSecurityAudit(c *gin.Context, reqLog *zap.Logger, coordinator *securitya return &decision } +func securityAuditWSTurn(c *gin.Context) (int, bool) { + turn, exists := c.Get(securityAuditWSTurnContextKey) + if !exists { + return 0, false + } + turnNo, ok := turn.(int) + return turnNo, ok +} + func buildSecurityAuditRequest(c *gin.Context, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol, model string, body []byte, stage string) securityaudit.Request { legacy := buildContentModerationInput(c, apiKey, subject, protocol, model, body) request := securityaudit.Request{ diff --git a/backend/internal/handler/security_audit_helper_test.go b/backend/internal/handler/security_audit_helper_test.go index fa3e0eeca7..f8a8b3c2b0 100644 --- a/backend/internal/handler/security_audit_helper_test.go +++ b/backend/internal/handler/security_audit_helper_test.go @@ -18,6 +18,9 @@ func TestCachesSecurityAuditCompletionSkipsWebSocketStages(t *testing.T) { require.True(t, cachesSecurityAuditCompletion("")) require.False(t, cachesSecurityAuditCompletion("first_turn")) require.False(t, cachesSecurityAuditCompletion("subsequent_turn")) + require.True(t, isSecurityAuditWebSocketStage("first_turn")) + require.True(t, isSecurityAuditWebSocketStage("subsequent_turn")) + require.False(t, isSecurityAuditWebSocketStage("http")) } func TestRunSecurityAuditDoesNotSkipSubsequentWebSocketTurns(t *testing.T) { @@ -48,9 +51,91 @@ func TestRunSecurityAuditDoesNotSkipSubsequentWebSocketTurns(t *testing.T) { require.Equal(t, int64(2), engine.enqueues.Load(), "subsequent WebSocket turns must be audited again") } +func TestRunSecurityAuditDeduplicatesRepeatedPayloadWithinWebSocketTurn(t *testing.T) { + gin.SetMode(gin.TestMode) + engine := &turnCountingEngine{mode: securityaudit.ModeBlocking} + coordinator := securityaudit.NewCoordinator(nil, engine) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + payload := []byte(`{"type":"response.create","response":{"input":"same turn"}}`) + c.Set(securityAuditWSTurnContextKey, 2) + first := runSecurityAudit(c, nil, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn") + second := runSecurityAudit(c, nil, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn") + require.NotNil(t, first) + require.NotNil(t, second) + require.True(t, first.AllowNextStage) + require.True(t, second.AllowNextStage) + require.Equal(t, int64(1), engine.evaluates.Load()) + + // The cache holds only one successful same-turn result. + entry, exists := c.Get(securityAuditWSDedupeContextKey) + require.True(t, exists) + require.IsType(t, securityAuditWSDedupeEntry{}, entry) + + c.Set(securityAuditWSTurnContextKey, 3) + runSecurityAudit(c, nil, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn") + require.Equal(t, int64(2), engine.evaluates.Load()) +} + +func TestRunSecurityAuditDoesNotCacheFailedWebSocketDecision(t *testing.T) { + gin.SetMode(gin.TestMode) + engine := &turnCountingEngine{ + mode: securityaudit.ModeBlocking, + decisions: []*securityaudit.PromptDecision{ + {Kind: securityaudit.DecisionUnavailable, AllowNextStage: false}, + {Kind: securityaudit.DecisionAllow, AllowNextStage: true}, + }, + } + coordinator := securityaudit.NewCoordinator(nil, engine) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + c.Set(securityAuditWSTurnContextKey, 2) + payload := []byte(`{"type":"response.create","response":{"input":"retry me"}}`) + + first := runSecurityAudit(c, nil, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn") + _, cachedAfterFailure := c.Get(securityAuditWSDedupeContextKey) + second := runSecurityAudit(c, nil, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn") + + require.False(t, first.AllowNextStage) + require.False(t, cachedAfterFailure) + require.True(t, second.AllowNextStage) + require.Equal(t, int64(2), engine.evaluates.Load()) +} + +func TestRunSecurityAuditDoesNotCacheFlaggedWebSocketDecision(t *testing.T) { + gin.SetMode(gin.TestMode) + engine := &turnCountingEngine{ + mode: securityaudit.ModeBlocking, + decisions: []*securityaudit.PromptDecision{ + {Kind: securityaudit.DecisionFlag, AllowNextStage: true}, + {Kind: securityaudit.DecisionAllow, AllowNextStage: true}, + }, + } + coordinator := securityaudit.NewCoordinator(nil, engine) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + c.Set(securityAuditWSTurnContextKey, 2) + payload := []byte(`{"type":"response.create","response":{"input":"retry flagged"}}`) + + first := runSecurityAudit(c, nil, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn") + _, cachedAfterFlag := c.Get(securityAuditWSDedupeContextKey) + second := runSecurityAudit(c, nil, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn") + + require.Equal(t, securityaudit.DecisionFlag, first.Kind) + require.True(t, first.AllowNextStage) + require.False(t, cachedAfterFlag) + require.Equal(t, securityaudit.DecisionAllow, second.Kind) + require.Equal(t, int64(2), engine.evaluates.Load()) +} + type turnCountingEngine struct { - mode securityaudit.Mode - enqueues atomic.Int64 + mode securityaudit.Mode + enqueues atomic.Int64 + evaluates atomic.Int64 + decisions []*securityaudit.PromptDecision } func (e *turnCountingEngine) EffectiveMode() securityaudit.Mode { return e.mode } @@ -59,5 +144,9 @@ func (e *turnCountingEngine) Enqueue(context.Context, securityaudit.Request) err return nil } func (e *turnCountingEngine) Evaluate(context.Context, securityaudit.Request) (*securityaudit.PromptDecision, error) { + call := e.evaluates.Add(1) + if int(call) <= len(e.decisions) { + return e.decisions[call-1], nil + } return &securityaudit.PromptDecision{Kind: securityaudit.DecisionAllow, AllowNextStage: true}, nil }