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