//! Per-process credential hook and pinned-address HTTP transport.
//! Transport derives from Wasmtime 49's default sender, with a policy check before connect.
use crate::cred::{Caller, Cred, inject::Outcome, script::Script};
use http_body_util::{BodyExt, Full};
use std::sync::Arc;
use wasmtime_wasi_http::{Error, RequestOptions};
use wasmtime_wasi_http::{WasiBody, WasiHttpHooks};
type Io = Box<dyn Future<Output = Result<(), Error>> + Send>;
pub type Router = Arc<crate::api::Api>;
pub(super) struct Hooks {
    pub caller: Caller,
    pub cred: Arc<Cred>,
    pub script: Arc<dyn Script>,
    pub policy: Arc<super::net::NetPolicy>,
    pub router: Router,
}
fn response(response: Response<Bytes>) -> (Response<WasiBody>, Io) {
    (
        response.map(|b| Full::new(b).map_err(|n| match n {}).boxed_unsync()),
        Box::new(async { Ok(()) }),
    )
}
impl WasiHttpHooks for Hooks {
    // Admit every scheme until the URI supplies the host name.
    fn is_supported_scheme(&mut self, _: &http::uri::Scheme) -> bool {
        true
    }
    fn send_request(
        &mut self,
        request: Request<WasiBody>,
        options: Option<RequestOptions>,
        fut: Io,
    ) -> Box<dyn Future<Output = Result<(Response<WasiBody>, Io), Error>> + Send> {
        drop(fut);
        let caller = self.caller.clone();
        let cred = self.cred.clone();
        let script = self.script.clone();
        let policy = self.policy.clone();
        let router = self.router.clone();
        Box::new(async move {
            let (parts, mut body) = request.into_parts();
            if parts
                .uri
                .host()
                .is_some_and(|h| h.eq_ignore_ascii_case("host.sausage"))
            {
                let bytes = body.collect().await?.to_bytes();
                let request = Request::from_parts(parts, bytes);
                return Ok((
                    crate::api::dispatch(&router, caller, request).await,
                    Box::new(async { Ok(()) }) as Io,
                ));
            }
            if !matches!(parts.uri.scheme_str(), Some("http" | "https")) {
                while let Some(frame) = body.frame().await {
                    frame?;
                }
                return Err(Error::HttpRequestUriInvalid);
            }
            let result =
                tokio::task::spawn_blocking(move || cred.inject(parts, &caller, script.as_ref()))
                    .await;
            let mut parts = match result {
                Ok(Ok(Outcome::Forward(parts))) => parts,
                outcome => {
                    // Drain without retaining uploads, so the guest can finish its output stream.
                    while let Some(frame) = body.frame().await {
                        frame?;
                    }
                    let (status, value) = match outcome {
                        Ok(Ok(Outcome::Denied { item, request })) => (
                            403,
                            serde_json::json!({"message": format!("Access to {item} needs approval. Request {request} awaits a decision."), "item": item, "request": request}),
                        ),
                        Ok(Err(e)) => (502, serde_json::json!({"error": e.to_string()})),
                        _ => (
                            502,
                            serde_json::json!({"error": "credential worker failed"}),
                        ),
                    };
                    return Ok(response(
                        Response::builder()
                            .status(status)
                            .header("content-type", "application/json")
                            .body(Bytes::from(serde_json::to_vec(&value).unwrap()))
                            .unwrap(),
                    ));
                }
            };
            let authority = parts.uri.authority().ok_or(Error::HttpRequestUriInvalid)?;
            if authority.as_str().contains('@') {
                return Err(Error::HttpRequestUriInvalid);
            }
            parts.headers.insert(
                http::header::HOST,
                authority
                    .as_str()
                    .parse()
                    .map_err(|_| Error::HttpRequestUriInvalid)?,
            );
            for value in parts.headers.values_mut() {
                value.set_sensitive(true);
            }
            // Hyper otherwise discards unknown-length GET bodies before WASI finishes them.
            if !parts.headers.contains_key(http::header::CONTENT_LENGTH) {
                parts.headers.insert(
                    http::header::TRANSFER_ENCODING,
                    http::HeaderValue::from_static("chunked"),
                );
            }
            let (res, io) =
                send_request(Request::from_parts(parts, body), options, &policy).await?;
            Ok((res.map(BodyExt::boxed_unsync), Box::new(io) as Io))
        })
    }
}
use bytes::Bytes;
use core::future::poll_fn;
use core::pin::{Pin, pin};
use core::task::{Poll, ready};
use http::uri::Scheme;
use http::{Request, Response};
use http_body::Body;
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::net::TcpStream;

