Skip to main content

pv_sandbox/
pushed.rs

1//! Console-pushed policy, layered on top of the local sandbox config.
2//!
3//! The console publishes a policy bundle; `pvctl console-sync
4//! --pull-policies` hands it to the daemon, which compiles it into one
5//! [`PushedRules`] value and swaps it into the running [`crate::Sandbox`].
6//! Pushed rules only ever add restrictions:
7//!
8//! * budgets are merged with [`BudgetSpec::tightened_by`], so the
9//!   stricter cap wins on every dimension;
10//! * allowlist rules run after the local policy chain, so a call must
11//!   pass both;
12//! * schema rules run after the local schema, so arguments must satisfy
13//!   both.
14//!
15//! A compromised console can therefore block calls the host allows, but
16//! it can never let through a call the host's own config refuses.
17
18use std::collections::HashSet;
19
20use pv_state::BudgetSpec;
21use serde_json::Value;
22
23/// A tool-name pattern taken from a policy scope (`tool.*`,
24/// `tool.create_payment`, `tool.payments.*`). `*` matches any run of
25/// characters, including none.
26#[derive(Debug, Clone, PartialEq, Eq)]
27pub struct ToolPattern {
28    parts: Vec<String>,
29}
30
31impl ToolPattern {
32    /// Parse a policy scope. Returns `None` for scopes that do not name
33    /// tools (`llm.request`, `llm.egress`, …): the daemon cannot apply a
34    /// tool rule to them and reports the policy as unsupported instead of
35    /// guessing.
36    pub fn from_scope(scope: &str) -> Option<Self> {
37        let scope = scope.trim();
38        let rest = if scope == "tool" {
39            "*"
40        } else {
41            let rest = scope.strip_prefix("tool.")?;
42            if rest.is_empty() {
43                "*"
44            } else {
45                rest
46            }
47        };
48        Some(Self {
49            parts: rest.split('*').map(str::to_owned).collect(),
50        })
51    }
52
53    /// True when `tool` matches the pattern.
54    pub fn matches(&self, tool: &str) -> bool {
55        let Some((first, rest)) = self.parts.split_first() else {
56            return false;
57        };
58        let Some(mut remaining) = tool.strip_prefix(first.as_str()) else {
59            return false;
60        };
61        let Some((last, middle)) = rest.split_last() else {
62            // No `*` at all: the whole name must have matched.
63            return remaining.is_empty();
64        };
65        for part in middle {
66            match remaining.find(part.as_str()) {
67                Some(index) => remaining = remaining.get(index + part.len()..).unwrap_or_default(),
68                None => return false,
69            }
70        }
71        remaining.ends_with(last.as_str())
72    }
73}
74
75/// One pushed allowlist policy, limited to the tools its scope names.
76#[derive(Debug, Clone)]
77pub struct PushedAllowlist {
78    /// Console policy name, recorded as the blocking policy.
79    pub policy: String,
80    /// Which tools the rule applies to.
81    pub scope: ToolPattern,
82    /// When set, only these tools may run (within the scope).
83    pub allow_only: Option<HashSet<String>>,
84    /// These tools never run.
85    pub deny: HashSet<String>,
86}
87
88impl PushedAllowlist {
89    /// The reason a call to `tool` is refused, if it is.
90    pub fn refuses(&self, tool: &str) -> Option<String> {
91        if !self.scope.matches(tool) {
92            return None;
93        }
94        if self.deny.contains(tool) {
95            return Some(format!("tool {tool:?} is deny-listed"));
96        }
97        match &self.allow_only {
98            Some(allowed) if !allowed.contains(tool) => Some(format!("tool {tool:?} is not allow-listed")),
99            _ => None,
100        }
101    }
102}
103
104/// One pushed argument schema for one tool.
105pub struct PushedSchema {
106    /// Console policy name, recorded as the blocking policy.
107    pub policy: String,
108    /// The tool whose arguments are validated.
109    pub tool: String,
110    validator: jsonschema::Validator,
111}
112
113impl std::fmt::Debug for PushedSchema {
114    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
115        f.debug_struct("PushedSchema")
116            .field("policy", &self.policy)
117            .field("tool", &self.tool)
118            .finish_non_exhaustive()
119    }
120}
121
122/// Refuses every external `$ref`. A pushed schema comes from the console,
123/// over the network; with jsonschema's default retriever a reference such
124/// as `http://169.254.169.254/…` or `file:///etc/…` would make the daemon
125/// fetch an address inside the customer's network, or read a local file,
126/// and report what it found back to the console in the policy status.
127/// References inside the schema itself (`#/$defs/…`) need no retrieval
128/// and keep working.
129struct NoExternalRefs;
130
131impl jsonschema::Retrieve for NoExternalRefs {
132    fn retrieve(
133        &self,
134        uri: &jsonschema::Uri<String>,
135    ) -> Result<Value, Box<dyn std::error::Error + Send + Sync>> {
136        Err(format!(
137            "pushed schemas may not reference external documents ({})",
138            uri.scheme().as_str()
139        )
140        .into())
141    }
142}
143
144impl PushedSchema {
145    /// Compile a schema the same way local tool schemas are compiled
146    /// (draft 2020-12 with `format` enforced), except that external
147    /// references are refused (see `NoExternalRefs`).
148    pub fn compile(policy: &str, tool: &str, schema: &Value) -> Result<Self, String> {
149        let validator = jsonschema::options()
150            .should_validate_formats(true)
151            .with_retriever(NoExternalRefs)
152            .build(schema)
153            .map_err(|e| format!("json_schema does not compile: {e}"))?;
154        Ok(Self {
155            policy: policy.to_owned(),
156            tool: tool.to_owned(),
157            validator,
158        })
159    }
160
161    /// Up to three validation errors for `arguments` (bounded work).
162    pub fn errors(&self, arguments: &Value) -> Vec<String> {
163        self.validator
164            .iter_errors(arguments)
165            .take(3)
166            .map(|e| e.to_string())
167            .collect()
168    }
169}
170
171/// Everything the daemon applied from one console bundle.
172#[derive(Debug, Default)]
173pub struct PushedRules {
174    /// The bundle etag these rules came from; `None` before any bundle.
175    pub etag: Option<String>,
176    /// Pushed budget caps (merged with the local budget, stricter wins).
177    pub budget: BudgetSpec,
178    /// Console names of the budget policies merged into `budget`.
179    pub budget_policies: Vec<String>,
180    /// Pushed allowlist rules.
181    pub allowlists: Vec<PushedAllowlist>,
182    /// Pushed argument schemas.
183    pub schemas: Vec<PushedSchema>,
184}
185
186impl PushedRules {
187    /// Which pushed budget policy set the cap a refusal hit, when the
188    /// binding cap came from the console rather than the local config.
189    /// `limit` is the refusal's limit name (`max_tokens`,
190    /// `max_tool_calls.db_write`, …).
191    pub fn budget_policy_for(&self, local: &BudgetSpec, limit: &str, cap: u64) -> Option<String> {
192        let (pushed, local_cap) = match limit {
193            "max_tokens" => (self.budget.max_tokens, local.max_tokens),
194            "max_payout_usd_micros" => (self.budget.max_payout_usd_micros, local.max_payout_usd_micros),
195            "max_total_tool_calls" => (self.budget.max_total_tool_calls, local.max_total_tool_calls),
196            other => {
197                let tool = other.strip_prefix("max_tool_calls.")?;
198                (
199                    self.budget.max_tool_calls.get(tool).copied(),
200                    local.max_tool_calls.get(tool).copied(),
201                )
202            }
203        };
204        let from_console = pushed == Some(cap) && local_cap.is_none_or(|local_cap| local_cap > cap);
205        if from_console && !self.budget_policies.is_empty() {
206            Some(self.budget_policies.join(","))
207        } else {
208            None
209        }
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
216
217    use super::*;
218    use serde_json::json;
219
220    #[test]
221    fn pushed_schemas_never_fetch_external_references() {
222        // A listener that would see the daemon's request if it made one.
223        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
224        listener.set_nonblocking(true).unwrap();
225        let url = format!("http://{}/schema.json", listener.local_addr().unwrap());
226        let err = PushedSchema::compile("p", "t", &json!({"$ref": url})).unwrap_err();
227        assert!(err.contains("may not reference external documents"), "{err}");
228        assert!(listener.accept().is_err(), "the daemon must not connect anywhere");
229
230        // A custom meta-schema and `$dynamicRef` are resolved the same way.
231        let meta = format!("http://{}/meta.json", listener.local_addr().unwrap());
232        assert!(PushedSchema::compile("p", "t", &json!({"$schema": meta})).is_err());
233        assert!(PushedSchema::compile("p", "t", &json!({"$dynamicRef": url})).is_err());
234        // The standard meta-schemas are built in and need no retrieval.
235        for draft in [
236            "https://json-schema.org/draft/2020-12/schema",
237            "https://json-schema.org/draft/2019-09/schema",
238            "http://json-schema.org/draft-07/schema#",
239            "http://json-schema.org/draft-07/schema",
240            "http://json-schema.org/draft-06/schema#",
241            "http://json-schema.org/draft-04/schema#",
242        ] {
243            PushedSchema::compile("p", "t", &json!({"$schema": draft, "type": "object"})).unwrap();
244        }
245        assert!(listener.accept().is_err(), "the daemon must not connect anywhere");
246
247        let file = std::env::temp_dir().join(format!("pv-pushed-ref-{}.json", std::process::id()));
248        std::fs::write(&file, r#"{"type":"object"}"#).unwrap();
249        // A valid file URL on every platform: `file:///C:/…` on Windows.
250        let path = file.display().to_string().replace('\\', "/");
251        let file_ref = if path.starts_with('/') {
252            format!("file://{path}")
253        } else {
254            format!("file:///{path}")
255        };
256        let err = PushedSchema::compile("p", "t", &json!({"$ref": file_ref})).unwrap_err();
257        let _ = std::fs::remove_file(&file);
258        assert!(err.contains("may not reference external documents"), "{err}");
259
260        // References inside the schema need no retrieval.
261        let local = PushedSchema::compile(
262            "p",
263            "t",
264            &json!({"$defs": {"amount": {"type": "integer", "maximum": 10}}, "properties": {"n": {"$ref": "#/$defs/amount"}}}),
265        )
266        .unwrap();
267        assert!(local.errors(&json!({"n": 5})).is_empty());
268        assert!(!local.errors(&json!({"n": 50})).is_empty());
269    }
270
271    #[test]
272    fn scopes_become_tool_patterns() {
273        let all = ToolPattern::from_scope("tool.*").unwrap();
274        assert!(all.matches("anything") && all.matches(""));
275        assert!(ToolPattern::from_scope("tool").unwrap().matches("x"));
276        let exact = ToolPattern::from_scope("tool.create_payment").unwrap();
277        assert!(exact.matches("create_payment"));
278        assert!(!exact.matches("create_payment_v2") && !exact.matches("create"));
279        let family = ToolPattern::from_scope("tool.payments.*").unwrap();
280        assert!(family.matches("payments.refund") && family.matches("payments."));
281        assert!(!family.matches("payment.refund"));
282        let middle = ToolPattern::from_scope("tool.db_*_write").unwrap();
283        assert!(middle.matches("db_orders_write") && middle.matches("db__write"));
284        assert!(!middle.matches("db_orders_read"));
285        for scope in ["llm.request", "llm.egress", "*", "", "tools.x"] {
286            assert!(ToolPattern::from_scope(scope).is_none(), "{scope}");
287        }
288    }
289
290    #[test]
291    fn allowlists_apply_only_inside_their_scope() {
292        let rule = PushedAllowlist {
293            policy: "runtime.tools".into(),
294            scope: ToolPattern::from_scope("tool.payments.*").unwrap(),
295            allow_only: Some(["payments.quote".to_owned()].into_iter().collect()),
296            deny: ["payments.wire".to_owned()].into_iter().collect(),
297        };
298        assert!(rule.refuses("search").is_none(), "outside the scope");
299        assert!(rule.refuses("payments.quote").is_none());
300        assert!(rule.refuses("payments.wire").unwrap().contains("deny-listed"));
301        assert!(rule
302            .refuses("payments.refund")
303            .unwrap()
304            .contains("not allow-listed"));
305    }
306
307    #[test]
308    fn pushed_schema_reports_bounded_errors() {
309        let schema = PushedSchema::compile(
310            "procurement.vendors",
311            "create_purchase_order",
312            &json!({"type": "object", "required": ["vendor"], "properties": {"vendor": {"enum": ["Contoso"]}}}),
313        )
314        .unwrap();
315        assert!(schema.errors(&json!({"vendor": "Contoso"})).is_empty());
316        assert!(!schema.errors(&json!({"vendor": "Initech"})).is_empty());
317        assert!(PushedSchema::compile("p", "t", &json!({"type": 12})).is_err());
318    }
319
320    #[test]
321    fn budget_refusals_name_the_console_policy_only_when_it_set_the_cap() {
322        let mut rules = PushedRules {
323            budget_policies: vec!["finance.cap".into()],
324            ..PushedRules::default()
325        };
326        rules.budget.max_total_tool_calls = Some(2);
327        rules.budget.max_tool_calls.insert("db_write".into(), 1);
328        let local = BudgetSpec {
329            max_total_tool_calls: Some(10),
330            ..BudgetSpec::default()
331        };
332        assert_eq!(
333            rules
334                .budget_policy_for(&local, "max_total_tool_calls", 2)
335                .as_deref(),
336            Some("finance.cap")
337        );
338        assert_eq!(
339            rules
340                .budget_policy_for(&local, "max_tool_calls.db_write", 1)
341                .as_deref(),
342            Some("finance.cap")
343        );
344        let tighter_local = BudgetSpec {
345            max_total_tool_calls: Some(1),
346            ..BudgetSpec::default()
347        };
348        assert!(rules
349            .budget_policy_for(&tighter_local, "max_total_tool_calls", 1)
350            .is_none());
351        assert!(rules.budget_policy_for(&local, "max_tokens", 5).is_none());
352    }
353}