sausage/spikes/s5-http/src/main.rs

//! Runs the HTTP hook and verifies real WASI guests against a local test server.
use anyhow::Result;
use bytes::Bytes;
use http_body_util::{BodyExt, Full};
use hyper::{server::conn::http1, service::service_fn};
use hyper_util::rt::TokioIo;
use std::sync::{
    Arc,
    atomic::{AtomicUsize, Ordering},
};
use wasmtime::component::{Component, Linker, ResourceTable};
use wasmtime::{Engine, Store};
use wasmtime_wasi::{WasiCtx, WasiCtxView, WasiView};
use wasmtime_wasi_http::{WasiBody, WasiHttpCtx, WasiHttpCtxView, WasiHttpHooks, WasiHttpView};

type Io = Box<dyn Future<Output = Result<(), wasmtime_wasi_http::Error>> + Send>;
struct Item {
    name: &'static str,
    scheme: &'static str,
    hosts: Vec<&'static str>,
    prefixes: Vec<&'static str>,
    methods: Option<Vec<http::Method>>,
    free: bool,
}
impl Item {
    fn matches(&self, request: &http::Request<WasiBody>) -> bool {
        let uri = request.uri();
        uri.scheme_str() == Some(self.scheme)
            && self
                .hosts
                .iter()
                .any(|pattern| host_matches(pattern, uri.host().unwrap_or("")))
            && self.prefixes.iter().any(|p| uri.path().starts_with(p))
            && self
                .methods
                .as_ref()
                .is_none_or(|m| m.contains(request.method()))
    }
}
fn host_matches(pattern: &str, host: &str) -> bool {
    match pattern.strip_prefix("*.") {
        Some(suffix) => host
            .strip_suffix(suffix)
            .is_some_and(|p| p.ends_with('.') && p.len() > 1),
        None => pattern == host,
    }
}
struct Hooks {
    id: String,
    grant: bool,
    token: Arc<String>,
    items: Arc<Vec<Item>>,
    outgoing: Arc<AtomicUsize>,
}
fn response(
    status: u16,
    text: String,
) -> wasmtime_wasi_http::Result<(http::Response<WasiBody>, Io)> {
    let body = Full::new(Bytes::from(text))
        .map_err(|never| match never {})
        .boxed_unsync();
    Ok((
        http::Response::builder().status(status).body(body).unwrap(),
        Box::new(async { Ok(()) }),
    ))
}
impl WasiHttpHooks for Hooks {
    fn send_request(
        &mut self,
        mut request: http::Request<WasiBody>,
        options: Option<wasmtime_wasi_http::RequestOptions>,
        fut: Io,
    ) -> Box<dyn Future<Output = wasmtime_wasi_http::Result<(http::Response<WasiBody>, Io)>> + Send>
    {
        drop(fut);
        if request.uri().host() == Some("host.sausage") {
            let result = response(200, format!("caller={}", self.id));
            return Box::new(async move {
                request.into_body().collect().await?;
                result
            });
        }
        let matches: Vec<_> = self.items.iter().filter(|i| i.matches(&request)).collect();
        if matches.len() > 1 {
            return Box::new(async move {
                request.into_body().collect().await?;
                Err(wasmtime_wasi_http::Error::InternalError(Some(
                    "two credential items match".into(),
                )))
            });
        }
        if let Some(item) = matches.first() {
            if !self.grant && !item.free {
                let result = response(403, format!("grant required: {}", item.name));
                return Box::new(async move {
                    request.into_body().collect().await?;
                    result
                });
            }
            // Stand-in for the isolated JavaScript transform. Only the host owns the token.
            let transform = |r: &mut http::Request<WasiBody>| {
                r.headers_mut().insert(
                    http::header::AUTHORIZATION,
                    format!("Bearer {}", self.token).parse().unwrap(),
                );
            };
            transform(&mut request);
        }
        self.outgoing.fetch_add(1, Ordering::SeqCst);
        Box::new(async move {
            // Hyper drops unknown-length GET bodies before the guest can finish its WASI stream.
            // Complete the body first. Collected preserves its data and trailers.
            let (parts, body) = request.into_parts();
            let body = body.collect().await?.map_err(|never| match never {});
            let request = http::Request::from_parts(parts, body);
            let (res, io) = wasmtime_wasi_http::default_send_request(request, options).await?;
            Ok((res.map(BodyExt::boxed_unsync), Box::new(io) as Io))
        })
    }
}
struct State {
    ctx: WasiCtx,
    table: ResourceTable,
    http: WasiHttpCtx,
    hooks: Hooks,
}
impl WasiView for State {
    fn ctx(&mut self) -> WasiCtxView<'_> {
        WasiCtxView {
            ctx: &mut self.ctx,
            table: &mut self.table,
        }
    }
}
impl WasiHttpView for State {
    fn http(&mut self) -> WasiHttpCtxView<'_> {
        WasiHttpCtxView {
            ctx: &mut self.http,
            table: &mut self.table,
            hooks: &mut self.hooks,
        }
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    let component_path = std::env::args().nth(1).expect("guest component path");
    std::fs::create_dir_all("/tmp/s5-http/world")?;
    std::fs::write("/tmp/s5-http/world/public.txt", "public fixture")?;
    // A fresh host-only token is absent from guest source, argv, env, and preopened files.
    let token = Arc::new(format!(
        "s5-secret-{}",
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)?
            .as_nanos()
    ));
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
    let port = listener.local_addr()?.port();
    let received = Arc::new(AtomicUsize::new(0));
    let count = received.clone();
    let server_token = token.clone();
    let server = tokio::spawn(async move {
        loop {
            let (socket, _) = listener.accept().await.unwrap();
            let count = count.clone();
            let token = server_token.clone();
            let service = service_fn(move |request: hyper::Request<hyper::body::Incoming>| {
                let count = count.clone();
                let token = token.clone();
                async move {
                    let (parts, body) = request.into_parts();
                    // Read through the terminal chunk before answering and closing the connection.
                    body.collect().await.unwrap();
                    assert_eq!(parts.headers["x-guest"], "unchanged");
                    let text = match parts.headers.get(http::header::AUTHORIZATION) {
                        Some(value) => {
                            assert_eq!(value.to_str().unwrap(), format!("Bearer {}", token));
                            "authorization=present; token=correct; x-guest=unchanged"
                        }
                        None => "authorization=absent; x-guest=unchanged",
                    };
                    count.fetch_add(1, Ordering::SeqCst);
                    Ok::<_, std::convert::Infallible>(http::Response::new(Full::new(
                        Bytes::from_static(text.as_bytes()),
                    )))
                }
            });
            http1::Builder::new()
                .keep_alive(false)
                .serve_connection(TokioIo::new(socket), service)
                .await
                .unwrap();
        }
    });
    let items = Arc::new(vec![
        Item {
            name: "test-token",
            scheme: "http",
            hosts: vec!["127.0.0.1"],
            prefixes: vec!["/private"],
            methods: Some(vec![http::Method::GET]),
            free: false,
        },
        Item {
            name: "collision",
            scheme: "http",
            hosts: vec!["127.0.0.1"],
            prefixes: vec!["/private/overlap"],
            methods: None,
            free: false,
        },
        Item {
            name: "free-token",
            scheme: "http",
            hosts: vec!["127.0.0.1"],
            prefixes: vec!["/free"],
            methods: None,
            free: true,
        },
    ]);
    let engine = Engine::default();
    let component = Component::from_file(&engine, component_path)?;
    let mut linker = Linker::<State>::new(&engine);
    wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
    wasmtime_wasi_http::p2::add_only_http_to_linker_async(&mut linker)?;
    let outgoing = Arc::new(AtomicUsize::new(0));
    let cases = [
        (
            "identity-a",
            "http://host.sausage/whoami".to_string(),
            "GET",
            false,
            "caller=instance-1",
            0,
        ),
        (
            "identity-b",
            "http://host.sausage/whoami".to_string(),
            "GET",
            false,
            "caller=instance-2",
            0,
        ),
        (
            "granted",
            format!("http://127.0.0.1:{port}/private"),
            "GET",
            true,
            "authorization=present; token=correct",
            1,
        ),
        (
            "unmatched",
            format!("http://127.0.0.1:{port}/public"),
            "GET",
            false,
            "authorization=absent",
            1,
        ),
        (
            "denied",
            format!("http://127.0.0.1:{port}/private"),
            "GET",
            false,
            "status=403\nresponse-headers={}\nbody=grant required: test-token",
            0,
        ),
        (
            "ambiguous",
            format!("http://127.0.0.1:{port}/private/overlap"),
            "GET",
            true,
            "two credential items match",
            0,
        ),
        (
            "method-mismatch",
            format!("http://127.0.0.1:{port}/private"),
            "POST",
            false,
            "authorization=absent",
            1,
        ),
        (
            "free-use",
            format!("http://127.0.0.1:{port}/free"),
            "POST",
            false,
            "authorization=present; token=correct",
            1,
        ),
    ];
    for (index, (name, url, method, grant, expected, sockets)) in cases.into_iter().enumerate() {
        let before = outgoing.load(Ordering::SeqCst);
        let stdout = wasmtime_wasi::p2::pipe::MemoryOutputPipe::new(16384);
        let stderr = wasmtime_wasi::p2::pipe::MemoryOutputPipe::new(16384);
        let mut ctx = WasiCtx::builder();
        ctx.args(&["s5-curl", &url, method])
            .env("PUBLIC", "visible")
            .stdout(stdout.clone())
            .stderr(stderr.clone());
        ctx.preopened_dir("/tmp/s5-http/world", "/", wasmtime_wasi::FsPerms::ReadOnly)?;
        let state = State {
            ctx: ctx.build(),
            table: ResourceTable::new(),
            http: WasiHttpCtx::new(),
            hooks: Hooks {
                id: format!("instance-{}", index + 1),
                grant,
                token: token.clone(),
                items: items.clone(),
                outgoing: outgoing.clone(),
            },
        };
        let mut store = Store::new(&engine, state);
        let command = wasmtime_wasi::p2::bindings::Command::instantiate_async(
            &mut store, &component, &linker,
        )
        .await?;
        assert!(command.wasi_cli_run().call_run(&mut store).await?.is_ok());
        let out = String::from_utf8(stdout.contents().to_vec())?;
        let err = String::from_utf8(stderr.contents().to_vec())?;
        assert!(
            !out.contains(token.as_str()) && !err.contains(token.as_str()),
            "secret leaked"
        );
        assert!(
            out.contains("environment=[(\"PUBLIC\", \"visible\")]"),
            "{out}"
        );
        assert!(out.contains("files=[\"public.txt\"]"), "{out}");
        assert!(out.contains("file=public fixture"), "{out}");
        assert!(
            out.contains("request-headers={\"x-guest\": \"unchanged\"}"),
            "{out}"
        );
        assert!(out.contains(expected), "{name}: {out}\n{err}");
        assert_eq!(outgoing.load(Ordering::SeqCst) - before, sockets, "{name}");
        println!("PASS {name}: {expected}; transport calls={sockets}; guest secret absent");
    }
    assert_eq!(received.load(Ordering::SeqCst), 4);
    server.abort();
    println!("PASS all 8 guest cases; local server received exactly 4 requests");
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn rule_dimensions() {
        let item = Item {
            name: "test",
            scheme: "http",
            hosts: vec!["*.example.test"],
            prefixes: vec!["/api"],
            methods: Some(vec![http::Method::GET]),
            free: false,
        };
        for (url, method, expected) in [
            ("http://api.example.test/api/v1", "GET", true),
            ("https://api.example.test/api/v1", "GET", false),
            ("http://example.test/api/v1", "GET", false),
            ("http://api.example.test/public", "GET", false),
            ("http://api.example.test/api/v1", "POST", false),
        ] {
            let body = Full::new(Bytes::new())
                .map_err(|n| match n {})
                .boxed_unsync();
            let request = http::Request::builder()
                .uri(url)
                .method(method)
                .body(body)
                .unwrap();
            assert_eq!(item.matches(&request), expected, "{url} {method}");
        }
    }
    #[test]
    fn host_patterns() {
        assert!(host_matches("*.example.test", "api.example.test"));
        assert!(host_matches("*.example.test", "a.b.example.test"));
        assert!(!host_matches("*.example.test", "example.test"));
        assert!(!host_matches("*.example.test", "badexample.test"));
        assert!(host_matches("example.test", "example.test"));
        assert!(!host_matches("example.test", "api.example.test"));
    }
}