trait TokioStream: AsyncRead + AsyncWrite + Send + Sync + Unpin + 'static {
    fn boxed(self) -> Box<dyn TokioStream>
    where
        Self: Sized,
    {
        Box::new(self)
    }
}
impl<T> TokioStream for T where T: AsyncRead + AsyncWrite + Send + Sync + Unpin + 'static {}

/// Preserve Wasmtime's HTTP/1 transport and request deadlines, but connect only to checked addresses.
/// The returned future drives the connection after the response headers arrive.
async fn send_request<T: Body<Data = Bytes, Error = Error> + Send + 'static>(
    mut req: Request<T>,
    options: Option<RequestOptions>,
    policy: &super::net::NetPolicy,
) -> Result<
    (
        Response<impl Body<Data = Bytes, Error = Error> + use<T>>,
        impl Future<Output = Result<(), Error>> + Send + use<T>,
    ),
    Error,
> {
    let uri = req.uri().clone();
    let authority = uri.authority().ok_or(Error::HttpRequestUriInvalid)?;
    let use_tls = uri.scheme() == Some(&Scheme::HTTPS);
    let authority = if authority.port().is_some() {
        authority.to_string()
    } else {
        let port = if use_tls { 443 } else { 80 };
        format!("{authority}:{port}")
    };

    let connect_timeout = options
        .and_then(|r| r.connect_timeout)
        .unwrap_or(Duration::from_secs(600));

    let first_byte_timeout = options
        .and_then(|r| r.first_byte_timeout)
        .unwrap_or(Duration::from_secs(600));

    let between_bytes_timeout = options
        .and_then(|r| r.between_bytes_timeout)
        .unwrap_or(Duration::from_secs(600));

    let stream = match tokio::time::timeout(connect_timeout, async {
        let addresses = policy.resolve(&uri).await?;
        TcpStream::connect(addresses.as_slice())
            .await
            .map_err(Error::Connect)
    })
    .await
    {
        Ok(stream) => stream?,
        Err(..) => return Err(Error::ConnectionTimeout),
    };
    let stream = if use_tls {
        // derived from https://github.com/rustls/rustls/blob/main/examples/src/bin/simpleclient.rs
        let root_cert_store = rustls::RootCertStore {
            roots: webpki_roots::TLS_SERVER_ROOTS.into(),
        };
        let config = rustls::ClientConfig::builder()
            .with_root_certificates(root_cert_store)
            .with_no_client_auth();
        let connector = tokio_rustls::TlsConnector::from(std::sync::Arc::new(config));
        let domain = tls_server_name(&authority)?;
        let stream = connector
            .connect(domain, stream)
            .await
            .map_err(Error::Tls)?;
        stream.boxed()
    } else {
        stream.boxed()
    };
    let (mut sender, conn) = tokio::time::timeout(
        connect_timeout,
        hyper::client::conn::http1::Builder::new().handshake(hyper_util::rt::TokioIo::new(stream)),
    )
    .await
    .map_err(|_| Error::ConnectionTimeout)??;

    // Send origin-form URIs, because this connection does not use a proxy.
    *req.uri_mut() = http::Uri::builder()
        .path_and_query(
            req.uri()
                .path_and_query()
                .map(|p| p.as_str())
                .unwrap_or("/"),
        )
        .build()
        .expect("comes from valid request");

    let send = async move {
        use core::task::Context;

        /// Wrapper around [hyper::body::Incoming] used to
        /// account for request option timeout configuration
        struct IncomingResponseBody {
            incoming: hyper::body::Incoming,
            timeout: tokio::time::Interval,
        }
        impl http_body::Body for IncomingResponseBody {
            type Data = <hyper::body::Incoming as http_body::Body>::Data;
            type Error = Error;

            fn poll_frame(
                mut self: Pin<&mut Self>,
                cx: &mut Context<'_>,
            ) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
                match Pin::new(&mut self.as_mut().incoming).poll_frame(cx) {
                    Poll::Ready(None) => Poll::Ready(None),
                    Poll::Ready(Some(Err(err))) => {
                        let err = if err.is_timeout() {
                            Error::HttpResponseTimeout
                        } else {
                            Error::from(err)
                        };
                        Poll::Ready(Some(Err(err)))
                    }
                    Poll::Ready(Some(Ok(frame))) => {
                        self.timeout.reset();
                        Poll::Ready(Some(Ok(frame)))
                    }
                    Poll::Pending => {
                        ready!(self.timeout.poll_tick(cx));
                        Poll::Ready(Some(Err(Error::ConnectionReadTimeout)))
                    }
                }
            }
            fn is_end_stream(&self) -> bool {
                self.incoming.is_end_stream()
            }
            fn size_hint(&self) -> http_body::SizeHint {
                self.incoming.size_hint()
            }
        }

        let res = tokio::time::timeout(first_byte_timeout, sender.send_request(req))
            .await
            .map_err(|_| Error::ConnectionReadTimeout)?
            .map_err(Error::from)?;
        // Tokio panics on an `interval` with a zero duration, so clamp it to a
        // 1ns timeout which is sort of effectively zero anyway.
        let between_bytes_timeout = if between_bytes_timeout == Duration::ZERO {
            Duration::new(0, 1)
        } else {
            between_bytes_timeout
        };
        let mut timeout = tokio::time::interval(between_bytes_timeout);
        timeout.reset();
        Ok(res.map(|incoming| IncomingResponseBody { incoming, timeout }))
    };
    let mut send = pin!(send);
    let mut conn = Some(conn);
    // Wait for response while driving connection I/O
    let res = poll_fn(|cx| match send.as_mut().poll(cx) {
        Poll::Ready(Ok(res)) => Poll::Ready(Ok(res)),
        Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
        Poll::Pending => {
            // Response is not ready, poll `hyper` connection to drive I/O if it has not completed yet
            let Some(fut) = conn.as_mut() else {
                // `hyper` connection already completed
                return Poll::Pending;
            };
            let res = ready!(Pin::new(fut).poll(cx));
            // `hyper` connection completed, record that to prevent repeated poll
            conn = None;
            match res {
                // `hyper` connection has successfully completed, optimistically poll for response
                Ok(()) => send.as_mut().poll(cx),
                // `hyper` connection has failed, return the error
                Err(err) => Poll::Ready(Err(Error::from(err))),
            }
        }
    })
    .await?;
    Ok((res, async move {
        let Some(conn) = conn.take() else {
            // `hyper` connection has already completed
            return Ok(());
        };
        if let Err(err) = conn.await {
            if err.is_timeout() {
                return Err(Error::HttpResponseTimeout);
            }
            return Err(err.into());
        }
        Ok(())
    }))
}

