Sub2API v1.0 - AI API 网关(二开初始版本,基于上游 Wei-Shaw/sub2api)
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

This commit is contained in:
李建琦
2026-08-21 18:30:13 +08:00
commit 6d655c9903
3584 changed files with 1270640 additions and 0 deletions
@@ -0,0 +1,139 @@
package securityaudit
import (
"context"
"errors"
"net/http"
"sync"
)
type LegacyEngine interface {
Check(ctx context.Context, req Request) (*LegacyDecision, error)
}
type PromptEngine interface {
EffectiveMode() Mode
Enqueue(ctx context.Context, req Request) error
Evaluate(ctx context.Context, req Request) (*PromptDecision, error)
}
type Coordinator struct {
legacy LegacyEngine
prompt PromptEngine
}
func NewCoordinator(legacy LegacyEngine, prompt PromptEngine) *Coordinator {
return &Coordinator{legacy: legacy, prompt: prompt}
}
func (c *Coordinator) Check(ctx context.Context, req Request) Decision {
if c == nil {
return allowDecision(nil, nil)
}
mode := ModeOff
if c.prompt != nil {
mode = c.prompt.EffectiveMode()
}
switch mode {
case ModeAsync:
// Enqueue is deliberately best-effort. The implementation owns a bounded
// context and copies request memory before it can outlive the Handler.
_ = c.prompt.Enqueue(ctx, req.Clone())
legacy, _ := c.checkLegacy(ctx, req)
return prioritize(legacy, nil)
case ModeBlocking:
return c.checkBlocking(ctx, req)
default:
legacy, _ := c.checkLegacy(ctx, req)
return prioritize(legacy, nil)
}
}
func (c *Coordinator) checkBlocking(ctx context.Context, req Request) Decision {
var wg sync.WaitGroup
wg.Add(2)
var legacy *LegacyDecision
var prompt *PromptDecision
go func() {
defer wg.Done()
legacy, _ = c.checkLegacy(ctx, req)
}()
go func() {
defer wg.Done()
if c.prompt == nil {
prompt = unavailablePromptDecision(ErrorCodeUnavailable)
return
}
result, err := c.prompt.Evaluate(ctx, req.Clone())
if err != nil {
var guardErr *GuardError
if errors.As(err, &guardErr) && guardErr.Code == ErrorCodeInvalidResponse {
prompt = unavailablePromptDecision(ErrorCodeInvalidResponse)
return
}
prompt = unavailablePromptDecision(ErrorCodeUnavailable)
return
}
if result == nil {
prompt = unavailablePromptDecision(ErrorCodeUnavailable)
return
}
prompt = result
}()
wg.Wait()
return prioritize(legacy, prompt)
}
func (c *Coordinator) checkLegacy(ctx context.Context, req Request) (*LegacyDecision, error) {
if c.legacy == nil {
return nil, nil
}
return c.legacy.Check(ctx, req)
}
func prioritize(legacy *LegacyDecision, prompt *PromptDecision) Decision {
if legacy != nil && legacy.Blocked {
status := legacy.StatusCode
if status < 400 || status > 599 {
status = http.StatusForbidden
}
code := legacy.ErrorCode
if code == "" {
code = "content_policy_violation"
}
return Decision{
Kind: DecisionBlock, HTTPStatus: status, ErrorCode: code, ClientMessage: legacy.Message,
Legacy: legacy, Prompt: prompt, AllowNextStage: false,
}
}
if prompt == nil {
return allowDecision(legacy, nil)
}
switch prompt.Kind {
case DecisionBlock:
return Decision{Kind: DecisionBlock, HTTPStatus: http.StatusForbidden, ErrorCode: ErrorCodeBlocked,
ClientMessage: "提示词安全审计拒绝了该请求,请调整输入后重试", Legacy: legacy, Prompt: prompt}
case DecisionInvalid:
return Decision{Kind: DecisionInvalid, HTTPStatus: http.StatusServiceUnavailable, ErrorCode: ErrorCodeInvalidResponse,
ClientMessage: "提示词安全审计暂时不可用,请稍后重试", Legacy: legacy, Prompt: prompt}
case DecisionUnavailable:
return Decision{Kind: DecisionUnavailable, HTTPStatus: http.StatusServiceUnavailable, ErrorCode: ErrorCodeUnavailable,
ClientMessage: "提示词安全审计暂时不可用,请稍后重试", Legacy: legacy, Prompt: prompt}
case DecisionFlag:
return Decision{Kind: DecisionFlag, HTTPStatus: http.StatusOK, Legacy: legacy, Prompt: prompt, AllowNextStage: true}
default:
return allowDecision(legacy, prompt)
}
}
func allowDecision(legacy *LegacyDecision, prompt *PromptDecision) Decision {
return Decision{Kind: DecisionAllow, HTTPStatus: http.StatusOK, Legacy: legacy, Prompt: prompt, AllowNextStage: true}
}
func unavailablePromptDecision(code string) *PromptDecision {
kind := DecisionUnavailable
if code == ErrorCodeInvalidResponse {
kind = DecisionInvalid
}
return &PromptDecision{Kind: kind, ErrorCode: code, AllowNextStage: false}
}
@@ -0,0 +1,35 @@
package securityaudit
import (
"context"
"github.com/Wei-Shaw/sub2api/internal/service"
)
type LegacyModerationAdapter struct {
service *service.ContentModerationService
}
func NewLegacyModerationAdapter(svc *service.ContentModerationService) LegacyEngine {
return &LegacyModerationAdapter{service: svc}
}
func (a *LegacyModerationAdapter) Check(ctx context.Context, req Request) (*LegacyDecision, error) {
if a == nil || a.service == nil {
return nil, nil
}
decision, err := a.service.Check(ctx, service.ContentModerationCheckInput{
RequestID: req.RequestID, UserID: req.UserID, UserEmail: req.UserEmail,
APIKeyID: req.APIKeyID, APIKeyName: req.APIKeyName, GroupID: cloneInt64Ptr(req.GroupID),
GroupName: req.GroupName, Endpoint: req.Endpoint, Provider: req.Provider,
Model: req.Model, Protocol: req.Protocol, Body: req.Body,
})
if err != nil || decision == nil {
return nil, err
}
return &LegacyDecision{
Allowed: decision.Allowed, Blocked: decision.Blocked, Flagged: decision.Flagged,
Message: decision.Message, StatusCode: decision.StatusCode,
ErrorCode: "content_policy_violation", Action: decision.Action,
}, nil
}
@@ -0,0 +1,176 @@
package securityaudit
import (
"context"
"errors"
"fmt"
"net/http"
"sync/atomic"
"testing"
"github.com/stretchr/testify/require"
)
type fakeLegacyEngine struct {
decision *LegacyDecision
err error
calls atomic.Int64
}
func (f *fakeLegacyEngine) Check(context.Context, Request) (*LegacyDecision, error) {
f.calls.Add(1)
return f.decision, f.err
}
type fakePromptEngine struct {
mode Mode
decision *PromptDecision
err error
enqueues atomic.Int64
evaluates atomic.Int64
}
func (f *fakePromptEngine) EffectiveMode() Mode { return f.mode }
func (f *fakePromptEngine) Enqueue(context.Context, Request) error {
f.enqueues.Add(1)
return f.err
}
func (f *fakePromptEngine) Evaluate(context.Context, Request) (*PromptDecision, error) {
f.evaluates.Add(1)
return f.decision, f.err
}
func TestCoordinatorModesAndPriority(t *testing.T) {
tests := []struct {
name string
mode Mode
legacy *LegacyDecision
prompt *PromptDecision
promptErr error
wantKind DecisionKind
wantCode string
wantEnqueue int64
wantEvaluation int64
}{
{name: "off", mode: ModeOff, wantKind: DecisionAllow},
{name: "async only enqueues", mode: ModeAsync, wantKind: DecisionAllow, wantEnqueue: 1},
{name: "prompt block", mode: ModeBlocking, prompt: &PromptDecision{Kind: DecisionBlock}, wantKind: DecisionBlock, wantCode: ErrorCodeBlocked, wantEvaluation: 1},
{name: "prompt unavailable", mode: ModeBlocking, promptErr: errors.New("down"), wantKind: DecisionUnavailable, wantCode: ErrorCodeUnavailable, wantEvaluation: 1},
{name: "legacy wins both block", mode: ModeBlocking,
legacy: &LegacyDecision{Blocked: true, StatusCode: http.StatusForbidden, ErrorCode: "content_policy_violation", Message: "legacy"},
prompt: &PromptDecision{Kind: DecisionBlock}, wantKind: DecisionBlock, wantCode: "content_policy_violation", wantEvaluation: 1},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
legacy := &fakeLegacyEngine{decision: tt.legacy}
prompt := &fakePromptEngine{mode: tt.mode, decision: tt.prompt, err: tt.promptErr}
decision := NewCoordinator(legacy, prompt).Check(context.Background(), Request{Body: []byte(`{}`)})
require.Equal(t, tt.wantKind, decision.Kind)
require.Equal(t, tt.wantCode, decision.ErrorCode)
require.Equal(t, int64(1), legacy.calls.Load())
require.Equal(t, tt.wantEnqueue, prompt.enqueues.Load())
require.Equal(t, tt.wantEvaluation, prompt.evaluates.Load())
})
}
}
func TestCoordinatorDoesNotMutateRequestBody(t *testing.T) {
body := []byte(`{"messages":[{"role":"user","content":"hello"}]}`)
original := append([]byte(nil), body...)
prompt := &fakePromptEngine{mode: ModeAsync}
decision := NewCoordinator(&fakeLegacyEngine{}, prompt).Check(context.Background(), Request{Body: body})
require.True(t, decision.AllowNextStage)
require.Equal(t, original, body)
}
func TestCoordinatorBlockingPriorityCoversBothEngineDecisionMatrix(t *testing.T) {
legacyCases := []struct {
name string
decision *LegacyDecision
}{
{name: "allow", decision: &LegacyDecision{Allowed: true, StatusCode: http.StatusOK, Action: "allow"}},
{name: "flag", decision: &LegacyDecision{Allowed: true, Flagged: true, StatusCode: http.StatusOK, Action: "flag"}},
{name: "block", decision: &LegacyDecision{Blocked: true, StatusCode: http.StatusForbidden, ErrorCode: "legacy_exact_code", Message: "legacy exact message", Action: "block"}},
}
promptCases := []struct {
name string
decision *PromptDecision
wantKind DecisionKind
wantCode string
}{
{name: "allow", decision: &PromptDecision{Kind: DecisionAllow, AllowNextStage: true}, wantKind: DecisionAllow},
{name: "flag", decision: &PromptDecision{Kind: DecisionFlag, AllowNextStage: true}, wantKind: DecisionFlag},
{name: "block", decision: &PromptDecision{Kind: DecisionBlock}, wantKind: DecisionBlock, wantCode: ErrorCodeBlocked},
{name: "unavailable", decision: &PromptDecision{Kind: DecisionUnavailable, ErrorCode: ErrorCodeUnavailable}, wantKind: DecisionUnavailable, wantCode: ErrorCodeUnavailable},
{name: "invalid", decision: &PromptDecision{Kind: DecisionInvalid, ErrorCode: ErrorCodeInvalidResponse}, wantKind: DecisionInvalid, wantCode: ErrorCodeInvalidResponse},
}
for _, legacyCase := range legacyCases {
for _, promptCase := range promptCases {
t.Run(fmt.Sprintf("legacy_%s_prompt_%s", legacyCase.name, promptCase.name), func(t *testing.T) {
legacy := &fakeLegacyEngine{decision: legacyCase.decision}
prompt := &fakePromptEngine{mode: ModeBlocking, decision: promptCase.decision}
decision := NewCoordinator(legacy, prompt).Check(context.Background(), Request{})
require.Same(t, legacyCase.decision, decision.Legacy)
require.Same(t, promptCase.decision, decision.Prompt)
require.Equal(t, int64(1), legacy.calls.Load())
require.Equal(t, int64(1), prompt.evaluates.Load())
if legacyCase.name == "block" {
require.Equal(t, DecisionBlock, decision.Kind)
require.Equal(t, "legacy_exact_code", decision.ErrorCode)
require.Equal(t, "legacy exact message", decision.ClientMessage)
require.False(t, decision.AllowNextStage)
return
}
require.Equal(t, promptCase.wantKind, decision.Kind)
require.Equal(t, promptCase.wantCode, decision.ErrorCode)
require.Equal(t, promptCase.decision.AllowNextStage, decision.AllowNextStage)
})
}
}
}
func TestCoordinatorPreservesIndependentEngineFactsAndMapsOnlyGatewayOutcome(t *testing.T) {
legacyDecision := &LegacyDecision{
Allowed: true, Flagged: true, Message: "legacy finding", StatusCode: http.StatusAccepted,
ErrorCode: "legacy_observation", Action: "legacy_action",
}
promptResult := &NormalizedResult{
Decision: EventCritical, RiskLevel: RiskCritical, Action: ActionBlock,
Categories: []string{"pii"}, ScannerScores: map[string]float64{"pii": 1},
}
promptDecision := &PromptDecision{Kind: DecisionBlock, Result: promptResult}
decision := NewCoordinator(
&fakeLegacyEngine{decision: legacyDecision},
&fakePromptEngine{mode: ModeBlocking, decision: promptDecision},
).Check(context.Background(), Request{})
require.Same(t, legacyDecision, decision.Legacy)
require.Same(t, promptDecision, decision.Prompt)
require.Same(t, promptResult, decision.Prompt.Result)
require.Equal(t, "legacy finding", decision.Legacy.Message)
require.Equal(t, []string{"pii"}, decision.Prompt.Result.Categories)
require.Equal(t, ErrorCodeBlocked, decision.ErrorCode)
}
func TestCoordinatorAsyncEnqueueFailuresNeverChangeResponseOrDownstreamDispatch(t *testing.T) {
for _, enqueueErr := range []error{ErrQueueFull, ErrQueueAdmissionBusy, errors.New("redis unavailable"), errors.New("publish failed")} {
prompt := &fakePromptEngine{mode: ModeAsync, err: enqueueErr}
decision := NewCoordinator(&fakeLegacyEngine{decision: &LegacyDecision{Allowed: true}}, prompt).Check(context.Background(), Request{})
downstreamDispatches := 0
status := http.StatusOK
responseBody := "unchanged-upstream-response"
if decision.AllowNextStage {
downstreamDispatches++
} else {
status = decision.HTTPStatus
responseBody = decision.ClientMessage
}
require.Equal(t, http.StatusOK, status)
require.Equal(t, "unchanged-upstream-response", responseBody)
require.Equal(t, 1, downstreamDispatches)
require.Equal(t, int64(1), prompt.enqueues.Load())
require.Zero(t, prompt.evaluates.Load())
}
}
@@ -0,0 +1,512 @@
package securityaudit
import (
"context"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"sort"
"strings"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
)
const (
DefaultWorkerCount = 4
MaxWorkerCount = 32
DefaultQueueCapacity = 32768
MaxQueueCapacity = 100000
DefaultTimeoutMS = 3000
MinTimeoutMS = 100
MaxTimeoutMS = 30000
DefaultInputLimit = 4000
MinInputLimit = 128
MaxInputLimit = 100000
DefaultPayloadTTL = 30 * time.Minute
)
type SecretEncryptor interface {
Encrypt(plaintext string) (string, error)
Decrypt(ciphertext string) (string, error)
}
// ConfigStore is the injectable boundary between hot-path prompt auditing and
// the concrete settings/PostgreSQL/Redis-backed configuration manager.
type ConfigStore interface {
Start(ctx context.Context) error
Shutdown(ctx context.Context) error
Active() (ActiveConfig, bool)
EffectiveMode() Mode
// BlockingActivationDegraded is true when storage intent requires blocking
// but no usable blocking snapshot is active (cold start or failed reload).
// It must stay false when blocking is not intended, even if config is
// untrusted—otherwise default-off deployments fail closed for all traffic.
BlockingActivationDegraded() bool
Public() (PublicConfig, error)
Save(ctx context.Context, req UpdateConfigRequest, actorID int64) (PublicConfig, error)
RuntimeState() (expected int64, active int64, loadedAt *time.Time, loadError string)
Encrypt(value string) (string, error)
Decrypt(value string) (string, error)
}
type StorageEndpoint struct {
ID string `json:"id"`
Name string `json:"name"`
Protocol string `json:"protocol"`
BaseURL string `json:"base_url"`
Model string `json:"model"`
TokenCiphertext string `json:"token_ciphertext,omitempty"`
TimeoutMS int `json:"timeout_ms"`
InputLimit int `json:"input_limit"`
Enabled bool `json:"enabled"`
}
type storageConfig struct {
Enabled bool `json:"enabled"`
BlockingEnabled bool `json:"blocking_enabled"`
BlockingLatestTurnOnly bool `json:"blocking_latest_turn_only"`
StorePassEvents bool `json:"store_pass_events"`
Strategy string `json:"strategy"`
WorkerCount int `json:"worker_count"`
QueueCapacity int `json:"queue_capacity"`
Scanners []string `json:"scanners"`
AllGroups bool `json:"all_groups"`
GroupIDs []int64 `json:"group_ids"`
Endpoints []StorageEndpoint `json:"endpoints"`
ConfigVersion int64 `json:"config_version"`
UpdatedAt time.Time `json:"updated_at"`
UpdatedBy int64 `json:"updated_by"`
ChangeSummary string `json:"change_summary"`
}
type ActiveEndpoint struct {
ID string
Name string
Protocol string
BaseURL string
Model string
Token string
TimeoutMS int
InputLimit int
Enabled bool
// TokenInvalid marks an endpoint whose persisted token ciphertext cannot be
// decrypted with the current encryption key (key changed or auto-generated
// on restart). The endpoint is kept visible for admins but excluded from
// runtime use until the token is re-entered or cleared (issue #4887).
TokenInvalid bool
}
type ActiveConfig struct {
RiskControlEnabled bool
Enabled bool
BlockingEnabled bool
BlockingLatestTurnOnly bool
StorePassEvents bool
Strategy string
WorkerCount int
QueueCapacity int
Scanners []string
AllGroups bool
GroupIDs []int64
Endpoints []ActiveEndpoint
ConfigVersion int64
UpdatedAt time.Time
UpdatedBy int64
ChangeSummary string
}
type PublicEndpoint struct {
ID string `json:"id"`
Name string `json:"name"`
Protocol string `json:"protocol"`
BaseURL string `json:"base_url"`
Model string `json:"model"`
TimeoutMS int `json:"timeout_ms"`
InputLimit int `json:"input_limit"`
Enabled bool `json:"enabled"`
HasToken bool `json:"has_token"`
TokenStatus string `json:"token_status"`
}
type PublicConfig struct {
Enabled bool `json:"enabled"`
BlockingEnabled bool `json:"blocking_enabled"`
BlockingLatestTurnOnly bool `json:"blocking_latest_turn_only"`
StorePassEvents bool `json:"store_pass_events"`
EffectiveMode Mode `json:"effective_mode"`
Strategy string `json:"strategy"`
WorkerCount int `json:"worker_count"`
QueueCapacity int `json:"queue_capacity"`
Scanners []string `json:"scanners"`
AllGroups bool `json:"all_groups"`
GroupIDs []int64 `json:"group_ids"`
Endpoints []PublicEndpoint `json:"endpoints"`
ConfigVersion int64 `json:"config_version"`
UpdatedAt time.Time `json:"updated_at"`
UpdatedBy int64 `json:"updated_by"`
ChangeSummary string `json:"change_summary"`
}
type UpdateEndpoint struct {
ID string `json:"id" binding:"required"`
Name string `json:"name" binding:"required"`
Protocol string `json:"protocol"`
BaseURL string `json:"base_url" binding:"required"`
Model string `json:"model"`
Token string `json:"token,omitempty"`
ClearToken bool `json:"clear_token"`
TimeoutMS int `json:"timeout_ms"`
InputLimit int `json:"input_limit"`
Enabled bool `json:"enabled"`
}
type UpdateConfigRequest struct {
ExpectedConfigVersion int64 `json:"expected_config_version" binding:"required"`
Enabled bool `json:"enabled"`
BlockingEnabled bool `json:"blocking_enabled"`
BlockingLatestTurnOnly bool `json:"blocking_latest_turn_only"`
StorePassEvents bool `json:"store_pass_events"`
Strategy string `json:"strategy"`
WorkerCount int `json:"worker_count"`
QueueCapacity int `json:"queue_capacity"`
Scanners []string `json:"scanners"`
AllGroups bool `json:"all_groups"`
GroupIDs []int64 `json:"group_ids"`
Endpoints []UpdateEndpoint `json:"endpoints"`
}
func DefaultStorageConfig() storageConfig {
return storageConfig{
Enabled: false,
BlockingEnabled: false,
BlockingLatestTurnOnly: false,
StorePassEvents: false,
Strategy: "priority",
WorkerCount: DefaultWorkerCount,
QueueCapacity: DefaultQueueCapacity,
Scanners: append([]string(nil), AllScannerIDs...),
AllGroups: true,
GroupIDs: []int64{},
Endpoints: []StorageEndpoint{},
ConfigVersion: 1,
}
}
func ParseStorageConfig(raw string) (storageConfig, error) {
cfg := DefaultStorageConfig()
if strings.TrimSpace(raw) == "" {
return cfg, nil
}
if err := json.Unmarshal([]byte(raw), &cfg); err != nil {
return storageConfig{}, fmt.Errorf("decode prompt audit config: %w", err)
}
normalizeStorageConfig(&cfg)
if err := validateStorageConfig(cfg); err != nil {
return storageConfig{}, err
}
return cfg, nil
}
func normalizeStorageConfig(cfg *storageConfig) {
if cfg == nil {
return
}
if cfg.ConfigVersion < 1 {
cfg.ConfigVersion = 1
}
if strings.TrimSpace(cfg.Strategy) == "" {
cfg.Strategy = "priority"
}
if cfg.WorkerCount == 0 {
cfg.WorkerCount = DefaultWorkerCount
}
if cfg.QueueCapacity == 0 {
cfg.QueueCapacity = DefaultQueueCapacity
}
if len(cfg.Scanners) == 0 {
cfg.Scanners = append([]string(nil), AllScannerIDs...)
}
cfg.Scanners = canonicalScannerIDs(cfg.Scanners)
cfg.GroupIDs = canonicalInt64s(cfg.GroupIDs)
// Preserve an invalid blocking-without-audit combination so validation can
// reject it instead of silently changing administrator intent.
for i := range cfg.Endpoints {
ep := &cfg.Endpoints[i]
ep.ID = strings.TrimSpace(ep.ID)
ep.Name = strings.TrimSpace(ep.Name)
ep.Protocol = strings.TrimSpace(ep.Protocol)
if ep.Protocol == "" {
ep.Protocol = "openai_compatible"
}
ep.BaseURL = strings.TrimSpace(ep.BaseURL)
ep.Model = strings.TrimSpace(ep.Model)
if ep.Model == "" {
ep.Model = DefaultGuardModel
}
if ep.TimeoutMS == 0 {
ep.TimeoutMS = DefaultTimeoutMS
}
if ep.InputLimit == 0 {
ep.InputLimit = DefaultInputLimit
}
}
}
func validateStorageConfig(cfg storageConfig) error {
if cfg.BlockingEnabled && !cfg.Enabled {
return infraerrors.BadRequest(ErrorCodeRequiresEnabled, "开启同步阻止前必须先启用提示词审计")
}
if cfg.Strategy != "priority" {
return infraerrors.BadRequest("prompt_audit_invalid_strategy", "提示词审计策略仅支持 priority")
}
if cfg.WorkerCount < 1 || cfg.WorkerCount > MaxWorkerCount {
return infraerrors.BadRequest("prompt_audit_invalid_worker_count", "Worker 数量超出允许范围")
}
if cfg.QueueCapacity < 1 || cfg.QueueCapacity > MaxQueueCapacity {
return infraerrors.BadRequest("prompt_audit_invalid_queue_capacity", "队列容量超出允许范围")
}
if !cfg.AllGroups && len(cfg.GroupIDs) == 0 {
return infraerrors.BadRequest("prompt_audit_groups_required", "指定分组模式至少需要选择一个分组")
}
if len(cfg.Scanners) == 0 {
return infraerrors.BadRequest("prompt_audit_scanners_required", "至少需要启用一个风险分类")
}
seen := make(map[string]struct{}, len(cfg.Endpoints))
enabled := 0
for _, ep := range cfg.Endpoints {
if ep.ID == "" || ep.Name == "" {
return infraerrors.BadRequest("prompt_audit_invalid_endpoint", "审计节点 ID 和名称不能为空")
}
if _, ok := seen[ep.ID]; ok {
return infraerrors.BadRequest("prompt_audit_duplicate_endpoint", "审计节点 ID 不能重复")
}
seen[ep.ID] = struct{}{}
if ep.Protocol != "openai_compatible" {
return infraerrors.BadRequest("prompt_audit_invalid_endpoint_protocol", "审计节点仅支持 OpenAI 兼容协议")
}
if _, err := NormalizeBaseURL(ep.BaseURL); err != nil {
return err
}
if ep.TimeoutMS < MinTimeoutMS || ep.TimeoutMS > MaxTimeoutMS {
return infraerrors.BadRequest("prompt_audit_invalid_timeout", "审计节点超时超出允许范围")
}
if ep.InputLimit < MinInputLimit || ep.InputLimit > MaxInputLimit {
return infraerrors.BadRequest("prompt_audit_invalid_input_limit", "审计节点输入上限超出允许范围")
}
if ep.Enabled {
enabled++
}
}
if cfg.Enabled && enabled == 0 {
return infraerrors.BadRequest("prompt_audit_endpoint_required", "启用提示词审计前至少需要启用一个审计节点")
}
return nil
}
func validateUpdateConfigRequest(req UpdateConfigRequest) error {
if strings.TrimSpace(req.Strategy) != "priority" {
return infraerrors.BadRequest("prompt_audit_invalid_strategy", "提示词审计策略仅支持 priority")
}
if req.WorkerCount < 1 || req.WorkerCount > MaxWorkerCount {
return infraerrors.BadRequest("prompt_audit_invalid_worker_count", "Worker 数量超出允许范围")
}
if req.QueueCapacity < 1 || req.QueueCapacity > MaxQueueCapacity {
return infraerrors.BadRequest("prompt_audit_invalid_queue_capacity", "队列容量超出允许范围")
}
if len(req.Scanners) == 0 {
return infraerrors.BadRequest("prompt_audit_scanners_required", "至少需要启用一个风险分类")
}
for _, scanner := range req.Scanners {
if _, ok := ScannerCatalog[NormalizeCategory(scanner)]; !ok {
return infraerrors.BadRequest("prompt_audit_invalid_scanner", "提示词审计风险分类无效")
}
}
if !req.AllGroups {
if len(req.GroupIDs) == 0 {
return infraerrors.BadRequest("prompt_audit_groups_required", "指定分组模式至少需要选择一个分组")
}
for _, groupID := range req.GroupIDs {
if groupID <= 0 {
return infraerrors.BadRequest("prompt_audit_invalid_group", "提示词审计分组 ID 无效")
}
}
}
for _, endpoint := range req.Endpoints {
if endpoint.TimeoutMS < MinTimeoutMS || endpoint.TimeoutMS > MaxTimeoutMS {
return infraerrors.BadRequest("prompt_audit_invalid_timeout", "审计节点超时超出允许范围")
}
if endpoint.InputLimit < MinInputLimit || endpoint.InputLimit > MaxInputLimit {
return infraerrors.BadRequest("prompt_audit_invalid_input_limit", "审计节点输入上限超出允许范围")
}
}
return nil
}
func (cfg ActiveConfig) EffectiveMode() Mode {
if !cfg.RiskControlEnabled || !cfg.Enabled {
return ModeOff
}
if cfg.BlockingEnabled {
return ModeBlocking
}
return ModeAsync
}
func (cfg ActiveConfig) IncludesGroup(groupID *int64) bool {
if cfg.AllGroups {
return true
}
if groupID == nil {
return false
}
i := sort.Search(len(cfg.GroupIDs), func(i int) bool { return cfg.GroupIDs[i] >= *groupID })
return i < len(cfg.GroupIDs) && cfg.GroupIDs[i] == *groupID
}
func (cfg ActiveConfig) EnabledEndpoints() []ActiveEndpoint {
result := make([]ActiveEndpoint, 0, len(cfg.Endpoints))
for _, ep := range cfg.Endpoints {
if ep.Enabled {
result = append(result, ep)
}
}
return result
}
// InvalidTokenEndpointIDs lists endpoints whose stored token could not be
// decrypted with the current encryption key.
func (cfg ActiveConfig) InvalidTokenEndpointIDs() []string {
ids := make([]string, 0)
for _, ep := range cfg.Endpoints {
if ep.TokenInvalid {
ids = append(ids, ep.ID)
}
}
return ids
}
func PublicFromStorage(cfg storageConfig, riskControlEnabled bool, invalidTokenEndpointIDs []string) PublicConfig {
invalid := make(map[string]struct{}, len(invalidTokenEndpointIDs))
for _, id := range invalidTokenEndpointIDs {
invalid[id] = struct{}{}
}
scanners := append([]string{}, cfg.Scanners...)
groupIDs := append([]int64{}, cfg.GroupIDs...)
endpoints := make([]PublicEndpoint, 0, len(cfg.Endpoints))
for _, ep := range cfg.Endpoints {
hasToken := strings.TrimSpace(ep.TokenCiphertext) != ""
status := "missing"
if hasToken {
status = "configured"
if _, ok := invalid[ep.ID]; ok {
status = "invalid"
}
}
endpoints = append(endpoints, PublicEndpoint{
ID: ep.ID, Name: ep.Name, Protocol: ep.Protocol, BaseURL: ep.BaseURL,
Model: ep.Model, TimeoutMS: ep.TimeoutMS, InputLimit: ep.InputLimit,
Enabled: ep.Enabled, HasToken: hasToken, TokenStatus: status,
})
}
active := ActiveConfig{RiskControlEnabled: riskControlEnabled, Enabled: cfg.Enabled, BlockingEnabled: cfg.BlockingEnabled}
return PublicConfig{
Enabled: cfg.Enabled, BlockingEnabled: cfg.BlockingEnabled, BlockingLatestTurnOnly: cfg.BlockingLatestTurnOnly, StorePassEvents: cfg.StorePassEvents,
EffectiveMode: active.EffectiveMode(), Strategy: cfg.Strategy, WorkerCount: cfg.WorkerCount,
QueueCapacity: cfg.QueueCapacity, Scanners: scanners, AllGroups: cfg.AllGroups,
GroupIDs: groupIDs, Endpoints: endpoints, ConfigVersion: cfg.ConfigVersion,
UpdatedAt: cfg.UpdatedAt, UpdatedBy: cfg.UpdatedBy, ChangeSummary: cfg.ChangeSummary,
}
}
func ActiveFromStorage(cfg storageConfig, riskControlEnabled bool, encryptor SecretEncryptor) (ActiveConfig, error) {
active := ActiveConfig{
RiskControlEnabled: riskControlEnabled, Enabled: cfg.Enabled, BlockingEnabled: cfg.BlockingEnabled,
BlockingLatestTurnOnly: cfg.BlockingLatestTurnOnly,
StorePassEvents: cfg.StorePassEvents, Strategy: cfg.Strategy, WorkerCount: cfg.WorkerCount,
QueueCapacity: cfg.QueueCapacity, Scanners: append([]string(nil), cfg.Scanners...), AllGroups: cfg.AllGroups,
GroupIDs: append([]int64(nil), cfg.GroupIDs...), ConfigVersion: cfg.ConfigVersion,
UpdatedAt: cfg.UpdatedAt, UpdatedBy: cfg.UpdatedBy, ChangeSummary: cfg.ChangeSummary,
Endpoints: make([]ActiveEndpoint, 0, len(cfg.Endpoints)),
}
for _, ep := range cfg.Endpoints {
token := ""
tokenInvalid := false
if ep.TokenCiphertext != "" {
if encryptor == nil {
return ActiveConfig{}, fmt.Errorf("prompt audit secret encryptor unavailable")
}
plain, err := encryptor.Decrypt(ep.TokenCiphertext)
if err != nil {
// An undecryptable token (encryption key changed or regenerated)
// must not take the whole config down: admins would otherwise be
// locked out of the real config version and unable to recover
// (issue #4887). Keep the ciphertext persisted, but exclude the
// endpoint from runtime use until the token is re-entered.
tokenInvalid = true
} else {
token = plain
}
}
active.Endpoints = append(active.Endpoints, ActiveEndpoint{
ID: ep.ID, Name: ep.Name, Protocol: ep.Protocol, BaseURL: ep.BaseURL, Model: ep.Model,
Token: token, TimeoutMS: ep.TimeoutMS, InputLimit: ep.InputLimit,
Enabled: ep.Enabled && !tokenInvalid, TokenInvalid: tokenInvalid,
})
}
return active, nil
}
func changeSummary(cfg storageConfig) string {
summary := struct {
Enabled bool `json:"enabled"`
BlockingEnabled bool `json:"blocking_enabled"`
BlockingLatestTurnOnly bool `json:"blocking_latest_turn_only"`
StorePassEvents bool `json:"store_pass_events"`
EndpointCount int `json:"endpoint_count"`
ScannerCount int `json:"scanner_count"`
AllGroups bool `json:"all_groups"`
GroupCount int `json:"group_count"`
GroupHash string `json:"group_hash"`
}{cfg.Enabled, cfg.BlockingEnabled, cfg.BlockingLatestTurnOnly, cfg.StorePassEvents, len(cfg.Endpoints), len(cfg.Scanners), cfg.AllGroups, len(cfg.GroupIDs), ""}
rawGroups, _ := json.Marshal(cfg.GroupIDs)
digest := sha256.Sum256(rawGroups)
summary.GroupHash = hex.EncodeToString(digest[:])
raw, _ := json.Marshal(summary)
return string(raw)
}
func canonicalInt64s(values []int64) []int64 {
seen := make(map[int64]struct{}, len(values))
result := make([]int64, 0, len(values))
for _, value := range values {
if value <= 0 {
continue
}
if _, ok := seen[value]; ok {
continue
}
seen[value] = struct{}{}
result = append(result, value)
}
sort.Slice(result, func(i, j int) bool { return result[i] < result[j] })
return result
}
func canonicalScannerIDs(values []string) []string {
seen := make(map[string]struct{}, len(values))
for _, value := range values {
id := NormalizeCategory(value)
if _, ok := ScannerCatalog[id]; ok {
seen[id] = struct{}{}
}
}
result := make([]string, 0, len(seen))
for _, id := range AllScannerIDs {
if _, ok := seen[id]; ok {
result = append(result, id)
}
}
return result
}
@@ -0,0 +1,244 @@
package securityaudit
import (
"context"
"database/sql"
"encoding/json"
"errors"
"os"
"strings"
"sync"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/repository"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/lib/pq"
"github.com/redis/go-redis/v9"
"github.com/stretchr/testify/require"
)
const promptAuditRedisTestEnv = "PROMPT_AUDIT_TEST_REDIS_ADDR"
type postgresPromptAuditSettingRepository struct{ db *sql.DB }
func (r postgresPromptAuditSettingRepository) Get(ctx context.Context, key string) (*service.Setting, error) {
var value string
var updated time.Time
err := r.db.QueryRowContext(ctx, `SELECT value,updated_at FROM settings WHERE key=$1`, key).Scan(&value, &updated)
if errors.Is(err, sql.ErrNoRows) {
return nil, service.ErrSettingNotFound
}
if err != nil {
return nil, err
}
return &service.Setting{Key: key, Value: value, UpdatedAt: updated}, nil
}
func (r postgresPromptAuditSettingRepository) GetValue(ctx context.Context, key string) (string, error) {
setting, err := r.Get(ctx, key)
if err != nil {
return "", err
}
return setting.Value, nil
}
func (r postgresPromptAuditSettingRepository) Set(ctx context.Context, key, value string) error {
_, err := r.db.ExecContext(ctx, `INSERT INTO settings(key,value,updated_at) VALUES($1,$2,NOW())
ON CONFLICT(key) DO UPDATE SET value=EXCLUDED.value,updated_at=EXCLUDED.updated_at`, key, value)
return err
}
func (r postgresPromptAuditSettingRepository) GetMultiple(ctx context.Context, keys []string) (map[string]string, error) {
result := make(map[string]string, len(keys))
for _, key := range keys {
result[key] = ""
}
rows, err := r.db.QueryContext(ctx, `SELECT key,value FROM settings WHERE key=ANY($1)`, pq.Array(keys))
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
for rows.Next() {
var key, value string
if err := rows.Scan(&key, &value); err != nil {
return nil, err
}
result[key] = value
}
return result, rows.Err()
}
func (r postgresPromptAuditSettingRepository) SetMultiple(ctx context.Context, values map[string]string) error {
tx, err := r.db.BeginTx(ctx, nil)
if err != nil {
return err
}
defer func() { _ = tx.Rollback() }()
for key, value := range values {
if _, err := tx.ExecContext(ctx, `INSERT INTO settings(key,value,updated_at) VALUES($1,$2,NOW())
ON CONFLICT(key) DO UPDATE SET value=EXCLUDED.value,updated_at=EXCLUDED.updated_at`, key, value); err != nil {
return err
}
}
return tx.Commit()
}
func (r postgresPromptAuditSettingRepository) GetAll(ctx context.Context) (map[string]string, error) {
rows, err := r.db.QueryContext(ctx, `SELECT key,value FROM settings`)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
result := map[string]string{}
for rows.Next() {
var key, value string
if err := rows.Scan(&key, &value); err != nil {
return nil, err
}
result[key] = value
}
return result, rows.Err()
}
func (r postgresPromptAuditSettingRepository) Delete(ctx context.Context, key string) error {
_, err := r.db.ExecContext(ctx, `DELETE FROM settings WHERE key=$1`, key)
return err
}
func promptAuditTestEncryptor(t *testing.T) service.SecretEncryptor {
t.Helper()
encryptor, err := repository.NewAESEncryptor(&config.Config{Totp: config.TotpConfig{EncryptionKey: strings.Repeat("42", 32)}})
require.NoError(t, err)
return encryptor
}
func promptAuditUpdateRequest(version int64, workerCount int, token string) UpdateConfigRequest {
return UpdateConfigRequest{
ExpectedConfigVersion: version, Enabled: true, BlockingEnabled: false, StorePassEvents: false,
Strategy: "priority", WorkerCount: workerCount, QueueCapacity: 64, Scanners: []string{"pii", "jailbreak"},
AllGroups: true, Endpoints: []UpdateEndpoint{{
ID: "guard-one", Name: "Guard One", Protocol: "openai_compatible",
BaseURL: "http://127.0.0.1:18080", Model: "", Token: token,
TimeoutMS: 1000, InputLimit: 1024, Enabled: true,
}},
}
}
func waitForConfigVersion(t *testing.T, manager *ConfigManager, version int64, timeout time.Duration) {
t.Helper()
require.Eventually(t, func() bool {
active, ok := manager.Active()
return ok && active.ConfigVersion == version
}, timeout, 20*time.Millisecond)
}
func TestPromptAuditConfigCASSecretRoundTripInvalidationAndTTL(t *testing.T) {
redisAddress := strings.TrimSpace(os.Getenv(promptAuditRedisTestEnv))
if redisAddress == "" {
t.Skip(promptAuditRedisTestEnv + " is not set")
}
db := openPromptAuditIntegrationDB(t)
settingRepo := postgresPromptAuditSettingRepository{db: db}
require.NoError(t, settingRepo.Set(context.Background(), SettingKeyRiskControl, "true"))
encryptor := promptAuditTestEncryptor(t)
redisClient := redis.NewClient(&redis.Options{Addr: redisAddress})
t.Cleanup(func() { require.NoError(t, redisClient.Close()) })
require.NoError(t, redisClient.Ping(context.Background()).Err())
managerOne := NewConfigManager(db, settingRepo, redisClient, encryptor, testTotpKeyConfig())
managerTwo := NewConfigManager(db, settingRepo, redisClient, encryptor, testTotpKeyConfig())
ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)
require.NoError(t, managerOne.Start(ctx))
require.NoError(t, managerTwo.Start(ctx))
t.Cleanup(func() {
require.NoError(t, managerOne.Shutdown(context.Background()))
require.NoError(t, managerTwo.Shutdown(context.Background()))
})
require.Eventually(t, func() bool {
return redisClient.PubSubNumSub(context.Background(), ConfigInvalidationChannel).Val()[ConfigInvalidationChannel] >= 2
}, 2*time.Second, 20*time.Millisecond)
const canary = "GUARD_TOKEN_CANARY_SECRET_4_CONFIG"
public, err := managerOne.Save(context.Background(), promptAuditUpdateRequest(1, 1, canary), 101)
require.NoError(t, err)
require.Equal(t, int64(2), public.ConfigVersion)
require.True(t, public.Endpoints[0].HasToken)
publicJSON, err := json.Marshal(public)
require.NoError(t, err)
require.NotContains(t, string(publicJSON), canary)
waitForConfigVersion(t, managerTwo, 2, 2*time.Second)
raw, err := settingRepo.GetValue(context.Background(), SettingKeyPromptAuditConfig)
require.NoError(t, err)
require.NotContains(t, raw, canary)
stored, err := ParseStorageConfig(raw)
require.NoError(t, err)
require.NotEmpty(t, stored.Endpoints[0].TokenCiphertext)
plain, err := encryptor.Decrypt(stored.Endpoints[0].TokenCiphertext)
require.NoError(t, err)
require.Equal(t, canary, plain)
require.NotContains(t, stored.ChangeSummary, canary)
require.NotContains(t, stored.ChangeSummary, stored.Endpoints[0].BaseURL)
type saveResult struct {
config PublicConfig
err error
}
start := make(chan struct{})
results := make(chan saveResult, 2)
var wg sync.WaitGroup
for index, manager := range []*ConfigManager{managerOne, managerTwo} {
wg.Add(1)
go func(index int, manager *ConfigManager) {
defer wg.Done()
<-start
cfg, saveErr := manager.Save(context.Background(), promptAuditUpdateRequest(2, index+2, ""), int64(201+index))
results <- saveResult{config: cfg, err: saveErr}
}(index, manager)
}
close(start)
wg.Wait()
close(results)
succeeded, conflicted := 0, 0
for result := range results {
if result.err == nil {
succeeded++
require.Equal(t, int64(3), result.config.ConfigVersion)
continue
}
conflicted++
require.Equal(t, ErrorCodeConfigConflict, infraerrors.Reason(result.err))
}
require.Equal(t, 1, succeeded)
require.Equal(t, 1, conflicted)
waitForConfigVersion(t, managerOne, 3, 2*time.Second)
waitForConfigVersion(t, managerTwo, 3, 2*time.Second)
// A manager without Redis subscriptions must still converge through the
// bounded five-second refresh loop.
ttlManager := NewConfigManager(db, settingRepo, nil, encryptor, testTotpKeyConfig())
require.NoError(t, ttlManager.Start(ctx))
t.Cleanup(func() { require.NoError(t, ttlManager.Shutdown(context.Background())) })
waitForConfigVersion(t, ttlManager, 3, time.Second)
updated, err := managerOne.Save(context.Background(), promptAuditUpdateRequest(3, 5, ""), 301)
require.NoError(t, err)
require.Equal(t, int64(4), updated.ConfigVersion)
waitForConfigVersion(t, ttlManager, 4, 7*time.Second)
// Redis publication failure is observable degradation, not a rollback of a
// successfully committed PostgreSQL config.
deadRedis := redis.NewClient(&redis.Options{Addr: "127.0.0.1:1", MaxRetries: 0, DialTimeout: 30 * time.Millisecond, ReadTimeout: 30 * time.Millisecond, WriteTimeout: 30 * time.Millisecond})
t.Cleanup(func() { _ = deadRedis.Close() })
degraded := NewConfigManager(db, settingRepo, deadRedis, encryptor, testTotpKeyConfig())
require.NoError(t, degraded.Reload(context.Background()))
degradedSaved, err := degraded.Save(context.Background(), promptAuditUpdateRequest(4, 6, ""), 401)
require.NoError(t, err)
require.Equal(t, int64(5), degradedSaved.ConfigVersion)
active, ok := degraded.Active()
require.True(t, ok)
require.Equal(t, int64(5), active.ConfigVersion)
}
@@ -0,0 +1,512 @@
package securityaudit
import (
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/redis/go-redis/v9"
)
type activeConfigSnapshot struct {
storage storageConfig
active ActiveConfig
loadedAt time.Time
}
type ConfigManager struct {
db *sql.DB
settings service.SettingRepository
redis *redis.Client
encryptor SecretEncryptor
clock Clock
// encryptionKeyConfigured mirrors cfg.Totp.EncryptionKeyConfigured. With an
// auto-generated (per-boot) key, newly saved endpoint tokens would become
// undecryptable after the next restart, so Save rejects them (issue #4887).
encryptionKeyConfigured bool
snapshot atomic.Pointer[activeConfigSnapshot]
expected atomic.Int64
// expectedBlocking records the last storage intent that could be decoded,
// independently of whether endpoint credentials or the full config could be
// activated. A config version alone cannot distinguish async from blocking.
expectedBlocking atomic.Bool
// configUntrusted is set when a load/reload fails before a trustworthy
// snapshot is installed. Combined with expectedBlocking, EffectiveMode
// fails closed so a persisted blocking policy cannot be silently skipped
// after startup or invalidation errors. Without blocking intent, untrusted
// alone must not force ModeBlocking—Prompt Audit is default-off and must
// not take the gateway down for every API request (see issue #4560).
configUntrusted atomic.Bool
stateMu sync.RWMutex
lastLoadError string
lastErrorAt *time.Time
lifecycleMu sync.Mutex
cancel context.CancelFunc
wg sync.WaitGroup
}
func NewConfigManager(db *sql.DB, settings service.SettingRepository, redisClient *redis.Client, encryptor service.SecretEncryptor, cfg *config.Config) *ConfigManager {
return &ConfigManager{
db: db, settings: settings, redis: redisClient, encryptor: encryptor, clock: realClock{},
encryptionKeyConfigured: cfg != nil && cfg.Totp.EncryptionKeyConfigured,
}
}
func (m *ConfigManager) Start(ctx context.Context) error {
if m == nil {
return errors.New("prompt audit config manager unavailable")
}
m.lifecycleMu.Lock()
if m.cancel != nil {
m.lifecycleMu.Unlock()
return nil
}
runCtx, cancel := context.WithCancel(ctx)
m.cancel = cancel
m.lifecycleMu.Unlock()
loadErr := m.Reload(runCtx)
if loadErr != nil {
m.markConfigUntrusted()
}
m.wg.Add(1)
go m.refreshLoop(runCtx)
if m.redis != nil {
m.wg.Add(1)
go m.subscribeLoop(runCtx)
}
return loadErr
}
func (m *ConfigManager) Shutdown(_ context.Context) error {
if m == nil {
return nil
}
m.lifecycleMu.Lock()
cancel := m.cancel
m.cancel = nil
m.lifecycleMu.Unlock()
if cancel != nil {
cancel()
}
m.wg.Wait()
return nil
}
func (m *ConfigManager) Reload(ctx context.Context) error {
if m == nil || m.settings == nil {
m.markUntrustedIfNoActiveSnapshot()
return errors.New("prompt audit setting repository unavailable")
}
values, err := m.settings.GetMultiple(ctx, []string{SettingKeyPromptAuditConfig, SettingKeyRiskControl})
if err != nil {
m.recordLoadError(err)
m.markUntrustedIfNoActiveSnapshot()
return err
}
m.observeExpectedState(values[SettingKeyPromptAuditConfig], values[SettingKeyRiskControl] == "true")
storage, err := ParseStorageConfig(values[SettingKeyPromptAuditConfig])
if err != nil {
m.recordLoadError(err)
m.markUntrustedIfNoActiveSnapshot()
return err
}
m.expected.Store(storage.ConfigVersion)
m.expectedBlocking.Store(values[SettingKeyRiskControl] == "true" && storage.Enabled && storage.BlockingEnabled)
active, err := ActiveFromStorage(storage, values[SettingKeyRiskControl] == "true", m.encryptor)
if err != nil {
m.recordLoadError(err)
// expectedBlocking may already require fail-closed via BlockingActivationDegraded.
m.markUntrustedIfNoActiveSnapshot()
return err
}
now := m.clock.Now()
previous := m.snapshot.Load()
m.snapshot.Store(&activeConfigSnapshot{storage: cloneStorageConfig(storage), active: cloneActiveConfig(active), loadedAt: now})
m.configUntrusted.Store(false)
m.clearLoadError()
m.logInvalidTokenEndpoints(previous, active)
LogInfo(EventConfigLoaded, map[string]any{
"config_version": storage.ConfigVersion, "status": "loaded",
})
return nil
}
// logInvalidTokenEndpoints warns once per change (not on every 5s refresh)
// when stored endpoint tokens cannot be decrypted with the current key.
func (m *ConfigManager) logInvalidTokenEndpoints(previous *activeConfigSnapshot, active ActiveConfig) {
invalid := active.InvalidTokenEndpointIDs()
if len(invalid) == 0 {
return
}
if previous != nil {
prior := previous.active.InvalidTokenEndpointIDs()
if len(prior) == len(invalid) {
same := true
for i := range invalid {
if prior[i] != invalid[i] {
same = false
break
}
}
if same && previous.active.ConfigVersion == active.ConfigVersion {
return
}
}
}
LogWarn(EventConfigTokenInvalid, map[string]any{
"config_version": active.ConfigVersion, "status": "degraded",
"error_code": "endpoint_token_undecryptable", "guard_endpoint_id": strings.Join(invalid, ","),
})
}
func (m *ConfigManager) Active() (ActiveConfig, bool) {
if m == nil {
return ActiveConfig{}, false
}
snapshot := m.snapshot.Load()
if snapshot == nil {
return ActiveConfig{}, false
}
return cloneActiveConfig(snapshot.active), true
}
func (m *ConfigManager) BlockingActivationDegraded() bool {
if m == nil {
return false
}
// Fail closed only when storage intent requires blocking. Untrusted config
// without blocking intent must remain ModeOff so administrators can still
// operate the gateway and turn Prompt Audit off after a failed reload.
if !m.expectedBlocking.Load() {
return false
}
if m.configUntrusted.Load() {
return true
}
active, ok := m.Active()
if !ok {
return true
}
// A still-active weaker snapshot after a failed blocking activation must not
// keep serving allow decisions under the old off/async mode.
return active.EffectiveMode() != ModeBlocking
}
func (m *ConfigManager) EffectiveMode() Mode {
if m != nil && m.BlockingActivationDegraded() {
return ModeBlocking
}
active, ok := m.Active()
if !ok {
return ModeOff
}
return active.EffectiveMode()
}
func (m *ConfigManager) markConfigUntrusted() {
if m == nil {
return
}
m.configUntrusted.Store(true)
}
func (m *ConfigManager) markUntrustedIfNoActiveSnapshot() {
if m == nil {
return
}
if _, ok := m.Active(); !ok {
m.markConfigUntrusted()
}
}
func (m *ConfigManager) Public() (PublicConfig, error) {
if m == nil {
return PublicConfig{}, infraerrors.ServiceUnavailable(ErrorCodeConfigUnavailable, "提示词审计配置暂不可用")
}
snapshot := m.snapshot.Load()
if snapshot == nil {
return PublicConfig{}, infraerrors.ServiceUnavailable(ErrorCodeConfigUnavailable, "提示词审计配置暂不可用")
}
return PublicFromStorage(cloneStorageConfig(snapshot.storage), snapshot.active.RiskControlEnabled, snapshot.active.InvalidTokenEndpointIDs()), nil
}
func (m *ConfigManager) Save(ctx context.Context, req UpdateConfigRequest, actorID int64) (PublicConfig, error) {
if m == nil || m.db == nil || m.encryptor == nil {
return PublicConfig{}, errors.New("prompt audit config persistence unavailable")
}
if req.ExpectedConfigVersion < 1 {
return PublicConfig{}, infraerrors.BadRequest("prompt_audit_expected_config_version_required", "必须提供有效的配置版本")
}
tx, err := m.db.BeginTx(ctx, &sql.TxOptions{Isolation: sql.LevelReadCommitted})
if err != nil {
return PublicConfig{}, err
}
defer func() { _ = tx.Rollback() }()
if _, err := tx.ExecContext(ctx, `SELECT pg_advisory_xact_lock($1)`, promptAuditConfigLockKey); err != nil {
return PublicConfig{}, err
}
current := DefaultStorageConfig()
var raw string
err = tx.QueryRowContext(ctx, `SELECT value FROM settings WHERE key=$1 FOR UPDATE`, SettingKeyPromptAuditConfig).Scan(&raw)
if err != nil && !errors.Is(err, sql.ErrNoRows) {
return PublicConfig{}, err
}
if err == nil {
current, err = ParseStorageConfig(raw)
if err != nil {
return PublicConfig{}, err
}
}
if current.ConfigVersion != req.ExpectedConfigVersion {
return PublicConfig{}, infraerrors.Conflict(ErrorCodeConfigConflict, "提示词审计配置已被其他管理员更新")
}
next, err := m.buildNextStorage(current, req, actorID)
if err != nil {
return PublicConfig{}, err
}
next.ConfigVersion = current.ConfigVersion + 1
next.UpdatedAt = m.clock.Now()
next.UpdatedBy = actorID
next.ChangeSummary = changeSummary(next)
rawNext, err := json.Marshal(next)
if err != nil {
return PublicConfig{}, err
}
if _, err := tx.ExecContext(ctx, `
INSERT INTO settings (key,value,updated_at) VALUES ($1,$2,NOW())
ON CONFLICT (key) DO UPDATE SET value=EXCLUDED.value, updated_at=EXCLUDED.updated_at`,
SettingKeyPromptAuditConfig, string(rawNext)); err != nil {
return PublicConfig{}, err
}
if err := tx.Commit(); err != nil {
return PublicConfig{}, err
}
// Install the snapshot with the current global gate, not merely the value
// cached when this process last reloaded Prompt Audit configuration.
riskControlEnabled := m.currentRiskControlEnabled()
if values, getErr := m.settings.GetMultiple(ctx, []string{SettingKeyRiskControl}); getErr == nil {
riskControlEnabled = values[SettingKeyRiskControl] == "true"
}
active, err := ActiveFromStorage(next, riskControlEnabled, m.encryptor)
if err != nil {
return PublicConfig{}, err
}
m.expected.Store(next.ConfigVersion)
m.expectedBlocking.Store(active.RiskControlEnabled && next.Enabled && next.BlockingEnabled)
previous := m.snapshot.Load()
m.snapshot.Store(&activeConfigSnapshot{storage: cloneStorageConfig(next), active: cloneActiveConfig(active), loadedAt: m.clock.Now()})
// A successful admin save installs a trustworthy snapshot; clear any prior
// fail-closed degradation so disabling audit actually takes effect.
m.configUntrusted.Store(false)
m.clearLoadError()
m.logInvalidTokenEndpoints(previous, active)
LogInfo(EventConfigUpdated, map[string]any{
"config_version": next.ConfigVersion, "status": "updated",
})
if m.redis != nil {
if err := m.redis.Publish(ctx, ConfigInvalidationChannel, strconv.FormatInt(next.ConfigVersion, 10)).Err(); err != nil {
LogWarn(EventConfigReloadDegraded, map[string]any{
"config_version": next.ConfigVersion, "status": "degraded", "error_code": "config_invalidation_publish_failed",
})
}
}
return PublicFromStorage(next, active.RiskControlEnabled, active.InvalidTokenEndpointIDs()), nil
}
func (m *ConfigManager) buildNextStorage(current storageConfig, req UpdateConfigRequest, actorID int64) (storageConfig, error) {
if err := validateUpdateConfigRequest(req); err != nil {
return storageConfig{}, err
}
currentByID := make(map[string]StorageEndpoint, len(current.Endpoints))
for _, endpoint := range current.Endpoints {
currentByID[endpoint.ID] = endpoint
}
next := storageConfig{
Enabled: req.Enabled, BlockingEnabled: req.BlockingEnabled, BlockingLatestTurnOnly: req.BlockingLatestTurnOnly, StorePassEvents: req.StorePassEvents,
Strategy: strings.TrimSpace(req.Strategy), WorkerCount: req.WorkerCount,
QueueCapacity: req.QueueCapacity, Scanners: append([]string(nil), req.Scanners...),
AllGroups: req.AllGroups, GroupIDs: append([]int64(nil), req.GroupIDs...),
ConfigVersion: current.ConfigVersion, UpdatedBy: actorID,
Endpoints: make([]StorageEndpoint, 0, len(req.Endpoints)),
}
for _, endpoint := range req.Endpoints {
baseURL, err := NormalizeBaseURL(endpoint.BaseURL)
if err != nil {
return storageConfig{}, err
}
stored := StorageEndpoint{
ID: strings.TrimSpace(endpoint.ID), Name: strings.TrimSpace(endpoint.Name),
Protocol: strings.TrimSpace(endpoint.Protocol), BaseURL: baseURL, Model: strings.TrimSpace(endpoint.Model),
TimeoutMS: endpoint.TimeoutMS, InputLimit: endpoint.InputLimit, Enabled: endpoint.Enabled,
}
old, hadOld := currentByID[stored.ID]
switch {
case endpoint.ClearToken:
stored.TokenCiphertext = ""
case strings.TrimSpace(endpoint.Token) != "":
if !m.encryptionKeyConfigured {
return storageConfig{}, infraerrors.BadRequest(ErrorCodeEncryptionKeyRequired,
"未配置固定加密密钥,审计节点 Token 将在服务重启后失效。请先设置 TOTP_ENCRYPTION_KEY 环境变量(64 位十六进制)并重启服务")
}
ciphertext, err := m.encryptor.Encrypt(strings.TrimSpace(endpoint.Token))
if err != nil {
return storageConfig{}, fmt.Errorf("encrypt prompt audit endpoint token: %w", err)
}
stored.TokenCiphertext = ciphertext
case hadOld:
stored.TokenCiphertext = old.TokenCiphertext
}
next.Endpoints = append(next.Endpoints, stored)
}
normalizeStorageConfig(&next)
if err := validateStorageConfig(next); err != nil {
return storageConfig{}, err
}
return next, nil
}
func (m *ConfigManager) RuntimeState() (expected int64, active int64, loadedAt *time.Time, loadError string) {
if m == nil {
return 1, 0, nil, "config_manager_unavailable"
}
expected = m.expected.Load()
if expected < 1 {
expected = 1
}
if snapshot := m.snapshot.Load(); snapshot != nil {
active = snapshot.active.ConfigVersion
value := snapshot.loadedAt
loadedAt = &value
}
m.stateMu.RLock()
loadError = m.lastLoadError
m.stateMu.RUnlock()
return
}
func (m *ConfigManager) Encrypt(value string) (string, error) { return m.encryptor.Encrypt(value) }
func (m *ConfigManager) Decrypt(value string) (string, error) { return m.encryptor.Decrypt(value) }
func (m *ConfigManager) currentRiskControlEnabled() bool {
if snapshot := m.snapshot.Load(); snapshot != nil {
return snapshot.active.RiskControlEnabled
}
return false
}
func (m *ConfigManager) observeExpectedState(raw string, riskControlEnabled bool) {
if m == nil {
return
}
if strings.TrimSpace(raw) == "" {
m.expected.Store(1)
m.expectedBlocking.Store(false)
return
}
var intent struct {
Enabled bool `json:"enabled"`
BlockingEnabled bool `json:"blocking_enabled"`
ConfigVersion int64 `json:"config_version"`
}
if err := json.Unmarshal([]byte(raw), &intent); err != nil {
return
}
if intent.ConfigVersion < 1 {
intent.ConfigVersion = 1
}
m.expected.Store(intent.ConfigVersion)
m.expectedBlocking.Store(riskControlEnabled && intent.Enabled && intent.BlockingEnabled)
}
func (m *ConfigManager) refreshLoop(ctx context.Context) {
defer m.wg.Done()
ticker := time.NewTicker(5 * time.Second)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
if err := m.Reload(ctx); err != nil {
LogWarn(EventConfigReloadDegraded, map[string]any{"status": "degraded", "error_code": "config_ttl_reload_failed"})
}
}
}
}
func (m *ConfigManager) subscribeLoop(ctx context.Context) {
defer m.wg.Done()
pubsub := m.redis.Subscribe(ctx, ConfigInvalidationChannel)
defer func() { _ = pubsub.Close() }()
channel := pubsub.Channel()
for {
select {
case <-ctx.Done():
return
case message, ok := <-channel:
if !ok {
return
}
version, err := strconv.ParseInt(strings.TrimSpace(message.Payload), 10, 64)
if err != nil || version < 1 {
continue
}
m.expected.Store(version)
if err := m.Reload(ctx); err != nil {
// A newer published version failed to activate. Until reload
// succeeds, do not keep serving a potentially stale weaker mode.
if active, ok := m.Active(); !ok || active.ConfigVersion < version {
m.markConfigUntrusted()
}
LogWarn(EventConfigReloadDegraded, map[string]any{
"config_version": version, "status": "degraded", "error_code": "config_invalidation_reload_failed",
})
}
}
}
}
func (m *ConfigManager) recordLoadError(_ error) {
if m == nil {
return
}
now := m.clock.Now()
m.stateMu.Lock()
m.lastLoadError = stableErrorMessage("config_load_failed")
m.lastErrorAt = &now
m.stateMu.Unlock()
}
func (m *ConfigManager) clearLoadError() {
m.stateMu.Lock()
m.lastLoadError = ""
m.lastErrorAt = nil
m.stateMu.Unlock()
}
func cloneStorageConfig(cfg storageConfig) storageConfig {
cfg.Scanners = append([]string(nil), cfg.Scanners...)
cfg.GroupIDs = append([]int64(nil), cfg.GroupIDs...)
cfg.Endpoints = append([]StorageEndpoint(nil), cfg.Endpoints...)
return cfg
}
func cloneActiveConfig(cfg ActiveConfig) ActiveConfig {
cfg.Scanners = append([]string(nil), cfg.Scanners...)
cfg.GroupIDs = append([]int64(nil), cfg.GroupIDs...)
cfg.Endpoints = append([]ActiveEndpoint(nil), cfg.Endpoints...)
return cfg
}
@@ -0,0 +1,456 @@
package securityaudit
import (
"context"
"encoding/json"
"errors"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/stretchr/testify/require"
)
type prefixEncryptor struct{}
func (prefixEncryptor) Encrypt(value string) (string, error) { return "enc:" + value, nil }
func (prefixEncryptor) Decrypt(value string) (string, error) {
if !strings.HasPrefix(value, "enc:") {
return "", errors.New("cipher: message authentication failed")
}
return value[4:], nil
}
// testTotpKeyConfig mirrors a deployment with a fixed TOTP_ENCRYPTION_KEY so
// unit tests may persist endpoint tokens.
func testTotpKeyConfig() *config.Config {
return &config.Config{Totp: config.TotpConfig{EncryptionKeyConfigured: true}}
}
func TestDefaultConfigIsOff(t *testing.T) {
storage, err := ParseStorageConfig("")
require.NoError(t, err)
require.False(t, storage.Enabled)
require.False(t, storage.BlockingLatestTurnOnly)
active, err := ActiveFromStorage(storage, true, prefixEncryptor{})
require.NoError(t, err)
require.Equal(t, ModeOff, active.EffectiveMode())
require.Equal(t, AllScannerIDs, storage.Scanners)
publicJSON, err := json.Marshal(PublicFromStorage(storage, true, nil))
require.NoError(t, err)
require.Contains(t, string(publicJSON), `"group_ids":[]`)
require.Contains(t, string(publicJSON), `"endpoints":[]`)
}
func TestBlockingLatestTurnOnlyConfigRoundTrip(t *testing.T) {
manager := &ConfigManager{encryptor: prefixEncryptor{}, encryptionKeyConfigured: true}
request := UpdateConfigRequest{
ExpectedConfigVersion: 1, Enabled: true, BlockingEnabled: true, BlockingLatestTurnOnly: true,
Strategy: "priority", WorkerCount: 1, QueueCapacity: 10, Scanners: []string{"pii"}, AllGroups: true,
Endpoints: []UpdateEndpoint{{
ID: "guard-1", Name: "Guard", Protocol: "openai_compatible", BaseURL: "http://127.0.0.1:8080",
Model: DefaultGuardModel, TimeoutMS: 1000, InputLimit: 1000, Enabled: true,
}},
}
next, err := manager.buildNextStorage(DefaultStorageConfig(), request, 9)
require.NoError(t, err)
require.True(t, next.BlockingLatestTurnOnly)
require.Contains(t, changeSummary(next), `"blocking_latest_turn_only":true`)
active, err := ActiveFromStorage(next, true, prefixEncryptor{})
require.NoError(t, err)
require.True(t, active.BlockingLatestTurnOnly)
public := PublicFromStorage(next, true, nil)
require.True(t, public.BlockingLatestTurnOnly)
}
func TestConfigRejectsBlockingWithoutAudit(t *testing.T) {
storage := DefaultStorageConfig()
storage.BlockingEnabled = true
require.Error(t, validateStorageConfig(storage))
}
func TestPublicConfigNeverMarshalsToken(t *testing.T) {
storage := DefaultStorageConfig()
storage.Endpoints = []StorageEndpoint{{ID: "one", Name: "One", Protocol: "openai_compatible", BaseURL: "http://127.0.0.1:8080", Model: DefaultGuardModel, TokenCiphertext: "GUARD_TOKEN_CANARY_SECRET", TimeoutMS: 1000, InputLimit: 1000, Enabled: true}}
public := PublicFromStorage(storage, true, nil)
raw, err := json.Marshal(public)
require.NoError(t, err)
require.NotContains(t, string(raw), "GUARD_TOKEN_CANARY_SECRET")
require.NotContains(t, string(raw), "ciphertext")
require.True(t, public.Endpoints[0].HasToken)
}
func TestConfigRuntimeLoadErrorIsStableBoundedAndSecretFree(t *testing.T) {
const canary = "CONFIG_LOAD_CANARY_SECRET"
manager := &ConfigManager{clock: fixedClock{}}
manager.recordLoadError(errors.New("decrypt failed for token " + canary + " Authorization: Bearer " + canary))
_, _, _, message := manager.RuntimeState()
require.Equal(t, stableErrorMessage("config_load_failed"), message)
require.NotContains(t, message, canary)
require.LessOrEqual(t, len([]rune(message)), 160)
}
func TestConfigManagerPublicRequiresSuccessfullyLoadedSnapshot(t *testing.T) {
t.Run("absent persisted setting is legitimate default", func(t *testing.T) {
manager := NewConfigManager(nil, staticSettingRepository{values: map[string]string{
SettingKeyPromptAuditConfig: "",
SettingKeyRiskControl: "false",
}}, nil, prefixEncryptor{}, testTotpKeyConfig())
require.NoError(t, manager.Reload(context.Background()))
public, err := manager.Public()
require.NoError(t, err)
require.Equal(t, int64(1), public.ConfigVersion)
require.False(t, public.Enabled)
})
t.Run("unparseable persisted config is unavailable", func(t *testing.T) {
const canary = "persisted-token-canary"
manager := NewConfigManager(nil, staticSettingRepository{values: map[string]string{
// Endpoint without id/name fails validation, so no trustworthy
// snapshot can be installed from this raw value.
SettingKeyPromptAuditConfig: `{"enabled":true,"config_version":9,"endpoints":[{"token_ciphertext":"` + canary + `"}]}`,
SettingKeyRiskControl: "true",
}}, nil, prefixEncryptor{}, testTotpKeyConfig())
require.Error(t, manager.Reload(context.Background()))
public, err := manager.Public()
require.Error(t, err)
require.Empty(t, public)
require.Equal(t, ErrorCodeConfigUnavailable, infraerrors.Reason(err))
require.NotContains(t, err.Error(), canary)
})
t.Run("reload failure preserves last successfully loaded snapshot", func(t *testing.T) {
storage := DefaultStorageConfig()
storage.ConfigVersion = 4
storage.ChangeSummary = "trusted snapshot"
raw, err := json.Marshal(storage)
require.NoError(t, err)
repository := &switchableSettingRepository{staticSettingRepository: staticSettingRepository{values: map[string]string{
SettingKeyPromptAuditConfig: string(raw),
SettingKeyRiskControl: "false",
}}}
manager := NewConfigManager(nil, repository, nil, prefixEncryptor{}, testTotpKeyConfig())
require.NoError(t, manager.Reload(context.Background()))
repository.loadErr = errors.New("settings unavailable")
require.Error(t, manager.Reload(context.Background()))
public, err := manager.Public()
require.NoError(t, err)
require.Equal(t, int64(4), public.ConfigVersion)
require.Equal(t, "trusted snapshot", public.ChangeSummary)
})
}
// Regression coverage for issue #4887: a persisted config whose endpoint token
// can no longer be decrypted (encryption key changed or auto-generated per
// boot) must stay visible and editable for admins instead of falling back to a
// default v1 config that makes every save fail the CAS version check.
func TestConfigManagerUndecryptableTokenKeepsConfigVisibleAndRecoverable(t *testing.T) {
const canary = "persisted-token-canary"
persisted := `{"enabled":true,"blocking_enabled":false,"config_version":9,"endpoints":[{"id":"g1","name":"Guard","protocol":"openai_compatible","base_url":"http://127.0.0.1:8080","model":"m","token_ciphertext":"` + canary + `","timeout_ms":1000,"input_limit":1000,"enabled":true}]}`
manager := NewConfigManager(nil, staticSettingRepository{values: map[string]string{
SettingKeyPromptAuditConfig: persisted,
SettingKeyRiskControl: "true",
}}, nil, prefixEncryptor{}, testTotpKeyConfig())
require.NoError(t, manager.Reload(context.Background()), "an undecryptable token must not fail the whole config load")
public, err := manager.Public()
require.NoError(t, err)
require.Equal(t, int64(9), public.ConfigVersion, "admins must see the real persisted version so CAS saves can succeed")
require.Len(t, public.Endpoints, 1)
require.True(t, public.Endpoints[0].HasToken)
require.Equal(t, "invalid", public.Endpoints[0].TokenStatus)
raw, err := json.Marshal(public)
require.NoError(t, err)
require.NotContains(t, string(raw), canary)
active, ok := manager.Active()
require.True(t, ok)
require.Len(t, active.Endpoints, 1)
require.False(t, active.Endpoints[0].Enabled, "an endpoint with an undecryptable token must not be used at runtime")
require.True(t, active.Endpoints[0].TokenInvalid)
require.Empty(t, active.Endpoints[0].Token)
require.Empty(t, active.EnabledEndpoints())
require.Equal(t, []string{"g1"}, active.InvalidTokenEndpointIDs())
expected, activeVersion, _, _ := manager.RuntimeState()
require.Equal(t, int64(9), expected)
require.Equal(t, int64(9), activeVersion)
}
func TestConfigManagerUndecryptableTokenStillFailsClosedForBlockingIntent(t *testing.T) {
persisted := `{"enabled":true,"blocking_enabled":true,"config_version":9,"endpoints":[{"id":"g1","name":"Guard","protocol":"openai_compatible","base_url":"http://127.0.0.1:8080","model":"m","token_ciphertext":"undecryptable","timeout_ms":1000,"input_limit":1000,"enabled":true}]}`
manager := NewConfigManager(nil, staticSettingRepository{values: map[string]string{
SettingKeyPromptAuditConfig: persisted,
SettingKeyRiskControl: "true",
}}, nil, prefixEncryptor{}, testTotpKeyConfig())
require.NoError(t, manager.Reload(context.Background()))
require.Equal(t, ModeBlocking, manager.EffectiveMode())
service := &PromptService{config: manager, evaluator: NewGuardEvaluator(NewOpenAICompatibleScanner(), nil, nil)}
decision, err := service.Evaluate(context.Background(), Request{
Protocol: "openai_chat_completions",
Body: []byte(`{"messages":[{"role":"user","content":"hi"}]}`),
})
require.Error(t, err, "blocking intent with no usable endpoint must not let requests pass unaudited")
require.Nil(t, decision)
var guardErr *GuardError
require.ErrorAs(t, err, &guardErr)
require.Equal(t, ErrorCodeUnavailable, guardErr.Code)
}
func TestBuildNextStoragePreserveReplaceAndClearToken(t *testing.T) {
manager := &ConfigManager{encryptor: prefixEncryptor{}, encryptionKeyConfigured: true}
current := DefaultStorageConfig()
current.Endpoints = []StorageEndpoint{{ID: "one", Name: "One", Protocol: "openai_compatible", BaseURL: "http://127.0.0.1:8080", Model: DefaultGuardModel, TokenCiphertext: "enc:old", TimeoutMS: 1000, InputLimit: 1000}}
base := UpdateConfigRequest{ExpectedConfigVersion: 1, Strategy: "priority", WorkerCount: 1, QueueCapacity: 10, Scanners: []string{"PII"}, AllGroups: true,
Endpoints: []UpdateEndpoint{{ID: "one", Name: "One", Protocol: "openai_compatible", BaseURL: "http://127.0.0.1:8080", TimeoutMS: 1000, InputLimit: 1000}}}
preserved, err := manager.buildNextStorage(current, base, 9)
require.NoError(t, err)
require.Equal(t, "enc:old", preserved.Endpoints[0].TokenCiphertext)
replacedReq := base
replacedReq.Endpoints = append([]UpdateEndpoint(nil), base.Endpoints...)
replacedReq.Endpoints[0].Token = "new"
replaced, err := manager.buildNextStorage(current, replacedReq, 9)
require.NoError(t, err)
require.Equal(t, "enc:new", replaced.Endpoints[0].TokenCiphertext)
clearedReq := base
clearedReq.Endpoints = append([]UpdateEndpoint(nil), base.Endpoints...)
clearedReq.Endpoints[0].ClearToken = true
cleared, err := manager.buildNextStorage(current, clearedReq, 9)
require.NoError(t, err)
require.Empty(t, cleared.Endpoints[0].TokenCiphertext)
}
// Without a fixed encryption key the per-boot auto-generated key would make a
// freshly saved token undecryptable after the next restart (issue #4887), so
// saving a new token must be rejected with an actionable error. Preserving or
// clearing an existing ciphertext stays allowed so admins can still edit or
// disable the feature.
func TestBuildNextStorageRejectsNewTokenWithoutConfiguredEncryptionKey(t *testing.T) {
manager := &ConfigManager{encryptor: prefixEncryptor{}, encryptionKeyConfigured: false}
current := DefaultStorageConfig()
current.Endpoints = []StorageEndpoint{{ID: "one", Name: "One", Protocol: "openai_compatible", BaseURL: "http://127.0.0.1:8080", Model: DefaultGuardModel, TokenCiphertext: "enc:old", TimeoutMS: 1000, InputLimit: 1000}}
base := UpdateConfigRequest{ExpectedConfigVersion: 1, Strategy: "priority", WorkerCount: 1, QueueCapacity: 10, Scanners: []string{"PII"}, AllGroups: true,
Endpoints: []UpdateEndpoint{{ID: "one", Name: "One", Protocol: "openai_compatible", BaseURL: "http://127.0.0.1:8080", TimeoutMS: 1000, InputLimit: 1000}}}
newTokenReq := base
newTokenReq.Endpoints = append([]UpdateEndpoint(nil), base.Endpoints...)
newTokenReq.Endpoints[0].Token = "fresh-token"
_, err := manager.buildNextStorage(current, newTokenReq, 9)
require.Error(t, err)
require.Equal(t, ErrorCodeEncryptionKeyRequired, infraerrors.Reason(err))
preserved, err := manager.buildNextStorage(current, base, 9)
require.NoError(t, err)
require.Equal(t, "enc:old", preserved.Endpoints[0].TokenCiphertext)
clearedReq := base
clearedReq.Endpoints = append([]UpdateEndpoint(nil), base.Endpoints...)
clearedReq.Endpoints[0].ClearToken = true
cleared, err := manager.buildNextStorage(current, clearedReq, 9)
require.NoError(t, err)
require.Empty(t, cleared.Endpoints[0].TokenCiphertext)
}
func TestEffectiveModeTruthTable(t *testing.T) {
tests := []struct {
risk, enabled, blocking bool
want Mode
}{
{false, false, false, ModeOff}, {false, true, true, ModeOff}, {true, false, false, ModeOff},
{true, true, false, ModeAsync}, {true, true, true, ModeBlocking},
}
for _, tt := range tests {
cfg := ActiveConfig{RiskControlEnabled: tt.risk, Enabled: tt.enabled, BlockingEnabled: tt.blocking}
require.Equal(t, tt.want, cfg.EffectiveMode())
}
}
func TestConfigManagerColdStartOnlyFailsClosedForExplicitBlockingIntent(t *testing.T) {
manager := &ConfigManager{}
manager.observeExpectedState(`{"enabled":true,"blocking_enabled":false,"config_version":42}`, true)
require.Equal(t, int64(42), manager.expected.Load())
require.Equal(t, ModeOff, manager.EffectiveMode(), "an async config version must not imply blocking")
require.False(t, manager.BlockingActivationDegraded())
manager.observeExpectedState(`{"enabled":true,"blocking_enabled":true,"config_version":43}`, false)
require.Equal(t, ModeOff, manager.EffectiveMode(), "the global risk-control switch still gates blocking")
manager.observeExpectedState(`{"enabled":true,"blocking_enabled":true,"config_version":44}`, true)
require.Equal(t, ModeBlocking, manager.EffectiveMode())
require.True(t, manager.BlockingActivationDegraded())
manager.observeExpectedState(`{"enabled":true`, true)
require.Equal(t, ModeBlocking, manager.EffectiveMode(), "undecodable storage must not erase the last known strict intent")
}
func TestConfigManagerStaleWeakerSnapshotFailsClosedWhenBlockingExpected(t *testing.T) {
manager := &ConfigManager{}
async := ActiveConfig{RiskControlEnabled: true, Enabled: true, BlockingEnabled: false, ConfigVersion: 1}
manager.snapshot.Store(&activeConfigSnapshot{active: async, storage: DefaultStorageConfig(), loadedAt: fixedClock{}.Now()})
manager.expected.Store(2)
manager.expectedBlocking.Store(true)
require.True(t, manager.BlockingActivationDegraded())
require.Equal(t, ModeBlocking, manager.EffectiveMode())
service := &PromptService{config: manager, evaluator: NewGuardEvaluator(nil, nil, nil)}
decision, err := service.Evaluate(context.Background(), Request{Protocol: "openai_chat_completions", Body: []byte(`{"messages":[{"role":"user","content":"hi"}]}`)})
require.Error(t, err)
require.Nil(t, decision)
var guardErr *GuardError
require.ErrorAs(t, err, &guardErr)
require.Equal(t, ErrorCodeUnavailable, guardErr.Code)
}
type errorSettingRepository struct{ staticSettingRepository }
func (errorSettingRepository) GetMultiple(context.Context, []string) (map[string]string, error) {
return nil, errors.New("settings unavailable")
}
type switchableSettingRepository struct {
staticSettingRepository
loadErr error
}
func (r *switchableSettingRepository) GetMultiple(ctx context.Context, keys []string) (map[string]string, error) {
if r.loadErr != nil {
return nil, r.loadErr
}
return r.staticSettingRepository.GetMultiple(ctx, keys)
}
func TestConfigManagerStartupLoadFailureDoesNotBlockWhenBlockingNotIntended(t *testing.T) {
// Settings unavailable and no prior blocking intent: stay ModeOff so the
// gateway remains usable and admins can still disable/configure Prompt Audit.
manager := NewConfigManager(nil, errorSettingRepository{}, nil, prefixEncryptor{}, testTotpKeyConfig())
err := manager.Start(context.Background())
require.Error(t, err)
require.True(t, manager.configUntrusted.Load())
require.False(t, manager.BlockingActivationDegraded())
require.Equal(t, ModeOff, manager.EffectiveMode())
service := &PromptService{config: manager, evaluator: NewGuardEvaluator(nil, nil, nil)}
decision, evalErr := service.Evaluate(context.Background(), Request{
Protocol: "openai_chat_completions",
Body: []byte(`{"messages":[{"role":"user","content":"hi"}]}`),
})
require.NoError(t, evalErr)
require.NotNil(t, decision)
require.Equal(t, DecisionAllow, decision.Kind)
require.NoError(t, manager.Shutdown(context.Background()))
}
func TestConfigManagerStartupLoadFailureFailsClosedWhenBlockingIntended(t *testing.T) {
manager := NewConfigManager(nil, errorSettingRepository{}, nil, prefixEncryptor{}, testTotpKeyConfig())
// Simulate intent observed before a later load failure (e.g. decrypt error).
manager.observeExpectedState(`{"enabled":true,"blocking_enabled":true,"config_version":3}`, true)
manager.markConfigUntrusted()
require.True(t, manager.BlockingActivationDegraded())
require.Equal(t, ModeBlocking, manager.EffectiveMode())
service := &PromptService{config: manager, evaluator: NewGuardEvaluator(nil, nil, nil)}
decision, err := service.Evaluate(context.Background(), Request{
Protocol: "openai_chat_completions",
Body: []byte(`{"messages":[{"role":"user","content":"hi"}]}`),
})
require.Error(t, err)
require.Nil(t, decision)
var guardErr *GuardError
require.ErrorAs(t, err, &guardErr)
require.Equal(t, ErrorCodeUnavailable, guardErr.Code)
}
func TestConfigManagerUntrustedClearsOnSuccessfulDisable(t *testing.T) {
// After a degraded fail-closed period, saving disabled config must restore ModeOff.
manager := &ConfigManager{encryptor: prefixEncryptor{}, clock: fixedClock{}}
manager.observeExpectedState(`{"enabled":true,"blocking_enabled":true,"config_version":5}`, true)
manager.markConfigUntrusted()
require.Equal(t, ModeBlocking, manager.EffectiveMode())
// Install a trusted disabled snapshot the same way Save does after commit.
disabled := DefaultStorageConfig()
disabled.ConfigVersion = 6
disabled.Enabled = false
disabled.BlockingEnabled = false
active, err := ActiveFromStorage(disabled, true, manager.encryptor)
require.NoError(t, err)
manager.expected.Store(disabled.ConfigVersion)
manager.expectedBlocking.Store(false)
manager.snapshot.Store(&activeConfigSnapshot{storage: disabled, active: active, loadedAt: manager.clock.Now()})
manager.configUntrusted.Store(false)
require.False(t, manager.BlockingActivationDegraded())
require.Equal(t, ModeOff, manager.EffectiveMode())
service := &PromptService{config: manager, evaluator: NewGuardEvaluator(nil, nil, nil)}
decision, evalErr := service.Evaluate(context.Background(), Request{
Protocol: "openai_chat_completions",
Body: []byte(`{"messages":[{"role":"user","content":"hi"}]}`),
})
require.NoError(t, evalErr)
require.Equal(t, DecisionAllow, decision.Kind)
}
func TestConfigManagerUntrustedWithoutBlockingDoesNotForceBlockingMode(t *testing.T) {
manager := &ConfigManager{}
manager.observeExpectedState(`{"enabled":true,"blocking_enabled":false,"config_version":2}`, true)
manager.markConfigUntrusted()
require.False(t, manager.expectedBlocking.Load())
require.False(t, manager.BlockingActivationDegraded())
require.Equal(t, ModeOff, manager.EffectiveMode(), "async intent + untrusted must not force blocking unavailable")
}
func TestParseLegacyConfigDefaultsMissingFieldsWithoutEnablingBlocking(t *testing.T) {
storage, err := ParseStorageConfig(`{"enabled":false,"config_version":9}`)
require.NoError(t, err)
require.False(t, storage.BlockingEnabled)
require.Equal(t, "priority", storage.Strategy)
require.Equal(t, DefaultWorkerCount, storage.WorkerCount)
require.Equal(t, DefaultQueueCapacity, storage.QueueCapacity)
require.Equal(t, AllScannerIDs, storage.Scanners)
require.True(t, storage.AllGroups)
}
func TestUpdateConfigStrictBoundsAndKnownValues(t *testing.T) {
valid := promptAuditUpdateRequest(1, 1, "")
require.NoError(t, validateUpdateConfigRequest(valid))
tests := []struct {
name string
mutate func(*UpdateConfigRequest)
reason string
}{
{name: "strategy", mutate: func(req *UpdateConfigRequest) { req.Strategy = "round_robin" }, reason: "prompt_audit_invalid_strategy"},
{name: "worker low", mutate: func(req *UpdateConfigRequest) { req.WorkerCount = 0 }, reason: "prompt_audit_invalid_worker_count"},
{name: "worker high", mutate: func(req *UpdateConfigRequest) { req.WorkerCount = MaxWorkerCount + 1 }, reason: "prompt_audit_invalid_worker_count"},
{name: "capacity low", mutate: func(req *UpdateConfigRequest) { req.QueueCapacity = 0 }, reason: "prompt_audit_invalid_queue_capacity"},
{name: "capacity high", mutate: func(req *UpdateConfigRequest) { req.QueueCapacity = MaxQueueCapacity + 1 }, reason: "prompt_audit_invalid_queue_capacity"},
{name: "unknown scanner", mutate: func(req *UpdateConfigRequest) { req.Scanners = []string{"made_up"} }, reason: "prompt_audit_invalid_scanner"},
{name: "group required", mutate: func(req *UpdateConfigRequest) { req.AllGroups = false; req.GroupIDs = nil }, reason: "prompt_audit_groups_required"},
{name: "group positive", mutate: func(req *UpdateConfigRequest) { req.AllGroups = false; req.GroupIDs = []int64{0} }, reason: "prompt_audit_invalid_group"},
{name: "timeout low", mutate: func(req *UpdateConfigRequest) { req.Endpoints[0].TimeoutMS = MinTimeoutMS - 1 }, reason: "prompt_audit_invalid_timeout"},
{name: "timeout high", mutate: func(req *UpdateConfigRequest) { req.Endpoints[0].TimeoutMS = MaxTimeoutMS + 1 }, reason: "prompt_audit_invalid_timeout"},
{name: "input low", mutate: func(req *UpdateConfigRequest) { req.Endpoints[0].InputLimit = MinInputLimit - 1 }, reason: "prompt_audit_invalid_input_limit"},
{name: "input high", mutate: func(req *UpdateConfigRequest) { req.Endpoints[0].InputLimit = MaxInputLimit + 1 }, reason: "prompt_audit_invalid_input_limit"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
req := valid
req.Scanners = append([]string(nil), valid.Scanners...)
req.GroupIDs = append([]int64(nil), valid.GroupIDs...)
req.Endpoints = append([]UpdateEndpoint(nil), valid.Endpoints...)
tt.mutate(&req)
err := validateUpdateConfigRequest(req)
require.Error(t, err)
require.Equal(t, tt.reason, infraerrors.Reason(err))
})
}
}
@@ -0,0 +1,99 @@
package securityaudit
import (
"context"
"errors"
)
type Enqueuer struct {
config ConfigStore
repo JobRepository
payload PayloadStore
metrics Metrics
}
func NewEnqueuer(config ConfigStore, repo JobRepository, payload PayloadStore, metrics ...Metrics) *Enqueuer {
var metric Metrics
if len(metrics) > 0 {
metric = metrics[0]
}
return &Enqueuer{config: config, repo: repo, payload: payload, metrics: metric}
}
func (e *Enqueuer) Enqueue(ctx context.Context, req Request) error {
if e == nil || e.config == nil || e.repo == nil || e.payload == nil {
return errors.New("prompt audit enqueuer unavailable")
}
cfg, ok := e.config.Active()
baseFields := requestLogFields(req)
if !ok || cfg.EffectiveMode() != ModeAsync {
LogInfo(EventEnqueueSkipped, mergeLogFields(baseFields, map[string]any{"status": "skipped", "error_code": "mode_not_async"}))
return nil
}
baseFields["config_version"] = cfg.ConfigVersion
if !cfg.IncludesGroup(req.GroupID) {
LogInfo(EventEnqueueSkipped, mergeLogFields(baseFields, map[string]any{"status": "skipped", "error_code": "group_out_of_scope"}))
return nil
}
if len(cfg.EnabledEndpoints()) == 0 {
e.recordDropped()
LogWarn(EventEnqueueDropped, mergeLogFields(baseFields, map[string]any{"status": "dropped", "error_code": "no_enabled_endpoint"}))
return nil
}
snapshot, err := ExtractPromptSnapshot(req)
if errors.Is(err, ErrNoPromptText) {
LogInfo(EventEnqueueSkipped, mergeLogFields(baseFields, map[string]any{"status": "skipped", "error_code": "no_user_text"}))
return nil
}
if err != nil {
e.recordDropped()
LogWarn(EventEnqueueDropped, mergeLogFields(baseFields, map[string]any{"status": "dropped", "error_code": "snapshot_invalid"}))
return nil
}
job, err := e.repo.CreateStagingWithCapacity(ctx, snapshot.Redacted(), cfg.ConfigVersion, 3, cfg.QueueCapacity)
if err != nil {
code := "database_unavailable"
if errors.Is(err, ErrQueueFull) {
code = "queue_full"
}
if errors.Is(err, ErrQueueAdmissionBusy) {
code = "queue_admission_busy"
}
LogWarn(EventEnqueueDropped, mergeLogFields(baseFields, map[string]any{
"queue_capacity": cfg.QueueCapacity, "status": "dropped", "error_code": code,
}))
e.recordDropped()
return err
}
if err := e.payload.Set(ctx, job.ID, snapshot.ScanText, DefaultPayloadTTL); err != nil {
_ = e.repo.MarkStagingFailed(ctx, job.ID, "payload_store_failed", "payload store unavailable")
LogWarn(EventEnqueueDropped, mergeLogFields(baseFields, map[string]any{
"job_id": job.ID, "status": "dropped", "error_code": "payload_store_failed",
}))
e.recordDropped()
return err
}
if err := e.repo.PublishQueued(ctx, job.ID); err != nil {
_ = e.payload.Delete(ctx, job.ID)
_ = e.repo.MarkStagingFailed(ctx, job.ID, "queue_publish_failed", "queue publish failed")
LogWarn(EventEnqueueDropped, mergeLogFields(baseFields, map[string]any{
"job_id": job.ID, "status": "dropped", "error_code": "queue_publish_failed",
}))
e.recordDropped()
return err
}
LogInfo(EventJobEnqueued, mergeLogFields(baseFields, map[string]any{
"job_id": job.ID,
"queue_capacity": cfg.QueueCapacity, "status": "queued",
}))
if e.metrics != nil {
e.metrics.IncEnqueued()
}
return nil
}
func (e *Enqueuer) recordDropped() {
if e != nil && e.metrics != nil {
e.metrics.IncDropped()
}
}
@@ -0,0 +1,387 @@
package securityaudit
import (
"context"
"crypto/sha256"
"database/sql"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"strings"
"time"
"github.com/lib/pq"
)
type EventFilter struct {
Decision string `json:"decision,omitempty"`
RiskLevel string `json:"risk_level,omitempty"`
Endpoint string `json:"endpoint,omitempty"`
GroupID *int64 `json:"group_id,omitempty"`
UserID *int64 `json:"user_id,omitempty"`
APIKeyID *int64 `json:"api_key_id,omitempty"`
RequestID string `json:"request_id,omitempty"`
PromptHash string `json:"prompt_hash,omitempty"`
Keyword string `json:"keyword,omitempty"`
StartAt *time.Time `json:"start_at,omitempty"`
EndAt *time.Time `json:"end_at,omitempty"`
}
type EventPage struct {
Items []*Event `json:"items"`
Total int64 `json:"total"`
Page int `json:"page"`
PageSize int `json:"page_size"`
Pages int `json:"pages"`
}
type DeletePreview struct {
MatchedCount int64 `json:"matched_count"`
FilterSummary EventFilter `json:"filter_summary"`
SnapshotMaxID int64 `json:"snapshot_max_id"`
FilterHash string `json:"filter_hash"`
ConfirmationToken string `json:"confirmation_token,omitempty"`
ExpiresAt time.Time `json:"expires_at,omitempty"`
}
type DeleteResult struct {
DeletedEvents int64 `json:"deleted_events"`
DeletedJobs int64 `json:"deleted_jobs"`
JobIDs []int64 `json:"-"`
}
type EventRepository interface {
ListEvents(ctx context.Context, filter EventFilter, page, pageSize int) (*EventPage, error)
GetEvent(ctx context.Context, id int64) (*Event, error)
DeleteEvent(ctx context.Context, id int64) (*DeleteResult, error)
DeleteEventsByIDs(ctx context.Context, ids []int64) (*DeleteResult, error)
PreviewDelete(ctx context.Context, filter EventFilter) (*DeletePreview, error)
DeleteEventsByFilter(ctx context.Context, filter EventFilter, snapshotMaxID int64, batchSize int) (*DeleteResult, error)
}
func (r *PostgreSQLRepository) ListEvents(ctx context.Context, filter EventFilter, page, pageSize int) (*EventPage, error) {
if page < 1 {
page = 1
}
if pageSize < 1 {
pageSize = 20
}
if pageSize > 100 {
pageSize = 100
}
where, args := buildEventWhere(filter, 1)
var total int64
if err := r.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM prompt_audit_events e`+where, args...).Scan(&total); err != nil {
return nil, err
}
queryArgs := append([]any(nil), args...)
limitIndex := len(queryArgs) + 1
queryArgs = append(queryArgs, pageSize, (page-1)*pageSize)
rows, err := r.db.QueryContext(ctx, `SELECT `+eventColumns("e")+` FROM prompt_audit_events e`+where+
fmt.Sprintf(` ORDER BY e.created_at DESC, e.id DESC LIMIT $%d OFFSET $%d`, limitIndex, limitIndex+1), queryArgs...)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
items := make([]*Event, 0, pageSize)
for rows.Next() {
event, err := scanEvent(rows)
if err != nil {
return nil, err
}
items = append(items, event)
}
if err := rows.Err(); err != nil {
return nil, err
}
pages := 0
if total > 0 {
pages = int((total + int64(pageSize) - 1) / int64(pageSize))
}
return &EventPage{Items: items, Total: total, Page: page, PageSize: pageSize, Pages: pages}, nil
}
func (r *PostgreSQLRepository) GetEvent(ctx context.Context, id int64) (*Event, error) {
event, err := scanEvent(r.db.QueryRowContext(ctx, `SELECT `+eventDetailColumns("e")+` FROM prompt_audit_events e WHERE e.id=$1`, id), true)
if errors.Is(err, sql.ErrNoRows) {
return nil, ErrEventNotFound
}
return event, err
}
func (r *PostgreSQLRepository) DeleteEvent(ctx context.Context, id int64) (*DeleteResult, error) {
return r.DeleteEventsByIDs(ctx, []int64{id})
}
func (r *PostgreSQLRepository) DeleteEventsByIDs(ctx context.Context, ids []int64) (*DeleteResult, error) {
ids = canonicalInt64s(ids)
if len(ids) == 0 {
return &DeleteResult{}, nil
}
if len(ids) > 500 {
return nil, errors.New("prompt audit delete batch exceeds 500 events")
}
tx, err := r.db.BeginTx(ctx, nil)
if err != nil {
return nil, err
}
defer func() { _ = tx.Rollback() }()
rows, err := tx.QueryContext(ctx, `DELETE FROM prompt_audit_events WHERE id=ANY($1) RETURNING job_id`, pq.Array(ids))
if err != nil {
return nil, err
}
jobIDs, err := scanReturnedJobIDs(rows)
if err != nil {
return nil, err
}
deletedJobs, err := deleteOrphanJobs(ctx, tx, jobIDs)
if err != nil {
return nil, err
}
if err := tx.Commit(); err != nil {
return nil, err
}
return &DeleteResult{DeletedEvents: int64(len(jobIDs)), DeletedJobs: deletedJobs, JobIDs: canonicalInt64s(jobIDs)}, nil
}
func (r *PostgreSQLRepository) PreviewDelete(ctx context.Context, filter EventFilter) (*DeletePreview, error) {
if err := validateDeleteFilter(filter); err != nil {
return nil, err
}
tx, err := r.db.BeginTx(ctx, &sql.TxOptions{Isolation: sql.LevelRepeatableRead, ReadOnly: true})
if err != nil {
return nil, err
}
defer func() { _ = tx.Rollback() }()
where, args := buildEventWhere(filter, 1)
var count, maxID int64
if err := tx.QueryRowContext(ctx, `SELECT COUNT(*), COALESCE(MAX(e.id),0) FROM prompt_audit_events e`+where, args...).Scan(&count, &maxID); err != nil {
return nil, err
}
if err := tx.Commit(); err != nil {
return nil, err
}
canonical := canonicalEventFilter(filter)
return &DeletePreview{MatchedCount: count, FilterSummary: canonical, SnapshotMaxID: maxID, FilterHash: FilterHash(canonical, maxID)}, nil
}
func (r *PostgreSQLRepository) DeleteEventsByFilter(ctx context.Context, filter EventFilter, snapshotMaxID int64, batchSize int) (*DeleteResult, error) {
if err := validateDeleteFilter(filter); err != nil {
return nil, err
}
if snapshotMaxID <= 0 {
return &DeleteResult{}, nil
}
if batchSize < 1 || batchSize > 1000 {
batchSize = 200
}
total := &DeleteResult{}
jobSet := map[int64]struct{}{}
for {
tx, err := r.db.BeginTx(ctx, nil)
if err != nil {
return nil, err
}
where, args := buildEventWhere(filter, 1)
maxIndex := len(args) + 1
limitIndex := maxIndex + 1
args = append(args, snapshotMaxID, batchSize)
rows, err := tx.QueryContext(ctx, `
WITH selected AS (
SELECT e.id FROM prompt_audit_events e`+where+
fmt.Sprintf(` AND e.id <= $%d ORDER BY e.id LIMIT $%d FOR UPDATE SKIP LOCKED`, maxIndex, limitIndex)+`
), deleted AS (
DELETE FROM prompt_audit_events e USING selected s WHERE e.id=s.id RETURNING e.job_id
) SELECT job_id FROM deleted`, args...)
if err != nil {
_ = tx.Rollback()
return nil, err
}
jobIDs, err := scanReturnedJobIDs(rows)
if err != nil {
_ = tx.Rollback()
return nil, err
}
deletedJobs, err := deleteOrphanJobs(ctx, tx, jobIDs)
if err != nil {
_ = tx.Rollback()
return nil, err
}
if err := tx.Commit(); err != nil {
return nil, err
}
total.DeletedEvents += int64(len(jobIDs))
total.DeletedJobs += deletedJobs
for _, id := range jobIDs {
jobSet[id] = struct{}{}
}
if len(jobIDs) < batchSize {
break
}
}
for id := range jobSet {
total.JobIDs = append(total.JobIDs, id)
}
total.JobIDs = canonicalInt64s(total.JobIDs)
return total, nil
}
func FilterHash(filter EventFilter, snapshotMaxID int64) string {
payload := struct {
Filter EventFilter `json:"filter"`
SnapshotMaxID int64 `json:"snapshot_max_id"`
}{canonicalEventFilter(filter), snapshotMaxID}
raw, _ := json.Marshal(payload)
digest := sha256.Sum256(raw)
return hex.EncodeToString(digest[:])
}
func validateDeleteFilter(filter EventFilter) error {
if filter.StartAt == nil || filter.EndAt == nil || !filter.StartAt.Before(*filter.EndAt) {
return errors.New("prompt audit filter delete requires a valid explicit time range")
}
return nil
}
func canonicalEventFilter(filter EventFilter) EventFilter {
filter.Decision = strings.TrimSpace(strings.ToLower(filter.Decision))
filter.RiskLevel = strings.TrimSpace(strings.ToLower(filter.RiskLevel))
filter.Endpoint = strings.TrimSpace(filter.Endpoint)
filter.RequestID = strings.TrimSpace(filter.RequestID)
filter.PromptHash = strings.ToLower(strings.TrimSpace(filter.PromptHash))
filter.Keyword = strings.TrimSpace(filter.Keyword)
if filter.StartAt != nil {
value := filter.StartAt.UTC()
filter.StartAt = &value
}
if filter.EndAt != nil {
value := filter.EndAt.UTC()
filter.EndAt = &value
}
return filter
}
func buildEventWhere(filter EventFilter, firstIndex int) (string, []any) {
filter = canonicalEventFilter(filter)
clauses := []string{" WHERE TRUE"}
args := make([]any, 0, 12)
add := func(clause string, value any) {
clauses = append(clauses, fmt.Sprintf(clause, firstIndex+len(args)))
args = append(args, value)
}
if filter.Decision != "" {
add(" AND e.decision=$%d", filter.Decision)
}
if filter.RiskLevel != "" {
add(" AND e.risk_level=$%d", filter.RiskLevel)
}
if filter.Endpoint != "" {
add(" AND e.endpoint=$%d", filter.Endpoint)
}
if filter.GroupID != nil {
add(" AND e.group_id=$%d", *filter.GroupID)
}
if filter.UserID != nil {
add(" AND e.user_id=$%d", *filter.UserID)
}
if filter.APIKeyID != nil {
add(" AND e.api_key_id=$%d", *filter.APIKeyID)
}
if filter.RequestID != "" {
add(" AND e.request_id=$%d", filter.RequestID)
}
if filter.PromptHash != "" {
add(" AND e.prompt_hash=$%d", filter.PromptHash)
}
if filter.Keyword != "" {
add(` AND (e.request_id ILIKE $%d OR e.prompt_hash ILIKE $%d OR e.redacted_preview ILIKE $%d
OR e.username_snapshot ILIKE $%d OR e.user_email_snapshot ILIKE $%d OR e.api_key_name_snapshot ILIKE $%d)`, "%"+TrimRunes(filter.Keyword, 128)+"%")
// The clause has six placeholders but add only supplied one. Rebuild it with one shared placeholder.
clauses[len(clauses)-1] = fmt.Sprintf(` AND (e.request_id ILIKE $%[1]d OR e.prompt_hash ILIKE $%[1]d OR e.redacted_preview ILIKE $%[1]d
OR e.username_snapshot ILIKE $%[1]d OR e.user_email_snapshot ILIKE $%[1]d OR e.api_key_name_snapshot ILIKE $%[1]d)`, firstIndex+len(args)-1)
}
if filter.StartAt != nil {
add(" AND e.created_at >= $%d", filter.StartAt.UTC())
}
if filter.EndAt != nil {
add(" AND e.created_at <= $%d", filter.EndAt.UTC())
}
return strings.Join(clauses, ""), args
}
func eventColumns(alias string) string {
return fmt.Sprintf(`%[1]s.id,%[1]s.job_id,%[1]s.request_id,%[1]s.user_id,%[1]s.username_snapshot,
%[1]s.user_email_snapshot,%[1]s.api_key_id,%[1]s.api_key_name_snapshot,%[1]s.group_id,%[1]s.group_name,
%[1]s.provider,%[1]s.endpoint,%[1]s.protocol,%[1]s.model,%[1]s.prompt_hash,%[1]s.redacted_preview,
%[1]s.stage,%[1]s.decision,%[1]s.risk_level,%[1]s.action,%[1]s.categories,%[1]s.matched_scanners,
%[1]s.scanner_scores,%[1]s.scanner_evidence,%[1]s.scanner_backend,%[1]s.scanner_version,
%[1]s.guard_endpoint_id,%[1]s.policy_id,%[1]s.policy_version,%[1]s.config_version,
%[1]s.chunk_total,%[1]s.latency_ms,%[1]s.created_at`, alias)
}
// eventDetailColumns adds the full prompt, which can be large, so it is only
// loaded for single-event detail reads and never for list pages.
func eventDetailColumns(alias string) string {
return eventColumns(alias) + fmt.Sprintf(",%[1]s.full_prompt", alias)
}
func scanEvent(row rowScanner, withFullPrompt ...bool) (*Event, error) {
event := &Event{}
var userID, apiKeyID, groupID sql.NullInt64
var categories, matched, scores, evidence []byte
dest := []any{&event.ID, &event.JobID, &event.Snapshot.RequestID, &userID,
&event.Snapshot.UsernameSnapshot, &event.Snapshot.UserEmailSnapshot, &apiKeyID,
&event.Snapshot.APIKeyNameSnapshot, &groupID, &event.Snapshot.GroupName,
&event.Snapshot.Provider, &event.Snapshot.Endpoint, &event.Snapshot.Protocol, &event.Snapshot.Model,
&event.Snapshot.PromptHash, &event.Snapshot.RedactedPreview, &event.Snapshot.Stage, &event.Decision,
&event.RiskLevel, &event.Action, &categories, &matched, &scores, &evidence, &event.ScannerBackend,
&event.ScannerVersion, &event.GuardEndpointID, &event.PolicyID, &event.PolicyVersion,
&event.ConfigVersion, &event.ChunkTotal, &event.LatencyMS, &event.CreatedAt}
if len(withFullPrompt) > 0 && withFullPrompt[0] {
dest = append(dest, &event.Snapshot.FullPrompt)
}
err := row.Scan(dest...)
if err != nil {
return nil, err
}
event.Snapshot.UserID = nullableInt64Value(userID)
event.Snapshot.APIKeyID = nullableInt64Value(apiKeyID)
event.Snapshot.GroupID = nullableInt64Ptr(groupID)
_ = json.Unmarshal(categories, &event.Categories)
_ = json.Unmarshal(matched, &event.MatchedScanners)
_ = json.Unmarshal(scores, &event.ScannerScores)
_ = json.Unmarshal(evidence, &event.ScannerEvidence)
result := NormalizedResult{Decision: event.Decision, RiskLevel: event.RiskLevel, Action: event.Action,
Categories: event.Categories, MatchedScanners: event.MatchedScanners, ScannerScores: event.ScannerScores,
ScannerEvidence: event.ScannerEvidence}
event.IssueSummaries = BuildIssueSummaries(result)
return event, nil
}
func scanReturnedJobIDs(rows *sql.Rows) ([]int64, error) {
defer func() { _ = rows.Close() }()
result := make([]int64, 0)
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
return nil, err
}
result = append(result, id)
}
return result, rows.Err()
}
func deleteOrphanJobs(ctx context.Context, tx *sql.Tx, jobIDs []int64) (int64, error) {
jobIDs = canonicalInt64s(jobIDs)
if len(jobIDs) == 0 {
return 0, nil
}
result, err := tx.ExecContext(ctx, `DELETE FROM prompt_audit_jobs j
WHERE j.id=ANY($1) AND j.status <> 'processing'
AND NOT EXISTS (SELECT 1 FROM prompt_audit_events e WHERE e.job_id=j.id)`, pq.Array(jobIDs))
if err != nil {
return 0, err
}
return result.RowsAffected()
}
@@ -0,0 +1,279 @@
package securityaudit
import (
"context"
"errors"
"sync"
"time"
)
type GuardEvaluator struct {
scanner PromptScanner
repo JobRepository
metrics Metrics
clock Clock
global chan struct{}
perNodeLimit int
nodeMu sync.Mutex
nodes map[string]chan struct{}
}
func NewGuardEvaluator(scanner PromptScanner, repo JobRepository, metrics Metrics) *GuardEvaluator {
return newGuardEvaluator(scanner, repo, metrics, 64, 16)
}
func newGuardEvaluator(scanner PromptScanner, repo JobRepository, metrics Metrics, globalLimit, perNodeLimit int) *GuardEvaluator {
if globalLimit < 1 {
globalLimit = 64
}
if perNodeLimit < 1 {
perNodeLimit = 16
}
return &GuardEvaluator{scanner: scanner, repo: repo, metrics: metrics, clock: realClock{},
global: make(chan struct{}, globalLimit), perNodeLimit: perNodeLimit, nodes: map[string]chan struct{}{}}
}
func (g *GuardEvaluator) Evaluate(ctx context.Context, cfg ActiveConfig, snapshot PromptSnapshot) (*PromptDecision, error) {
if g == nil || g.scanner == nil {
if g != nil && g.metrics != nil {
g.metrics.Observe(DecisionUnavailable, 0)
}
logGuardFailure(snapshot, cfg, DecisionUnavailable, ErrorCodeUnavailable, "", 0)
return nil, &GuardError{Code: ErrorCodeUnavailable}
}
start := g.clock.Now()
baseFields := snapshotLogFields(snapshot)
baseFields["config_version"] = cfg.ConfigVersion
endpoints := cfg.EnabledEndpoints()
if len(endpoints) == 0 {
if g.metrics != nil {
g.metrics.Observe(DecisionUnavailable, g.clock.Now().Sub(start))
}
logGuardFailure(snapshot, cfg, DecisionUnavailable, ErrorCodeUnavailable, "", g.clock.Now().Sub(start))
return nil, &GuardError{Code: ErrorCodeUnavailable}
}
select {
case g.global <- struct{}{}:
defer func() { <-g.global }()
default:
if g.metrics != nil {
g.metrics.IncBulkheadFull()
g.metrics.Observe(DecisionUnavailable, g.clock.Now().Sub(start))
}
logGuardFailure(snapshot, cfg, DecisionUnavailable, ErrorCodeUnavailable, "", g.clock.Now().Sub(start))
return nil, &GuardError{Code: ErrorCodeUnavailable}
}
timeout := time.Duration(endpoints[0].TimeoutMS) * time.Millisecond
if timeout <= 0 {
timeout = DefaultTimeoutMS * time.Millisecond
}
evalCtx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
inputLimit := minimumInputLimit(endpoints)
chunks := SplitRunes(snapshot.ScanText, inputLimit)
if len(chunks) == 0 {
if g.metrics != nil {
g.metrics.Observe(DecisionAllow, g.clock.Now().Sub(start))
}
return &PromptDecision{Kind: DecisionAllow, AllowNextStage: true}, nil
}
LogInfo(EventEvaluationStarted, mergeLogFields(baseFields, map[string]any{"chunk_total": len(chunks), "status": "started"}))
results := make([]*NormalizedResult, 0, len(chunks))
for index, chunk := range chunks {
chunkStarted := g.clock.Now()
LogInfo(EventChunkStarted, mergeLogFields(baseFields, map[string]any{
"chunk_index": index + 1, "chunk_total": len(chunks),
"chunk_chars": len([]rune(chunk)), "input_chars": snapshot.PromptLength, "input_limit": inputLimit,
"status": "started",
}))
result, err := g.scanChunk(evalCtx, cfg, endpoints, chunk)
if err != nil {
code := guardErrorCode(err)
LogWarn(EventChunkFailed, mergeLogFields(baseFields, map[string]any{
"chunk_index": index + 1, "chunk_total": len(chunks),
"chunk_chars": len([]rune(chunk)), "input_chars": snapshot.PromptLength, "input_limit": inputLimit,
"latency_ms": g.clock.Now().Sub(chunkStarted).Milliseconds(), "error_code": code, "status": "failed",
}))
kind := DecisionUnavailable
if code == ErrorCodeInvalidResponse {
kind = DecisionInvalid
}
if g.metrics != nil {
g.metrics.Observe(kind, g.clock.Now().Sub(start))
var guardErr *GuardError
if errors.As(err, &guardErr) && guardErr.Timeout {
g.metrics.IncTimeout()
}
}
logGuardFailure(snapshot, cfg, kind, code, "", g.clock.Now().Sub(start))
return nil, err
}
result.ChunkTotal = len(chunks)
results = append(results, result)
LogInfo(EventChunkCompleted, mergeLogFields(baseFields, map[string]any{
"chunk_index": index + 1, "chunk_total": len(chunks),
"chunk_chars": len([]rune(chunk)), "input_chars": snapshot.PromptLength, "input_limit": inputLimit,
"guard_endpoint_id": result.GuardEndpointID, "action": result.Action,
"latency_ms": g.clock.Now().Sub(chunkStarted).Milliseconds(), "status": "completed",
}))
if result.Action == ActionBlock {
break
}
}
aggregated, err := AggregateResults(results, g.clock.Now().Sub(start))
if err != nil {
if g.metrics != nil {
g.metrics.Observe(DecisionInvalid, g.clock.Now().Sub(start))
}
logGuardFailure(snapshot, cfg, DecisionInvalid, ErrorCodeInvalidResponse, "", g.clock.Now().Sub(start))
return nil, &GuardError{Code: ErrorCodeInvalidResponse, Cause: err}
}
aggregated.ChunkTotal = len(chunks)
kind := DecisionAllow
if aggregated.Action == ActionWarn {
kind = DecisionFlag
}
if aggregated.Action == ActionBlock {
kind = DecisionBlock
}
decision := &PromptDecision{Kind: kind, Result: aggregated, AllowNextStage: kind == DecisionAllow || kind == DecisionFlag}
if kind == DecisionBlock {
decision.ErrorCode = ErrorCodeBlocked
}
if g.metrics != nil {
g.metrics.Observe(kind, g.clock.Now().Sub(start))
}
LogInfo(EventChunksAggregated, mergeLogFields(baseFields, map[string]any{
"decision": kind,
"risk_level": aggregated.RiskLevel, "action": aggregated.Action, "chunk_total": aggregated.ChunkTotal,
"latency_ms": aggregated.LatencyMS, "guard_endpoint_id": aggregated.GuardEndpointID, "stage": snapshot.Stage,
"status": "completed",
}))
if g.repo != nil {
if _, recordErr := g.repo.RecordBlocking(ctx, snapshot.Redacted(), cfg.ConfigVersion, aggregated, cfg.StorePassEvents); recordErr != nil {
if g.metrics != nil {
g.metrics.IncRecordFailed()
}
LogWarn(EventResultRecordFailed, mergeLogFields(baseFields, map[string]any{
"decision": kind, "error_code": "result_record_failed", "stage": snapshot.Stage,
"status": "failed",
}))
}
}
if kind == DecisionBlock {
LogWarn(EventGuardBlocked, mergeLogFields(baseFields, map[string]any{
"guard_endpoint_id": aggregated.GuardEndpointID,
"decision": kind, "risk_level": aggregated.RiskLevel, "action": aggregated.Action, "chunk_total": aggregated.ChunkTotal,
"latency_ms": aggregated.LatencyMS, "status": "blocked", "error_code": ErrorCodeBlocked,
"stage": snapshot.Stage, "upstream_dispatched": false, "billing_preconsumed": false,
}))
} else {
LogInfo(EventGuardAllowed, mergeLogFields(baseFields, map[string]any{
"decision": kind, "risk_level": aggregated.RiskLevel, "action": aggregated.Action,
"guard_endpoint_id": aggregated.GuardEndpointID, "chunk_total": aggregated.ChunkTotal,
"latency_ms": aggregated.LatencyMS, "stage": snapshot.Stage, "status": "allowed",
}))
}
return decision, nil
}
func logGuardFailure(snapshot PromptSnapshot, cfg ActiveConfig, kind DecisionKind, code, guardEndpointID string, latency time.Duration) {
fields := snapshotLogFields(snapshot)
fields["config_version"] = cfg.ConfigVersion
LogWarn(EventGuardFailed, mergeLogFields(fields, map[string]any{
"decision": kind, "guard_endpoint_id": guardEndpointID, "latency_ms": latency.Milliseconds(),
"status": "failed", "error_code": code, "upstream_dispatched": false, "billing_preconsumed": false,
}))
}
func (g *GuardEvaluator) scanChunk(ctx context.Context, cfg ActiveConfig, endpoints []ActiveEndpoint, chunk string) (*NormalizedResult, error) {
var lastErr error
for index, endpoint := range endpoints {
semaphore := g.nodeSemaphore(endpoint.ID)
select {
case semaphore <- struct{}{}:
case <-ctx.Done():
return nil, &GuardError{Code: ErrorCodeUnavailable, Retryable: true, Timeout: errors.Is(ctx.Err(), context.DeadlineExceeded), Cause: ctx.Err()}
default:
if g.metrics != nil {
g.metrics.IncBulkheadFull()
}
lastErr = &GuardError{Code: ErrorCodeUnavailable, Retryable: true}
if index < len(endpoints)-1 && g.metrics != nil {
g.metrics.IncFailover()
}
continue
}
result, err := callPromptScanner(ctx, g.scanner, endpoint, chunk, cfg.Scanners)
<-semaphore
if err == nil && result != nil {
return result, nil
}
if err == nil {
err = &GuardError{Code: ErrorCodeInvalidResponse, Retryable: false}
}
lastErr = err
var guardErr *GuardError
if !errors.As(err, &guardErr) || !guardErr.Retryable {
return nil, err
}
if index < len(endpoints)-1 && g.metrics != nil {
g.metrics.IncFailover()
}
}
if lastErr == nil {
lastErr = &GuardError{Code: ErrorCodeUnavailable}
}
return nil, lastErr
}
func callPromptScanner(ctx context.Context, scanner PromptScanner, endpoint ActiveEndpoint, chunk string, scanners []string) (result *NormalizedResult, err error) {
defer func() {
if recover() != nil {
result = nil
err = &GuardError{Code: ErrorCodeUnavailable, Retryable: false}
}
}()
return scanner.Scan(ctx, endpoint, chunk, scanners)
}
func (g *GuardEvaluator) nodeSemaphore(id string) chan struct{} {
g.nodeMu.Lock()
defer g.nodeMu.Unlock()
semaphore := g.nodes[id]
if semaphore == nil {
semaphore = make(chan struct{}, g.perNodeLimit)
g.nodes[id] = semaphore
}
return semaphore
}
func minimumInputLimit(endpoints []ActiveEndpoint) int {
limit := DefaultInputLimit
for index, endpoint := range endpoints {
value := endpoint.InputLimit
if value <= 0 {
value = DefaultInputLimit
}
if index == 0 || value < limit {
limit = value
}
}
return limit
}
func guardErrorCode(err error) string {
var guardErr *GuardError
if errors.As(err, &guardErr) && guardErr.Code != "" {
return guardErr.Code
}
return ErrorCodeUnavailable
}
func pointerLogID(value *int64) int64 {
if value == nil {
return 0
}
return *value
}
@@ -0,0 +1,300 @@
package securityaudit
import (
"context"
"errors"
"strings"
"sync"
"testing"
"time"
"github.com/stretchr/testify/require"
)
type scriptedScanner struct {
mu sync.Mutex
calls []string
block <-chan struct{}
entered chan<- struct{}
}
func (s *scriptedScanner) Scan(ctx context.Context, endpoint ActiveEndpoint, _ string, _ []string) (*NormalizedResult, error) {
s.mu.Lock()
s.calls = append(s.calls, endpoint.ID)
s.mu.Unlock()
if s.entered != nil {
select {
case s.entered <- struct{}{}:
default:
}
}
if s.block != nil {
select {
case <-s.block:
case <-ctx.Done():
return nil, &GuardError{Code: ErrorCodeUnavailable, Retryable: true, Timeout: true, Cause: ctx.Err()}
}
}
if endpoint.ID == "bad" {
return nil, &GuardError{Code: ErrorCodeUnavailable, Retryable: true}
}
if endpoint.ID == "invalid" {
return nil, &GuardError{Code: ErrorCodeInvalidResponse}
}
return &NormalizedResult{Decision: EventPass, RiskLevel: RiskLow, Action: ActionAllow, Safety: "Safe", ScannerScores: map[string]float64{}, ScannerEvidence: map[string]string{}, GuardEndpointID: endpoint.ID}, nil
}
func guardConfig(endpoints ...ActiveEndpoint) ActiveConfig {
return ActiveConfig{RiskControlEnabled: true, Enabled: true, BlockingEnabled: true, ConfigVersion: 2, Scanners: AllScannerIDs, Endpoints: endpoints}
}
func TestGuardEvaluatorOrderedFailoverAndInvalidTerminal(t *testing.T) {
scanner := &scriptedScanner{}
metrics := NewAtomicMetrics()
evaluator := newGuardEvaluator(scanner, nil, metrics, 4, 2)
snapshot := PromptSnapshot{RequestID: "r", ScanText: "hello", PromptLength: 5}
decision, err := evaluator.Evaluate(context.Background(), guardConfig(
ActiveEndpoint{ID: "bad", Enabled: true, TimeoutMS: 1000, InputLimit: 100},
ActiveEndpoint{ID: "good", Enabled: true, TimeoutMS: 1000, InputLimit: 100},
), snapshot)
require.NoError(t, err)
require.Equal(t, DecisionAllow, decision.Kind)
require.Equal(t, int64(1), metrics.Snapshot().Failovers)
_, err = evaluator.Evaluate(context.Background(), guardConfig(
ActiveEndpoint{ID: "invalid", Enabled: true, TimeoutMS: 1000, InputLimit: 100},
ActiveEndpoint{ID: "good", Enabled: true, TimeoutMS: 1000, InputLimit: 100},
), snapshot)
var guardErr *GuardError
require.ErrorAs(t, err, &guardErr)
require.Equal(t, ErrorCodeInvalidResponse, guardErr.Code)
snapshotMetrics := metrics.Snapshot()
require.Equal(t, int64(2), snapshotMetrics.Total)
require.Equal(t, int64(1), snapshotMetrics.Allowed)
require.Equal(t, int64(1), snapshotMetrics.Invalid)
}
func TestGuardEvaluatorGlobalBulkheadIsNonBlocking(t *testing.T) {
release := make(chan struct{})
entered := make(chan struct{}, 1)
scanner := &scriptedScanner{block: release, entered: entered}
metrics := NewAtomicMetrics()
evaluator := newGuardEvaluator(scanner, nil, metrics, 1, 1)
cfg := guardConfig(ActiveEndpoint{ID: "good", Enabled: true, TimeoutMS: 2000, InputLimit: 100})
done := make(chan error, 1)
go func() {
_, err := evaluator.Evaluate(context.Background(), cfg, PromptSnapshot{ScanText: "one", PromptLength: 3})
done <- err
}()
select {
case <-entered:
case <-time.After(time.Second):
t.Fatal("first evaluation did not enter scanner")
}
start := time.Now()
_, err := evaluator.Evaluate(context.Background(), cfg, PromptSnapshot{ScanText: "two", PromptLength: 3})
require.Error(t, err)
require.Less(t, time.Since(start), 200*time.Millisecond)
require.Equal(t, int64(1), metrics.Snapshot().BulkheadFull)
close(release)
require.NoError(t, <-done)
snapshotMetrics := metrics.Snapshot()
require.Equal(t, int64(2), snapshotMetrics.Total)
require.Equal(t, int64(1), snapshotMetrics.Allowed)
require.Equal(t, int64(1), snapshotMetrics.Unavailable)
}
func TestGuardEvaluatorPerNodeBulkheadIsNonBlocking(t *testing.T) {
release := make(chan struct{})
entered := make(chan struct{}, 1)
scanner := &scriptedScanner{block: release, entered: entered}
metrics := NewAtomicMetrics()
evaluator := newGuardEvaluator(scanner, nil, metrics, 2, 1)
cfg := guardConfig(ActiveEndpoint{ID: "same-node", Enabled: true, TimeoutMS: 2000, InputLimit: 100})
done := make(chan error, 1)
go func() {
_, err := evaluator.Evaluate(context.Background(), cfg, PromptSnapshot{ScanText: "one", PromptLength: 3})
done <- err
}()
select {
case <-entered:
case <-time.After(time.Second):
t.Fatal("first evaluation did not enter scanner")
}
started := time.Now()
_, err := evaluator.Evaluate(context.Background(), cfg, PromptSnapshot{ScanText: "two", PromptLength: 3})
require.Error(t, err)
require.Less(t, time.Since(started), 200*time.Millisecond)
require.GreaterOrEqual(t, metrics.Snapshot().BulkheadFull, int64(1))
close(release)
require.NoError(t, <-done)
}
func TestGuardEvaluatorLastChunkFailureNeverAllows(t *testing.T) {
call := 0
scanner := PromptScannerFunc(func(context.Context, ActiveEndpoint, string, []string) (*NormalizedResult, error) {
call++
if call == 2 {
return nil, &GuardError{Code: ErrorCodeUnavailable, Retryable: true, Cause: errors.New("down")}
}
return &NormalizedResult{Decision: EventPass, RiskLevel: RiskLow, Action: ActionAllow, ScannerScores: map[string]float64{}, ScannerEvidence: map[string]string{}}, nil
})
metrics := NewAtomicMetrics()
evaluator := newGuardEvaluator(scanner, nil, metrics, 2, 2)
_, err := evaluator.Evaluate(context.Background(), guardConfig(ActiveEndpoint{ID: "one", Enabled: true, TimeoutMS: 1000, InputLimit: 3}), PromptSnapshot{ScanText: "abcdef", PromptLength: 6})
require.Error(t, err)
}
func TestGuardEvaluatorScansLatestUserPromptAsIndependentFirstChunk(t *testing.T) {
latest := "请帮我编写一篇黄色小说 名字你来取"
history := strings.Repeat("# AGENTS.md instructions 项目安全规则。", 30)
seen := make([]string, 0, 4)
scanner := PromptScannerFunc(func(_ context.Context, _ ActiveEndpoint, prompt string, _ []string) (*NormalizedResult, error) {
seen = append(seen, prompt)
return &NormalizedResult{Decision: EventPass, RiskLevel: RiskLow, Action: ActionAllow, ScannerScores: map[string]float64{}, ScannerEvidence: map[string]string{}}, nil
})
evaluator := newGuardEvaluator(scanner, nil, NewAtomicMetrics(), 2, 2)
_, err := evaluator.Evaluate(context.Background(), guardConfig(
ActiveEndpoint{ID: "one", Enabled: true, TimeoutMS: 1000, InputLimit: 128},
), PromptSnapshot{ScanText: latest + promptAuditPrioritySeparator + history, PromptLength: len([]rune(latest + history))})
require.NoError(t, err)
require.Greater(t, len(seen), 1)
require.Equal(t, latest, seen[0])
require.Equal(t, history, strings.Join(seen[1:], ""))
}
func TestGuardEvaluatorBlockStopsRemainingChunksButReportsPlannedTotal(t *testing.T) {
calls := 0
scanner := PromptScannerFunc(func(context.Context, ActiveEndpoint, string, []string) (*NormalizedResult, error) {
calls++
return &NormalizedResult{
Decision: EventCritical, RiskLevel: RiskCritical, Action: ActionBlock, Safety: "Unsafe",
Categories: []string{"jailbreak"}, MatchedScanners: []string{"jailbreak"},
ScannerScores: map[string]float64{"jailbreak": 1}, ScannerEvidence: map[string]string{"jailbreak": "Jailbreak"},
}, nil
})
metrics := NewAtomicMetrics()
evaluator := newGuardEvaluator(scanner, nil, metrics, 2, 2)
decision, err := evaluator.Evaluate(context.Background(), guardConfig(
ActiveEndpoint{ID: "one", Enabled: true, TimeoutMS: 1000, InputLimit: 3},
), PromptSnapshot{ScanText: "abcdefghi", PromptLength: 9})
require.NoError(t, err)
require.Equal(t, DecisionBlock, decision.Kind)
require.Equal(t, 1, calls)
require.Equal(t, 3, decision.Result.ChunkTotal)
require.Equal(t, int64(1), metrics.Snapshot().Blocked)
}
func TestGuardEvaluatorFlagSharedDeadlineFailClosedAndContextCancel(t *testing.T) {
t.Run("flag allows next stage", func(t *testing.T) {
metrics := NewAtomicMetrics()
evaluator := newGuardEvaluator(PromptScannerFunc(func(context.Context, ActiveEndpoint, string, []string) (*NormalizedResult, error) {
return &NormalizedResult{Decision: EventFlag, RiskLevel: RiskMedium, Action: ActionWarn, Safety: "Controversial", Categories: []string{"violent"}, MatchedScanners: []string{"violent"}, ScannerScores: map[string]float64{"violent": .5}, ScannerEvidence: map[string]string{"violent": "Violent"}}, nil
}), nil, metrics, 2, 2)
decision, err := evaluator.Evaluate(context.Background(), guardConfig(ActiveEndpoint{ID: "one", Enabled: true, TimeoutMS: 1000, InputLimit: 100}), PromptSnapshot{ScanText: "review", PromptLength: 6})
require.NoError(t, err)
require.Equal(t, DecisionFlag, decision.Kind)
require.True(t, decision.AllowNextStage)
require.Equal(t, int64(1), metrics.Snapshot().Flagged)
})
t.Run("all failovers share first endpoint deadline", func(t *testing.T) {
calls := 0
scanner := PromptScannerFunc(func(ctx context.Context, endpoint ActiveEndpoint, _ string, _ []string) (*NormalizedResult, error) {
calls++
if endpoint.ID == "first" {
select {
case <-time.After(35 * time.Millisecond):
return nil, &GuardError{Code: ErrorCodeUnavailable, Retryable: true}
case <-ctx.Done():
return nil, &GuardError{Code: ErrorCodeUnavailable, Retryable: true, Timeout: true, Cause: ctx.Err()}
}
}
<-ctx.Done()
return nil, &GuardError{Code: ErrorCodeUnavailable, Retryable: true, Timeout: true, Cause: ctx.Err()}
})
metrics := NewAtomicMetrics()
evaluator := newGuardEvaluator(scanner, nil, metrics, 2, 2)
started := time.Now()
_, err := evaluator.Evaluate(context.Background(), guardConfig(
ActiveEndpoint{ID: "first", Enabled: true, TimeoutMS: 70, InputLimit: 100},
ActiveEndpoint{ID: "second", Enabled: true, TimeoutMS: 500, InputLimit: 100},
), PromptSnapshot{ScanText: "deadline", PromptLength: 8})
elapsed := time.Since(started)
require.Error(t, err)
require.Equal(t, 2, calls)
// The bound only has to prove the failover shared the first endpoint's
// 70ms deadline instead of taking the second endpoint's own 500ms one.
// An unshared deadline lands at ~535ms, so 350ms still fails loudly
// while leaving room for scheduler delay on a busy CI machine. A
// tighter bound made this test flaky, not stricter.
require.Less(t, elapsed, 350*time.Millisecond)
require.GreaterOrEqual(t, elapsed, 50*time.Millisecond)
require.Equal(t, int64(1), metrics.Snapshot().Failovers)
require.Equal(t, int64(1), metrics.Snapshot().Timeouts)
})
t.Run("canceled parent never allows", func(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
cancel()
evaluator := newGuardEvaluator(PromptScannerFunc(func(ctx context.Context, _ ActiveEndpoint, _ string, _ []string) (*NormalizedResult, error) {
<-ctx.Done()
return nil, &GuardError{Code: ErrorCodeUnavailable, Retryable: true, Cause: ctx.Err()}
}), nil, NewAtomicMetrics(), 2, 2)
decision, err := evaluator.Evaluate(ctx, guardConfig(ActiveEndpoint{ID: "one", Enabled: true, TimeoutMS: 1000, InputLimit: 100}), PromptSnapshot{ScanText: "cancel", PromptLength: 6})
require.Error(t, err)
require.Nil(t, decision)
})
}
func TestGuardEvaluatorRecordsExistingResultOnceAndRecordFailureDoesNotChangeDecision(t *testing.T) {
for _, recordErr := range []error{nil, errors.New("database unavailable")} {
repo := &fakeJobRepository{recordBlockingErr: recordErr}
metrics := NewAtomicMetrics()
scannerCalls := 0
evaluator := newGuardEvaluator(PromptScannerFunc(func(context.Context, ActiveEndpoint, string, []string) (*NormalizedResult, error) {
scannerCalls++
return &NormalizedResult{Decision: EventCritical, RiskLevel: RiskCritical, Action: ActionBlock, Safety: "Unsafe", Categories: []string{"pii"}, MatchedScanners: []string{"pii"}, ScannerScores: map[string]float64{"pii": 1}, ScannerEvidence: map[string]string{"pii": "PII"}}, nil
}), repo, metrics, 2, 2)
decision, err := evaluator.Evaluate(context.Background(), guardConfig(ActiveEndpoint{ID: "one", Enabled: true, TimeoutMS: 1000, InputLimit: 100}), PromptSnapshot{ScanText: "raw prompt", RedactedPreview: "raw***", PromptLength: 10})
require.NoError(t, err)
require.Equal(t, DecisionBlock, decision.Kind)
require.Equal(t, 1, scannerCalls)
require.Equal(t, 1, repo.recordBlockingCalls)
require.Empty(t, repo.recordBlockingSnapshot.ScanText)
require.Same(t, decision.Result, repo.recordBlockingResult)
if recordErr != nil {
require.Equal(t, int64(1), metrics.Snapshot().RecordFailed)
} else {
require.Zero(t, metrics.Snapshot().RecordFailed)
}
}
}
func TestGuardEvaluatorNilResultAndScannerPanicBecomeStableFailures(t *testing.T) {
tests := []struct {
name string
scan PromptScannerFunc
code string
}{
{name: "nil result", scan: func(context.Context, ActiveEndpoint, string, []string) (*NormalizedResult, error) { return nil, nil }, code: ErrorCodeInvalidResponse},
{name: "panic", scan: func(context.Context, ActiveEndpoint, string, []string) (*NormalizedResult, error) {
panic("raw prompt canary")
}, code: ErrorCodeUnavailable},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
evaluator := newGuardEvaluator(tt.scan, nil, NewAtomicMetrics(), 2, 2)
_, err := evaluator.Evaluate(context.Background(), guardConfig(ActiveEndpoint{ID: "one", Enabled: true, TimeoutMS: 1000, InputLimit: 100}), PromptSnapshot{ScanText: "input", PromptLength: 5})
var guardErr *GuardError
require.ErrorAs(t, err, &guardErr)
require.Equal(t, tt.code, guardErr.Code)
require.NotContains(t, err.Error(), "canary")
})
}
}
type PromptScannerFunc func(context.Context, ActiveEndpoint, string, []string) (*NormalizedResult, error)
func (f PromptScannerFunc) Scan(ctx context.Context, endpoint ActiveEndpoint, chunk string, scanners []string) (*NormalizedResult, error) {
return f(ctx, endpoint, chunk, scanners)
}
@@ -0,0 +1,313 @@
package securityaudit
import (
"context"
"errors"
"strconv"
"strings"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/gin-gonic/gin"
)
type PromptAdminService interface {
GetConfig() (PublicConfig, error)
SaveConfig(context.Context, UpdateConfigRequest, int64) (PublicConfig, error)
Probe(context.Context, ProbeRequest) ProbeResult
Runtime(context.Context) RuntimeSnapshot
ListEvents(context.Context, EventFilter, int, int) (*EventPage, error)
GetEvent(context.Context, int64) (*Event, error)
DeleteEvent(context.Context, int64) (*DeleteResult, error)
DeleteEventsByIDs(context.Context, []int64) (*DeleteResult, error)
PreviewDelete(context.Context, EventFilter, int64) (*DeletePreview, error)
DeleteByFilter(context.Context, DeleteByFilterRequest, int64) (*DeleteResult, error)
}
type PromptAdminHandler struct{ service PromptAdminService }
func NewPromptAdminHandler(service PromptAdminService) *PromptAdminHandler {
return &PromptAdminHandler{service: service}
}
func (h *PromptAdminHandler) GetConfig(c *gin.Context) {
config, err := h.service.GetConfig()
if err != nil {
response.ErrorFrom(c, err)
return
}
response.Success(c, config)
}
func (h *PromptAdminHandler) UpdateConfig(c *gin.Context) {
var request UpdateConfigRequest
if err := c.ShouldBindJSON(&request); err != nil {
setPromptAdminAudit(c, "failed", "prompt_audit_invalid_config_request", nil)
response.ErrorFrom(c, infraerrors.BadRequest("prompt_audit_invalid_config_request", "提示词审计配置请求无效"))
return
}
config, err := h.service.SaveConfig(c.Request.Context(), request, adminID(c))
if err != nil {
setPromptAdminAudit(c, "failed", infraerrors.Reason(err), configAuditFields(request, nil))
response.ErrorFrom(c, err)
return
}
setPromptAdminAudit(c, "success", "", configAuditFields(request, &config))
response.Success(c, config)
}
func (h *PromptAdminHandler) ProbeEndpoint(c *gin.Context) {
var request ProbeRequest
if err := c.ShouldBindJSON(&request); err != nil {
setPromptAdminAudit(c, "failed", "prompt_audit_invalid_probe_request", nil)
response.ErrorFrom(c, infraerrors.BadRequest("prompt_audit_invalid_probe_request", "审计节点探测请求无效"))
return
}
result := h.service.Probe(c.Request.Context(), request)
status := "failed"
if result.OK {
status = "success"
}
setPromptAdminAudit(c, status, result.ErrorCode, map[string]any{
"guard_endpoint_id": request.Endpoint.ID, "http_status": result.HTTPStatus,
"latency_ms": result.LatencyMS, "token_applied": result.TokenApplied, "retryable": result.Retryable,
})
response.Success(c, result)
}
func (h *PromptAdminHandler) GetRuntime(c *gin.Context) {
response.Success(c, h.service.Runtime(c.Request.Context()))
}
func (h *PromptAdminHandler) ListEvents(c *gin.Context) {
page, err := positiveIntQuery(c, "page", 1, 0)
if err != nil {
response.ErrorFrom(c, err)
return
}
pageSize, err := positiveIntQuery(c, "page_size", 20, 100)
if err != nil {
response.ErrorFrom(c, err)
return
}
filter, err := eventFilterFromQuery(c)
if err != nil {
response.ErrorFrom(c, err)
return
}
result, err := h.service.ListEvents(c.Request.Context(), filter, page, pageSize)
if err != nil {
response.ErrorFrom(c, err)
return
}
response.Success(c, result)
}
func (h *PromptAdminHandler) GetEvent(c *gin.Context) {
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil || id <= 0 {
response.ErrorFrom(c, infraerrors.BadRequest("prompt_audit_invalid_event_id", "事件 ID 无效"))
return
}
event, err := h.service.GetEvent(c.Request.Context(), id)
if errors.Is(err, ErrEventNotFound) {
response.ErrorFrom(c, infraerrors.NotFound("prompt_audit_event_not_found", "提示词审计事件不存在"))
return
}
if err != nil {
response.ErrorFrom(c, err)
return
}
response.Success(c, event)
}
func (h *PromptAdminHandler) DeleteEvent(c *gin.Context) {
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil || id <= 0 {
setPromptAdminAudit(c, "failed", "prompt_audit_invalid_event_id", nil)
response.ErrorFrom(c, infraerrors.BadRequest("prompt_audit_invalid_event_id", "事件 ID 无效"))
return
}
result, err := h.service.DeleteEvent(c.Request.Context(), id)
if err != nil {
setPromptAdminAudit(c, "failed", infraerrors.Reason(err), map[string]any{"event_id": id})
response.ErrorFrom(c, err)
return
}
setPromptAdminAudit(c, "success", "", deleteAuditFields(result, map[string]any{"event_id": id}))
LogWarn(EventEventDeleted, map[string]any{"user_id": adminID(c), "event_id": id, "status": "deleted"})
response.Success(c, result)
}
type batchDeleteRequest struct {
IDs []int64 `json:"ids" binding:"required"`
}
func (h *PromptAdminHandler) BatchDelete(c *gin.Context) {
var request batchDeleteRequest
if err := c.ShouldBindJSON(&request); err != nil || len(request.IDs) == 0 || len(request.IDs) > 500 {
setPromptAdminAudit(c, "failed", "prompt_audit_invalid_delete_batch", nil)
response.ErrorFrom(c, infraerrors.BadRequest("prompt_audit_invalid_delete_batch", "批量删除必须包含 1-500 个事件 ID"))
return
}
for _, id := range request.IDs {
if id <= 0 {
setPromptAdminAudit(c, "failed", "prompt_audit_invalid_event_id", map[string]any{"requested_count": len(request.IDs)})
response.ErrorFrom(c, infraerrors.BadRequest("prompt_audit_invalid_event_id", "事件 ID 无效"))
return
}
}
result, err := h.service.DeleteEventsByIDs(c.Request.Context(), request.IDs)
if err != nil {
setPromptAdminAudit(c, "failed", infraerrors.Reason(err), map[string]any{"requested_count": len(request.IDs)})
response.ErrorFrom(c, err)
return
}
setPromptAdminAudit(c, "success", "", deleteAuditFields(result, map[string]any{"requested_count": len(request.IDs)}))
LogWarn(EventEventsDeleted, map[string]any{"user_id": adminID(c), "status": "deleted"})
response.Success(c, result)
}
func (h *PromptAdminHandler) DeletePreview(c *gin.Context) {
var filter EventFilter
if err := c.ShouldBindJSON(&filter); err != nil {
setPromptAdminAudit(c, "failed", "prompt_audit_delete_preview_invalid", nil)
response.ErrorFrom(c, infraerrors.BadRequest("prompt_audit_delete_preview_invalid", "删除预览筛选无效"))
return
}
preview, err := h.service.PreviewDelete(c.Request.Context(), filter, adminID(c))
if err != nil {
setPromptAdminAudit(c, "failed", "prompt_audit_delete_preview_invalid", nil)
response.ErrorFrom(c, infraerrors.BadRequest("prompt_audit_delete_preview_invalid", "删除预览筛选无效"))
return
}
setPromptAdminAudit(c, "success", "", map[string]any{
"matched_count": preview.MatchedCount, "snapshot_max_id": preview.SnapshotMaxID, "filter_hash": preview.FilterHash,
})
response.Success(c, preview)
}
func (h *PromptAdminHandler) DeleteByFilter(c *gin.Context) {
var request DeleteByFilterRequest
if err := c.ShouldBindJSON(&request); err != nil {
setPromptAdminAudit(c, "failed", "prompt_audit_delete_confirmation_invalid", nil)
response.ErrorFrom(c, infraerrors.BadRequest("prompt_audit_delete_confirmation_invalid", "删除确认无效或已过期"))
return
}
result, err := h.service.DeleteByFilter(c.Request.Context(), request, adminID(c))
if err != nil {
setPromptAdminAudit(c, "failed", "prompt_audit_delete_confirmation_invalid", map[string]any{
"snapshot_max_id": request.SnapshotMaxID, "filter_hash": request.FilterHash, "confirm": request.Confirm,
})
response.ErrorFrom(c, infraerrors.BadRequest("prompt_audit_delete_confirmation_invalid", "删除确认无效或已过期"))
return
}
setPromptAdminAudit(c, "success", "", deleteAuditFields(result, map[string]any{
"snapshot_max_id": request.SnapshotMaxID, "filter_hash": request.FilterHash, "confirm": request.Confirm,
}))
response.Success(c, result)
}
func setPromptAdminAudit(c *gin.Context, result, errorCode string, fields map[string]any) {
details := make(map[string]any, len(fields)+2)
details["result"] = result
if strings.TrimSpace(errorCode) != "" {
details["error_code"] = errorCode
}
for key, value := range fields {
details[key] = value
}
middleware.SetAuditExtra(c, details)
}
func configAuditFields(request UpdateConfigRequest, saved *PublicConfig) map[string]any {
version := request.ExpectedConfigVersion
if saved != nil {
version = saved.ConfigVersion
}
return map[string]any{
"enabled": request.Enabled, "blocking_enabled": request.BlockingEnabled,
"blocking_latest_turn_only": request.BlockingLatestTurnOnly,
"config_version": version, "endpoint_count": len(request.Endpoints),
"scanner_count": len(request.Scanners), "all_groups": request.AllGroups,
"group_count": len(request.GroupIDs),
}
}
func deleteAuditFields(result *DeleteResult, base map[string]any) map[string]any {
fields := make(map[string]any, len(base)+2)
for key, value := range base {
fields[key] = value
}
if result != nil {
fields["deleted_events"] = result.DeletedEvents
fields["deleted_jobs"] = result.DeletedJobs
}
return fields
}
func adminID(c *gin.Context) int64 {
subject, ok := middleware.GetAuthSubjectFromContext(c)
if !ok {
return 0
}
return subject.UserID
}
func eventFilterFromQuery(c *gin.Context) (EventFilter, error) {
groupID, err := optionalPositiveInt64Query(c, "group_id")
if err != nil {
return EventFilter{}, err
}
userID, err := optionalPositiveInt64Query(c, "user_id")
if err != nil {
return EventFilter{}, err
}
apiKeyID, err := optionalPositiveInt64Query(c, "api_key_id")
if err != nil {
return EventFilter{}, err
}
filter := EventFilter{
Decision: c.Query("decision"), RiskLevel: c.Query("risk_level"), Endpoint: c.Query("endpoint"),
GroupID: groupID, UserID: userID, APIKeyID: apiKeyID, RequestID: c.Query("request_id"),
PromptHash: c.Query("prompt_hash"), Keyword: c.Query("keyword"),
}
if value := strings.TrimSpace(c.Query("start_at")); value != "" {
filter.StartAt = parseTimeQuery(value)
if filter.StartAt == nil {
return EventFilter{}, infraerrors.BadRequest("prompt_audit_invalid_time", "开始时间无效")
}
}
if value := strings.TrimSpace(c.Query("end_at")); value != "" {
filter.EndAt = parseTimeQuery(value)
if filter.EndAt == nil {
return EventFilter{}, infraerrors.BadRequest("prompt_audit_invalid_time", "结束时间无效")
}
}
return filter, nil
}
func optionalPositiveInt64Query(c *gin.Context, key string) (*int64, error) {
value := strings.TrimSpace(c.Query(key))
if value == "" {
return nil, nil
}
parsed, err := strconv.ParseInt(value, 10, 64)
if err != nil || parsed <= 0 {
return nil, infraerrors.BadRequest("prompt_audit_invalid_filter_id", "事件筛选 ID 无效")
}
return &parsed, nil
}
func positiveIntQuery(c *gin.Context, key string, defaultValue, maxValue int) (int, error) {
value := strings.TrimSpace(c.Query(key))
if value == "" {
return defaultValue, nil
}
parsed, err := strconv.Atoi(value)
if err != nil || parsed <= 0 || (maxValue > 0 && parsed > maxValue) {
return 0, infraerrors.BadRequest("prompt_audit_invalid_pagination", "分页参数无效")
}
return parsed, nil
}
@@ -0,0 +1,249 @@
package securityaudit
import (
"bytes"
"context"
"encoding/json"
"errors"
"net/http"
"net/http/httptest"
"strings"
"testing"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type fakePromptAdminService struct {
config PublicConfig
save func(context.Context, UpdateConfigRequest, int64) (PublicConfig, error)
probe func(context.Context, ProbeRequest) ProbeResult
runtime RuntimeSnapshot
list func(context.Context, EventFilter, int, int) (*EventPage, error)
get func(context.Context, int64) (*Event, error)
deleteOne func(context.Context, int64) (*DeleteResult, error)
deleteIDs func(context.Context, []int64) (*DeleteResult, error)
preview func(context.Context, EventFilter, int64) (*DeletePreview, error)
deleteFilter func(context.Context, DeleteByFilterRequest, int64) (*DeleteResult, error)
}
func (s *fakePromptAdminService) GetConfig() (PublicConfig, error) {
return s.config, nil
}
func (s *fakePromptAdminService) SaveConfig(ctx context.Context, req UpdateConfigRequest, actorID int64) (PublicConfig, error) {
if s.save == nil {
return PublicConfig{}, errors.New("unexpected SaveConfig call")
}
return s.save(ctx, req, actorID)
}
func (s *fakePromptAdminService) Probe(ctx context.Context, req ProbeRequest) ProbeResult {
if s.probe == nil {
return ProbeResult{}
}
return s.probe(ctx, req)
}
func (s *fakePromptAdminService) Runtime(context.Context) RuntimeSnapshot { return s.runtime }
func (s *fakePromptAdminService) ListEvents(ctx context.Context, filter EventFilter, page, pageSize int) (*EventPage, error) {
if s.list == nil {
return &EventPage{}, nil
}
return s.list(ctx, filter, page, pageSize)
}
func (s *fakePromptAdminService) GetEvent(ctx context.Context, id int64) (*Event, error) {
if s.get == nil {
return nil, ErrEventNotFound
}
return s.get(ctx, id)
}
func (s *fakePromptAdminService) DeleteEvent(ctx context.Context, id int64) (*DeleteResult, error) {
if s.deleteOne == nil {
return &DeleteResult{}, nil
}
return s.deleteOne(ctx, id)
}
func (s *fakePromptAdminService) DeleteEventsByIDs(ctx context.Context, ids []int64) (*DeleteResult, error) {
if s.deleteIDs == nil {
return &DeleteResult{}, nil
}
return s.deleteIDs(ctx, ids)
}
func (s *fakePromptAdminService) PreviewDelete(ctx context.Context, filter EventFilter, actorID int64) (*DeletePreview, error) {
if s.preview == nil {
return &DeletePreview{}, nil
}
return s.preview(ctx, filter, actorID)
}
func (s *fakePromptAdminService) DeleteByFilter(ctx context.Context, req DeleteByFilterRequest, actorID int64) (*DeleteResult, error) {
if s.deleteFilter == nil {
return &DeleteResult{}, nil
}
return s.deleteFilter(ctx, req, actorID)
}
func promptAdminRouter(service PromptAdminService) *gin.Engine {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(string(servermiddleware.ContextKeyUser), servermiddleware.AuthSubject{UserID: 42})
c.Set(string(servermiddleware.ContextKeyUserRole), "admin")
c.Next()
})
handler := NewPromptAdminHandler(service)
group := router.Group("/admin/prompt-audit")
group.GET("/config", handler.GetConfig)
group.PUT("/config", handler.UpdateConfig)
group.POST("/endpoints/probe", handler.ProbeEndpoint)
group.GET("/runtime", handler.GetRuntime)
group.GET("/events", handler.ListEvents)
group.GET("/events/:id", handler.GetEvent)
group.DELETE("/events/:id", handler.DeleteEvent)
group.POST("/events/batch-delete", handler.BatchDelete)
group.POST("/events/delete-preview", handler.DeletePreview)
group.POST("/events/delete-by-filter", handler.DeleteByFilter)
return router
}
func promptAdminRequest(t *testing.T, router http.Handler, method, path string, body any) *httptest.ResponseRecorder {
t.Helper()
var reader *bytes.Reader
if body == nil {
reader = bytes.NewReader(nil)
} else {
raw, err := json.Marshal(body)
require.NoError(t, err)
reader = bytes.NewReader(raw)
}
req := httptest.NewRequest(method, path, reader)
req.Header.Set("Content-Type", "application/json")
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, req)
return recorder
}
func TestPromptAdminConfigRequiresVersionMapsConflictAndNeverEchoesToken(t *testing.T) {
const canary = "prompt-admin-token-canary"
t.Run("missing expected version", func(t *testing.T) {
router := promptAdminRouter(&fakePromptAdminService{})
response := promptAdminRequest(t, router, http.MethodPut, "/admin/prompt-audit/config", map[string]any{})
require.Equal(t, http.StatusBadRequest, response.Code)
require.Contains(t, response.Body.String(), "prompt_audit_invalid_config_request")
})
t.Run("CAS conflict", func(t *testing.T) {
service := &fakePromptAdminService{save: func(context.Context, UpdateConfigRequest, int64) (PublicConfig, error) {
return PublicConfig{}, infraerrors.Conflict(ErrorCodeConfigConflict, "配置已被更新")
}}
response := promptAdminRequest(t, promptAdminRouter(service), http.MethodPut, "/admin/prompt-audit/config", validHandlerUpdateRequest(canary))
require.Equal(t, http.StatusConflict, response.Code)
require.Contains(t, response.Body.String(), ErrorCodeConfigConflict)
require.NotContains(t, response.Body.String(), canary)
})
t.Run("success public DTO", func(t *testing.T) {
service := &fakePromptAdminService{save: func(_ context.Context, req UpdateConfigRequest, actorID int64) (PublicConfig, error) {
require.Equal(t, int64(42), actorID)
require.Equal(t, canary, req.Endpoints[0].Token)
return PublicConfig{ConfigVersion: 8, Endpoints: []PublicEndpoint{{ID: "guard-1", HasToken: true, TokenStatus: "configured"}}}, nil
}}
response := promptAdminRequest(t, promptAdminRouter(service), http.MethodPut, "/admin/prompt-audit/config", validHandlerUpdateRequest(canary))
require.Equal(t, http.StatusOK, response.Code)
body := response.Body.String()
require.NotContains(t, body, canary)
require.NotContains(t, body, "token_ciphertext")
require.NotContains(t, body, `"token":`)
require.Contains(t, body, `"has_token":true`)
})
}
func TestPromptAdminGetConfigReturnsSecretFreeUnavailableError(t *testing.T) {
const canary = "persisted-config-secret-canary"
repository := &switchableSettingRepository{loadErr: errors.New("failed to load token " + canary)}
manager := NewConfigManager(nil, repository, nil, prefixEncryptor{}, testTotpKeyConfig())
require.Error(t, manager.Reload(context.Background()))
service := &PromptService{config: manager}
response := promptAdminRequest(t, promptAdminRouter(service), http.MethodGet, "/admin/prompt-audit/config", nil)
require.Equal(t, http.StatusServiceUnavailable, response.Code)
require.Contains(t, response.Body.String(), ErrorCodeConfigUnavailable)
require.NotContains(t, response.Body.String(), canary)
require.NotContains(t, response.Body.String(), `"config_version"`)
require.NotContains(t, response.Body.String(), `"token"`)
}
func TestPromptAdminProbeSupportsTemporaryOrSavedTokenWithoutEcho(t *testing.T) {
const canary = "probe-token-canary"
for _, tc := range []struct {
name string
token string
tokenApplied bool
}{
{name: "temporary token", token: canary, tokenApplied: true},
{name: "saved token", token: "", tokenApplied: true},
} {
t.Run(tc.name, func(t *testing.T) {
service := &fakePromptAdminService{probe: func(_ context.Context, req ProbeRequest) ProbeResult {
require.Equal(t, tc.token, req.Endpoint.Token)
return ProbeResult{OK: true, Status: "healthy", Message: "ok", TokenApplied: tc.tokenApplied}
}}
endpoint := validHandlerUpdateRequest(tc.token).Endpoints[0]
response := promptAdminRequest(t, promptAdminRouter(service), http.MethodPost, "/admin/prompt-audit/endpoints/probe", ProbeRequest{Endpoint: endpoint})
require.Equal(t, http.StatusOK, response.Code)
require.NotContains(t, response.Body.String(), canary)
require.NotContains(t, response.Body.String(), `"token":`)
require.Contains(t, response.Body.String(), `"token_applied":true`)
})
}
}
func TestPromptAdminRejectsInvalidEventIDsTimesAndPagination(t *testing.T) {
router := promptAdminRouter(&fakePromptAdminService{})
for _, tc := range []struct {
method string
path string
body any
reason string
}{
{http.MethodGet, "/admin/prompt-audit/events/not-a-number", nil, "prompt_audit_invalid_event_id"},
{http.MethodDelete, "/admin/prompt-audit/events/-1", nil, "prompt_audit_invalid_event_id"},
{http.MethodGet, "/admin/prompt-audit/events?group_id=bad", nil, "prompt_audit_invalid_filter_id"},
{http.MethodGet, "/admin/prompt-audit/events?start_at=not-time", nil, "prompt_audit_invalid_time"},
{http.MethodGet, "/admin/prompt-audit/events?page=0", nil, "prompt_audit_invalid_pagination"},
{http.MethodPost, "/admin/prompt-audit/events/batch-delete", map[string]any{"ids": []int64{1, -2}}, "prompt_audit_invalid_event_id"},
} {
response := promptAdminRequest(t, router, tc.method, tc.path, tc.body)
require.Equalf(t, http.StatusBadRequest, response.Code, "%s %s", tc.method, tc.path)
require.Contains(t, response.Body.String(), tc.reason)
}
}
func validHandlerUpdateRequest(token string) UpdateConfigRequest {
return UpdateConfigRequest{
ExpectedConfigVersion: 7,
Strategy: "priority",
WorkerCount: 1,
QueueCapacity: 10,
Scanners: []string{"pii"},
AllGroups: true,
Endpoints: []UpdateEndpoint{{
ID: "guard-1", Name: "Guard One", Protocol: "openai_compatible",
BaseURL: "http://127.0.0.1:18080", Model: DefaultGuardModel, Token: token,
TimeoutMS: 1000, InputLimit: 1024, Enabled: true,
}},
}
}
func TestPromptAdminDeleteConfirmationErrorsStayGeneric(t *testing.T) {
service := &fakePromptAdminService{deleteFilter: func(context.Context, DeleteByFilterRequest, int64) (*DeleteResult, error) {
return nil, errors.New("sensitive-token-or-filter-detail")
}}
response := promptAdminRequest(t, promptAdminRouter(service), http.MethodPost, "/admin/prompt-audit/events/delete-by-filter", DeleteByFilterRequest{
SnapshotMaxID: 3, FilterHash: strings.Repeat("a", 64), ConfirmationToken: "secret-confirmation", Confirm: true,
})
require.Equal(t, http.StatusBadRequest, response.Code)
require.Contains(t, response.Body.String(), "prompt_audit_delete_confirmation_invalid")
require.NotContains(t, response.Body.String(), "sensitive-token")
require.NotContains(t, response.Body.String(), "secret-confirmation")
}
@@ -0,0 +1,69 @@
package securityaudit
import (
"crypto/sha256"
"encoding/hex"
)
func BuildIssueSummaries(result NormalizedResult) []IssueSummary {
resultCategories := result.Categories
if len(resultCategories) == 0 {
resultCategories = result.MatchedScanners
}
summaries := make([]IssueSummary, 0, len(resultCategories)+len(result.UnknownCategories))
for _, category := range resultCategories {
definition, ok := ScannerCatalog[category]
if !ok {
continue
}
evidence := RedactPreview(result.ScannerEvidence[category], 160)
if evidence == "" {
evidence = definition.Label
}
digest := sha256.Sum256([]byte(evidence))
summaries = append(summaries, IssueSummary{
Category: category, ScannerID: category, Title: definition.LabelZH,
Description: definition.Description, Severity: string(result.RiskLevel),
SeverityLabel: riskLabelZH(result.RiskLevel), Action: string(result.Action),
ActionLabel: actionLabelZH(result.Action), Code: "prompt_audit_" + category,
Score: result.ScannerScores[category], Evidence: evidence,
EvidenceHash: hex.EncodeToString(digest[:]),
})
}
for _, category := range result.UnknownCategories {
evidence := "unknown_unsafe"
digest := sha256.Sum256([]byte(evidence + ":" + category))
summaries = append(summaries, IssueSummary{
Category: category, ScannerID: "unknown_unsafe", Title: "未知高风险分类",
Description: "审计节点返回了未知但不可忽略的高风险分类", Severity: string(RiskCritical),
SeverityLabel: riskLabelZH(RiskCritical), Action: string(ActionBlock),
ActionLabel: actionLabelZH(ActionBlock), Code: "prompt_audit_unknown_unsafe",
Score: 1, Evidence: evidence, EvidenceHash: hex.EncodeToString(digest[:]),
})
}
return summaries
}
func riskLabelZH(risk RiskLevel) string {
switch risk {
case RiskCritical:
return "严重"
case RiskHigh:
return "高"
case RiskMedium:
return "中"
default:
return "低"
}
}
func actionLabelZH(action Action) string {
switch action {
case ActionBlock:
return "阻止"
case ActionWarn:
return "警告"
default:
return "允许"
}
}
@@ -0,0 +1,173 @@
package securityaudit
import (
"context"
"log/slog"
"strings"
)
const (
EventConfigUpdated = "prompt_audit.config_updated"
EventConfigLoaded = "prompt_guard.config_loaded"
EventConfigReloadDegraded = "prompt_guard.config_reload_degraded"
EventConfigTokenInvalid = "prompt_guard.config_token_invalid"
EventProbeStarted = "prompt_audit.endpoint_probe_started"
EventProbeFinished = "prompt_audit.endpoint_probe_finished"
EventProbeFailed = "prompt_audit.endpoint_probe_failed"
EventJobEnqueued = "prompt_audit.job_enqueued"
EventEnqueueSkipped = "prompt_audit.enqueue_skipped"
EventEnqueueDropped = "prompt_audit.enqueue_dropped"
EventAuditStarted = "prompt_audit.started"
EventProcessingReclaimed = "prompt_audit.processing_reclaimed"
EventProcessed = "prompt_audit.processed"
EventProcessFailed = "prompt_audit.process_failed"
EventFindingRecorded = "prompt_audit.finding_recorded"
EventChunkStarted = "prompt_audit.scan_chunk_started"
EventChunkCompleted = "prompt_audit.scan_chunk_completed"
EventChunkFailed = "prompt_audit.scan_chunk_failed"
EventChunksAggregated = "prompt_audit.scan_chunks_aggregated"
EventEvaluationStarted = "prompt_guard.evaluation_started"
EventGuardAllowed = "prompt_guard.allowed"
EventGuardBlocked = "prompt_guard.blocked"
EventGuardFailed = "prompt_guard.failed"
EventResultRecordFailed = "prompt_guard.result_record_failed"
EventEventDeleted = "prompt_audit.event_deleted"
EventEventsDeleted = "prompt_audit.events_deleted"
EventDeletePreviewed = "prompt_audit.events_delete_previewed"
EventEventsFilterDeleted = "prompt_audit.events_filter_deleted"
)
var knownLogEvents = map[string]struct{}{
EventConfigUpdated: {}, EventConfigLoaded: {}, EventConfigReloadDegraded: {}, EventConfigTokenInvalid: {},
EventProbeStarted: {}, EventProbeFinished: {}, EventProbeFailed: {},
EventJobEnqueued: {}, EventEnqueueSkipped: {}, EventEnqueueDropped: {},
EventAuditStarted: {}, EventProcessingReclaimed: {}, EventProcessed: {}, EventProcessFailed: {}, EventFindingRecorded: {},
EventChunkStarted: {}, EventChunkCompleted: {}, EventChunkFailed: {}, EventChunksAggregated: {},
EventEvaluationStarted: {}, EventGuardAllowed: {}, EventGuardBlocked: {}, EventGuardFailed: {}, EventResultRecordFailed: {},
EventEventDeleted: {}, EventEventsDeleted: {}, EventDeletePreviewed: {}, EventEventsFilterDeleted: {},
}
var allowedLogFields = map[string]struct{}{
"request_id": {}, "user_id": {}, "api_key_id": {}, "group_id": {}, "provider": {},
"protocol": {}, "endpoint": {}, "model": {}, "job_id": {}, "event_id": {},
"config_version": {}, "guard_endpoint_id": {}, "decision": {}, "risk_level": {},
"action": {}, "chunk_index": {}, "chunk_total": {}, "chunk_chars": {}, "input_chars": {},
"input_limit": {}, "latency_ms": {}, "status": {}, "error_code": {}, "error_kind": {},
"queue_length": {}, "queue_capacity": {}, "stage": {}, "upstream_dispatched": {},
"billing_preconsumed": {}, "worker_id": {}, "reclaimed_total": {}, "attempts": {},
"max_attempts": {}, "claim_version": {}, "http_status": {}, "retryable": {},
}
func LogInfo(event string, fields map[string]any) {
if _, ok := knownLogEvents[event]; !ok {
return
}
slog.LogAttrs(context.Background(), slog.LevelInfo, event, safeAttrs(fields)...)
}
func LogWarn(event string, fields map[string]any) {
if _, ok := knownLogEvents[event]; !ok {
return
}
slog.LogAttrs(context.Background(), slog.LevelWarn, event, safeAttrs(fields)...)
}
func LogError(event string, fields map[string]any) {
if _, ok := knownLogEvents[event]; !ok {
return
}
slog.LogAttrs(context.Background(), slog.LevelError, event, safeAttrs(fields)...)
}
func safeAttrs(fields map[string]any) []slog.Attr {
attrs := make([]slog.Attr, 0, len(fields))
for key, value := range fields {
key = strings.TrimSpace(key)
if _, allowed := allowedLogFields[key]; !allowed {
continue
}
if text, ok := value.(string); ok {
if key == "error_kind" || key == "error_code" {
value = stableErrorCode(text)
} else {
value = TrimRunes(strings.TrimSpace(text), 256)
}
}
attrs = append(attrs, slog.Any(key, value))
}
return attrs
}
func mergeLogFields(base map[string]any, extra map[string]any) map[string]any {
result := make(map[string]any, len(base)+len(extra))
for key, value := range base {
result[key] = value
}
for key, value := range extra {
result[key] = value
}
return result
}
func requestLogFields(req Request) map[string]any {
return map[string]any{
"request_id": req.RequestID, "user_id": req.UserID, "api_key_id": req.APIKeyID,
"group_id": pointerLogID(req.GroupID), "provider": req.Provider, "protocol": req.Protocol,
"endpoint": req.Endpoint, "model": req.Model, "stage": req.Stage,
}
}
func snapshotLogFields(snapshot PromptSnapshot) map[string]any {
return map[string]any{
"request_id": snapshot.RequestID, "user_id": snapshot.UserID, "api_key_id": snapshot.APIKeyID,
"group_id": pointerLogID(snapshot.GroupID), "provider": snapshot.Provider, "protocol": snapshot.Protocol,
"endpoint": snapshot.Endpoint, "model": snapshot.Model, "stage": snapshot.Stage,
}
}
func jobLogFields(job *Job) map[string]any {
if job == nil {
return map[string]any{}
}
fields := snapshotLogFields(job.Snapshot)
fields["job_id"] = job.ID
fields["config_version"] = job.ConfigVersion
fields["claim_version"] = job.ClaimVersion
return fields
}
func stableErrorCode(code string) string {
code = strings.ToLower(strings.TrimSpace(code))
if code == "" {
return "unknown_error"
}
for _, char := range code {
if (char >= 'a' && char <= 'z') || (char >= '0' && char <= '9') || char == '_' || char == '-' || char == '.' {
continue
}
return "redacted_error"
}
return TrimRunes(code, 64)
}
func stableErrorMessage(code string) string {
switch stableErrorCode(code) {
case ErrorCodeBlocked:
return "Prompt Guard blocked the request"
case ErrorCodeUnavailable, "payload_store_unavailable", "payload_missing":
return "Prompt Audit dependency is unavailable"
case ErrorCodeInvalidResponse:
return "Prompt Guard returned an invalid response"
case "queue_full", "queue_admission_busy":
return "Prompt Audit queue is unavailable"
case "worker_panic":
return "Prompt Audit worker failed"
case "config_load_failed", "config_ttl_reload_failed", "config_invalidation_reload_failed":
return "Prompt Audit configuration could not be loaded"
default:
return "Prompt Audit operation failed"
}
}
func sanitizeStoredError(code string) (string, string) {
stableCode := stableErrorCode(code)
return stableCode, TrimRunes(stableErrorMessage(stableCode), 160)
}
@@ -0,0 +1,66 @@
package securityaudit
import (
"bytes"
"encoding/json"
"log/slog"
"testing"
"time"
"github.com/stretchr/testify/require"
)
func TestPromptAuditLogAllowlistAndErrorsDoNotLeakCanarySecrets(t *testing.T) {
const canary = "PROMPT_AUDIT_CANARY_SECRET_DO_NOT_PERSIST"
var output bytes.Buffer
previous := slog.Default()
slog.SetDefault(slog.New(slog.NewJSONHandler(&output, nil)))
t.Cleanup(func() { slog.SetDefault(previous) })
LogWarn(EventConfigReloadDegraded, map[string]any{
"status": "degraded",
"error_code": "config_reload_failed",
"error_kind": "Authorization: Bearer " + canary,
"token": canary,
"body": canary,
"base_url": "https://guard.example.test/path?api_key=" + canary,
"raw_prompt": "prompt " + canary,
})
require.NotContains(t, output.String(), canary)
require.NotContains(t, output.String(), "api_key=")
require.Contains(t, output.String(), EventConfigReloadDegraded)
beforeUnknown := output.Len()
LogWarn("prompt_audit.typo_event", map[string]any{"status": "failed"})
require.Equal(t, beforeUnknown, output.Len(), "events outside the stable dictionary must not be emitted")
require.Len(t, knownLogEvents, 28)
_, err := NormalizeBaseURL("https://guard.example.test/path?token=" + canary)
require.Error(t, err)
require.NotContains(t, err.Error(), canary)
}
func TestPromptGuardFailureLogUsesCompleteAllowlistedContextAndNoSideEffects(t *testing.T) {
var output bytes.Buffer
previous := slog.Default()
slog.SetDefault(slog.New(slog.NewJSONHandler(&output, nil)))
t.Cleanup(func() { slog.SetDefault(previous) })
groupID := int64(9)
snapshot := PromptSnapshot{
RequestID: "req-1", UserID: 2, APIKeyID: 3, GroupID: &groupID,
Provider: "openai", Protocol: "openai_chat", Endpoint: "/v1/chat/completions",
Model: "gpt-test", Stage: "http",
}
logGuardFailure(snapshot, ActiveConfig{ConfigVersion: 7}, DecisionUnavailable, ErrorCodeUnavailable, "guard-1", 25*time.Millisecond)
var entry map[string]any
require.NoError(t, json.Unmarshal(output.Bytes(), &entry))
for key := range snapshotLogFields(snapshot) {
require.Contains(t, entry, key)
}
require.EqualValues(t, 7, entry["config_version"])
require.Equal(t, ErrorCodeUnavailable, entry["error_code"])
require.Equal(t, false, entry["upstream_dispatched"])
require.Equal(t, false, entry["billing_preconsumed"])
require.EqualValues(t, 25, entry["latency_ms"])
}
@@ -0,0 +1,145 @@
package securityaudit
import (
"sort"
"sync"
"sync/atomic"
"time"
)
const latencySampleCapacity = 2048
type AtomicMetrics struct {
total atomic.Int64
allowed atomic.Int64
flagged atomic.Int64
blocked atomic.Int64
unavailable atomic.Int64
invalid atomic.Int64
timeouts atomic.Int64
failovers atomic.Int64
bulkheadFull atomic.Int64
recordFailed atomic.Int64
latencyTotal atomic.Int64
latencyMax atomic.Int64
enqueued atomic.Int64
dropped atomic.Int64
latencyMu sync.RWMutex
latencies []int64
latencyNext int
}
func NewAtomicMetrics() *AtomicMetrics { return &AtomicMetrics{} }
func (m *AtomicMetrics) Snapshot() GuardMetricsSnapshot {
if m == nil {
return GuardMetricsSnapshot{}
}
snapshot := GuardMetricsSnapshot{
Total: m.total.Load(), Allowed: m.allowed.Load(), Flagged: m.flagged.Load(),
Blocked: m.blocked.Load(), Unavailable: m.unavailable.Load(), Invalid: m.invalid.Load(),
Timeouts: m.timeouts.Load(), Failovers: m.failovers.Load(), BulkheadFull: m.bulkheadFull.Load(),
RecordFailed: m.recordFailed.Load(), LatencyCount: m.total.Load(), LatencyMaxMS: m.latencyMax.Load(),
}
if snapshot.LatencyCount > 0 {
snapshot.LatencyAvgMS = m.latencyTotal.Load() / snapshot.LatencyCount
}
m.latencyMu.RLock()
samples := append([]int64(nil), m.latencies...)
m.latencyMu.RUnlock()
if len(samples) > 0 {
sort.Slice(samples, func(i, j int) bool { return samples[i] < samples[j] })
snapshot.LatencyP50MS = percentile(samples, 0.50)
snapshot.LatencyP95MS = percentile(samples, 0.95)
snapshot.LatencyP99MS = percentile(samples, 0.99)
}
return snapshot
}
func (m *AtomicMetrics) AuditSnapshot() AuditMetricsSnapshot {
if m == nil {
return AuditMetricsSnapshot{}
}
return AuditMetricsSnapshot{Enqueued: m.enqueued.Load(), Dropped: m.dropped.Load()}
}
func (m *AtomicMetrics) Observe(kind DecisionKind, latency time.Duration) {
if m == nil {
return
}
m.total.Add(1)
latencyMS := latency.Milliseconds()
if latencyMS < 0 {
latencyMS = 0
}
m.latencyTotal.Add(latencyMS)
for current := m.latencyMax.Load(); latencyMS > current && !m.latencyMax.CompareAndSwap(current, latencyMS); current = m.latencyMax.Load() {
}
m.latencyMu.Lock()
if len(m.latencies) < latencySampleCapacity {
m.latencies = append(m.latencies, latencyMS)
} else {
m.latencies[m.latencyNext] = latencyMS
m.latencyNext = (m.latencyNext + 1) % latencySampleCapacity
}
m.latencyMu.Unlock()
switch kind {
case DecisionFlag:
m.flagged.Add(1)
case DecisionBlock:
m.blocked.Add(1)
case DecisionUnavailable:
m.unavailable.Add(1)
case DecisionInvalid:
m.invalid.Add(1)
default:
m.allowed.Add(1)
}
}
func percentile(sorted []int64, quantile float64) int64 {
if len(sorted) == 0 {
return 0
}
index := int(float64(len(sorted)-1) * quantile)
if index < 0 {
index = 0
}
if index >= len(sorted) {
index = len(sorted) - 1
}
return sorted[index]
}
func (m *AtomicMetrics) IncEnqueued() {
if m != nil {
m.enqueued.Add(1)
}
}
func (m *AtomicMetrics) IncDropped() {
if m != nil {
m.dropped.Add(1)
}
}
func (m *AtomicMetrics) IncTimeout() {
if m != nil {
m.timeouts.Add(1)
}
}
func (m *AtomicMetrics) IncFailover() {
if m != nil {
m.failovers.Add(1)
}
}
func (m *AtomicMetrics) IncBulkheadFull() {
if m != nil {
m.bulkheadFull.Add(1)
}
}
func (m *AtomicMetrics) IncRecordFailed() {
if m != nil {
m.recordFailed.Add(1)
}
}
@@ -0,0 +1,52 @@
package securityaudit
import (
"sync"
"testing"
"time"
"github.com/stretchr/testify/require"
)
func TestAtomicMetricsExposeCountsLatencyDistributionAndAsyncDelivery(t *testing.T) {
metrics := NewAtomicMetrics()
latencies := []time.Duration{10, 20, 30, 40, 100}
kinds := []DecisionKind{DecisionAllow, DecisionFlag, DecisionBlock, DecisionUnavailable, DecisionInvalid}
for index := range latencies {
metrics.Observe(kinds[index], latencies[index]*time.Millisecond)
}
metrics.IncTimeout()
metrics.IncFailover()
metrics.IncBulkheadFull()
metrics.IncRecordFailed()
metrics.IncEnqueued()
metrics.IncDropped()
snapshot := metrics.Snapshot()
require.Equal(t, int64(5), snapshot.Total)
require.Equal(t, int64(5), snapshot.LatencyCount)
require.Equal(t, int64(40), snapshot.LatencyAvgMS)
require.Equal(t, int64(30), snapshot.LatencyP50MS)
require.Equal(t, int64(40), snapshot.LatencyP95MS)
require.Equal(t, int64(40), snapshot.LatencyP99MS)
require.Equal(t, int64(100), snapshot.LatencyMaxMS)
require.Equal(t, AuditMetricsSnapshot{Enqueued: 1, Dropped: 1}, metrics.AuditSnapshot())
}
func TestAtomicMetricsConcurrentObservationIsBoundedAndRaceSafe(t *testing.T) {
metrics := NewAtomicMetrics()
const observations = 4096
var wg sync.WaitGroup
for index := 0; index < observations; index++ {
wg.Add(1)
go func(value int) {
defer wg.Done()
metrics.Observe(DecisionAllow, time.Duration(value%250)*time.Millisecond)
}(index)
}
wg.Wait()
require.Equal(t, int64(observations), metrics.Snapshot().Total)
metrics.latencyMu.RLock()
require.LessOrEqual(t, len(metrics.latencies), latencySampleCapacity)
metrics.latencyMu.RUnlock()
}
@@ -0,0 +1,23 @@
package securityaudit
import "github.com/google/wire"
var ProviderSet = wire.NewSet(
NewPostgreSQLRepository,
wire.Bind(new(JobRepository), new(*PostgreSQLRepository)),
wire.Bind(new(EventRepository), new(*PostgreSQLRepository)),
NewRedisPayloadStore,
wire.Bind(new(PayloadStore), new(*RedisPayloadStore)),
NewOpenAICompatibleScanner,
wire.Bind(new(PromptScanner), new(*OpenAICompatibleScanner)),
NewAtomicMetrics,
wire.Bind(new(Metrics), new(*AtomicMetrics)),
NewConfigManager,
wire.Bind(new(ConfigStore), new(*ConfigManager)),
NewPromptService,
wire.Bind(new(PromptEngine), new(*PromptService)),
wire.Bind(new(PromptAdminService), new(*PromptService)),
NewLegacyModerationAdapter,
NewCoordinator,
NewPromptAdminHandler,
)
@@ -0,0 +1,89 @@
package securityaudit
import (
"crypto/tls"
"net"
"net/http"
"net/url"
"strings"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
)
const maxGuardResponseBytes int64 = 256 * 1024
func NormalizeBaseURL(raw string) (string, error) {
raw = strings.TrimSpace(raw)
parsed, err := url.Parse(raw)
if err != nil || parsed.Scheme == "" || parsed.Host == "" {
return "", infraerrors.BadRequest("prompt_audit_invalid_base_url", "审计节点地址无效")
}
parsed.Scheme = strings.ToLower(parsed.Scheme)
if parsed.Scheme != "http" && parsed.Scheme != "https" {
return "", infraerrors.BadRequest("prompt_audit_invalid_base_url_scheme", "审计节点仅支持 HTTP(S)")
}
if parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" {
return "", infraerrors.BadRequest("prompt_audit_unsafe_base_url", "审计节点地址不能包含凭据、查询参数或片段")
}
host := strings.TrimSpace(parsed.Hostname())
if host == "" {
return "", infraerrors.BadRequest("prompt_audit_invalid_base_url", "审计节点地址无效")
}
path := strings.TrimRight(parsed.EscapedPath(), "/")
if strings.EqualFold(path, "/v1") {
path = ""
}
parsed.Path = path
parsed.RawPath = ""
return strings.TrimRight(parsed.String(), "/"), nil
}
func ChatCompletionsURL(base string) (string, error) {
normalized, err := NormalizeBaseURL(base)
if err != nil {
return "", err
}
return normalized + "/v1/chat/completions", nil
}
func ModelsURL(base string) (string, error) {
normalized, err := NormalizeBaseURL(base)
if err != nil {
return "", err
}
return normalized + "/v1/models", nil
}
func NewSecureHTTPClient(endpoint ActiveEndpoint) (*http.Client, error) {
_, err := NormalizeBaseURL(endpoint.BaseURL)
if err != nil {
return nil, err
}
dialer := &net.Dialer{Timeout: 3 * time.Second, KeepAlive: 30 * time.Second}
transport := &http.Transport{
// Do not inherit HTTP(S)_PROXY. A proxy would move the actual destination
// dial outside secureDialContext and bypass this module's DNS/IP validation.
Proxy: nil,
ForceAttemptHTTP2: true,
MaxIdleConns: 64,
MaxIdleConnsPerHost: 16,
IdleConnTimeout: 90 * time.Second,
TLSHandshakeTimeout: 5 * time.Second,
ResponseHeaderTimeout: time.Duration(endpoint.TimeoutMS) * time.Millisecond,
ExpectContinueTimeout: time.Second,
TLSClientConfig: &tls.Config{MinVersion: tls.VersionTLS12},
}
// Endpoint ownership and destination trust are administrator concerns.
// Use the standard dialer so configured private, loopback, reserved, and
// DNS-resolved addresses are all reachable from the service environment.
transport.DialContext = dialer.DialContext
timeout := time.Duration(endpoint.TimeoutMS) * time.Millisecond
if timeout <= 0 {
timeout = DefaultTimeoutMS * time.Millisecond
}
return &http.Client{
Transport: transport,
Timeout: timeout,
}, nil
}
@@ -0,0 +1,236 @@
package securityaudit
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/stretchr/testify/require"
)
func TestNormalizeBaseURLAllowsAdministratorConfiguredDestinations(t *testing.T) {
allowed := []string{
"https://guard.example.com", "https://guard.example.com/v1", "http://guard.example.com",
"http://127.0.0.1:8080", "http://10.0.0.8:8080", "https://172.16.0.5",
"http://169.254.169.254", "https://metadata.google.internal", "https://192.0.2.1",
"http://internal-admin.local", "http://guard.local:8080",
}
for _, raw := range allowed {
_, err := NormalizeBaseURL(raw)
require.NoError(t, err, raw)
}
blocked := []string{
"ftp://guard.example.com", "https://user:pass@guard.example.com",
"https://guard.example.com?q=secret", "https://guard.example.com/#fragment",
}
for _, raw := range blocked {
_, err := NormalizeBaseURL(raw)
require.Error(t, err, raw)
}
url, err := ChatCompletionsURL("https://guard.example.com/v1")
require.NoError(t, err)
require.Equal(t, "https://guard.example.com/v1/chat/completions", url)
}
func TestHTTPClientUsesDirectStandardDialer(t *testing.T) {
client, err := NewSecureHTTPClient(ActiveEndpoint{BaseURL: "https://guard.example.com", TimeoutMS: 1000})
require.NoError(t, err)
transport, ok := client.Transport.(*http.Transport)
require.True(t, ok)
require.Nil(t, transport.Proxy)
require.NotNil(t, transport.DialContext)
}
func TestOpenAICompatibleScannerRequestContract(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
require.Equal(t, "/v1/chat/completions", r.URL.Path)
require.Equal(t, "Bearer token", r.Header.Get("Authorization"))
var payload map[string]any
require.NoError(t, json.NewDecoder(r.Body).Decode(&payload))
require.Equal(t, DefaultGuardModel, payload["model"])
require.Equal(t, float64(0), payload["temperature"])
require.Equal(t, float64(64), payload["max_tokens"])
require.Equal(t, float64(42), payload["seed"])
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"Safety: Safe\nCategories: None"}}]}`))
}))
defer server.Close()
scanner := NewOpenAICompatibleScanner()
result, err := scanner.Scan(context.Background(), ActiveEndpoint{ID: "one", BaseURL: server.URL, Model: DefaultGuardModel, Token: "token", TimeoutMS: 1000}, "hello", AllScannerIDs)
require.NoError(t, err)
require.Equal(t, EventPass, result.Decision)
}
func TestOpenAICompatibleScannerFollowsRedirectAndRejectsOversize(t *testing.T) {
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"Safety: Safe\nCategories: None"}}]}`))
}))
defer target.Close()
redirect := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Redirect(w, r, target.URL, http.StatusFound) }))
defer redirect.Close()
result, err := NewOpenAICompatibleScanner().Scan(context.Background(), ActiveEndpoint{ID: "redirect", BaseURL: redirect.URL, Model: DefaultGuardModel, TimeoutMS: 1000}, "hello", AllScannerIDs)
require.NoError(t, err)
require.Equal(t, EventPass, result.Decision)
oversize := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
_, _ = w.Write([]byte(strings.Repeat("x", int(maxGuardResponseBytes)+1)))
}))
defer oversize.Close()
_, err = NewOpenAICompatibleScanner().Scan(context.Background(), ActiveEndpoint{ID: "large", BaseURL: oversize.URL, Model: DefaultGuardModel, TimeoutMS: 1000}, "hello", AllScannerIDs)
require.Error(t, err)
}
func TestOpenAICompatibleScannerClassifiesHTTPConnectionAndTimeoutFailures(t *testing.T) {
tests := []struct {
name string
status int
retryable bool
}{
{name: "authentication", status: http.StatusUnauthorized, retryable: false},
{name: "forbidden", status: http.StatusForbidden, retryable: false},
{name: "rate limited", status: http.StatusTooManyRequests, retryable: true},
{name: "server failure", status: http.StatusBadGateway, retryable: true},
{name: "other client error", status: http.StatusBadRequest, retryable: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(tt.status)
}))
defer server.Close()
_, err := NewOpenAICompatibleScanner().Scan(context.Background(), ActiveEndpoint{ID: "status", BaseURL: server.URL, Model: DefaultGuardModel, TimeoutMS: 1000}, "hello", AllScannerIDs)
var guardErr *GuardError
require.ErrorAs(t, err, &guardErr)
require.Equal(t, ErrorCodeUnavailable, guardErr.Code)
require.Equal(t, tt.status, guardErr.HTTPStatus)
require.Equal(t, tt.retryable, guardErr.Retryable)
require.NotContains(t, err.Error(), server.URL)
})
}
closed := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
closedURL := closed.URL
closed.Close()
_, err := NewOpenAICompatibleScanner().Scan(context.Background(), ActiveEndpoint{ID: "closed", BaseURL: closedURL, Model: DefaultGuardModel, TimeoutMS: 100}, "hello", AllScannerIDs)
var connectionErr *GuardError
require.ErrorAs(t, err, &connectionErr)
require.True(t, connectionErr.Retryable)
timeout := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
time.Sleep(100 * time.Millisecond)
w.WriteHeader(http.StatusOK)
}))
defer timeout.Close()
_, err = NewOpenAICompatibleScanner().Scan(context.Background(), ActiveEndpoint{ID: "timeout", BaseURL: timeout.URL, Model: DefaultGuardModel, TimeoutMS: 20}, "hello", AllScannerIDs)
var timeoutErr *GuardError
require.ErrorAs(t, err, &timeoutErr)
require.True(t, timeoutErr.Retryable)
require.True(t, timeoutErr.Timeout)
}
func TestPromptAuditProbeModelsFallbackAndResponseSafety(t *testing.T) {
t.Run("models contains configured model", func(t *testing.T) {
var chatCalls atomic.Int64
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
require.Equal(t, "Bearer temporary-token", r.Header.Get("Authorization"))
if r.URL.Path == "/v1/models" {
_, _ = w.Write([]byte(`{"data":[{"id":"` + DefaultGuardModel + `"}]}`))
return
}
chatCalls.Add(1)
}))
defer server.Close()
result := newProbeTestService().Probe(context.Background(), ProbeRequest{Endpoint: probeEndpoint(server.URL, "temporary-token")})
require.True(t, result.OK)
require.True(t, result.TokenApplied)
require.Equal(t, http.StatusOK, result.HTTPStatus)
require.Zero(t, chatCalls.Load())
})
t.Run("invalid models response performs real guard fallback", func(t *testing.T) {
var chatCalls atomic.Int64
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/v1/models" {
_, _ = w.Write([]byte(`{"unexpected":true}`))
return
}
chatCalls.Add(1)
_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"Safety: Safe\nCategories: None"}}]}`))
}))
defer server.Close()
result := newProbeTestService().Probe(context.Background(), ProbeRequest{Endpoint: probeEndpoint(server.URL, "temporary-token")})
require.True(t, result.OK)
require.Equal(t, int64(1), chatCalls.Load())
})
t.Run("fallback authentication failure is stable", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/v1/models" {
w.WriteHeader(http.StatusNotFound)
return
}
w.WriteHeader(http.StatusUnauthorized)
}))
defer server.Close()
result := newProbeTestService().Probe(context.Background(), ProbeRequest{Endpoint: probeEndpoint(server.URL, "temporary-token")})
require.False(t, result.OK)
require.Equal(t, ErrorCodeUnavailable, result.ErrorCode)
require.Equal(t, http.StatusUnauthorized, result.HTTPStatus)
require.False(t, result.Retryable)
})
t.Run("oversized models response is rejected without fallback", func(t *testing.T) {
var chatCalls atomic.Int64
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/v1/models" {
chatCalls.Add(1)
}
_, _ = w.Write([]byte(strings.Repeat("x", int(maxGuardResponseBytes)+1)))
}))
defer server.Close()
result := newProbeTestService().Probe(context.Background(), ProbeRequest{Endpoint: probeEndpoint(server.URL, "temporary-token")})
require.False(t, result.OK)
require.Equal(t, "response_too_large", result.ErrorCode)
require.Zero(t, chatCalls.Load())
})
}
func TestResolveProbeEndpointReusesTokenOnlyForMatchingBaseURL(t *testing.T) {
manager := &ConfigManager{}
manager.snapshot.Store(&activeConfigSnapshot{active: ActiveConfig{Endpoints: []ActiveEndpoint{{
ID: "guard-1", BaseURL: "https://guard.example.com", Token: "STORED_GUARD_TOKEN", TimeoutMS: 1000, InputLimit: 1024, Enabled: true,
}}}})
service := &PromptService{config: manager}
matched, applied, err := service.resolveProbeEndpoint(UpdateEndpoint{
ID: "guard-1", BaseURL: "https://guard.example.com/v1", TimeoutMS: 1000, InputLimit: 1024,
})
require.NoError(t, err)
require.True(t, applied)
require.Equal(t, "STORED_GUARD_TOKEN", matched.Token)
mismatched, applied, err := service.resolveProbeEndpoint(UpdateEndpoint{
ID: "guard-1", BaseURL: "https://attacker.example.com", TimeoutMS: 1000, InputLimit: 1024,
})
require.NoError(t, err)
require.False(t, applied)
require.Empty(t, mismatched.Token)
}
func newProbeTestService() *PromptService {
return &PromptService{
config: &ConfigManager{}, scanner: NewOpenAICompatibleScanner(), clock: realClock{},
probes: map[string]ProbeResult{},
}
}
func probeEndpoint(baseURL, token string) UpdateEndpoint {
return UpdateEndpoint{
ID: "probe-one", Name: "Probe One", Protocol: "openai_compatible", BaseURL: baseURL,
Model: DefaultGuardModel, Token: token, TimeoutMS: 1000, InputLimit: 1024, Enabled: true,
}
}
@@ -0,0 +1,63 @@
package securityaudit
import (
"context"
"fmt"
"strconv"
"time"
"github.com/redis/go-redis/v9"
)
type PayloadStore interface {
Set(ctx context.Context, jobID int64, scanText string, ttl time.Duration) error
Get(ctx context.Context, jobID int64) (string, error)
Delete(ctx context.Context, jobID int64) error
Ping(ctx context.Context) error
}
type RedisPayloadStore struct {
client *redis.Client
}
func NewRedisPayloadStore(client *redis.Client) *RedisPayloadStore {
return &RedisPayloadStore{client: client}
}
func (s *RedisPayloadStore) Set(ctx context.Context, jobID int64, scanText string, ttl time.Duration) error {
if s == nil || s.client == nil {
return fmt.Errorf("prompt audit payload store unavailable")
}
if jobID <= 0 || scanText == "" {
return fmt.Errorf("prompt audit payload input invalid")
}
if ttl <= 0 || ttl > DefaultPayloadTTL {
ttl = DefaultPayloadTTL
}
return s.client.Set(ctx, payloadKey(jobID), scanText, ttl).Err()
}
func (s *RedisPayloadStore) Get(ctx context.Context, jobID int64) (string, error) {
if s == nil || s.client == nil {
return "", fmt.Errorf("prompt audit payload store unavailable")
}
return s.client.Get(ctx, payloadKey(jobID)).Result()
}
func (s *RedisPayloadStore) Delete(ctx context.Context, jobID int64) error {
if s == nil || s.client == nil {
return fmt.Errorf("prompt audit payload store unavailable")
}
return s.client.Del(ctx, payloadKey(jobID)).Err()
}
func (s *RedisPayloadStore) Ping(ctx context.Context) error {
if s == nil || s.client == nil {
return fmt.Errorf("prompt audit payload store unavailable")
}
return s.client.Ping(ctx).Err()
}
func payloadKey(jobID int64) string {
return PayloadKeyPrefix + strconv.FormatInt(jobID, 10)
}
@@ -0,0 +1,85 @@
package securityaudit
import (
"context"
"os"
"strings"
"testing"
"time"
"github.com/redis/go-redis/v9"
"github.com/stretchr/testify/require"
)
func TestRedisPayloadStoreRoundTripTTLNamespaceAndDelete(t *testing.T) {
address := strings.TrimSpace(os.Getenv(promptAuditRedisTestEnv))
if address == "" {
t.Skip(promptAuditRedisTestEnv + " is not set")
}
client := redis.NewClient(&redis.Options{Addr: address})
t.Cleanup(func() { require.NoError(t, client.Close()) })
store := NewRedisPayloadStore(client)
ctx := context.Background()
const jobID int64 = 987654321
const canary = "PROMPT_CANARY_REDIS_ONLY_PAYLOAD"
_ = store.Delete(ctx, jobID)
require.NoError(t, store.Set(ctx, jobID, canary, 2*DefaultPayloadTTL))
require.Equal(t, PayloadKeyPrefix+"987654321", payloadKey(jobID))
value, err := store.Get(ctx, jobID)
require.NoError(t, err)
require.Equal(t, canary, value)
ttl, err := client.TTL(ctx, payloadKey(jobID)).Result()
require.NoError(t, err)
require.Greater(t, ttl, time.Duration(0))
require.LessOrEqual(t, ttl, DefaultPayloadTTL)
require.NoError(t, store.Delete(ctx, jobID))
_, err = store.Get(ctx, jobID)
require.ErrorIs(t, err, redis.Nil)
}
func TestPromptRuntimeAggregatesConfigWorkersQueueRedisEndpointsAndGuardMetrics(t *testing.T) {
address := strings.TrimSpace(os.Getenv(promptAuditRedisTestEnv))
if address == "" {
t.Skip(promptAuditRedisTestEnv + " is not set")
}
db := openPromptAuditIntegrationDB(t)
client := redis.NewClient(&redis.Options{Addr: address})
t.Cleanup(func() { require.NoError(t, client.Close()) })
config := &fakeConfigStore{active: true, cfg: ActiveConfig{
RiskControlEnabled: true, Enabled: true, WorkerCount: 3, QueueCapacity: 123,
ConfigVersion: 9, AllGroups: true,
}}
metrics := NewAtomicMetrics()
metrics.Observe(DecisionBlock, 25*time.Millisecond)
metrics.IncFailover()
metrics.IncEnqueued()
metrics.IncDropped()
service := NewPromptService(
config,
NewPostgreSQLRepository(db),
NewRedisPayloadStore(client),
NewOpenAICompatibleScanner(),
metrics,
)
service.probes["guard-1"] = ProbeResult{OK: true, Status: "healthy", HTTPStatus: 200}
runtime := service.Runtime(context.Background())
require.Equal(t, ModeAsync, runtime.EffectiveMode)
require.Equal(t, int64(9), runtime.ExpectedConfigVersion)
require.Equal(t, int64(9), runtime.ActiveConfigVersion)
require.Equal(t, 3, runtime.WorkerTotal)
require.Equal(t, 123, runtime.QueueCapacity)
require.Equal(t, "ok", runtime.DatabaseStatus)
require.Equal(t, "ok", runtime.RedisStatus)
require.Contains(t, runtime.Endpoints, "guard-1")
require.Equal(t, int64(1), runtime.GuardMetrics.Total)
require.Equal(t, int64(1), runtime.GuardMetrics.Blocked)
require.Equal(t, int64(1), runtime.GuardMetrics.Failovers)
require.Equal(t, int64(25), runtime.GuardMetrics.LatencyP95MS)
require.Equal(t, int64(1), runtime.EnqueuedTotal)
require.Equal(t, int64(1), runtime.DroppedTotal)
// The runner has not been started in this integration test, so the honest
// process status is degraded rather than a fabricated running heartbeat.
require.Equal(t, "degraded", runtime.ProcessStatus)
}
@@ -0,0 +1,327 @@
package securityaudit
import (
"bytes"
"context"
"crypto/sha256"
"encoding/json"
"errors"
"fmt"
"io"
"net"
"net/http"
"sort"
"strings"
"sync"
)
type ScannerDefinition struct {
ID string `json:"id"`
Label string `json:"label"`
LabelZH string `json:"label_zh"`
Description string `json:"description"`
}
var AllScannerIDs = []string{
"violent",
"non_violent_illegal_acts",
"sexual_content_or_sexual_acts",
"pii",
"suicide_and_self_harm",
"unethical_acts",
"politically_sensitive_topics",
"copyright_violation",
"jailbreak",
}
var ScannerCatalog = map[string]ScannerDefinition{
"violent": {ID: "violent", Label: "Violent", LabelZH: "暴力", Description: "Violence or threats of violence"},
"non_violent_illegal_acts": {ID: "non_violent_illegal_acts", Label: "Non-violent Illegal Acts", LabelZH: "非暴力违法行为", Description: "Non-violent illegal activity"},
"sexual_content_or_sexual_acts": {ID: "sexual_content_or_sexual_acts", Label: "Sexual Content or Sexual Acts", LabelZH: "性内容或性行为", Description: "Sexual content or sexual acts"},
"pii": {ID: "pii", Label: "PII", LabelZH: "个人敏感信息", Description: "Personal identifying information"},
"suicide_and_self_harm": {ID: "suicide_and_self_harm", Label: "Suicide & Self-Harm", LabelZH: "自杀与自残", Description: "Suicide or self-harm"},
"unethical_acts": {ID: "unethical_acts", Label: "Unethical Acts", LabelZH: "不道德行为", Description: "Unethical behavior"},
"politically_sensitive_topics": {ID: "politically_sensitive_topics", Label: "Politically Sensitive Topics", LabelZH: "政治敏感话题", Description: "Politically sensitive topics"},
"copyright_violation": {ID: "copyright_violation", Label: "Copyright Violation", LabelZH: "版权侵权", Description: "Copyright infringement"},
"jailbreak": {ID: "jailbreak", Label: "Jailbreak", LabelZH: "越狱攻击", Description: "Prompt injection or jailbreak attempt"},
}
var categoryAliases = map[string]string{
"violent": "violent", "violence": "violent",
"non violent illegal acts": "non_violent_illegal_acts", "non-violent illegal acts": "non_violent_illegal_acts",
"sexual content or sexual acts": "sexual_content_or_sexual_acts", "sexual": "sexual_content_or_sexual_acts",
"pii": "pii", "personal identifying information": "pii", "personal identifiable information": "pii",
"suicide self harm": "suicide_and_self_harm", "suicide and self harm": "suicide_and_self_harm", "suicide & self-harm": "suicide_and_self_harm",
"unethical acts": "unethical_acts", "unethical": "unethical_acts",
"politically sensitive topics": "politically_sensitive_topics", "political": "politically_sensitive_topics",
"copyright violation": "copyright_violation", "copyright": "copyright_violation",
"jailbreak": "jailbreak", "prompt injection": "jailbreak",
}
type GuardError struct {
Code string
HTTPStatus int
Retryable bool
Timeout bool
Cause error
}
func (e *GuardError) Error() string {
if e == nil {
return "<nil>"
}
return e.Code
}
func (e *GuardError) Unwrap() error { return e.Cause }
func NormalizeCategory(value string) string {
normalized := strings.ToLower(strings.TrimSpace(value))
normalized = strings.NewReplacer("_", " ", "&", " and ", "/", " ", "-", " ", "", " ", "—", " ").Replace(normalized)
normalized = strings.Join(strings.Fields(normalized), " ")
if canonical, ok := categoryAliases[normalized]; ok {
return canonical
}
return strings.ReplaceAll(normalized, " ", "_")
}
func ParseQwen3Guard(content string, enabledScanners []string) (*NormalizedResult, error) {
var safety string
var categoryLine string
for _, line := range strings.Split(strings.ReplaceAll(content, "\r\n", "\n"), "\n") {
line = strings.TrimSpace(line)
if line == "" {
continue
}
lower := strings.ToLower(line)
switch {
case strings.HasPrefix(lower, "safety:"):
if safety != "" {
return nil, &GuardError{Code: ErrorCodeInvalidResponse}
}
safety = strings.TrimSpace(line[len("safety:"):])
case strings.HasPrefix(lower, "categories:"):
if categoryLine != "" {
return nil, &GuardError{Code: ErrorCodeInvalidResponse}
}
categoryLine = strings.TrimSpace(line[len("categories:"):])
default:
// Auxiliary Guard fields, such as Refusal, do not affect audit decisions.
}
}
switch strings.ToLower(safety) {
case "safe":
safety = "Safe"
case "controversial":
safety = "Controversial"
case "unsafe":
safety = "Unsafe"
default:
return nil, &GuardError{Code: ErrorCodeInvalidResponse}
}
if categoryLine == "" {
return nil, &GuardError{Code: ErrorCodeInvalidResponse}
}
enabled := make(map[string]struct{}, len(enabledScanners))
for _, scanner := range enabledScanners {
enabled[NormalizeCategory(scanner)] = struct{}{}
}
known := map[string]struct{}{}
unknown := map[string]struct{}{}
for _, raw := range strings.Split(categoryLine, ",") {
raw = strings.TrimSpace(raw)
if raw == "" || strings.EqualFold(raw, "none") || strings.EqualFold(raw, "n/a") {
continue
}
category := NormalizeCategory(raw)
if _, ok := ScannerCatalog[category]; ok {
known[category] = struct{}{}
} else {
unknown[unknownCategoryID(category)] = struct{}{}
}
}
knownList := orderedScannerKeys(known)
unknownList := sortedKeys(unknown)
matched := make([]string, 0, len(knownList))
for _, category := range knownList {
if _, ok := enabled[category]; ok {
matched = append(matched, category)
}
}
result := &NormalizedResult{
Safety: safety, Categories: knownList, MatchedScanners: matched, UnknownCategories: unknownList,
ScannerScores: map[string]float64{}, ScannerEvidence: map[string]string{},
ScannerBackend: "qwen3guard-openai", ScannerVersion: "qwen3guard",
PolicyID: "priority", PolicyVersion: 1,
Decision: EventPass, RiskLevel: RiskLow, Action: ActionAllow,
}
score := 0.0
if safety == "Controversial" {
score = 0.5
result.Decision, result.RiskLevel, result.Action = EventFlag, RiskMedium, ActionWarn
}
if safety == "Unsafe" {
score = 1
if len(matched) > 0 || len(unknownList) > 0 || len(knownList) == 0 {
result.Decision, result.RiskLevel, result.Action = EventCritical, RiskCritical, ActionBlock
} else {
result.Decision, result.RiskLevel, result.Action = EventFlag, RiskHigh, ActionWarn
}
}
for _, category := range matched {
result.ScannerScores[category] = score
result.ScannerEvidence[category] = ScannerCatalog[category].Label
if safety == "Controversial" && isElevatedControversial(category) {
result.Decision, result.RiskLevel, result.Action = EventCritical, RiskCritical, ActionBlock
}
}
return result, nil
}
func unknownCategoryID(value string) string {
digest := sha256.Sum256([]byte(strings.TrimSpace(strings.ToLower(value))))
return fmt.Sprintf("unknown:%x", digest[:8])
}
func isElevatedControversial(category string) bool {
return category == "jailbreak" || category == "pii" || category == "suicide_and_self_harm"
}
type OpenAICompatibleScanner struct {
clients sync.Map
}
func NewOpenAICompatibleScanner() *OpenAICompatibleScanner { return &OpenAICompatibleScanner{} }
func (s *OpenAICompatibleScanner) Scan(ctx context.Context, endpoint ActiveEndpoint, chunk string, enabledScanners []string) (*NormalizedResult, error) {
client, err := s.clientFor(endpoint)
if err != nil {
return nil, &GuardError{Code: ErrorCodeUnavailable, Cause: err}
}
requestURL, err := ChatCompletionsURL(endpoint.BaseURL)
if err != nil {
return nil, &GuardError{Code: ErrorCodeUnavailable, Cause: err}
}
payload := map[string]any{
"model": endpoint.Model,
"messages": []map[string]string{{"role": "user", "content": chunk}},
"temperature": 0,
"max_tokens": 64,
"seed": 42,
}
body, err := json.Marshal(payload)
if err != nil {
return nil, &GuardError{Code: ErrorCodeInvalidResponse, Cause: err}
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, requestURL, bytes.NewReader(body))
if err != nil {
return nil, &GuardError{Code: ErrorCodeUnavailable, Cause: err}
}
req.Header.Set("Content-Type", "application/json")
if endpoint.Token != "" {
req.Header.Set("Authorization", "Bearer "+endpoint.Token)
}
resp, err := client.Do(req)
if err != nil {
timeout := errors.Is(err, context.DeadlineExceeded)
var netErr net.Error
if errors.As(err, &netErr) && netErr.Timeout() {
timeout = true
}
return nil, &GuardError{Code: ErrorCodeUnavailable, Retryable: true, Timeout: timeout, Cause: err}
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
retryable := resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500
return nil, &GuardError{Code: ErrorCodeUnavailable, HTTPStatus: resp.StatusCode, Retryable: retryable}
}
limited := io.LimitReader(resp.Body, maxGuardResponseBytes+1)
responseBody, err := io.ReadAll(limited)
if err != nil {
return nil, &GuardError{Code: ErrorCodeUnavailable, Retryable: true, Cause: err}
}
if int64(len(responseBody)) > maxGuardResponseBytes {
return nil, &GuardError{Code: ErrorCodeInvalidResponse}
}
content, err := extractOpenAIContent(responseBody)
if err != nil {
return nil, &GuardError{Code: ErrorCodeInvalidResponse, Cause: err}
}
result, err := ParseQwen3Guard(content, enabledScanners)
if err != nil {
return nil, err
}
result.GuardEndpointID = endpoint.ID
result.ScannerVersion = endpoint.Model
return result, nil
}
func (s *OpenAICompatibleScanner) clientFor(endpoint ActiveEndpoint) (*http.Client, error) {
key := fmt.Sprintf("%s|%s|%d", endpoint.ID, endpoint.BaseURL, endpoint.TimeoutMS)
if cached, ok := s.clients.Load(key); ok {
client, valid := cached.(*http.Client)
if !valid {
s.clients.Delete(key)
return nil, errors.New("prompt guard client cache invalid")
}
return client, nil
}
client, err := NewSecureHTTPClient(endpoint)
if err != nil {
return nil, err
}
actual, _ := s.clients.LoadOrStore(key, client)
actualClient, ok := actual.(*http.Client)
if !ok {
s.clients.Delete(key)
return nil, errors.New("prompt guard client cache invalid")
}
return actualClient, nil
}
func extractOpenAIContent(body []byte) (string, error) {
var response struct {
Choices []struct {
Message struct {
Content any `json:"content"`
} `json:"message"`
} `json:"choices"`
}
if err := json.Unmarshal(body, &response); err != nil || len(response.Choices) == 0 {
return "", errors.New("prompt guard response envelope invalid")
}
content := response.Choices[0].Message.Content
switch typed := content.(type) {
case string:
if strings.TrimSpace(typed) == "" {
return "", errors.New("prompt guard response content empty")
}
return typed, nil
case []any:
parts := make([]string, 0, len(typed))
for _, item := range typed {
object, ok := item.(map[string]any)
if !ok {
continue
}
if text, ok := object["text"].(string); ok && strings.TrimSpace(text) != "" {
parts = append(parts, text)
}
}
if len(parts) == 0 {
return "", errors.New("prompt guard response content empty")
}
return strings.Join(parts, "\n"), nil
default:
return "", errors.New("prompt guard response content invalid")
}
}
func ScannerDefinitions() []ScannerDefinition {
result := make([]ScannerDefinition, 0, len(AllScannerIDs))
for _, id := range AllScannerIDs {
result = append(result, ScannerCatalog[id])
}
sort.SliceStable(result, func(i, j int) bool { return i < j })
return result
}
@@ -0,0 +1,146 @@
package securityaudit
import (
"encoding/json"
"testing"
"time"
"github.com/stretchr/testify/require"
)
func TestParseQwen3GuardStrictAndPolicy(t *testing.T) {
tests := []struct {
name, output string
enabled []string
decision EventDecision
action Action
wantErr bool
}{
{"safe", "Safety: Safe\nCategories: None", AllScannerIDs, EventPass, ActionAllow, false},
{"controversial", "Safety: Controversial\nCategories: Violent", AllScannerIDs, EventFlag, ActionWarn, false},
{"controversial pii escalates", "Safety: Controversial\nCategories: PII", AllScannerIDs, EventCritical, ActionBlock, false},
{"unsafe", "Safety: Unsafe\nCategories: Jailbreak", AllScannerIDs, EventCritical, ActionBlock, false},
{"unknown unsafe", "Safety: Unsafe\nCategories: Future Risk", AllScannerIDs, EventCritical, ActionBlock, false},
{"disabled unsafe warns", "Safety: Unsafe\nCategories: Violent", []string{"PII"}, EventFlag, ActionWarn, false},
{"extra explanation", "Safety: Safe\nCategories: None\nThis is safe", AllScannerIDs, EventPass, ActionAllow, false},
{"duplicate", "Safety: Safe\nSafety: Safe", AllScannerIDs, "", "", true},
{"duplicate categories", "Safety: Safe\nCategories: None\nCategories: PII", AllScannerIDs, "", "", true},
{"missing categories", "Safety: Safe\n", AllScannerIDs, "", "", true},
{"unknown safety", "Safety: Maybe\nCategories: PII", AllScannerIDs, "", "", true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result, err := ParseQwen3Guard(tt.output, tt.enabled)
if tt.wantErr {
require.Error(t, err)
return
}
require.NoError(t, err)
require.Equal(t, tt.decision, result.Decision)
require.Equal(t, tt.action, result.Action)
})
}
}
func TestParseQwen3GuardIgnoresAuxiliaryResponseFields(t *testing.T) {
result, err := ParseQwen3Guard("Safety: Unsafe\nCategories: Jailbreak\nRefusal: No", AllScannerIDs)
require.NoError(t, err)
require.Equal(t, "Unsafe", result.Safety)
require.Equal(t, []string{"jailbreak"}, result.Categories)
serialized, err := json.Marshal(result)
require.NoError(t, err)
require.NotContains(t, string(serialized), "Refusal")
require.NotContains(t, string(serialized), "No")
}
func TestQwen3GuardOfficialCategoriesAliasesAndUnknownAreStable(t *testing.T) {
official := "Violent, Non-violent Illegal Acts, Sexual Content or Sexual Acts, PII, Suicide & Self-Harm, Unethical Acts, Politically Sensitive Topics, Copyright Violation, Jailbreak"
result, err := ParseQwen3Guard("Safety: Unsafe\nCategories: "+official, AllScannerIDs)
require.NoError(t, err)
require.Equal(t, AllScannerIDs, result.MatchedScanners)
require.Empty(t, result.UnknownCategories)
require.Equal(t, "priority", result.PolicyID)
require.Equal(t, 1, result.PolicyVersion)
aliases := map[string]string{
"violence": "violent", "non_violent_illegal_acts": "non_violent_illegal_acts",
"sexual": "sexual_content_or_sexual_acts", "personal identifiable information": "pii",
"suicide/self harm": "suicide_and_self_harm", "unethical": "unethical_acts",
"political": "politically_sensitive_topics", "copyright": "copyright_violation",
"prompt injection": "jailbreak",
}
for alias, canonical := range aliases {
require.Equal(t, canonical, NormalizeCategory(alias), alias)
}
const canary = "PROMPT_CANARY_RAW_UNKNOWN_CATEGORY"
unknown, err := ParseQwen3Guard("Safety: Unsafe\nCategories: "+canary, AllScannerIDs)
require.NoError(t, err)
require.Len(t, unknown.UnknownCategories, 1)
require.NotContains(t, unknown.UnknownCategories[0], "canary")
require.NotContains(t, unknown.UnknownCategories[0], "raw")
require.Contains(t, unknown.UnknownCategories[0], "unknown:")
}
func TestExtractOpenAIContentSupportsStringAndTextBlocks(t *testing.T) {
content, err := extractOpenAIContent([]byte(`{"choices":[{"message":{"content":"Safety: Safe\nCategories: None"}}]}`))
require.NoError(t, err)
require.Equal(t, "Safety: Safe\nCategories: None", content)
content, err = extractOpenAIContent([]byte(`{"choices":[{"message":{"content":[{"type":"text","text":"Safety: Safe"},{"type":"text","text":"Categories: None"}]}}]}`))
require.NoError(t, err)
require.Equal(t, "Safety: Safe\nCategories: None", content)
for _, body := range []string{`{}`, `{"choices":[]}`, `{"choices":[{"message":{"content":null}}]}`} {
_, err := extractOpenAIContent([]byte(body))
require.Error(t, err)
}
}
func TestAggregateRequiresEveryResult(t *testing.T) {
_, err := AggregateResults([]*NormalizedResult{{Decision: EventPass, Action: ActionAllow}, nil}, 0)
require.Error(t, err)
result, err := AggregateResults([]*NormalizedResult{
{Decision: EventPass, RiskLevel: RiskLow, Action: ActionAllow, Categories: []string{"pii"}},
{Decision: EventCritical, RiskLevel: RiskCritical, Action: ActionBlock, Categories: []string{"jailbreak"}},
}, 0)
require.NoError(t, err)
require.Equal(t, EventCritical, result.Decision)
require.Equal(t, ActionBlock, result.Action)
require.Equal(t, []string{"pii", "jailbreak"}, result.Categories)
}
func TestAggregateDeduplicatesFactsAndUsesMostSevereEndpointMetadata(t *testing.T) {
result, err := AggregateResults([]*NormalizedResult{
{Decision: EventPass, RiskLevel: RiskLow, Action: ActionAllow, Safety: "Safe", Categories: []string{"pii"}, MatchedScanners: []string{"pii"}, ScannerScores: map[string]float64{"pii": 0}, ScannerEvidence: map[string]string{"pii": "first"}, GuardEndpointID: "safe-node", ScannerVersion: "safe-version", PolicyID: "priority", PolicyVersion: 1},
{Decision: EventCritical, RiskLevel: RiskCritical, Action: ActionBlock, Safety: "Unsafe", Categories: []string{"pii", "jailbreak"}, MatchedScanners: []string{"pii", "jailbreak"}, ScannerScores: map[string]float64{"pii": 1, "jailbreak": 1}, ScannerEvidence: map[string]string{"pii": "second", "jailbreak": "blocked"}, GuardEndpointID: "block-node", ScannerVersion: "block-version", PolicyID: "priority", PolicyVersion: 2},
}, 7*time.Millisecond)
require.NoError(t, err)
require.Equal(t, []string{"pii", "jailbreak"}, result.Categories)
require.Equal(t, []string{"pii", "jailbreak"}, result.MatchedScanners)
require.Equal(t, "first", result.ScannerEvidence["pii"], "evidence is deterministically first-seen")
require.Equal(t, "block-node", result.GuardEndpointID)
require.Equal(t, "block-version", result.ScannerVersion)
require.Equal(t, 2, result.PolicyVersion)
require.Equal(t, 7, result.LatencyMS)
}
func TestIssueSummariesAreDeterministicRedactedDerivedDTOs(t *testing.T) {
const canary = "PROMPT_CANARY_EVIDENCE_SECRET"
result := NormalizedResult{
Decision: EventCritical, RiskLevel: RiskCritical, Action: ActionBlock,
Categories: []string{"jailbreak", "pii"}, MatchedScanners: []string{"pii"},
ScannerScores: map[string]float64{"pii": 1}, ScannerEvidence: map[string]string{"pii": canary},
UnknownCategories: []string{unknownCategoryID("future risk")},
}
summaries := BuildIssueSummaries(result)
require.Len(t, summaries, 3, "known categories are not hidden merely because policy disabled one")
raw, err := json.Marshal(summaries)
require.NoError(t, err)
require.NotContains(t, string(raw), canary)
for _, summary := range summaries {
require.NotEmpty(t, summary.Title)
require.NotEmpty(t, summary.Description)
require.NotEmpty(t, summary.Code)
require.NotEmpty(t, summary.EvidenceHash)
}
}
@@ -0,0 +1,450 @@
package securityaudit
import (
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"strings"
"time"
)
const (
promptAuditAdmissionLockKey int64 = 579147893221901921
promptAuditConfigLockKey int64 = 579147893221901922
)
var (
ErrQueueFull = errors.New("prompt audit queue full")
ErrQueueAdmissionBusy = errors.New("prompt audit queue admission busy")
ErrLeaseLost = errors.New("prompt audit worker lease lost")
ErrEventNotFound = errors.New("prompt audit event not found")
)
type Job struct {
ID int64
Snapshot PromptSnapshot
ExecutionMode Mode
ConfigVersion int64
Status string
Attempts int
MaxAttempts int
ClaimVersion int64
NextAttemptAt time.Time
ProcessingStartedAt *time.Time
ProcessedAt *time.Time
LastErrorCode string
LastErrorMessage string
CreatedAt time.Time
UpdatedAt time.Time
}
type Event struct {
ID int64 `json:"id"`
JobID int64 `json:"job_id"`
Snapshot PromptSnapshot `json:"snapshot"`
Decision EventDecision `json:"decision"`
RiskLevel RiskLevel `json:"risk_level"`
Action Action `json:"action"`
Categories []string `json:"categories"`
MatchedScanners []string `json:"matched_scanners"`
ScannerScores map[string]float64 `json:"scanner_scores"`
ScannerEvidence map[string]string `json:"scanner_evidence"`
ScannerBackend string `json:"scanner_backend"`
ScannerVersion string `json:"scanner_version"`
GuardEndpointID string `json:"guard_endpoint_id"`
PolicyID string `json:"policy_id"`
PolicyVersion int `json:"policy_version"`
ConfigVersion int64 `json:"config_version"`
ChunkTotal int `json:"chunk_total"`
LatencyMS int `json:"latency_ms"`
IssueSummaries []IssueSummary `json:"issue_summaries"`
CreatedAt time.Time `json:"created_at"`
}
type JobRepository interface {
CreateStagingWithCapacity(ctx context.Context, snapshot PromptSnapshot, configVersion int64, maxAttempts, capacity int) (*Job, error)
PublishQueued(ctx context.Context, jobID int64) error
MarkStagingFailed(ctx context.Context, jobID int64, code, message string) error
ClaimNextJob(ctx context.Context, now time.Time) (*Job, bool, error)
RefreshLease(ctx context.Context, jobID, claimVersion int64, now time.Time) error
Complete(ctx context.Context, job *Job, result *NormalizedResult, storePassEvents bool) (*Event, error)
Retry(ctx context.Context, jobID, claimVersion int64, next time.Time, code, message string) error
Fail(ctx context.Context, jobID, claimVersion int64, code, message string) error
ReclaimStale(ctx context.Context, stagingBefore, processingBefore time.Time, limit int) (int64, error)
QueueStats(ctx context.Context) (QueueStats, error)
RecordBlocking(ctx context.Context, snapshot PromptSnapshot, configVersion int64, result *NormalizedResult, storePassEvents bool) (*Event, error)
}
type PostgreSQLRepository struct {
db *sql.DB
clock Clock
}
func NewPostgreSQLRepository(db *sql.DB) *PostgreSQLRepository {
return &PostgreSQLRepository{db: db, clock: realClock{}}
}
func (r *PostgreSQLRepository) CreateStagingWithCapacity(ctx context.Context, snapshot PromptSnapshot, configVersion int64, maxAttempts, capacity int) (*Job, error) {
if r == nil || r.db == nil {
return nil, errors.New("prompt audit database unavailable")
}
tx, err := r.db.BeginTx(ctx, &sql.TxOptions{Isolation: sql.LevelReadCommitted})
if err != nil {
return nil, err
}
defer func() { _ = tx.Rollback() }()
var locked bool
if err := tx.QueryRowContext(ctx, `SELECT pg_try_advisory_xact_lock($1)`, promptAuditAdmissionLockKey).Scan(&locked); err != nil {
return nil, err
}
if !locked {
return nil, ErrQueueAdmissionBusy
}
var active int
if err := tx.QueryRowContext(ctx, `
SELECT COUNT(*) FROM prompt_audit_jobs
WHERE status IN ('staging','queued','processing','retry')`).Scan(&active); err != nil {
return nil, err
}
if capacity <= 0 || active >= capacity {
return nil, ErrQueueFull
}
if maxAttempts <= 0 {
maxAttempts = 3
}
job, err := insertJob(ctx, tx, snapshot.Redacted(), ModeAsync, configVersion, "staging", maxAttempts)
if err != nil {
return nil, err
}
if err := tx.Commit(); err != nil {
return nil, err
}
return job, nil
}
func (r *PostgreSQLRepository) PublishQueued(ctx context.Context, jobID int64) error {
result, err := r.db.ExecContext(ctx, `
UPDATE prompt_audit_jobs SET status='queued', next_attempt_at=NOW(), updated_at=NOW()
WHERE id=$1 AND status='staging'`, jobID)
return requireOneRow(result, err, ErrLeaseLost)
}
func (r *PostgreSQLRepository) MarkStagingFailed(ctx context.Context, jobID int64, code, _ string) error {
code, message := sanitizeStoredError(code)
result, err := r.db.ExecContext(ctx, `
UPDATE prompt_audit_jobs
SET status='failed', processed_at=NOW(), updated_at=NOW(), last_error_code=$2, last_error_message=$3
WHERE id=$1 AND status='staging'`, jobID, code, message)
return requireOneRow(result, err, ErrLeaseLost)
}
func (r *PostgreSQLRepository) ClaimNextJob(ctx context.Context, now time.Time) (*Job, bool, error) {
row := r.db.QueryRowContext(ctx, `
WITH candidate AS (
SELECT id FROM prompt_audit_jobs
WHERE status IN ('queued','retry') AND next_attempt_at <= $1
ORDER BY next_attempt_at, id
FOR UPDATE SKIP LOCKED
LIMIT 1
)
UPDATE prompt_audit_jobs AS j
SET status='processing', attempts=j.attempts+1, claim_version=j.claim_version+1,
processing_started_at=$1, updated_at=$1
FROM candidate
WHERE j.id=candidate.id
RETURNING `+jobColumns("j"), now.UTC())
job, err := scanJob(row)
if errors.Is(err, sql.ErrNoRows) {
return nil, false, nil
}
return job, err == nil, err
}
func (r *PostgreSQLRepository) RefreshLease(ctx context.Context, jobID, claimVersion int64, now time.Time) error {
result, err := r.db.ExecContext(ctx, `
UPDATE prompt_audit_jobs SET processing_started_at=$3, updated_at=$3
WHERE id=$1 AND status='processing' AND claim_version=$2`, jobID, claimVersion, now.UTC())
return requireOneRow(result, err, ErrLeaseLost)
}
func (r *PostgreSQLRepository) Complete(ctx context.Context, job *Job, result *NormalizedResult, storePassEvents bool) (*Event, error) {
if job == nil || result == nil {
return nil, errors.New("prompt audit completion requires job and result")
}
tx, err := r.db.BeginTx(ctx, nil)
if err != nil {
return nil, err
}
defer func() { _ = tx.Rollback() }()
updateResult, err := tx.ExecContext(ctx, `
UPDATE prompt_audit_jobs SET status='done', processed_at=NOW(), updated_at=NOW(),
last_error_code='', last_error_message=''
WHERE id=$1 AND status='processing' AND claim_version=$2`, job.ID, job.ClaimVersion)
if err := requireOneRow(updateResult, err, ErrLeaseLost); err != nil {
return nil, err
}
var event *Event
if shouldStorePromptAuditEvent(result.Decision, storePassEvents) {
event, err = insertEvent(ctx, tx, job.ID, job.Snapshot.Redacted(), job.ConfigVersion, result)
if err != nil {
return nil, err
}
}
if err := tx.Commit(); err != nil {
return nil, err
}
return event, nil
}
func (r *PostgreSQLRepository) Retry(ctx context.Context, jobID, claimVersion int64, next time.Time, code, _ string) error {
code, message := sanitizeStoredError(code)
result, err := r.db.ExecContext(ctx, `
UPDATE prompt_audit_jobs SET status='retry', next_attempt_at=$3, processing_started_at=NULL,
updated_at=NOW(), last_error_code=$4, last_error_message=$5
WHERE id=$1 AND status='processing' AND claim_version=$2`,
jobID, claimVersion, next.UTC(), code, message)
return requireOneRow(result, err, ErrLeaseLost)
}
func (r *PostgreSQLRepository) Fail(ctx context.Context, jobID, claimVersion int64, code, _ string) error {
code, message := sanitizeStoredError(code)
result, err := r.db.ExecContext(ctx, `
UPDATE prompt_audit_jobs SET status='failed', processed_at=NOW(), processing_started_at=NULL,
updated_at=NOW(), last_error_code=$3, last_error_message=$4
WHERE id=$1 AND status='processing' AND claim_version=$2`,
jobID, claimVersion, code, message)
return requireOneRow(result, err, ErrLeaseLost)
}
func (r *PostgreSQLRepository) ReclaimStale(ctx context.Context, stagingBefore, processingBefore time.Time, limit int) (int64, error) {
if limit <= 0 || limit > 1000 {
limit = 100
}
result, err := r.db.ExecContext(ctx, `
WITH stale AS (
SELECT id FROM prompt_audit_jobs
WHERE (status='staging' AND updated_at < $1)
OR (status='processing' AND processing_started_at < $2)
ORDER BY updated_at, id FOR UPDATE SKIP LOCKED LIMIT $3
)
UPDATE prompt_audit_jobs AS j
SET status=CASE
WHEN j.status='staging' THEN 'failed'
WHEN j.attempts < j.max_attempts THEN 'retry'
ELSE 'failed' END,
next_attempt_at=CASE WHEN j.status='processing' AND j.attempts < j.max_attempts THEN NOW() ELSE j.next_attempt_at END,
processing_started_at=NULL,
processed_at=CASE WHEN j.status='staging' OR j.attempts >= j.max_attempts THEN NOW() ELSE NULL END,
last_error_code=CASE WHEN j.status='staging' THEN 'staging_timeout' ELSE 'processing_lease_expired' END,
last_error_message='', updated_at=NOW()
FROM stale WHERE j.id=stale.id`, stagingBefore.UTC(), processingBefore.UTC(), limit)
if err != nil {
return 0, err
}
return result.RowsAffected()
}
func (r *PostgreSQLRepository) QueueStats(ctx context.Context) (QueueStats, error) {
rows, err := r.db.QueryContext(ctx, `SELECT status, COUNT(*) FROM prompt_audit_jobs GROUP BY status`)
if err != nil {
return QueueStats{}, err
}
defer func() { _ = rows.Close() }()
var stats QueueStats
for rows.Next() {
var status string
var count int64
if err := rows.Scan(&status, &count); err != nil {
return QueueStats{}, err
}
switch status {
case "staging":
stats.Staging = count
case "queued":
stats.Queued = count
case "processing":
stats.Processing = count
case "retry":
stats.Retry = count
case "done":
stats.Done = count
case "failed":
stats.Failed = count
}
}
stats.Active = stats.Staging + stats.Queued + stats.Processing + stats.Retry
return stats, rows.Err()
}
func (r *PostgreSQLRepository) RecordBlocking(ctx context.Context, snapshot PromptSnapshot, configVersion int64, result *NormalizedResult, storePassEvents bool) (*Event, error) {
if result == nil {
return nil, errors.New("prompt guard result required")
}
tx, err := r.db.BeginTx(ctx, nil)
if err != nil {
return nil, err
}
defer func() { _ = tx.Rollback() }()
job, err := insertJob(ctx, tx, snapshot.Redacted(), ModeBlocking, configVersion, "done", 1)
if err != nil {
return nil, err
}
var event *Event
if shouldStorePromptAuditEvent(result.Decision, storePassEvents) {
event, err = insertEvent(ctx, tx, job.ID, snapshot.Redacted(), configVersion, result)
if err != nil {
return nil, err
}
}
if err := tx.Commit(); err != nil {
return nil, err
}
return event, nil
}
// shouldStorePromptAuditEvent keeps store_pass_events scoped to safe results.
// Risk events are always persisted while prompt auditing itself is enabled.
func shouldStorePromptAuditEvent(decision EventDecision, storePassEvents bool) bool {
return decision != EventPass || storePassEvents
}
type sqlQueryer interface {
QueryRowContext(context.Context, string, ...any) *sql.Row
}
func insertJob(ctx context.Context, queryer sqlQueryer, snapshot PromptSnapshot, mode Mode, configVersion int64, status string, maxAttempts int) (*Job, error) {
processedExpr := "NULL"
if status == "done" || status == "failed" {
processedExpr = "NOW()"
}
row := queryer.QueryRowContext(ctx, `
INSERT INTO prompt_audit_jobs (
request_id,user_id,username_snapshot,user_email_snapshot,api_key_id,api_key_name_snapshot,
group_id,group_name,provider,endpoint,protocol,model,prompt_hash,redacted_preview,
prompt_length,message_count,stage,execution_mode,config_version,status,max_attempts,processed_at
) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,`+processedExpr+`)
RETURNING `+jobColumns("prompt_audit_jobs"),
snapshot.RequestID, nullableID(snapshot.UserID), snapshot.UsernameSnapshot, snapshot.UserEmailSnapshot,
nullableID(snapshot.APIKeyID), snapshot.APIKeyNameSnapshot, snapshot.GroupID, snapshot.GroupName,
snapshot.Provider, snapshot.Endpoint, snapshot.Protocol, snapshot.Model, snapshot.PromptHash,
snapshot.RedactedPreview, snapshot.PromptLength, snapshot.MessageCount, normalizeStage(snapshot.Stage),
string(mode), configVersion, status, maxAttempts)
return scanJob(row)
}
func insertEvent(ctx context.Context, queryer sqlQueryer, jobID int64, snapshot PromptSnapshot, configVersion int64, result *NormalizedResult) (*Event, error) {
categories, _ := json.Marshal(result.Categories)
matched, _ := json.Marshal(result.MatchedScanners)
scores, _ := json.Marshal(result.ScannerScores)
evidence := make(map[string]string, len(result.ScannerEvidence))
for key, value := range result.ScannerEvidence {
evidence[key] = RedactPreview(value, 160)
}
evidenceJSON, _ := json.Marshal(evidence)
row := queryer.QueryRowContext(ctx, `
INSERT INTO prompt_audit_events (
job_id,request_id,user_id,username_snapshot,user_email_snapshot,api_key_id,api_key_name_snapshot,
group_id,group_name,provider,endpoint,protocol,model,prompt_hash,redacted_preview,stage,
decision,risk_level,action,categories,matched_scanners,scanner_scores,scanner_evidence,
scanner_backend,scanner_version,guard_endpoint_id,policy_id,policy_version,config_version,chunk_total,latency_ms,
full_prompt
) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,
$20::jsonb,$21::jsonb,$22::jsonb,$23::jsonb,$24,$25,$26,$27,$28,$29,$30,$31,$32)
RETURNING `+eventDetailColumns("prompt_audit_events"),
jobID, snapshot.RequestID, nullableID(snapshot.UserID), snapshot.UsernameSnapshot, snapshot.UserEmailSnapshot,
nullableID(snapshot.APIKeyID), snapshot.APIKeyNameSnapshot, snapshot.GroupID, snapshot.GroupName,
snapshot.Provider, snapshot.Endpoint, snapshot.Protocol, snapshot.Model, snapshot.PromptHash,
snapshot.RedactedPreview, normalizeStage(snapshot.Stage), string(result.Decision), string(result.RiskLevel),
string(result.Action), categories, matched, scores, evidenceJSON, result.ScannerBackend, result.ScannerVersion,
result.GuardEndpointID, result.PolicyID, result.PolicyVersion, configVersion, result.ChunkTotal, result.LatencyMS,
snapshot.FullPrompt)
return scanEvent(row, true)
}
type rowScanner interface{ Scan(...any) error }
func scanJob(row rowScanner) (*Job, error) {
job := &Job{}
var userID, apiKeyID, groupID sql.NullInt64
var processingStarted, processed sql.NullTime
err := row.Scan(
&job.ID, &job.Snapshot.RequestID, &userID, &job.Snapshot.UsernameSnapshot, &job.Snapshot.UserEmailSnapshot,
&apiKeyID, &job.Snapshot.APIKeyNameSnapshot, &groupID, &job.Snapshot.GroupName, &job.Snapshot.Provider,
&job.Snapshot.Endpoint, &job.Snapshot.Protocol, &job.Snapshot.Model, &job.Snapshot.PromptHash,
&job.Snapshot.RedactedPreview, &job.Snapshot.PromptLength, &job.Snapshot.MessageCount, &job.Snapshot.Stage,
&job.ExecutionMode, &job.ConfigVersion, &job.Status, &job.Attempts, &job.MaxAttempts, &job.ClaimVersion,
&job.NextAttemptAt, &processingStarted, &processed, &job.LastErrorCode, &job.LastErrorMessage,
&job.CreatedAt, &job.UpdatedAt,
)
if err != nil {
return nil, err
}
job.Snapshot.UserID = nullableInt64Value(userID)
job.Snapshot.APIKeyID = nullableInt64Value(apiKeyID)
job.Snapshot.GroupID = nullableInt64Ptr(groupID)
if processingStarted.Valid {
value := processingStarted.Time
job.ProcessingStartedAt = &value
}
if processed.Valid {
value := processed.Time
job.ProcessedAt = &value
}
return job, nil
}
func jobColumns(alias string) string {
return fmt.Sprintf(`%[1]s.id,%[1]s.request_id,%[1]s.user_id,%[1]s.username_snapshot,%[1]s.user_email_snapshot,
%[1]s.api_key_id,%[1]s.api_key_name_snapshot,%[1]s.group_id,%[1]s.group_name,%[1]s.provider,
%[1]s.endpoint,%[1]s.protocol,%[1]s.model,%[1]s.prompt_hash,%[1]s.redacted_preview,
%[1]s.prompt_length,%[1]s.message_count,%[1]s.stage,%[1]s.execution_mode,%[1]s.config_version,%[1]s.status,
%[1]s.attempts,%[1]s.max_attempts,%[1]s.claim_version,%[1]s.next_attempt_at,
%[1]s.processing_started_at,%[1]s.processed_at,%[1]s.last_error_code,%[1]s.last_error_message,
%[1]s.created_at,%[1]s.updated_at`, alias)
}
func normalizeStage(stage string) string {
stage = strings.TrimSpace(stage)
if stage == "" {
return "http"
}
return stage
}
func requireOneRow(result sql.Result, err error, missing error) error {
if err != nil {
return err
}
rows, err := result.RowsAffected()
if err != nil {
return err
}
if rows != 1 {
return missing
}
return nil
}
func nullableID(value int64) any {
if value <= 0 {
return nil
}
return value
}
func nullableInt64Value(value sql.NullInt64) int64 {
if !value.Valid {
return 0
}
return value.Int64
}
func nullableInt64Ptr(value sql.NullInt64) *int64 {
if !value.Valid {
return nil
}
result := value.Int64
return &result
}
@@ -0,0 +1,462 @@
package securityaudit
import (
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"os"
"path/filepath"
"sort"
"strings"
"sync"
"testing"
"time"
_ "github.com/lib/pq"
"github.com/stretchr/testify/require"
)
const promptAuditPostgresTestEnv = "PROMPT_AUDIT_TEST_POSTGRES_DSN"
func openPromptAuditIntegrationDB(t *testing.T) *sql.DB {
t.Helper()
dsn := strings.TrimSpace(os.Getenv(promptAuditPostgresTestEnv))
if dsn == "" {
t.Skip(promptAuditPostgresTestEnv + " is not set")
}
db, err := sql.Open("postgres", dsn)
require.NoError(t, err)
db.SetMaxOpenConns(16)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
require.NoError(t, db.PingContext(ctx))
_, err = db.ExecContext(ctx, `
CREATE TABLE IF NOT EXISTS users (id BIGSERIAL PRIMARY KEY);
CREATE TABLE IF NOT EXISTS groups (id BIGSERIAL PRIMARY KEY);
CREATE TABLE IF NOT EXISTS api_keys (id BIGSERIAL PRIMARY KEY);
CREATE TABLE IF NOT EXISTS settings (
key VARCHAR(255) PRIMARY KEY,
value TEXT NOT NULL DEFAULT '',
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
`)
require.NoError(t, err)
for _, name := range []string{"181_prompt_audit.sql", "182_prompt_audit_full_prompt.sql"} {
migration, err := os.ReadFile(filepath.Join("..", "..", "migrations", name))
require.NoError(t, err)
// The migration runner can retry an interrupted deployment; the migration
// must therefore be safe to execute more than once.
_, err = db.ExecContext(ctx, string(migration))
require.NoError(t, err)
_, err = db.ExecContext(ctx, string(migration))
require.NoError(t, err)
}
t.Cleanup(func() { require.NoError(t, db.Close()) })
resetPromptAuditIntegrationDB(t, db)
return db
}
func resetPromptAuditIntegrationDB(t *testing.T, db *sql.DB) {
t.Helper()
_, err := db.Exec(`TRUNCATE TABLE prompt_audit_events, prompt_audit_jobs, api_keys, users, groups, settings RESTART IDENTITY CASCADE`)
require.NoError(t, err)
}
func insertIdentity(t *testing.T, db *sql.DB, table string) int64 {
t.Helper()
var id int64
require.NoError(t, db.QueryRow(`INSERT INTO `+table+` DEFAULT VALUES RETURNING id`).Scan(&id))
return id
}
func integrationSnapshot(seed string) PromptSnapshot {
return PromptSnapshot{
RequestID: "request-" + seed, UsernameSnapshot: "user-" + seed,
UserEmailSnapshot: "user-" + seed + "@example.test", APIKeyNameSnapshot: "key-" + seed,
GroupName: "group-" + seed, Provider: "openai", Endpoint: "/v1/chat/completions",
Protocol: "openai_chat", Model: "gpt-test", PromptHash: strings.Repeat(seed[:1], 64),
RedactedPreview: "redacted-" + seed, PromptLength: len([]rune(seed)), MessageCount: 1,
}
}
func integrationResult(decision EventDecision) *NormalizedResult {
result := &NormalizedResult{
Decision: decision, RiskLevel: RiskLow, Action: ActionAllow, Safety: "Safe",
Categories: []string{}, MatchedScanners: []string{}, ScannerScores: map[string]float64{},
ScannerEvidence: map[string]string{}, ScannerBackend: "qwen3guard-openai",
ScannerVersion: "test", GuardEndpointID: "guard-1", PolicyID: "priority",
PolicyVersion: 1, ChunkTotal: 1, LatencyMS: 2,
}
if decision != EventPass {
result.RiskLevel = RiskCritical
result.Action = ActionBlock
result.Safety = "Unsafe"
result.Categories = []string{"pii"}
result.MatchedScanners = []string{"pii"}
result.ScannerScores["pii"] = 1
result.ScannerEvidence["pii"] = "redacted evidence"
}
return result
}
func TestPromptAuditMigrationSchemaAndLeakageGate(t *testing.T) {
db := openPromptAuditIntegrationDB(t)
ctx := context.Background()
rows, err := db.QueryContext(ctx, `SELECT table_name, column_name FROM information_schema.columns
WHERE table_schema='public' AND table_name IN ('prompt_audit_jobs','prompt_audit_events')`)
require.NoError(t, err)
defer func() { _ = rows.Close() }()
forbidden := []string{"raw_prompt", "raw_request", "payload", "token", "authorization", "credential", "ciphertext"}
for rows.Next() {
var tableName, columnName string
require.NoError(t, rows.Scan(&tableName, &columnName))
lower := strings.ToLower(columnName)
for _, word := range forbidden {
require.NotContainsf(t, lower, word, "%s.%s is a forbidden raw/credential column", tableName, columnName)
}
}
require.NoError(t, rows.Err())
indexRows, err := db.QueryContext(ctx, `SELECT indexname FROM pg_indexes
WHERE schemaname='public' AND tablename IN ('prompt_audit_jobs','prompt_audit_events')`)
require.NoError(t, err)
defer func() { _ = indexRows.Close() }()
indexes := map[string]bool{}
for indexRows.Next() {
var name string
require.NoError(t, indexRows.Scan(&name))
indexes[name] = true
}
for _, name := range []string{
"idx_prompt_audit_jobs_schedule", "idx_prompt_audit_jobs_request", "idx_prompt_audit_jobs_user_created",
"idx_prompt_audit_jobs_api_key_created", "idx_prompt_audit_jobs_group_created", "idx_prompt_audit_jobs_prompt_hash",
"idx_prompt_audit_jobs_created", "idx_prompt_audit_events_job", "idx_prompt_audit_events_request",
"idx_prompt_audit_events_decision_created", "idx_prompt_audit_events_risk_created",
"idx_prompt_audit_events_user_created", "idx_prompt_audit_events_api_key_created",
"idx_prompt_audit_events_group_created", "idx_prompt_audit_events_prompt_hash", "idx_prompt_audit_events_created",
} {
require.Truef(t, indexes[name], "missing index %s", name)
}
_, err = db.ExecContext(ctx, `INSERT INTO prompt_audit_jobs(status) VALUES ('unknown')`)
require.Error(t, err)
_, err = db.ExecContext(ctx, `INSERT INTO prompt_audit_jobs(prompt_length) VALUES (-1)`)
require.Error(t, err)
var jobID int64
require.NoError(t, db.QueryRowContext(ctx, `INSERT INTO prompt_audit_jobs DEFAULT VALUES RETURNING id`).Scan(&jobID))
_, err = db.ExecContext(ctx, `INSERT INTO prompt_audit_events(job_id,chunk_total) VALUES ($1,-1)`, jobID)
require.Error(t, err)
}
func TestPromptAuditDatabasePersistsFullPromptOnEventsOnly(t *testing.T) {
db := openPromptAuditIntegrationDB(t)
repo := NewPostgreSQLRepository(db)
ctx := context.Background()
const promptCanary = "PROMPT_AUDIT_CANARY_SECRET_DO_NOT_PERSIST"
request := Request{
RequestID: "canary-request", Provider: "openai",
Endpoint: "/v1/chat/completions", Protocol: "openai_chat", Model: "gpt-test", Stage: "http",
Body: []byte(`{"messages":[{"role":"user","content":"` + promptCanary + `"}]}`),
}
snapshot, err := ExtractPromptSnapshot(request)
require.NoError(t, err)
require.NotContains(t, snapshot.RedactedPreview, promptCanary)
require.Contains(t, snapshot.FullPrompt, promptCanary)
event, err := repo.RecordBlocking(ctx, snapshot.Redacted(), 1, integrationResult(EventCritical), true)
require.NoError(t, err)
// The event intentionally retains the full prompt for admin review; the
// redacted preview and transient job row still never contain it.
adminJSON, err := json.Marshal(event)
require.NoError(t, err)
require.Contains(t, string(adminJSON), promptCanary)
require.NotContains(t, event.Snapshot.RedactedPreview, promptCanary)
var storedFullPrompt string
require.NoError(t, db.QueryRow(`SELECT full_prompt FROM prompt_audit_events WHERE id=$1`, event.ID).Scan(&storedFullPrompt))
require.Contains(t, storedFullPrompt, promptCanary)
detail, err := repo.GetEvent(ctx, event.ID)
require.NoError(t, err)
require.Contains(t, detail.Snapshot.FullPrompt, promptCanary)
var jobJSON string
require.NoError(t, db.QueryRow(`SELECT row_to_json(j)::text FROM prompt_audit_jobs j WHERE id=$1`, event.JobID).Scan(&jobJSON))
require.NotContains(t, jobJSON, promptCanary)
failedJob, err := repo.CreateStagingWithCapacity(ctx, integrationSnapshot("error"), 1, 3, 10)
require.NoError(t, err)
const errorCanary = "GUARD_RAW_RESPONSE_CANARY_SECRET"
require.NoError(t, repo.MarkStagingFailed(ctx, failedJob.ID, "payload_store_failed", "raw guard body: "+errorCanary))
var code, message string
require.NoError(t, db.QueryRow(`SELECT last_error_code,last_error_message FROM prompt_audit_jobs WHERE id=$1`, failedJob.ID).Scan(&code, &message))
require.Equal(t, "payload_store_failed", code)
require.Equal(t, stableErrorMessage(code), message)
require.NotContains(t, message, errorCanary)
require.LessOrEqual(t, len([]rune(message)), 160)
}
func TestPromptAuditRepositoryAdmissionClaimFencingAndEventTransaction(t *testing.T) {
db := openPromptAuditIntegrationDB(t)
repo := NewPostgreSQLRepository(db)
ctx := context.Background()
start := make(chan struct{})
type admissionResult struct {
job *Job
err error
}
results := make(chan admissionResult, 2)
var wg sync.WaitGroup
for i := 0; i < 2; i++ {
wg.Add(1)
go func(index int) {
defer wg.Done()
<-start
job, err := repo.CreateStagingWithCapacity(ctx, integrationSnapshot(string(rune('a'+index))), 1, 3, 1)
results <- admissionResult{job: job, err: err}
}(i)
}
close(start)
wg.Wait()
close(results)
var accepted *Job
rejected := 0
for result := range results {
if result.err == nil {
require.Nil(t, accepted)
accepted = result.job
continue
}
require.True(t, errors.Is(result.err, ErrQueueFull) || errors.Is(result.err, ErrQueueAdmissionBusy))
rejected++
}
require.NotNil(t, accepted)
require.Equal(t, 1, rejected)
stats, err := repo.QueueStats(ctx)
require.NoError(t, err)
require.Equal(t, int64(1), stats.Active)
require.NoError(t, repo.PublishQueued(ctx, accepted.ID))
claimStart := make(chan struct{})
claims := make(chan *Job, 2)
for i := 0; i < 2; i++ {
wg.Add(1)
go func() {
defer wg.Done()
<-claimStart
job, claimed, claimErr := repo.ClaimNextJob(ctx, time.Now().Add(time.Second))
require.NoError(t, claimErr)
if claimed {
claims <- job
}
}()
}
close(claimStart)
wg.Wait()
close(claims)
claimedJobs := make([]*Job, 0, 1)
for job := range claims {
claimedJobs = append(claimedJobs, job)
}
require.Len(t, claimedJobs, 1)
firstClaim := claimedJobs[0]
require.Equal(t, int64(1), firstClaim.ClaimVersion)
reclaimed, err := repo.ReclaimStale(ctx, time.Now().Add(time.Hour), time.Now().Add(time.Hour), 10)
require.NoError(t, err)
require.Equal(t, int64(1), reclaimed)
secondClaim, claimed, err := repo.ClaimNextJob(ctx, time.Now().Add(time.Second))
require.NoError(t, err)
require.True(t, claimed)
require.Greater(t, secondClaim.ClaimVersion, firstClaim.ClaimVersion)
require.ErrorIs(t, repo.RefreshLease(ctx, firstClaim.ID, firstClaim.ClaimVersion, time.Now()), ErrLeaseLost)
_, err = repo.Complete(ctx, firstClaim, integrationResult(EventCritical), true)
require.ErrorIs(t, err, ErrLeaseLost)
event, err := repo.Complete(ctx, secondClaim, integrationResult(EventCritical), true)
require.NoError(t, err)
require.NotNil(t, event)
var status string
var eventCount int
require.NoError(t, db.QueryRow(`SELECT status FROM prompt_audit_jobs WHERE id=$1`, secondClaim.ID).Scan(&status))
require.Equal(t, "done", status)
require.NoError(t, db.QueryRow(`SELECT COUNT(*) FROM prompt_audit_events WHERE job_id=$1`, secondClaim.ID).Scan(&eventCount))
require.Equal(t, 1, eventCount)
staging, err := repo.CreateStagingWithCapacity(ctx, integrationSnapshot("stale"), 1, 3, 10)
require.NoError(t, err)
reclaimed, err = repo.ReclaimStale(ctx, time.Now().Add(time.Hour), time.Now().Add(time.Hour), 10)
require.NoError(t, err)
require.Equal(t, int64(1), reclaimed)
require.NoError(t, db.QueryRow(`SELECT status FROM prompt_audit_jobs WHERE id=$1`, staging.ID).Scan(&status))
require.Equal(t, "failed", status)
}
func TestPromptAuditRepositoryForeignKeysFiltersAndStableIdentitySnapshots(t *testing.T) {
db := openPromptAuditIntegrationDB(t)
repo := NewPostgreSQLRepository(db)
ctx := context.Background()
userID := insertIdentity(t, db, "users")
apiKeyID := insertIdentity(t, db, "api_keys")
groupID := insertIdentity(t, db, "groups")
snapshot := integrationSnapshot("identity")
snapshot.UserID, snapshot.APIKeyID, snapshot.GroupID = userID, apiKeyID, &groupID
event, err := repo.RecordBlocking(ctx, snapshot, 7, integrationResult(EventCritical), true)
require.NoError(t, err)
require.NotNil(t, event)
start, end := time.Now().Add(-time.Hour), time.Now().Add(time.Hour)
page, err := repo.ListEvents(ctx, EventFilter{
Decision: string(EventCritical), RiskLevel: string(RiskCritical), Endpoint: snapshot.Endpoint,
GroupID: &groupID, UserID: &userID, APIKeyID: &apiKeyID, RequestID: snapshot.RequestID,
PromptHash: snapshot.PromptHash, Keyword: snapshot.UsernameSnapshot, StartAt: &start, EndAt: &end,
}, 1, 10)
require.NoError(t, err)
require.Equal(t, int64(1), page.Total)
require.Len(t, page.Items, 1)
require.NotEmpty(t, page.Items[0].IssueSummaries)
require.Equal(t, snapshot.UsernameSnapshot, page.Items[0].Snapshot.UsernameSnapshot)
require.Equal(t, snapshot.UserEmailSnapshot, page.Items[0].Snapshot.UserEmailSnapshot)
require.Equal(t, snapshot.APIKeyNameSnapshot, page.Items[0].Snapshot.APIKeyNameSnapshot)
_, err = db.Exec(`DELETE FROM users WHERE id=$1`, userID)
require.NoError(t, err)
_, err = db.Exec(`DELETE FROM api_keys WHERE id=$1`, apiKeyID)
require.NoError(t, err)
_, err = db.Exec(`DELETE FROM groups WHERE id=$1`, groupID)
require.NoError(t, err)
stored, err := repo.GetEvent(ctx, event.ID)
require.NoError(t, err)
require.Zero(t, stored.Snapshot.UserID)
require.Zero(t, stored.Snapshot.APIKeyID)
require.Nil(t, stored.Snapshot.GroupID)
require.Equal(t, snapshot.UsernameSnapshot, stored.Snapshot.UsernameSnapshot)
require.Equal(t, snapshot.UserEmailSnapshot, stored.Snapshot.UserEmailSnapshot)
require.Equal(t, snapshot.APIKeyNameSnapshot, stored.Snapshot.APIKeyNameSnapshot)
_, err = db.Exec(`DELETE FROM prompt_audit_jobs WHERE id=$1`, event.JobID)
require.NoError(t, err)
_, err = repo.GetEvent(ctx, event.ID)
require.ErrorIs(t, err, ErrEventNotFound)
}
func TestPromptAuditRepositoryHighWaterAndSafeDeletion(t *testing.T) {
db := openPromptAuditIntegrationDB(t)
repo := NewPostgreSQLRepository(db)
ctx := context.Background()
first, err := repo.RecordBlocking(ctx, integrationSnapshot("first"), 1, integrationResult(EventCritical), true)
require.NoError(t, err)
second, err := repo.RecordBlocking(ctx, integrationSnapshot("second"), 1, integrationResult(EventCritical), true)
require.NoError(t, err)
start, end := time.Now().Add(-time.Hour), time.Now().Add(time.Hour)
filter := EventFilter{Decision: string(EventCritical), StartAt: &start, EndAt: &end}
preview, err := repo.PreviewDelete(ctx, filter)
require.NoError(t, err)
require.Equal(t, int64(2), preview.MatchedCount)
require.Equal(t, second.ID, preview.SnapshotMaxID)
require.Equal(t, FilterHash(preview.FilterSummary, preview.SnapshotMaxID), preview.FilterHash)
newer, err := repo.RecordBlocking(ctx, integrationSnapshot("newer"), 1, integrationResult(EventCritical), true)
require.NoError(t, err)
result, err := repo.DeleteEventsByFilter(ctx, filter, preview.SnapshotMaxID, 1)
require.NoError(t, err)
require.Equal(t, int64(2), result.DeletedEvents)
require.Equal(t, int64(2), result.DeletedJobs)
_, err = repo.GetEvent(ctx, first.ID)
require.ErrorIs(t, err, ErrEventNotFound)
_, err = repo.GetEvent(ctx, second.ID)
require.ErrorIs(t, err, ErrEventNotFound)
_, err = repo.GetEvent(ctx, newer.ID)
require.NoError(t, err, "an event created after preview must survive high-water deletion")
processingEvent, err := repo.RecordBlocking(ctx, integrationSnapshot("processing"), 1, integrationResult(EventCritical), true)
require.NoError(t, err)
_, err = db.Exec(`UPDATE prompt_audit_jobs SET status='processing' WHERE id=$1`, processingEvent.JobID)
require.NoError(t, err)
deleteResult, err := repo.DeleteEvent(ctx, processingEvent.ID)
require.NoError(t, err)
require.Equal(t, int64(1), deleteResult.DeletedEvents)
require.Zero(t, deleteResult.DeletedJobs)
var remaining int
require.NoError(t, db.QueryRow(`SELECT COUNT(*) FROM prompt_audit_jobs WHERE id=$1`, processingEvent.JobID).Scan(&remaining))
require.Equal(t, 1, remaining, "processing jobs must not be deleted as orphans")
batchOne, err := repo.RecordBlocking(ctx, integrationSnapshot("batch-one"), 1, integrationResult(EventCritical), true)
require.NoError(t, err)
batchTwo, err := repo.RecordBlocking(ctx, integrationSnapshot("batch-two"), 1, integrationResult(EventCritical), true)
require.NoError(t, err)
ids := []int64{batchTwo.ID, batchOne.ID, batchOne.ID}
sort.Slice(ids, func(i, j int) bool { return ids[i] > ids[j] })
batchResult, err := repo.DeleteEventsByIDs(ctx, ids)
require.NoError(t, err)
require.Equal(t, int64(2), batchResult.DeletedEvents)
}
func TestPromptAuditServiceConfirmationKeepsPostPreviewEventsAndConcurrentDeletesAreSafe(t *testing.T) {
db := openPromptAuditIntegrationDB(t)
repo := NewPostgreSQLRepository(db)
ctx := context.Background()
now := time.Now().UTC()
start, end := now.Add(-time.Hour), now.Add(time.Hour)
filter := EventFilter{Decision: string(EventCritical), StartAt: &start, EndAt: &end}
for i := 0; i < 12; i++ {
_, err := repo.RecordBlocking(ctx, integrationSnapshot(fmt.Sprintf("event-%02d", i)), 1, integrationResult(EventCritical), true)
require.NoError(t, err)
}
service := &PromptService{
config: &fakeConfigStore{}, repo: repo, payload: NewRedisPayloadStore(nil), clock: fixedClock{now: now},
}
preview, err := service.PreviewDelete(ctx, filter, 77)
require.NoError(t, err)
require.Equal(t, int64(12), preview.MatchedCount)
newer, err := repo.RecordBlocking(ctx, integrationSnapshot("post-preview"), 1, integrationResult(EventCritical), true)
require.NoError(t, err)
result, err := service.DeleteByFilter(ctx, DeleteByFilterRequest{
Filter: filter, SnapshotMaxID: preview.SnapshotMaxID, FilterHash: preview.FilterHash,
ConfirmationToken: preview.ConfirmationToken, Confirm: true,
}, 77)
require.NoError(t, err)
require.Equal(t, int64(12), result.DeletedEvents)
_, err = repo.GetEvent(ctx, newer.ID)
require.NoError(t, err, "events created after delete-preview must survive")
resetPromptAuditIntegrationDB(t, db)
for i := 0; i < 24; i++ {
_, err := repo.RecordBlocking(ctx, integrationSnapshot(fmt.Sprintf("race-%02d", i)), 1, integrationResult(EventCritical), true)
require.NoError(t, err)
}
preview, err = repo.PreviewDelete(ctx, filter)
require.NoError(t, err)
type deleteOutcome struct {
result *DeleteResult
err error
}
outcomes := make(chan deleteOutcome, 2)
var wg sync.WaitGroup
for i := 0; i < 2; i++ {
wg.Add(1)
go func() {
defer wg.Done()
deleted, deleteErr := repo.DeleteEventsByFilter(ctx, filter, preview.SnapshotMaxID, 1)
outcomes <- deleteOutcome{result: deleted, err: deleteErr}
}()
}
wg.Wait()
close(outcomes)
var deletedTotal int64
for outcome := range outcomes {
require.NoError(t, outcome.err)
require.NotNil(t, outcome.result)
deletedTotal += outcome.result.DeletedEvents
}
require.Equal(t, int64(24), deletedTotal, "concurrent deleters must neither double-count nor strand matching events")
remaining, err := repo.ListEvents(ctx, filter, 1, 100)
require.NoError(t, err)
require.Zero(t, remaining.Total)
}
@@ -0,0 +1,25 @@
package securityaudit
import "testing"
func TestShouldStorePromptAuditEvent(t *testing.T) {
tests := []struct {
name string
storePassEvents bool
decision EventDecision
want bool
}{
{name: "pass disabled", storePassEvents: false, decision: EventPass, want: false},
{name: "flag disabled", storePassEvents: false, decision: EventFlag, want: true},
{name: "critical disabled", storePassEvents: false, decision: EventCritical, want: true},
{name: "pass enabled", storePassEvents: true, decision: EventPass, want: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := shouldStorePromptAuditEvent(tt.decision, tt.storePassEvents); got != tt.want {
t.Fatalf("shouldStorePromptAuditEvent(%q, %t) = %t, want %t", tt.decision, tt.storePassEvents, got, tt.want)
}
})
}
}
@@ -0,0 +1,122 @@
package securityaudit
import (
"errors"
"sort"
"strings"
"time"
)
func SplitRunes(value string, limit int) []string {
if limit <= 0 {
return nil
}
segments := strings.Split(value, promptAuditPrioritySeparator)
chunks := make([]string, 0, len(segments))
for _, segment := range segments {
runes := []rune(segment)
for start := 0; start < len(runes); start += limit {
end := start + limit
if end > len(runes) {
end = len(runes)
}
chunks = append(chunks, string(runes[start:end]))
}
}
return chunks
}
func AggregateResults(results []*NormalizedResult, latency time.Duration) (*NormalizedResult, error) {
if len(results) == 0 {
return nil, errors.New("prompt guard produced no complete result")
}
aggregated := &NormalizedResult{
Decision: EventPass, RiskLevel: RiskLow, Action: ActionAllow,
ScannerBackend: "qwen3guard-openai", Categories: []string{}, MatchedScanners: []string{},
ScannerScores: map[string]float64{}, ScannerEvidence: map[string]string{}, ChunkTotal: len(results),
LatencyMS: int(latency.Milliseconds()),
}
categories := map[string]struct{}{}
matched := map[string]struct{}{}
unknown := map[string]struct{}{}
for _, result := range results {
if result == nil {
return nil, errors.New("prompt guard partial result is not allowed")
}
if resultSeverity(result.Decision) > resultSeverity(aggregated.Decision) {
aggregated.Decision = result.Decision
aggregated.RiskLevel = result.RiskLevel
aggregated.Action = result.Action
aggregated.Safety = result.Safety
aggregated.GuardEndpointID = result.GuardEndpointID
aggregated.ScannerVersion = result.ScannerVersion
aggregated.PolicyID = result.PolicyID
aggregated.PolicyVersion = result.PolicyVersion
}
if aggregated.GuardEndpointID == "" {
aggregated.GuardEndpointID = result.GuardEndpointID
aggregated.ScannerVersion = result.ScannerVersion
aggregated.PolicyID = result.PolicyID
aggregated.PolicyVersion = result.PolicyVersion
}
for _, category := range result.Categories {
categories[category] = struct{}{}
}
for _, scanner := range result.MatchedScanners {
matched[scanner] = struct{}{}
}
for scanner, score := range result.ScannerScores {
if score > aggregated.ScannerScores[scanner] {
aggregated.ScannerScores[scanner] = score
}
}
for scanner, evidence := range result.ScannerEvidence {
if _, exists := aggregated.ScannerEvidence[scanner]; !exists {
aggregated.ScannerEvidence[scanner] = RedactPreview(evidence, 160)
}
}
for _, category := range result.UnknownCategories {
unknown[category] = struct{}{}
}
}
aggregated.Categories = orderedScannerKeys(categories)
aggregated.MatchedScanners = orderedScannerKeys(matched)
aggregated.UnknownCategories = sortedKeys(unknown)
return aggregated, nil
}
func resultSeverity(decision EventDecision) int {
switch decision {
case EventCritical:
return 3
case EventFlag:
return 2
default:
return 1
}
}
func sortedKeys(values map[string]struct{}) []string {
result := make([]string, 0, len(values))
for key := range values {
result = append(result, key)
}
sort.Strings(result)
return result
}
func orderedScannerKeys(values map[string]struct{}) []string {
result := make([]string, 0, len(values))
remaining := make(map[string]struct{}, len(values))
for key := range values {
remaining[key] = struct{}{}
}
for _, scannerID := range AllScannerIDs {
if _, ok := remaining[scannerID]; ok {
result = append(result, scannerID)
delete(remaining, scannerID)
}
}
result = append(result, sortedKeys(remaining)...)
return result
}
@@ -0,0 +1,478 @@
package securityaudit
import (
"context"
"encoding/json"
"errors"
"io"
"net"
"net/http"
"strings"
"sync"
"time"
)
type PromptService struct {
config ConfigStore
repo *PostgreSQLRepository
payload *RedisPayloadStore
enqueuer *Enqueuer
runner *Runner
evaluator *GuardEvaluator
scanner *OpenAICompatibleScanner
metrics *AtomicMetrics
clock Clock
lifecycleMu sync.Mutex
cancel context.CancelFunc
background context.Context
enqueueWG sync.WaitGroup
enqueueSlots chan struct{}
probeMu sync.RWMutex
probes map[string]ProbeResult
}
func NewPromptService(
config ConfigStore,
repo *PostgreSQLRepository,
payload *RedisPayloadStore,
scanner *OpenAICompatibleScanner,
metrics *AtomicMetrics,
) *PromptService {
enqueuer := NewEnqueuer(config, repo, payload, metrics)
evaluator := NewGuardEvaluator(scanner, repo, metrics)
runner := NewRunner(config, repo, payload, scanner, metrics)
return &PromptService{
config: config, repo: repo, payload: payload, scanner: scanner, metrics: metrics,
enqueuer: enqueuer, evaluator: evaluator, runner: runner, clock: realClock{},
enqueueSlots: make(chan struct{}, 128), probes: map[string]ProbeResult{},
}
}
func (s *PromptService) Start(ctx context.Context) error {
if s == nil || s.config == nil || s.runner == nil {
return errors.New("prompt audit service unavailable")
}
s.lifecycleMu.Lock()
if s.cancel != nil {
s.lifecycleMu.Unlock()
return nil
}
background, cancel := context.WithCancel(ctx)
s.background, s.cancel = background, cancel
s.lifecycleMu.Unlock()
configErr := s.config.Start(background)
workerErr := s.runner.Start(background)
return errors.Join(configErr, workerErr)
}
func (s *PromptService) Shutdown(ctx context.Context) error {
if s == nil {
return nil
}
s.lifecycleMu.Lock()
cancel := s.cancel
s.cancel = nil
s.lifecycleMu.Unlock()
if cancel != nil {
cancel()
}
var workerErr error
if s.runner != nil {
workerErr = s.runner.Shutdown(ctx)
}
done := make(chan struct{})
go func() { s.enqueueWG.Wait(); close(done) }()
select {
case <-done:
case <-ctx.Done():
if workerErr == nil {
workerErr = ctx.Err()
}
}
var configErr error
if s.config != nil {
configErr = s.config.Shutdown(ctx)
}
if workerErr != nil {
return workerErr
}
return configErr
}
func (s *PromptService) EffectiveMode() Mode {
if s == nil || s.config == nil {
return ModeOff
}
return s.config.EffectiveMode()
}
func (s *PromptService) Enqueue(_ context.Context, req Request) error {
if s == nil || s.enqueuer == nil || s.EffectiveMode() != ModeAsync {
return nil
}
select {
case s.enqueueSlots <- struct{}{}:
default:
if s.metrics != nil {
s.metrics.IncDropped()
}
LogWarn(EventEnqueueDropped, map[string]any{"request_id": req.RequestID, "status": "dropped", "error_code": "local_enqueue_busy"})
return nil
}
s.lifecycleMu.Lock()
background := s.background
s.lifecycleMu.Unlock()
if background == nil {
<-s.enqueueSlots
return errors.New("prompt audit service not started")
}
requestCopy := req.Clone()
s.enqueueWG.Add(1)
go func() {
defer s.enqueueWG.Done()
defer func() { <-s.enqueueSlots }()
ctx, cancel := context.WithTimeout(background, 2*time.Second)
defer cancel()
_ = s.enqueuer.Enqueue(ctx, requestCopy)
}()
return nil
}
func (s *PromptService) Evaluate(ctx context.Context, req Request) (*PromptDecision, error) {
if s == nil || s.config == nil || s.evaluator == nil {
return nil, &GuardError{Code: ErrorCodeUnavailable}
}
if s.config.BlockingActivationDegraded() {
return nil, &GuardError{Code: ErrorCodeUnavailable}
}
cfg, ok := s.config.Active()
if !ok {
if s.config.EffectiveMode() == ModeBlocking {
return nil, &GuardError{Code: ErrorCodeUnavailable}
}
return &PromptDecision{Kind: DecisionAllow, AllowNextStage: true}, nil
}
if cfg.EffectiveMode() != ModeBlocking || !cfg.IncludesGroup(req.GroupID) {
return &PromptDecision{Kind: DecisionAllow, AllowNextStage: true}, nil
}
snapshot, err := ExtractBlockingPromptSnapshot(req, cfg.BlockingLatestTurnOnly)
if errors.Is(err, ErrNoPromptText) {
return &PromptDecision{Kind: DecisionAllow, AllowNextStage: true}, nil
}
if err != nil {
return nil, &GuardError{Code: ErrorCodeInvalidResponse, Cause: err}
}
return s.evaluator.Evaluate(ctx, cfg, snapshot)
}
func (s *PromptService) GetConfig() (PublicConfig, error) { return s.config.Public() }
func (s *PromptService) SaveConfig(ctx context.Context, req UpdateConfigRequest, actorID int64) (PublicConfig, error) {
return s.config.Save(ctx, req, actorID)
}
func (s *PromptService) Runtime(ctx context.Context) RuntimeSnapshot {
expected, activeVersion, loadedAt, loadError := s.config.RuntimeState()
cfg, hasConfig := s.config.Active()
mode := s.EffectiveMode()
workerTotal, queueCapacity := 0, 0
if hasConfig {
workerTotal, queueCapacity = cfg.WorkerCount, cfg.QueueCapacity
}
runtime := RuntimeSnapshot{
ProcessStatus: "disabled", EffectiveMode: mode, ExpectedConfigVersion: expected,
ActiveConfigVersion: activeVersion, ConfigLoadedAt: loadedAt, ConfigLoadError: loadError,
WorkerTotal: workerTotal, QueueCapacity: queueCapacity, DatabaseStatus: "ok", RedisStatus: "ok",
Endpoints: s.probeSnapshot(), GuardMetrics: s.metrics.Snapshot(),
}
if s.repo != nil {
stats, err := s.repo.QueueStats(ctx)
if err != nil {
runtime.DatabaseStatus = "error"
runtime.LastErrorCode = "database_unavailable"
} else {
runtime.Queue = stats
}
} else {
runtime.DatabaseStatus = "error"
}
if s.payload == nil || s.payload.Ping(ctx) != nil {
runtime.RedisStatus = "error"
if runtime.LastErrorCode == "" {
runtime.LastErrorCode = "payload_store_unavailable"
}
}
activeWorkers, processed, failed, heartbeat, lastProcessed, workerCode, workerMessage := s.runner.Snapshot()
runtime.WorkerActive, runtime.ProcessedTotal, runtime.FailedTotal = activeWorkers, processed, failed
if s.metrics != nil {
auditMetrics := s.metrics.AuditSnapshot()
runtime.EnqueuedTotal, runtime.DroppedTotal = auditMetrics.Enqueued, auditMetrics.Dropped
}
runtime.WorkerHeartbeatAt, runtime.LastProcessedAt = heartbeat, lastProcessed
if workerCode != "" {
runtime.LastErrorCode, runtime.LastErrorMessage = workerCode, workerMessage
}
if mode != ModeOff {
runtime.ProcessStatus = "running"
if loadError != "" || runtime.DatabaseStatus != "ok" || runtime.RedisStatus != "ok" || activeVersion != expected {
runtime.ProcessStatus = "degraded"
}
if heartbeat == nil || s.clock.Now().Sub(*heartbeat) > 10*time.Second {
runtime.ProcessStatus = "degraded"
}
}
return runtime
}
type ProbeRequest struct {
Endpoint UpdateEndpoint `json:"endpoint"`
}
func (s *PromptService) Probe(ctx context.Context, request ProbeRequest) ProbeResult {
started := s.clock.Now()
endpoint, tokenApplied, err := s.resolveProbeEndpoint(request.Endpoint)
if err != nil {
return s.finishProbe(request.Endpoint.ID, started, ProbeResult{Status: "failed", ErrorCode: "endpoint_invalid", Message: "审计节点配置无效"})
}
LogInfo(EventProbeStarted, map[string]any{"guard_endpoint_id": endpoint.ID, "status": "started"})
client, err := NewSecureHTTPClient(endpoint)
if err != nil {
return s.finishProbe(endpoint.ID, started, ProbeResult{Status: "failed", ErrorCode: "endpoint_unsafe", Message: "审计节点地址不在允许范围", TokenApplied: tokenApplied})
}
modelsURL, _ := ModelsURL(endpoint.BaseURL)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, modelsURL, nil)
if err != nil {
return s.finishProbe(endpoint.ID, started, ProbeResult{Status: "failed", ErrorCode: "probe_request_invalid", Message: "无法创建探测请求", TokenApplied: tokenApplied})
}
if endpoint.Token != "" {
req.Header.Set("Authorization", "Bearer "+endpoint.Token)
}
resp, err := client.Do(req)
if err != nil {
code := "connection_failed"
var netErr net.Error
if errors.Is(err, context.DeadlineExceeded) || (errors.As(err, &netErr) && netErr.Timeout()) {
code = "timeout"
}
return s.finishProbe(endpoint.ID, started, ProbeResult{Status: "failed", ErrorCode: code, Message: "无法连接审计节点", Retryable: true, TokenApplied: tokenApplied})
}
responseBody, readErr := io.ReadAll(io.LimitReader(resp.Body, maxGuardResponseBytes+1))
_ = resp.Body.Close()
if readErr != nil {
return s.finishProbe(endpoint.ID, started, ProbeResult{Status: "failed", ErrorCode: "response_read_failed", Message: "审计节点响应读取失败", HTTPStatus: resp.StatusCode, Retryable: true, TokenApplied: tokenApplied})
}
if int64(len(responseBody)) > maxGuardResponseBytes {
return s.finishProbe(endpoint.ID, started, ProbeResult{Status: "failed", ErrorCode: "response_too_large", Message: "审计节点响应无效", HTTPStatus: resp.StatusCode, TokenApplied: tokenApplied})
}
if resp.StatusCode >= 200 && resp.StatusCode < 300 && modelsResponseReady(responseBody, endpoint.Model) {
return s.finishProbe(endpoint.ID, started, ProbeResult{OK: true, Status: "healthy", Message: "审计节点连接正常", HTTPStatus: resp.StatusCode, TokenApplied: tokenApplied})
}
if (resp.StatusCode >= 200 && resp.StatusCode < 300) || resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusMethodNotAllowed {
result, scanErr := s.scanner.Scan(ctx, endpoint, "Hello", AllScannerIDs)
if scanErr == nil && result != nil {
return s.finishProbe(endpoint.ID, started, ProbeResult{OK: true, Status: "healthy", Message: "审计节点模型调用正常", HTTPStatus: http.StatusOK, TokenApplied: tokenApplied})
}
code, status, retryable := guardErrorCode(scanErr), 0, false
var guardErr *GuardError
if errors.As(scanErr, &guardErr) {
status, retryable = guardErr.HTTPStatus, guardErr.Retryable
}
if code == "" {
code = ErrorCodeInvalidResponse
}
return s.finishProbe(endpoint.ID, started, ProbeResult{Status: "failed", ErrorCode: code, Message: "审计节点模型调用失败", HTTPStatus: status, Retryable: retryable, TokenApplied: tokenApplied})
}
code, retryable := "probe_http_error", resp.StatusCode == 429 || resp.StatusCode >= 500
if resp.StatusCode == 401 || resp.StatusCode == 403 {
code = "authentication_failed"
}
return s.finishProbe(endpoint.ID, started, ProbeResult{Status: "failed", ErrorCode: code, Message: "审计节点探测失败", HTTPStatus: resp.StatusCode, Retryable: retryable, TokenApplied: tokenApplied})
}
func modelsResponseReady(body []byte, model string) bool {
var response struct {
Data []struct {
ID string `json:"id"`
} `json:"data"`
}
if json.Unmarshal(body, &response) != nil || response.Data == nil {
return false
}
model = strings.TrimSpace(model)
if model == "" {
return true
}
for _, item := range response.Data {
if strings.TrimSpace(item.ID) == model {
return true
}
}
return false
}
func (s *PromptService) resolveProbeEndpoint(input UpdateEndpoint) (ActiveEndpoint, bool, error) {
baseURL, err := NormalizeBaseURL(input.BaseURL)
if err != nil {
return ActiveEndpoint{}, false, err
}
token := strings.TrimSpace(input.Token)
if token == "" {
if cfg, ok := s.config.Active(); ok {
for _, endpoint := range cfg.Endpoints {
if endpoint.ID != strings.TrimSpace(input.ID) {
continue
}
// Reuse a stored credential only when the probe targets the same
// normalized base URL. Otherwise an admin probe could exfiltrate
// the Guard token to an attacker-controlled HTTPS host.
if endpoint.BaseURL == baseURL {
token = endpoint.Token
}
break
}
}
}
model := strings.TrimSpace(input.Model)
if model == "" {
model = DefaultGuardModel
}
timeout := input.TimeoutMS
if timeout == 0 {
timeout = DefaultTimeoutMS
}
limit := input.InputLimit
if limit == 0 {
limit = DefaultInputLimit
}
storage := storageConfig{Enabled: false, Strategy: "priority", WorkerCount: DefaultWorkerCount, QueueCapacity: DefaultQueueCapacity, Scanners: append([]string(nil), AllScannerIDs...), AllGroups: true,
Endpoints: []StorageEndpoint{{ID: strings.TrimSpace(input.ID), Name: strings.TrimSpace(input.Name), Protocol: "openai_compatible", BaseURL: baseURL, Model: model, TimeoutMS: timeout, InputLimit: limit}}}
if storage.Endpoints[0].ID == "" {
storage.Endpoints[0].ID = "probe"
}
if storage.Endpoints[0].Name == "" {
storage.Endpoints[0].Name = "Probe"
}
if err := validateStorageConfig(storage); err != nil {
return ActiveEndpoint{}, false, err
}
return ActiveEndpoint{ID: storage.Endpoints[0].ID, Name: storage.Endpoints[0].Name, Protocol: "openai_compatible", BaseURL: baseURL, Model: model, Token: token, TimeoutMS: timeout, InputLimit: limit, Enabled: true}, token != "", nil
}
func (s *PromptService) finishProbe(id string, started time.Time, result ProbeResult) ProbeResult {
result.CheckedAt = s.clock.Now()
result.LatencyMS = int(result.CheckedAt.Sub(started).Milliseconds())
if result.OK {
LogInfo(EventProbeFinished, map[string]any{"guard_endpoint_id": id, "status": result.Status, "latency_ms": result.LatencyMS, "http_status": result.HTTPStatus})
} else {
LogWarn(EventProbeFailed, map[string]any{"guard_endpoint_id": id, "status": result.Status, "latency_ms": result.LatencyMS, "http_status": result.HTTPStatus, "error_code": result.ErrorCode, "retryable": result.Retryable})
}
s.probeMu.Lock()
s.probes[id] = result
s.probeMu.Unlock()
return result
}
func (s *PromptService) probeSnapshot() map[string]ProbeResult {
s.probeMu.RLock()
defer s.probeMu.RUnlock()
result := make(map[string]ProbeResult, len(s.probes))
for id, probe := range s.probes {
result[id] = probe
}
return result
}
func (s *PromptService) ListEvents(ctx context.Context, filter EventFilter, page, pageSize int) (*EventPage, error) {
return s.repo.ListEvents(ctx, filter, page, pageSize)
}
func (s *PromptService) GetEvent(ctx context.Context, id int64) (*Event, error) {
return s.repo.GetEvent(ctx, id)
}
func (s *PromptService) DeleteEvent(ctx context.Context, id int64) (*DeleteResult, error) {
result, err := s.repo.DeleteEvent(ctx, id)
if err == nil {
s.deletePayloads(ctx, result.JobIDs)
}
return result, err
}
func (s *PromptService) DeleteEventsByIDs(ctx context.Context, ids []int64) (*DeleteResult, error) {
result, err := s.repo.DeleteEventsByIDs(ctx, ids)
if err == nil {
s.deletePayloads(ctx, result.JobIDs)
}
return result, err
}
type deleteClaims struct {
FilterHash string `json:"filter_hash"`
SnapshotMaxID int64 `json:"snapshot_max_id"`
AdminID int64 `json:"admin_id"`
IssuedAt time.Time `json:"issued_at"`
ExpiresAt time.Time `json:"expires_at"`
}
func (s *PromptService) PreviewDelete(ctx context.Context, filter EventFilter, adminID int64) (*DeletePreview, error) {
preview, err := s.repo.PreviewDelete(ctx, filter)
if err != nil {
return nil, err
}
now := s.clock.Now()
expires := now.Add(5 * time.Minute)
claimsRaw, _ := json.Marshal(deleteClaims{FilterHash: preview.FilterHash, SnapshotMaxID: preview.SnapshotMaxID, AdminID: adminID, IssuedAt: now, ExpiresAt: expires})
token, err := s.config.Encrypt(string(claimsRaw))
if err != nil {
return nil, err
}
preview.ConfirmationToken, preview.ExpiresAt = token, expires
LogInfo(EventDeletePreviewed, map[string]any{"user_id": adminID, "status": "previewed"})
return preview, nil
}
type DeleteByFilterRequest struct {
Filter EventFilter `json:"filter"`
SnapshotMaxID int64 `json:"snapshot_max_id"`
FilterHash string `json:"filter_hash"`
ConfirmationToken string `json:"confirmation_token"`
Confirm bool `json:"confirm"`
}
func (s *PromptService) DeleteByFilter(ctx context.Context, request DeleteByFilterRequest, adminID int64) (*DeleteResult, error) {
if !request.Confirm {
return nil, errors.New("prompt audit filter delete requires confirm=true")
}
plain, err := s.config.Decrypt(strings.TrimSpace(request.ConfirmationToken))
if err != nil {
return nil, errors.New("prompt audit confirmation token invalid")
}
var claims deleteClaims
if json.Unmarshal([]byte(plain), &claims) != nil {
return nil, errors.New("prompt audit confirmation token invalid")
}
computed := FilterHash(request.Filter, request.SnapshotMaxID)
if claims.AdminID != adminID || claims.SnapshotMaxID != request.SnapshotMaxID || claims.FilterHash != request.FilterHash || request.FilterHash != computed || !s.clock.Now().Before(claims.ExpiresAt) {
return nil, errors.New("prompt audit confirmation token does not match deletion request")
}
result, err := s.repo.DeleteEventsByFilter(ctx, request.Filter, request.SnapshotMaxID, 200)
if err == nil {
s.deletePayloads(ctx, result.JobIDs)
LogWarn(EventEventsFilterDeleted, map[string]any{"user_id": adminID, "status": "deleted"})
}
return result, err
}
func (s *PromptService) deletePayloads(ctx context.Context, jobIDs []int64) {
for _, id := range jobIDs {
_ = s.payload.Delete(ctx, id)
}
}
func parseTimeQuery(value string) *time.Time {
parsed, err := time.Parse(time.RFC3339, strings.TrimSpace(value))
if err != nil {
return nil
}
parsed = parsed.UTC()
return &parsed
}
@@ -0,0 +1,145 @@
package securityaudit
import (
"context"
"encoding/json"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
)
type staticSettingRepository struct {
values map[string]string
}
func (r staticSettingRepository) Get(context.Context, string) (*service.Setting, error) {
return nil, service.ErrSettingNotFound
}
func (r staticSettingRepository) GetValue(context.Context, string) (string, error) {
return "", service.ErrSettingNotFound
}
func (r staticSettingRepository) Set(context.Context, string, string) error { return nil }
func (r staticSettingRepository) GetMultiple(_ context.Context, keys []string) (map[string]string, error) {
result := make(map[string]string, len(keys))
for _, key := range keys {
result[key] = r.values[key]
}
return result, nil
}
func (r staticSettingRepository) SetMultiple(context.Context, map[string]string) error { return nil }
func (r staticSettingRepository) GetAll(context.Context) (map[string]string, error) {
return r.values, nil
}
func (r staticSettingRepository) Delete(context.Context, string) error { return nil }
func TestPromptServiceHasExplicitIdempotentLifecycle(t *testing.T) {
config := NewConfigManager(nil, staticSettingRepository{values: map[string]string{
SettingKeyPromptAuditConfig: "",
SettingKeyRiskControl: "false",
}}, nil, prefixEncryptor{}, testTotpKeyConfig())
service := NewPromptService(
config,
NewPostgreSQLRepository(nil),
NewRedisPayloadStore(nil),
NewOpenAICompatibleScanner(),
NewAtomicMetrics(),
)
require.Nil(t, service.cancel, "construction must not start background work")
require.NoError(t, service.Start(context.Background()))
require.NotNil(t, service.cancel)
require.NoError(t, service.Start(context.Background()), "Start must be idempotent")
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
require.NoError(t, service.Shutdown(ctx))
require.Nil(t, service.cancel)
require.NoError(t, service.Shutdown(ctx), "Shutdown must be idempotent")
}
func TestPromptServiceStartReportsDependencyFailureWithoutPanic(t *testing.T) {
service := &PromptService{}
require.Error(t, service.Start(context.Background()))
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
require.NoError(t, service.Shutdown(ctx))
}
func TestPromptServiceBlockingLatestTurnOnlyUsesNarrowSnapshot(t *testing.T) {
seen := make([]string, 0, 2)
evaluator := newGuardEvaluator(PromptScannerFunc(func(_ context.Context, _ ActiveEndpoint, chunk string, _ []string) (*NormalizedResult, error) {
seen = append(seen, chunk)
return &NormalizedResult{Decision: EventPass, RiskLevel: RiskLow, Action: ActionAllow, ScannerScores: map[string]float64{}, ScannerEvidence: map[string]string{}}, nil
}), nil, NewAtomicMetrics(), 2, 2)
service := &PromptService{
config: &fakeConfigStore{active: true, cfg: ActiveConfig{
RiskControlEnabled: true, Enabled: true, BlockingEnabled: true, BlockingLatestTurnOnly: true, AllGroups: true,
Scanners: AllScannerIDs, Endpoints: []ActiveEndpoint{{ID: "guard-1", Enabled: true, TimeoutMS: 1000, InputLimit: 4096}},
}},
evaluator: evaluator,
}
decision, err := service.Evaluate(context.Background(), Request{Protocol: "openai_chat_completions", Body: []byte(`{"messages":[{"role":"system","content":"system instruction"},{"role":"user","content":"older user input"},{"role":"assistant","content":"previous output"},{"role":"user","content":"latest user input"}]}`)})
require.NoError(t, err)
require.Equal(t, DecisionAllow, decision.Kind)
require.Equal(t, []string{"latest user input", "previous output"}, seen)
}
func TestPromptServiceRejectsInvalidDeleteConfirmationClaims(t *testing.T) {
now := time.Date(2026, 7, 16, 10, 0, 0, 0, time.UTC)
start, end := now.Add(-time.Hour), now.Add(time.Hour)
filter := EventFilter{Decision: string(EventCritical), StartAt: &start, EndAt: &end}
const snapshotMaxID int64 = 10
filterHash := FilterHash(filter, snapshotMaxID)
validClaims := deleteClaims{
FilterHash: filterHash, SnapshotMaxID: snapshotMaxID, AdminID: 7,
IssuedAt: now, ExpiresAt: now.Add(5 * time.Minute),
}
claimsToken := func(claims deleteClaims) string {
raw, err := json.Marshal(claims)
require.NoError(t, err)
return string(raw)
}
validRequest := DeleteByFilterRequest{
Filter: filter, SnapshotMaxID: snapshotMaxID, FilterHash: filterHash,
ConfirmationToken: claimsToken(validClaims), Confirm: true,
}
tests := []struct {
name string
request DeleteByFilterRequest
adminID int64
}{
{name: "confirm false", request: func() DeleteByFilterRequest { value := validRequest; value.Confirm = false; return value }(), adminID: 7},
{name: "malformed token", request: func() DeleteByFilterRequest {
value := validRequest
value.ConfirmationToken = "not-json"
return value
}(), adminID: 7},
{name: "different administrator", request: validRequest, adminID: 8},
{name: "filter hash mismatch", request: func() DeleteByFilterRequest {
value := validRequest
value.FilterHash = strings.Repeat("b", 64)
return value
}(), adminID: 7},
{name: "snapshot mismatch", request: func() DeleteByFilterRequest { value := validRequest; value.SnapshotMaxID++; return value }(), adminID: 7},
{name: "expired", request: func() DeleteByFilterRequest {
value := validRequest
claims := validClaims
claims.ExpiresAt = now
value.ConfirmationToken = claimsToken(claims)
return value
}(), adminID: 7},
}
service := &PromptService{config: &fakeConfigStore{}, clock: fixedClock{now: now}}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
result, err := service.DeleteByFilter(context.Background(), test.request, test.adminID)
require.Error(t, err)
require.Nil(t, result)
})
}
}
@@ -0,0 +1,659 @@
package securityaudit
import (
"crypto/sha256"
"encoding/hex"
"encoding/json"
"errors"
"regexp"
"sort"
"strings"
"unicode/utf8"
)
var (
ErrNoPromptText = errors.New("prompt audit request contains no user text")
bearerPattern = regexp.MustCompile(`(?i)\bBearer\s+[A-Za-z0-9._~+\-/]+=*`)
apiKeyPattern = regexp.MustCompile(`(?i)\b(sk|rk|pk|api[_-]?key|token|secret|password)[-_:=\s]+[A-Za-z0-9._~+\-/]{8,}`)
canaryPattern = regexp.MustCompile(`(?i)([A-Z]+_CANARY_)[A-Za-z0-9_-]+`)
emailPattern = regexp.MustCompile(`(?i)\b[A-Z0-9._%+\-]+@[A-Z0-9.\-]+\.[A-Z]{2,}\b`)
phonePattern = regexp.MustCompile(`(?:\+?\d[\d\s().-]{8,}\d)`)
)
const promptAuditPrioritySeparator = "\x00SUB2API_PROMPT_AUDIT_PRIORITY_END\x00"
type promptSegment struct {
text string
user bool
role string
}
func ExtractPromptSnapshot(req Request) (PromptSnapshot, error) {
return extractPromptSnapshot(req, false)
}
// ExtractBlockingPromptSnapshot builds the narrow, low-latency blocking input
// when configured. Asynchronous auditing always uses ExtractPromptSnapshot so
// the complete client-controlled transcript is retained for review.
func ExtractBlockingPromptSnapshot(req Request, latestTurnOnly bool) (PromptSnapshot, error) {
return extractPromptSnapshot(req, latestTurnOnly)
}
func extractPromptSnapshot(req Request, latestTurnOnly bool) (PromptSnapshot, error) {
var document any
if err := json.Unmarshal(req.Body, &document); err != nil {
return PromptSnapshot{}, errors.New("prompt audit request JSON is invalid")
}
extracted := extractProtocolSegments(req.Protocol, document)
segments := normalizeSegmentsLatestUserFirst(extracted)
if latestTurnOnly {
segments = blockingSegmentsLatestUserAndPreviousOutput(extracted)
}
if len(segments) == 0 {
return PromptSnapshot{}, ErrNoPromptText
}
scanText, metadataText := buildPrioritizedScanText(segments)
digest := sha256.Sum256([]byte(metadataText))
stage := strings.TrimSpace(req.Stage)
if stage == "" {
stage = "http"
}
return PromptSnapshot{
RequestID: req.RequestID, UserID: req.UserID, UsernameSnapshot: req.Username,
UserEmailSnapshot: req.UserEmail, APIKeyID: req.APIKeyID, APIKeyNameSnapshot: req.APIKeyName,
GroupID: cloneInt64Ptr(req.GroupID), GroupName: req.GroupName, Provider: req.Provider,
Endpoint: req.Endpoint, Protocol: req.Protocol, Model: req.Model,
PromptHash: hex.EncodeToString(digest[:]), RedactedPreview: BuildPromptPreview(metadataText, DefaultPromptPreviewMaxRunes),
FullPrompt: BuildFullPrompt(metadataText, DefaultFullPromptMaxRunes),
PromptLength: utf8.RuneCountInString(metadataText), MessageCount: len(segments), Stage: stage,
ScanText: scanText,
}, nil
}
// DefaultPromptPreviewMaxRunes caps how much sanitized prompt text may be
// considered before BuildPromptPreview withholds the majority for storage/UI.
const DefaultPromptPreviewMaxRunes = 96
// DefaultFullPromptMaxRunes caps how much unredacted prompt text is persisted
// on an audit event for admin review. It is deliberately generous so realistic
// prompts are kept intact while bounding per-row storage.
const DefaultFullPromptMaxRunes = 65536
func extractProtocolSegments(protocol string, document any) []promptSegment {
root, _ := document.(map[string]any)
protocol = strings.ToLower(strings.TrimSpace(protocol))
switch protocol {
case "openai_chat_completions", "openai_chat", "chat_completions":
return extractChatLikeSegments(root)
case "anthropic_messages", "claude_messages", "messages":
return append(extractAnthropicSystem(root["system"]), extractMessages(root["messages"], clientInstructionRoles...)...)
case "gemini", "gemini_generate_content":
return extractGeminiRoot(root)
case "openai_responses", "responses", "responses_websocket":
if frameType := stringValue(root["type"]); frameType != "" || protocol == "responses_websocket" {
if frameType != "response.create" {
return nil
}
if input, exists := root["input"]; exists && input != nil {
return append(extractInstructions(root["instructions"]), extractResponses(input)...)
}
if response, ok := root["response"].(map[string]any); ok {
return append(extractInstructions(response["instructions"]), extractResponses(response["input"])...)
}
return extractInstructions(root["instructions"])
}
return append(extractInstructions(root["instructions"]), extractResponses(root["input"])...)
case "openai_images", "grok_media", "media", "images":
return userPromptSegments(extractMediaPrompts(root))
default:
if segments := extractChatLikeSegments(root); len(segments) > 0 {
return segments
}
if responses := append(extractInstructions(root["instructions"]), extractResponses(root["input"])...); len(responses) > 0 {
return responses
}
if gemini := extractGeminiRoot(root); len(gemini) > 0 {
return gemini
}
return userPromptSegments(extractMediaPrompts(root))
}
}
// clientInstructionRoles are roles a client may freely populate. Attackers can
// place jailbreak/PII text in assistant/tool turns, so blocking audit must scan
// them too—not only user/system/developer instructions.
var clientInstructionRoles = []string{"user", "system", "developer", "assistant", "tool"}
func extractChatLikeSegments(root map[string]any) []promptSegment {
if root == nil {
return nil
}
return extractMessages(root["messages"], clientInstructionRoles...)
}
func extractMessages(value any, wantedRoles ...string) []promptSegment {
items, ok := value.([]any)
if !ok {
return nil
}
wanted := make(map[string]struct{}, len(wantedRoles))
for _, role := range wantedRoles {
wanted[strings.ToLower(strings.TrimSpace(role))] = struct{}{}
}
result := make([]promptSegment, 0, len(items))
for _, item := range items {
message, ok := item.(map[string]any)
if !ok {
continue
}
role := strings.ToLower(stringValue(message["role"]))
if _, match := wanted[role]; !match {
continue
}
texts := contentTexts(message["content"])
for _, text := range texts {
result = append(result, promptSegment{text: text, user: role == "user", role: role})
}
}
return result
}
func extractInstructions(value any) []promptSegment {
switch typed := value.(type) {
case string:
if text := strings.TrimSpace(typed); text != "" {
return []promptSegment{{text: text, role: "system"}}
}
case []any:
return systemPromptSegments(contentTexts(typed))
case map[string]any:
return systemPromptSegments(contentTexts(typed))
}
return nil
}
func extractAnthropicSystem(value any) []promptSegment {
switch typed := value.(type) {
case string:
if text := strings.TrimSpace(typed); text != "" {
return []promptSegment{{text: text, role: "system"}}
}
case []any:
return systemPromptSegments(contentTexts(typed))
case map[string]any:
return systemPromptSegments(contentTexts(typed))
}
return nil
}
func extractResponses(value any) []promptSegment {
switch typed := value.(type) {
case string:
return []promptSegment{{text: typed, user: true, role: "user"}}
case []any:
result := make([]promptSegment, 0, len(typed))
for _, item := range typed {
switch entry := item.(type) {
case string:
result = append(result, promptSegment{text: entry, user: true, role: "user"})
case map[string]any:
role := strings.ToLower(stringValue(entry["role"]))
if role != "" && !isClientInstructionRole(role) {
continue
}
if content, exists := entry["content"]; exists {
for _, text := range contentTexts(content) {
result = append(result, promptSegment{text: text, user: role == "" || role == "user", role: role})
}
} else if text := stringValue(entry["text"]); text != "" {
result = append(result, promptSegment{text: text, user: role == "" || role == "user", role: role})
}
}
}
return result
case map[string]any:
role := strings.ToLower(stringValue(typed["role"]))
if role != "" && !isClientInstructionRole(role) {
return nil
}
return promptSegmentsForRole(contentTexts(typed["content"]), role)
default:
return nil
}
}
func isClientInstructionRole(role string) bool {
switch strings.ToLower(strings.TrimSpace(role)) {
case "user", "system", "developer", "assistant", "tool", "model":
return true
default:
return false
}
}
func extractGemini(value any) []promptSegment {
var contents []any
switch typed := value.(type) {
case []any:
contents = typed
case map[string]any:
contents = []any{typed}
default:
return nil
}
result := make([]promptSegment, 0, len(contents))
for _, item := range contents {
content, ok := item.(map[string]any)
if !ok {
continue
}
role := strings.ToLower(stringValue(content["role"]))
if role != "" && !isClientInstructionRole(role) {
continue
}
parts, _ := content["parts"].([]any)
for _, part := range parts {
if object, ok := part.(map[string]any); ok {
if text := stringValue(object["text"]); text != "" {
result = append(result, promptSegment{text: text, user: role == "" || role == "user", role: role})
}
}
}
}
return result
}
func extractGeminiRoot(root map[string]any) []promptSegment {
if root == nil {
return nil
}
result := extractGeminiSystemInstruction(root["systemInstruction"])
result = append(result, extractGeminiSystemInstruction(root["system_instruction"])...)
result = append(result, extractGemini(root["contents"])...)
result = append(result, extractGemini(root["content"])...)
result = append(result, extractGeminiInstances(root["instances"])...)
if requests, ok := root["requests"].([]any); ok {
for _, item := range requests {
request, ok := item.(map[string]any)
if !ok {
continue
}
result = append(result, extractGeminiSystemInstruction(request["systemInstruction"])...)
result = append(result, extractGeminiSystemInstruction(request["system_instruction"])...)
result = append(result, extractGemini(request["contents"])...)
result = append(result, extractGemini(request["content"])...)
result = append(result, extractGeminiInstances(request["instances"])...)
}
}
return result
}
func extractGeminiSystemInstruction(value any) []promptSegment {
switch typed := value.(type) {
case string:
if text := strings.TrimSpace(typed); text != "" {
return []promptSegment{{text: text, role: "system"}}
}
case map[string]any:
if parts, ok := typed["parts"].([]any); ok {
result := make([]promptSegment, 0, len(parts))
for _, part := range parts {
if object, ok := part.(map[string]any); ok {
if text := stringValue(object["text"]); text != "" {
result = append(result, promptSegment{text: text, role: "system"})
}
}
}
return result
}
return systemPromptSegments(contentTexts(typed))
case []any:
segments := extractGemini(typed)
for index := range segments {
segments[index].user = false
segments[index].role = "system"
}
return segments
}
return nil
}
func extractGeminiInstances(value any) []promptSegment {
instances, ok := value.([]any)
if !ok {
return nil
}
result := make([]promptSegment, 0, len(instances))
for _, item := range instances {
if instance, ok := item.(map[string]any); ok {
if prompt := stringValue(instance["prompt"]); prompt != "" {
result = append(result, promptSegment{text: prompt, user: true, role: "user"})
}
}
}
return result
}
func extractMediaPrompts(root map[string]any) []string {
if root == nil {
return nil
}
result := make([]string, 0, 4)
seen := map[string]struct{}{}
var walk func(any, string)
walk = func(value any, key string) {
switch typed := value.(type) {
case map[string]any:
keys := make([]string, 0, len(typed))
for childKey := range typed {
keys = append(keys, childKey)
}
sort.Strings(keys)
for _, childKey := range keys {
walk(typed[childKey], childKey)
}
case []any:
for _, item := range typed {
walk(item, key)
}
case string:
if !isMediaPromptKey(key) || looksLikeMediaPayload(typed) {
return
}
text := strings.TrimSpace(typed)
if text == "" {
return
}
if _, duplicate := seen[text]; duplicate {
return
}
seen[text] = struct{}{}
result = append(result, text)
}
}
walk(root, "")
return result
}
func isMediaPromptKey(key string) bool {
normalized := strings.NewReplacer("_", "", "-", "").Replace(strings.ToLower(strings.TrimSpace(key)))
switch normalized {
case "prompt", "inputprompt", "textprompt", "description", "query", "lyrics", "negativeprompt",
"positiveprompt", "gptdescriptionprompt", "prompten", "finalprompt", "finalzhprompt",
"origprompt", "actualprompt", "imageprompt", "input":
return true
default:
return false
}
}
func looksLikeMediaPayload(value string) bool {
trimmed := strings.TrimSpace(value)
lower := strings.ToLower(trimmed)
if strings.HasPrefix(lower, "data:image/") || strings.HasPrefix(lower, "data:video/") ||
strings.HasPrefix(lower, "http://") || strings.HasPrefix(lower, "https://") {
return true
}
if len(trimmed) >= 256 {
for _, r := range trimmed {
alphaNumeric := (r >= 'A' && r <= 'Z') || (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9')
if !alphaNumeric && r != '+' && r != '/' && r != '=' {
return false
}
}
return true
}
return false
}
func contentTexts(value any) []string {
switch typed := value.(type) {
case string:
return []string{typed}
case []any:
result := make([]string, 0, len(typed))
for _, part := range typed {
object, ok := part.(map[string]any)
if !ok {
continue
}
typeName := strings.ToLower(stringValue(object["type"]))
if typeName != "" && typeName != "text" && typeName != "input_text" && typeName != "output_text" {
continue
}
if text := stringValue(object["text"]); text != "" {
result = append(result, text)
}
}
return result
case map[string]any:
if text := stringValue(typed["text"]); text != "" {
return []string{text}
}
}
return nil
}
func normalizeSegmentsLatestUserFirst(values []promptSegment) []string {
normalized := normalizedPromptSegments(values)
if len(normalized) == 0 {
return nil
}
priorityIndex := len(normalized) - 1
for index := len(normalized) - 1; index >= 0; index-- {
if isUserSegment(normalized[index]) {
priorityIndex = index
break
}
}
result := make([]string, 0, len(normalized))
result = append(result, normalized[priorityIndex].text)
for index, segment := range normalized {
if index != priorityIndex {
result = append(result, segment.text)
}
}
return result
}
// blockingSegmentsLatestUserAndPreviousOutput limits synchronous guard input to
// the current user turn and the nearest preceding assistant/model turn. It is
// deliberately opt-in because full transcript scanning remains stronger at
// finding client-controlled content placed in older or non-user messages.
func blockingSegmentsLatestUserAndPreviousOutput(values []promptSegment) []string {
normalized := normalizedPromptSegments(values)
latestUserStart := latestUserSegmentStart(normalized)
if latestUserStart < 0 {
// A request without user content cannot be narrowed safely. Preserve the
// established full-snapshot behavior for unusual protocol payloads.
return normalizeSegmentsLatestUserFirst(values)
}
latestUserEnd := latestUserStart
for latestUserEnd < len(normalized) && isUserSegment(normalized[latestUserEnd]) {
latestUserEnd++
}
currentUserText := make([]string, 0, latestUserEnd-latestUserStart)
for _, segment := range normalized[latestUserStart:latestUserEnd] {
currentUserText = append(currentUserText, segment.text)
}
// A single client turn may have several text content parts. Keep it in one
// priority segment so every part of the latest input is scanned before the
// prior output begins.
selected := []promptSegment{{text: strings.Join(currentUserText, "\n\n"), user: true, role: "user"}}
for index := latestUserStart - 1; index >= 0; index-- {
if !isAssistantOutputSegment(normalized[index]) {
continue
}
start := index
for start > 0 && isAssistantOutputSegment(normalized[start-1]) {
start--
}
selected = append(selected, normalized[start:index+1]...)
break
}
return promptSegmentTexts(selected)
}
func normalizedPromptSegments(values []promptSegment) []promptSegment {
normalized := make([]promptSegment, 0, len(values))
for _, value := range values {
value.text = strings.TrimSpace(value.text)
if value.text != "" {
normalized = append(normalized, value)
}
}
return normalized
}
func latestUserSegmentStart(values []promptSegment) int {
latest := -1
for index := len(values) - 1; index >= 0; index-- {
if isUserSegment(values[index]) {
latest = index
break
}
}
for latest > 0 && isUserSegment(values[latest-1]) {
latest--
}
return latest
}
func isUserSegment(segment promptSegment) bool {
return segment.user || segment.role == "user"
}
func isAssistantOutputSegment(segment promptSegment) bool {
return segment.role == "assistant" || segment.role == "model"
}
func promptSegmentTexts(values []promptSegment) []string {
result := make([]string, 0, len(values))
for _, value := range values {
result = append(result, value.text)
}
return result
}
func buildPrioritizedScanText(segments []string) (scanText string, metadataText string) {
metadataText = strings.Join(segments, "\n\n")
if len(segments) <= 1 {
return metadataText, metadataText
}
return segments[0] + promptAuditPrioritySeparator + strings.Join(segments[1:], "\n\n"), metadataText
}
func promptSegmentsForRole(texts []string, role string) []promptSegment {
result := make([]promptSegment, 0, len(texts))
for _, text := range texts {
result = append(result, promptSegment{text: text, user: role == "" || role == "user", role: role})
}
return result
}
func userPromptSegments(texts []string) []promptSegment {
return promptSegmentsForRole(texts, "user")
}
func systemPromptSegments(texts []string) []promptSegment {
return promptSegmentsForRole(texts, "system")
}
func RedactPreview(value string, maxRunes int) string {
value = bearerPattern.ReplaceAllString(value, "Bearer ***")
value = apiKeyPattern.ReplaceAllStringFunc(value, func(match string) string {
if index := strings.IndexAny(match, ":= \t"); index >= 0 {
return match[:index+1] + "***"
}
return "***"
})
value = canaryPattern.ReplaceAllString(value, "${1}***")
value = emailPattern.ReplaceAllString(value, "***@***")
value = phonePattern.ReplaceAllString(value, "***PHONE***")
return TrimRunes(value, maxRunes)
}
// BuildPromptPreview stores only a short, non-recoverable head of sanitized
// input. Ordinary confidential prompts must not land nearly intact in PostgreSQL
// or the admin UI merely because no secret regex matched.
func BuildPromptPreview(value string, maxRunes int) string {
if maxRunes <= 0 {
maxRunes = DefaultPromptPreviewMaxRunes
}
redacted := strings.TrimSpace(RedactPreview(value, maxRunes))
if redacted == "" {
return ""
}
runes := []rune(redacted)
hadTruncation := strings.HasSuffix(redacted, "…")
if hadTruncation && len(runes) > 0 {
runes = runes[:len(runes)-1]
}
if len(runes) == 0 {
return "***…"
}
// Short unlabelled secrets would otherwise leak a recoverable prefix (e.g.
// 20 runes → 5 visible). Fully withhold anything below the keep threshold.
const minLengthForPartialPreview = 32
if len(runes) < minLengthForPartialPreview {
if hadTruncation {
return "***…"
}
return "***"
}
// Keep at most a quarter of the already-truncated text, and never more than
// 24 runes, so the majority of prompt content is withheld by default.
keep := len(runes) / 4
if keep > 24 {
keep = 24
}
preview := string(runes[:keep]) + "***"
if hadTruncation || keep < len(runes) {
preview += "…"
}
return preview
}
// BuildFullPrompt returns the complete prompt text for audit-event storage and
// admin review, without redaction. NUL bytes are stripped because PostgreSQL
// TEXT rejects them, and the result is capped at maxRunes.
func BuildFullPrompt(value string, maxRunes int) string {
if maxRunes <= 0 {
maxRunes = DefaultFullPromptMaxRunes
}
value = strings.ReplaceAll(value, "\x00", "")
return TrimRunes(strings.TrimSpace(value), maxRunes)
}
// FullPromptFromScanText reconstructs the display prompt from the worker scan
// payload. buildPrioritizedScanText inserts exactly one priority separator
// between the prioritized segment and the remainder, so replacing it with the
// metadata joiner yields the original multi-segment text.
func FullPromptFromScanText(scanText string) string {
return BuildFullPrompt(strings.ReplaceAll(scanText, promptAuditPrioritySeparator, "\n\n"), DefaultFullPromptMaxRunes)
}
func TrimRunes(value string, limit int) string {
if limit <= 0 {
return ""
}
runes := []rune(value)
if len(runes) <= limit {
return value
}
return string(runes[:limit]) + "…"
}
func stringValue(value any) string {
text, _ := value.(string)
return strings.TrimSpace(text)
}
func cloneInt64Ptr(value *int64) *int64 {
if value == nil {
return nil
}
cloned := *value
return &cloned
}
@@ -0,0 +1,414 @@
package securityaudit
import (
"crypto/sha256"
"encoding/hex"
"encoding/json"
"errors"
"strings"
"testing"
"unicode/utf8"
"github.com/stretchr/testify/require"
)
func TestExtractPromptSnapshotProtocols(t *testing.T) {
tests := []struct {
protocol, body, first string
count int
}{
{"openai_chat_completions", `{"messages":[{"role":"user","content":"old"},{"role":"assistant","content":"assistant turn"},{"role":"user","content":[{"type":"text","text":"最新😀"}]}]}`, "最新😀", 3},
{"openai_responses", `{"input":[{"role":"user","content":[{"type":"input_text","text":"response text"}]}]}`, "response text", 1},
{"anthropic_messages", `{"messages":[{"role":"user","content":[{"type":"text","text":"claude"}]}]}`, "claude", 1},
{"gemini", `{"contents":[{"role":"user","parts":[{"text":"gemini"},{"inline_data":{"data":"BASE64"}}]}]}`, "gemini", 1},
{"openai_images", `{"prompt":"draw a cat","image":"BASE64SECRET"}`, "draw a cat", 1},
{"responses_websocket", `{"type":"response.create","response":{"input":"turn two"}}`, "turn two", 1},
}
for _, tt := range tests {
t.Run(tt.protocol, func(t *testing.T) {
snapshot, err := ExtractPromptSnapshot(Request{Protocol: tt.protocol, Body: []byte(tt.body), Stage: "http"})
require.NoError(t, err)
require.True(t, strings.HasPrefix(snapshot.ScanText, tt.first))
require.Equal(t, tt.count, snapshot.MessageCount)
require.Equal(t, utf8.RuneCountInString(metadataTextForTest(snapshot.ScanText)), snapshot.PromptLength)
require.NotEmpty(t, snapshot.PromptHash)
require.NotContains(t, snapshot.ScanText, "BASE64SECRET")
})
}
}
func TestSnapshotRedactsCanariesAndPreservesHashOfScanText(t *testing.T) {
body := `{"messages":[{"role":"user","content":"PROMPT_CANARY_ABC123 email@example.com +86 138 0013 8000 Bearer AUTH_CANARY_XYZ sk-secretvalue123 password=supersecret123"}]}`
snapshot, err := ExtractPromptSnapshot(Request{Protocol: "openai_chat_completions", Body: []byte(body)})
require.NoError(t, err)
require.NotContains(t, snapshot.RedactedPreview, "ABC123")
require.NotContains(t, snapshot.RedactedPreview, "email@example.com")
require.NotContains(t, snapshot.RedactedPreview, "AUTH_CANARY_XYZ")
require.NotContains(t, snapshot.RedactedPreview, "secretvalue123")
require.NotContains(t, snapshot.RedactedPreview, "supersecret123")
require.NotContains(t, snapshot.RedactedPreview, "138 0013 8000")
require.Contains(t, snapshot.ScanText, "PROMPT_CANARY_ABC123")
require.NotEqual(t, snapshot.ScanText, snapshot.RedactedPreview)
digest := sha256.Sum256([]byte(metadataTextForTest(snapshot.ScanText)))
require.Equal(t, hex.EncodeToString(digest[:]), snapshot.PromptHash)
require.Empty(t, snapshot.Redacted().ScanText)
}
func TestSnapshotFullPromptKeepsUnredactedText(t *testing.T) {
body := `{"messages":[{"role":"user","content":"PROMPT_CANARY_ABC123 email@example.com sk-secretvalue123"}]}`
snapshot, err := ExtractPromptSnapshot(Request{Protocol: "openai_chat_completions", Body: []byte(body)})
require.NoError(t, err)
// The full prompt is stored verbatim for admin review, unlike the preview.
require.Contains(t, snapshot.FullPrompt, "PROMPT_CANARY_ABC123 email@example.com sk-secretvalue123")
require.NotContains(t, snapshot.RedactedPreview, "PROMPT_CANARY_ABC123")
require.Equal(t, snapshot.FullPrompt, snapshot.Redacted().FullPrompt)
}
func TestBuildFullPromptStripsNULAndTruncates(t *testing.T) {
require.Equal(t, "abcd", BuildFullPrompt("ab\x00cd", 0))
long := strings.Repeat("长", DefaultFullPromptMaxRunes+10)
trimmed := BuildFullPrompt(long, DefaultFullPromptMaxRunes)
require.Equal(t, DefaultFullPromptMaxRunes+1, utf8.RuneCountInString(trimmed))
require.True(t, strings.HasSuffix(trimmed, "…"))
}
func TestFullPromptFromScanTextRestoresMultiSegmentLayout(t *testing.T) {
scanText, metadataText := buildPrioritizedScanText([]string{"latest user", "system policy", "earlier user"})
require.Contains(t, scanText, promptAuditPrioritySeparator)
require.Equal(t, metadataText, FullPromptFromScanText(scanText))
singleScan, singleMeta := buildPrioritizedScanText([]string{"only"})
require.NotContains(t, singleScan, promptAuditPrioritySeparator)
require.Equal(t, singleMeta, FullPromptFromScanText(singleScan))
}
func TestSplitRunesDoesNotSplitUTF8(t *testing.T) {
chunks := SplitRunes("中文😀éabc", 2)
require.Equal(t, []string{"中文", "😀e", "́a", "bc"}, chunks)
for _, chunk := range chunks {
require.True(t, utf8.ValidString(chunk))
}
require.Equal(t, "中文😀éabc", strings.Join(chunks, ""))
}
func TestSplitRunesKeepsPrioritySegmentIndependent(t *testing.T) {
latest := "请帮我编写一篇黄色小说 名字你来取"
history := strings.Repeat("AGENTS.md 项目约束。", 40)
chunks := SplitRunes(latest+promptAuditPrioritySeparator+history, 128)
require.Greater(t, len(chunks), 2)
require.Equal(t, latest, chunks[0])
require.Equal(t, history, strings.Join(chunks[1:], ""))
for _, chunk := range chunks {
require.NotContains(t, chunk, promptAuditPrioritySeparator)
}
}
func TestPromptSnapshotLatestUserTextBlockIsOnePrioritizedSegment(t *testing.T) {
body := []byte(`{
"messages":[
{"role":"user","content":"历史输入"},
{"role":"assistant","content":"assistant client injection"},
{"role":"tool","content":"tool client injection"},
{"role":"user","content":[
{"type":"text","text":"最新第一块😀"},
{"type":"image_url","image_url":{"url":"data:image/png;base64,IMAGE_CANARY_BASE64"}},
{"type":"text","text":"最新第二块é"}
]}
]
}`)
snapshot, err := ExtractPromptSnapshot(Request{Protocol: "openai_chat_completions", Body: body})
require.NoError(t, err)
require.Equal(t, 5, snapshot.MessageCount)
require.True(t, strings.HasPrefix(snapshot.ScanText, "最新第二块é"+promptAuditPrioritySeparator))
require.Contains(t, snapshot.ScanText, "最新第一块😀")
require.Contains(t, snapshot.ScanText, "历史输入")
require.Contains(t, snapshot.ScanText, "assistant client injection")
require.Contains(t, snapshot.ScanText, "tool client injection")
require.NotContains(t, snapshot.ScanText, "IMAGE_CANARY_BASE64")
require.Equal(t, utf8.RuneCountInString(metadataTextForTest(snapshot.ScanText)), snapshot.PromptLength)
}
func TestPromptSnapshotSeparatesAnthropicUserPromptFromHarnessBlocks(t *testing.T) {
latest := "请帮我编写一篇黄色小说 名字你来取"
agents := "# AGENTS.md instructions\n<INSTRUCTIONS>" + strings.Repeat("安全约束。", 80) + "</INSTRUCTIONS>"
environment := "<environment_context><cwd>/workspace</cwd></environment_context>"
body := []byte(`{"system":"system policy","messages":[{"role":"user","content":[` +
`{"type":"text","text":` + string(mustJSON(t, agents)) + `},` +
`{"type":"text","text":` + string(mustJSON(t, environment)) + `},` +
`{"type":"text","text":` + string(mustJSON(t, latest)) + `}` +
`]}]}`)
snapshot, err := ExtractPromptSnapshot(Request{Protocol: "anthropic_messages", Body: body})
require.NoError(t, err)
require.Equal(t, 4, snapshot.MessageCount)
require.True(t, strings.HasPrefix(snapshot.ScanText, latest+promptAuditPrioritySeparator))
require.True(t, strings.HasPrefix(snapshot.RedactedPreview, "请帮我编写一篇黄色小说"))
chunks := SplitRunes(snapshot.ScanText, 128)
require.Equal(t, latest, chunks[0])
require.Contains(t, strings.Join(chunks[1:], ""), "# AGENTS.md instructions")
require.Contains(t, strings.Join(chunks[1:], ""), "<environment_context>")
require.NotContains(t, strings.Join(chunks, ""), promptAuditPrioritySeparator)
}
func TestPromptSnapshotResponsesShapes(t *testing.T) {
tests := []struct {
name string
body string
want string
}{
{name: "string", body: `{"input":"plain response input"}`, want: "plain response input"},
{name: "message array", body: `{"input":[{"role":"assistant","content":"assistant turn"},{"role":"user","content":[{"type":"input_text","text":"message block"}]}]}`, want: "message block\n\nassistant turn"},
{name: "direct input text", body: `{"input":[{"type":"input_text","text":"direct block"}]}`, want: "direct block"},
{name: "single object", body: `{"input":{"role":"user","content":[{"type":"input_text","text":"single object"}]}}`, want: "single object"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
snapshot, err := ExtractPromptSnapshot(Request{Protocol: "openai_responses", Body: []byte(tt.body)})
require.NoError(t, err)
require.Equal(t, tt.want, metadataTextForTest(snapshot.ScanText))
})
}
}
func TestPromptSnapshotGeminiBatchShapesAndMediaExclusion(t *testing.T) {
body := []byte(`{
"contents":{"role":"user","parts":[{"text":"root content"},{"inlineData":{"data":"ROOT_BASE64"}}]},
"instances":[{"prompt":"instance prompt"}],
"requests":[
{"contents":[{"role":"model","parts":[{"text":"ignore model"}]},{"role":"user","parts":[{"text":"nested user"}]}]},
{"instances":[{"prompt":"nested instance"}]}
]
}`)
snapshot, err := ExtractPromptSnapshot(Request{Protocol: "gemini", Body: body})
require.NoError(t, err)
require.True(t, strings.HasPrefix(snapshot.ScanText, "nested instance"))
for _, expected := range []string{"root content", "instance prompt", "nested user", "nested instance"} {
require.Contains(t, snapshot.ScanText, expected)
}
require.NotContains(t, snapshot.ScanText, "ROOT_BASE64")
require.Contains(t, snapshot.ScanText, "ignore model")
}
func TestPromptSnapshotMediaOnlyExtractsDeterministicTextPrompts(t *testing.T) {
body := []byte(`{
"prompt":"draw a lighthouse",
"image":"data:image/png;base64,IMAGE_CANARY",
"input":{"negative_prompt":"no fog","image_prompt":"https://example.test/input.png","prompt":"draw a lighthouse"},
"request":{"lyrics":"ocean song","input":"` + strings.Repeat("A", 300) + `"},
"images":[{"description":"nested textual direction","image_url":"https://example.test/image.png"}]
}`)
snapshot, err := ExtractPromptSnapshot(Request{Protocol: "grok_media", Body: body})
require.NoError(t, err)
require.Equal(t, 4, snapshot.MessageCount)
for _, expected := range []string{"draw a lighthouse", "no fog", "ocean song", "nested textual direction"} {
require.Contains(t, snapshot.ScanText, expected)
}
require.Equal(t, 1, strings.Count(snapshot.ScanText, "draw a lighthouse"))
require.NotContains(t, snapshot.ScanText, "IMAGE_CANARY")
require.NotContains(t, snapshot.ScanText, "example.test")
require.NotContains(t, snapshot.ScanText, strings.Repeat("A", 100))
}
func TestResponsesWebSocketOnlyAuditsResponseCreateAndPreservesStage(t *testing.T) {
for _, stage := range []string{"first_turn", "subsequent_turn"} {
snapshot, err := ExtractPromptSnapshot(Request{
Protocol: "openai_responses", Stage: stage,
Body: []byte(`{"type":"response.create","response":{"model":"gpt-test","input":[{"role":"user","content":[{"type":"input_text","text":"ws turn"}]}]}}`),
})
require.NoError(t, err)
require.Equal(t, "ws turn", snapshot.ScanText)
require.Equal(t, stage, snapshot.Stage)
}
_, err := ExtractPromptSnapshot(Request{
Protocol: "openai_responses", Stage: "subsequent_turn",
Body: []byte(`{"type":"conversation.item.create","response":{"input":"must not scan this frame"}}`),
})
require.True(t, errors.Is(err, ErrNoPromptText))
}
func TestPromptSnapshotEmptyAndLongUnicodeInput(t *testing.T) {
_, err := ExtractPromptSnapshot(Request{Protocol: "openai_chat_completions", Body: []byte(`{"messages":[{"role":"function","content":"not audited role"},{"role":"user","content":" "}]}`)})
require.True(t, errors.Is(err, ErrNoPromptText))
latest := strings.Repeat("最新😀é", 80)
history := strings.Repeat("历史中文", 80)
body := []byte(`{"messages":[{"role":"user","content":` + string(mustJSON(t, history)) + `},{"role":"user","content":` + string(mustJSON(t, latest)) + `}]}`)
snapshot, err := ExtractPromptSnapshot(Request{Protocol: "openai_chat_completions", Body: body})
require.NoError(t, err)
require.True(t, strings.HasPrefix(snapshot.ScanText, latest))
chunks := SplitRunes(snapshot.ScanText, 127)
require.Equal(t, strings.Replace(snapshot.ScanText, promptAuditPrioritySeparator, "", 1), strings.Join(chunks, ""))
require.Equal(t, latest, chunks[0]+strings.Join(chunks[1:len(SplitRunes(latest, 127))], ""))
for _, chunk := range chunks {
require.LessOrEqual(t, len([]rune(chunk)), 127)
require.True(t, utf8.ValidString(chunk))
}
}
func TestPromptSnapshotIncludesClientControlledInstructions(t *testing.T) {
tests := []struct {
name, protocol, body string
want []string
}{
{
name: "openai system developer assistant tool",
protocol: "openai_chat_completions",
body: `{"messages":[{"role":"system","content":"system jailbreak"},{"role":"developer","content":"developer policy"},{"role":"assistant","content":"assistant jailbreak"},{"role":"tool","content":"tool payload"},{"role":"user","content":"hello"}]}`,
want: []string{"system jailbreak", "developer policy", "assistant jailbreak", "tool payload", "hello"},
},
{
name: "openai system only",
protocol: "openai_chat_completions",
body: `{"messages":[{"role":"system","content":"only system instruction"}]}`,
want: []string{"only system instruction"},
},
{
name: "responses instructions",
protocol: "openai_responses",
body: `{"instructions":"response instructions","input":[{"role":"user","content":[{"type":"input_text","text":"user turn"}]}]}`,
want: []string{"response instructions", "user turn"},
},
{
name: "anthropic system",
protocol: "anthropic_messages",
body: `{"system":"claude system","messages":[{"role":"user","content":[{"type":"text","text":"claude user"}]}]}`,
want: []string{"claude system", "claude user"},
},
{
name: "gemini systemInstruction",
protocol: "gemini",
body: `{"systemInstruction":{"parts":[{"text":"gemini system"}]},"contents":[{"role":"user","parts":[{"text":"gemini user"}]}]}`,
want: []string{"gemini system", "gemini user"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
snapshot, err := ExtractPromptSnapshot(Request{Protocol: tt.protocol, Body: []byte(tt.body)})
require.NoError(t, err)
for _, expected := range tt.want {
require.Contains(t, snapshot.ScanText, expected)
}
})
}
}
func TestBlockingPromptSnapshotLimitsInputToLatestUserAndPreviousOutput(t *testing.T) {
tests := []struct {
name, protocol, body, want string
omitted []string
}{
{
name: "chat keeps multipart latest user and prior assistant",
protocol: "openai_chat_completions",
body: `{"messages":[
{"role":"system","content":"system instruction"},
{"role":"user","content":"older user input"},
{"role":"assistant","content":"older assistant output"},
{"role":"tool","content":"tool payload"},
{"role":"assistant","content":"previous assistant output"},
{"role":"user","content":[{"type":"text","text":"latest user first part"},{"type":"text","text":"latest user second part"}]}
]}`,
want: "latest user first part\n\nlatest user second part" + promptAuditPrioritySeparator + "previous assistant output",
omitted: []string{"system instruction", "older user input", "older assistant output", "tool payload"},
},
{
name: "gemini keeps prior model output",
protocol: "gemini",
body: `{"systemInstruction":{"parts":[{"text":"system instruction"}]},"contents":[
{"role":"user","parts":[{"text":"older user input"}]},
{"role":"model","parts":[{"text":"previous model output"}]},
{"role":"user","parts":[{"text":"latest user input"}]}
]}`,
want: "latest user input" + promptAuditPrioritySeparator + "previous model output",
omitted: []string{"system instruction", "older user input"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
snapshot, err := ExtractBlockingPromptSnapshot(Request{Protocol: tt.protocol, Body: []byte(tt.body)}, true)
require.NoError(t, err)
require.Equal(t, tt.want, snapshot.ScanText)
for _, omitted := range tt.omitted {
require.NotContains(t, snapshot.ScanText, omitted)
}
})
}
}
func TestContentTextsIncludesSupportedTextTypes(t *testing.T) {
value := []any{
map[string]any{"type": "text", "text": "plain text"},
map[string]any{"type": "input_text", "text": "input text"},
map[string]any{"type": "output_text", "text": "output text"},
map[string]any{"type": "image_url", "text": "ignored text"},
}
require.Equal(t, []string{"plain text", "input text", "output text"}, contentTexts(value))
}
func TestResponsesOutputTextIncludedInFullAndLatestTurnSnapshots(t *testing.T) {
body := []byte(`{"input":[
{"type":"message","role":"user","content":[{"type":"input_text","text":"earlier user input"}]},
{"type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","annotations":[],"text":"captured previous assistant output"}]},
{"type":"message","role":"user","content":[{"type":"input_text","text":"captured latest user input"}]}
]}`)
req := Request{Protocol: "openai_responses", Body: body}
full, err := ExtractPromptSnapshot(req)
require.NoError(t, err)
require.Contains(t, full.ScanText, "captured previous assistant output")
require.Contains(t, full.FullPrompt, "captured previous assistant output")
require.Equal(t, 3, full.MessageCount)
latestTurn, err := ExtractBlockingPromptSnapshot(req, true)
require.NoError(t, err)
require.Equal(t, "captured latest user input"+promptAuditPrioritySeparator+"captured previous assistant output", latestTurn.ScanText)
require.Equal(t, 2, latestTurn.MessageCount)
require.NotContains(t, latestTurn.ScanText, "earlier user input")
}
func TestBlockingPromptSnapshotPreservesFullScopeByDefaultAndWithoutUserInput(t *testing.T) {
req := Request{Protocol: "openai_chat_completions", Body: []byte(`{"messages":[{"role":"system","content":"system instruction"},{"role":"user","content":"older user input"},{"role":"assistant","content":"previous output"},{"role":"user","content":"latest user input"}]}`)}
full, err := ExtractPromptSnapshot(req)
require.NoError(t, err)
defaultBlocking, err := ExtractBlockingPromptSnapshot(req, false)
require.NoError(t, err)
require.Equal(t, full, defaultBlocking)
noUser := Request{Protocol: "openai_chat_completions", Body: []byte(`{"messages":[{"role":"system","content":"system instruction"},{"role":"assistant","content":"assistant output"}]}`)}
fullWithoutUser, err := ExtractPromptSnapshot(noUser)
require.NoError(t, err)
narrowWithoutUser, err := ExtractBlockingPromptSnapshot(noUser, true)
require.NoError(t, err)
require.Equal(t, fullWithoutUser, narrowWithoutUser)
}
func TestBuildPromptPreviewWithholdsMajorityOfOrdinaryText(t *testing.T) {
prompt := strings.Repeat("机密业务提示词内容", 40)
preview := BuildPromptPreview(prompt, DefaultPromptPreviewMaxRunes)
require.NotEmpty(t, preview)
require.Contains(t, preview, "***")
require.LessOrEqual(t, utf8.RuneCountInString(strings.TrimSuffix(strings.TrimSuffix(preview, "…"), "***")), 24)
require.Less(t, utf8.RuneCountInString(preview), utf8.RuneCountInString(prompt)/2)
require.NotContains(t, preview, prompt)
}
func TestBuildPromptPreviewFullyMasksShortUnlabelledSecrets(t *testing.T) {
require.Equal(t, "***", BuildPromptPreview("short-secret-value!!", DefaultPromptPreviewMaxRunes))
require.Equal(t, "***", BuildPromptPreview(strings.Repeat("a", 31), DefaultPromptPreviewMaxRunes))
partial := BuildPromptPreview(strings.Repeat("b", 32), DefaultPromptPreviewMaxRunes)
require.True(t, strings.HasPrefix(partial, "b"))
require.Contains(t, partial, "***")
}
func mustJSON(t *testing.T, value string) []byte {
t.Helper()
raw, err := json.Marshal(value)
require.NoError(t, err)
return raw
}
func metadataTextForTest(scanText string) string {
return strings.Replace(scanText, promptAuditPrioritySeparator, "\n\n", 1)
}
@@ -0,0 +1,279 @@
package securityaudit
import (
"context"
"time"
)
const (
SettingKeyPromptAuditConfig = "prompt_audit_config"
SettingKeyRiskControl = "risk_control_enabled"
ConfigInvalidationChannel = "sub2api:prompt_guard:config:invalidate"
PayloadKeyPrefix = "sub2api:prompt_audit:payload:"
ErrorCodeBlocked = "prompt_guard_blocked"
ErrorCodeUnavailable = "prompt_guard_unavailable"
ErrorCodeInvalidResponse = "prompt_guard_invalid_response"
ErrorCodeConfigConflict = "prompt_audit_config_conflict"
ErrorCodeConfigUnavailable = "prompt_audit_config_unavailable"
ErrorCodeEncryptionKeyRequired = "prompt_audit_encryption_key_required"
ErrorCodeRequiresEnabled = "prompt_guard_requires_audit_enabled"
DefaultGuardModel = "sileader/qwen3guard:0.6b"
)
type Mode string
const (
ModeOff Mode = "off"
ModeAsync Mode = "async_audit"
ModeBlocking Mode = "blocking"
)
type DecisionKind string
const (
DecisionAllow DecisionKind = "allow"
DecisionFlag DecisionKind = "flag"
DecisionBlock DecisionKind = "block"
DecisionUnavailable DecisionKind = "unavailable"
DecisionInvalid DecisionKind = "invalid"
)
type EventDecision string
const (
EventPass EventDecision = "pass"
EventFlag EventDecision = "flag"
EventCritical EventDecision = "critical"
)
type RiskLevel string
const (
RiskLow RiskLevel = "low"
RiskMedium RiskLevel = "medium"
RiskHigh RiskLevel = "high"
RiskCritical RiskLevel = "critical"
)
type Action string
const (
ActionAllow Action = "Allow"
ActionWarn Action = "Warn"
ActionBlock Action = "Block"
)
type Request struct {
RequestID string
UserID int64
Username string
UserEmail string
APIKeyID int64
APIKeyName string
GroupID *int64
GroupName string
Provider string
Endpoint string
Protocol string
Model string
Body []byte
Stage string
}
func (r Request) Clone() Request {
r.Body = append([]byte(nil), r.Body...)
if r.GroupID != nil {
id := *r.GroupID
r.GroupID = &id
}
return r
}
type PromptSnapshot struct {
RequestID string `json:"request_id"`
UserID int64 `json:"user_id"`
UsernameSnapshot string `json:"username"`
UserEmailSnapshot string `json:"user_email"`
APIKeyID int64 `json:"api_key_id"`
APIKeyNameSnapshot string `json:"api_key_name"`
GroupID *int64 `json:"group_id,omitempty"`
GroupName string `json:"group_name"`
Provider string `json:"provider"`
Endpoint string `json:"endpoint"`
Protocol string `json:"protocol"`
Model string `json:"model"`
PromptHash string `json:"prompt_hash"`
RedactedPreview string `json:"redacted_preview"`
FullPrompt string `json:"full_prompt"`
PromptLength int `json:"prompt_length"`
MessageCount int `json:"message_count"`
Stage string `json:"stage"`
ScanText string `json:"-"`
}
func (s PromptSnapshot) Redacted() PromptSnapshot {
s.ScanText = ""
return s
}
type NormalizedResult struct {
Decision EventDecision `json:"decision"`
RiskLevel RiskLevel `json:"risk_level"`
Action Action `json:"action"`
Safety string `json:"safety"`
Categories []string `json:"categories"`
MatchedScanners []string `json:"matched_scanners"`
ScannerScores map[string]float64 `json:"scanner_scores"`
ScannerEvidence map[string]string `json:"scanner_evidence"`
ScannerBackend string `json:"scanner_backend"`
ScannerVersion string `json:"scanner_version"`
GuardEndpointID string `json:"guard_endpoint_id"`
PolicyID string `json:"policy_id"`
PolicyVersion int `json:"policy_version"`
ChunkTotal int `json:"chunk_total"`
LatencyMS int `json:"latency_ms"`
UnknownCategories []string `json:"unknown_categories,omitempty"`
}
type PromptDecision struct {
Kind DecisionKind `json:"kind"`
ErrorCode string `json:"error_code,omitempty"`
Result *NormalizedResult `json:"result,omitempty"`
AllowNextStage bool `json:"allow_next_stage"`
}
type LegacyDecision struct {
Allowed bool `json:"allowed"`
Blocked bool `json:"blocked"`
Flagged bool `json:"flagged"`
Message string `json:"message"`
StatusCode int `json:"status_code"`
ErrorCode string `json:"error_code"`
Action string `json:"action"`
}
type Decision struct {
Kind DecisionKind `json:"kind"`
HTTPStatus int `json:"http_status"`
ErrorCode string `json:"error_code,omitempty"`
ClientMessage string `json:"client_message,omitempty"`
Legacy *LegacyDecision `json:"legacy,omitempty"`
Prompt *PromptDecision `json:"prompt,omitempty"`
AllowNextStage bool `json:"allow_next_stage"`
}
type IssueSummary struct {
Category string `json:"category"`
ScannerID string `json:"scanner_id"`
Title string `json:"title"`
Description string `json:"description"`
Severity string `json:"severity"`
SeverityLabel string `json:"severity_label"`
Action string `json:"action"`
ActionLabel string `json:"action_label"`
Code string `json:"code"`
Score float64 `json:"score"`
Evidence string `json:"evidence"`
EvidenceHash string `json:"evidence_hash"`
StartRune *int `json:"start_rune,omitempty"`
EndRune *int `json:"end_rune,omitempty"`
}
type ProbeResult struct {
OK bool `json:"ok"`
Status string `json:"status"`
ErrorCode string `json:"error_code,omitempty"`
Message string `json:"message"`
LatencyMS int `json:"latency_ms"`
HTTPStatus int `json:"http_status"`
Retryable bool `json:"retryable"`
CheckedAt time.Time `json:"checked_at"`
TokenApplied bool `json:"token_applied"`
}
type GuardMetricsSnapshot struct {
Total int64 `json:"total"`
Allowed int64 `json:"allowed"`
Flagged int64 `json:"flagged"`
Blocked int64 `json:"blocked"`
Unavailable int64 `json:"unavailable"`
Invalid int64 `json:"invalid"`
Timeouts int64 `json:"timeouts"`
Failovers int64 `json:"failovers"`
BulkheadFull int64 `json:"bulkhead_full"`
RecordFailed int64 `json:"record_failed"`
LatencyCount int64 `json:"latency_count"`
LatencyAvgMS int64 `json:"latency_avg_ms"`
LatencyP50MS int64 `json:"latency_p50_ms"`
LatencyP95MS int64 `json:"latency_p95_ms"`
LatencyP99MS int64 `json:"latency_p99_ms"`
LatencyMaxMS int64 `json:"latency_max_ms"`
}
type AuditMetricsSnapshot struct {
Enqueued int64 `json:"enqueued"`
Dropped int64 `json:"dropped"`
}
type QueueStats struct {
Staging int64 `json:"staging"`
Queued int64 `json:"queued"`
Processing int64 `json:"processing"`
Retry int64 `json:"retry"`
Done int64 `json:"done"`
Failed int64 `json:"failed"`
Active int64 `json:"active"`
}
type RuntimeSnapshot struct {
ProcessStatus string `json:"process_status"`
EffectiveMode Mode `json:"effective_mode"`
ExpectedConfigVersion int64 `json:"expected_config_version"`
ActiveConfigVersion int64 `json:"active_config_version"`
ConfigLoadedAt *time.Time `json:"config_loaded_at,omitempty"`
ConfigLoadError string `json:"config_load_error,omitempty"`
WorkerTotal int `json:"worker_total"`
WorkerActive int64 `json:"worker_active"`
WorkerHeartbeatAt *time.Time `json:"worker_heartbeat_at,omitempty"`
QueueCapacity int `json:"queue_capacity"`
Queue QueueStats `json:"queue"`
ProcessedTotal int64 `json:"processed_total"`
FailedTotal int64 `json:"failed_total"`
EnqueuedTotal int64 `json:"enqueued_total"`
DroppedTotal int64 `json:"dropped_total"`
LastProcessedAt *time.Time `json:"last_processed_at,omitempty"`
LastErrorCode string `json:"last_error_code,omitempty"`
LastErrorMessage string `json:"last_error_message,omitempty"`
DatabaseStatus string `json:"database_status"`
RedisStatus string `json:"redis_status"`
Endpoints map[string]ProbeResult `json:"endpoints"`
GuardMetrics GuardMetricsSnapshot `json:"guard_metrics"`
}
type Clock interface {
Now() time.Time
}
type realClock struct{}
func (realClock) Now() time.Time { return time.Now().UTC() }
type Metrics interface {
Snapshot() GuardMetricsSnapshot
AuditSnapshot() AuditMetricsSnapshot
Observe(kind DecisionKind, latency time.Duration)
IncEnqueued()
IncDropped()
IncTimeout()
IncFailover()
IncBulkheadFull()
IncRecordFailed()
}
type PromptScanner interface {
Scan(ctx context.Context, endpoint ActiveEndpoint, chunk string, enabledScanners []string) (*NormalizedResult, error)
}
@@ -0,0 +1,350 @@
package securityaudit
import (
"context"
"errors"
"sync"
"sync/atomic"
"time"
)
type WorkerRuntime struct {
active atomic.Int64
processed atomic.Int64
failed atomic.Int64
heartbeatNS atomic.Int64
lastProcessedNS atomic.Int64
lastErrorMu sync.RWMutex
lastErrorCode string
lastErrorMessage string
}
type Runner struct {
config ConfigStore
repo JobRepository
payload PayloadStore
scanner PromptScanner
metrics Metrics
clock Clock
runtime WorkerRuntime
mu sync.Mutex
cancel context.CancelFunc
wg sync.WaitGroup
}
func NewRunner(config ConfigStore, repo JobRepository, payload PayloadStore, scanner PromptScanner, metrics Metrics) *Runner {
return &Runner{config: config, repo: repo, payload: payload, scanner: scanner, metrics: metrics, clock: realClock{}}
}
func (r *Runner) Start(ctx context.Context) error {
if r == nil || r.config == nil || r.repo == nil || r.payload == nil || r.scanner == nil {
return errors.New("prompt audit worker dependencies unavailable")
}
r.mu.Lock()
if r.cancel != nil {
r.mu.Unlock()
return nil
}
runCtx, cancel := context.WithCancel(ctx)
r.cancel = cancel
r.mu.Unlock()
if err := r.payload.Ping(runCtx); err != nil {
r.setLastError("payload_store_unavailable", err.Error())
}
for workerID := 0; workerID < MaxWorkerCount; workerID++ {
r.wg.Add(1)
go r.worker(runCtx, workerID)
}
r.wg.Add(1)
go r.reclaimer(runCtx)
return nil
}
func (r *Runner) Shutdown(ctx context.Context) error {
if r == nil {
return nil
}
r.mu.Lock()
cancel := r.cancel
r.cancel = nil
r.mu.Unlock()
if cancel != nil {
cancel()
}
done := make(chan struct{})
go func() { r.wg.Wait(); close(done) }()
select {
case <-done:
return nil
case <-ctx.Done():
LogWarn(EventProcessFailed, map[string]any{"status": "shutdown_timeout", "error_code": "worker_shutdown_timeout"})
return ctx.Err()
}
}
func (r *Runner) worker(ctx context.Context, workerID int) {
defer r.wg.Done()
ticker := time.NewTicker(500 * time.Millisecond)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
r.runtime.heartbeatNS.Store(r.clock.Now().UnixNano())
cfg, ok := r.config.Active()
if !ok || !cfg.RiskControlEnabled || !cfg.Enabled || workerID >= cfg.WorkerCount {
continue
}
for {
job, claimed, err := r.repo.ClaimNextJob(ctx, r.clock.Now())
if err != nil {
r.setLastError("claim_job_failed", err.Error())
break
}
if !claimed {
break
}
r.runtime.active.Add(1)
r.processSafely(ctx, workerID, cfg, job)
r.runtime.active.Add(-1)
}
}
}
}
func (r *Runner) processSafely(ctx context.Context, workerID int, cfg ActiveConfig, job *Job) {
defer func() {
if recovered := recover(); recovered != nil {
r.runtime.failed.Add(1)
// Panic values may contain scanner response fragments or prompt data.
// Keep only a stable generic message in runtime state and logs.
r.setLastError("worker_panic", "worker panic recovered")
_ = r.repo.Fail(ctx, job.ID, job.ClaimVersion, "worker_panic", "worker panic recovered")
LogError(EventProcessFailed, mergeLogFields(jobLogFields(job), map[string]any{"worker_id": workerID, "status": "failed", "error_code": "worker_panic"}))
}
}()
if err := r.processJob(ctx, workerID, cfg, job); err != nil {
r.runtime.failed.Add(1)
} else {
r.runtime.processed.Add(1)
r.runtime.lastProcessedNS.Store(r.clock.Now().UnixNano())
}
}
func (r *Runner) processJob(ctx context.Context, workerID int, cfg ActiveConfig, job *Job) error {
baseFields := jobLogFields(job)
LogInfo(EventAuditStarted, mergeLogFields(baseFields, map[string]any{"worker_id": workerID, "attempts": job.Attempts, "status": "processing"}))
scanText, err := r.payload.Get(ctx, job.ID)
if err != nil {
return r.finishFailure(ctx, job, &GuardError{Code: "payload_missing", Retryable: false, Cause: err})
}
// The job row only carries redacted metadata; the full prompt for the audit
// event is reconstructed here from the transient scan payload.
job.Snapshot.FullPrompt = FullPromptFromScanText(scanText)
endpoints := cfg.EnabledEndpoints()
if len(endpoints) == 0 {
return r.finishFailure(ctx, job, &GuardError{Code: "no_enabled_endpoint", Retryable: true})
}
chunks := SplitRunes(scanText, minimumInputLimit(endpoints))
results := make([]*NormalizedResult, 0, len(chunks))
started := r.clock.Now()
for index, chunk := range chunks {
if err := r.repo.RefreshLease(ctx, job.ID, job.ClaimVersion, r.clock.Now()); err != nil {
return err
}
chunkStarted := r.clock.Now()
LogInfo(EventChunkStarted, mergeLogFields(baseFields, map[string]any{"worker_id": workerID, "chunk_index": index + 1, "chunk_total": len(chunks), "chunk_chars": len([]rune(chunk)), "input_chars": job.Snapshot.PromptLength, "input_limit": minimumInputLimit(endpoints), "status": "started"}))
result, scanErr := scanWithFailover(ctx, r.scanner, cfg.Scanners, endpoints, chunk, r.metrics)
if scanErr != nil {
LogWarn(EventChunkFailed, mergeLogFields(baseFields, map[string]any{
"worker_id": workerID, "chunk_index": index + 1, "chunk_total": len(chunks),
"chunk_chars": len([]rune(chunk)), "input_chars": job.Snapshot.PromptLength,
"input_limit": minimumInputLimit(endpoints), "latency_ms": r.clock.Now().Sub(chunkStarted).Milliseconds(),
"error_code": guardErrorCode(scanErr), "status": "failed",
}))
r.observeAsyncFailure(scanErr, r.clock.Now().Sub(started))
return r.finishFailure(ctx, job, scanErr)
}
results = append(results, result)
LogInfo(EventChunkCompleted, mergeLogFields(baseFields, map[string]any{"worker_id": workerID, "chunk_index": index + 1, "chunk_total": len(chunks), "guard_endpoint_id": result.GuardEndpointID, "action": result.Action, "latency_ms": r.clock.Now().Sub(chunkStarted).Milliseconds(), "status": "completed"}))
if result.Action == ActionBlock {
break
}
}
aggregated, err := AggregateResults(results, r.clock.Now().Sub(started))
if err != nil {
if r.metrics != nil {
r.metrics.Observe(DecisionInvalid, r.clock.Now().Sub(started))
}
return r.finishFailure(ctx, job, &GuardError{Code: ErrorCodeInvalidResponse, Cause: err})
}
aggregated.ChunkTotal = len(chunks)
if r.metrics != nil {
r.metrics.Observe(decisionKindForResult(aggregated), r.clock.Now().Sub(started))
}
LogInfo(EventChunksAggregated, mergeLogFields(baseFields, map[string]any{
"worker_id": workerID, "decision": aggregated.Decision, "risk_level": aggregated.RiskLevel,
"action": aggregated.Action, "chunk_total": aggregated.ChunkTotal,
"latency_ms": aggregated.LatencyMS, "guard_endpoint_id": aggregated.GuardEndpointID, "status": "completed",
}))
event, err := r.repo.Complete(ctx, job, aggregated, cfg.StorePassEvents)
if err != nil {
return err
}
if deleteErr := r.payload.Delete(ctx, job.ID); deleteErr != nil {
LogWarn(EventProcessFailed, mergeLogFields(baseFields, map[string]any{"worker_id": workerID, "status": "payload_delete_deferred", "error_code": "payload_delete_failed"}))
}
LogInfo(EventProcessed, mergeLogFields(baseFields, map[string]any{"worker_id": workerID, "event_id": eventID(event), "decision": aggregated.Decision, "risk_level": aggregated.RiskLevel, "action": aggregated.Action, "guard_endpoint_id": aggregated.GuardEndpointID, "latency_ms": aggregated.LatencyMS, "status": "done"}))
if event != nil && aggregated.Decision != EventPass {
LogWarn(EventFindingRecorded, mergeLogFields(baseFields, map[string]any{"worker_id": workerID, "event_id": event.ID, "decision": aggregated.Decision, "risk_level": aggregated.RiskLevel, "action": aggregated.Action, "guard_endpoint_id": aggregated.GuardEndpointID, "status": "recorded"}))
}
return nil
}
func (r *Runner) observeAsyncFailure(err error, latency time.Duration) {
if r == nil || r.metrics == nil {
return
}
kind := DecisionUnavailable
if guardErrorCode(err) == ErrorCodeInvalidResponse {
kind = DecisionInvalid
}
r.metrics.Observe(kind, latency)
var guardErr *GuardError
if errors.As(err, &guardErr) && guardErr.Timeout {
r.metrics.IncTimeout()
}
}
func decisionKindForResult(result *NormalizedResult) DecisionKind {
if result == nil {
return DecisionInvalid
}
switch result.Action {
case ActionBlock:
return DecisionBlock
case ActionWarn:
return DecisionFlag
default:
return DecisionAllow
}
}
func (r *Runner) finishFailure(ctx context.Context, job *Job, err error) error {
baseFields := jobLogFields(job)
code := guardErrorCode(err)
retryable := false
var guardErr *GuardError
if errors.As(err, &guardErr) {
retryable = guardErr.Retryable
}
if retryable && job.Attempts < job.MaxAttempts {
next := r.clock.Now().Add(retryBackoff(job.Attempts))
if updateErr := r.repo.Retry(ctx, job.ID, job.ClaimVersion, next, code, "prompt guard temporarily unavailable"); updateErr != nil {
return updateErr
}
LogWarn(EventProcessFailed, mergeLogFields(baseFields, map[string]any{"attempts": job.Attempts, "max_attempts": job.MaxAttempts, "status": "retry", "error_code": code, "retryable": true}))
} else {
if updateErr := r.repo.Fail(ctx, job.ID, job.ClaimVersion, code, "prompt guard processing failed"); updateErr != nil {
return updateErr
}
_ = r.payload.Delete(ctx, job.ID)
LogError(EventProcessFailed, mergeLogFields(baseFields, map[string]any{"attempts": job.Attempts, "max_attempts": job.MaxAttempts, "status": "failed", "error_code": code, "retryable": false}))
}
r.setLastError(code, err.Error())
return err
}
func (r *Runner) reclaimer(ctx context.Context) {
defer r.wg.Done()
ticker := time.NewTicker(time.Minute)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
now := r.clock.Now()
count, err := r.repo.ReclaimStale(ctx, now.Add(-2*time.Minute), now.Add(-90*time.Second), 100)
if err != nil {
r.setLastError("reclaim_failed", err.Error())
continue
}
if count > 0 {
LogWarn(EventProcessingReclaimed, map[string]any{"reclaimed_total": count, "status": "reclaimed"})
}
}
}
}
func (r *Runner) Snapshot() (active, processed, failed int64, heartbeat, lastProcessed *time.Time, code, message string) {
if r == nil {
return
}
active, processed, failed = r.runtime.active.Load(), r.runtime.processed.Load(), r.runtime.failed.Load()
if ns := r.runtime.heartbeatNS.Load(); ns > 0 {
value := time.Unix(0, ns).UTC()
heartbeat = &value
}
if ns := r.runtime.lastProcessedNS.Load(); ns > 0 {
value := time.Unix(0, ns).UTC()
lastProcessed = &value
}
r.runtime.lastErrorMu.RLock()
code, message = r.runtime.lastErrorCode, r.runtime.lastErrorMessage
r.runtime.lastErrorMu.RUnlock()
return
}
func (r *Runner) setLastError(code, _ string) {
code, message := sanitizeStoredError(code)
r.runtime.lastErrorMu.Lock()
r.runtime.lastErrorCode = code
r.runtime.lastErrorMessage = message
r.runtime.lastErrorMu.Unlock()
}
func scanWithFailover(ctx context.Context, scanner PromptScanner, scanners []string, endpoints []ActiveEndpoint, chunk string, metrics Metrics) (*NormalizedResult, error) {
var lastErr error
for index, endpoint := range endpoints {
result, err := scanner.Scan(ctx, endpoint, chunk, scanners)
if err == nil && result != nil {
return result, nil
}
if err == nil {
err = &GuardError{Code: ErrorCodeInvalidResponse, Retryable: false}
}
lastErr = err
var guardErr *GuardError
if !errors.As(err, &guardErr) || !guardErr.Retryable {
return nil, err
}
if index < len(endpoints)-1 && metrics != nil {
metrics.IncFailover()
}
}
if lastErr == nil {
lastErr = &GuardError{Code: ErrorCodeUnavailable}
}
return nil, lastErr
}
func retryBackoff(attempt int) time.Duration {
switch attempt {
case 1:
return 5 * time.Second
case 2:
return 30 * time.Second
default:
return 2 * time.Minute
}
}
func eventID(event *Event) int64 {
if event == nil {
return 0
}
return event.ID
}
@@ -0,0 +1,595 @@
package securityaudit
import (
"context"
"errors"
"fmt"
"reflect"
"strings"
"sync"
"testing"
"time"
"github.com/stretchr/testify/require"
)
type fixedClock struct{ now time.Time }
func (c fixedClock) Now() time.Time { return c.now }
type advancingClock struct {
mu sync.Mutex
now time.Time
step time.Duration
}
func (c *advancingClock) Now() time.Time {
c.mu.Lock()
defer c.mu.Unlock()
c.now = c.now.Add(c.step)
return c.now
}
type fakeConfigStore struct {
cfg ActiveConfig
active bool
}
func (s *fakeConfigStore) Start(context.Context) error { return nil }
func (s *fakeConfigStore) Shutdown(context.Context) error { return nil }
func (s *fakeConfigStore) Active() (ActiveConfig, bool) { return cloneActiveConfig(s.cfg), s.active }
func (s *fakeConfigStore) EffectiveMode() Mode {
if s.BlockingActivationDegraded() {
return ModeBlocking
}
if !s.active {
return ModeOff
}
return s.cfg.EffectiveMode()
}
func (s *fakeConfigStore) BlockingActivationDegraded() bool { return false }
func (s *fakeConfigStore) Public() (PublicConfig, error) { return PublicConfig{}, nil }
func (s *fakeConfigStore) Save(context.Context, UpdateConfigRequest, int64) (PublicConfig, error) {
return PublicConfig{}, nil
}
func (s *fakeConfigStore) RuntimeState() (int64, int64, *time.Time, string) {
return s.cfg.ConfigVersion, s.cfg.ConfigVersion, nil, ""
}
func (s *fakeConfigStore) Encrypt(value string) (string, error) { return value, nil }
func (s *fakeConfigStore) Decrypt(value string) (string, error) { return value, nil }
type fakeJobRepository struct {
mu sync.Mutex
trace *[]string
createJob *Job
createErr error
publishErr error
refreshErr error
completeErr error
retryErr error
failErr error
createdSnapshot PromptSnapshot
markedCode string
completedResult *NormalizedResult
completedStore bool
completeCount int
eventCount int
retryAt time.Time
retryCode string
retried int
failedCode string
failed int
refreshes int
claimQueue []*Job
recordBlockingCalls int
recordBlockingSnapshot PromptSnapshot
recordBlockingResult *NormalizedResult
recordBlockingErr error
}
func (r *fakeJobRepository) record(value string) {
if r.trace != nil {
*r.trace = append(*r.trace, value)
}
}
func (r *fakeJobRepository) CreateStagingWithCapacity(_ context.Context, snapshot PromptSnapshot, _ int64, _, _ int) (*Job, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.record("create_staging")
r.createdSnapshot = snapshot
if r.createErr != nil {
return nil, r.createErr
}
if r.createJob == nil {
r.createJob = &Job{ID: 1, Snapshot: snapshot}
}
return r.createJob, nil
}
func (r *fakeJobRepository) PublishQueued(context.Context, int64) error {
r.mu.Lock()
defer r.mu.Unlock()
r.record("publish_queued")
return r.publishErr
}
func (r *fakeJobRepository) MarkStagingFailed(_ context.Context, _ int64, code, _ string) error {
r.mu.Lock()
defer r.mu.Unlock()
r.record("mark_staging_failed")
r.markedCode = code
return nil
}
func (r *fakeJobRepository) ClaimNextJob(context.Context, time.Time) (*Job, bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
if len(r.claimQueue) == 0 {
return nil, false, nil
}
job := r.claimQueue[0]
r.claimQueue = r.claimQueue[1:]
return job, true, nil
}
func (r *fakeJobRepository) RefreshLease(context.Context, int64, int64, time.Time) error {
r.mu.Lock()
defer r.mu.Unlock()
r.refreshes++
return r.refreshErr
}
func (r *fakeJobRepository) Complete(_ context.Context, _ *Job, result *NormalizedResult, storePass bool) (*Event, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.completeCount++
r.completedResult, r.completedStore = result, storePass
if r.completeErr != nil {
return nil, r.completeErr
}
if result.Decision == EventPass && !storePass {
return nil, nil
}
r.eventCount++
return &Event{ID: 99, Decision: result.Decision}, nil
}
func (r *fakeJobRepository) Retry(_ context.Context, _, _ int64, next time.Time, code, _ string) error {
r.mu.Lock()
defer r.mu.Unlock()
r.retried++
r.retryAt, r.retryCode = next, code
return r.retryErr
}
func (r *fakeJobRepository) Fail(_ context.Context, _, _ int64, code, _ string) error {
r.mu.Lock()
defer r.mu.Unlock()
r.failed++
r.failedCode = code
return r.failErr
}
func (r *fakeJobRepository) ReclaimStale(context.Context, time.Time, time.Time, int) (int64, error) {
return 0, nil
}
func (r *fakeJobRepository) QueueStats(context.Context) (QueueStats, error) { return QueueStats{}, nil }
func (r *fakeJobRepository) RecordBlocking(_ context.Context, snapshot PromptSnapshot, _ int64, result *NormalizedResult, _ bool) (*Event, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.recordBlockingCalls++
r.recordBlockingSnapshot, r.recordBlockingResult = snapshot, result
return nil, r.recordBlockingErr
}
type fakePayloadStore struct {
mu sync.Mutex
trace *[]string
values map[int64]string
setErr error
getErr error
deleteErr error
pingErr error
setTTL time.Duration
deleted []int64
}
func (s *fakePayloadStore) Set(_ context.Context, jobID int64, value string, ttl time.Duration) error {
s.mu.Lock()
defer s.mu.Unlock()
if s.trace != nil {
*s.trace = append(*s.trace, "payload_set")
}
if s.setErr != nil {
return s.setErr
}
if s.values == nil {
s.values = map[int64]string{}
}
s.values[jobID], s.setTTL = value, ttl
return nil
}
func (s *fakePayloadStore) Get(_ context.Context, jobID int64) (string, error) {
s.mu.Lock()
defer s.mu.Unlock()
if s.getErr != nil {
return "", s.getErr
}
value, ok := s.values[jobID]
if !ok {
return "", errors.New("missing")
}
return value, nil
}
func (s *fakePayloadStore) Delete(_ context.Context, jobID int64) error {
s.mu.Lock()
defer s.mu.Unlock()
if s.trace != nil {
*s.trace = append(*s.trace, "payload_delete")
}
s.deleted = append(s.deleted, jobID)
delete(s.values, jobID)
return s.deleteErr
}
func (s *fakePayloadStore) Ping(context.Context) error { return s.pingErr }
func asyncConfig() ActiveConfig {
return ActiveConfig{
RiskControlEnabled: true, Enabled: true, BlockingEnabled: false, Strategy: "priority",
WorkerCount: 1, QueueCapacity: 8, Scanners: []string{"pii"}, AllGroups: true, ConfigVersion: 7,
Endpoints: []ActiveEndpoint{{ID: "guard", Enabled: true, TimeoutMS: 1000, InputLimit: 3}},
}
}
func asyncRequest() Request {
return Request{RequestID: "request-async", Protocol: "openai_chat_completions", Body: []byte(`{"messages":[{"role":"user","content":"payload canary text"}]}`)}
}
func TestEnqueuerStagingPayloadPublishProtocolAndFailureCleanup(t *testing.T) {
t.Run("success", func(t *testing.T) {
trace := []string{}
repo := &fakeJobRepository{trace: &trace, createJob: &Job{ID: 41}}
payload := &fakePayloadStore{trace: &trace, values: map[int64]string{}}
enqueuer := NewEnqueuer(&fakeConfigStore{cfg: asyncConfig(), active: true}, repo, payload)
require.NoError(t, enqueuer.Enqueue(context.Background(), asyncRequest()))
require.Equal(t, []string{"create_staging", "payload_set", "publish_queued"}, trace)
require.Empty(t, repo.createdSnapshot.ScanText)
require.Equal(t, "payload canary text", payload.values[41])
require.Equal(t, DefaultPayloadTTL, payload.setTTL)
})
t.Run("queue admission failures never touch payload", func(t *testing.T) {
for _, createErr := range []error{ErrQueueFull, ErrQueueAdmissionBusy, errors.New("database down")} {
trace := []string{}
repo := &fakeJobRepository{trace: &trace, createErr: createErr}
payload := &fakePayloadStore{trace: &trace, values: map[int64]string{}}
err := NewEnqueuer(&fakeConfigStore{cfg: asyncConfig(), active: true}, repo, payload).Enqueue(context.Background(), asyncRequest())
require.ErrorIs(t, err, createErr)
require.Equal(t, []string{"create_staging"}, trace)
}
})
t.Run("payload failure marks staging failed", func(t *testing.T) {
trace := []string{}
repo := &fakeJobRepository{trace: &trace, createJob: &Job{ID: 42}}
payload := &fakePayloadStore{trace: &trace, values: map[int64]string{}, setErr: errors.New("redis down")}
err := NewEnqueuer(&fakeConfigStore{cfg: asyncConfig(), active: true}, repo, payload).Enqueue(context.Background(), asyncRequest())
require.Error(t, err)
require.Equal(t, []string{"create_staging", "payload_set", "mark_staging_failed"}, trace)
require.Equal(t, "payload_store_failed", repo.markedCode)
})
t.Run("publish failure removes payload and marks staging failed", func(t *testing.T) {
trace := []string{}
repo := &fakeJobRepository{trace: &trace, createJob: &Job{ID: 43}, publishErr: errors.New("publish down")}
payload := &fakePayloadStore{trace: &trace, values: map[int64]string{}}
err := NewEnqueuer(&fakeConfigStore{cfg: asyncConfig(), active: true}, repo, payload).Enqueue(context.Background(), asyncRequest())
require.Error(t, err)
require.Equal(t, []string{"create_staging", "payload_set", "publish_queued", "payload_delete", "mark_staging_failed"}, trace)
require.Equal(t, "queue_publish_failed", repo.markedCode)
require.NotContains(t, payload.values, int64(43))
})
}
func TestEnqueuerSkipsOffOutOfScopeAndNoText(t *testing.T) {
tests := []struct {
name string
cfg ActiveConfig
req Request
}{
{name: "off", cfg: ActiveConfig{}, req: asyncRequest()},
{name: "out of scope", cfg: func() ActiveConfig {
cfg := asyncConfig()
cfg.AllGroups = false
cfg.GroupIDs = []int64{9}
return cfg
}(), req: asyncRequest()},
{name: "no user text", cfg: asyncConfig(), req: Request{Protocol: "openai_chat_completions", Body: []byte(`{"messages":[{"role":"function","content":"not audited"}]}`)}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
repo := &fakeJobRepository{}
err := NewEnqueuer(&fakeConfigStore{cfg: tt.cfg, active: true}, repo, &fakePayloadStore{}).Enqueue(context.Background(), tt.req)
require.NoError(t, err)
require.Zero(t, repo.createdSnapshot.MessageCount)
})
}
}
func TestEnqueuerRecordsAcceptedDroppedAndSkippedMetrics(t *testing.T) {
t.Run("accepted increments enqueued", func(t *testing.T) {
metrics := NewAtomicMetrics()
repo := &fakeJobRepository{createJob: &Job{ID: 44}}
payload := &fakePayloadStore{values: map[int64]string{}}
require.NoError(t, NewEnqueuer(
&fakeConfigStore{cfg: asyncConfig(), active: true},
repo,
payload,
metrics,
).Enqueue(context.Background(), asyncRequest()))
require.Equal(t, AuditMetricsSnapshot{Enqueued: 1}, metrics.AuditSnapshot())
})
t.Run("queue full increments dropped", func(t *testing.T) {
metrics := NewAtomicMetrics()
repo := &fakeJobRepository{createErr: ErrQueueFull}
err := NewEnqueuer(
&fakeConfigStore{cfg: asyncConfig(), active: true},
repo,
&fakePayloadStore{},
metrics,
).Enqueue(context.Background(), asyncRequest())
require.ErrorIs(t, err, ErrQueueFull)
require.Equal(t, AuditMetricsSnapshot{Dropped: 1}, metrics.AuditSnapshot())
})
t.Run("skipped request does not increment dropped", func(t *testing.T) {
metrics := NewAtomicMetrics()
require.NoError(t, NewEnqueuer(
&fakeConfigStore{cfg: ActiveConfig{}, active: true},
&fakeJobRepository{},
&fakePayloadStore{},
metrics,
).Enqueue(context.Background(), asyncRequest()))
require.Equal(t, AuditMetricsSnapshot{}, metrics.AuditSnapshot())
})
}
func workerJob(attempts, maxAttempts int) *Job {
return &Job{ID: 51, ClaimVersion: 3, Attempts: attempts, MaxAttempts: maxAttempts, ConfigVersion: 7,
Snapshot: PromptSnapshot{RequestID: "worker-request", PromptLength: 6, RedactedPreview: "red***"}}
}
func TestWorkerCompletesPassWithoutEventRefreshesEveryChunkAndDeletesPayload(t *testing.T) {
repo := &fakeJobRepository{}
payload := &fakePayloadStore{values: map[int64]string{51: "abcdef"}}
scannerCalls := 0
scanner := PromptScannerFunc(func(_ context.Context, endpoint ActiveEndpoint, chunk string, _ []string) (*NormalizedResult, error) {
scannerCalls++
return &NormalizedResult{Decision: EventPass, RiskLevel: RiskLow, Action: ActionAllow, Safety: "Safe", Categories: []string{}, MatchedScanners: []string{}, ScannerScores: map[string]float64{}, ScannerEvidence: map[string]string{}, GuardEndpointID: endpoint.ID}, nil
})
metrics := NewAtomicMetrics()
runner := NewRunner(&fakeConfigStore{cfg: asyncConfig(), active: true}, repo, payload, scanner, metrics)
runner.clock = fixedClock{now: time.Unix(100, 0).UTC()}
require.NoError(t, runner.processJob(context.Background(), 0, asyncConfig(), workerJob(1, 3)))
require.Equal(t, 2, scannerCalls)
require.Equal(t, 2, repo.refreshes)
require.NotNil(t, repo.completedResult)
require.Equal(t, EventPass, repo.completedResult.Decision)
require.False(t, repo.completedStore)
require.Equal(t, []int64{51}, payload.deleted)
require.Equal(t, int64(1), metrics.Snapshot().Total)
require.Equal(t, int64(1), metrics.Snapshot().Allowed)
}
func TestWorkerRetryBackoffTerminalFailureAndFailover(t *testing.T) {
now := time.Unix(200, 0).UTC()
for _, tt := range []struct {
name string
attempts int
maxAttempts int
err *GuardError
wantRetry bool
wantBackoff time.Duration
}{
{name: "first retry", attempts: 1, maxAttempts: 3, err: &GuardError{Code: ErrorCodeUnavailable, Retryable: true}, wantRetry: true, wantBackoff: 5 * time.Second},
{name: "second retry", attempts: 2, maxAttempts: 3, err: &GuardError{Code: ErrorCodeUnavailable, Retryable: true}, wantRetry: true, wantBackoff: 30 * time.Second},
{name: "third retry", attempts: 3, maxAttempts: 4, err: &GuardError{Code: ErrorCodeUnavailable, Retryable: true}, wantRetry: true, wantBackoff: 2 * time.Minute},
{name: "max attempts", attempts: 3, maxAttempts: 3, err: &GuardError{Code: ErrorCodeUnavailable, Retryable: true}},
{name: "invalid terminal", attempts: 1, maxAttempts: 3, err: &GuardError{Code: ErrorCodeInvalidResponse, Retryable: false}},
} {
t.Run(tt.name, func(t *testing.T) {
repo := &fakeJobRepository{}
payload := &fakePayloadStore{values: map[int64]string{51: "abc"}}
metrics := NewAtomicMetrics()
runner := NewRunner(&fakeConfigStore{cfg: asyncConfig(), active: true}, repo, payload, PromptScannerFunc(func(context.Context, ActiveEndpoint, string, []string) (*NormalizedResult, error) {
return nil, tt.err
}), metrics)
runner.clock = fixedClock{now: now}
err := runner.processJob(context.Background(), 0, asyncConfig(), workerJob(tt.attempts, tt.maxAttempts))
require.Error(t, err)
if tt.wantRetry {
require.Equal(t, 1, repo.retried)
require.Equal(t, now.Add(tt.wantBackoff), repo.retryAt)
require.Empty(t, payload.deleted)
} else {
require.Equal(t, 1, repo.failed)
require.Equal(t, tt.err.Code, repo.failedCode)
require.Equal(t, []int64{51}, payload.deleted)
}
snapshot := metrics.Snapshot()
require.Equal(t, int64(1), snapshot.Total)
if tt.err.Code == ErrorCodeInvalidResponse {
require.Equal(t, int64(1), snapshot.Invalid)
} else {
require.Equal(t, int64(1), snapshot.Unavailable)
}
})
}
repo := &fakeJobRepository{}
payload := &fakePayloadStore{values: map[int64]string{51: "abc"}}
metrics := NewAtomicMetrics()
scanner := PromptScannerFunc(func(_ context.Context, endpoint ActiveEndpoint, _ string, _ []string) (*NormalizedResult, error) {
if endpoint.ID == "first" {
return nil, &GuardError{Code: ErrorCodeUnavailable, Retryable: true}
}
return integrationResult(EventPass), nil
})
cfg := asyncConfig()
cfg.Endpoints = []ActiveEndpoint{{ID: "first", Enabled: true, InputLimit: 10}, {ID: "second", Enabled: true, InputLimit: 10}}
runner := NewRunner(&fakeConfigStore{cfg: cfg, active: true}, repo, payload, scanner, metrics)
require.NoError(t, runner.processJob(context.Background(), 0, cfg, workerJob(1, 3)))
require.Equal(t, int64(1), metrics.Snapshot().Failovers)
}
func TestWorkerPanicLeaseLossAndLifecycleAreContained(t *testing.T) {
t.Run("panic", func(t *testing.T) {
repo := &fakeJobRepository{}
payload := &fakePayloadStore{values: map[int64]string{51: "abc"}}
runner := NewRunner(&fakeConfigStore{cfg: asyncConfig(), active: true}, repo, payload, PromptScannerFunc(func(context.Context, ActiveEndpoint, string, []string) (*NormalizedResult, error) {
panic("scanner panic canary")
}), NewAtomicMetrics())
require.NotPanics(t, func() { runner.processSafely(context.Background(), 0, asyncConfig(), workerJob(1, 3)) })
_, _, failed, _, _, code, message := runner.Snapshot()
require.Equal(t, int64(1), failed)
require.Equal(t, "worker_panic", code)
require.NotContains(t, message, "canary")
require.Equal(t, 1, repo.failed)
})
t.Run("lease loss", func(t *testing.T) {
repo := &fakeJobRepository{refreshErr: ErrLeaseLost}
payload := &fakePayloadStore{values: map[int64]string{51: "abc"}}
calls := 0
runner := NewRunner(&fakeConfigStore{cfg: asyncConfig(), active: true}, repo, payload, PromptScannerFunc(func(context.Context, ActiveEndpoint, string, []string) (*NormalizedResult, error) {
calls++
return integrationResult(EventPass), nil
}), NewAtomicMetrics())
require.ErrorIs(t, runner.processJob(context.Background(), 0, asyncConfig(), workerJob(1, 3)), ErrLeaseLost)
require.Zero(t, calls)
require.Zero(t, repo.retried)
require.Zero(t, repo.failed)
})
t.Run("start and shutdown", func(t *testing.T) {
cfg := asyncConfig()
cfg.Enabled = false
configStore := &fakeConfigStore{cfg: cfg, active: true}
repo := &fakeJobRepository{}
payload := &fakePayloadStore{pingErr: errors.New("redis unavailable")}
runner := NewRunner(configStore, repo, payload, PromptScannerFunc(func(context.Context, ActiveEndpoint, string, []string) (*NormalizedResult, error) {
return integrationResult(EventPass), nil
}), NewAtomicMetrics())
require.NoError(t, runner.Start(context.Background()))
require.NoError(t, runner.Start(context.Background()))
_, _, _, _, _, code, _ := runner.Snapshot()
require.Equal(t, "payload_store_unavailable", code)
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
require.NoError(t, runner.Shutdown(ctx))
require.NoError(t, runner.Shutdown(ctx))
})
t.Run("shutdown timeout is bounded", func(t *testing.T) {
runner := &Runner{}
release := make(chan struct{})
runner.wg.Add(1)
go func() {
defer runner.wg.Done()
<-release
}()
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
defer cancel()
require.ErrorIs(t, runner.Shutdown(ctx), context.DeadlineExceeded)
close(release)
ctx2, cancel2 := context.WithTimeout(context.Background(), time.Second)
defer cancel2()
require.NoError(t, runner.Shutdown(ctx2))
})
}
func TestPromptAuditSyntheticAsyncBaseline(t *testing.T) {
const totalRequests = 100
cfg := asyncConfig()
cfg.Endpoints[0].InputLimit = 256
cfg.StorePassEvents = false
repo := &fakeJobRepository{}
payload := &fakePayloadStore{values: make(map[int64]string, totalRequests)}
metrics := NewAtomicMetrics()
knownBenignFindings := 0
knownMaliciousBlocked := 0
scanner := PromptScannerFunc(func(_ context.Context, endpoint ActiveEndpoint, chunk string, _ []string) (*NormalizedResult, error) {
switch {
case strings.HasPrefix(chunk, "benign"):
return &NormalizedResult{Decision: EventPass, RiskLevel: RiskLow, Action: ActionAllow, Safety: "Safe", GuardEndpointID: endpoint.ID}, nil
case strings.HasPrefix(chunk, "flag"):
return &NormalizedResult{Decision: EventFlag, RiskLevel: RiskMedium, Action: ActionWarn, Safety: "Controversial", Categories: []string{"politically_sensitive_topics"}, GuardEndpointID: endpoint.ID}, nil
case strings.HasPrefix(chunk, "block"):
knownMaliciousBlocked++
return &NormalizedResult{Decision: EventCritical, RiskLevel: RiskCritical, Action: ActionBlock, Safety: "Unsafe", Categories: []string{"jailbreak"}, GuardEndpointID: endpoint.ID}, nil
case strings.HasPrefix(chunk, "invalid"):
return nil, &GuardError{Code: ErrorCodeInvalidResponse}
default:
return nil, &GuardError{Code: ErrorCodeUnavailable, Retryable: true, Timeout: true}
}
})
runner := NewRunner(&fakeConfigStore{cfg: cfg, active: true}, repo, payload, scanner, metrics)
runner.clock = &advancingClock{now: time.Unix(1_000, 0).UTC(), step: time.Millisecond}
for index := 1; index <= totalRequests; index++ {
text := fmt.Sprintf("benign-%03d", index)
switch {
case index > 90 && index <= 95:
text = fmt.Sprintf("flag-%03d", index)
case index > 95 && index <= 98:
text = fmt.Sprintf("block-%03d", index)
case index == 99:
text = "invalid-099"
case index == 100:
text = "timeout-100"
}
jobID := int64(index)
payload.values[jobID] = text
job := &Job{ID: jobID, ClaimVersion: 1, Attempts: 1, MaxAttempts: 1, ConfigVersion: cfg.ConfigVersion,
Snapshot: PromptSnapshot{RequestID: fmt.Sprintf("baseline-%03d", index), PromptLength: len([]rune(text)), RedactedPreview: "synthetic"}}
err := runner.processJob(context.Background(), 0, cfg, job)
if index <= 98 {
require.NoError(t, err)
} else {
require.Error(t, err)
}
}
snapshot := metrics.Snapshot()
require.Equal(t, int64(totalRequests), snapshot.Total)
require.Equal(t, int64(90), snapshot.Allowed)
require.Equal(t, int64(5), snapshot.Flagged)
require.Equal(t, int64(3), snapshot.Blocked)
require.Equal(t, int64(1), snapshot.Invalid)
require.Equal(t, int64(1), snapshot.Unavailable)
require.Equal(t, int64(1), snapshot.Timeouts)
require.Zero(t, knownBenignFindings)
require.Equal(t, 3, knownMaliciousBlocked)
require.Equal(t, 98, repo.completeCount)
require.Equal(t, 8, repo.eventCount, "store_pass_events=false only grows events for flag/block fixtures")
require.Positive(t, snapshot.LatencyP50MS)
require.LessOrEqual(t, snapshot.LatencyP50MS, snapshot.LatencyP95MS)
require.LessOrEqual(t, snapshot.LatencyP95MS, snapshot.LatencyP99MS)
t.Logf("synthetic async baseline: p50=%dms p95=%dms p99=%dms failure_rate=2%% false_positive_rate=0%% event_growth=8/100", snapshot.LatencyP50MS, snapshot.LatencyP95MS, snapshot.LatencyP99MS)
}
func TestRequestCloneOwnsMutableInputs(t *testing.T) {
groupID := int64(7)
req := Request{Body: []byte("original"), GroupID: &groupID}
clone := req.Clone()
clone.Body[0] = 'X'
*clone.GroupID = 8
require.Equal(t, []byte("original"), req.Body)
require.Equal(t, int64(7), *req.GroupID)
require.False(t, reflect.ValueOf(req.Body).Pointer() == reflect.ValueOf(clone.Body).Pointer())
}