From: Stefan Gasser Date: Tue, 23 Jun 2026 06:04:47 +0000 (+0200) Subject: Add configurable masking denylist and regex whitelist (#101) X-Git-Tag: v0.6.0~1 X-Git-Url: http://git.99rst.org/?a=commitdiff_plain;h=4e0f56aab51ba42688096886bc6c9713f78de667;p=sgasser-llm-shield.git Add configurable masking denylist and regex whitelist (#101) * Add configurable masking denylist * Align whitelist pattern config * Harden denylist/whitelist masking - Merge denylist matches additively so they never shrink detector coverage - Skip denylist matches inside existing placeholders to avoid corrupting secret/PII masks - Anchor regex whitelist to the full entity so a partial match can't unmask larger PII - Skip detection when PII detection is off and no denylist is configured - Reject regex patterns that match the empty string at config load - Reuse the conflict-resolver overlap helper; update docs and tests * Fix streaming unmask when a placeholder delimiter splits across chunks findPartialPlaceholderStart only buffered when the full "[[" delimiter appeared within one chunk. If a stream chunk ended with a lone "[" (the first half of "[["), it was emitted as safe, so the placeholder was never reassembled and leaked to the client un-restored. Buffer a trailing partial of the start delimiter too. Restores PII and secrets placeholders that the upstream model tokenizes across the "[[" boundary. * Remove ReDoS caveat from PII detection docs * Simplify partial-placeholder detection and cover the closing-delimiter split Replace the single-iteration loop in findPartialPlaceholderStart with a direct trailing-bracket check, and add tests for a placeholder whose closing "]]" is split across stream chunks. * Exclude denylist matches by known placeholders; de-magic the match score Replace the placeholder-shape regex heuristic with exact exclusion against the real placeholders carried over from secrets masking: secretPlaceholders() is threaded through detectPII/analyzeRequest and passed in /api/mask, so a denylist pattern can no longer match (and corrupt) the internals of an existing secret/PII placeholder. Also replace the vestigial denylist match score 2 with a named DENYLIST_MATCH_SCORE constant. * Require knownPlaceholders on detectPII so routes can't skip placeholder protection * Clarify whitelist/denylist docs and drop Claude Code wording from the default --- diff --git a/config.example.yaml b/config.example.yaml index 575cfc1..e4a8cdb 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -52,7 +52,17 @@ masking: # Text patterns that are never masked (protects against false positives) # whitelist: - # - "Company Name Inc." + # - pattern: "Company Name Inc." + # - pattern: 'TEST-\d+' + # regex: true + + # Text or regex patterns that are always masked, even if the detector misses them + # denylist: + # - pattern: "ProjectX" + # type: PROJECT_NAME + # - pattern: 'CUST-\d{6}' + # type: CUSTOMER_ID + # regex: true # PII Detection settings (detector service, /analyze contract) pii_detection: diff --git a/docs/concepts/mask-mode.mdx b/docs/concepts/mask-mode.mdx index 79e8a54..8558b5d 100644 --- a/docs/concepts/mask-mode.mdx +++ b/docs/concepts/mask-mode.mdx @@ -46,12 +46,24 @@ providers: masking: show_markers: false marker_text: "[protected]" + whitelist: + - pattern: "Company Name Inc." + - pattern: 'TEST-\d+' + regex: true + denylist: + - pattern: "ProjectX" + type: PROJECT_NAME + - pattern: 'CUST-\d{6}' + type: CUSTOMER_ID + regex: true ``` | Option | Default | Description | |--------|---------|-------------| | `show_markers` | `false` | Add visual markers around unmasked values | | `marker_text` | `[protected]` | Marker text if enabled | +| `whitelist` | `[]` | Text patterns that are never masked; set `regex: true` for regex patterns | +| `denylist` | `[]` | Text patterns that are always masked with the configured `type`; set `regex: true` for regex patterns | ## Response Headers diff --git a/docs/configuration/pii-detection.mdx b/docs/configuration/pii-detection.mdx index 7961e1a..92cda1e 100644 --- a/docs/configuration/pii-detection.mdx +++ b/docs/configuration/pii-detection.mdx @@ -92,11 +92,29 @@ Exclude specific text patterns from PII masking. Useful for preventing false pos ```yaml masking: whitelist: - - "Acme Corp" - - "Product XYZ" + - pattern: "Acme Corp" + - pattern: "Product XYZ" + - pattern: 'TEST-\d+' + regex: true ``` -Patterns match bidirectionally - detected text containing a whitelist entry (or vice versa) is excluded. +A literal entry is matched as a substring: a detected value is left unmasked if it contains the entry, or the entry contains it. Set `regex: true` for JavaScript regex syntax — a regex entry must match the **entire** detected value, so `\d{4}` won't unmask a longer number that merely contains four digits. + +## Denylist + +Force specific text or regex patterns to be masked, even when the detector does not report them or PII detection is disabled. Each entry needs a `type`, which is used for the placeholder name. + +```yaml +masking: + denylist: + - pattern: "ProjectX" + type: PROJECT_NAME + - pattern: 'CUST-\d{6}' + type: CUSTOMER_ID + regex: true +``` + +Patterns are matched literally by default. Set `regex: true` for JavaScript regex syntax — use single quotes in YAML when the pattern contains backslashes. A regex pattern must not match the empty string (use `\d+`, not `\d*`); empty-matching patterns are rejected at startup because they would mask nothing. ## Scan Roles diff --git a/src/config.test.ts b/src/config.test.ts index 1d53c2c..1b62691 100644 --- a/src/config.test.ts +++ b/src/config.test.ts @@ -55,4 +55,107 @@ pii_detection: cleanupConfig(path); } }); + + test("accepts masking whitelist and denylist patterns", () => { + const path = writeConfig(` +mode: mask +providers: + openai: {} + anthropic: {} +masking: + whitelist: + - "Acme Corp" + - pattern: 'TEST-\\d+' + regex: true + denylist: + - pattern: "ProjectX" + type: PROJECT_NAME + - pattern: 'CUST-\\d{6}' + type: CUSTOMER_ID + regex: true +pii_detection: + detector_url: http://localhost:5002 +`); + + try { + const config = loadConfig(path); + + expect(config.masking.whitelist).toEqual([ + { pattern: "You are Claude Code, Anthropic's official CLI for Claude.", regex: false }, + { pattern: "Acme Corp", regex: false }, + { pattern: "TEST-\\d+", regex: true }, + ]); + expect(config.masking.denylist).toEqual([ + { pattern: "ProjectX", type: "PROJECT_NAME", regex: false }, + { pattern: "CUST-\\d{6}", type: "CUSTOMER_ID", regex: true }, + ]); + } finally { + cleanupConfig(path); + } + }); + + test("rejects invalid masking whitelist regex patterns", () => { + const path = writeConfig(` +mode: mask +providers: + openai: {} + anthropic: {} +masking: + whitelist: + - pattern: "[Acme" + regex: true +pii_detection: + detector_url: http://localhost:5002 +`); + + try { + expect(() => loadConfig(path)).toThrow("Invalid configuration"); + } finally { + cleanupConfig(path); + } + }); + + test("rejects invalid masking denylist regex patterns", () => { + const path = writeConfig(` +mode: mask +providers: + openai: {} + anthropic: {} +masking: + denylist: + - pattern: "[ProjectX" + type: PROJECT_NAME + regex: true +pii_detection: + detector_url: http://localhost:5002 +`); + + try { + expect(() => loadConfig(path)).toThrow("Invalid configuration"); + } finally { + cleanupConfig(path); + } + }); + + test("rejects denylist regex patterns that match the empty string", () => { + const path = writeConfig(` +mode: mask +providers: + openai: {} + anthropic: {} +masking: + denylist: + - pattern: 'x*' + type: NUM + regex: true +pii_detection: + detector_url: http://localhost:5002 +`); + + try { + expect(() => loadConfig(path)).toThrow("Invalid configuration"); + } finally { + cleanupConfig(path); + } + }); }); diff --git a/src/config.ts b/src/config.ts index 19f0642..60bf1c4 100644 --- a/src/config.ts +++ b/src/config.ts @@ -30,15 +30,73 @@ const CodexProviderSchema = z.object({ base_url: z.string().url().default("https://chatgpt.com/backend-api/codex"), }); -const DEFAULT_WHITELIST = ["You are Claude Code, Anthropic's official CLI for Claude."]; +const DEFAULT_WHITELIST = [ + { pattern: "You are Claude Code, Anthropic's official CLI for Claude.", regex: false }, +]; + +function validateRegexPattern( + pattern: string, + regex: boolean, + ctx: z.RefinementCtx, + message: string, +): void { + if (!regex) return; + + let compiled: RegExp; + try { + compiled = new RegExp(pattern, "g"); + } catch { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + path: ["pattern"], + message, + }); + return; + } + + // Reject patterns that match the empty string: zero-length matches are skipped, so they mask nothing. + if (compiled.test("")) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + path: ["pattern"], + message: `${message} (must not match the empty string)`, + }); + } +} + +const WhitelistPatternSchema = z.union([ + z + .string() + .min(1) + .transform((pattern) => ({ pattern, regex: false })), + z + .object({ + pattern: z.string().min(1), + regex: z.boolean().default(false), + }) + .superRefine((entry, ctx) => { + validateRegexPattern(entry.pattern, entry.regex, ctx, "Invalid whitelist regex pattern"); + }), +]); + +const DenylistPatternSchema = z + .object({ + pattern: z.string().min(1), + type: z.string().min(1), + regex: z.boolean().default(false), + }) + .superRefine((entry, ctx) => { + validateRegexPattern(entry.pattern, entry.regex, ctx, "Invalid denylist regex pattern"); + }); const MaskingSchema = z.object({ show_markers: z.boolean().default(false), marker_text: z.string().default("[protected]"), whitelist: z - .array(z.string()) + .array(WhitelistPatternSchema) .default([]) .transform((arr) => [...DEFAULT_WHITELIST, ...arr]), + denylist: z.array(DenylistPatternSchema).default([]), }); const LanguageEnum = z.enum(SUPPORTED_LANGUAGES); @@ -170,6 +228,8 @@ export type AnthropicProviderConfig = z.infer; export type CodexProviderConfig = z.infer; export type LocalProviderConfig = z.infer; export type MaskingConfig = z.infer; +export type WhitelistPattern = z.infer; +export type DenylistPattern = z.infer; export type SecretsDetectionConfig = z.infer; export type ServerConfig = z.infer; diff --git a/src/masking/conflict-resolver.ts b/src/masking/conflict-resolver.ts index 1ae4148..32d4f9a 100644 --- a/src/masking/conflict-resolver.ts +++ b/src/masking/conflict-resolver.ts @@ -17,7 +17,7 @@ export interface EntityWithScore extends Span { entity_type: string; } -function overlaps(a: Span, b: Span): boolean { +export function overlaps(a: Span, b: Span): boolean { return a.start < b.end && b.start < a.end; } diff --git a/src/masking/placeholders.test.ts b/src/masking/placeholders.test.ts index d60a084..5ad5be5 100644 --- a/src/masking/placeholders.test.ts +++ b/src/masking/placeholders.test.ts @@ -83,8 +83,15 @@ describe("findPartialPlaceholderStart", () => { expect(findPartialPlaceholderStart(text)).toBe(6); }); - test("handles text ending with single bracket", () => { - // Single [ is not a placeholder start, so should return -1 - expect(findPartialPlaceholderStart("Hello [")).toBe(-1); + test("buffers a trailing single bracket (first half of a split [[)", () => { + expect(findPartialPlaceholderStart("Hello [")).toBe(6); + }); + + test("buffers a trailing single bracket after a complete placeholder", () => { + expect(findPartialPlaceholderStart("[[PERSON_1]] [")).toBe(13); + }); + + test("buffers an incomplete placeholder whose closing ]] is split", () => { + expect(findPartialPlaceholderStart("Hi [[PERSON_1]")).toBe(3); }); }); diff --git a/src/masking/placeholders.ts b/src/masking/placeholders.ts index 5669436..af1a84a 100644 --- a/src/masking/placeholders.ts +++ b/src/masking/placeholders.ts @@ -35,19 +35,18 @@ export function generateSecretPlaceholder(type: string, count: number): string { * Returns the position where it's safe to split, or -1 if entire string is safe */ export function findPartialPlaceholderStart(text: string): number { - const placeholderStart = text.lastIndexOf(PLACEHOLDER_DELIMITERS.start); + const { start, end } = PLACEHOLDER_DELIMITERS; + const placeholderStart = text.lastIndexOf(start); - if (placeholderStart === -1) { - return -1; // No potential placeholder, entire string is safe + // An opened "[[" with no closing "]]" after it is an incomplete placeholder. + if (placeholderStart !== -1 && !text.slice(placeholderStart).includes(end)) { + return placeholderStart; } - // Check if there's a complete placeholder after the last [[ - const afterStart = text.slice(placeholderStart); - const hasCompletePlaceholder = afterStart.includes(PLACEHOLDER_DELIMITERS.end); - - if (hasCompletePlaceholder) { - return -1; // Placeholder is complete, entire string is safe + // A trailing "[" may be the first half of a "[[" that completes in the next chunk. + if (text.endsWith(start.slice(0, 1))) { + return text.length - 1; } - return placeholderStart; // Return position where partial placeholder starts + return -1; // Entire string is safe to emit } diff --git a/src/pii/detect.test.ts b/src/pii/detect.test.ts index 8629ef5..7fb6e4c 100644 --- a/src/pii/detect.test.ts +++ b/src/pii/detect.test.ts @@ -1,7 +1,13 @@ import { afterEach, describe, expect, mock, test } from "bun:test"; +import { getConfig } from "../config"; import { openaiExtractor } from "../masking/extractors/openai"; import type { OpenAIMessage, OpenAIRequest } from "../providers/openai/types"; -import { filterWhitelistedEntities, PIIDetector } from "./detect"; +import { + filterWhitelistedEntities, + findDenylistedEntities, + mergeDenylistEntities, + PIIDetector, +} from "./detect"; const originalFetch = globalThis.fetch; @@ -164,6 +170,100 @@ describe("PIIDetector", () => { // First message (empty string) has no entities expect(result.spanEntities[0]).toHaveLength(0); }); + + test("adds denylist entities when detector returns none", async () => { + const config = getConfig(); + const previousDenylist = config.masking.denylist; + config.masking.denylist = [{ pattern: "ProjectX", type: "PROJECT_NAME", regex: false }]; + mockDetector({}); + + try { + const detector = new PIIDetector(); + const request = createRequest([{ role: "user", content: "Launch ProjectX" }]); + + const result = await detector.analyzeRequest(request, openaiExtractor); + + expect(result.hasPII).toBe(true); + expect(result.spanEntities[0]).toEqual([ + { entity_type: "PROJECT_NAME", start: 7, end: 15, score: 1 }, + ]); + } finally { + config.masking.denylist = previousDenylist; + } + }); + + test("applies denylist when PII detection is disabled", async () => { + const config = getConfig(); + const previousEnabled = config.pii_detection.enabled; + const previousDenylist = config.masking.denylist; + config.pii_detection.enabled = false; + config.masking.denylist = [{ pattern: "ProjectX", type: "PROJECT_NAME", regex: false }]; + const fetchMock = mock(async () => { + throw new Error("Detector should not be called"); + }); + globalThis.fetch = fetchMock as unknown as typeof fetch; + + try { + const detector = new PIIDetector(); + const request = createRequest([{ role: "user", content: "Launch ProjectX" }]); + + const result = await detector.analyzeRequest(request, openaiExtractor); + + expect(result.hasPII).toBe(true); + expect(fetchMock).not.toHaveBeenCalled(); + } finally { + config.pii_detection.enabled = previousEnabled; + config.masking.denylist = previousDenylist; + } + }); + + test("does not let a denylist substring shrink an overlapping detector entity", async () => { + const config = getConfig(); + const previousDenylist = config.masking.denylist; + config.masking.denylist = [{ pattern: "ProjectX", type: "PROJECT_NAME", regex: false }]; + mockDetector({ + "ProjectX@corp.com": [{ entity_type: "EMAIL_ADDRESS", start: 6, end: 23, score: 0.95 }], + }); + + try { + const detector = new PIIDetector(); + const request = createRequest([{ role: "user", content: "Email ProjectX@corp.com" }]); + + const result = await detector.analyzeRequest(request, openaiExtractor); + + expect(result.spanEntities[0]).toEqual([ + { entity_type: "EMAIL_ADDRESS", start: 6, end: 23, score: 1 }, + ]); + } finally { + config.masking.denylist = previousDenylist; + } + }); + + test("skips detection entirely when disabled and no denylist is configured", async () => { + const config = getConfig(); + const previousEnabled = config.pii_detection.enabled; + const previousDenylist = config.masking.denylist; + config.pii_detection.enabled = false; + config.masking.denylist = []; + const fetchMock = mock(async () => { + throw new Error("Detector should not be called"); + }); + globalThis.fetch = fetchMock as unknown as typeof fetch; + + try { + const detector = new PIIDetector(); + const request = createRequest([{ role: "user", content: "Launch ProjectX" }]); + + const result = await detector.analyzeRequest(request, openaiExtractor); + + expect(result.hasPII).toBe(false); + expect(result.spanEntities).toEqual([]); + expect(fetchMock).not.toHaveBeenCalled(); + } finally { + config.pii_detection.enabled = previousEnabled; + config.masking.denylist = previousDenylist; + } + }); }); describe("detectPII", () => { @@ -215,7 +315,9 @@ describe("PIIDetector", () => { test("filters entities matching whitelist pattern", () => { const text = "You are Claude Code, Anthropic's official CLI for Claude."; const entities = [{ entity_type: "PERSON", start: 8, end: 14, score: 0.9 }]; - const whitelist = ["You are Claude Code, Anthropic's official CLI for Claude."]; + const whitelist = [ + { pattern: "You are Claude Code, Anthropic's official CLI for Claude.", regex: false }, + ]; const result = filterWhitelistedEntities(text, entities, whitelist); @@ -228,7 +330,7 @@ describe("PIIDetector", () => { { entity_type: "PERSON", start: 8, end: 16, score: 0.9 }, { entity_type: "EMAIL_ADDRESS", start: 20, end: 36, score: 0.95 }, ]; - const whitelist = ["Claude"]; + const whitelist = [{ pattern: "Claude", regex: false }]; const result = filterWhitelistedEntities(text, entities, whitelist); @@ -238,7 +340,7 @@ describe("PIIDetector", () => { test("filters when entity text is contained in whitelist pattern", () => { const text = "Hello Claude, how are you?"; const entities = [{ entity_type: "PERSON", start: 6, end: 12, score: 0.85 }]; - const whitelist = ["You are Claude Code"]; + const whitelist = [{ pattern: "You are Claude Code", regex: false }]; const result = filterWhitelistedEntities(text, entities, whitelist); @@ -256,5 +358,122 @@ describe("PIIDetector", () => { expect(result).toHaveLength(2); }); + + test("filters entities matching regex whitelist pattern", () => { + const text = "Reference TEST-1234 is public"; + const entities = [{ entity_type: "CUSTOMER_ID", start: 10, end: 19, score: 0.9 }]; + const whitelist = [{ pattern: "TEST-\\d+", regex: true }]; + + const result = filterWhitelistedEntities(text, entities, whitelist); + + expect(result).toHaveLength(0); + }); + + test("does not filter when a regex whitelist only partially matches the entity", () => { + const text = "card 1234567890123456 end"; + const entities = [{ entity_type: "CREDIT_CARD", start: 5, end: 21, score: 0.99 }]; + const whitelist = [{ pattern: "\\d{4}", regex: true }]; + + const result = filterWhitelistedEntities(text, entities, whitelist); + + expect(result).toHaveLength(1); + }); + }); + + describe("findDenylistedEntities", () => { + test("finds literal denylist patterns", () => { + const result = findDenylistedEntities("ProjectX uses ProjectX-API", [ + { pattern: "ProjectX", type: "PROJECT_NAME", regex: false }, + ]); + + expect(result).toEqual([ + { entity_type: "PROJECT_NAME", start: 0, end: 8, score: 1 }, + { entity_type: "PROJECT_NAME", start: 14, end: 22, score: 1 }, + ]); + }); + + test("finds regex denylist patterns", () => { + const result = findDenylistedEntities("Customers CUST-123456 and CUST-654321", [ + { pattern: "CUST-\\d{6}", type: "CUSTOMER_ID", regex: true }, + ]); + + expect(result).toEqual([ + { entity_type: "CUSTOMER_ID", start: 10, end: 21, score: 1 }, + { entity_type: "CUSTOMER_ID", start: 26, end: 37, score: 1 }, + ]); + }); + + test("matches regex syntax literally unless regex is enabled", () => { + const result = findDenylistedEntities("Internal [ProjectX", [ + { pattern: "[ProjectX", type: "PROJECT_NAME", regex: false }, + ]); + + expect(result).toEqual([{ entity_type: "PROJECT_NAME", start: 9, end: 18, score: 1 }]); + }); + + test("matches regex patterns containing escaped non-syntax characters", () => { + const result = findDenylistedEntities("Customer CUST-123456 onboarded", [ + { pattern: "CUST\\-\\d{6}", type: "CUSTOMER_ID", regex: true }, + ]); + + expect(result).toEqual([{ entity_type: "CUSTOMER_ID", start: 9, end: 20, score: 1 }]); + }); + + test("ignores matches that fall inside a known placeholder", () => { + const result = findDenylistedEntities( + "conn [[CONNECTION_STRING_1]] ProjectX", + [ + { pattern: "\\d+", type: "NUM", regex: true }, + { pattern: "ProjectX", type: "PROJECT_NAME", regex: false }, + ], + ["[[CONNECTION_STRING_1]]"], + ); + + expect(result).toEqual([{ entity_type: "PROJECT_NAME", start: 29, end: 37, score: 1 }]); + }); + }); + + describe("mergeDenylistEntities", () => { + test("returns detector entities unchanged when there is no denylist", () => { + const detected = [{ entity_type: "EMAIL_ADDRESS", start: 0, end: 16, score: 0.9 }]; + + expect(mergeDenylistEntities(detected, [])).toBe(detected); + }); + + test("adds non-overlapping denylist matches", () => { + const detected = [{ entity_type: "EMAIL_ADDRESS", start: 0, end: 5, score: 0.9 }]; + const denylisted = [{ entity_type: "PROJECT_NAME", start: 10, end: 18, score: 1 }]; + + expect(mergeDenylistEntities(detected, denylisted)).toEqual([ + { entity_type: "EMAIL_ADDRESS", start: 0, end: 5, score: 0.9 }, + { entity_type: "PROJECT_NAME", start: 10, end: 18, score: 1 }, + ]); + }); + + test("keeps the full detector span when a denylist match is contained within it", () => { + const detected = [{ entity_type: "EMAIL_ADDRESS", start: 0, end: 17, score: 0.95 }]; + const denylisted = [{ entity_type: "PROJECT_NAME", start: 0, end: 8, score: 1 }]; + + expect(mergeDenylistEntities(detected, denylisted)).toEqual([ + { entity_type: "EMAIL_ADDRESS", start: 0, end: 17, score: 1 }, + ]); + }); + + test("unions a partial overlap so no covered region is left unmasked", () => { + const detected = [{ entity_type: "PERSON", start: 0, end: 10, score: 0.9 }]; + const denylisted = [{ entity_type: "PROJECT_NAME", start: 5, end: 15, score: 1 }]; + + expect(mergeDenylistEntities(detected, denylisted)).toEqual([ + { entity_type: "PERSON", start: 0, end: 15, score: 1 }, + ]); + }); + + test("returns denylist-only matches when the detector found nothing", () => { + const denylisted = [{ entity_type: "PROJECT_NAME", start: 0, end: 8, score: 1 }]; + + expect(mergeDenylistEntities([], denylisted)).toEqual([ + { entity_type: "PROJECT_NAME", start: 0, end: 8, score: 1 }, + ]); + }); }); }); diff --git a/src/pii/detect.ts b/src/pii/detect.ts index edf13e4..ee4410e 100644 --- a/src/pii/detect.ts +++ b/src/pii/detect.ts @@ -1,5 +1,6 @@ -import { getConfig } from "../config"; +import { type DenylistPattern, getConfig, type WhitelistPattern } from "../config"; import { HEALTH_CHECK_TIMEOUT_MS } from "../constants/timeouts"; +import { overlaps, resolveConflicts } from "../masking/conflict-resolver"; import type { RequestExtractor } from "../masking/types"; import { getLanguageDetector, type SupportedLanguage } from "../services/language-detector"; @@ -10,18 +11,121 @@ export interface PIIEntity { score: number; } +// Denylist matches are exact (operator-configured), so they carry full confidence. +const DENYLIST_MATCH_SCORE = 1; + +function findLiteralMatches(text: string, pattern: string, type: string): PIIEntity[] { + const matches: PIIEntity[] = []; + let index = text.indexOf(pattern); + + while (index !== -1) { + matches.push({ + entity_type: type, + start: index, + end: index + pattern.length, + score: DENYLIST_MATCH_SCORE, + }); + index = text.indexOf(pattern, index + pattern.length); + } + + return matches; +} + +function findRegexMatches(text: string, pattern: string, type: string): PIIEntity[] { + const regex = new RegExp(pattern, "g"); + const matches: PIIEntity[] = []; + + for (const match of text.matchAll(regex)) { + if (match.index === undefined || match[0].length === 0) continue; + matches.push({ + entity_type: type, + start: match.index, + end: match.index + match[0].length, + score: DENYLIST_MATCH_SCORE, + }); + } + + return matches; +} + +function placeholderSpans( + text: string, + placeholders: readonly string[], +): Array<{ start: number; end: number }> { + const spans: Array<{ start: number; end: number }> = []; + for (const placeholder of placeholders) { + let index = text.indexOf(placeholder); + while (index !== -1) { + spans.push({ start: index, end: index + placeholder.length }); + index = text.indexOf(placeholder, index + placeholder.length); + } + } + return spans; +} + +export function findDenylistedEntities( + text: string, + denylist: DenylistPattern[], + knownPlaceholders: readonly string[] = [], +): PIIEntity[] { + if (denylist.length === 0 || !text) return []; + + const matches = denylist.flatMap(({ pattern, type, regex }) => + regex ? findRegexMatches(text, pattern, type) : findLiteralMatches(text, pattern, type), + ); + if (matches.length === 0 || knownPlaceholders.length === 0) return matches; + + // Drop matches inside an already-masked placeholder; re-masking its internals would corrupt the earlier mask. + const masked = placeholderSpans(text, knownPlaceholders); + if (masked.length === 0) return matches; + return matches.filter((m) => !masked.some((p) => overlaps(m, p))); +} + +// Additive merge: a denylist match extends coverage but never shrinks an overlapping detector span; overlaps become their union and the detector type wins. +export function mergeDenylistEntities(detected: PIIEntity[], denylisted: PIIEntity[]): PIIEntity[] { + if (denylisted.length === 0) return detected; + + const resolvedDetector = resolveConflicts(detected); + const tagged = [ + ...resolvedDetector.map((e) => ({ e, forced: false })), + ...denylisted.map((e) => ({ e, forced: true })), + ].sort((a, b) => a.e.start - b.e.start); + + const result: { e: PIIEntity; forced: boolean }[] = []; + for (const item of tagged) { + const last = result[result.length - 1]; + if (last && overlaps(item.e, last.e)) { + last.e = { + entity_type: last.forced && !item.forced ? item.e.entity_type : last.e.entity_type, + start: last.e.start, + end: Math.max(last.e.end, item.e.end), + score: Math.max(last.e.score, item.e.score), + }; + last.forced = last.forced && item.forced; + } else { + result.push({ e: { ...item.e }, forced: item.forced }); + } + } + + return result.map((r) => r.e); +} + export function filterWhitelistedEntities( text: string, entities: PIIEntity[], - whitelist: string[], + whitelist: WhitelistPattern[], ): PIIEntity[] { if (whitelist.length === 0) return entities; return entities.filter((entity) => { const detectedText = text.slice(entity.start, entity.end); - return !whitelist.some( - (pattern) => pattern.includes(detectedText) || detectedText.includes(pattern), - ); + return !whitelist.some(({ pattern, regex }) => { + if (regex) { + // Anchor to the whole entity so a partial match can't un-mask a larger detected span. + return new RegExp(`^(?:${pattern})$`).test(detectedText); + } + return pattern.includes(detectedText) || detectedText.includes(pattern); + }); }); } @@ -101,10 +205,23 @@ export class PIIDetector { async analyzeRequest( request: TRequest, extractor: RequestExtractor, + knownPlaceholders: readonly string[] = [], ): Promise { const startTime = Date.now(); const config = getConfig(); + // Pure pass-through: detection off and no denylist, so skip extraction and language detection. + if (!config.pii_detection.enabled && config.masking.denylist.length === 0) { + return { + hasPII: false, + spanEntities: [], + allEntities: [], + scanTimeMs: 0, + language: config.pii_detection.fallback_language, + languageFallback: true, + }; + } + // Extract all text spans from request const spans = extractor.extractTexts(request); @@ -120,15 +237,22 @@ export class PIIDetector { ? new Set(config.pii_detection.scan_roles) : null; const whitelist = config.masking.whitelist; + const denylist = config.masking.denylist; const spanEntities: PIIEntity[][] = await Promise.all( spans.map(async (span) => { + if (!span.text) return []; + const denylistedEntities = findDenylistedEntities(span.text, denylist, knownPlaceholders); + if (scanRoles && span.role && !scanRoles.has(span.role)) { - return []; + return mergeDenylistEntities([], denylistedEntities); } - if (!span.text) return []; - const entities = await this.detectPII(span.text, langResult.language); - return filterWhitelistedEntities(span.text, entities, whitelist); + + const detectedEntities = config.pii_detection.enabled + ? await this.detectPII(span.text, langResult.language) + : []; + const filteredEntities = filterWhitelistedEntities(span.text, detectedEntities, whitelist); + return mergeDenylistEntities(filteredEntities, denylistedEntities); }), ); diff --git a/src/pii/mask.test.ts b/src/pii/mask.test.ts index ff54a2a..da7ceff 100644 --- a/src/pii/mask.test.ts +++ b/src/pii/mask.test.ts @@ -18,12 +18,14 @@ const defaultConfig: MaskingConfig = { show_markers: false, marker_text: "[protected]", whitelist: [], + denylist: [], }; const configWithMarkers: MaskingConfig = { show_markers: true, marker_text: "[protected]", whitelist: [], + denylist: [], }; /** Helper to create a minimal request from messages */ diff --git a/src/providers/anthropic/stream-transformer.test.ts b/src/providers/anthropic/stream-transformer.test.ts index 84d430d..28bfe16 100644 --- a/src/providers/anthropic/stream-transformer.test.ts +++ b/src/providers/anthropic/stream-transformer.test.ts @@ -7,6 +7,7 @@ const defaultConfig: MaskingConfig = { show_markers: false, marker_text: "[protected]", whitelist: [], + denylist: [], }; /** diff --git a/src/providers/openai/stream-transformer.test.ts b/src/providers/openai/stream-transformer.test.ts index c586147..d4dcd73 100644 --- a/src/providers/openai/stream-transformer.test.ts +++ b/src/providers/openai/stream-transformer.test.ts @@ -7,6 +7,7 @@ const defaultConfig: MaskingConfig = { show_markers: false, marker_text: "[protected]", whitelist: [], + denylist: [], }; /** @@ -101,6 +102,40 @@ describe("createUnmaskingStream", () => { expect(result).toContain("a@b.com"); }); + test("buffers a placeholder split between the two opening brackets", async () => { + const context = createMaskingContext(); + context.mapping["[[EMAIL_ADDRESS_1]]"] = "a@b.com"; + + const chunks = [ + `data: {"choices":[{"delta":{"content":"Hello ["}}]}\n\n`, + `data: {"choices":[{"delta":{"content":"[EMAIL_ADDRESS_1]] world"}}]}\n\n`, + ]; + const source = createSSEStream(chunks); + + const unmaskedStream = createUnmaskingStream(source, context, defaultConfig); + const result = await consumeStream(unmaskedStream); + + expect(result).toContain("a@b.com"); + expect(result).not.toContain("[[EMAIL_ADDRESS_1]]"); + }); + + test("buffers a placeholder split between the two closing brackets", async () => { + const context = createMaskingContext(); + context.mapping["[[EMAIL_ADDRESS_1]]"] = "a@b.com"; + + const chunks = [ + `data: {"choices":[{"delta":{"content":"Hello [[EMAIL_ADDRESS_1]"}}]}\n\n`, + `data: {"choices":[{"delta":{"content":"] world"}}]}\n\n`, + ]; + const source = createSSEStream(chunks); + + const unmaskedStream = createUnmaskingStream(source, context, defaultConfig); + const result = await consumeStream(unmaskedStream); + + expect(result).toContain("a@b.com"); + expect(result).not.toContain("[[EMAIL_ADDRESS_1]"); + }); + test("flushes remaining buffer on stream end", async () => { const context = createMaskingContext(); context.mapping["[[EMAIL_ADDRESS_1]]"] = "test@test.com"; diff --git a/src/routes/anthropic.ts b/src/routes/anthropic.ts index 23da843..4cbaa76 100644 --- a/src/routes/anthropic.ts +++ b/src/routes/anthropic.ts @@ -31,7 +31,11 @@ import { callLocalAnthropic } from "../providers/local"; import { unmaskSecretsResponse } from "../secrets/mask"; import { logRequest } from "../services/logger"; import { detectPII, maskPII, type PIIDetectResult } from "../services/pii"; -import { processSecretsRequest, type SecretsProcessResult } from "../services/secrets"; +import { + processSecretsRequest, + type SecretsProcessResult, + secretPlaceholders, +} from "../services/secrets"; import { createLogData, errorFormats, @@ -110,27 +114,13 @@ anthropicRoutes.post( request = secretsResult.request; } - // Step 2: Detect PII (skip if disabled) + // Step 2: Detect PII and configured denylist terms let piiResult: PIIDetectResult; - if (!config.pii_detection.enabled) { - piiResult = { - detection: { - hasPII: false, - spanEntities: [], - allEntities: [], - scanTimeMs: 0, - language: "en", - languageFallback: false, - }, - hasPII: false, - }; - } else { - try { - piiResult = await detectPII(request, anthropicExtractor); - } catch (error) { - console.error("PII detection error:", error); - return respondDetectionError(c, request, secretsResult, startTime); - } + try { + piiResult = await detectPII(request, anthropicExtractor, secretPlaceholders(secretsResult)); + } catch (error) { + console.error("PII detection error:", error); + return respondDetectionError(c, request, secretsResult, startTime); } // Step 3: Route mode - send to local if PII or secrets detected diff --git a/src/routes/api.test.ts b/src/routes/api.test.ts index d053e9d..ffdc738 100644 --- a/src/routes/api.test.ts +++ b/src/routes/api.test.ts @@ -1,6 +1,11 @@ import { describe, expect, mock, test } from "bun:test"; import { Hono } from "hono"; -import { filterWhitelistedEntities, type PIIEntity } from "../pii/detect"; +import { + filterWhitelistedEntities, + findDenylistedEntities, + mergeDenylistEntities, + type PIIEntity, +} from "../pii/detect"; // Mock the PII detector to avoid needing the detector running const mockDetectPII = mock<(text: string, language: string) => Promise>(() => @@ -12,6 +17,8 @@ mock.module("../pii/detect", () => ({ healthCheck: mock(() => Promise.resolve(true)), }), filterWhitelistedEntities, + findDenylistedEntities, + mergeDenylistEntities, })); // Mock the logger to avoid database operations @@ -24,6 +31,8 @@ const realConfig = await import("../config"); const baseConfig = realConfig.getConfig(); const testConfig = { ...baseConfig, + // Pin detection on so the detector mock is consumed regardless of config.yaml. + pii_detection: { ...baseConfig.pii_detection, enabled: true }, secrets_detection: { ...baseConfig.secrets_detection, enabled: true, @@ -133,6 +142,69 @@ describe("POST /api/mask", () => { expect(body.entities[0].type).toBe("EMAIL_ADDRESS"); }); + test("masks configured denylist patterns", async () => { + const previousDenylist = testConfig.masking.denylist; + testConfig.masking.denylist = [ + { pattern: "ProjectX", type: "PROJECT_NAME", regex: false }, + { pattern: "CUST-\\d{6}", type: "CUSTOMER_ID", regex: true }, + ]; + mockDetectPII.mockResolvedValueOnce([]); + + try { + const res = await app.request("/api/mask", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ text: "ProjectX customer CUST-123456" }), + }); + + expect(res.status).toBe(200); + const body = (await res.json()) as { + masked: string; + context: Record; + entities: { type: string; placeholder: string }[]; + }; + expect(body.masked).toBe("[[PROJECT_NAME_1]] customer [[CUSTOMER_ID_1]]"); + expect(body.context["[[PROJECT_NAME_1]]"]).toBe("ProjectX"); + expect(body.context["[[CUSTOMER_ID_1]]"]).toBe("CUST-123456"); + expect(body.entities).toEqual([ + { type: "PROJECT_NAME", placeholder: "[[PROJECT_NAME_1]]" }, + { type: "CUSTOMER_ID", placeholder: "[[CUSTOMER_ID_1]]" }, + ]); + } finally { + testConfig.masking.denylist = previousDenylist; + } + }); + + test("denylist match inside a detected entity does not leak the rest of it", async () => { + const previousDenylist = testConfig.masking.denylist; + testConfig.masking.denylist = [{ pattern: "ProjectX", type: "PROJECT_NAME", regex: false }]; + mockDetectPII.mockResolvedValueOnce([ + { entity_type: "EMAIL_ADDRESS", start: 6, end: 23, score: 0.95 }, + ]); + + try { + const res = await app.request("/api/mask", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ text: "Email ProjectX@corp.com" }), + }); + + expect(res.status).toBe(200); + const body = (await res.json()) as { + masked: string; + context: Record; + entities: { type: string; placeholder: string }[]; + }; + expect(body.masked).toBe("Email [[EMAIL_ADDRESS_1]]"); + expect(body.context["[[EMAIL_ADDRESS_1]]"]).toBe("ProjectX@corp.com"); + expect(body.entities).toEqual([ + { type: "EMAIL_ADDRESS", placeholder: "[[EMAIL_ADDRESS_1]]" }, + ]); + } finally { + testConfig.masking.denylist = previousDenylist; + } + }); + test("respects startFrom counters", async () => { mockDetectPII.mockResolvedValueOnce([ { entity_type: "EMAIL_ADDRESS", start: 0, end: 16, score: 0.9 }, @@ -282,6 +354,36 @@ describe("POST /api/mask", () => { expect(body.entities.some((e) => e.type === "EMAIL_ADDRESS")).toBe(false); }); + test("denylist does not corrupt an existing secret placeholder", async () => { + const previousDenylist = testConfig.masking.denylist; + testConfig.masking.denylist = [{ pattern: "\\d+", type: "NUM", regex: true }]; + mockDetectPII.mockResolvedValueOnce([]); + + try { + const res = await app.request("/api/mask", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + text: "Connection: postgres://admin:S3cretPass@db.example.com:5432/appdb", + }), + }); + + expect(res.status).toBe(200); + const body = (await res.json()) as { + masked: string; + context: Record; + entities: { type: string }[]; + }; + expect(body.masked).toContain("[[CONNECTION_STRING_1]]"); + expect(body.masked).not.toContain("[[NUM"); + expect(body.masked).not.toContain("STRING_[["); + expect(body.entities.some((e) => e.type === "NUM")).toBe(false); + expect(body.context["[[CONNECTION_STRING_1]]"]).toContain("postgres://"); + } finally { + testConfig.masking.denylist = previousDenylist; + } + }); + test("returns 400 for malformed JSON", async () => { const res = await app.request("/api/mask", { method: "POST", diff --git a/src/routes/api.ts b/src/routes/api.ts index 5806bca..8c00f33 100644 --- a/src/routes/api.ts +++ b/src/routes/api.ts @@ -9,7 +9,12 @@ import { Hono } from "hono"; import { z } from "zod"; import { getConfig, type SecretsDetectionConfig } from "../config"; import { createPlaceholderContext, type PlaceholderContext } from "../masking/context"; -import { filterWhitelistedEntities, getPIIDetector } from "../pii/detect"; +import { + filterWhitelistedEntities, + findDenylistedEntities, + getPIIDetector, + mergeDenylistEntities, +} from "../pii/detect"; import { mask as maskPII } from "../pii/mask"; import { detectSecrets } from "../secrets/detect"; import { maskSecrets } from "../secrets/mask"; @@ -198,7 +203,9 @@ apiRoutes.post("/mask", async (c) => { try { const piiStartTime = Date.now(); const detector = getPIIDetector(); - const piiEntities = await detector.detectPII(maskedText, language); + const piiEntities = config.pii_detection.enabled + ? await detector.detectPII(maskedText, language) + : []; scanTimeMs = Date.now() - piiStartTime; // Apply whitelist filtering @@ -207,15 +214,19 @@ apiRoutes.post("/mask", async (c) => { piiEntities, config.masking.whitelist, ); + const entitiesToMask = mergeDenylistEntities( + filteredEntities, + findDenylistedEntities(maskedText, config.masking.denylist, Object.keys(context.mapping)), + ); // Capture counters before masking to track new entities const countersBefore = { ...context.counters }; - const piiResult = maskPII(maskedText, filteredEntities, context); + const piiResult = maskPII(maskedText, entitiesToMask, context); maskedText = piiResult.masked; allEntities.push(...extractEntities(countersBefore, piiResult.context)); // Collect unique entity types for logging - for (const entity of filteredEntities) { + for (const entity of entitiesToMask) { if (!piiEntityTypes.includes(entity.entity_type)) { piiEntityTypes.push(entity.entity_type); } diff --git a/src/routes/codex.ts b/src/routes/codex.ts index be85074..dae9585 100644 --- a/src/routes/codex.ts +++ b/src/routes/codex.ts @@ -23,7 +23,11 @@ import { } from "../secrets/mask"; import { logRequest } from "../services/logger"; import { detectPII, maskPII, type PIIDetectResult } from "../services/pii"; -import { processSecretsRequest, type SecretsProcessResult } from "../services/secrets"; +import { + processSecretsRequest, + type SecretsProcessResult, + secretPlaceholders, +} from "../services/secrets"; import { createLogData, errorFormats, @@ -81,25 +85,11 @@ codexRoutes.post( } let piiResult: PIIDetectResult; - if (!config.pii_detection.enabled) { - piiResult = { - detection: { - hasPII: false, - spanEntities: [], - allEntities: [], - scanTimeMs: 0, - language: config.pii_detection.fallback_language, - languageFallback: false, - }, - hasPII: false, - }; - } else { - try { - piiResult = await detectPII(request, codexExtractor); - } catch (error) { - console.error("PII detection error:", error); - return respondDetectionError(c, request, startTime); - } + try { + piiResult = await detectPII(request, codexExtractor, secretPlaceholders(secretsResult)); + } catch (error) { + console.error("PII detection error:", error); + return respondDetectionError(c, request, startTime); } const shouldBlockRouteMode = diff --git a/src/routes/openai.ts b/src/routes/openai.ts index 51ca97d..1f66250 100644 --- a/src/routes/openai.ts +++ b/src/routes/openai.ts @@ -31,7 +31,11 @@ import { import { unmaskSecretsResponse } from "../secrets/mask"; import { logRequest } from "../services/logger"; import { detectPII, maskPII, type PIIDetectResult } from "../services/pii"; -import { processSecretsRequest, type SecretsProcessResult } from "../services/secrets"; +import { + processSecretsRequest, + type SecretsProcessResult, + secretPlaceholders, +} from "../services/secrets"; import { extractTextContent } from "../utils/content"; import { createLogData, @@ -80,27 +84,13 @@ openaiRoutes.post( request = secretsResult.request; } - // Step 2: Detect PII (skip if disabled) + // Step 2: Detect PII and configured denylist terms let piiResult: PIIDetectResult; - if (!config.pii_detection.enabled) { - piiResult = { - detection: { - hasPII: false, - spanEntities: [], - allEntities: [], - scanTimeMs: 0, - language: "en", - languageFallback: false, - }, - hasPII: false, - }; - } else { - try { - piiResult = await detectPII(request, openaiExtractor); - } catch (error) { - console.error("PII detection error:", error); - return respondDetectionError(c, request, startTime); - } + try { + piiResult = await detectPII(request, openaiExtractor, secretPlaceholders(secretsResult)); + } catch (error) { + console.error("PII detection error:", error); + return respondDetectionError(c, request, startTime); } // Step 3: Process based on mode diff --git a/src/services/pii.ts b/src/services/pii.ts index 46e9382..d042e55 100644 --- a/src/services/pii.ts +++ b/src/services/pii.ts @@ -23,9 +23,11 @@ export interface PIIMaskResult { export async function detectPII( request: TRequest, extractor: RequestExtractor, + // Required (no default) so a new route can't silently skip the secrets placeholders. + knownPlaceholders: readonly string[], ): Promise { const detector = getPIIDetector(); - const detection = await detector.analyzeRequest(request, extractor); + const detection = await detector.analyzeRequest(request, extractor, knownPlaceholders); return { detection, diff --git a/src/services/secrets.ts b/src/services/secrets.ts index ca44eac..6b0fa2d 100644 --- a/src/services/secrets.ts +++ b/src/services/secrets.ts @@ -18,6 +18,11 @@ export interface SecretsProcessResult { masked: boolean; } +/** Placeholder strings already inserted by secrets masking (so later PII/denylist passes skip them). */ +export function secretPlaceholders(result: SecretsProcessResult): string[] { + return result.maskingContext ? Object.keys(result.maskingContext.mapping) : []; +} + /** * Process a request for secrets detection */