//! Outbound connections: resolve in the enclave, tunnel to the IP through the
//! parent, then TLS with the real host name (and Encrypted Client Hello when
//! the upstream offers it).

use crate::dns::Resolver;
use crate::relay::{Net, Purpose};
use crate::transport::Stream;
use anyhow::{anyhow, bail, Context, Result};
use http_body_util::{BodyExt, Full};
use hyper::Request;
use hyper_util::rt::TokioIo;
use rustls::client::{EchConfig, EchMode, EchStatus};
use std::net::IpAddr;
use std::sync::Arc;
use std::time::Duration;

/// Typed category attached (as anyhow context) to upstream failures so they
/// can be reported without inspecting error text.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ErrorKind {
    Dns,
    Connect,
    Tls,
    Timeout,
    RelayDown,
}

impl std::fmt::Display for ErrorKind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            ErrorKind::Dns => "dns",
            ErrorKind::Connect => "connect",
            ErrorKind::Tls => "tls",
            ErrorKind::Timeout => "timeout",
            ErrorKind::RelayDown => "relay down",
        })
    }
}

impl std::error::Error for ErrorKind {}

pub struct Dialer {
    net: Arc<Net>,
    resolver: Arc<Resolver>,
    roots: Arc<rustls::RootCertStore>,
    plain: Arc<rustls::ClientConfig>,
    ech_enabled: bool,
}

pub struct Upstream {
    pub stream: tokio_rustls::client::TlsStream<Stream>,
    pub addr: IpAddr,
    pub ech: EchStatus,
}

impl Dialer {
    pub fn new(net: Arc<Net>, resolver: Arc<Resolver>, ech_enabled: bool) -> Self {
        let mut roots = rustls::RootCertStore::empty();
        roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
        let roots = Arc::new(roots);
        let plain = rustls::ClientConfig::builder()
            .with_root_certificates(roots.clone())
            .with_no_client_auth();
        Self {
            net,
            resolver,
            roots,
            plain: Arc::new(plain),
            ech_enabled,
        }
    }

    pub fn resolver(&self) -> &Resolver {
        &self.resolver
    }

    fn ech_config(&self, list: &[u8]) -> Result<Arc<rustls::ClientConfig>> {
        let cfg = EchConfig::new(
            rustls::pki_types::EchConfigListBytes::from(list.to_vec()),
            rustls::crypto::aws_lc_rs::hpke::ALL_SUPPORTED_SUITES,
        )
        .map_err(|e| anyhow!("unusable ECH config: {e}"))?;
        let provider = rustls::crypto::CryptoProvider::get_default()
            .cloned()
            .ok_or_else(|| anyhow!("no crypto provider"))?;
        let client = rustls::ClientConfig::builder_with_provider(provider)
            .with_ech(EchMode::from(cfg))
            .map_err(|e| anyhow!("ech builder: {e}"))?
            .with_root_certificates(self.roots.clone())
            .with_no_client_auth();
        Ok(Arc::new(client))
    }

    /// TCP (via the parent) + TLS to `host:port`, certificate verified against
    /// `host`. The parent sees only the IP and, unless ECH is in use, the SNI.
    pub async fn connect_tls(&self, host: &str, port: u16, purpose: Purpose) -> Result<Upstream> {
        match self.connect_tls_once(host, port, purpose, true).await {
            Err(e) if is_ech_error(&e) => {
                // The upstream rejected our ECH config (rustls reports this as
                // a handshake error). Most likely a stale key in DNS: forget
                // it and retry once without ECH.
                crate::log("upstream rejected ECH config; retrying without ECH");
                self.resolver.forget(host);
                self.connect_tls_once(host, port, purpose, false).await
            }
            other => other,
        }
    }

    async fn connect_tls_once(&self, host: &str, port: u16, purpose: Purpose, try_ech: bool) -> Result<Upstream> {
        let resolved = self.resolver.resolve(host, purpose).await.context(ErrorKind::Dns)?;
        let name = rustls::pki_types::ServerName::try_from(host.to_string())
            .map_err(|e| anyhow!("invalid server name: {e}"))
            .context(ErrorKind::Dns)?;
        let config = match (&resolved.ech_config, self.ech_enabled && try_ech) {
            (Some(list), true) => match self.ech_config(list) {
                Ok(c) => c,
                Err(_) => {
                    crate::log("unusable ECH config; connecting without ECH");
                    self.plain.clone()
                }
            },
            _ => self.plain.clone(),
        };
        let connector = tokio_rustls::TlsConnector::from(config);
        let mut last = anyhow!("no addresses");
        for ip in &resolved.addrs {
            let addr = IpAddr::V4(*ip);
            if purpose == Purpose::Upstream && !public_destination(addr) {
                bail!("non-public upstream address refused");
            }
            let raw = match self.net.connect_ip(addr, port, purpose).await {
                Ok(s) => s,
                Err(e) => {
                    let kind = if e.downcast_ref::<crate::relay::RelayUnavailable>().is_some() {
                        ErrorKind::RelayDown
                    } else {
                        ErrorKind::Connect
                    };
                    last = e.context(kind);
                    continue;
                }
            };
            let tls = tokio::time::timeout(Duration::from_secs(30), connector.connect(name.clone(), raw))
                .await
                .map_err(|_| anyhow!("tls handshake timeout"))
                .context(ErrorKind::Timeout)?
                .with_context(|| format!("tls to {host} ({addr})"))
                .context(ErrorKind::Tls)?;
            let ech = tls.get_ref().1.ech_status();
            return Ok(Upstream { stream: tls, addr, ech });
        }
        Err(last.context(format!("connecting to {host}:{port}")))
    }

