From: Stefan Gasser Date: Sun, 26 Jul 2026 07:51:19 +0000 (+0200) Subject: Preserve sensitive fragments around placeholders (#154) X-Git-Tag: v0.8.3~1 X-Git-Url: http://git.99rst.org/?a=commitdiff_plain;h=f194b0f90a7d741bc313e0d7541103a9ddcf9d45;p=sgasser-llm-shield.git Preserve sensitive fragments around placeholders (#154) --- diff --git a/detector/detector/merge.py b/detector/detector/merge.py index 97815e0..b58da78 100644 --- a/detector/detector/merge.py +++ b/detector/detector/merge.py @@ -1,7 +1,8 @@ """Merge the deterministic and fuzzy layers into the final entity list. Rules: - * deterministic spans (score 1.0) always outrank fuzzy spans on overlap; + * deterministic spans (score 1.0) always outrank fuzzy spans on overlap, + while uncovered fuzzy fragments remain eligible; * among fuzzy spans, the longer one wins, then the higher score; * fuzzy spans below `score_threshold` are dropped (deterministic = 1.0 always passes); @@ -15,6 +16,25 @@ from collections.abc import Iterable from .entities import Span, overlaps +def _subtract_overlaps(span: Span, blockers: list[Span]) -> list[Span]: + """Return the fragments of `span` not covered by any blocker.""" + intervals = [(span.start, span.end)] + + for blocker in blockers: + uncovered: list[tuple[int, int]] = [] + for start, end in intervals: + if blocker.end <= start or end <= blocker.start: + uncovered.append((start, end)) + continue + if start < blocker.start: + uncovered.append((start, blocker.start)) + if blocker.end < end: + uncovered.append((blocker.end, end)) + intervals = uncovered + + return [Span(span.entity_type, start, end, span.score) for start, end in intervals] + + def merge( deterministic: list[Span], fuzzy: list[Span], @@ -32,8 +52,13 @@ def merge( # Deterministic spans are pre-resolved (non-overlapping) and take precedence. accepted: list[Span] = [s for s in deterministic if s.score >= score_threshold] + deterministic_blockers = list(accepted) + uncovered_fuzzy = [ + fragment for span in fuzzy for fragment in _subtract_overlaps(span, deterministic_blockers) + ] + # Longer fuzzy spans first, then higher score, for stable overlap resolution. - for span in sorted(fuzzy, key=lambda s: (-s.length, -s.score)): + for span in sorted(uncovered_fuzzy, key=lambda s: (-s.length, -s.score)): if span.score < score_threshold: continue if any(overlaps(span, a) for a in accepted): diff --git a/detector/tests/test_merge.py b/detector/tests/test_merge.py index 4904572..cb41610 100644 --- a/detector/tests/test_merge.py +++ b/detector/tests/test_merge.py @@ -4,11 +4,36 @@ from detector.entities import IBAN_CODE, LOCATION, PERSON, Span from detector.merge import merge -def test_deterministic_precedence_over_overlapping_fuzzy(): - det = [Span(IBAN_CODE, 0, 27, 1.0)] - fuzzy = [Span(LOCATION, 0, 4, 0.9)] # e.g. "DE89" mis-tagged +def test_deterministic_full_overlap_removes_fuzzy(): + det = [Span(IBAN_CODE, 10, 20, 1.0)] + fuzzy = [Span(PERSON, 12, 18, 0.9)] out = merge(det, fuzzy, None, 0.7) - assert out == [Span(IBAN_CODE, 0, 27, 1.0)] + assert out == [Span(IBAN_CODE, 10, 20, 1.0)] + + +def test_deterministic_left_overlap_preserves_fuzzy_left_fragment(): + det = [Span(IBAN_CODE, 10, 20, 1.0)] + fuzzy = [Span(PERSON, 0, 15, 0.9)] + out = merge(det, fuzzy, None, 0.7) + assert out == [Span(PERSON, 0, 10, 0.9), Span(IBAN_CODE, 10, 20, 1.0)] + + +def test_deterministic_right_overlap_preserves_fuzzy_right_fragment(): + det = [Span(IBAN_CODE, 10, 20, 1.0)] + fuzzy = [Span(PERSON, 15, 30, 0.9)] + out = merge(det, fuzzy, None, 0.7) + assert out == [Span(IBAN_CODE, 10, 20, 1.0), Span(PERSON, 20, 30, 0.9)] + + +def test_deterministic_split_overlap_preserves_both_fuzzy_fragments(): + det = [Span(IBAN_CODE, 10, 20, 1.0)] + fuzzy = [Span(PERSON, 0, 30, 0.9)] + out = merge(det, fuzzy, None, 0.7) + assert out == [ + Span(PERSON, 0, 10, 0.9), + Span(IBAN_CODE, 10, 20, 1.0), + Span(PERSON, 20, 30, 0.9), + ] def test_fuzzy_below_threshold_dropped(): diff --git a/src/pii/detect.test.ts b/src/pii/detect.test.ts index 9876245..b49f331 100644 --- a/src/pii/detect.test.ts +++ b/src/pii/detect.test.ts @@ -8,6 +8,7 @@ import { findDenylistedEntities, mergeDenylistEntities, PIIDetector, + subtractPlaceholderOverlaps, } from "./detect"; const originalFetch = globalThis.fetch; @@ -560,6 +561,62 @@ describe("PIIDetector", () => { }); }); + describe("subtractPlaceholderOverlaps", () => { + const text = "left [[TOKEN_1]] right"; + const placeholders = ["[[TOKEN_1]]"]; + + test("removes an entity fully inside a placeholder", () => { + const entities = [{ entity_type: "PERSON", start: 7, end: 14, score: 0.91 }]; + + expect(subtractPlaceholderOverlaps(text, entities, placeholders)).toEqual([]); + }); + + test("leaves non-overlapping entities unchanged", () => { + const entities = [ + { entity_type: "PERSON", start: 0, end: 4, score: 0.91 }, + { entity_type: "LOCATION", start: 17, end: 22, score: 0.82 }, + ]; + + expect(subtractPlaceholderOverlaps(text, entities, placeholders)).toEqual(entities); + }); + + test("keeps the uncovered left fragment of a partial overlap", () => { + const entities = [{ entity_type: "PERSON", start: 0, end: 10, score: 0.91 }]; + + expect(subtractPlaceholderOverlaps(text, entities, placeholders)).toEqual([ + { entity_type: "PERSON", start: 0, end: 5, score: 0.91 }, + ]); + }); + + test("keeps the uncovered right fragment of a partial overlap", () => { + const entities = [{ entity_type: "PERSON", start: 10, end: 22, score: 0.91 }]; + + expect(subtractPlaceholderOverlaps(text, entities, placeholders)).toEqual([ + { entity_type: "PERSON", start: 16, end: 22, score: 0.91 }, + ]); + }); + + test("splits an entity around a placeholder", () => { + const entities = [{ entity_type: "PERSON", start: 0, end: 22, score: 0.91 }]; + + expect(subtractPlaceholderOverlaps(text, entities, placeholders)).toEqual([ + { entity_type: "PERSON", start: 0, end: 5, score: 0.91 }, + { entity_type: "PERSON", start: 16, end: 22, score: 0.91 }, + ]); + }); + + test("subtracts multiple placeholders from one spanning entity", () => { + const multiText = "aa [[ONE_1]] bb [[TWO_2]] cc"; + const entities = [{ entity_type: "LOCATION", start: 1, end: 27, score: 0.83 }]; + + expect(subtractPlaceholderOverlaps(multiText, entities, ["[[ONE_1]]", "[[TWO_2]]"])).toEqual([ + { entity_type: "LOCATION", start: 1, end: 3, score: 0.83 }, + { entity_type: "LOCATION", start: 12, end: 16, score: 0.83 }, + { entity_type: "LOCATION", start: 25, end: 27, score: 0.83 }, + ]); + }); + }); + describe("filterAllowlistedEntities", () => { test("filters entities matching allowlist pattern", () => { const text = "You are Claude Code, Anthropic's official CLI for Claude."; diff --git a/src/pii/detect.ts b/src/pii/detect.ts index f4ccf70..028aa87 100644 --- a/src/pii/detect.ts +++ b/src/pii/detect.ts @@ -59,7 +59,48 @@ function placeholderSpans( index = text.indexOf(placeholder, index + placeholder.length); } } - return spans; + return spans.sort((a, b) => a.start - b.start || a.end - b.end); +} + +/** + * Remove placeholder-covered intervals from detector entities. + * + * Entity offsets stay relative to the original text. A partial overlap can + * therefore produce both a left and a right fragment with the original type + * and score. + */ +export function subtractPlaceholderOverlaps( + text: string, + entities: PIIEntity[], + knownPlaceholders: readonly string[], +): PIIEntity[] { + if (entities.length === 0 || knownPlaceholders.length === 0) return entities; + + const placeholders = placeholderSpans(text, knownPlaceholders); + if (placeholders.length === 0) return entities; + + return entities.flatMap((entity) => { + let fragments = [entity]; + + for (const placeholder of placeholders) { + fragments = fragments.flatMap((fragment) => { + if (!overlaps(fragment, placeholder)) return [fragment]; + + const uncovered: PIIEntity[] = []; + if (fragment.start < placeholder.start) { + uncovered.push({ ...fragment, end: placeholder.start }); + } + if (placeholder.end < fragment.end) { + uncovered.push({ ...fragment, start: placeholder.end }); + } + return uncovered; + }); + + if (fragments.length === 0) break; + } + + return fragments; + }); } export function findDenylistedEntities( @@ -72,12 +113,7 @@ export function findDenylistedEntities( 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))); + return subtractPlaceholderOverlaps(text, matches, knownPlaceholders); } // Additive merge: a denylist match extends coverage but never shrinks an overlapping detector span; overlaps become their union and the detector type wins. @@ -242,7 +278,12 @@ export class PIIDetector { const denylistedEntities = findDenylistedEntities(span.text, denylist, knownPlaceholders); const detectedEntities = config.pii_detection.enabled ? await this.detectPII(span.text) : []; - const filteredEntities = filterAllowlistedEntities(span.text, detectedEntities, allowlist); + const uncoveredEntities = subtractPlaceholderOverlaps( + span.text, + detectedEntities, + knownPlaceholders, + ); + const filteredEntities = filterAllowlistedEntities(span.text, uncoveredEntities, allowlist); spanEntities.push(mergeDenylistEntities(filteredEntities, denylistedEntities)); } diff --git a/src/privacy/pipeline.test.ts b/src/privacy/pipeline.test.ts index 72faae6..427942c 100644 --- a/src/privacy/pipeline.test.ts +++ b/src/privacy/pipeline.test.ts @@ -77,6 +77,44 @@ describe("processPrivacyPipeline", () => { expect(knownPlaceholders).toEqual(["[[API_KEY_SK_1]]"]); }); + test("masks sensitive text beside a secret placeholder without remasking the placeholder", async () => { + const connectionString = "postgres://admin:S3cretPass@db.example.com:5432/appdb"; + const detector = new PIIDetector(); + detector.detectPII = mock(async (text: string) => { + expect(text).toBe("hunter2 [[CONNECTION_STRING_1]]"); + return [{ entity_type: "PERSON", start: 0, end: text.length, score: 0.93 }]; + }); + mockAnalyzeRequest.mockImplementationOnce((maskedRequest, extractor, knownPlaceholders) => + detector.analyzeRequest( + maskedRequest as OpenAIRequest, + extractor as typeof openaiExtractor, + knownPlaceholders, + ), + ); + + const result = await processPrivacyPipeline( + request(`hunter2 ${connectionString}`), + { + ...baseConfig, + secrets_detection: { + ...baseConfig.secrets_detection, + entities: ["CONNECTION_STRING"], + }, + }, + openaiExtractor, + ); + + expect(result.requestAfterSecrets.messages[0].content).toBe("hunter2 [[CONNECTION_STRING_1]]"); + expect(result.piiResult?.detection.spanEntities[0]).toEqual([ + { entity_type: "PERSON", start: 0, end: 8, score: 0.93 }, + ]); + expect(result.request.messages[0].content).toBe("[[PERSON_1]][[CONNECTION_STRING_1]]"); + expect(result.piiMaskingContext?.mapping["[[PERSON_1]]"]).toBe("hunter2 "); + expect(result.secretsResult.maskingContext?.mapping["[[CONNECTION_STRING_1]]"]).toBe( + connectionString, + ); + }); + test("returns privacy facts without route decisions", async () => { const result = await processPrivacyPipeline(request("Hello"), baseConfig, openaiExtractor); diff --git a/src/routes/api.test.ts b/src/routes/api.test.ts index 907028f..54e921f 100644 --- a/src/routes/api.test.ts +++ b/src/routes/api.test.ts @@ -371,6 +371,42 @@ describe("POST /api/mask", () => { } }); + test("masks a detector fragment beside a secret placeholder without corrupting it", async () => { + const connectionString = "postgres://admin:S3cretPass@db.example.com:5432/appdb"; + mockDetectPII.mockImplementationOnce((text: string) => { + expect(text).toBe("hunter2 [[CONNECTION_STRING_1]]"); + return Promise.resolve([ + { + entity_type: "PERSON", + start: 0, + end: text.length, + score: 0.93, + }, + ]); + }); + + const res = await app.request("/api/mask", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ text: `hunter2 ${connectionString}` }), + }); + + expect(res.status).toBe(200); + const body = (await res.json()) as { + masked: string; + context: Record; + entities: { type: string; placeholder: string }[]; + }; + expect(body.masked).toBe("[[PERSON_1]][[CONNECTION_STRING_1]]"); + expect(body.context["[[PERSON_1]]"]).toBe("hunter2 "); + expect(body.context["[[CONNECTION_STRING_1]]"]).toBe(connectionString); + expect(body.entities).toContainEqual({ + type: "CONNECTION_STRING", + placeholder: "[[CONNECTION_STRING_1]]", + }); + expect(body.entities).toContainEqual({ type: "PERSON", placeholder: "[[PERSON_1]]" }); + }); + 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 f7633d6..854b8cb 100644 --- a/src/routes/api.ts +++ b/src/routes/api.ts @@ -15,6 +15,7 @@ import { findDenylistedEntities, getPIIDetector, mergeDenylistEntities, + subtractPlaceholderOverlaps, } from "../pii/detect"; import { mask as maskPII } from "../pii/mask"; import { detectSecrets } from "../secrets/detect"; @@ -189,9 +190,14 @@ async function maskHandler(c: Context, getDetector: DetectorProvider) { const piiEntities = config.pii_detection.enabled ? await detector.detectPII(maskedText) : []; scanTimeMs = Date.now() - piiStartTime; - const filteredEntities = filterAllowlistedEntities( + const uncoveredEntities = subtractPlaceholderOverlaps( maskedText, piiEntities, + Object.keys(context.mapping), + ); + const filteredEntities = filterAllowlistedEntities( + maskedText, + uncoveredEntities, config.masking.allowlist, ); const entitiesToMask = mergeDenylistEntities(