use crate::tools::RegisteredTool;
pub const CAPABILITIES: [&str; 5] = ["fs", "net", "env", "session", "all"];
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Grants {
pub filesystem: bool,
pub process: bool,
pub network: bool,
pub environment: bool,
pub session: bool,
}
impl Default for Grants {
fn default() -> Grants {
Grants {
filesystem: true,
process: true,
network: true,
environment: true,
session: true,
}
}
}
impl std::fmt::Display for Grants {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let names = self.names();
if names.is_empty() {
return formatter.write_str("none");
}
formatter.write_str(&names.join(", "))
}
}
impl Grants {
pub fn none() -> Grants {
Grants {
filesystem: false,
process: false,
network: false,
environment: false,
session: false,
}
}
pub fn is_all(&self) -> bool {
*self == Grants::default()
}
fn of(capability: &str) -> Option<Grants> {
let mut grants = Grants::none();
match capability {
"fs" => grants.filesystem = true,
"net" => grants.network = true,
"env" => grants.environment = true,
"session" => grants.session = true,
"all" => grants = Grants::default(),
_ => return None,
}
Some(grants)
}
pub fn covers(&self, capability: &str) -> bool {
Grants::of(capability).is_some_and(|needed| self.intersection(needed) == needed)
}
pub fn names(&self) -> Vec<&'static str> {
if self.is_all() {
return vec!["all"];
}
let mut names: Vec<&'static str> = CAPABILITIES
.into_iter()
.filter(|capability| self.covers(capability))
.collect();
if self.process {
names.push("process");
}
names
}
pub fn parse(values: &[String]) -> Result<Grants, String> {
let mut grants = Grants::none();
for value in values {
match value.as_str() {
"none" if values.len() == 1 => {}
"none" => return Err("\"none\" cannot be combined with other capabilities".into()),
"process" => {
return Err(
"\"process\" is not a capability: a process runs as you, outside the sandbox, so only \"all\" allows it".into(),
);
}
other => match Grants::of(other) {
Some(granted) => grants = grants.union(granted),
None => return Err(format!("unknown capability \"{other}\"")),
},
}
}
Ok(grants)
}
pub fn union(self, other: Grants) -> Grants {
Grants {
filesystem: self.filesystem || other.filesystem,
process: self.process || other.process,
network: self.network || other.network,
environment: self.environment || other.environment,
session: self.session || other.session,
}
}
pub fn intersection(self, other: Grants) -> Grants {
Grants {
filesystem: self.filesystem && other.filesystem,
process: self.process && other.process,
network: self.network && other.network,
environment: self.environment && other.environment,
session: self.session && other.session,
}
}
}
pub fn tool_granted(grants: &Grants, tool: &RegisteredTool) -> bool {
let all = grants.is_all();
match tool.name() {
_ if tool.source.is_some() => all,
"bash" => grants.process,
"grep" | "find" => grants.filesystem && grants.process && grants.network,
"read" | "write" | "edit" | "ls" => grants.filesystem,
_ => all,
}
}
pub fn receives(kind: &str, grants: &Grants) -> bool {
match kind {
"before_provider_request" | "before_provider_headers" => grants.network,
_ => true,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn capabilities_parse() {
let fs = Grants {
filesystem: true,
..Grants::none()
};
let all_but_process = Grants {
process: false,
..Grants::default()
};
let process = "\"process\" is not a capability: a process runs as you, outside the sandbox, so only \"all\" allows it";
let cases: &[(&[&str], Result<Grants, &str>)] = &[
(&[], Ok(Grants::none())),
(&["none"], Ok(Grants::none())),
(&["all"], Ok(Grants::default())),
(&["fs", "all"], Ok(Grants::default())),
(&["fs"], Ok(fs)),
(&["fs", "net", "env", "session"], Ok(all_but_process)),
(&["process"], Err(process)),
(
&["none", "fs"],
Err("\"none\" cannot be combined with other capabilities"),
),
(&["exec"], Err("unknown capability \"exec\"")),
];
for (values, expected) in cases {
let values: Vec<String> = values.iter().map(|value| (*value).to_owned()).collect();
let expected = expected.map_err(str::to_owned);
assert_eq!(Grants::parse(&values), expected, "{values:?}");
}
let no_net = Grants {
network: false,
..Grants::default()
};
for (grants, name) in [
(Grants::default(), "all"),
(Grants::none(), "none"),
(fs, "fs"),
(all_but_process, "fs, net, env, session"),
(no_net, "fs, env, session, process"),
] {
assert_eq!(grants.to_string(), name);
for capability in CAPABILITIES {
let named = grants.names().contains(&capability) || grants.is_all();
assert_eq!(grants.covers(capability), named, "{name}: {capability}");
}
}
assert!(!Grants::default().covers("process"));
assert!(!Grants::default().covers("exec"));
let net = Grants {
network: true,
..Grants::none()
};
assert_eq!(
fs.union(net),
Grants {
filesystem: true,
network: true,
..Grants::none()
}
);
}
#[test]
fn grants_gate_tools() {
let env = crate::tools::ToolEnv {
cwd: std::path::PathBuf::new(),
runtime: std::sync::Arc::default(),
bin_dir: std::path::PathBuf::new(),
};
let fs = Grants {
filesystem: true,
..Grants::none()
};
let all_but_process = Grants {
process: false,
..Grants::default()
};
let cases = [
("read", false, [false, true, true, true]),
("write", false, [false, true, true, true]),
("edit", false, [false, true, true, true]),
("ls", false, [false, true, true, true]),
("bash", false, [false, false, false, true]),
("grep", false, [false, false, false, true]),
("find", false, [false, false, false, true]),
("read", true, [false, false, false, true]),
];
for (name, extension, expected) in cases {
let mut tool = crate::tools::builtin(name, &env).unwrap();
if extension {
tool.source = Some(yapi_types::rpc::SourceInfo::default());
}
let granted = [Grants::none(), fs, all_but_process, Grants::default()]
.map(|grants| tool_granted(&grants, &tool));
assert_eq!(granted, expected, "{name}, extension {extension}");
}
}
#[test]
fn provider_events_need_network() {
let net = Grants {
network: true,
..Grants::none()
};
let no_net = Grants {
network: false,
..Grants::default()
};
for (kind, heard_without_network) in [
("before_provider_request", false),
("before_provider_headers", false),
("after_provider_response", true),
("provider_stream_event", true),
("tool_call", true),
] {
assert!(receives(kind, &net), "{kind}");
assert_eq!(receives(kind, &no_net), heard_without_network, "{kind}");
}
}
}