Skip to content

Commit 62ed13b

Browse files
chiruu12chiruu12
committed
Fix safe-prefix cache boundary bypass
1 parent 18dee37 commit 62ed13b

6 files changed

Lines changed: 224 additions & 16 deletions

File tree

‎sdk/src/unplug/api/cache.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,23 +7,29 @@
77
from __future__ import annotations
88

99
from unplug.core.runtime.cache import (
10+
DEFAULT_PREFIX_OVERLAP_CHARS,
1011
CacheKeyParts,
1112
SafePrefixState,
1213
ScanCache,
1314
cache_key_parts,
15+
chunk_storage_key,
16+
effective_prefix_skip,
1417
merge_suffix_result,
1518
offset_findings,
1619
prefix_storage_key,
1720
)
1821
from unplug.core.runtime.versions import MODEL_VERSION_LOCAL, NORMALIZER_VERSION
1922

2023
__all__ = [
24+
"DEFAULT_PREFIX_OVERLAP_CHARS",
2125
"MODEL_VERSION_LOCAL",
2226
"NORMALIZER_VERSION",
2327
"CacheKeyParts",
2428
"SafePrefixState",
2529
"ScanCache",
2630
"cache_key_parts",
31+
"chunk_storage_key",
32+
"effective_prefix_skip",
2733
"merge_suffix_result",
2834
"offset_findings",
2935
"prefix_storage_key",

‎sdk/src/unplug/config/cache.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,9 @@
44

55
from pydantic import BaseModel, Field
66

7+
# Keep in sync with unplug.core.runtime.cache.DEFAULT_PREFIX_OVERLAP_CHARS.
8+
_DEFAULT_PREFIX_OVERLAP_CHARS = 256
9+
710

811
class CacheConfig(BaseModel):
912
"""Safe-prefix and chunk LRU settings."""
@@ -13,3 +16,5 @@ class CacheConfig(BaseModel):
1316
enabled: bool = True
1417
max_chunk_entries: int = Field(default=256, ge=1)
1518
advance_prefix_on_redact: bool = True
19+
# Re-scan this many chars at the safe-prefix boundary (StreamScanner-aligned).
20+
prefix_overlap_chars: int = Field(default=_DEFAULT_PREFIX_OVERLAP_CHARS, ge=0)

‎sdk/src/unplug/core/runtime/cache.py‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,9 @@
1111
from unplug.api.types import Finding, ScanResult
1212
from unplug.core.runtime.versions import MODEL_VERSION_LOCAL, NORMALIZER_VERSION
1313

14+
# Match StreamScanner default so append-only scans re-check the boundary.
15+
DEFAULT_PREFIX_OVERLAP_CHARS = 256
16+
1417

1518
def _sha256(text: str) -> str:
1619
return hashlib.sha256(text.encode("utf-8")).hexdigest()
@@ -39,6 +42,8 @@ class CacheKeyParts(BaseModel):
3942
full_hash: str
4043
normalizer_version: str
4144
model_version: str
45+
source: str = ""
46+
policy_fingerprint: str = ""
4247

4348

4449
def cache_key_parts(
@@ -47,18 +52,40 @@ def cache_key_parts(
4752
document_id: str | None,
4853
normalizer_version: str = NORMALIZER_VERSION,
4954
model_version: str = MODEL_VERSION_LOCAL,
55+
source: str = "",
56+
policy_fingerprint: str = "",
5057
) -> CacheKeyParts:
5158
doc_key = f"doc:{document_id}" if document_id else f"hash:{_sha256(text)[:16]}"
5259
return CacheKeyParts(
5360
doc_key=doc_key,
5461
full_hash=_sha256(text),
5562
normalizer_version=normalizer_version,
5663
model_version=model_version,
64+
source=source,
65+
policy_fingerprint=policy_fingerprint,
5766
)
5867

5968

6069
def prefix_storage_key(parts: CacheKeyParts) -> str:
61-
return f"{parts.doc_key}|{parts.normalizer_version}|{parts.model_version}"
70+
return (
71+
f"{parts.doc_key}|{parts.source}|{parts.policy_fingerprint}|"
72+
f"{parts.normalizer_version}|{parts.model_version}"
73+
)
74+
75+
76+
def chunk_storage_key(parts: CacheKeyParts) -> str:
77+
return (
78+
f"{parts.full_hash}|{parts.source}|{parts.policy_fingerprint}|"
79+
f"{parts.normalizer_version}|{parts.model_version}"
80+
)
81+
82+
83+
def effective_prefix_skip(prefix_len: int, overlap_chars: int) -> int:
84+
"""Return the scan start offset, re-scanning ``overlap_chars`` at the boundary."""
85+
if prefix_len <= 0:
86+
return 0
87+
overlap = max(0, overlap_chars)
88+
return max(0, prefix_len - overlap)
6289

6390

6491
class ScanCache:
@@ -76,17 +103,24 @@ def cache_key_parts(
76103
document_id: str | None,
77104
normalizer_version: str = NORMALIZER_VERSION,
78105
model_version: str = MODEL_VERSION_LOCAL,
106+
source: str = "",
107+
policy_fingerprint: str = "",
79108
) -> CacheKeyParts:
80109
return cache_key_parts(
81110
text,
82111
document_id=document_id,
83112
normalizer_version=normalizer_version,
84113
model_version=model_version,
114+
source=source,
115+
policy_fingerprint=policy_fingerprint,
85116
)
86117

87118
def prefix_storage_key(self, parts: CacheKeyParts) -> str:
88119
return prefix_storage_key(parts)
89120

121+
def chunk_storage_key(self, parts: CacheKeyParts) -> str:
122+
return chunk_storage_key(parts)
123+
90124
def get_safe_prefix(self, parts: CacheKeyParts) -> SafePrefixState | None:
91125
return self._prefixes.get(prefix_storage_key(parts))
92126

@@ -102,6 +136,12 @@ def set_chunk(self, content_hash: str, result: ScanResult) -> None:
102136
while len(self._chunks) > self._max_chunk_entries:
103137
self._chunks.popitem(last=False)
104138

139+
def get_chunk_for_parts(self, parts: CacheKeyParts) -> ScanResult | None:
140+
return self.get_chunk(chunk_storage_key(parts))
141+
142+
def set_chunk_for_parts(self, parts: CacheKeyParts, result: ScanResult) -> None:
143+
self.set_chunk(chunk_storage_key(parts), result)
144+
105145
@staticmethod
106146
def should_advance_prefix(action: Action, *, advance_on_redact: bool) -> bool:
107147
if action == Action.BLOCK:

‎sdk/src/unplug/guard.py‎

Lines changed: 51 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,12 @@
2222
from unplug.core.policy import policy_from_request
2323
from unplug.core.privacy import NullPrivacyFilter, PrivacyFilterService
2424
from unplug.core.privacy.secrets import SecretsRegistry, SecretsSanitizer
25-
from unplug.core.runtime.cache import SafePrefixState, ScanCache, merge_suffix_result
25+
from unplug.core.runtime.cache import (
26+
SafePrefixState,
27+
ScanCache,
28+
effective_prefix_skip,
29+
merge_suffix_result,
30+
)
2631
from unplug.core.runtime.logging import correlation_scope, get_logger
2732
from unplug.core.runtime.model_runtime import (
2833
load_active_model_provider,
@@ -418,6 +423,7 @@ def stream_scanner(
418423
self,
419424
source=source,
420425
scan_every_chars=scan_every_chars,
426+
overlap_chars=self._config.cache.prefix_overlap_chars,
421427
document_id=document_id,
422428
)
423429

@@ -484,6 +490,37 @@ def _request_context(self, request: ScanRequest, *, isolated: bool) -> Execution
484490
def _model_version_for_cache(self) -> str:
485491
return self._model_cache_version
486492

493+
def _cache_policy_fingerprint(self, request: ScanRequest) -> str:
494+
"""Fingerprint policy + scanner set so cache entries are not reused across configs."""
495+
policy = self._config.policy
496+
scanners = request.scanners if request.scanners is not None else list(self._config.scanners)
497+
parts = [
498+
",".join(scanners),
499+
str(
500+
request.block_threshold
501+
if request.block_threshold is not None
502+
else policy.block_threshold
503+
),
504+
str(
505+
request.redact_threshold
506+
if request.redact_threshold is not None
507+
else policy.redact_threshold
508+
),
509+
str(
510+
request.review_threshold
511+
if request.review_threshold is not None
512+
else policy.review_threshold
513+
),
514+
str(
515+
request.block_coverage_ratio
516+
if request.block_coverage_ratio is not None
517+
else policy.block_coverage_ratio
518+
),
519+
str(request.redact),
520+
str(request.redaction_mode.value if request.redaction_mode is not None else ""),
521+
]
522+
return "|".join(parts)
523+
487524
def _run_input_with_cache(self, request: ScanRequest, ctx: ExecutionContext) -> ScanResult:
488525
cache = ctx.scan_cache
489526
if cache is None or not self._config.cache.enabled:
@@ -498,8 +535,10 @@ def _run_input_with_cache(self, request: ScanRequest, ctx: ExecutionContext) ->
498535
document_id=request.document_id,
499536
normalizer_version=NORMALIZER_VERSION,
500537
model_version=self._model_version_for_cache(),
538+
source=str(request.source),
539+
policy_fingerprint=self._cache_policy_fingerprint(request),
501540
)
502-
hit = cache.get_chunk(parts.full_hash)
541+
hit = cache.get_chunk_for_parts(parts)
503542
if hit is not None:
504543
return hit
505544

@@ -508,13 +547,20 @@ def _run_input_with_cache(self, request: ScanRequest, ctx: ExecutionContext) ->
508547
if prefix_state is not None and prefix_state.verify(request.text):
509548
prefix_len = prefix_state.prefix_len
510549

550+
scan_start = 0
511551
if 0 < prefix_len < len(request.text):
552+
scan_start = effective_prefix_skip(
553+
prefix_len,
554+
self._config.cache.prefix_overlap_chars,
555+
)
556+
557+
if 0 < scan_start < len(request.text):
512558
suffix_result = self._input_pipeline.run(
513-
request.text[prefix_len:],
559+
request.text[scan_start:],
514560
source=request.source,
515561
context=ctx,
516562
)
517-
result = merge_suffix_result(suffix_result, prefix_len)
563+
result = merge_suffix_result(suffix_result, scan_start)
518564
else:
519565
result = self._input_pipeline.run(
520566
request.text,
@@ -530,7 +576,7 @@ def _run_input_with_cache(self, request: ScanRequest, ctx: ExecutionContext) ->
530576
parts,
531577
SafePrefixState.from_text(request.text, len(request.text)),
532578
)
533-
cache.set_chunk(parts.full_hash, result)
579+
cache.set_chunk_for_parts(parts, result)
534580
return result
535581

536582
def scan_output(self, text: str | TaintedText) -> ScanResult:

‎sdk/src/unplug/streaming.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,17 @@
77

88
from unplug.api.enums import Action, Source
99
from unplug.api.types import ScanResult
10-
from unplug.core.runtime.cache import merge_suffix_result
10+
from unplug.core.runtime.cache import (
11+
DEFAULT_PREFIX_OVERLAP_CHARS,
12+
effective_prefix_skip,
13+
merge_suffix_result,
14+
)
1115

1216
if TYPE_CHECKING:
1317
from unplug.guard import Guard
1418

1519
_DEFAULT_SCAN_EVERY = 1024
16-
_DEFAULT_OVERLAP = 256
20+
_DEFAULT_OVERLAP = DEFAULT_PREFIX_OVERLAP_CHARS
1721

1822

1923
class StreamScanner:
@@ -73,7 +77,7 @@ def _scan_accumulated(self) -> ScanResult:
7377
and self._last_result.action in (Action.ALLOW, Action.REVIEW)
7478
and self._last_result.safe
7579
):
76-
prefix_len = max(0, self._safe_prefix_len - self._overlap)
80+
prefix_len = effective_prefix_skip(self._safe_prefix_len, self._overlap)
7781

7882
scan_body = body[prefix_len:] if prefix_len else body
7983
request = self._guard._build_scan_request(scan_body, self._source)

0 commit comments

Comments
 (0)