kapiti 0.0.3

use std::net::{Shutdown, SocketAddr, TcpStream, ToSocketAddrs};
use std::pin::Pin;
use std::str::FromStr;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;

use anyhow::{anyhow, bail, Context as _, Error, Result};
use async_lock::{Barrier, Mutex};
use async_native_tls_alpn::{Protocol, TlsConnector, TlsStream};
use http::Uri;
use hyper::{Body, Client};
use smol::{io, prelude::*, Async, Task};
use tokio::io::ReadBuf;
use tracing::{trace, warn};

use crate::resolver;
use crate::timeout;

/// Returns a new Hyper HTTP client:
/// - Using smol for the connection and async runtime
/// - Using the provided Kapiti Resolver for resolving any hostnames
/// This is for use when the HTTP endpoint may need "bootstrap" resolving of its own.
pub fn client_kapiti(
    mut resolver: resolver::Resolver,
    http2: bool,
    get_ipv6: bool,
    udp_size: u16,
    connect_timeout: Duration,
) -> Client<SmolConnector> {
    let (resolver_tx, resolver_rx): (
        async_channel::Sender<ResolverQuery>,
        async_channel::Receiver<ResolverQuery>,
    ) = async_channel::bounded(32);

    // Create a separate task that will perform lookups on behalf of the SmolConnector.
    // This is mainly to get around locking + async issues across the spawned tasks for each query.
    // The only other alternative would be to create a whole new Resolver with each query.
    // Also, we don't worry about tracking the handle for this task because it should expire automatically when resolver_tx is dropped.
    let resolver_task = Arc::new(smol::spawn(async move {
        trace!("Internal resolver waiting for requests");
        // Returns Err when channel is closed and has no more messages
        while let Ok(msg) = resolver_rx.recv().await {
            trace!("Internal resolver: {}", msg.host);
            // If the "host" appears to already be an IP, return it as-is rather than trying to resolve it.
            // This effectively mirrors the behavior of the system resolver via client_system().
            // This is not in the user query path, and should only come up if e.g. a filter URL is at an IP.
            let endpoint_str = format!("{}:{}", msg.host, msg.port);
            if let Ok(lookup_result) = SocketAddr::from_str(endpoint_str.as_str()) {
                trace!(
                    "Internal resolver IP shortcut: {} = {:?}",
                    endpoint_str,
                    lookup_result
                );
                // Store the result, then notify the barrier
                msg.result.lock().await.replace(Ok(lookup_result));
                msg.result_barrier.wait().await;
            } else {
                // It doesn't look like a socket address, so do the resolve.
                let lookup_result = resolver
                    .resolve_str(&msg.host, msg.port, get_ipv6, udp_size)
                    .await;
                if let Err(e) = &lookup_result {
                    warn!("Internal resolver lookup failed: {:?}", e);
                } else {
                    trace!("Internal resolver: {} = {:?}", endpoint_str, lookup_result);
                }
                // Store the result, then notify the barrier
                msg.result.lock().await.replace(lookup_result);
                msg.result_barrier.wait().await;
            }
        }
        trace!("Internal resolver exiting");
    }));

    Client::builder()
        .executor(SmolExecutor)
        .http2_only(http2)
        .build::<_, Body>(SmolConnector {
            _resolver_task: resolver_task,
            http2,
            resolver_tx,
            connect_timeout,
        })
}

/// Returns a new Hyper HTTP client:
/// - Using smol for the connection and async runtime
/// - Using no resolver for resolving any hostnames (hosts MUST be provided as IPs)
/// This is for use when the HTTP endpoint will be provided as an IP.
pub fn client_iponly(http2: bool, connect_timeout: Duration) -> Client<SmolConnector> {
    let (resolver_tx, resolver_rx): (
        async_channel::Sender<ResolverQuery>,
        async_channel::Receiver<ResolverQuery>,
    ) = async_channel::bounded(32);

    // Create a separate task that will perform lookups on behalf of the SmolConnector.
    // This is for cases where an IP endpoint was provided, in which case any "lookups"
    // must be for converting an IP string to a SocketAddr.
    // Technically the separate channel+task is overkill here, but lets keep things consistent.
    let resolver_task = Arc::new(smol::spawn(async move {
        trace!("Internal IP converter waiting for requests");
        // Returns Err when channel is closed and has no more messages
        while let Ok(msg) = resolver_rx.recv().await {
            trace!("Internal IP converter: {}", msg.host);
            // If the "host" appears to already be an IP, return it as-is rather than trying to resolve it.
            // This effectively mirrors the behavior of the system resolver via client_system().
            // This is not in the user query path, and should only come up if e.g. a filter URL is at an IP.
            let endpoint_str = format!("{}:{}", msg.host, msg.port);
            if let Ok(lookup_result) = SocketAddr::from_str(endpoint_str.as_str()) {
                trace!(
                    "Internal IP converted: {} = {:?}",
                    endpoint_str,
                    lookup_result
                );
                // Store the result, then notify the barrier
                msg.result.lock().await.replace(Ok(lookup_result));
                msg.result_barrier.wait().await;
            } else {
                // Invalid endpoint - we don't support actual host resolution!
                // Return an immediate error.
                msg.result
                    .lock()
                    .await
                    .replace(Err(anyhow!("Invalid IP endpoint: {}", endpoint_str)));
                msg.result_barrier.wait().await;
            }
        }
        trace!("Internal IP converter exiting");
    }));

    Client::builder()
        .executor(SmolExecutor)
        .http2_only(http2)
        .build::<_, Body>(SmolConnector {
            _resolver_task: resolver_task,
            http2,
            resolver_tx,
            connect_timeout,
        })
}

