httpmock 0.9.0

use std::{future::pending, io, net::SocketAddr, pin::Pin, sync::Arc};

use http::{Request, StatusCode};
use http_body_util::{BodyExt, Empty, Full, combinators::BoxBody};
use hyper::{
    Method, Response,
    body::{Bytes, Incoming},
    service::service_fn,
};
use hyper_util::rt::tokio::TokioIo;
#[cfg(feature = "http2")]
use hyper_util::server::conn::auto::Builder as ServerBuilder;
#[cfg(feature = "https")]
use rustls::ServerConfig;
use thiserror::Error;
use tokio::{
    net::{TcpListener, TcpStream},
    sync::oneshot::Sender,
    task::spawn,
};
#[cfg(feature = "https")]
use tokio_rustls::TlsAcceptor;

use crate::server::{
    self, handler,
    transport::Error::{BufferError, LocalSocketAddrError, PublishSocketAddrError, RouterError, SocketBindError},
};

#[derive(Error, Debug)]
pub enum Error {
    #[error("cannot bind to socket addr {0}: {1}")]
    SocketBindError(SocketAddr, std::io::Error),
    #[error("cannot parse socket address: {0}")]
    SocketAddrParseError(#[from] std::net::AddrParseError),
    #[error("cannot obtain local error: {0}")]
    LocalSocketAddrError(std::io::Error),
    #[error("cannot send reserved TCP address to test thread {0}")]
    PublishSocketAddrError(SocketAddr),
    #[error("cannot create response: {0}")]
    ResponseConstructionError(http::Error),
    #[error("buffering error: {0}")]
    BufferError(hyper::Error),
    #[error("HTTP error: {0}")]
    HTTPError(#[from] http::Error),
    #[error("cannot process request: {0}")]
    RouterError(#[from] server::handler::Error),
    #[error("HTTPS error: {0}")]
    TlsError(String),
    #[error("Server configuration error: {0}")]
    ConfigurationError(String),
    #[error("Server I/O error: {0}")]
    IOError(io::Error),
    #[error("Server error: {0}")]
    ServerError(#[from] hyper::Error),
    #[error("Server error: {0}")]
    ServerConnectionError(Box<dyn std::error::Error + Send + Sync>),
    #[error("unknown data store error")]
    Unknown,
}

#[cfg(feature = "https")]
pub struct MockServerHttpsConfig {
    pub cert_resolver_factory: Arc<dyn CertificateResolverFactory + Send + Sync>,
}

pub struct MockServerConfig {
    pub static_port: Option<u16>,
    pub expose: bool,
    pub print_access_log: bool,
    #[cfg(feature = "https")]
    pub https: MockServerHttpsConfig,
}

/// The `HttpMockServer` struct represents a mock server that can handle incoming HTTP requests.
pub struct HttpMockServer {
    handler: handler::Handler,
    config: MockServerConfig,
}

impl HttpMockServer {
    /// Creates a new `HttpMockServer` instance with the given handler and configuration.
    ///
    /// # Parameters
    /// - `handler`: The request handler.
    /// - `config`: The configuration settings for the mock server.
    pub(crate) fn new(handler: handler::Handler, config: MockServerConfig) -> Self {
        HttpMockServer { handler, config }
    }

    /// Starts the mock server asynchronously.
    pub async fn start(self) -> Result<(), Error> {
        self.start_with_signals(None, pending()).await
    }

    /// Starts the mock server asynchronously with support for handling external shutdown signals.
    ///
    /// # Parameters
    /// - `socket_addr_sender`: An optional `Sender` to send the server's socket address once it's bound.
    /// - `shutdown`: A future that resolves when the server should shut down.
    ///
    pub async fn start_with_signals<F>(
        self,
        socket_addr_sender: Option<Sender<SocketAddr>>,
        shutdown: F,
    ) -> Result<(), Error>
    where
        F: Future<Output = ()>,
    {
        let host = if self.config.expose { "0.0.0.0" } else { "127.0.0.1" };
        let addr: SocketAddr = format!("{}:{}", host, self.config.static_port.unwrap_or(0)).parse()?;
        let listener = TcpListener::bind(addr).await.map_err(|e| SocketBindError(addr, e))?;

        if let Some(sender) = socket_addr_sender {
            let addr = listener.local_addr().map_err(LocalSocketAddrError)?;
            sender.send(addr).map_err(PublishSocketAddrError)?;
        }

        // ****************************************************************************************
        // SERVER START
        tracing::info!("Listening on {}", addr);
        self.run_accept_loop(listener, shutdown).await
    }

    pub async fn run_accept_loop<F>(self, listener: TcpListener, shutdown: F) -> Result<(), Error>
    where
        F: Future<Output = ()>,
    {
        let mut shutdown = std::pin::pin!(shutdown);
        let server = Arc::new(self);

        loop {
            tokio::select! {
                accepted = listener.accept() => {
                    match accepted {
                        Ok((tcp_stream, remote_address)) => {
                            let server = server.clone();
                            spawn(async move {
                               if let Err(err) = server.handle_tcp_stream(tcp_stream, remote_address).await {
                                    tracing::error!("{:?}", err);
                                }
                            });
                        },
                        Err(err) =>  {
                            tracing::error!("TCP error: {:?}", err);
                        },
                    };
                }
                _ = &mut shutdown => {
                    break;
                }
            }
        }

        Ok(())
    }

    async fn service(self: Arc<Self>, req: Request<Incoming>) -> Result<Response<BoxBody<Bytes, hyper::Error>>, Error> {
        tracing::trace!("New HTTP request received: {}", req.uri());

        if req.method() == Method::CONNECT {
            #[cfg(feature = "https")]
            {
                // CONNECT handling (HTTP proxy tunneling):
                //
                // For a CONNECT request the client asks us to turn the HTTP connection into a
                // raw, bidirectional byte tunnel (e.g., for TLS-over-HTTP proxies).
                //
                // In Hyper, the upgrade to this raw mode only happens AFTER we return the
                // 200 Connection Established response so the server can flush it and switch
                // the underlying socket into “upgraded” mode. The future from
                // `hyper::upgrade::on(req)` (captured as `on_upgrade`) will only resolve once
                // that response is sent and the connection is upgraded.
                //
                // If we awaited `on_upgrade` inline here, we would deadlock/block the upgrade:
                //   - We would wait for the upgrade to complete, but Hyper can’t complete it
                //     until we first return 200.
                //
                // Therefore we must detach the upgrade handling from this request future. We
                // do that by spawning a background task which:
                //   1) awaits `on_upgrade` once Hyper has upgraded the connection, and
                //   2) runs the tunnel logic in `handle_connect_upgraded`.
                //
                // Meanwhile, we immediately return `200 Connection Established` (with an empty
                // body) to let Hyper proceed with the upgrade.

                let authority = req.uri().authority().map(|a| a.to_string());
                let on_upgrade = hyper::upgrade::on(req);
                let server = self.clone();

                spawn(async move {
                    match on_upgrade.await {
                        Ok(upgraded) => {
                            let io = TokioIo::new(upgraded);
                            if let Err(e) = serve_tls_connection(server, io, authority).await {
                                tracing::warn!("failed to serve upgraded TLS connection: {:?}", e);
                            }
                        }
                        Err(err) => {
                            let e = Error::ServerConnectionError(Box::new(err));
                            tracing::warn!("CONNECT upgraded handling failed: {:?}", e);
                        }
                    }
                });
            }

            // Respond with 200 Connection Established and empty body
            return Ok(Response::builder().status(StatusCode::OK).body(empty())?);
        }

        let mut req = match buffer_request(req).await {
            Ok(req) => req,
            Err(err) => {
                return error_response(StatusCode::INTERNAL_SERVER_ERROR, BufferError(err));
            }
        };

        // Normalize the request URI to absolute-form for both HTTP and HTTPS.
        // The forms can be different and depends on how the client is talking to us and which
        // role our server plays (origin or proxy) and the protocol version.
        //
        // By normalizing it here, the rest of our servers application logic can uniformly
        // expect absolute-form URIs (e.g., matchers can read the correct host/authority
        // from the URI without falling back to reading the HOST header, etc.).
        //
        // Clients typically send origin-form ("/path") with a Host header; after a
        // CONNECT+TLS MITM this is also what the browser sends to us. We normalize to
        // absolute-form (scheme://authority/path?query) so that:
        // - matchers and the recorder can reliably read scheme/host/port from req.uri();
        // - we can avoid ad-hoc Host-header parsing throughout the codebase.
        //
        // See handler::proxy() for the inverse step where we convert back to
        // origin-form right before sending the request upstream, since most origin
        // servers (HTTP/1.1 and HTTP/2) expect origin-form on the wire.
        if let Err(err) = to_absolute_form_uri(&mut req) {
            return error_response(StatusCode::INTERNAL_SERVER_ERROR, err);
        }

        let access_log_req_data = self
            .config
            .print_access_log
            .then_some((req.method().clone(), req.uri().clone()));

        let resp = match self.handler.handle(req).await {
            Ok(response) => to_service_response(response),
            Err(err) => error_response(StatusCode::INTERNAL_SERVER_ERROR, RouterError(err)),
        };

        // print access log if enabled
        if let Some((method, uri)) = access_log_req_data
            && let Ok(resp) = &resp
        {
            tracing::info!("{} {} -> {}", method, uri, resp.status());
        }

        resp
    }

    async fn handle_tcp_stream(
        self: Arc<Self>,
        tcp_stream: TcpStream,
        _remote_address: SocketAddr,
    ) -> Result<(), Error> {
        tracing::trace!("new TCP connection incoming");

        #[cfg(feature = "https")]
        {
            let is_tls = match starts_with_tls_handshake(&tcp_stream).await {
                Ok(is_tls) => is_tls,
                // The connection broke before the client sent anything, so there is nothing to
                // serve. This is usually a client that resets right away, like a health check or
                // a port scanner, which isn't worth an error.
                Err(err) => {
                    tracing::debug!("TCP connection failed before the client sent any data: {:?}", err);
                    return Ok(());
                }
            };

            if is_tls {
                tracing::trace!("TCP connection seems to be TLS encrypted");

                // Since we get a request via HTTPS, the target host for this request is this server.
                // This is important for SNI and selecting the right certificate.
                // We are not yet in the tunneling case here, so we need to use the servers
                // local address. The tunneling case is handled in the CONNECT branch of `service()`.
                let tcp_address = tcp_stream.local_addr().map_err(Error::IOError)?;
                return serve_tls_connection(self, tcp_stream, Some(tcp_address.to_string())).await;
            }
        }

        tracing::trace!("TCP connection is not TLS encrypted");
        serve_connection(self.clone(), tcp_stream, "http").await
    }
}

#[cfg(feature = "https")]
async fn serve_tls_connection<S>(
    server: Arc<HttpMockServer>,
    stream: S,
    authority: Option<String>, // The target host:port for SNI and cert selection
) -> Result<(), Error>
where
    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
    // Build the TLS acceptor for this connection.
    let cert_resolver = server.config.https.cert_resolver_factory.build(authority);
    // Respect the application's provider. An explicit fallback avoids ambiguous
    // crate-feature auto-detection when both ring and aws-lc-rs are enabled.
    let provider = rustls::crypto::CryptoProvider::get_default()
        .cloned()
        .unwrap_or_else(|| Arc::new(rustls::crypto::ring::default_provider()));
    let mut server_config = ServerConfig::builder_with_provider(provider)
        .with_safe_default_protocol_versions()
        .map_err(|e| Error::TlsError(format!("cannot build TLS server config: {:?}", e)))?
        .with_no_client_auth()
        .with_cert_resolver(cert_resolver);

    server_config.alpn_protocols = vec![
        #[cfg(feature = "http2")]
        b"h2".to_vec(),
        b"http/1.1".to_vec(),
        b"http/1.0".to_vec(),
    ];

    let tls_acceptor = TlsAcceptor::from(Arc::new(server_config));
    let tls_stream = tls_acceptor
        .accept(stream)
        .await
        .map_err(|e| Error::TlsError(format!("TLS accept failed: {:?}", e)))?;

    serve_connection(server, tls_stream, "https").await
}

// `serve_connection` cannot be a plain `async fn` (nor return an implicit/explicit
// `impl Future`): `service_fn`'s closure below calls `HttpMockServer::service`, which can spawn a
// task that calls `serve_tls_connection` -> `serve_connection` again for CONNECT/upgrade
// tunneling. Any opaque `impl Future` return type here would need to resolve itself through
// that indirect cycle, which rustc cannot do ("cycle detected when computing type of opaque
// type"), surfacing as cascading bogus "is not Send" errors. Returning a boxed future erases
// the concrete type and breaks the cycle. This runs once per TCP connection (not per
// request, unlike `service`), so the extra heap allocation is negligible.
fn serve_connection<S>(
    server: Arc<HttpMockServer>,
    stream: S,
    scheme: &'static str,
) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + 'static>>
where
    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
    Box::pin(async move {
        let service = service_fn(|mut req| {
            // We pass authority None here since we don't know it for non-CONNECT requests
            // yet. We only know it when the full request has been buffered in `service()`.
            // Here, we only the scheme is known from the connection type.
            req.extensions_mut().insert(RequestMetadata::new(scheme));
            server.clone().service(req)
        });

        #[cfg(feature = "http2")]
        {
            let mut server_builder = ServerBuilder::new(TokioExecutor::new());

            server_builder.http1().preserve_header_case(true);
            server_builder.http2();

            server_builder
                .serve_connection_with_upgrades(TokioIo::new(stream), service)
                .await
                .map_err(ServerConnectionError)
        }

        #[cfg(not(feature = "http2"))]
        {
            // Without the `http2` feature we use hyper's plain HTTP/1 connection builder
            // instead of hyper-util's `auto` builder: `hyper-util/server-auto`
            // unconditionally enables `hyper/http2`, which would pull the `h2` dependency
            // subtree into builds that opted out of HTTP/2.
            let mut server_builder = hyper::server::conn::http1::Builder::new();

            server_builder.preserve_header_case(true);

            server_builder
                .serve_connection(TokioIo::new(stream), service)
                .with_upgrades()
                .await
                .map_err(|e| ServerConnectionError(Box::new(e)))
        }
    })
}

async fn buffer_request(req: Request<Incoming>) -> Result<Request<Bytes>, hyper::Error> {
    let (parts, body) = req.into_parts();
    let body = body.collect().await?.to_bytes();
    Ok(Request::from_parts(parts, body))
}

fn full<T: Into<Bytes>>(chunk: T) -> BoxBody<Bytes, hyper::Error> {
    Full::new(chunk.into()).map_err(|never| match never {}).boxed()
}

fn empty() -> BoxBody<Bytes, hyper::Error> {
    Empty::<Bytes>::new().map_err(|never| match never {}).boxed()
}

fn error_response(code: StatusCode, err: Error) -> Result<Response<BoxBody<Bytes, hyper::Error>>, Error> {
    tracing::error!("failed to process request: {}", err.to_string());
    Ok(Response::builder().status(code).body(full(err.to_string()))?)
}

fn to_service_response(response: Response<Bytes>) -> Result<Response<BoxBody<Bytes, hyper::Error>>, Error> {
    let (parts, body) = response.into_parts();
    Ok(Response::from_parts(parts, full(body)))
}

#[cfg(feature = "http2")]
use hyper_util::rt::TokioExecutor;
use tokio::io::{AsyncRead, AsyncWrite};

#[cfg(feature = "https")]
use crate::server::tls::{CertificateResolverFactory, starts_with_tls_handshake};
use crate::server::{Error::ServerConnectionError, RequestMetadata};

fn to_absolute_form_uri(req: &mut Request<Bytes>) -> Result<(), Error> {
    let default_scheme = req
        .extensions()
        .get::<RequestMetadata>()
        .map(|m| m.scheme)
        .unwrap_or("http");

    // If already absolute-form (scheme + authority), leave as-is
    let uri = req.uri().clone();
    if uri.scheme().is_some() && uri.authority().is_some() {
        return Ok(());
    }

    // Use Host header to reconstruct the absolute URI
    let host = req
        .headers()
        .get(http::header::HOST)
        .and_then(|v| v.to_str().ok())
        .ok_or_else(|| Error::ConfigurationError("Missing Host header on origin-form request".into()))?;

    let path_and_query = uri.path_and_query().map(|pq| pq.as_str()).unwrap_or("/");

    let abs = format!("{}://{}{}", default_scheme, host, path_and_query);
    let new_uri: http::Uri = abs
        .parse()
        .map_err(|e| Error::ConfigurationError(format!("Invalid absolute URI constructed from Host+path: {}", e)))?;

    *req.uri_mut() = new_uri;
    Ok(())
}

#[cfg(all(test, feature = "https"))]
mod tests {
    use std::{sync::Arc, time::Duration};

    use tokio::{
        net::{TcpListener, TcpStream},
        time::timeout,
    };

    use crate::server::builder::HttpMockServerBuilder;

    #[tokio::test]
    async fn connection_reset_before_any_data_is_not_an_error() {
        let server = Arc::new(HttpMockServerBuilder::new().build().unwrap());
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let client = TcpStream::connect(listener.local_addr().unwrap()).await.unwrap();
        let (tcp_stream, remote_address) = listener.accept().await.unwrap();

        // With a zero linger timeout, closing the socket resets the connection, as health checks
        // and port scanners often do.
        client.set_zero_linger().unwrap();
        drop(client);

        let handling = server.handle_tcp_stream(tcp_stream, remote_address);
        timeout(Duration::from_secs(5), handling)
            .await
            .expect("the reset connection should be handled right away")
            .unwrap();
    }
}