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

backend/internal/handler/security_audit_order_test.go

@@ -0,0 +1,85 @@
+package handler
+
+import (
+	"go/ast"
+	"go/parser"
+	"go/token"
+	"os"
+	"regexp"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+type promptAuditOrderCase struct {
+	file       string
+	function   string
+	auditToken string
+}
+
+func TestPromptAuditGatePrecedesAccountBillingAndUpstreamSideEffects(t *testing.T) {
+	tests := []promptAuditOrderCase{
+		{file: "gateway_handler.go", function: "Messages", auditToken: "checkSecurityAudit"},
+		{file: "gateway_handler_chat_completions.go", function: "ChatCompletions", auditToken: "checkSecurityAudit"},
+		{file: "gateway_handler_responses.go", function: "Responses", auditToken: "checkSecurityAudit"},
+		{file: "gemini_v1beta_handler.go", function: "GeminiV1BetaModels", auditToken: "checkSecurityAudit"},
+		{file: "openai_gateway_handler.go", function: "Responses", auditToken: "checkSecurityAudit"},
+		{file: "openai_gateway_handler.go", function: "Messages", auditToken: "checkSecurityAudit"},
+		{file: "openai_chat_completions.go", function: "ChatCompletions", auditToken: "checkSecurityAudit"},
+		{file: "openai_images.go", function: "Images", auditToken: "checkSecurityAudit"},
+		{file: "grok_media.go", function: "handleGrokMedia", auditToken: "checkSecurityAudit"},
+		{file: "openai_embeddings.go", function: "Embeddings", auditToken: "checkSecurityAudit"},
+		{file: "openai_alpha_search.go", function: "AlphaSearch", auditToken: "checkSecurityAudit"},
+		{file: "image_task_handler.go", function: "Submit", auditToken: "checkSecurityAuditBeforeSubmit"},
+		{file: "batch_image_handler.go", function: "Submit", auditToken: "checkSecurityAuditBeforeSubmit"},
+	}
+	sideEffectTokens := []string{
+		"CheckBillingEligibility(", "SelectAccount", ".Forward", "acquireResponsesUserSlot(",
+		"AcquireUserSlot", "TryAcquireUserSlot", "acquireImageGenerationSlot(",
+		"h.tasks.Create(", "h.service.Submit(",
+	}
+	for _, tt := range tests {
+		t.Run(tt.file+"/"+tt.function, func(t *testing.T) {
+			functionSource := stripGoComments(goFunctionSource(t, tt.file, tt.function))
+			auditIndex := strings.Index(functionSource, tt.auditToken)
+			require.NotEqual(t, -1, auditIndex, "missing Prompt Audit gate")
+			foundSideEffect := false
+			for _, sideEffect := range sideEffectTokens {
+				index := strings.Index(functionSource, sideEffect)
+				if index < 0 {
+					continue
+				}
+				foundSideEffect = true
+				require.Lessf(t, auditIndex, index, "%s must run before %s", tt.auditToken, sideEffect)
+			}
+			require.True(t, foundSideEffect, "coverage case must contain a downstream side effect")
+		})
+	}
+}
+
+func stripGoComments(source string) string {
+	source = regexp.MustCompile(`(?s)/\*.*?\*/`).ReplaceAllString(source, "")
+	return regexp.MustCompile(`(?m)//.*$`).ReplaceAllString(source, "")
+}
+
+func goFunctionSource(t *testing.T, filename, functionName string) string {
+	t.Helper()
+	raw, err := os.ReadFile(filename)
+	require.NoError(t, err)
+	files := token.NewFileSet()
+	parsed, err := parser.ParseFile(files, filename, raw, 0)
+	require.NoError(t, err)
+	for _, declaration := range parsed.Decls {
+		function, ok := declaration.(*ast.FuncDecl)
+		if !ok || function.Name.Name != functionName || function.Body == nil {
+			continue
+		}
+		start := files.Position(function.Pos()).Offset
+		end := files.Position(function.End()).Offset
+		require.Greater(t, end, start)
+		return string(raw[start:end])
+	}
+	t.Fatalf("function %s not found in %s", functionName, filename)
+	return ""
+}