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
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:
@@ -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())
|
||||
}
|
||||
Reference in New Issue
Block a user