/// Returns a new Hyper HTTP client:
/// - Using smol for the connection and async runtime
/// - Using the system resolver for resolving any hostnames
/// This is only meant for use in internal tooling, not for client requests.
pub fn client_system(http2: bool, connect_timeout: Duration) -> Client<SmolConnector> {
    let (resolver_tx, resolver_rx): (
        async_channel::Sender<ResolverQuery>,
        async_channel::Receiver<ResolverQuery>,
    ) = async_channel::bounded(32);

    // Create a separate task that will perform lookups on behalf of the SmolConnector.
    // This isn't strictly needed for the system resolver, but keeps things in line with client_kapiti().
    let resolver_task = Arc::new(smol::spawn(async move {
        trace!("System resolver waiting for requests");
        // Returns Err when channel is closed and has no more messages
        while let Ok(msg) = resolver_rx.recv().await {
            trace!("System resolver: {}", msg.host);
            let host = msg.host.clone();
            let port = msg.port;
            let lookup_result =
                match smol::unblock(move || (host.as_str(), port).to_socket_addrs()).await {
                    Ok(mut socket_addrs) => match socket_addrs.next() {
                        Some(socket_addr) => Ok(socket_addr),
                        None => Err(anyhow!("No results for hostname {}", msg.host)),
                    },
                    Err(e) => {
                        Err(e).with_context(|| format!("Failed to query for hostname {}", msg.host))
                    }
                };

            trace!("System resolver: {} = {:?}", msg.host, lookup_result);
            // Store the result, then notify the barrier
            msg.result.lock().await.replace(lookup_result);
            msg.result_barrier.wait().await;
        }
        trace!("System resolver exiting");
    }));

    Client::builder()
        .executor(SmolExecutor)
        .http2_only(http2)
        .build::<_, Body>(SmolConnector {
            _resolver_task: resolver_task,
            http2,
            resolver_tx,
            connect_timeout,
        })
}

/// Spawns futures.
#[derive(Clone)]
struct SmolExecutor;

impl<F: Future + Send + 'static> hyper::rt::Executor<F> for SmolExecutor {
    fn execute(&self, fut: F) {
        smol::spawn(async { drop(fut.await) }).detach();
    }
}

/// The request for a host to be resolved, along with an output for returning the response.
#[derive(Debug)]
struct ResolverQuery {
    /// The hostname to look up
    host: String,
    /// The port to include in the resolved result
    port: u16,
    /// Barrier to wait for the result to appear. The requestor should wait on this before accessing result.
    result_barrier: Arc<Barrier>,
    /// Where the result should go. Should be an error if the hostname could not be resolved (e.g. not found).
    result: Arc<Mutex<Option<Result<SocketAddr>>>>,
}

/// Connects to URLs.
#[derive(Clone)]
pub struct SmolConnector {
    /// Handle to keep the resolver task from dying prematurely
    _resolver_task: Arc<Task<()>>,
    /// Whether to enable http2 features
    http2: bool,
    /// Channel for sending requests to the resolver task
    resolver_tx: async_channel::Sender<ResolverQuery>,
    /// Timeout for connection/handshake calls to complete
    connect_timeout: Duration,
}

