| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176 |
- 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())
- }
- }
|