subako-function-sdk 0.2.0

Write Subako function packages in Rust: typed functions become tools, and the program prints its own manifest.
Documentation
//! HTTP from a guest. The worker makes the request on the guest's behalf,
//! under the invocation's egress policy: which hosts it may reach and which
//! credential headers ride along to which origins. A file body or a file
//! response moves between the socket and the volume without passing through
//! guest memory.
//!
//! # ABI
//!
//! The import module `subako_http` holds two functions:
//!
//! - `request(ptr: i32, len: i32) -> i32` reads `len` bytes at `ptr` as a
//!   [`Request`] in JSON, performs it, and holds the answer. It returns the
//!   answer's length in bytes.
//! - `read_response(ptr: i32, len: i32) -> i32` copies the held answer to
//!   `ptr`, where `len` must be the length `request` returned, and forgets
//!   it. It returns 0.
//!
//! The answer is [`Answer`] in JSON: `{"ok": Response}` or
//! `{"error": HttpError}`. A negative return is a misuse of the ABI itself:
//! [`ABI_OUT_OF_BOUNDS`], [`ABI_NOTHING_HELD`], or [`ABI_LENGTH_MISMATCH`].

use std::collections::BTreeMap;

use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

use crate::ToolError;

/// The guest's memory range lies outside its memory.
pub const ABI_OUT_OF_BOUNDS: i32 = -1;
/// `read_response` ran with no answer held.
pub const ABI_NOTHING_HELD: i32 = -2;
/// `read_response`'s length is not the answer's.
pub const ABI_LENGTH_MISMATCH: i32 = -3;

/// The largest response body returned inline, in bytes.
pub const INLINE_MAX_BYTES: usize = 256 * 1024;

/// The most redirects one request follows.
pub const MAX_REDIRECTS: usize = 10;

/// The largest request JSON the worker reads from guest memory, in bytes. A
/// larger body travels as a file.
pub const REQUEST_MAX_BYTES: usize = 8 * 1024 * 1024;

/// One HTTP request. The `schemars` descriptions are what a model reads when
/// the `http` package serves this as its input.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct Request {
    #[serde(default)]
    pub method: Method,
    #[schemars(
        description = "An https URL; the host must be one the function's network setting allows."
    )]
    pub url: String,
    /// Sent as given on the first hop. Connection-level headers such as
    /// `host` and `content-length` are the worker's own and are dropped, and
    /// a header a credential sets for the hop's origin is refused.
    #[serde(default)]
    #[schemars(
        description = "Request headers by name. A header a configured credential sets for the \
                       URL's origin is refused; host, content-length and other connection \
                       headers are ignored."
    )]
    pub headers: BTreeMap<String, String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub body: Option<RequestBody>,
    #[serde(default)]
    pub output: ResponseOutput,
    /// Up to [`MAX_REDIRECTS`] hops, each checked like the first. A stored
    /// credential rides only to its own origin; on a hop to another origin the
    /// guest's own authorization, cookie, and proxy-authorization headers are
    /// dropped and the rest are sent as given.
    #[serde(default = "Request::default_follow_redirects")]
    #[schemars(description = format!(
        "Follow up to {MAX_REDIRECTS} redirects, each checked like the first request. Stored \
         credentials are attached per hop for that hop's origin only. On a redirect to \
         another origin, authorization, cookie and proxy-authorization headers you set are \
         dropped; other headers are sent as given."
    ))]
    pub follow_redirects: bool,
    /// Bounds the whole request, every hop and the body included. The
    /// invocation's own deadline bounds it regardless.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    #[schemars(
        range(min = 1),
        description = "Bound on the whole request; the call's own time limit applies regardless."
    )]
    pub timeout_ms: Option<u64>,
}

impl Request {
    /// A GET with every other field at its default.
    pub fn get(url: impl Into<String>) -> Self {
        Request {
            method: Method::Get,
            url: url.into(),
            headers: BTreeMap::new(),
            body: None,
            output: ResponseOutput::Inline,
            follow_redirects: Self::default_follow_redirects(),
            timeout_ms: None,
        }
    }

    fn default_follow_redirects() -> bool {
        true
    }
}

/// The methods a request may use. CONNECT, TRACE, and TRACK have no variant:
/// the first asks for a tunnel, and the others echo the request, credential
/// headers included, back to the guest.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "UPPERCASE")]
#[schemars(description = "HTTP method.")]
pub enum Method {
    #[default]
    Get,
    Head,
    Post,
    Put,
    Patch,
    Delete,
    Options,
}

impl Method {
    pub fn as_str(&self) -> &'static str {
        match self {
            Method::Get => "GET",
            Method::Head => "HEAD",
            Method::Post => "POST",
            Method::Put => "PUT",
            Method::Patch => "PATCH",
            Method::Delete => "DELETE",
            Method::Options => "OPTIONS",
        }
    }
}

