wasm_proxy.rs

use anyhow::{Context, Result, anyhow};
use wasi::http::types::{
    ErrorCode, IncomingBody, IncomingRequest, OutgoingBody, OutgoingRequest, OutgoingResponse,
    ResponseOutparam, Scheme,
};
use wasi::io::streams::StreamError;

wasi::http::proxy::export!(TestComponent);

struct TestComponent;

fn forward_proxy(incoming: IncomingRequest) -> Result<OutgoingResponse> {
    let mut target = incoming.headers().get("X-Forward-To");
    if target.is_empty() {
        return Err(anyhow!("X-Forward-To is not set"));
    }
    let authority =
        String::from_utf8(target.pop().unwrap()).context("forwarding target is not a string")?;

    let req = OutgoingRequest::new(incoming.headers());

    req.set_authority(Some(&authority))
        .map_err(|()| anyhow!("set authority"))?;
    req.set_method(&incoming.method())
        .map_err(|()| anyhow!("set method"))?;
    req.set_scheme(Some(&Scheme::Http))
        .map_err(|()| anyhow!("set scheme"))?;
    req.set_path_with_query(incoming.path_with_query().as_deref())
        .map_err(|()| anyhow!("set path_with_query"))?;

    forward_body(
        incoming
            .consume()
            .map_err(|()| anyhow!("get incoming req body"))?,
        req.body().map_err(|()| anyhow!("get outgoing req body"))?,
    )?;

    let future_response = wasi::http::outgoing_handler::handle(req, None)?;

    let in_response = loop {
        if let Some(response) = future_response.get() {
            break response
                .map_err(|()| anyhow!("future response get"))?
                .context("incoming response error")?;
        } else {
            future_response.subscribe().block()
        }
    };

    let out_response = OutgoingResponse::new(in_response.headers());
    out_response
        .set_status_code(in_response.status())
        .map_err(|()| anyhow!("out response set status"))?;

    forward_body(
        in_response
            .consume()
            .map_err(|()| anyhow!("get incoming resp body"))?,
        out_response
            .body()
            .map_err(|()| anyhow!("get outgoing resp body"))?,
    )?;

    Ok(out_response)
}

fn forward_body(incoming: IncomingBody, outgoing: OutgoingBody) -> Result<()> {
    let outgoing_stream = outgoing.write().map_err(|_| anyhow!("out body's stream"))?;
    let incoming_stream = incoming
        .stream()
        .map_err(|()| anyhow!("in body's stream"))?;
    loop {
        match outgoing_stream.blocking_splice(&incoming_stream, u64::MAX) {
            Err(StreamError::Closed) => break,
            Err(StreamError::LastOperationFailed(err)) => return Err(err.into()),
            Ok(_) => continue,
        }
    }
    drop(incoming_stream);
    drop(outgoing_stream);
    IncomingBody::finish(incoming);
    OutgoingBody::finish(outgoing, None)?;
    Ok(())
}

impl wasi::exports::http::incoming_handler::Guest for TestComponent {
    fn handle(request: IncomingRequest, outparam: ResponseOutparam) {
        ResponseOutparam::set(
            outparam,
            forward_proxy(request)
                .map_err(|err| ErrorCode::InternalError(Some(format!("{err:?}")))),
        )
    }
}

fn main() {}