coordinator_test.go 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176
  1. package securityaudit
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "net/http"
  7. "sync/atomic"
  8. "testing"
  9. "github.com/stretchr/testify/require"
  10. )
  11. type fakeLegacyEngine struct {
  12. decision *LegacyDecision
  13. err error
  14. calls atomic.Int64
  15. }
  16. func (f *fakeLegacyEngine) Check(context.Context, Request) (*LegacyDecision, error) {
  17. f.calls.Add(1)
  18. return f.decision, f.err
  19. }
  20. type fakePromptEngine struct {
  21. mode Mode
  22. decision *PromptDecision
  23. err error
  24. enqueues atomic.Int64
  25. evaluates atomic.Int64
  26. }
  27. func (f *fakePromptEngine) EffectiveMode() Mode { return f.mode }
  28. func (f *fakePromptEngine) Enqueue(context.Context, Request) error {
  29. f.enqueues.Add(1)
  30. return f.err
  31. }
  32. func (f *fakePromptEngine) Evaluate(context.Context, Request) (*PromptDecision, error) {
  33. f.evaluates.Add(1)
  34. return f.decision, f.err
  35. }
  36. func TestCoordinatorModesAndPriority(t *testing.T) {
  37. tests := []struct {
  38. name string
  39. mode Mode
  40. legacy *LegacyDecision
  41. prompt *PromptDecision
  42. promptErr error
  43. wantKind DecisionKind
  44. wantCode string
  45. wantEnqueue int64
  46. wantEvaluation int64
  47. }{
  48. {name: "off", mode: ModeOff, wantKind: DecisionAllow},
  49. {name: "async only enqueues", mode: ModeAsync, wantKind: DecisionAllow, wantEnqueue: 1},
  50. {name: "prompt block", mode: ModeBlocking, prompt: &PromptDecision{Kind: DecisionBlock}, wantKind: DecisionBlock, wantCode: ErrorCodeBlocked, wantEvaluation: 1},
  51. {name: "prompt unavailable", mode: ModeBlocking, promptErr: errors.New("down"), wantKind: DecisionUnavailable, wantCode: ErrorCodeUnavailable, wantEvaluation: 1},
  52. {name: "legacy wins both block", mode: ModeBlocking,
  53. legacy: &LegacyDecision{Blocked: true, StatusCode: http.StatusForbidden, ErrorCode: "content_policy_violation", Message: "legacy"},
  54. prompt: &PromptDecision{Kind: DecisionBlock}, wantKind: DecisionBlock, wantCode: "content_policy_violation", wantEvaluation: 1},
  55. }
  56. for _, tt := range tests {
  57. t.Run(tt.name, func(t *testing.T) {
  58. legacy := &fakeLegacyEngine{decision: tt.legacy}
  59. prompt := &fakePromptEngine{mode: tt.mode, decision: tt.prompt, err: tt.promptErr}
  60. decision := NewCoordinator(legacy, prompt).Check(context.Background(), Request{Body: []byte(`{}`)})
  61. require.Equal(t, tt.wantKind, decision.Kind)
  62. require.Equal(t, tt.wantCode, decision.ErrorCode)
  63. require.Equal(t, int64(1), legacy.calls.Load())
  64. require.Equal(t, tt.wantEnqueue, prompt.enqueues.Load())
  65. require.Equal(t, tt.wantEvaluation, prompt.evaluates.Load())
  66. })
  67. }
  68. }
  69. func TestCoordinatorDoesNotMutateRequestBody(t *testing.T) {
  70. body := []byte(`{"messages":[{"role":"user","content":"hello"}]}`)
  71. original := append([]byte(nil), body...)
  72. prompt := &fakePromptEngine{mode: ModeAsync}
  73. decision := NewCoordinator(&fakeLegacyEngine{}, prompt).Check(context.Background(), Request{Body: body})
  74. require.True(t, decision.AllowNextStage)
  75. require.Equal(t, original, body)
  76. }
  77. func TestCoordinatorBlockingPriorityCoversBothEngineDecisionMatrix(t *testing.T) {
  78. legacyCases := []struct {
  79. name string
  80. decision *LegacyDecision
  81. }{
  82. {name: "allow", decision: &LegacyDecision{Allowed: true, StatusCode: http.StatusOK, Action: "allow"}},
  83. {name: "flag", decision: &LegacyDecision{Allowed: true, Flagged: true, StatusCode: http.StatusOK, Action: "flag"}},
  84. {name: "block", decision: &LegacyDecision{Blocked: true, StatusCode: http.StatusForbidden, ErrorCode: "legacy_exact_code", Message: "legacy exact message", Action: "block"}},
  85. }
  86. promptCases := []struct {
  87. name string
  88. decision *PromptDecision
  89. wantKind DecisionKind
  90. wantCode string
  91. }{
  92. {name: "allow", decision: &PromptDecision{Kind: DecisionAllow, AllowNextStage: true}, wantKind: DecisionAllow},
  93. {name: "flag", decision: &PromptDecision{Kind: DecisionFlag, AllowNextStage: true}, wantKind: DecisionFlag},
  94. {name: "block", decision: &PromptDecision{Kind: DecisionBlock}, wantKind: DecisionBlock, wantCode: ErrorCodeBlocked},
  95. {name: "unavailable", decision: &PromptDecision{Kind: DecisionUnavailable, ErrorCode: ErrorCodeUnavailable}, wantKind: DecisionUnavailable, wantCode: ErrorCodeUnavailable},
  96. {name: "invalid", decision: &PromptDecision{Kind: DecisionInvalid, ErrorCode: ErrorCodeInvalidResponse}, wantKind: DecisionInvalid, wantCode: ErrorCodeInvalidResponse},
  97. }
  98. for _, legacyCase := range legacyCases {
  99. for _, promptCase := range promptCases {
  100. t.Run(fmt.Sprintf("legacy_%s_prompt_%s", legacyCase.name, promptCase.name), func(t *testing.T) {
  101. legacy := &fakeLegacyEngine{decision: legacyCase.decision}
  102. prompt := &fakePromptEngine{mode: ModeBlocking, decision: promptCase.decision}
  103. decision := NewCoordinator(legacy, prompt).Check(context.Background(), Request{})
  104. require.Same(t, legacyCase.decision, decision.Legacy)
  105. require.Same(t, promptCase.decision, decision.Prompt)
  106. require.Equal(t, int64(1), legacy.calls.Load())
  107. require.Equal(t, int64(1), prompt.evaluates.Load())
  108. if legacyCase.name == "block" {
  109. require.Equal(t, DecisionBlock, decision.Kind)
  110. require.Equal(t, "legacy_exact_code", decision.ErrorCode)
  111. require.Equal(t, "legacy exact message", decision.ClientMessage)
  112. require.False(t, decision.AllowNextStage)
  113. return
  114. }
  115. require.Equal(t, promptCase.wantKind, decision.Kind)
  116. require.Equal(t, promptCase.wantCode, decision.ErrorCode)
  117. require.Equal(t, promptCase.decision.AllowNextStage, decision.AllowNextStage)
  118. })
  119. }
  120. }
  121. }
  122. func TestCoordinatorPreservesIndependentEngineFactsAndMapsOnlyGatewayOutcome(t *testing.T) {
  123. legacyDecision := &LegacyDecision{
  124. Allowed: true, Flagged: true, Message: "legacy finding", StatusCode: http.StatusAccepted,
  125. ErrorCode: "legacy_observation", Action: "legacy_action",
  126. }
  127. promptResult := &NormalizedResult{
  128. Decision: EventCritical, RiskLevel: RiskCritical, Action: ActionBlock,
  129. Categories: []string{"pii"}, ScannerScores: map[string]float64{"pii": 1},
  130. }
  131. promptDecision := &PromptDecision{Kind: DecisionBlock, Result: promptResult}
  132. decision := NewCoordinator(
  133. &fakeLegacyEngine{decision: legacyDecision},
  134. &fakePromptEngine{mode: ModeBlocking, decision: promptDecision},
  135. ).Check(context.Background(), Request{})
  136. require.Same(t, legacyDecision, decision.Legacy)
  137. require.Same(t, promptDecision, decision.Prompt)
  138. require.Same(t, promptResult, decision.Prompt.Result)
  139. require.Equal(t, "legacy finding", decision.Legacy.Message)
  140. require.Equal(t, []string{"pii"}, decision.Prompt.Result.Categories)
  141. require.Equal(t, ErrorCodeBlocked, decision.ErrorCode)
  142. }
  143. func TestCoordinatorAsyncEnqueueFailuresNeverChangeResponseOrDownstreamDispatch(t *testing.T) {
  144. for _, enqueueErr := range []error{ErrQueueFull, ErrQueueAdmissionBusy, errors.New("redis unavailable"), errors.New("publish failed")} {
  145. prompt := &fakePromptEngine{mode: ModeAsync, err: enqueueErr}
  146. decision := NewCoordinator(&fakeLegacyEngine{decision: &LegacyDecision{Allowed: true}}, prompt).Check(context.Background(), Request{})
  147. downstreamDispatches := 0
  148. status := http.StatusOK
  149. responseBody := "unchanged-upstream-response"
  150. if decision.AllowNextStage {
  151. downstreamDispatches++
  152. } else {
  153. status = decision.HTTPStatus
  154. responseBody = decision.ClientMessage
  155. }
  156. require.Equal(t, http.StatusOK, status)
  157. require.Equal(t, "unchanged-upstream-response", responseBody)
  158. require.Equal(t, 1, downstreamDispatches)
  159. require.Equal(t, int64(1), prompt.enqueues.Load())
  160. require.Zero(t, prompt.evaluates.Load())
  161. }
  162. }