Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/handler/security_audit_helper_test.go

@@ -0,0 +1,182 @@
+package handler
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+
+	"github.com/Wei-Shaw/sub2api/internal/securityaudit"
+	middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
+	"go.uber.org/zap/zaptest/observer"
+)
+
+func TestCachesSecurityAuditCompletionSkipsWebSocketStages(t *testing.T) {
+	require.True(t, cachesSecurityAuditCompletion("http"))
+	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) {
+	gin.SetMode(gin.TestMode)
+	engine := &turnCountingEngine{mode: securityaudit.ModeAsync}
+	coordinator := securityaudit.NewCoordinator(nil, engine)
+
+	recorder := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(recorder)
+	c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
+
+	subject := middleware2.AuthSubject{UserID: 7, Concurrency: 1}
+	first := runSecurityAudit(c, nil, coordinator, nil, nil, subject, "openai_responses", "gpt-test",
+		[]byte(`{"type":"response.create","response":{"input":"benign"}}`), "first_turn")
+	require.NotNil(t, first)
+	require.True(t, first.AllowNextStage)
+	require.Equal(t, int64(1), engine.enqueues.Load())
+	_, cached := c.Get(securityAuditCompletedContextKey)
+	require.False(t, cached, "WebSocket stages must not set the HTTP completion cache")
+
+	// Even if an HTTP path previously cached completion on this Context, WS turns
+	// must still audit every response.create payload.
+	c.Set(securityAuditCompletedContextKey, true)
+
+	second := runSecurityAudit(c, nil, coordinator, nil, nil, subject, "openai_responses", "gpt-test",
+		[]byte(`{"type":"response.create","response":{"input":"malicious follow-up"}}`), "subsequent_turn")
+	require.NotNil(t, second)
+	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())
+}
+
+func TestRunSecurityAuditLogsWebSocketChecksAndCacheHits(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	engine := &turnCountingEngine{mode: securityaudit.ModeBlocking}
+	coordinator := securityaudit.NewCoordinator(nil, engine)
+	core, logs := observer.New(zap.InfoLevel)
+	reqLog := zap.New(core)
+	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":"same turn"}}`)
+
+	runSecurityAudit(c, reqLog, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn")
+	runSecurityAudit(c, reqLog, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn")
+
+	startLogs := logs.FilterMessage("security_audit.gateway_check_start").All()
+	require.Len(t, startLogs, 1)
+	require.Equal(t, false, startLogs[0].ContextMap()["cached"])
+
+	doneLogs := logs.FilterMessage("security_audit.gateway_check_done").All()
+	require.Len(t, doneLogs, 2)
+	require.Equal(t, false, doneLogs[0].ContextMap()["cached"])
+	require.Equal(t, true, doneLogs[1].ContextMap()["cached"])
+	require.Equal(t, "allow", doneLogs[1].ContextMap()["decision"])
+	require.Equal(t, "subsequent_turn", doneLogs[1].ContextMap()["stage"])
+	require.Equal(t, int64(1), engine.evaluates.Load())
+}
+
+type turnCountingEngine struct {
+	mode      securityaudit.Mode
+	enqueues  atomic.Int64
+	evaluates atomic.Int64
+	decisions []*securityaudit.PromptDecision
+}
+
+func (e *turnCountingEngine) EffectiveMode() securityaudit.Mode { return e.mode }
+func (e *turnCountingEngine) Enqueue(context.Context, securityaudit.Request) error {
+	e.enqueues.Add(1)
+	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
+}