impl hyper::service::Service<Uri> for SmolConnector {
    type Response = SmolStream;
    type Error = Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Poll::Ready(Ok(()))
    }

    fn call(&mut self, uri: Uri) -> Self::Future {
        // Get copies for async move:
        let http2 = self.http2.clone();
        let resolver_tx_copy = self.resolver_tx.clone();
        let connect_timeout = self.connect_timeout.clone();

        Box::pin(async move {
            let host = uri
                .host()
                .with_context(|| format!("Cannot parse host: {:?}", uri))?;
            // Release when both the requestor and requestee have called wait
            let result_barrier = Arc::new(Barrier::new(2));
            // Where the requestee will store the result before calling result_barrier.wait
            let result = Arc::new(Mutex::new(None));
            match uri.scheme_str() {
                Some("http") => {
                    // Send lookup, with place to give us back the result:
                    trace!("HTTP lookup: {}", host);
                    let query = ResolverQuery {
                        host: host.to_string(),
                        port: uri.port_u16().unwrap_or(80),
                        result_barrier: result_barrier.clone(),
                        result: result.clone(),
                    };
                    resolver_tx_copy
                        .send(query)
                        .await
                        .context("Failed to send HTTP resolver query")?;
                    trace!("HTTP lookup sent");
                    // Wait on the barrier to complete
                    result_barrier.wait().await;
                    // Barrier has completed, get the stored result.
                    // Jump through weird reference hoops to get the SocketAddr value out of the mutex.
                    match result
                        .lock()
                        .await
                        .as_ref()
                        .expect("Missing resolve result following barrier")
                    {
                        Ok(socket_addr) => {
                            let stream = timeout::timeout(
                                Async::<TcpStream>::connect(socket_addr.clone()),
                                &connect_timeout,
                            )
                            .await
                            .with_context(|| {
                                format!("HTTP TCP connect timed out: {:?}", socket_addr)
                            })?
                            .with_context(|| {
                                format!("HTTP TCP connect failed: {:?}", socket_addr)
                            })?;
                            Ok(SmolStream::Plain(stream))
                        }
                        // e is a reference, and anyhow doesn't like using with_context with it. So just give up and create a new error.
                        Err(e) => Err(anyhow!("Failed to resolve host {:?}: {}", uri, e)),
                    }
                }
                Some("https") => {
                    // Send lookup, with place to give us back the result:
                    trace!("HTTPS lookup: {}", host);
                    let query = ResolverQuery {
                        host: host.to_string(),
                        port: uri.port_u16().unwrap_or(443),
                        result_barrier: result_barrier.clone(),
                        result: result.clone(),
                    };
                    resolver_tx_copy
                        .send(query)
                        .await
                        .context("Failed to send HTTPS resolver query")?;
                    trace!("HTTPS lookup sent");
                    // Wait on the barrier to complete
                    result_barrier.wait().await;
                    // Barrier has completed, get the stored result.
                    // Jump through weird reference hoops to get the SocketAddr value out of the mutex.
                    match result
                        .lock()
                        .await
                        .as_ref()
                        .expect("Missing resolve result following barrier")
                    {
                        Ok(socket_addr) => {
                            let stream = timeout::timeout(
                                Async::<TcpStream>::connect(socket_addr.clone()),
                                &connect_timeout,
                            )
                            .await
                            .with_context(|| {
                                format!("HTTPS TCP connect timed out: {:?}", socket_addr)
                            })?
                            .with_context(|| {
                                format!("HTTPS TCP connect failed: {:?}", socket_addr)
                            })?;
                            // Min protocol: If things don't have at least TLS1.2 by now, we should just name and shame.
                            let mut connector =
                                TlsConnector::new().min_protocol_version(Some(Protocol::Tlsv12));
                            // ALPN: Required for http2/DoH, otherwise we get 'http2 error: protocol error: frame with invalid size'
                            // Meanwhile, required OFF for most other hosts, otherwise we get 'connection closed before message completed'
                            if http2 {
                                connector = connector.request_alpns(&["h2"]);
                            }
                            let stream =
                                timeout::timeout(connector.connect(host, stream), &connect_timeout)
                                    .await
                                    .with_context(|| {
                                        format!(
                                            "HTTPS TCP session connect timed out: {:?}",
                                            socket_addr
                                        )
                                    })?
                                    .with_context(|| {
                                        format!(
                                            "HTTPS TCP session connect failed: {:?}",
                                            socket_addr
                                        )
                                    })?;
                            Ok(SmolStream::Tls(stream))
                        }
                        // e is a reference, and anyhow doesn't like using with_context with it. So just give up and create a new error.
                        Err(e) => Err(anyhow!("Failed to resolve host {:?}: {}", uri, e)),
                    }
                }
                scheme => bail!("Unsupported scheme: {:?}", scheme),
            }
        })
    }
}

/// A TCP or TCP+TLS connection.
pub enum SmolStream {
    /// A plain TCP connection.
    Plain(Async<TcpStream>),

    /// A TCP connection secured by TLS.
    Tls(TlsStream<Async<TcpStream>>),
}

impl hyper::client::connect::Connection for SmolStream {
    fn connected(&self) -> hyper::client::connect::Connected {
        hyper::client::connect::Connected::new()
    }
}

impl tokio::io::AsyncRead for SmolStream {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        match &mut *self {
            SmolStream::Plain(s) => {
                Pin::new(s)
                    .poll_read(cx, buf.initialize_unfilled())
                    .map_ok(|size| {
                        buf.advance(size);
                        ()
                    })
            }
            SmolStream::Tls(s) => {
                Pin::new(s)
                    .poll_read(cx, buf.initialize_unfilled())
                    .map_ok(|size| {
                        buf.advance(size);
                        ()
                    })
            }
        }
    }
}

impl tokio::io::AsyncWrite for SmolStream {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        match &mut *self {
            SmolStream::Plain(s) => Pin::new(s).poll_write(cx, buf),
            SmolStream::Tls(s) => Pin::new(s).poll_write(cx, buf),
        }
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        match &mut *self {
            SmolStream::Plain(s) => Pin::new(s).poll_flush(cx),
            SmolStream::Tls(s) => Pin::new(s).poll_flush(cx),
        }
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        match &mut *self {
            SmolStream::Plain(s) => {
                s.get_ref().shutdown(Shutdown::Write)?;
                Poll::Ready(Ok(()))
            }
            SmolStream::Tls(s) => Pin::new(s).poll_close(cx),
        }
    }
}