//! TLS for client and server (spec §8.6). The server makes a certificate
//! of its own; the instance key vouches for it with an `instance.tls`
//! event; a client trusts that one certificate, pinned by the SHA-256 of
//! its SubjectPublicKeyInfo, and never a certificate authority.

use anyhow::{bail, Context, Result};
use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};
use rustls::pki_types::{CertificateDer, PrivateKeyDer, ServerName, UnixTime};
use rustls::{DigitallySignedStruct, SignatureScheme};
use std::net::SocketAddr;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::net::{TcpListener, TcpStream};

/// The crypto provider: ring, as OxygenHub's.
pub fn provider() -> Arc<rustls::crypto::CryptoProvider> {
    Arc::new(rustls::crypto::ring::default_provider())
}

/// SHA-256 of a certificate's SubjectPublicKeyInfo: what a client pins.
pub fn spki_sha256(cert_der: &[u8]) -> Result<[u8; 32]> {
    let (_, cert) = x509_parser::parse_x509_certificate(cert_der).map_err(|e| anyhow::anyhow!("not an X.509 certificate: {e}"))?;
    let d = ring::digest::digest(&ring::digest::SHA256, cert.public_key().raw);
    Ok(d.as_ref().try_into().expect("SHA-256 is 32 bytes"))
}

/// Where a server's certificate comes from (`[server] tls`).
#[derive(Clone, Debug)]
pub enum CertSource {
    /// Made on first start in `<data>/keys/tls.{key,crt}`, naming these names.
    Own { names: Vec<String> },
    /// Files from elsewhere (PEM).
    Files { cert: PathBuf, key: PathBuf },
}

/// A server's TLS: its rustls configuration and the certificate's pin.
pub struct ServerTls {
    pub config: Arc<rustls::ServerConfig>,
    pub spki_sha256: [u8; 32],
    pub names: Vec<String>,
}

impl ServerTls {
    pub fn load(data: &Path, source: &CertSource) -> Result<ServerTls> {
        let (cert_pem, key_pem, names) = match source {
            CertSource::Own { names } => {
                let keys = data.join("keys");
                let (crt, key) = (keys.join("tls.crt"), keys.join("tls.key"));
                if !crt.exists() || !key.exists() {
                    make_own(&crt, &key, names)?;
                }
                (std::fs::read_to_string(&crt)?, std::fs::read_to_string(&key)?, names.clone())
            }
            CertSource::Files { cert, key } => (
                std::fs::read_to_string(cert).with_context(|| format!("reading {}", cert.display()))?,
                std::fs::read_to_string(key).with_context(|| format!("reading {}", key.display()))?,
                vec![],
            ),
        };
        let chain: Vec<CertificateDer<'static>> = rustls_pemfile::certs(&mut cert_pem.as_bytes()).collect::<Result<_, _>>()?;
        let leaf = chain.first().context("no certificate in the PEM")?;
        let spki = spki_sha256(leaf)?;
        let key: PrivateKeyDer<'static> = rustls_pemfile::private_key(&mut key_pem.as_bytes())?.context("no private key in the PEM")?;
        let mut config = rustls::ServerConfig::builder_with_provider(provider())
            .with_safe_default_protocol_versions()?
            .with_no_client_auth()
            .with_single_cert(chain, key)?;
        config.alpn_protocols = vec![b"http/1.1".to_vec()];
        Ok(ServerTls { config: Arc::new(config), spki_sha256: spki, names })
    }
}

fn make_own(crt: &Path, key: &Path, names: &[String]) -> Result<()> {
    if names.is_empty() {
        bail!("a certificate needs at least one name ([server] names)");
    }
    let ck = rcgen::generate_simple_self_signed(names.to_vec())?;
    if let Some(dir) = crt.parent() {
        std::fs::create_dir_all(dir)?;
    }
    write_private(key, ck.signing_key.serialize_pem().as_bytes())?;
    std::fs::write(crt, ck.cert.pem())?;
    Ok(())
}

fn write_private(path: &Path, bytes: &[u8]) -> Result<()> {
    #[cfg(unix)]
    {
        use std::io::Write;
        use std::os::unix::fs::OpenOptionsExt;
        let mut f = std::fs::OpenOptions::new().write(true).create(true).truncate(true).mode(0o600).open(path)?;
        f.write_all(bytes)?;
        Ok(())
    }
    #[cfg(not(unix))]
    {
        std::fs::write(path, bytes)?;
        Ok(())
    }
}

/// A TCP listener that hands axum connections once their TLS handshake is
/// done; handshakes run concurrently, so a slow one holds up nobody.
pub struct TlsListener {
    rx: tokio::sync::mpsc::Receiver<(tokio_rustls::server::TlsStream<TcpStream>, SocketAddr)>,
    addr: SocketAddr,
}

