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/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..de40571 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,37 @@ 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. +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; @@ -36,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 { @@ -50,6 +85,7 @@ impl Gateway { extra: None, suggestions: VecDeque::new(), metrics, + dicts: Mutex::new(HashMap::new()), } } @@ -116,6 +152,37 @@ 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 + } + + /// 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 } @@ -213,7 +280,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) { @@ -273,7 +350,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 { @@ -285,9 +362,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 } @@ -317,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 = DictionaryDetector::for_scope(&self.store, scope); - for seg in outbound { + let dict = self.dictionary(scope); + 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()); } @@ -346,6 +435,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/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/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 905df14..f3274a7 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(())), } } } @@ -267,10 +271,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"), @@ -533,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/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/hardening.rs b/tests/hardening.rs index 702f99d..115f60e 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() { @@ -164,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"); +} 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")); +} 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); +}