//! TLS 1.3 responder identity for QUIC.

use std::sync::Arc;
#[cfg(test)]
use std::{
    collections::HashMap,
    sync::{
        LazyLock,
        Mutex,
        Weak,
        atomic::{AtomicUsize, Ordering},
    },
    thread::{self, ThreadId},
};

use lanspread_proto::{ALPN_PROTOCOL, PeerId};
#[cfg(test)]
use rustls::sign::{CertifiedKey, SingleCertAndKey};
use rustls::{
    CertificateError,
    ClientConfig,
    DigitallySignedStruct,
    Error as RustlsError,
    ServerConfig,
    SignatureScheme,
    client::{
        Resumption,
        danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier},
    },
    crypto::WebPkiSupportedAlgorithms,
    pki_types::{CertificateDer, ServerName, UnixTime},
    server::{NoServerSessionStorage, ParsedCertificate},
};
use s2n_quic::provider::tls::rustls::{
    client::Client as S2nRustlsClient,
    server::Server as S2nRustlsServer,
};

use crate::identity::{PeerIdentity, peer_id_from_spki, server_name_for_peer};

pub(crate) fn protocol_alpn() -> Vec<u8> {
    ALPN_PROTOCOL.to_vec()
}

#[derive(Debug)]
struct PeerIdServerCertVerifier {
    supported_algorithms: WebPkiSupportedAlgorithms,
}