impl TlsListener {
    pub fn new(tcp: TcpListener, config: Arc<rustls::ServerConfig>) -> Result<TlsListener> {
        let addr = tcp.local_addr()?;
        let acceptor = tokio_rustls::TlsAcceptor::from(config);
        let (tx, rx) = tokio::sync::mpsc::channel(64);
        tokio::spawn(async move {
            loop {
                let (stream, peer) = match tcp.accept().await {
                    Ok(x) => x,
                    Err(e) => {
                        tracing::warn!("accept: {e}");
                        tokio::time::sleep(Duration::from_millis(50)).await;
                        continue;
                    }
                };
                let (acceptor, out) = (acceptor.clone(), tx.clone());
                tokio::spawn(async move {
                    match tokio::time::timeout(Duration::from_secs(10), acceptor.accept(stream)).await {
                        Ok(Ok(tls)) => {
                            let _ = out.send((tls, peer)).await;
                        }
                        Ok(Err(e)) => tracing::debug!(%peer, "TLS handshake failed: {e}"),
                        Err(_) => tracing::debug!(%peer, "TLS handshake timed out"),
                    }
                });
                if tx.is_closed() {
                    return;
                }
            }
        });
        Ok(TlsListener { rx, addr })
    }
}

impl axum::serve::Listener for TlsListener {
    type Io = tokio_rustls::server::TlsStream<TcpStream>;
    type Addr = SocketAddr;

    async fn accept(&mut self) -> (Self::Io, Self::Addr) {
        match self.rx.recv().await {
            Some(x) => x,
            // The accept loop never ends while we hold the receiver.
            None => std::future::pending().await,
        }
    }

    fn local_addr(&self) -> std::io::Result<Self::Addr> {
        Ok(self.addr)
    }
}

/// A server certificate verifier that accepts exactly one public key, or,
/// with no pin, any key, remembering it for the caller to check
/// (`spectra connect`). The handshake's signature is always checked, so
/// the server must hold the certificate's private key.
#[derive(Debug)]
pub struct Pinned {
    pin: Option<[u8; 32]>,
    seen: Mutex<Option<(Vec<u8>, [u8; 32])>>,
    algs: rustls::crypto::WebPkiSupportedAlgorithms,
}

impl Pinned {
    pub fn new(pin: Option<[u8; 32]>) -> Arc<Pinned> {
        Arc::new(Pinned { pin, seen: Mutex::new(None), algs: provider().signature_verification_algorithms })
    }

    /// The certificate (DER) and its pin from the last handshake.
    pub fn seen(&self) -> Option<(Vec<u8>, [u8; 32])> {
        self.seen.lock().unwrap().clone()
    }

    pub fn client_config(self: &Arc<Self>) -> Result<Arc<rustls::ClientConfig>> {
        let mut c = rustls::ClientConfig::builder_with_provider(provider())
            .with_safe_default_protocol_versions()?
            .dangerous()
            .with_custom_certificate_verifier(self.clone())
            .with_no_client_auth();
        c.alpn_protocols = vec![b"http/1.1".to_vec()];
        Ok(Arc::new(c))
    }
}

impl ServerCertVerifier for Pinned {
    fn verify_server_cert(
        &self,
        end_entity: &CertificateDer<'_>,
        _intermediates: &[CertificateDer<'_>],
        _server_name: &ServerName<'_>,
        _ocsp: &[u8],
        _now: UnixTime,
    ) -> Result<ServerCertVerified, rustls::Error> {
        let spki = spki_sha256(end_entity).map_err(|_| rustls::Error::InvalidCertificate(rustls::CertificateError::BadEncoding))?;
        *self.seen.lock().unwrap() = Some((end_entity.to_vec(), spki));
        match self.pin {
            Some(pin) if pin != spki => Err(rustls::Error::InvalidCertificate(rustls::CertificateError::ApplicationVerificationFailure)),
            _ => Ok(ServerCertVerified::assertion()),
        }
    }

    fn verify_tls12_signature(
        &self,
        message: &[u8],
        cert: &CertificateDer<'_>,
        dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, rustls::Error> {
        rustls::crypto::verify_tls12_signature(message, cert, dss, &self.algs)
    }

    fn verify_tls13_signature(
        &self,
        message: &[u8],
        cert: &CertificateDer<'_>,
        dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, rustls::Error> {
        rustls::crypto::verify_tls13_signature(message, cert, dss, &self.algs)
    }

    fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
        self.algs.supported_schemes()
    }
}

/// Connect to `host:port` over TLS with `config`.
pub async fn connect(addr: &str, config: Arc<rustls::ClientConfig>) -> Result<tokio_rustls::client::TlsStream<TcpStream>> {
    let tcp = tokio::time::timeout(Duration::from_secs(10), TcpStream::connect(addr))
        .await
        .with_context(|| format!("connecting to {addr}: timed out"))?
        .with_context(|| format!("connecting to {addr}"))?;
    tcp.set_nodelay(true).ok();
    let host = addr.rsplit_once(':').map(|(h, _)| h.trim_start_matches('[').trim_end_matches(']')).unwrap_or(addr);
    let name = ServerName::try_from(host.to_string()).unwrap_or_else(|_| ServerName::try_from("spectra").unwrap());
    let tls = tokio_rustls::TlsConnector::from(config).connect(name, tcp).await.with_context(|| format!("TLS with {addr}"))?;
    Ok(tls)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn own_certificate_is_made_once_and_pinned() {
        let dir = tempfile::tempdir().unwrap();
        let src = CertSource::Own { names: vec!["localhost".into()] };
        let a = ServerTls::load(dir.path(), &src).unwrap();
        let b = ServerTls::load(dir.path(), &src).unwrap();
        assert_eq!(a.spki_sha256, b.spki_sha256, "the certificate is kept between starts");
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let mode = std::fs::metadata(dir.path().join("keys/tls.key")).unwrap().permissions().mode();
            assert_eq!(mode & 0o777, 0o600);
        }
    }
}