prompt_qwen3guard_test.go 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146
  1. package securityaudit
  2. import (
  3. "encoding/json"
  4. "testing"
  5. "time"
  6. "github.com/stretchr/testify/require"
  7. )
  8. func TestParseQwen3GuardStrictAndPolicy(t *testing.T) {
  9. tests := []struct {
  10. name, output string
  11. enabled []string
  12. decision EventDecision
  13. action Action
  14. wantErr bool
  15. }{
  16. {"safe", "Safety: Safe\nCategories: None", AllScannerIDs, EventPass, ActionAllow, false},
  17. {"controversial", "Safety: Controversial\nCategories: Violent", AllScannerIDs, EventFlag, ActionWarn, false},
  18. {"controversial pii escalates", "Safety: Controversial\nCategories: PII", AllScannerIDs, EventCritical, ActionBlock, false},
  19. {"unsafe", "Safety: Unsafe\nCategories: Jailbreak", AllScannerIDs, EventCritical, ActionBlock, false},
  20. {"unknown unsafe", "Safety: Unsafe\nCategories: Future Risk", AllScannerIDs, EventCritical, ActionBlock, false},
  21. {"disabled unsafe warns", "Safety: Unsafe\nCategories: Violent", []string{"PII"}, EventFlag, ActionWarn, false},
  22. {"extra explanation", "Safety: Safe\nCategories: None\nThis is safe", AllScannerIDs, EventPass, ActionAllow, false},
  23. {"duplicate", "Safety: Safe\nSafety: Safe", AllScannerIDs, "", "", true},
  24. {"duplicate categories", "Safety: Safe\nCategories: None\nCategories: PII", AllScannerIDs, "", "", true},
  25. {"missing categories", "Safety: Safe\n", AllScannerIDs, "", "", true},
  26. {"unknown safety", "Safety: Maybe\nCategories: PII", AllScannerIDs, "", "", true},
  27. }
  28. for _, tt := range tests {
  29. t.Run(tt.name, func(t *testing.T) {
  30. result, err := ParseQwen3Guard(tt.output, tt.enabled)
  31. if tt.wantErr {
  32. require.Error(t, err)
  33. return
  34. }
  35. require.NoError(t, err)
  36. require.Equal(t, tt.decision, result.Decision)
  37. require.Equal(t, tt.action, result.Action)
  38. })
  39. }
  40. }
  41. func TestParseQwen3GuardIgnoresAuxiliaryResponseFields(t *testing.T) {
  42. result, err := ParseQwen3Guard("Safety: Unsafe\nCategories: Jailbreak\nRefusal: No", AllScannerIDs)
  43. require.NoError(t, err)
  44. require.Equal(t, "Unsafe", result.Safety)
  45. require.Equal(t, []string{"jailbreak"}, result.Categories)
  46. serialized, err := json.Marshal(result)
  47. require.NoError(t, err)
  48. require.NotContains(t, string(serialized), "Refusal")
  49. require.NotContains(t, string(serialized), "No")
  50. }
  51. func TestQwen3GuardOfficialCategoriesAliasesAndUnknownAreStable(t *testing.T) {
  52. official := "Violent, Non-violent Illegal Acts, Sexual Content or Sexual Acts, PII, Suicide & Self-Harm, Unethical Acts, Politically Sensitive Topics, Copyright Violation, Jailbreak"
  53. result, err := ParseQwen3Guard("Safety: Unsafe\nCategories: "+official, AllScannerIDs)
  54. require.NoError(t, err)
  55. require.Equal(t, AllScannerIDs, result.MatchedScanners)
  56. require.Empty(t, result.UnknownCategories)
  57. require.Equal(t, "priority", result.PolicyID)
  58. require.Equal(t, 1, result.PolicyVersion)
  59. aliases := map[string]string{
  60. "violence": "violent", "non_violent_illegal_acts": "non_violent_illegal_acts",
  61. "sexual": "sexual_content_or_sexual_acts", "personal identifiable information": "pii",
  62. "suicide/self harm": "suicide_and_self_harm", "unethical": "unethical_acts",
  63. "political": "politically_sensitive_topics", "copyright": "copyright_violation",
  64. "prompt injection": "jailbreak",
  65. }
  66. for alias, canonical := range aliases {
  67. require.Equal(t, canonical, NormalizeCategory(alias), alias)
  68. }
  69. const canary = "PROMPT_CANARY_RAW_UNKNOWN_CATEGORY"
  70. unknown, err := ParseQwen3Guard("Safety: Unsafe\nCategories: "+canary, AllScannerIDs)
  71. require.NoError(t, err)
  72. require.Len(t, unknown.UnknownCategories, 1)
  73. require.NotContains(t, unknown.UnknownCategories[0], "canary")
  74. require.NotContains(t, unknown.UnknownCategories[0], "raw")
  75. require.Contains(t, unknown.UnknownCategories[0], "unknown:")
  76. }
  77. func TestExtractOpenAIContentSupportsStringAndTextBlocks(t *testing.T) {
  78. content, err := extractOpenAIContent([]byte(`{"choices":[{"message":{"content":"Safety: Safe\nCategories: None"}}]}`))
  79. require.NoError(t, err)
  80. require.Equal(t, "Safety: Safe\nCategories: None", content)
  81. content, err = extractOpenAIContent([]byte(`{"choices":[{"message":{"content":[{"type":"text","text":"Safety: Safe"},{"type":"text","text":"Categories: None"}]}}]}`))
  82. require.NoError(t, err)
  83. require.Equal(t, "Safety: Safe\nCategories: None", content)
  84. for _, body := range []string{`{}`, `{"choices":[]}`, `{"choices":[{"message":{"content":null}}]}`} {
  85. _, err := extractOpenAIContent([]byte(body))
  86. require.Error(t, err)
  87. }
  88. }
  89. func TestAggregateRequiresEveryResult(t *testing.T) {
  90. _, err := AggregateResults([]*NormalizedResult{{Decision: EventPass, Action: ActionAllow}, nil}, 0)
  91. require.Error(t, err)
  92. result, err := AggregateResults([]*NormalizedResult{
  93. {Decision: EventPass, RiskLevel: RiskLow, Action: ActionAllow, Categories: []string{"pii"}},
  94. {Decision: EventCritical, RiskLevel: RiskCritical, Action: ActionBlock, Categories: []string{"jailbreak"}},
  95. }, 0)
  96. require.NoError(t, err)
  97. require.Equal(t, EventCritical, result.Decision)
  98. require.Equal(t, ActionBlock, result.Action)
  99. require.Equal(t, []string{"pii", "jailbreak"}, result.Categories)
  100. }
  101. func TestAggregateDeduplicatesFactsAndUsesMostSevereEndpointMetadata(t *testing.T) {
  102. result, err := AggregateResults([]*NormalizedResult{
  103. {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},
  104. {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},
  105. }, 7*time.Millisecond)
  106. require.NoError(t, err)
  107. require.Equal(t, []string{"pii", "jailbreak"}, result.Categories)
  108. require.Equal(t, []string{"pii", "jailbreak"}, result.MatchedScanners)
  109. require.Equal(t, "first", result.ScannerEvidence["pii"], "evidence is deterministically first-seen")
  110. require.Equal(t, "block-node", result.GuardEndpointID)
  111. require.Equal(t, "block-version", result.ScannerVersion)
  112. require.Equal(t, 2, result.PolicyVersion)
  113. require.Equal(t, 7, result.LatencyMS)
  114. }
  115. func TestIssueSummariesAreDeterministicRedactedDerivedDTOs(t *testing.T) {
  116. const canary = "PROMPT_CANARY_EVIDENCE_SECRET"
  117. result := NormalizedResult{
  118. Decision: EventCritical, RiskLevel: RiskCritical, Action: ActionBlock,
  119. Categories: []string{"jailbreak", "pii"}, MatchedScanners: []string{"pii"},
  120. ScannerScores: map[string]float64{"pii": 1}, ScannerEvidence: map[string]string{"pii": canary},
  121. UnknownCategories: []string{unknownCategoryID("future risk")},
  122. }
  123. summaries := BuildIssueSummaries(result)
  124. require.Len(t, summaries, 3, "known categories are not hidden merely because policy disabled one")
  125. raw, err := json.Marshal(summaries)
  126. require.NoError(t, err)
  127. require.NotContains(t, string(raw), canary)
  128. for _, summary := range summaries {
  129. require.NotEmpty(t, summary.Title)
  130. require.NotEmpty(t, summary.Description)
  131. require.NotEmpty(t, summary.Code)
  132. require.NotEmpty(t, summary.EvidenceHash)
  133. }
  134. }