impl ServerCertVerifier for PeerIdServerCertVerifier {
    fn verify_server_cert(
        &self,
        end_entity: &CertificateDer<'_>,
        intermediates: &[CertificateDer<'_>],
        server_name: &ServerName<'_>,
        ocsp_response: &[u8],
        _now: UnixTime,
    ) -> Result<ServerCertVerified, RustlsError> {
        if !intermediates.is_empty() || !ocsp_response.is_empty() {
            return Err(application_verification_failure());
        }

        let expected_peer = expected_peer_from_server_name(server_name)?;
        let parsed = ParsedCertificate::try_from(end_entity)?;
        let spki = parsed.subject_public_key_info();
        let actual_peer =
            peer_id_from_spki(spki.as_ref()).map_err(|_| application_verification_failure())?;
        if actual_peer != expected_peer {
            return Err(application_verification_failure());
        }

        rustls::client::verify_server_name(&parsed, server_name)?;
        Ok(ServerCertVerified::assertion())
    }

    fn verify_tls12_signature(
        &self,
        _message: &[u8],
        _cert: &CertificateDer<'_>,
        _dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, RustlsError> {
        Err(RustlsError::General(
            "TLS 1.2 CertificateVerify is disabled".to_owned(),
        ))
    }

    fn verify_tls13_signature(
        &self,
        message: &[u8],
        cert: &CertificateDer<'_>,
        dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, RustlsError> {
        if dss.scheme != SignatureScheme::ED25519 {
            return Err(RustlsError::General(
                "unexpected TLS 1.3 CertificateVerify scheme".to_owned(),
            ));
        }
        rustls::crypto::verify_tls13_signature(message, cert, dss, &self.supported_algorithms)
    }

    fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
        vec![SignatureScheme::ED25519]
    }
}

pub(crate) fn client_provider() -> Result<S2nRustlsClient, RustlsError> {
    let provider = rustls::crypto::aws_lc_rs::default_provider();
    let verifier: Arc<dyn ServerCertVerifier> = Arc::new(PeerIdServerCertVerifier {
        supported_algorithms: provider.signature_verification_algorithms,
    });
    #[cfg(test)]
    let verifier = decorate_verifier_for_current_test(verifier);
    let mut config = ClientConfig::builder_with_provider(Arc::new(provider))
        .with_protocol_versions(&[&rustls::version::TLS13])?
        .dangerous()
        .with_custom_certificate_verifier(verifier)
        .with_no_client_auth();
    config.alpn_protocols = vec![protocol_alpn()];
    config.check_selected_alpn = true;
    config.enable_sni = true;
    config.resumption = Resumption::disabled();
    config.enable_early_data = false;
    Ok(S2nRustlsClient::from(config))
}

pub(crate) fn server_provider(identity: &PeerIdentity) -> Result<S2nRustlsServer, RustlsError> {
    let provider = Arc::new(rustls::crypto::aws_lc_rs::default_provider());
    let mut config = ServerConfig::builder_with_provider(provider)
        .with_protocol_versions(&[&rustls::version::TLS13])?
        .with_no_client_auth()
        .with_single_cert(vec![identity.certificate()], identity.private_key())?;
    harden_server_config(&mut config);
    Ok(S2nRustlsServer::from(config))
}

fn harden_server_config(config: &mut ServerConfig) {
    config.alpn_protocols = vec![protocol_alpn()];
    config.session_storage = Arc::new(NoServerSessionStorage {});
    config.send_tls13_tickets = 0;
    config.max_tls13_tickets = 0;
    config.max_early_data_size = 0;
    config.send_half_rtt_data = false;
}

#[cfg(test)]
pub(crate) fn hostile_mismatched_server_provider(
    certificate_identity: &PeerIdentity,
    signing_identity: &PeerIdentity,
) -> Result<S2nRustlsServer, RustlsError> {
    let provider = Arc::new(rustls::crypto::aws_lc_rs::default_provider());
    let signer =
        rustls::crypto::aws_lc_rs::sign::any_supported_type(&signing_identity.private_key())?;
    let certified_key = CertifiedKey::new(vec![certificate_identity.certificate()], signer);
    let resolver = Arc::new(SingleCertAndKey::from(certified_key));
    let mut config = ServerConfig::builder_with_provider(provider)
        .with_protocol_versions(&[&rustls::version::TLS13])?
        .with_no_client_auth()
        .with_cert_resolver(resolver);
    harden_server_config(&mut config);
    Ok(S2nRustlsServer::from(config))
}

#[cfg(test)]
#[derive(Debug, Default)]
struct VerificationCounts {
    certificate_calls: AtomicUsize,
    certificate_accepts: AtomicUsize,
    tls12_calls: AtomicUsize,
    tls13_calls: AtomicUsize,
    tls13_accepts: AtomicUsize,
}

#[cfg(test)]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct VerificationSnapshot {
    pub(crate) certificate_calls: usize,
    pub(crate) certificate_accepts: usize,
    pub(crate) tls12_calls: usize,
    pub(crate) tls13_calls: usize,
    pub(crate) tls13_accepts: usize,
}

#[cfg(test)]
impl VerificationCounts {
    fn snapshot(&self) -> VerificationSnapshot {
        VerificationSnapshot {
            certificate_calls: self.certificate_calls.load(Ordering::SeqCst),
            certificate_accepts: self.certificate_accepts.load(Ordering::SeqCst),
            tls12_calls: self.tls12_calls.load(Ordering::SeqCst),
            tls13_calls: self.tls13_calls.load(Ordering::SeqCst),
            tls13_accepts: self.tls13_accepts.load(Ordering::SeqCst),
        }
    }
}

#[cfg(test)]
#[derive(Debug)]
struct CountingVerifier {
    inner: Arc<dyn ServerCertVerifier>,
    counts: Arc<VerificationCounts>,
}

#[cfg(test)]
impl ServerCertVerifier for CountingVerifier {
    fn verify_server_cert(
        &self,
        end_entity: &CertificateDer<'_>,
        intermediates: &[CertificateDer<'_>],
        server_name: &ServerName<'_>,
        ocsp_response: &[u8],
        now: UnixTime,
    ) -> Result<ServerCertVerified, RustlsError> {
        self.counts.certificate_calls.fetch_add(1, Ordering::SeqCst);
        let result = self.inner.verify_server_cert(
            end_entity,
            intermediates,
            server_name,
            ocsp_response,
            now,
        );
        if result.is_ok() {
            self.counts
                .certificate_accepts
                .fetch_add(1, Ordering::SeqCst);
        }
        result
    }

    fn verify_tls12_signature(
        &self,
        message: &[u8],
        cert: &CertificateDer<'_>,
        dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, RustlsError> {
        self.counts.tls12_calls.fetch_add(1, Ordering::SeqCst);
        self.inner.verify_tls12_signature(message, cert, dss)
    }

    fn verify_tls13_signature(
        &self,
        message: &[u8],
        cert: &CertificateDer<'_>,
        dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, RustlsError> {
        self.counts.tls13_calls.fetch_add(1, Ordering::SeqCst);
        let result = self.inner.verify_tls13_signature(message, cert, dss);
        if result.is_ok() {
            self.counts.tls13_accepts.fetch_add(1, Ordering::SeqCst);
        }
        result
    }

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

#[cfg(test)]
static TEST_VERIFICATION_COUNTERS: LazyLock<Mutex<HashMap<ThreadId, Weak<VerificationCounts>>>> =
    LazyLock::new(|| Mutex::new(HashMap::new()));

/// Lexically scopes a verifier decorator to the client provider constructed on
/// the current test thread. The verifier retains its own counter after QUIC
/// starts; unrelated parallel test threads continue to use the undecorated
/// production verifier.
#[cfg(test)]
pub(crate) struct TestVerificationCounter {
    thread_id: ThreadId,
    counts: Arc<VerificationCounts>,
}

#[cfg(test)]
impl TestVerificationCounter {
    pub(crate) fn install() -> eyre::Result<Self> {
        let thread_id = thread::current().id();
        let counts = Arc::new(VerificationCounts::default());
        let mut installed = TEST_VERIFICATION_COUNTERS
            .lock()
            .map_err(|_| eyre::eyre!("TLS verification counter registry was poisoned"))?;
        installed.retain(|_, counter| counter.strong_count() != 0);
        eyre::ensure!(
            !installed.contains_key(&thread_id),
            "TLS verification counter is already installed on this test thread"
        );
        installed.insert(thread_id, Arc::downgrade(&counts));
        drop(installed);
        Ok(Self { thread_id, counts })
    }

    pub(crate) fn snapshot(&self) -> VerificationSnapshot {
        self.counts.snapshot()
    }
}

#[cfg(test)]
impl Drop for TestVerificationCounter {
    fn drop(&mut self) {
        let Ok(mut installed) = TEST_VERIFICATION_COUNTERS.lock() else {
            return;
        };
        let should_remove = installed
            .get(&self.thread_id)
            .and_then(Weak::upgrade)
            .is_some_and(|counts| Arc::ptr_eq(&counts, &self.counts));
        if should_remove {
            installed.remove(&self.thread_id);
        }
    }
}

#[cfg(test)]
fn decorate_verifier_for_current_test(
    verifier: Arc<dyn ServerCertVerifier>,
) -> Arc<dyn ServerCertVerifier> {
    let counts = TEST_VERIFICATION_COUNTERS
        .lock()
        .ok()
        .and_then(|installed| {
            installed
                .get(&thread::current().id())
                .and_then(Weak::upgrade)
        });
    counts.map_or(verifier.clone(), |counts| {
        Arc::new(CountingVerifier {
            inner: verifier,
            counts,
        })
    })
}

pub(crate) fn sni_for_peer(peer_id: PeerId) -> Result<String, RustlsError> {
    server_name_for_peer(peer_id).map_err(|_| application_verification_failure())
}

fn expected_peer_from_server_name(server_name: &ServerName<'_>) -> Result<PeerId, RustlsError> {
    let ServerName::DnsName(dns_name) = server_name else {
        return Err(application_verification_failure());
    };
    let encoded = dns_name.as_ref();
    let suffix = crate::identity::PEER_SNI_SUFFIX;
    let raw_peer_id = encoded
        .strip_suffix(suffix)
        .ok_or_else(application_verification_failure)?;
    let peer_id = raw_peer_id
        .parse::<PeerId>()
        .map_err(|_| application_verification_failure())?;
    if server_name_for_peer(peer_id).map_err(|_| application_verification_failure())? != encoded {
        return Err(application_verification_failure());
    }
    Ok(peer_id)
}

fn application_verification_failure() -> RustlsError {
    RustlsError::InvalidCertificate(CertificateError::ApplicationVerificationFailure)
}

#[cfg(test)]
mod tests {
    use eyre::ensure;
    use lanspread_proto::PeerId;
    use rustls::pki_types::ServerName;

    use super::{expected_peer_from_server_name, protocol_alpn, sni_for_peer};

    #[test]
    fn alpn_is_bound_to_the_current_wire_version() {
        assert_eq!(protocol_alpn(), lanspread_proto::ALPN_PROTOCOL);
        assert_ne!(protocol_alpn(), b"lanspread/7");
    }

    #[test]
    fn expected_peer_round_trips_only_through_exact_sni() -> eyre::Result<()> {
        let peer_id = PeerId::from_bytes([0x5a; 32]);
        let encoded = sni_for_peer(peer_id)?;
        let server_name = ServerName::try_from(encoded)?;
        ensure!(expected_peer_from_server_name(&server_name)? == peer_id);

        let wrong_suffix = ServerName::try_from(format!("{peer_id}.example.invalid"))?;
        ensure!(expected_peer_from_server_name(&wrong_suffix).is_err());
        Ok(())
    }
}