impl std::fmt::Display for Method {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.as_str())
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
pub enum RequestBody {
    /// Sent as UTF-8, `text/plain; charset=utf-8` unless a header says
    /// otherwise.
    #[schemars(description = "Send text, as text/plain unless a content-type header is given.")]
    Text { text: String },
    /// Sent serialized, `application/json` unless a header says otherwise.
    #[schemars(
        description = "Send this value as JSON, as application/json unless a content-type \
                       header is given."
    )]
    Json { json: serde_json::Value },
    /// The volume file at `path`, streamed.
    #[schemars(description = "Stream a volume file as the body.")]
    File {
        #[schemars(
            description = "Absolute volume path under /tmp, or /data when the agent enables it."
        )]
        path: String,
    },
}

/// Where the response body goes.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
pub enum ResponseOutput {
    /// Into the answer as text, up to [`INLINE_MAX_BYTES`].
    #[default]
    #[schemars(description = format!(
        "Return the body as text in the reply (the default), up to {} KiB of UTF-8.",
        INLINE_MAX_BYTES / 1024
    ))]
    Inline,
    /// Into the volume file at `path`, overwritten, parents created, counted
    /// against the invocation's write budget.
    #[schemars(
        description = "Save the body to a volume file, bounded only by the call's write \
                       budget, and reply with its size. Use this for large or binary responses."
    )]
    File {
        #[schemars(description = "Absolute volume path, overwritten, parents created.")]
        path: String,
    },
}

/// The final hop's response.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
pub struct Response {
    pub status: u16,
    /// Lowercase names; repeated headers joined with `, `.
    #[schemars(description = "Lowercase names; repeated headers joined with a comma.")]
    pub headers: BTreeMap<String, String>,
    pub body: ResponseBody,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ResponseBody {
    Inline { text: String },
    File { path: String, bytes: u64 },
}

/// Why a request produced no response.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct HttpError {
    pub kind: HttpErrorKind,
    pub message: String,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum HttpErrorKind {
    /// The request does not parse, or names a method, URL, or header no
    /// request can carry.
    InvalidRequest,
    /// The function's network policy does not admit the host.
    HostNotAllowed,
    /// The worker refuses the destination: plain http, or an address in a
    /// private or reserved range.
    ForbiddenDestination,
    /// A header the request sets is one a credential sets for that origin.
    CredentialHeader,
    /// The session's vaults hold no single usable credential for an origin
    /// the function's grant lists.
    CredentialUnavailable,
    TooManyRedirects,
    Timeout,
    /// The response body is over [`INLINE_MAX_BYTES`].
    TooLarge,
    /// The response body is not UTF-8, so it cannot be returned inline.
    NotText,
    /// Saving the response would pass the invocation's write budget.
    WriteBudget,
    /// A body or output path is not a usable volume file.
    BadPath,
    /// The connection failed or the server broke the exchange.
    Connection,
    /// The import itself misbehaved, or the program runs outside the worker.
    Abi,
}

impl std::fmt::Display for HttpError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.message)
    }
}

impl std::error::Error for HttpError {}

impl From<HttpError> for ToolError {
    fn from(error: HttpError) -> Self {
        ToolError::new(error.message)
    }
}

/// What `request` holds for `read_response`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Answer {
    Ok(Response),
    Error(HttpError),
}

impl From<Answer> for Result<Response, HttpError> {
    fn from(answer: Answer) -> Self {
        match answer {
            Answer::Ok(response) => Ok(response),
            Answer::Error(error) => Err(error),
        }
    }
}

/// Performs `request` through the worker.
pub fn request(request: &Request) -> Result<Response, HttpError> {
    let abi = |message: String| HttpError {
        kind: HttpErrorKind::Abi,
        message,
    };
    // COVERAGE(never): a request is plain data, so it always serializes.
    let body = serde_json::to_vec(request).map_err(|e| abi(format!("request is not JSON: {e}")))?;
    // COVERAGE(untestable): off the worker the import always fails, so only a guest under wasmtime gets past it.
    let answer = imports::call(&body).map_err(abi)?;
    // COVERAGE(untestable): guest code runs under wasmtime, where the host's llvm-cov cannot see it.
    match serde_json::from_slice::<Answer>(&answer) {
        Ok(answer) => answer.into(),
        Err(e) => Err(abi(format!("subako_http answered with no answer: {e}"))),
    }
}

#[cfg(target_family = "wasm")]
mod imports {
    #[link(wasm_import_module = "subako_http")]
    unsafe extern "C" {
        fn request(ptr: *const u8, len: usize) -> i32;
        fn read_response(ptr: *mut u8, len: usize) -> i32;
    }

    // COVERAGE(untestable): guest code runs under wasmtime, where the host's llvm-cov cannot see it.
    pub(super) fn call(body: &[u8]) -> Result<Vec<u8>, String> {
        let failed = |code: i32| format!("subako_http returned {code}");
        // SAFETY: the host reads `body.len()` bytes from `body`'s start.
        let len = unsafe { request(body.as_ptr(), body.len()) };
        let len = usize::try_from(len).map_err(|_| failed(len))?;
        let mut answer = vec![0u8; len];
        // SAFETY: the host writes exactly `len` bytes into `answer`.
        match unsafe { read_response(answer.as_mut_ptr(), len) } {
            0 => Ok(answer),
            code => Err(failed(code)),
        }
    }
}

#[cfg(not(target_family = "wasm"))]
mod imports {
    pub(super) fn call(_: &[u8]) -> Result<Vec<u8>, String> {
        Err("subako_http exists only inside the function worker".to_string())
    }
}