李建琦
6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
86 lines
3.3 KiB
Go
86 lines
3.3 KiB
Go
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 ""
|
|
}
|