    /// One HTTPS request on a fresh connection; returns the full response
    /// with the body collected.
    pub async fn request_full(
        &self,
        method: &str,
        host: &str,
        path: &str,
        headers: &[(&str, &str)],
        body: Vec<u8>,
        purpose: Purpose,
    ) -> Result<http::Response<bytes::Bytes>> {
        let up = self.connect_tls(host, 443, purpose).await?;
        let (mut sender, conn) = hyper::client::conn::http1::handshake(TokioIo::new(up.stream)).await?;
        tokio::spawn(async move {
            let _ = conn.await;
        });
        let mut req = Request::builder()
            .method(method)
            .uri(path)
            .header("host", host)
            .header("user-agent", "timelock-proxy-enclave");
        for (k, v) in headers {
            if k.eq_ignore_ascii_case("host") || k.eq_ignore_ascii_case("user-agent") {
                continue;
            }
            req = req.header(*k, *v);
        }
        if !body.is_empty() {
            req = req.header("content-length", body.len());
        }
        let req = req.body(Full::new(bytes::Bytes::from(body)))?;
        let resp = tokio::time::timeout(Duration::from_secs(30), sender.send_request(req))
            .await
            .map_err(|_| anyhow!("timeout"))??;
        let (parts, body) = resp.into_parts();
        // Infrastructure APIs must not supply an unbounded response body.
        let bytes = http_body_util::Limited::new(body, 1024 * 1024).collect().await
            .map_err(|_| anyhow!("infrastructure response exceeded limit or failed"))?.to_bytes();
        Ok(http::Response::from_parts(parts, bytes))
    }

    /// Minimal HTTPS GET, used for beacon fetches.
    pub async fn https_get(&self, host: &str, path: &str, purpose: Purpose) -> Result<Vec<u8>> {
        let resp = self.request_full("GET", host, path, &[], Vec::new(), purpose).await?;
        if !resp.status().is_success() {
            bail!("status {}", resp.status());
        }
        Ok(resp.into_body().to_vec())
    }

    /// JSON request/response helper for small APIs.
    pub async fn json<T: serde::de::DeserializeOwned>(
        &self,
        method: &str,
        host: &str,
        path: &str,
        bearer: Option<&str>,
        body: Option<serde_json::Value>,
        purpose: Purpose,
    ) -> Result<T> {
        let auth = bearer.map(|t| format!("Bearer {t}"));
        let mut headers: Vec<(&str, &str)> = vec![("accept", "application/json")];
        if let Some(a) = &auth {
            headers.push(("authorization", a.as_str()));
        }
        let bytes = match &body {
            Some(v) => {
                headers.push(("content-type", "application/json"));
                serde_json::to_vec(v)?
            }
            None => Vec::new(),
        };
        let resp = self.request_full(method, host, path, &headers, bytes, purpose).await?;
        let status = resp.status();
        let body = resp.into_body();
        if !status.is_success() {
            bail!("{method} {host}{path}: status {status}: {}", String::from_utf8_lossy(&body[..body.len().min(300)]));
        }
        if body.is_empty() {
            return Ok(serde_json::from_str("null")?);
        }
        serde_json::from_slice(&body).context("parsing JSON response")
    }
}

/// Enforced inside the enclave after DNS resolution, not by the hostile host.
pub fn public_destination(ip: IpAddr) -> bool {
    match ip {
        IpAddr::V4(ip) => {
            let [a,b,_,_] = ip.octets();
            !(ip.is_private() || ip.is_loopback() || ip.is_link_local() || ip.is_broadcast()
                || ip.is_documentation() || ip.is_unspecified() || ip.is_multicast()
                || a == 0 || a >= 240 || (a == 100 && (64..128).contains(&b))
                || (a == 198 && (b == 18 || b == 19)) || (a == 192 && b == 0))
        }
        IpAddr::V6(_) => false, // Current DoH transport supports IPv4 destinations only.
    }
}

fn is_ech_error(e: &anyhow::Error) -> bool {
    let t = format!("{e:#}");
    t.contains("EncryptedClientHello") || t.contains("Ech")
}

pub fn ech_status_str(s: EchStatus) -> &'static str {
    match s {
        EchStatus::NotOffered => "not-offered",
        EchStatus::Grease => "grease",
        EchStatus::Offered => "offered",
        EchStatus::Accepted => "accepted",
        EchStatus::Rejected => "rejected",
    }
}