From 4b094b572d6ae59776c4a3480212d7abf535a7c5 Mon Sep 17 00:00:00 2001 From: Vederis L Date: Fri, 2 Oct 2026 14:48:13 +0800 Subject: [PATCH 1/4] fix: ONNX detector offsets, visible ML failures, lock-free /healthz - The ONNX detector split words on a lowercased copy and returned those byte offsets for the original text, the same bug fixed earlier in the dictionary: characters whose case mapping changes length shifted every later span. Words are now located in the original text and lowercased only for the tokenizer. - Spans from any extra detector are widened to character boundaries before use, so a misaligned span redacts too much instead of panicking or too little. - Inference failures were swallowed silently, turning ML recall off unnoticed. They are now logged by stage (never text) and counted in portcullis_detector_errors_total. - /healthz read the term count under the gateway lock and queued behind scans; it now reads a lock-free gauge. - Split a six-element tuple into NerInput (clippy type_complexity under onnx). Co-Authored-By: Claude Sonnet 5.5 --- src/detect.rs | 7 +++ src/detect_onnx.rs | 128 +++++++++++++++++++++++++++++++++------------ src/metrics.rs | 23 ++++++++ src/pipeline.rs | 48 ++++++++++++++++- src/proxy/mod.rs | 7 ++- tests/hardening.rs | 22 ++++++++ 6 files changed, 198 insertions(+), 37 deletions(-) diff --git a/src/detect.rs b/src/detect.rs index 851868a..fea65b9 100644 --- a/src/detect.rs +++ b/src/detect.rs @@ -21,6 +21,13 @@ pub struct Span { /// tasks behind an `Arc>`. pub trait Detector: Send + Sync { fn detect(&self, text: &str) -> Vec; + + /// Inference failures since the last call, for detectors that can fail (the ML + /// one). The gateway turns this into a metric, so a detector that has quietly + /// stopped working is visible. + fn take_failures(&self) -> u64 { + 0 + } } /// Deterministic dictionary layer — the **guarantee**. Runs first, always. diff --git a/src/detect_onnx.rs b/src/detect_onnx.rs index 0b2b6ef..78361e1 100644 --- a/src/detect_onnx.rs +++ b/src/detect_onnx.rs @@ -10,7 +10,8 @@ use std::collections::HashMap; use std::path::{Path, PathBuf}; -use std::sync::Mutex; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Mutex, OnceLock}; use anyhow::{Context, Result}; use ndarray::{Array2, Array3, Axis, Ix2, Ix3}; @@ -87,6 +88,8 @@ pub struct OnnxDetector { max_width: usize, labels: Vec, threshold: f32, + /// Segments where inference failed since the last [`Detector::take_failures`]. + failures: AtomicU64, } impl OnnxDetector { @@ -137,6 +140,7 @@ impl OnnxDetector { max_width: config.max_width, labels: resolve_labels(), threshold: resolve_threshold(), + failures: AtomicU64::new(0), }) } @@ -188,25 +192,52 @@ impl Detector for OnnxDetector { fn detect(&self, text: &str) -> Vec { self.detect_with_threshold(text, self.threshold) } + + fn take_failures(&self) -> u64 { + self.failures.swap(0, Ordering::Relaxed) + } +} + +/// Everything the encoder needs for one text. +struct NerInput { + input_ids: Array2, + attention_mask: Array2, + e_positions: Vec, + word_offsets: Vec<(usize, usize)>, + text_start_idx: usize, + first_token_positions: Vec, } impl OnnxDetector { + /// Inference failed for this text. The segment still has dictionary and regex + /// coverage, but ML recall is silently gone for it — so say so. Only the stage is + /// logged, never the text or an error message that might quote it. + fn fail(&self, stage: &str, e: &anyhow::Error) -> Vec { + self.failures.fetch_add(1, Ordering::Relaxed); + tracing::warn!( + stage, + cause = %e.root_cause().to_string().chars().take(120).collect::(), + "ONNX detection failed; segment has dictionary and regex coverage only" + ); + Vec::new() + } + /// Run detection with a configurable threshold. pub fn detect_with_threshold(&self, text: &str, threshold: f32) -> Vec { if text.trim().is_empty() { return Vec::new(); } - let ( + let NerInput { input_ids, attention_mask, e_positions, word_offsets, text_start_idx, first_token_positions, - ) = match self.build_ner_input(text) { + } = match self.build_ner_input(text) { Ok(x) => x, - Err(_) => return Vec::new(), + Err(e) => return self.fail("tokenize", &e), }; let num_words = word_offsets.len(); @@ -217,7 +248,7 @@ impl OnnxDetector { // hidden: [1, seq_len, hidden] let hidden = match run_encoder(&self.encoder, input_ids, attention_mask) { Ok(h) => h, - Err(_) => return Vec::new(), + Err(e) => return self.fail("encoder", &e), }; let dim = hidden.shape()[2]; @@ -257,12 +288,12 @@ impl OnnxDetector { &token_span_end, ) { Ok(s) => s, - Err(_) => return Vec::new(), + Err(e) => return self.fail("span_rep", &e), }; let scores = match compute_scores(&self.count_embed, &span_rep, &label_embeddings) { Ok(s) => s, - Err(_) => return Vec::new(), + Err(e) => return self.fail("scores", &e), }; let mut entities: Vec = Vec::new(); @@ -297,33 +328,20 @@ impl OnnxDetector { } /// Build the NER token sequence and the word -> first-token map. - fn build_ner_input( - &self, - text: &str, - ) -> Result<( - Array2, - Array2, - Vec, - Vec<(usize, usize)>, - usize, - Vec, - )> { + fn build_ner_input(&self, text: &str) -> Result { let (mut tokens, e_positions) = self.build_schema_prefix(NER_TASK_NAME, &self.labels, TOKEN_E)?; let text_start_idx = tokens.len(); - let lower = text.to_lowercase(); - let word_re = word_pattern(); - let mut word_offsets: Vec<(usize, usize)> = Vec::new(); let mut first_token_positions: Vec = Vec::new(); let mut token_idx = 0usize; - for m in word_re.find_iter(&lower) { - word_offsets.push((m.start(), m.end())); + for (start, end) in word_spans(text) { + word_offsets.push((start, end)); first_token_positions.push(token_idx); - let ids = encode_token(&self.tokenizer, m.as_str())?; + let ids = encode_token(&self.tokenizer, &text[start..end].to_lowercase())?; for &id in &ids { tokens.push(id as i64); } @@ -334,14 +352,14 @@ impl OnnxDetector { let input_ids = Array2::from_shape_vec((1, seq_len), tokens)?; let attention_mask = Array2::::ones((1, seq_len)); - Ok(( + Ok(NerInput { input_ids, attention_mask, e_positions, word_offsets, text_start_idx, first_token_positions, - )) + }) } /// Build `( [P] task ( [E] label1 [E] label2 ... ) ) [SEP_TEXT]`. @@ -412,20 +430,36 @@ impl OnnxDetector { fn encode_token(tokenizer: &Tokenizer, text: &str) -> Result> { tokenizer .encode(text, false) - .map_err(|e| anyhow::anyhow!("tokenization failed for {text:?}: {e}")) + .map_err(|e| anyhow::anyhow!("tokenization failed: {e}")) .map(|enc| enc.get_ids().to_vec()) } -fn word_pattern() -> Regex { - Regex::new( - r"(?ix) +fn word_pattern() -> &'static Regex { + static RE: OnceLock = OnceLock::new(); + RE.get_or_init(|| { + Regex::new( + r"(?ix) (?:https?://[^\s]+|www\.[^\s]+) | [a-z0-9._%+-]+@[a-z0-9.-]+\.[a-z]{2,} | @[a-z0-9_]+ | \w+(?:[-_]\w+)* | \S", - ) - .expect("word regex is valid") + ) + .expect("word regex is valid") + }) +} + +/// Byte ranges of the words in `text`, **in `text`'s own coordinates**. +/// +/// The pattern is case-insensitive, so there is no need to lowercase first — and +/// lowercasing a copy changes byte lengths for some characters (`İ`), which would shift +/// every later offset and make the spans point at the wrong bytes of the original. +/// Callers lowercase each word only for the tokenizer. +fn word_spans(text: &str) -> Vec<(usize, usize)> { + word_pattern() + .find_iter(text) + .map(|m| (m.start(), m.end())) + .collect() } fn generate_spans(num_words: usize, max_width: usize) -> (Vec, Vec) { @@ -546,3 +580,33 @@ fn sigmoid(mut x: Array2) -> Array2 { }); x } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn word_spans_index_the_original_text_even_after_length_changing_case_maps() { + // `İ` lowercases to two code points (3 bytes), so offsets taken from a + // lowercased copy would drift by one byte per `İ`. + let text = "İİİ met Daniel Pratt"; + let words: Vec<&str> = word_spans(text).iter().map(|&(s, e)| &text[s..e]).collect(); + assert_eq!(words, ["İİİ", "met", "Daniel", "Pratt"]); + } + + #[test] + fn word_spans_keep_emails_and_urls_whole() { + let text = "mail dana@example.com or see https://example.com/x"; + let words: Vec<&str> = word_spans(text).iter().map(|&(s, e)| &text[s..e]).collect(); + assert_eq!( + words, + [ + "mail", + "dana@example.com", + "or", + "see", + "https://example.com/x" + ] + ); + } +} diff --git a/src/metrics.rs b/src/metrics.rs index 853ffa6..fc64a44 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -86,6 +86,7 @@ pub struct Metrics { cache_hits: AtomicU64, cache_misses: AtomicU64, store_terms: AtomicU64, + detector_errors: AtomicU64, } impl Default for Metrics { @@ -119,6 +120,7 @@ impl Metrics { cache_hits: AtomicU64::new(0), cache_misses: AtomicU64::new(0), store_terms: AtomicU64::new(0), + detector_errors: AtomicU64::new(0), } } @@ -173,6 +175,16 @@ impl Metrics { g.scan_seconds.observe(scan_seconds); } + /// The ML detector failed on `n` segments (they kept dictionary and regex coverage). + pub fn detector_errors(&self, n: u64) { + self.detector_errors.fetch_add(n, Ordering::Relaxed); + } + + /// Taught terms, for `/healthz` (read lock-free, so a probe never queues behind a scan). + pub fn store_terms(&self) -> u64 { + self.store_terms.load(Ordering::Relaxed) + } + pub fn set_store_terms(&self, n: usize) { self.store_terms.store(n as u64, Ordering::Relaxed); } @@ -284,6 +296,17 @@ impl Metrics { self.cache_misses.load(Ordering::Relaxed) ); + let _ = writeln!( + o, + "# HELP portcullis_detector_errors_total Segments where the ML detector failed and only dictionary and regex coverage applied." + ); + let _ = writeln!(o, "# TYPE portcullis_detector_errors_total counter"); + let _ = writeln!( + o, + "portcullis_detector_errors_total {}", + self.detector_errors.load(Ordering::Relaxed) + ); + for (name, help, h) in [ ( "portcullis_scan_duration_seconds", diff --git a/src/pipeline.rs b/src/pipeline.rs index 57f6e4d..0369bae 100644 --- a/src/pipeline.rs +++ b/src/pipeline.rs @@ -8,6 +8,21 @@ use crate::metrics::Metrics; use crate::store::Store; use crate::vault::Vault; +/// Widen `start..end` outward to character boundaries and clamp to the segment. +/// `None` only for an empty or out-of-range span. Widening, never dropping, is the +/// safe direction: it can only hide more. +fn snap_to_boundaries(seg: &str, start: usize, end: usize) -> Option<(usize, usize)> { + let mut start = start.min(seg.len()); + let mut end = end.min(seg.len()); + while !seg.is_char_boundary(start) { + start -= 1; + } + while !seg.is_char_boundary(end) { + end += 1; + } + (start < end).then_some((start, end)) +} + /// Delta-cache size at which it is flushed, so a long-lived proxy cannot grow it /// without bound. const CACHE_LIMIT: usize = 20_000; @@ -213,7 +228,17 @@ impl Gateway { let mut spans = dict.detect(seg); spans.extend(self.regexes.detect(seg)); if let Some(extra) = &self.extra { - spans.extend(extra.detect(seg)); + // A third-party detector's offsets are not trusted to land on character + // boundaries: slicing mid-character would panic, and a misaligned span + // would redact the wrong bytes. Widen to the enclosing boundaries. + spans.extend(extra.detect(seg).into_iter().filter_map(|mut s| { + (s.start, s.end) = snap_to_boundaries(seg, s.start, s.end)?; + Some(s) + })); + let failed = extra.take_failures(); + if failed > 0 { + self.metrics.detector_errors(failed); + } } let mut kept = Vec::new(); for s in merge(spans) { @@ -346,6 +371,27 @@ impl Gateway { mod tests { use super::*; + struct Misaligned; + impl Detector for Misaligned { + // Byte 1 is inside the 2-byte "é"; byte 4 is inside the 3-byte "€". + fn detect(&self, _text: &str) -> Vec { + vec![crate::detect::Span { + start: 1, + end: 4, + label: "PERSON".into(), + source: "onnx", + score: 0.9, + }] + } + } + + #[test] + fn a_misaligned_detector_span_is_widened_not_panicked_on_or_dropped() { + let mut gw = Gateway::new(Store::default()).with_detector(Box::new(Misaligned)); + let out = gw.process(&["é€x".to_string()]); + assert_eq!(out[0], "<>x", "{:?}", out[0]); + } + #[test] fn suggestions_are_deduplicated_and_bounded() { let mut gw = Gateway::new(Store::default()); diff --git a/src/proxy/mod.rs b/src/proxy/mod.rs index 905df14..5be932a 100644 --- a/src/proxy/mod.rs +++ b/src/proxy/mod.rs @@ -267,10 +267,9 @@ async fn metrics_handler(State(state): State, headers: HeaderMap) -> /// Liveness and readiness probe. Unauthenticated, and deliberately reports only /// counts — never a term, a value or a key. async fn healthz(State(state): State) -> Json { - let store_terms = { - let gw = state.gateway.lock().await; - gw.store.deny.len() - }; + // Lock-free on purpose: scans serialize and can take seconds, and a liveness probe + // that queues behind one gets a healthy process restarted. + let store_terms = state.metrics.store_terms(); Json(json!({ "status": "ok", "version": env!("CARGO_PKG_VERSION"), diff --git a/tests/hardening.rs b/tests/hardening.rs index 702f99d..13ef174 100644 --- a/tests/hardening.rs +++ b/tests/hardening.rs @@ -66,6 +66,28 @@ async fn healthz_is_available_and_leaks_nothing() { ); } +/// A liveness probe must not queue behind a long scan: hold the gateway lock (as a slow +/// detection does) and `/healthz` must still answer promptly. +#[tokio::test] +async fn healthz_does_not_wait_for_the_gateway_lock() { + let state = proxy_state( + Duration::from_secs(5), + "http://127.0.0.1:1/v1/chat/completions".into(), + ); + let gateway = state.gateway.clone(); + let base = spawn(app(state)).await; + + let _held = gateway.lock().await; + let resp = tokio::time::timeout( + Duration::from_secs(2), + reqwest::get(format!("{base}/healthz")), + ) + .await + .expect("/healthz blocked behind the gateway lock") + .unwrap(); + assert_eq!(resp.status(), 200); +} + /// An oversized request body is refused before it can reach the detector. #[tokio::test] async fn oversized_body_is_rejected() { From c09ddfef285acc30ef603135890dfc321c1246e6 Mon Sep 17 00:00:00 2001 From: Vederis L Date: Fri, 2 Oct 2026 14:49:06 +0800 Subject: [PATCH 2/4] perf: cache the compiled dictionary per scope, validated by a deny-list fingerprint Co-Authored-By: Claude Sonnet 5.5 --- src/pipeline.rs | 55 ++++++++++++++++++++++++++++++++++++++++----- tests/invariants.rs | 18 +++++++++++++++ 2 files changed, 67 insertions(+), 6 deletions(-) diff --git a/src/pipeline.rs b/src/pipeline.rs index 0369bae..5e60633 100644 --- a/src/pipeline.rs +++ b/src/pipeline.rs @@ -1,5 +1,6 @@ use std::collections::{HashMap, VecDeque}; -use std::sync::Arc; +use std::hash::{Hash, Hasher}; +use std::sync::{Arc, Mutex}; use sha2::{Digest, Sha256}; @@ -8,6 +9,22 @@ use crate::metrics::Metrics; use crate::store::Store; use crate::vault::Vault; +/// `scope -> (deny-list fingerprint, compiled dictionary)`. +type DictCache = HashMap, (u64, Arc)>; + +/// Distinct scopes whose compiled dictionaries are kept (the scope set is configured, so +/// this is only a backstop). +const DICT_CACHE_LIMIT: usize = 64; + +/// A hash of everything in the deny-list that affects matching. +fn deny_fingerprint(store: &Store) -> u64 { + let mut h = std::collections::hash_map::DefaultHasher::new(); + for e in &store.deny { + (&e.term, &e.label, &e.scope, &e.aliases, e.whole_word).hash(&mut h); + } + h.finish() +} + /// Widen `start..end` outward to character boundaries and clamp to the segment. /// `None` only for an empty or out-of-range span. Widening, never dropping, is the /// safe direction: it can only hide more. @@ -51,6 +68,9 @@ pub struct Gateway { /// Recent spans redacted by a non-dictionary detector (regex, ONNX, …). suggestions: VecDeque<(Option, String)>, metrics: Arc, + /// Compiled dictionaries by scope, each tagged with the fingerprint of the deny-list + /// it was built from. + dicts: Mutex, } impl Gateway { @@ -65,6 +85,7 @@ impl Gateway { extra: None, suggestions: VecDeque::new(), metrics, + dicts: Mutex::new(HashMap::new()), } } @@ -131,6 +152,29 @@ impl Gateway { self.metrics.clone() } + /// The compiled dictionary for `scope`, reused across requests. + /// + /// Compiling it is the expensive part of a dictionary scan (one regex over every + /// term), so it is cached. The cache is validated against a fingerprint of the + /// deny-list on every call, not invalidated by hooks: `store` is a public field and + /// can be edited directly, and a stale dictionary would silently miss a new term. + fn dictionary(&self, scope: Option<&str>) -> Arc { + let fingerprint = deny_fingerprint(&self.store); + let key = scope.map(str::to_string); + let mut cache = self.dicts.lock().unwrap_or_else(|p| p.into_inner()); + if let Some((fp, dict)) = cache.get(&key) { + if *fp == fingerprint { + return dict.clone(); + } + } + if cache.len() >= DICT_CACHE_LIMIT { + cache.clear(); + } + let dict = Arc::new(DictionaryDetector::for_scope(&self.store, scope)); + cache.insert(key, (fingerprint, dict.clone())); + dict + } + pub fn vault(&self) -> &Vault { &self.vault } @@ -298,7 +342,7 @@ impl Gateway { self.cache.clear(); } // Compiled once per call, and only if some segment misses the cache. - let mut dict: Option = None; + let mut dict: Option> = None; let (mut hits, mut misses) = (0u64, 0u64); let mut out = Vec::with_capacity(segments.len()); for seg in segments { @@ -310,9 +354,8 @@ impl Gateway { } None => { misses += 1; - let dict = dict - .get_or_insert_with(|| DictionaryDetector::for_scope(&self.store, scope)); - let spans = self.detect_spans(scope, dict, seg); + let dict = dict.get_or_insert_with(|| self.dictionary(scope)).clone(); + let spans = self.detect_spans(scope, &dict, seg); self.cache.insert(h, spans.clone()); spans } @@ -347,7 +390,7 @@ impl Gateway { scope: Option<&str>, outbound: &[String], ) -> Result<(), String> { - let dict = DictionaryDetector::for_scope(&self.store, scope); + let dict = self.dictionary(scope); for seg in outbound { for span in dict.detect(seg) { if !self.store.is_allowed_for(scope, &seg[span.start..span.end]) { diff --git a/tests/invariants.rs b/tests/invariants.rs index 4b269ef..aef1718 100644 --- a/tests/invariants.rs +++ b/tests/invariants.rs @@ -150,3 +150,21 @@ fn longest_form_wins_at_the_same_position() { let out = gw.process(&["About Project Loki today.".to_string()]); assert!(out[0].contains("<>"), "{}", out[0]); } + +/// The compiled-dictionary cache must never serve a stale list: `Gateway::store` is a +/// public field, and a term added by editing it directly has to take effect at once. +#[test] +fn directly_edited_store_is_never_served_from_a_stale_dictionary() { + let mut gw = gw_with(&[("Cartalian", "ORG")]); + let first = gw.process(&["Loki and Cartalian".to_string()]); + assert!(first[0].contains("Loki"), "{}", first[0]); + + gw.store.teach("Loki", "ORG", "global"); // bypasses Gateway::teach on purpose + let leaky = vec!["Loki is here".to_string()]; + assert!( + gw.assert_clean(&leaky).is_err(), + "assert_clean used a stale dictionary" + ); + let mut fresh = gw_with(&[("Cartalian", "ORG"), ("Loki", "ORG")]); + assert!(!fresh.process(&leaky)[0].contains("Loki")); +} From a937402aa56db4480f43b67204dd19b66c4d0f59 Mon Sep 17 00:00:00 2001 From: Vederis L Date: Fri, 2 Oct 2026 14:50:00 +0800 Subject: [PATCH 3/4] fix: HTTP teach re-reads the store file and saves off the gateway lock teach/unteach overwrote the file from the server's in-memory copy, silently dropping a term added by the CLI in the meantime. They now re-read the file (when it exists), apply the change, and save, serialized by a dedicated writer lock. The read and save run on the blocking pool, so Argon2 under encryption no longer stalls requests behind the gateway lock. An unreadable file fails the request rather than being overwritten. Co-Authored-By: Claude Sonnet 5.5 --- docs/TEACHING.md | 9 +++++ src/pipeline.rs | 8 +++++ src/proxy/admin.rs | 84 +++++++++++++++++++++++++++++++++++++++------- src/proxy/mod.rs | 4 +++ src/store.rs | 2 +- tests/m2_teach.rs | 40 ++++++++++++++++++++++ 6 files changed, 133 insertions(+), 14 deletions(-) diff --git a/docs/TEACHING.md b/docs/TEACHING.md index ba17c64..ef85cd8 100644 --- a/docs/TEACHING.md +++ b/docs/TEACHING.md @@ -208,3 +208,12 @@ Being precise here matters more than looking complete. The honest summary: **the guarantee layer works today; the convenience layers are roadmap.** + +## Teaching while the server runs + +Prefer teaching over HTTP (`/teach`) while `serve` is running: it takes effect on the next +request. The CLI (`portcullis teach`) edits the file directly, and a running server does **not** +re-read it — the term applies after a restart. It is no longer *lost*, though: every HTTP +`teach`/`unteach` re-reads the store file first and writes back on top of it, so a term added +by the CLI survives (and from that point the server has it too). The file read and write run +off the gateway lock, so a teach — even with encryption on — does not stall other requests. diff --git a/src/pipeline.rs b/src/pipeline.rs index 5e60633..493f30e 100644 --- a/src/pipeline.rs +++ b/src/pipeline.rs @@ -175,6 +175,14 @@ impl Gateway { dict } + /// Replace the whole store (e.g. with a fresh read from disk). A policy change, so + /// the delta cache is cleared (invariant 3). + pub fn replace_store(&mut self, store: Store) { + self.store = store; + self.cache.clear(); + self.metrics.set_store_terms(self.store.deny.len()); + } + pub fn vault(&self) -> &Vault { &self.vault } diff --git a/src/proxy/admin.rs b/src/proxy/admin.rs index e20eb6c..da2c4fa 100644 --- a/src/proxy/admin.rs +++ b/src/proxy/admin.rs @@ -95,6 +95,56 @@ fn check_admin(state: &ProxyState, headers: &HeaderMap) -> Result<(), (StatusCod Ok(()) } +/// Apply `op` to the gateway and persist the result. +/// +/// * **Read-modify-write against the file.** If the store file exists it is re-read +/// first and becomes the starting point, so a term added out of band (`portcullis +/// teach` while serving) is not silently overwritten by this write. +/// * **Disk work off the gateway lock.** Reading and saving can run Argon2 (64 MiB) when +/// encryption is on; doing that under the gateway lock would stall every request, so +/// it runs on the blocking pool and the lock is held only for the in-memory change. +/// * Writers are serialized by `store_write_lock`, so two teaches cannot interleave. +/// +/// An unreadable store file (wrong key, corruption) fails the request instead of being +/// overwritten. +async fn mutate_store( + state: &ProxyState, + op: impl FnOnce(&mut crate::Gateway) -> T, +) -> Result { + let internal = |e: String| (StatusCode::INTERNAL_SERVER_ERROR, e); + let _writer = state.store_write_lock.lock().await; + let path = store_path(state); + + let on_disk = { + let path = path.clone(); + tokio::task::spawn_blocking(move || { + if std::path::Path::new(&path).exists() { + crate::Store::load(&path).map(Some) + } else { + Ok(None) + } + }) + .await + .map_err(|e| internal(e.to_string()))? + .map_err(|e| internal(e.to_string()))? + }; + + let (result, snapshot) = { + let mut gw = state.gateway.lock().await; + if let Some(store) = on_disk { + gw.replace_store(store); + } + let result = op(&mut gw); + (result, gw.store.clone()) + }; + + tokio::task::spawn_blocking(move || snapshot.save(&path)) + .await + .map_err(|e| internal(e.to_string()))? + .map_err(|e| internal(e.to_string()))?; + Ok(result) +} + pub(super) async fn teach_handler( State(state): State, headers: HeaderMap, @@ -122,13 +172,21 @@ pub(super) async fn teach_handler( .into_response(); } - let path = store_path(&state); - let mut gw = state.gateway.lock().await; - gw.teach_with(&req.term, &req.label, &req.scope, req.whole_word); - if let Err(e) = gw.store.save(&path) { - return (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response(); - } - let total = gw.store.deny.len(); + let (term, label, scope, whole_word) = ( + req.term.clone(), + req.label.clone(), + req.scope.clone(), + req.whole_word, + ); + let total = match mutate_store(&state, move |gw| { + gw.teach_with(&term, &label, &scope, whole_word); + gw.store.deny.len() + }) + .await + { + Ok(total) => total, + Err(e) => return e.into_response(), + }; Json(json!({ "ok": true, "taught": req.term, "total": total })).into_response() } @@ -141,12 +199,12 @@ pub(super) async fn unteach_handler( return e.into_response(); } - let path = store_path(&state); - let mut gw = state.gateway.lock().await; - let removed = gw.unteach_in(&req.term, req.scope.as_deref()); - if let Err(e) = gw.store.save(&path) { - return (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response(); - } + let (term, scope) = (req.term.clone(), req.scope.clone()); + let removed = match mutate_store(&state, move |gw| gw.unteach_in(&term, scope.as_deref())).await + { + Ok(removed) => removed, + Err(e) => return e.into_response(), + }; Json(json!({ "ok": true, "removed": removed })).into_response() } diff --git a/src/proxy/mod.rs b/src/proxy/mod.rs index 5be932a..f0a57e8 100644 --- a/src/proxy/mod.rs +++ b/src/proxy/mod.rs @@ -73,6 +73,9 @@ pub struct ProxyState { pub metrics: Arc, /// Bearer token required by `/metrics`. `None`: open, like `/healthz`. pub metrics_token: Option, + /// Serializes teach/unteach so their read-modify-write of the store file cannot + /// interleave. Held across disk I/O, which is why it is not the gateway lock. + pub store_write_lock: Arc>, } impl ProxyState { @@ -100,6 +103,7 @@ impl ProxyState { scope_tokens: Vec::new(), metrics, metrics_token: None, + store_write_lock: Arc::new(Mutex::new(())), } } } diff --git a/src/store.rs b/src/store.rs index 765986f..00ebaf6 100644 --- a/src/store.rs +++ b/src/store.rs @@ -40,7 +40,7 @@ pub struct HiddenForm { /// (see docs/TEACHING.md, "What teaching does NOT do"). This file is the crown /// jewels: a map of everything you consider private. It must never leave the /// machine and must never be committed (see `.gitignore`). -#[derive(Debug, Default, Serialize, Deserialize)] +#[derive(Debug, Default, Clone, Serialize, Deserialize)] pub struct Store { #[serde(default)] pub deny: Vec, diff --git a/tests/m2_teach.rs b/tests/m2_teach.rs index 050d432..75c508a 100644 --- a/tests/m2_teach.rs +++ b/tests/m2_teach.rs @@ -317,3 +317,43 @@ async fn unteach_removes_term() { let _ = std::fs::remove_file(&store_path); } + +/// A term written to the store file out of band (the CLI, while the server runs) must +/// survive the next HTTP teach instead of being overwritten from stale memory. +#[tokio::test] +async fn http_teach_keeps_terms_added_to_the_file_out_of_band() { + let store_path = temp_store_path(); + let gw = Gateway::new(Store::default()); + let upstream = "http://127.0.0.1:1/v1/chat/completions".to_string(); + let mut state = ProxyState::new(gw, Client::new(), upstream, "key".into()); + state.store_path = Some(store_path.clone()); + state.admin_token = Some("secret-token".into()); + let proxy_url = spawn_server(portcullis::proxy::app(state)).await; + let base = base_url(&proxy_url); + let client = Client::new(); + + let teach = |term: &'static str| { + let req = client + .post(format!("{base}/teach")) + .bearer_auth("secret-token") + .json(&json!({"term": term, "label": "ORG"})); + async move { req.send().await.unwrap().status() } + }; + assert!(teach("Aerolith").await.is_success()); + + // The CLI adds a term straight to the file. + let mut on_disk = Store::load(&store_path).unwrap(); + on_disk.teach("FromTheCli", "ORG", "global"); + on_disk.save(&store_path).unwrap(); + + assert!(teach("Cartalian").await.is_success()); + let final_store = Store::load(&store_path).unwrap(); + let terms: Vec<_> = final_store.deny.iter().map(|e| e.term.as_str()).collect(); + assert!( + terms.contains(&"FromTheCli"), + "out-of-band term was overwritten: {terms:?}" + ); + assert!(terms.contains(&"Aerolith") && terms.contains(&"Cartalian")); + + let _ = std::fs::remove_file(&store_path); +} From 2f509f209dbd4cae4bfd097020a2a48ff924aeae Mon Sep 17 00:00:00 2001 From: Vederis L Date: Fri, 2 Oct 2026 14:50:48 +0800 Subject: [PATCH 4/4] fix: fail-closed check also covers decoded JSON strings and keys The outbound assertion ran only over the serialized body, where JSON escaping hides a taught term containing a quote, backslash or control character. Every decoded string value and object key is now checked as well, with the placeholder check kept to the whole body. Shared by both providers. Co-Authored-By: Claude Sonnet 5.5 --- src/pipeline.rs | 19 ++++++++++++++--- src/proxy/anthropic.rs | 12 +---------- src/proxy/mod.rs | 47 ++++++++++++++++++++++++++++++++++++++++++ src/proxy/openai.rs | 12 +---------- tests/hardening.rs | 47 ++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 112 insertions(+), 25 deletions(-) diff --git a/src/pipeline.rs b/src/pipeline.rs index 493f30e..de40571 100644 --- a/src/pipeline.rs +++ b/src/pipeline.rs @@ -393,18 +393,31 @@ impl Gateway { } /// [`Gateway::assert_clean`] for a request in `scope`. - pub fn assert_clean_scoped( + /// Only the taught-term half of [`Gateway::assert_clean_scoped`]. For checking the + /// individual strings of a JSON body, where "<<" without ">>" is just prose. + pub fn assert_no_terms_scoped( &self, scope: Option<&str>, - outbound: &[String], + segments: &[String], ) -> Result<(), String> { let dict = self.dictionary(scope); - for seg in outbound { + for seg in segments { for span in dict.detect(seg) { if !self.store.is_allowed_for(scope, &seg[span.start..span.end]) { return Err("residual protected term in outbound".into()); } } + } + Ok(()) + } + + pub fn assert_clean_scoped( + &self, + scope: Option<&str>, + outbound: &[String], + ) -> Result<(), String> { + self.assert_no_terms_scoped(scope, outbound)?; + for seg in outbound { if seg.contains("<<") && !seg.contains(">>") { return Err("malformed placeholder in outbound".into()); } diff --git a/src/proxy/anthropic.rs b/src/proxy/anthropic.rs index 55c1184..912a113 100644 --- a/src/proxy/anthropic.rs +++ b/src/proxy/anthropic.rs @@ -256,17 +256,7 @@ pub(super) async fn handle_anthropic( // Invariant 4: fail-closed outbound assertion over the ENTIRE assembled body. { let gw = state.gateway.lock().await; - let assembled = serde_json::to_string(&redacted_body) - .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?; - gw.assert_clean_scoped(scope.as_deref(), &[assembled]) - .map_err(|e| { - state.metrics.blocked(if e.contains("residual") { - "residual_term" - } else { - "malformed_placeholder" - }); - (StatusCode::BAD_GATEWAY, e) - })?; + assert_outbound_clean(&gw, &state.metrics, scope.as_deref(), &redacted_body)?; } let streaming = body diff --git a/src/proxy/mod.rs b/src/proxy/mod.rs index f0a57e8..f3274a7 100644 --- a/src/proxy/mod.rs +++ b/src/proxy/mod.rs @@ -536,3 +536,50 @@ pub(super) fn resolve_scope( } } } + +/// Every decoded string in `body` — values **and** object keys. +/// +/// The fail-closed assertion also runs over the serialized body, but JSON escapes some +/// characters (`"`, `\`, control characters, and optionally non-ASCII), so a taught +/// term containing one would not appear verbatim there. Checking the decoded strings +/// closes that gap. +pub(super) fn outbound_strings(body: &Value) -> Vec { + fn walk(v: &Value, out: &mut Vec) { + match v { + Value::String(s) => out.push(s.clone()), + Value::Array(a) => a.iter().for_each(|x| walk(x, out)), + Value::Object(m) => { + for (k, x) in m { + out.push(k.clone()); + walk(x, out); + } + } + _ => {} + } + } + let mut out = Vec::new(); + walk(body, &mut out); + out +} + +/// The fail-closed outbound check (invariant 4), shared by both providers' handlers. +/// Counts a block in the metrics by reason. +pub(super) fn assert_outbound_clean( + gw: &Gateway, + metrics: &crate::Metrics, + scope: Option<&str>, + body: &Value, +) -> Result<(), (StatusCode, String)> { + let assembled = + serde_json::to_string(body).map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?; + gw.assert_clean_scoped(scope, &[assembled]) + .and_then(|()| gw.assert_no_terms_scoped(scope, &outbound_strings(body))) + .map_err(|e| { + metrics.blocked(if e.contains("residual") { + "residual_term" + } else { + "malformed_placeholder" + }); + (StatusCode::BAD_GATEWAY, e) + }) +} diff --git a/src/proxy/openai.rs b/src/proxy/openai.rs index 021885f..f9c1fe6 100644 --- a/src/proxy/openai.rs +++ b/src/proxy/openai.rs @@ -194,17 +194,7 @@ pub(super) async fn handle( // collect a field that carries a protected term, this is what stops it leaving. { let gw = state.gateway.lock().await; - let assembled = serde_json::to_string(&upstream_body) - .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?; - gw.assert_clean_scoped(scope.as_deref(), &[assembled]) - .map_err(|e| { - state.metrics.blocked(if e.contains("residual") { - "residual_term" - } else { - "malformed_placeholder" - }); - (StatusCode::BAD_GATEWAY, e) - })?; + assert_outbound_clean(&gw, &state.metrics, scope.as_deref(), &upstream_body)?; } let streaming = body diff --git a/tests/hardening.rs b/tests/hardening.rs index 13ef174..115f60e 100644 --- a/tests/hardening.rs +++ b/tests/hardening.rs @@ -186,3 +186,50 @@ fn save_tightens_a_pre_existing_loose_file() { let mode = std::fs::metadata(&path).unwrap().permissions().mode() & 0o777; assert_eq!(mode, 0o600, "an existing 0644 store should be tightened"); } + +/// A taught term containing a character JSON escapes (`"`) never appears verbatim in the +/// serialized body, so the fail-closed check must also look at the decoded strings. Here +/// the term rides in a field the redaction walker does not collect (`metadata`). +#[tokio::test] +async fn fail_closed_catches_terms_that_json_escapes() { + use std::sync::Arc; + + let mut store = Store::default(); + store.teach("the \"Falcon\" project", "ORG", "global"); + let hits = Arc::new(tokio::sync::Mutex::new(0u32)); + let counter = hits.clone(); + let upstream = spawn(axum::Router::new().route( + "/v1/chat/completions", + axum::routing::post(move || { + let counter = counter.clone(); + async move { + *counter.lock().await += 1; + axum::Json(serde_json::json!({"choices": []})) + } + }), + )) + .await; + let mut state = proxy_state( + Duration::from_secs(5), + format!("{upstream}/v1/chat/completions"), + ); + state.gateway = Arc::new(tokio::sync::Mutex::new(portcullis::Gateway::new(store))); + let base = spawn(app(state)).await; + + let resp = reqwest::Client::new() + .post(format!("{base}/v1/chat/completions")) + .json(&serde_json::json!({ + "model": "m", + "messages": [{"role": "user", "content": "hello"}], + "metadata": {"note": "about the \"Falcon\" project"} + })) + .send() + .await + .unwrap(); + assert_eq!( + resp.status(), + 502, + "a term in an unwalked field must fail closed" + ); + assert_eq!(*hits.lock().await, 0, "nothing may reach the upstream"); +}