Skip to main content

hessra_context_token/
exposure.rs

1//! Exposure tracking operations for context tokens.
2//!
3//! Exposure labels are added as append-only Biscuit blocks. Each block contains
4//! `exposure({label})` facts. Labels accumulate and cannot be removed.
5
6extern crate biscuit_auth as biscuit;
7
8use biscuit::Biscuit;
9use biscuit::macros::block;
10use chrono::Utc;
11use hessra_token_core::{KeyPair, PublicKey, TokenError};
12use std::error::Error;
13
14/// Add exposure labels to a context token.
15///
16/// Creates a new Biscuit block containing `exposure({label})` facts for each
17/// provided label, and `exposure_source({source})` identifying where the exposure
18/// came from.
19///
20/// This operation is append-only: the resulting token has strictly more exposure
21/// than the input token.
22///
23/// # Arguments
24/// * `token` - The base64-encoded context token
25/// * `public_key` - The public key used to verify the token signature
26/// * `labels` - The exposure labels to add (e.g., `["PII:SSN", "PII:email"]`)
27/// * `source` - The data source that produced the exposure (e.g., `"data:user-ssn"`)
28///
29/// # Returns
30/// Updated base64-encoded context token with exposure labels appended
31pub fn add_exposure(
32    token: &str,
33    public_key: PublicKey,
34    labels: &[String],
35    source: String,
36) -> Result<String, Box<dyn Error>> {
37    if labels.is_empty() {
38        return Ok(token.to_string());
39    }
40
41    let biscuit = Biscuit::from_base64(token, public_key)?;
42
43    let now = Utc::now().timestamp();
44
45    // Build a block with exposure facts
46    let mut block_builder = block!(
47        r#"
48            exposure_source({source});
49            exposure_time({now});
50        "#
51    );
52
53    for label in labels {
54        let label_str = label.clone();
55        block_builder = block_builder.fact(biscuit::macros::fact!(r#"exposure({label_str});"#))?;
56    }
57
58    let new_biscuit = biscuit.append(block_builder)?;
59    let new_token = new_biscuit.to_base64()?;
60
61    Ok(new_token)
62}
63
64/// Extract all exposure labels from a context token by parsing its Biscuit blocks.
65///
66/// Iterates through all blocks in the token looking for `exposure("label")` facts
67/// and returns the deduplicated set of labels.
68///
69/// This is a diagnostic/inspection method. For authorization decisions, use
70/// `ContextVerifier::check_precluded_exposures` instead, which delegates to the
71/// Biscuit authorization engine.
72///
73/// # Arguments
74/// * `token` - The base64-encoded context token
75/// * `public_key` - The public key used to verify the token signature
76///
77/// # Returns
78/// Deduplicated list of exposure label strings
79pub fn extract_exposure_labels(
80    token: &str,
81    public_key: PublicKey,
82) -> Result<Vec<String>, TokenError> {
83    let biscuit = Biscuit::from_base64(token, public_key)?;
84
85    let mut labels = Vec::new();
86
87    // Iterate through all blocks looking for exposure facts
88    let block_count = biscuit.block_count();
89    for i in 0..block_count {
90        let block_source = biscuit.print_block_source(i).unwrap_or_default();
91        // Parse exposure facts from block source: lines like `exposure("PII:SSN");`
92        for line in block_source.lines() {
93            let trimmed = line.trim();
94            if let Some(rest) = trimmed.strip_prefix("exposure(") {
95                if let Some(label_str) = rest.strip_suffix(");") {
96                    // Remove quotes
97                    let label = label_str.trim_matches('"').to_string();
98                    if !labels.contains(&label) {
99                        labels.push(label);
100                    }
101                }
102            }
103        }
104    }
105
106    Ok(labels)
107}
108
109/// Fork a context token for a sub-agent, inheriting the parent's exposure.
110///
111/// Creates a fresh context token for the child subject, pre-populated with
112/// all of the parent's exposure labels. This prevents contamination laundering
113/// through delegation.
114///
115/// # Arguments
116/// * `parent_token` - The base64-encoded parent context token
117/// * `parent_public_key` - The public key used to verify the parent token
118/// * `child_subject` - The child subject identifier (e.g., "agent:openclaw:subtask-1")
119/// * `time_config` - Time configuration for the child context token
120/// * `keypair` - The keypair to sign the child token with
121///
122/// # Returns
123/// Base64-encoded child context token with inherited exposure
124pub fn fork_context(
125    parent_token: &str,
126    parent_public_key: PublicKey,
127    child_subject: String,
128    time_config: hessra_token_core::TokenTimeConfig,
129    keypair: &KeyPair,
130) -> Result<String, Box<dyn Error>> {
131    // Extract parent's exposure labels
132    let parent_labels = extract_exposure_labels(parent_token, parent_public_key)?;
133
134    // Create a fresh context for the child
135    let child_token = crate::mint::HessraContext::new(child_subject, time_config).issue(keypair)?;
136
137    // If parent has no exposure, just return the fresh child context
138    if parent_labels.is_empty() {
139        return Ok(child_token);
140    }
141
142    // Apply all parent exposure labels to the child
143    add_exposure(
144        &child_token,
145        keypair.public(),
146        &parent_labels,
147        "inherited".to_string(),
148    )
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154    use crate::mint::HessraContext;
155    use hessra_token_core::TokenTimeConfig;
156
157    #[test]
158    fn test_add_exposure_labels() {
159        let keypair = KeyPair::new();
160        let public_key = keypair.public();
161
162        let token = HessraContext::new("agent:test".to_string(), TokenTimeConfig::default())
163            .issue(&keypair)
164            .expect("Failed to create context token");
165
166        // No exposure initially
167        let labels = extract_exposure_labels(&token, public_key).expect("Failed to extract labels");
168        assert!(labels.is_empty());
169
170        // Add exposure
171        let exposed = add_exposure(
172            &token,
173            public_key,
174            &["PII:SSN".to_string()],
175            "data:user-ssn".to_string(),
176        )
177        .expect("Failed to add exposure");
178
179        let labels =
180            extract_exposure_labels(&exposed, public_key).expect("Failed to extract labels");
181        assert_eq!(labels, vec!["PII:SSN".to_string()]);
182    }
183
184    #[test]
185    fn test_add_empty_exposure_is_noop() {
186        let keypair = KeyPair::new();
187        let public_key = keypair.public();
188
189        let token = HessraContext::new("agent:test".to_string(), TokenTimeConfig::default())
190            .issue(&keypair)
191            .expect("Failed to create context token");
192
193        let result = add_exposure(&token, public_key, &[], "source".to_string())
194            .expect("Failed with empty exposure");
195
196        assert_eq!(result, token);
197    }
198
199    #[test]
200    fn test_multiple_exposure_labels() {
201        let keypair = KeyPair::new();
202        let public_key = keypair.public();
203
204        let token = HessraContext::new("agent:test".to_string(), TokenTimeConfig::default())
205            .issue(&keypair)
206            .expect("Failed to create context token");
207
208        let exposed = add_exposure(
209            &token,
210            public_key,
211            &["PII:email".to_string(), "PII:address".to_string()],
212            "data:user-profile".to_string(),
213        )
214        .expect("Failed to add exposure");
215
216        let labels =
217            extract_exposure_labels(&exposed, public_key).expect("Failed to extract labels");
218        assert_eq!(labels.len(), 2);
219        assert!(labels.contains(&"PII:email".to_string()));
220        assert!(labels.contains(&"PII:address".to_string()));
221    }
222
223    #[test]
224    fn test_cumulative_exposure() {
225        let keypair = KeyPair::new();
226        let public_key = keypair.public();
227
228        let token = HessraContext::new("agent:test".to_string(), TokenTimeConfig::default())
229            .issue(&keypair)
230            .expect("Failed to create context token");
231
232        // First exposure
233        let exposed = add_exposure(
234            &token,
235            public_key,
236            &["PII:email".to_string(), "PII:address".to_string()],
237            "data:user-profile".to_string(),
238        )
239        .expect("Failed to add first exposure");
240
241        // Second exposure
242        let more_exposed = add_exposure(
243            &exposed,
244            public_key,
245            &["PII:SSN".to_string()],
246            "data:user-ssn".to_string(),
247        )
248        .expect("Failed to add second exposure");
249
250        let labels =
251            extract_exposure_labels(&more_exposed, public_key).expect("Failed to extract labels");
252        assert_eq!(labels.len(), 3);
253        assert!(labels.contains(&"PII:email".to_string()));
254        assert!(labels.contains(&"PII:address".to_string()));
255        assert!(labels.contains(&"PII:SSN".to_string()));
256    }
257
258    #[test]
259    fn test_duplicate_exposure_labels_deduplicated() {
260        let keypair = KeyPair::new();
261        let public_key = keypair.public();
262
263        let token = HessraContext::new("agent:test".to_string(), TokenTimeConfig::default())
264            .issue(&keypair)
265            .expect("Failed to create context token");
266
267        let exposed = add_exposure(
268            &token,
269            public_key,
270            &["PII:SSN".to_string()],
271            "data:user-ssn".to_string(),
272        )
273        .expect("Failed to add first exposure");
274
275        // Add same label again
276        let double_exposed = add_exposure(
277            &exposed,
278            public_key,
279            &["PII:SSN".to_string()],
280            "another-source".to_string(),
281        )
282        .expect("Failed to add duplicate exposure");
283
284        let labels =
285            extract_exposure_labels(&double_exposed, public_key).expect("Failed to extract labels");
286        assert_eq!(labels.len(), 1);
287        assert_eq!(labels[0], "PII:SSN");
288    }
289
290    #[test]
291    fn test_fork_context_inherits_exposure() {
292        let keypair = KeyPair::new();
293        let public_key = keypair.public();
294
295        let parent = HessraContext::new("agent:parent".to_string(), TokenTimeConfig::default())
296            .issue(&keypair)
297            .expect("Failed to create parent context");
298
299        // Expose the parent
300        let exposed_parent = add_exposure(
301            &parent,
302            public_key,
303            &["PII:SSN".to_string()],
304            "data:user-ssn".to_string(),
305        )
306        .expect("Failed to add exposure to parent");
307
308        // Fork for child
309        let child = fork_context(
310            &exposed_parent,
311            public_key,
312            "agent:parent:child".to_string(),
313            TokenTimeConfig::default(),
314            &keypair,
315        )
316        .expect("Failed to fork context");
317
318        // Child should inherit parent's exposure
319        let child_labels =
320            extract_exposure_labels(&child, public_key).expect("Failed to extract child labels");
321        assert_eq!(child_labels, vec!["PII:SSN".to_string()]);
322    }
323
324    #[test]
325    fn test_fork_clean_context() {
326        let keypair = KeyPair::new();
327        let public_key = keypair.public();
328
329        let parent = HessraContext::new("agent:parent".to_string(), TokenTimeConfig::default())
330            .issue(&keypair)
331            .expect("Failed to create parent context");
332
333        // Fork without any exposure on parent
334        let child = fork_context(
335            &parent,
336            public_key,
337            "agent:parent:child".to_string(),
338            TokenTimeConfig::default(),
339            &keypair,
340        )
341        .expect("Failed to fork context");
342
343        let child_labels =
344            extract_exposure_labels(&child, public_key).expect("Failed to extract child labels");
345        assert!(child_labels.is_empty());
346    }
347
348    #[test]
349    fn test_fork_inherits_multiple_exposure_labels() {
350        let keypair = KeyPair::new();
351        let public_key = keypair.public();
352
353        let parent = HessraContext::new("agent:parent".to_string(), TokenTimeConfig::default())
354            .issue(&keypair)
355            .expect("Failed to create parent context");
356
357        // Add multiple exposure labels
358        let exposed = add_exposure(
359            &parent,
360            public_key,
361            &["PII:email".to_string(), "PII:address".to_string()],
362            "data:user-profile".to_string(),
363        )
364        .expect("Failed to add profile exposure");
365
366        let more_exposed = add_exposure(
367            &exposed,
368            public_key,
369            &["PII:SSN".to_string()],
370            "data:user-ssn".to_string(),
371        )
372        .expect("Failed to add SSN exposure");
373
374        // Fork
375        let child = fork_context(
376            &more_exposed,
377            public_key,
378            "agent:parent:child".to_string(),
379            TokenTimeConfig::default(),
380            &keypair,
381        )
382        .expect("Failed to fork context");
383
384        let child_labels =
385            extract_exposure_labels(&child, public_key).expect("Failed to extract child labels");
386        assert_eq!(child_labels.len(), 3);
387        assert!(child_labels.contains(&"PII:email".to_string()));
388        assert!(child_labels.contains(&"PII:address".to_string()));
389        assert!(child_labels.contains(&"PII:SSN".to_string()));
390    }
391}