From 58ba01c35862012d7dfd1fff454b876a268d0a99 Mon Sep 17 00:00:00 2001
From: Waleed Latif <walif6@gmail.com>
Date: Mon, 2 Nov 2026 12:38:48 -0800
Subject: [PATCH] fix(pii): bound custom-pattern matching time per request

Presidio's REGEX_TIMEOUT_SECONDS bounds one custom pattern on one text, so a
batch request could cost texts x patterns x that timeout (2,000 texts x 20
patterns x 2s) on a single handler thread. Charge every custom-pattern match to
a per-request 60s budget and reject the request with a 400 once it is spent.
---
 apps/pii/server.py                        | 39 ++++++++++++++++++++++-
 apps/sim/lib/guardrails/validate-regex.ts |  6 ++--
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/apps/pii/server.py b/apps/pii/server.py
index e0410c7b1b4..6934daf8355 100644
--- a/apps/pii/server.py
+++ b/apps/pii/server.py
@@ -391,6 +391,41 @@ def custom_operators(patterns: list[CustomPattern] | None) -> dict[str, dict[str
     }
 
 
+# Seconds one request may spend matching its custom patterns, summed over every
+# pattern and text. REGEX_TIMEOUT_SECONDS bounds only one pattern on one text, so
+# without this a batch could cost texts x patterns x that timeout.
+CUSTOM_PATTERN_BUDGET_SECONDS = 60
+
+
+class _MatchBudget:
+    """Custom-pattern matching seconds left for one request, shared by its
+    recognizers. A request is analyzed on a single thread, so no lock is needed."""
+
+    def __init__(self, seconds: float) -> None:
+        self.remaining = seconds
+
+
+class BudgetedPatternRecognizer(PatternRecognizer):
+    """Custom-pattern recognizer that charges its matching time to the request's
+    budget and fails the request with a 400 once the budget is spent."""
+
+    def __init__(self, budget: _MatchBudget, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._budget = budget
+
+    def analyze(self, *args: Any, **kwargs: Any) -> list[RecognizerResult]:
+        started = time.perf_counter()
+        results = super().analyze(*args, **kwargs)
+        self._budget.remaining -= time.perf_counter() - started
+        if self._budget.remaining < 0:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Custom patterns exceeded the {CUSTOM_PATTERN_BUDGET_SECONDS}s "
+                "matching budget for this request",
+            )
+        return results
+
+
 def build_custom_recognizers(
     patterns: list[CustomPattern] | None, language: str
 ) -> tuple[list[PatternRecognizer], list[str]]:
@@ -398,6 +433,7 @@ def build_custom_recognizers(
 
     Each regex is precompiled so a malformed pattern fails fast as a 400 rather
     than surfacing later as an opaque analyze-time 500."""
+    budget = _MatchBudget(CUSTOM_PATTERN_BUDGET_SECONDS)
     recognizers: list[PatternRecognizer] = []
     entity_ids: list[str] = []
     for i, p in enumerate(patterns or []):
@@ -409,7 +445,8 @@ def build_custom_recognizers(
             ) from exc
         entity = f"{CUSTOM_ENTITY_PREFIX}{i}"
         recognizers.append(
-            PatternRecognizer(
+            BudgetedPatternRecognizer(
+                budget,
                 supported_entity=entity,
                 # Score 1.0 so a user's explicit pattern wins any overlap with a
                 # built-in detector (e.g. spaCy tagging "EMP-123456" as ORGANIZATION
diff --git a/apps/sim/lib/guardrails/validate-regex.ts b/apps/sim/lib/guardrails/validate-regex.ts
index f8a90dc6e5e..4da2e099195 100644
--- a/apps/sim/lib/guardrails/validate-regex.ts
+++ b/apps/sim/lib/guardrails/validate-regex.ts
@@ -32,8 +32,10 @@ export interface RegexPatternValidation {
  *
  * Nor could a screen here be made sound: these patterns execute in Presidio's
  * Python engine, which backtracks on shapes RE2 accepts, so RE2-representability
- * says nothing about their runtime there. Presidio's own request timeout is the
- * real bound — note it fails open on timeout, leaving PII unredacted.
+ * says nothing about their runtime there. Presidio bounds them instead: a
+ * per-match timeout, which fails open (the timed-out pattern is skipped, leaving
+ * its PII unredacted), and a per-request matching budget, which rejects the
+ * request.
  *
  * Anything that matches a caller-supplied pattern *in this process* must use
  * `compileLinearRegex` from `@/lib/core/security/linear-regex` instead.