From: Stefan Gasser <redacted>
Date: Wed, 21 Jan 2026 17:42:58 +0000 (+0100)
Subject: Reuse filterWhitelistedEntities and add combined masking test
X-Git-Url: http://git.99rst.org/?a=commitdiff_plain;h=61b2080018112118aded7a320f2bad67a4067488;p=sgasser-llm-shield.git

Reuse filterWhitelistedEntities and add combined masking test

- Import filterWhitelistedEntities from pii/detect instead of duplicating logic
- Add test for combined PII + secrets masking (default behavior)
---

diff --git a/src/routes/api.test.ts b/src/routes/api.test.ts
index 0ad1117..6d6d7d0 100644
--- a/src/routes/api.test.ts
+++ b/src/routes/api.test.ts
@@ -1,6 +1,6 @@
 import { describe, expect, mock, test } from "bun:test";
 import { Hono } from "hono";
-import type { PIIEntity } from "../pii/detect";
+import { filterWhitelistedEntities, type PIIEntity } from "../pii/detect";
 
 // Mock the PII detector to avoid needing Presidio running
 const mockDetectPII = mock<(text: string, language: string) => Promise<PIIEntity[]>>(() =>
@@ -10,6 +10,7 @@ mock.module("../pii/detect", () => ({
   getPIIDetector: () => ({
     detectPII: mockDetectPII,
   }),
+  filterWhitelistedEntities,
 }));
 
 // Mock the logger to avoid database operations
@@ -185,4 +186,37 @@ describe("POST /api/mask", () => {
     expect(body.counters.PERSON).toBe(1);
     expect(body.counters.EMAIL_ADDRESS).toBe(1);
   });
+
+  test("masks both PII and secrets in single request (default behavior)", async () => {
+    mockDetectPII.mockResolvedValueOnce([
+      { entity_type: "EMAIL_ADDRESS", start: 8, end: 24, score: 0.9 },
+    ]);
+
+    const res = await app.request("/api/mask", {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify({
+        text: "Contact john@example.com with key -----BEGIN RSA PRIVATE KEY-----\nMIIE...\n-----END RSA PRIVATE KEY-----",
+      }),
+    });
+
+    expect(res.status).toBe(200);
+    const body = (await res.json()) as {
+      masked: string;
+      context: Record<string, string>;
+      entities: { type: string; placeholder: string }[];
+    };
+
+    // Both PII and secrets should be masked
+    expect(body.masked).toContain("[[EMAIL_ADDRESS_1]]");
+    expect(body.masked).toContain("[[SECRET_MASKED_PEM_PRIVATE_KEY_1]]");
+
+    // Context should contain mappings for both
+    expect(body.context["[[EMAIL_ADDRESS_1]]"]).toBe("john@example.com");
+    expect(body.context["[[SECRET_MASKED_PEM_PRIVATE_KEY_1]]"]).toBeDefined();
+
+    // Entities should include both types
+    expect(body.entities.some((e) => e.type === "EMAIL_ADDRESS")).toBe(true);
+    expect(body.entities.some((e) => e.type === "PEM_PRIVATE_KEY")).toBe(true);
+  });
 });
diff --git a/src/routes/api.ts b/src/routes/api.ts
index 79adbb0..46b382d 100644
--- a/src/routes/api.ts
+++ b/src/routes/api.ts
@@ -9,7 +9,7 @@ import { Hono } from "hono";
 import { z } from "zod";
 import { getConfig, type SecretsDetectionConfig } from "../config";
 import { createPlaceholderContext, type PlaceholderContext } from "../masking/context";
-import { getPIIDetector } from "../pii/detect";
+import { filterWhitelistedEntities, getPIIDetector } from "../pii/detect";
 import { mask as maskPII } from "../pii/mask";
 import { detectSecrets } from "../secrets/detect";
 import { maskSecrets } from "../secrets/mask";
@@ -139,13 +139,11 @@ apiRoutes.post("/mask", async (c) => {
       scanTimeMs = Date.now() - piiStartTime;
 
       // Apply whitelist filtering
-      const whitelist = config.masking.whitelist;
-      const filteredEntities = piiEntities.filter((entity) => {
-        const detectedText = maskedText.slice(entity.start, entity.end);
-        return !whitelist.some(
-          (pattern) => pattern.includes(detectedText) || detectedText.includes(pattern),
-        );
-      });
+      const filteredEntities = filterWhitelistedEntities(
+        maskedText,
+        piiEntities,
+        config.masking.whitelist,
+      );
 
       // Capture counters before masking to track new entities
       const countersBefore = { ...context.counters };