Skip to content

Commit d11bdb1

Browse files
xtt21625457xtt21625457
committed
feat(security-audit): add OpenAI-compatible prompt auditing
1 parent a2779cd commit d11bdb1

103 files changed

Lines changed: 17549 additions & 70 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎backend/cmd/server/main.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,13 @@ func runMainServer() {
153153
log.Fatalf("Failed to initialize application: %v", err)
154154
}
155155
defer app.Cleanup()
156+
if app.PromptAudit != nil {
157+
if err := app.PromptAudit.Start(context.Background()); err != nil {
158+
// Prompt Audit is default-off and isolated. Startup degradation must be
159+
// observable but must not take unrelated APIs down.
160+
log.Printf("Prompt Audit started in degraded state: %v", err)
161+
}
162+
}
156163

157164
// 启动服务器
158165
go func() {

‎backend/cmd/server/wire.go‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
"github.com/Wei-Shaw/sub2api/internal/handler"
1616
"github.com/Wei-Shaw/sub2api/internal/payment"
1717
"github.com/Wei-Shaw/sub2api/internal/repository"
18+
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
1819
"github.com/Wei-Shaw/sub2api/internal/server"
1920
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
2021
"github.com/Wei-Shaw/sub2api/internal/service"
@@ -24,8 +25,9 @@ import (
2425
)
2526

2627
type Application struct {
27-
Server *http.Server
28-
Cleanup func()
28+
Server *http.Server
29+
PromptAudit *securityaudit.PromptService
30+
Cleanup func()
2931
}
3032

3133
func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
@@ -36,6 +38,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
3638
// Business layer ProviderSets
3739
repository.ProviderSet,
3840
service.ProviderSet,
41+
securityaudit.ProviderSet,
3942
payment.ProviderSet,
4043
middleware.ProviderSet,
4144
handler.ProviderSet,
@@ -53,7 +56,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
5356
provideCleanup,
5457

5558
// Application struct
56-
wire.Struct(new(Application), "Server", "Cleanup"),
59+
wire.Struct(new(Application), "Server", "PromptAudit", "Cleanup"),
5760
)
5861
return nil, nil
5962
}
@@ -105,6 +108,7 @@ func provideCleanup(
105108
quotaFlusher *service.UserPlatformQuotaUsageFlusher,
106109
upstreamBillingProbe *service.UpstreamBillingProbeService,
107110
auditLog *service.AuditLogService,
111+
promptAudit *securityaudit.PromptService,
108112
) func() {
109113
return func() {
110114
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
@@ -117,6 +121,12 @@ func provideCleanup(
117121

118122
// 应用层清理步骤可并行执行，基础设施资源（Redis/Ent）最后按顺序关闭。
119123
parallelSteps := []cleanupStep{
124+
{"PromptAuditService", func() error {
125+
if promptAudit != nil {
126+
return promptAudit.Shutdown(ctx)
127+
}
128+
return nil
129+
}},
120130
{"OpsScheduledReportService", func() error {
121131
if opsScheduledReport != nil {
122132
opsScheduledReport.Stop()

‎backend/cmd/server/wire_gen.go‎

Lines changed: 28 additions & 9 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎backend/cmd/server/wire_gen_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) {
8585
nil, // quotaFlusher
8686
nil, // upstreamBillingProbe
8787
nil, // auditLog
88+
nil, // promptAudit
8889
)
8990

9091
require.NotPanics(t, func() {

‎backend/internal/handler/batch_image_handler.go‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package handler
22

33
import (
4+
"encoding/json"
45
"errors"
56
"io"
67
"net/http"
@@ -20,6 +21,7 @@ type BatchImageHandler struct {
2021
service *service.BatchImagePublicService
2122
download *service.BatchImageDownloadService
2223
cleanup *service.BatchImageCleanupService
24+
openAI *OpenAIGatewayHandler
2325
}
2426

2527
func NewBatchImageHandler(service *service.BatchImagePublicService, download *service.BatchImageDownloadService, cleanup *service.BatchImageCleanupService) *BatchImageHandler {
@@ -37,6 +39,9 @@ func (h *BatchImageHandler) Submit(c *gin.Context) {
3739
batchImageError(c, infraerrors.New(http.StatusUnauthorized, "API_KEY_REQUIRED", "API key is required"))
3840
return
3941
}
42+
if !h.checkSecurityAuditBeforeSubmit(c, &req) {
43+
return
44+
}
4045
got, err := h.service.Submit(c.Request.Context(), owner, req, c.GetHeader("Idempotency-Key"))
4146
if err != nil {
4247
batchImageError(c, err)
@@ -45,6 +50,44 @@ func (h *BatchImageHandler) Submit(c *gin.Context) {
4550
c.JSON(http.StatusOK, got)
4651
}
4752

53+
func (h *BatchImageHandler) checkSecurityAuditBeforeSubmit(c *gin.Context, req *service.BatchImageSubmitRequest) bool {
54+
if h == nil || h.openAI == nil || req == nil {
55+
return true
56+
}
57+
apiKey, ok := middleware.GetAPIKeyFromContext(c)
58+
if !ok || apiKey == nil {
59+
batchImageError(c, infraerrors.New(http.StatusUnauthorized, "API_KEY_REQUIRED", "API key is required"))
60+
return false
61+
}
62+
subject, ok := middleware.GetAuthSubjectFromContext(c)
63+
if !ok {
64+
batchImageError(c, infraerrors.New(http.StatusInternalServerError, "USER_CONTEXT_REQUIRED", "User context not found"))
65+
return false
66+
}
67+
items := make([]map[string]string, 0, len(req.Items))
68+
for _, item := range req.Items {
69+
if prompt := strings.TrimSpace(item.Prompt); prompt != "" {
70+
items = append(items, map[string]string{"prompt": prompt})
71+
}
72+
}
73+
if len(items) == 0 {
74+
return true
75+
}
76+
body, err := json.Marshal(map[string]any{"request": map[string]any{"items": items}})
77+
if err != nil {
78+
batchImageError(c, infraerrors.New(http.StatusBadRequest, "INVALID_BATCH_PROMPT", "batch prompts are invalid"))
79+
return false
80+
}
81+
reqLog := requestLogger(c, "handler.batch_image.security_audit",
82+
zap.Int64("user_id", subject.UserID), zap.Int64("api_key_id", apiKey.ID), zap.String("model", req.Model))
83+
decision := h.openAI.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIImages, req.Model, body)
84+
if decision != nil && !decision.AllowNextStage {
85+
h.openAI.openAISecurityAuditError(c, decision)
86+
return false
87+
}
88+
return true
89+
}
90+
4891
func (h *BatchImageHandler) Get(c *gin.Context) {
4992
owner, ok := batchImageOwnerFromContext(c)
5093
if !ok {

‎backend/internal/handler/content_moderation_helper.go‎

Lines changed: 0 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -12,13 +12,6 @@ import (
1212
"go.uber.org/zap"
1313
)
1414

15-
func (h *GatewayHandler) checkContentModeration(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol string, model string, body []byte) *service.ContentModerationDecision {
16-
if h == nil || h.contentModerationService == nil {
17-
return nil
18-
}
19-
return runContentModeration(c, reqLog, h.contentModerationService, apiKey, subject, protocol, model, body)
20-
}
21-
2215
func contentModerationStatus(decision *service.ContentModerationDecision) int {
2316
if decision == nil || decision.StatusCode < 400 || decision.StatusCode > 599 {
2417
return http.StatusForbidden
@@ -30,13 +23,6 @@ func contentModerationErrorCode(decision *service.ContentModerationDecision) str
3023
return "content_policy_violation"
3124
}
3225

33-
func (h *OpenAIGatewayHandler) checkContentModeration(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol string, model string, body []byte) *service.ContentModerationDecision {
34-
if h == nil || h.contentModerationService == nil {
35-
return nil
36-
}
37-
return runContentModeration(c, reqLog, h.contentModerationService, apiKey, subject, protocol, model, body)
38-
}
39-
4026
func runContentModeration(c *gin.Context, reqLog *zap.Logger, svc *service.ContentModerationService, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol string, model string, body []byte) *service.ContentModerationDecision {
4127
if svc == nil || c == nil || c.Request == nil {
4228
return nil

‎backend/internal/handler/gateway_handler.go‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import (
2525
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
2626
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
2727
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
28+
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
2829
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
2930
"github.com/Wei-Shaw/sub2api/internal/service"
3031

@@ -49,6 +50,7 @@ type GatewayHandler struct {
4950
usageRecordWorkerPool *service.UsageRecordWorkerPool
5051
errorPassthroughService *service.ErrorPassthroughService
5152
contentModerationService *service.ContentModerationService
53+
securityAuditCoordinator *securityaudit.Coordinator
5254
concurrencyHelper *ConcurrencyHelper
5355
userMsgQueueHelper *UserMsgQueueHelper
5456
maxAccountSwitches int
@@ -199,8 +201,8 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
199201
return
200202
}
201203

202-
if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolAnthropicMessages, reqModel, body); decision != nil && decision.Blocked {
203-
h.errorResponse(c, contentModerationStatus(decision), contentModerationErrorCode(decision), decision.Message)
204+
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolAnthropicMessages, reqModel, body); decision != nil && !decision.AllowNextStage {
205+
h.anthropicSecurityAuditError(c, decision)
204206
return
205207
}
206208

‎backend/internal/handler/gateway_handler_chat_completions.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -99,8 +99,8 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
9999
return
100100
}
101101

102-
if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIChat, reqModel, body); decision != nil && decision.Blocked {
103-
h.chatCompletionsErrorResponse(c, contentModerationStatus(decision), contentModerationErrorCode(decision), decision.Message)
102+
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIChat, reqModel, body); decision != nil && !decision.AllowNextStage {
103+
h.openAISecurityAuditError(c, decision)
104104
return
105105
}
106106

‎backend/internal/handler/gateway_handler_responses.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -104,8 +104,8 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
104104
return
105105
}
106106

107-
if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIResponses, reqModel, body); decision != nil && decision.Blocked {
108-
h.responsesErrorResponse(c, contentModerationStatus(decision), contentModerationErrorCode(decision), decision.Message)
107+
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIResponses, reqModel, body); decision != nil && !decision.AllowNextStage {
108+
h.responsesSecurityAuditError(c, decision)
109109
return
110110
}
111111

‎backend/internal/handler/gemini_v1beta_handler.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -187,8 +187,8 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
187187
setOpsRequestContext(c, modelName, stream)
188188
setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(stream, false)))
189189

190-
if decision := h.checkContentModeration(c, reqLog, apiKey, authSubject, service.ContentModerationProtocolGemini, modelName, body); decision != nil && decision.Blocked {
191-
googleError(c, contentModerationStatus(decision), decision.Message)
190+
if decision := h.checkSecurityAudit(c, reqLog, apiKey, authSubject, service.ContentModerationProtocolGemini, modelName, body); decision != nil && !decision.AllowNextStage {
191+
googleSecurityAuditError(c, decision)
192192
return
193193
}
194194

0 commit comments

Comments
 (0)