veil_engine/
walk.rs

1//! Walks every string in a JSON value, letting the caller skip structural keys.
2
3use serde_json::{Map, Value};
4
5use crate::Error;
6
7/// What to do with one key's value.
8#[derive(Clone, Copy, Debug, PartialEq, Eq)]
9pub enum KeyAction {
10    /// Leave this key's value untouched.
11    Skip,
12    /// Walk this key's value with the same policy.
13    Walk,
14    /// Walk every string under this key, ignoring the policy. Used for
15    /// free-form user data such as tool-call arguments.
16    WalkAll,
17    /// The value is a string that itself holds JSON, such as OpenAI's
18    /// `function.arguments`. Every string inside it is walked.
19    JsonString,
20}
21
22/// Decides, per object key, whether its value is redacted or restored.
23pub trait KeyPolicy: Sync {
24    /// What to do with `key`, given the object it is in.
25    fn action(&self, parent: &Map<String, Value>, key: &str) -> KeyAction;
26}
27
28/// A policy that walks everything.
29pub struct WalkAll;
30
31impl KeyPolicy for WalkAll {
32    fn action(&self, _: &Map<String, Value>, _: &str) -> KeyAction {
33        KeyAction::Walk
34    }
35}
36
37/// What kind of string the visitor is given.
38#[derive(Clone, Copy, Debug, PartialEq, Eq)]
39pub(crate) enum StrKind {
40    Plain,
41    Json,
42}
43
44pub(crate) type Visitor<'a> = dyn FnMut(&mut String, StrKind) -> Result<(), Error> + 'a;
45
46pub(crate) fn walk(value: &mut Value, policy: &dyn KeyPolicy, f: &mut Visitor<'_>) -> Result<(), Error> {
47    match value {
48        Value::String(s) => f(s, StrKind::Plain),
49        Value::Array(items) => items.iter_mut().try_for_each(|v| walk(v, policy, f)),
50        Value::Object(map) => {
51            let actions: Vec<(String, KeyAction)> = map.keys().map(|k| (k.clone(), policy.action(map, k))).collect();
52            for (key, action) in actions {
53                let child = map.get_mut(&key).expect("key exists");
54                match action {
55                    KeyAction::Skip => {}
56                    KeyAction::Walk => walk(child, policy, f)?,
57                    KeyAction::WalkAll => walk(child, &WalkAll, f)?,
58                    KeyAction::JsonString => match child {
59                        Value::String(s) => f(s, StrKind::Json)?,
60                        other => walk(other, &WalkAll, f)?,
61                    },
62                }
63            }
64            Ok(())
65        }
66        Value::Null | Value::Bool(_) | Value::Number(_) => Ok(()),
67    }
68}
veil_engine/walk.rs

veil_engine/
walk.rs
