//! Connects to a known set of `HOSTS`, captures the root certificates for the
//! certificate chain presented, and writes them to `src/anchors.rs` in a way
//! that is easy to consume for use as `extra_roots` in rustls-platform-verifier.

use std::{
    fs,
    sync::{Arc, Mutex},
};

use rustls::{
    DigitallySignedStruct, Error, RootCertStore, SignatureScheme,
    client::{
        WebPkiServerVerifier,
        danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier},
    },
    crypto::{CryptoProvider, aws_lc_rs},
    pki_types::{CertificateDer, ServerName, TrustAnchor, UnixTime},
};
use tokio::net::TcpStream;
use tokio_rustls::TlsConnector;
use webpki::{EndEntityCert, anchor_from_trusted_cert};
use webpki_root_certs::TLS_SERVER_ROOT_CERTS;

#[tokio::test]
async fn store_static_roots() -> anyhow::Result<()> {
    let provider = Arc::new(aws_lc_rs::default_provider());
    let mut root_store = RootCertStore::empty();
    let mut roots = Vec::with_capacity(TLS_SERVER_ROOT_CERTS.len());
    for cert_der in TLS_SERVER_ROOT_CERTS {
        let ta = anchor_from_trusted_cert(cert_der)?;
        roots.push((cert_der, ta.clone()));
        root_store.roots.push(ta);
    }

    let root_store = Arc::new(root_store);
    let inner = WebPkiServerVerifier::builder_with_provider(root_store.clone(), provider.clone())
        .build()?;

    let verifier = Arc::new(TrackRootVerifier {
        root: Mutex::default(),
        roots: root_store,
        inner,
        provider: provider.clone(),
    });

    let config = Arc::new(
        rustls::ClientConfig::builder_with_provider(provider)
            .with_safe_default_protocol_versions()?
            .dangerous()
            .with_custom_certificate_verifier(verifier.clone())
            .with_no_client_auth(),
    );

    let mut code = "// @generated by tests/static_roots.rs\n".to_string();
    code.push_str("use rustls::pki_types::CertificateDer;\n\n");
    code.push_str("pub(crate) const RUSTUP_TRUST_ANCHORS: &[CertificateDer<'static>] = &[\n");
    let connector = TlsConnector::from(config);
    for &host in HOSTS {
        connector
            .connect(
                ServerName::try_from(host)?,
                TcpStream::connect((host, 443)).await?,
            )
            .await?;

        let root = verifier.root.lock().unwrap().take().unwrap();
        let root_cert = roots
            .iter()
            .find_map(|(cert_der, ta)| (ta == &root).then_some(*cert_der))
            .unwrap();

        code.push_str(&format!("    // Root for host: {host}\n"));
        code.push_str("    CertificateDer::from_slice(\n        b\"\\\n");
        for chunk in root_cert.chunks(20) {
            code.push_str("        ");
            for byte in chunk {
                code.push_str(&format!("\\x{:02x}", byte));
            }
            code.push_str("\\\n");
        }
        code.push_str("        \",\n    ),\n");
    }
    code.push_str("];\n");

    let old = fs::read_to_string(PATH)?;
    if old != code {
        fs::write(PATH, &code)?;
        panic!("anchors.rs is outdated; updated it");
    }

    Ok(())
}

const PATH: &str = "src/anchors.rs";

#[derive(Debug)]
struct TrackRootVerifier {
    root: Mutex<Option<TrustAnchor<'static>>>,
    inner: Arc<WebPkiServerVerifier>,
    roots: Arc<RootCertStore>,
    provider: Arc<CryptoProvider>,
}

impl ServerCertVerifier for TrackRootVerifier {
    fn verify_server_cert(
        &self,
        end_entity: &CertificateDer<'_>,
        intermediates: &[CertificateDer<'_>],
        server_name: &ServerName<'_>,
        ocsp_response: &[u8],
        now: UnixTime,
    ) -> Result<ServerCertVerified, Error> {
        let verified = self.inner.verify_server_cert(
            end_entity,
            intermediates,
            server_name,
            ocsp_response,
            now,
        )?;

        let cert = EndEntityCert::try_from(end_entity)
            .map_err(|e| Error::General(format!("invalid end entity certificate: {e}")))?;

        let path = cert
            .verify_for_usage(
                self.provider.signature_verification_algorithms.all,
                &self.roots.roots,
                intermediates,
                now,
                webpki::KeyUsage::server_auth(),
                None,
                None,
            )
            .unwrap();

        let mut root = self.root.lock().unwrap();
        *root = Some(path.anchor().to_owned());
        Ok(verified)
    }

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

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

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

const HOSTS: &[&str] = &[
    "fastly-static.rust-lang.org",
    "cloudfront-static.rust-lang.org",
];
