boxmaker/crates/brokerd/src/args.rs

//! 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 serde::{Deserialize, Serialize};
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 {
        url: String,
        host: String,
    },
}

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 { url, .. } => {
                let value = HttpFetchArgs { url: 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
    }
}

/// The argument shape for one tool, decoded then re-serialised.
#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct ReadFileArgs {
    path: String,
}

#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct WriteFileArgs {
    path: String,
    content: String,
}

#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct ShellArgs {
    command: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    cwd: Option<String>,
}

#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct HttpFetchArgs {
    url: String,
}

/// 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 {
                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,
    }
}

/// A host label is 1 to 63 bytes of `a-z`, `0-9` or `-`, and neither starts nor ends with `-`.
fn valid_label(label: &str) -> bool {
    if !(1..=63).contains(&label.len()) {
        return false;
    }
    let bytes = label.as_bytes();
    if bytes.iter().next() == Some(&b'-') || bytes.iter().last() == Some(&b'-') {
        return false;
    }
    bytes
        .iter()
        .all(|&byte| matches!(byte, b'a'..=b'z' | b'0'..=b'9' | b'-'))
}

/// A host is 1 to 253 bytes of dot-separated labels, each a whole non-dot name, and the last label
/// starts with a letter. That last rule keeps out every spelling of an IPv4 address.
pub fn valid_host(host: &str) -> bool {
    if !(1..=253).contains(&host.len()) {
        return false;
    }
    let labels: Vec<&str> = host.split('.').collect();
    if labels.len() < 2 {
        return false;
    }
    for label in &labels {
        if !valid_label(label) {
            return false;
        }
    }
    match labels.iter().last() {
        Some(last) => matches!(last.bytes().next(), Some(byte) if byte.is_ascii_lowercase()),
        None => false,
    }
}

/// A host, or `*.` followed by a host. Nothing else.
pub fn valid_host_pattern(pattern: &str) -> bool {
    match pattern.strip_prefix("*.") {
        Some(base) => valid_host(base),
        None => valid_host(pattern),
    }
}

/// Does `pattern` match `host`? Without `*.` the strings must be equal; with `*.base` the host must
/// end in `.base` with something before the dot.
pub fn host_matches(pattern: &str, host: &str) -> bool {
    let base = match pattern.strip_prefix("*.") {
        Some(base) => base,
        None => return pattern == host,
    };
    match host.strip_suffix(base) {
        Some(before) => match before.strip_suffix('.') {
            Some(prefix) => !prefix.is_empty(),
            None => false,
        },
        None => false,
    }
}

/// 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
    }
}