/// Resolve the rustls [`ServerName`] used for TLS certificate verification from
/// an outbound request `authority`.
///
/// `authority` is in `host:port` form, where an IPv6 `host` is wrapped in
/// brackets (for example `[::1]:443`). An IP literal is recognized by parsing
/// the whole authority as a [`SocketAddr`]; this handles the bracketed IPv6
/// form, which splitting on the first `:` would truncate. Anything else is
/// treated as a host name, with the port stripped off before it is handed to
/// rustls.
///
/// [`ServerName`]: rustls::pki_types::ServerName
/// [`SocketAddr`]: std::net::SocketAddr
pub(crate) fn tls_server_name(
    authority: &str,
) -> Result<rustls::pki_types::ServerName<'static>, rustls::pki_types::InvalidDnsNameError> {
    use rustls::pki_types::ServerName;

    if let Ok(addr) = authority.parse::<std::net::SocketAddr>() {
        return Ok(ServerName::from(addr.ip()));
    }
    let host = match authority.split_once(':') {
        Some((host, _port)) => host,
        None => authority,
    };
    Ok(ServerName::try_from(host)?.to_owned())
}

#[cfg(test)]
mod tls_server_name_tests {
    use super::tls_server_name;
    use rustls::pki_types::ServerName;

    #[test]
    fn resolves_server_name_from_authority() {
        // Host names keep their host and drop the port.
        assert_eq!(
            tls_server_name("example.com:443").unwrap(),
            ServerName::try_from("example.com").unwrap()
        );
        assert_eq!(
            tls_server_name("example.com").unwrap(),
            ServerName::try_from("example.com").unwrap()
        );

        // IP literals resolve to an `IpAddress` server name. The bracketed IPv6
        // form must not be truncated at the first `:`.
        assert_eq!(
            tls_server_name("127.0.0.1:80").unwrap(),
            ServerName::from(std::net::Ipv4Addr::LOCALHOST)
        );
        assert_eq!(
            tls_server_name("[::1]:443").unwrap(),
            ServerName::from(std::net::Ipv6Addr::LOCALHOST)
        );
        assert_eq!(
            tls_server_name("[2001:db8::1]:8443").unwrap(),
            ServerName::from("2001:db8::1".parse::<std::net::Ipv6Addr>().unwrap())
        );
    }
}