//! 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);

        }

    }

}