Skip to main content

tokenfold_core/
safety.rs

1//! Fail-closed safety invariants (INTERFACES.md Part 2 "Safety Invariants on Order"). The
2//! pipeline checks these after every transform; a violation rolls that transform back
3//! (pre-transform bytes restored) rather than shipping a corrupted or unsafe output.
4
5use serde_json::Value;
6
7/// Invariant 1: after any JSON transform, the output must still parse as JSON.
8pub fn json_still_valid(bytes: &[u8]) -> bool {
9    serde_json::from_slice::<Value>(bytes).is_ok()
10}
11
12/// Invariant 2: after any JSON transform, every object's key order must be unchanged.
13/// Compares the depth-first sequence of object keys (not values, not array lengths — a
14/// transform may legitimately shorten an array while leaving every object's key order alone).
15pub fn json_key_order_preserved(before: &[u8], after: &[u8]) -> bool {
16    let (Ok(before_value), Ok(after_value)) = (
17        serde_json::from_slice::<Value>(before),
18        serde_json::from_slice::<Value>(after),
19    ) else {
20        return false;
21    };
22    let mut before_keys = Vec::new();
23    let mut after_keys = Vec::new();
24    collect_key_order(&before_value, &mut before_keys);
25    collect_key_order(&after_value, &mut after_keys);
26    before_keys == after_keys
27}
28
29fn collect_key_order(value: &Value, out: &mut Vec<String>) {
30    match value {
31        Value::Object(map) => {
32            for (key, nested) in map {
33                out.push(key.clone());
34                collect_key_order(nested, out);
35            }
36        }
37        Value::Array(items) => {
38            for item in items {
39                collect_key_order(item, out);
40            }
41        }
42        _ => {}
43    }
44}
45
46/// Invariant 3: every protected-content segment (system turns, latest user message, diff
47/// headers) must still be present, byte-for-byte, somewhere in the output. Segments are
48/// checked individually (see `budget::protected_segments`) because concatenating them into
49/// one blob would rarely be contiguous in the real document.
50pub fn protected_segments_present(segments: &[Vec<u8>], output: &[u8]) -> bool {
51    segments
52        .iter()
53        .all(|segment| contains_subslice(output, segment))
54}
55
56fn contains_subslice(haystack: &[u8], needle: &[u8]) -> bool {
57    if needle.is_empty() {
58        return true;
59    }
60    if needle.len() > haystack.len() {
61        return false;
62    }
63    haystack
64        .windows(needle.len())
65        .any(|window| window == needle)
66}
67
68/// Invariant 4: `secret_redaction` output must not contain any pattern that matched in the
69/// original. Re-running redaction over its own output and confirming no further match is
70/// found is equivalent to a fresh scan, without needing a separate "scan only" API.
71pub fn no_redaction_bypass(redacted_output: &[u8]) -> bool {
72    crate::transforms::redaction::redact(redacted_output).redacted_count == 0
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78
79    #[test]
80    fn json_still_valid_detects_broken_json() {
81        assert!(json_still_valid(b"{\"a\":1}"));
82        assert!(!json_still_valid(b"{not json"));
83    }
84
85    #[test]
86    fn key_order_preserved_across_whitespace_only_change() {
87        let before = br#"{"b": 1, "a": 2}"#;
88        let after = br#"{"b":1,"a":2}"#;
89        assert!(json_key_order_preserved(before, after));
90    }
91
92    #[test]
93    fn key_order_violation_detected() {
94        let before = br#"{"b": 1, "a": 2}"#;
95        let after = br#"{"a":2,"b":1}"#;
96        assert!(!json_key_order_preserved(before, after));
97    }
98
99    #[test]
100    fn key_order_preserved_when_nested_array_is_shortened() {
101        let before = br#"{"name":"x","examples":[1,2,3,4,5]}"#;
102        let after = br#"{"name":"x","examples":[1]}"#;
103        assert!(json_key_order_preserved(before, after));
104    }
105
106    #[test]
107    fn protected_segments_present_checks_each_segment_independently() {
108        let segments = vec![b"system prompt".to_vec(), b"latest question".to_vec()];
109        let output = b"preamble system prompt middle latest question tail".to_vec();
110        assert!(protected_segments_present(&segments, &output));
111    }
112
113    #[test]
114    fn protected_segments_present_fails_when_one_segment_is_missing() {
115        let segments = vec![b"system prompt".to_vec(), b"latest question".to_vec()];
116        let output = b"preamble system prompt middle tail".to_vec();
117        assert!(!protected_segments_present(&segments, &output));
118    }
119
120    #[test]
121    fn empty_segment_list_is_trivially_satisfied() {
122        assert!(protected_segments_present(&[], b"anything"));
123    }
124
125    #[test]
126    fn no_redaction_bypass_true_after_clean_redaction() {
127        let outcome = crate::transforms::redaction::redact(b"Bearer sk-abcdEFGH1234567890123456");
128        assert!(no_redaction_bypass(&outcome.bytes));
129    }
130}