1use crate::tools::RegisteredTool;
4
5pub const CAPABILITIES: [&str; 5] = ["fs", "net", "env", "session", "all"];
9
10#[derive(Clone, Copy, Debug, PartialEq, Eq)]
12pub struct Grants {
13 pub filesystem: bool,
15 pub process: bool,
17 pub network: bool,
19 pub environment: bool,
21 pub session: bool,
23}
24
25impl Default for Grants {
26 fn default() -> Grants {
27 Grants {
28 filesystem: true,
29 process: true,
30 network: true,
31 environment: true,
32 session: true,
33 }
34 }
35}
36
37impl std::fmt::Display for Grants {
40 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41 let names = self.names();
42 if names.is_empty() {
43 return formatter.write_str("none");
44 }
45 formatter.write_str(&names.join(", "))
46 }
47}
48
49impl Grants {
50 pub fn none() -> Grants {
52 Grants {
53 filesystem: false,
54 process: false,
55 network: false,
56 environment: false,
57 session: false,
58 }
59 }
60
61 pub fn is_all(&self) -> bool {
63 *self == Grants::default()
64 }
65
66 fn of(capability: &str) -> Option<Grants> {
68 let mut grants = Grants::none();
69 match capability {
70 "fs" => grants.filesystem = true,
71 "net" => grants.network = true,
72 "env" => grants.environment = true,
73 "session" => grants.session = true,
74 "all" => grants = Grants::default(),
75 _ => return None,
76 }
77 Some(grants)
78 }
79
80 pub fn covers(&self, capability: &str) -> bool {
83 Grants::of(capability).is_some_and(|needed| self.intersection(needed) == needed)
84 }
85
86 pub fn names(&self) -> Vec<&'static str> {
91 if self.is_all() {
92 return vec!["all"];
93 }
94 let mut names: Vec<&'static str> = CAPABILITIES
95 .into_iter()
96 .filter(|capability| self.covers(capability))
97 .collect();
98 if self.process {
99 names.push("process");
100 }
101 names
102 }
103
104 pub fn parse(values: &[String]) -> Result<Grants, String> {
109 let mut grants = Grants::none();
110 for value in values {
111 match value.as_str() {
112 "none" if values.len() == 1 => {}
113 "none" => return Err("\"none\" cannot be combined with other capabilities".into()),
114 "process" => {
115 return Err(
116 "\"process\" is not a capability: a process runs as you, outside the sandbox, so only \"all\" allows it".into(),
117 );
118 }
119 other => match Grants::of(other) {
120 Some(granted) => grants = grants.union(granted),
121 None => return Err(format!("unknown capability \"{other}\"")),
122 },
123 }
124 }
125 Ok(grants)
126 }
127
128 pub fn union(self, other: Grants) -> Grants {
130 Grants {
131 filesystem: self.filesystem || other.filesystem,
132 process: self.process || other.process,
133 network: self.network || other.network,
134 environment: self.environment || other.environment,
135 session: self.session || other.session,
136 }
137 }
138
139 pub fn intersection(self, other: Grants) -> Grants {
141 Grants {
142 filesystem: self.filesystem && other.filesystem,
143 process: self.process && other.process,
144 network: self.network && other.network,
145 environment: self.environment && other.environment,
146 session: self.session && other.session,
147 }
148 }
149}
150
151pub fn tool_granted(grants: &Grants, tool: &RegisteredTool) -> bool {
156 let all = grants.is_all();
157 match tool.name() {
158 _ if tool.source.is_some() => all,
159 "bash" => grants.process,
160 "grep" | "find" => grants.filesystem && grants.process && grants.network,
162 "read" | "write" | "edit" | "ls" => grants.filesystem,
163 _ => all,
164 }
165}
166
167pub fn receives(kind: &str, grants: &Grants) -> bool {
171 match kind {
172 "before_provider_request" | "before_provider_headers" => grants.network,
173 _ => true,
174 }
175}
176
177#[cfg(test)]
178mod tests {
179 use super::*;
180
181 #[test]
182 fn capabilities_parse() {
183 let fs = Grants {
184 filesystem: true,
185 ..Grants::none()
186 };
187 let all_but_process = Grants {
188 process: false,
189 ..Grants::default()
190 };
191 let process = "\"process\" is not a capability: a process runs as you, outside the sandbox, so only \"all\" allows it";
192 let cases: &[(&[&str], Result<Grants, &str>)] = &[
193 (&[], Ok(Grants::none())),
194 (&["none"], Ok(Grants::none())),
195 (&["all"], Ok(Grants::default())),
196 (&["fs", "all"], Ok(Grants::default())),
197 (&["fs"], Ok(fs)),
198 (&["fs", "net", "env", "session"], Ok(all_but_process)),
199 (&["process"], Err(process)),
200 (
201 &["none", "fs"],
202 Err("\"none\" cannot be combined with other capabilities"),
203 ),
204 (&["exec"], Err("unknown capability \"exec\"")),
205 ];
206 for (values, expected) in cases {
207 let values: Vec<String> = values.iter().map(|value| (*value).to_owned()).collect();
208 let expected = expected.map_err(str::to_owned);
209 assert_eq!(Grants::parse(&values), expected, "{values:?}");
210 }
211 let no_net = Grants {
213 network: false,
214 ..Grants::default()
215 };
216 for (grants, name) in [
217 (Grants::default(), "all"),
218 (Grants::none(), "none"),
219 (fs, "fs"),
220 (all_but_process, "fs, net, env, session"),
221 (no_net, "fs, env, session, process"),
222 ] {
223 assert_eq!(grants.to_string(), name);
224 for capability in CAPABILITIES {
225 let named = grants.names().contains(&capability) || grants.is_all();
226 assert_eq!(grants.covers(capability), named, "{name}: {capability}");
227 }
228 }
229 assert!(!Grants::default().covers("process"));
231 assert!(!Grants::default().covers("exec"));
232 let net = Grants {
233 network: true,
234 ..Grants::none()
235 };
236 assert_eq!(
237 fs.union(net),
238 Grants {
239 filesystem: true,
240 network: true,
241 ..Grants::none()
242 }
243 );
244 }
245
246 #[test]
249 fn grants_gate_tools() {
250 let env = crate::tools::ToolEnv {
251 cwd: std::path::PathBuf::new(),
252 runtime: std::sync::Arc::default(),
253 bin_dir: std::path::PathBuf::new(),
254 };
255 let fs = Grants {
256 filesystem: true,
257 ..Grants::none()
258 };
259 let all_but_process = Grants {
260 process: false,
261 ..Grants::default()
262 };
263 let cases = [
266 ("read", false, [false, true, true, true]),
267 ("write", false, [false, true, true, true]),
268 ("edit", false, [false, true, true, true]),
269 ("ls", false, [false, true, true, true]),
270 ("bash", false, [false, false, false, true]),
271 ("grep", false, [false, false, false, true]),
272 ("find", false, [false, false, false, true]),
273 ("read", true, [false, false, false, true]),
274 ];
275 for (name, extension, expected) in cases {
276 let mut tool = crate::tools::builtin(name, &env).unwrap();
277 if extension {
278 tool.source = Some(yapi_types::rpc::SourceInfo::default());
279 }
280 let granted = [Grants::none(), fs, all_but_process, Grants::default()]
281 .map(|grants| tool_granted(&grants, &tool));
282 assert_eq!(granted, expected, "{name}, extension {extension}");
283 }
284 }
285
286 #[test]
287 fn provider_events_need_network() {
288 let net = Grants {
289 network: true,
290 ..Grants::none()
291 };
292 let no_net = Grants {
293 network: false,
294 ..Grants::default()
295 };
296 for (kind, heard_without_network) in [
297 ("before_provider_request", false),
298 ("before_provider_headers", false),
299 ("after_provider_response", true),
300 ("provider_stream_event", true),
301 ("tool_call", true),
302 ] {
303 assert!(receives(kind, &net), "{kind}");
304 assert_eq!(receives(kind, &no_net), heard_without_network, "{kind}");
305 }
306 }
307}