302 lines
9.7 KiB
Rust
302 lines
9.7 KiB
Rust
//! Tool arguments and the form checks for paths, hosts and URLs.
|
|
//!
|
|
//! `parse` turns the model's argument string into a typed value; the free functions say whether a
|
|
//! path, a host or a URL is well formed. The module is pure: no I/O, no clock. Everything it reads
|
|
//! was written by the model, so it is treated as hostile.
|
|
|
|
use proto::tools::{HttpFetchArgs, ReadFileArgs, ShellArgs, WriteFileArgs};
|
|
use std::fmt;
|
|
|
|
/// Maximum length, in bytes, of a path.
|
|
pub const MAX_PATH: usize = 4096;
|
|
/// Maximum length, in bytes, of a URL.
|
|
pub const MAX_URL: usize = 2048;
|
|
|
|
/// The four tools the broker can dispatch to.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum ToolName {
|
|
ReadFile,
|
|
WriteFile,
|
|
Shell,
|
|
HttpFetch,
|
|
}
|
|
|
|
impl ToolName {
|
|
/// The four tools, in the order listed in the interface.
|
|
pub const ALL: [ToolName; 4] = [
|
|
ToolName::ReadFile,
|
|
ToolName::WriteFile,
|
|
ToolName::Shell,
|
|
ToolName::HttpFetch,
|
|
];
|
|
|
|
/// Parse a tool name as written by the model.
|
|
pub fn parse(name: &str) -> Option<ToolName> {
|
|
match name {
|
|
"read_file" => Some(ToolName::ReadFile),
|
|
"write_file" => Some(ToolName::WriteFile),
|
|
"shell" => Some(ToolName::Shell),
|
|
"http_fetch" => Some(ToolName::HttpFetch),
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
/// The name as written by the model.
|
|
pub fn as_str(self) -> &'static str {
|
|
match self {
|
|
ToolName::ReadFile => "read_file",
|
|
ToolName::WriteFile => "write_file",
|
|
ToolName::Shell => "shell",
|
|
ToolName::HttpFetch => "http_fetch",
|
|
}
|
|
}
|
|
}
|
|
|
|
/// A parsed set of tool arguments, one variant per tool.
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub enum ToolArgs {
|
|
ReadFile {
|
|
path: String,
|
|
},
|
|
WriteFile {
|
|
path: String,
|
|
content: String,
|
|
},
|
|
Shell {
|
|
command: String,
|
|
cwd: Option<String>,
|
|
},
|
|
HttpFetch(FetchUrl),
|
|
}
|
|
|
|
/// A URL that passed the checks, and its host. Only `parse` makes one, so the host policy matched
|
|
/// is always the host of the URL the tool fetches.
|
|
///
|
|
/// ```compile_fail
|
|
/// let _ = brokerd::args::FetchUrl {
|
|
/// url: "https://evil.example/".to_string(),
|
|
/// host: "example.com".to_string(),
|
|
/// };
|
|
/// ```
|
|
///
|
|
/// ```
|
|
/// let args = brokerd::args::parse(brokerd::args::ToolName::HttpFetch, r#"{"url":"https://example.com/a"}"#);
|
|
/// let Ok(brokerd::args::ToolArgs::HttpFetch(target)) = args else { panic!("{args:?}") };
|
|
/// assert_eq!((target.url(), target.host()), ("https://example.com/a", "example.com"));
|
|
/// ```
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub struct FetchUrl {
|
|
url: String, // private
|
|
host: String, // private
|
|
}
|
|
|
|
impl FetchUrl {
|
|
/// The URL that passed the checks.
|
|
pub fn url(&self) -> &str {
|
|
&self.url
|
|
}
|
|
/// The host of that URL, as `parse` derived it.
|
|
pub fn host(&self) -> &str {
|
|
&self.host
|
|
}
|
|
}
|
|
|
|
impl ToolArgs {
|
|
/// The tool these arguments belong to.
|
|
pub fn tool(&self) -> ToolName {
|
|
match self {
|
|
ToolArgs::ReadFile { .. } => ToolName::ReadFile,
|
|
ToolArgs::WriteFile { .. } => ToolName::WriteFile,
|
|
ToolArgs::Shell { .. } => ToolName::Shell,
|
|
ToolArgs::HttpFetch(_) => ToolName::HttpFetch,
|
|
}
|
|
}
|
|
|
|
/// The arguments written out again, so two spellings of one value look the same.
|
|
///
|
|
/// Fields come out in the table's order, an absent `cwd` is left out, and `host` is never
|
|
/// written because it is not an argument. A serialization failure yields `"{}"`.
|
|
pub fn canonical_json(&self) -> String {
|
|
match self {
|
|
ToolArgs::ReadFile { path } => {
|
|
let value = ReadFileArgs { path: path.clone() };
|
|
serde_json::to_string(&value).unwrap_or_else(|_| "{}".to_string())
|
|
}
|
|
ToolArgs::WriteFile { path, content } => {
|
|
let value = WriteFileArgs {
|
|
path: path.clone(),
|
|
content: content.clone(),
|
|
};
|
|
serde_json::to_string(&value).unwrap_or_else(|_| "{}".to_string())
|
|
}
|
|
ToolArgs::Shell { command, cwd } => {
|
|
let value = ShellArgs {
|
|
command: command.clone(),
|
|
cwd: cwd.clone(),
|
|
};
|
|
serde_json::to_string(&value).unwrap_or_else(|_| "{}".to_string())
|
|
}
|
|
ToolArgs::HttpFetch(target) => {
|
|
let value = HttpFetchArgs {
|
|
url: target.url.clone(),
|
|
};
|
|
serde_json::to_string(&value).unwrap_or_else(|_| "{}".to_string())
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Why `parse` refused a set of arguments.
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub enum ArgsError {
|
|
/// The arguments were not the right shape for the tool.
|
|
Shape(String),
|
|
/// A path was not well formed.
|
|
Path(String),
|
|
/// A URL gave no valid host.
|
|
Url(String),
|
|
}
|
|
|
|
impl fmt::Display for ArgsError {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
match self {
|
|
ArgsError::Shape(text) => write!(f, "invalid argument shape: {text}"),
|
|
ArgsError::Path(text) => write!(f, "invalid path: {text}"),
|
|
ArgsError::Url(text) => write!(f, "invalid url: {text}"),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl std::error::Error for ArgsError {
|
|
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// Decode `arguments` for `tool` into a typed value.
|
|
pub fn parse(tool: ToolName, arguments: &str) -> Result<ToolArgs, ArgsError> {
|
|
match tool {
|
|
ToolName::ReadFile => {
|
|
let value: ReadFileArgs = match serde_json::from_str(arguments) {
|
|
Ok(value) => value,
|
|
Err(error) => return Err(ArgsError::Shape(error.to_string())),
|
|
};
|
|
if !valid_path(&value.path) {
|
|
return Err(ArgsError::Path(value.path));
|
|
}
|
|
Ok(ToolArgs::ReadFile { path: value.path })
|
|
}
|
|
ToolName::WriteFile => {
|
|
let value: WriteFileArgs = match serde_json::from_str(arguments) {
|
|
Ok(value) => value,
|
|
Err(error) => return Err(ArgsError::Shape(error.to_string())),
|
|
};
|
|
if !valid_path(&value.path) {
|
|
return Err(ArgsError::Path(value.path));
|
|
}
|
|
Ok(ToolArgs::WriteFile {
|
|
path: value.path,
|
|
content: value.content,
|
|
})
|
|
}
|
|
ToolName::Shell => {
|
|
let value: ShellArgs = match serde_json::from_str(arguments) {
|
|
Ok(value) => value,
|
|
Err(error) => return Err(ArgsError::Shape(error.to_string())),
|
|
};
|
|
if let Some(cwd) = &value.cwd
|
|
&& !valid_path(cwd)
|
|
{
|
|
return Err(ArgsError::Path(cwd.clone()));
|
|
}
|
|
Ok(ToolArgs::Shell {
|
|
command: value.command,
|
|
cwd: value.cwd,
|
|
})
|
|
}
|
|
ToolName::HttpFetch => {
|
|
let value: HttpFetchArgs = match serde_json::from_str(arguments) {
|
|
Ok(value) => value,
|
|
Err(error) => return Err(ArgsError::Shape(error.to_string())),
|
|
};
|
|
let host = match url_host(&value.url) {
|
|
Some(host) => host.to_string(),
|
|
None => return Err(ArgsError::Url(value.url)),
|
|
};
|
|
Ok(ToolArgs::HttpFetch(FetchUrl {
|
|
url: value.url,
|
|
host,
|
|
}))
|
|
}
|
|
}
|
|
}
|
|
|
|
/// A path is well formed if it is at most `MAX_PATH` bytes, has no NUL, starts with `/`, and every
|
|
/// component is a whole, non-dot name. The root `/` alone is valid.
|
|
pub fn valid_path(path: &str) -> bool {
|
|
if !(1..=MAX_PATH).contains(&path.len()) || path.contains('\0') {
|
|
return false;
|
|
}
|
|
let rest = match path.strip_prefix('/') {
|
|
Some(rest) => rest,
|
|
None => return false,
|
|
};
|
|
if rest.is_empty() {
|
|
return true;
|
|
}
|
|
for component in rest.split('/') {
|
|
if component.is_empty() || component == "." || component == ".." {
|
|
return false;
|
|
}
|
|
}
|
|
true
|
|
}
|
|
|
|
/// Is `path` inside `grant_path`, by whole components rather than by bytes?
|
|
pub fn inside(grant_path: &str, path: &str) -> bool {
|
|
match path.strip_prefix(grant_path) {
|
|
Some(rest) => rest.is_empty() || rest.starts_with('/') || grant_path == "/",
|
|
None => false,
|
|
}
|
|
}
|
|
|
|
/// Moved to `proto::hosts` in M3b, so `toolkit` checks hosts with the same rules.
|
|
pub use proto::hosts::{host_matches, valid_host, valid_host_pattern};
|
|
|
|
/// A character that may appear in a host name.
|
|
fn host_char(c: char) -> bool {
|
|
matches!(c, 'a'..='z' | '0'..='9' | '.' | '-')
|
|
}
|
|
|
|
/// The host of a URL, if the URL is well formed. See the task's `url_host` rules for the order.
|
|
pub fn url_host(url: &str) -> Option<&str> {
|
|
if url.len() > MAX_URL {
|
|
return None;
|
|
}
|
|
let without_scheme = url.strip_prefix("https://")?;
|
|
// Uppercase is not allowed anywhere in a valid URL, so the host is read as written.
|
|
let end = match without_scheme.find(|c: char| !host_char(c)) {
|
|
Some(end) => end,
|
|
None => without_scheme.len(),
|
|
};
|
|
let (host, rest) = without_scheme.split_at_checked(end)?;
|
|
if !valid_host(host) {
|
|
return None;
|
|
}
|
|
// A trailing `:443` is allowed; any other port, or anything else, must be the end or a path.
|
|
let after = match rest.strip_prefix(":443") {
|
|
Some(after) => after,
|
|
None => rest,
|
|
};
|
|
let remainder = if after.is_empty() {
|
|
after
|
|
} else {
|
|
after.strip_prefix('/')?
|
|
};
|
|
if remainder.bytes().all(|byte| (0x21..=0x7e).contains(&byte)) {
|
|
Some(host)
|
|
} else {
|
|
None
|
|
}
|
|
}
|