1extern 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
14pub 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 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
64pub 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 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 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 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
109pub 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 let parent_labels = extract_exposure_labels(parent_token, parent_public_key)?;
133
134 let child_token = crate::mint::HessraContext::new(child_subject, time_config).issue(keypair)?;
136
137 if parent_labels.is_empty() {
139 return Ok(child_token);
140 }
141
142 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 let labels = extract_exposure_labels(&token, public_key).expect("Failed to extract labels");
168 assert!(labels.is_empty());
169
170 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 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 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 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 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 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 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 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 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 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}