VerseVDI-Data-Plane/core/src/tls.rs

use std::fmt;
use std::io::Cursor;
use std::sync::{Arc, Mutex};

use rustls::client::ResolvesClientCert;
use rustls::pki_types::CertificateDer;
use rustls::sign::{CertifiedKey, Signer, SigningKey};
use rustls::{ClientConfig, RootCertStore, SignatureAlgorithm, SignatureScheme};

use crate::error::{CoreError, Result};
use crate::wire::NativeTunnelCredential;

pub(crate) type SignCallback = dyn Fn(&[u8]) -> Result<[u8; 64]> + Send + Sync;

pub(crate) struct CallbackSigningKey {
    callback: Arc<Mutex<Option<Arc<SignCallback>>>>,
}

impl fmt::Debug for CallbackSigningKey {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("CallbackSigningKey")
    }
}

impl CallbackSigningKey {
    pub(crate) fn new(callback: Arc<SignCallback>) -> Self {
        Self {
            callback: Arc::new(Mutex::new(Some(callback))),
        }
    }
}

impl SigningKey for CallbackSigningKey {
    fn choose_scheme(&self, offered: &[SignatureScheme]) -> Option<Box<dyn Signer>> {
        offered
            .contains(&SignatureScheme::ED25519)
            .then(|| Box::new(CallbackSigner(Arc::clone(&self.callback))) as Box<dyn Signer>)
    }

    fn algorithm(&self) -> SignatureAlgorithm {
        SignatureAlgorithm::ED25519
    }
}

struct CallbackSigner(Arc<Mutex<Option<Arc<SignCallback>>>>);

impl fmt::Debug for CallbackSigner {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("CallbackSigner")
    }
}

impl Signer for CallbackSigner {
    fn sign(&self, message: &[u8]) -> std::result::Result<Vec<u8>, rustls::Error> {
        let callback = self
            .0
            .lock()
            .map_err(|_| rustls::Error::General("client signing state failed".to_owned()))?
            .take()
            .ok_or_else(|| rustls::Error::General("client signer already used".to_owned()))?;
        callback(message)
            .map(|signature| signature.to_vec())
            .map_err(|_| rustls::Error::General("client signing failed".to_owned()))
    }

    fn scheme(&self) -> SignatureScheme {
        SignatureScheme::ED25519
    }
}

#[derive(Debug)]
struct ClientIdentity(Arc<CertifiedKey>);

impl ResolvesClientCert for ClientIdentity {
    fn resolve(
        &self,
        _root_hint_subjects: &[&[u8]],
        sigschemes: &[SignatureScheme],
    ) -> Option<Arc<CertifiedKey>> {
        sigschemes
            .contains(&SignatureScheme::ED25519)
            .then(|| Arc::clone(&self.0))
    }

    fn has_certs(&self) -> bool {
        true
    }
}

pub(crate) fn client_config(
    credential: &NativeTunnelCredential,
    callback: Arc<SignCallback>,
) -> Result<ClientConfig> {
    let certificate_chain = rustls_pemfile::certs(&mut Cursor::new(
        credential.certificate_chain_pem().as_bytes(),
    ))
    .collect::<std::result::Result<Vec<CertificateDer<'static>>, _>>()
    .map_err(|_| CoreError::Tls)?;
    if certificate_chain.is_empty() {
        return Err(CoreError::Tls);
    }
    let mut roots = RootCertStore::empty();
    let trust_bundle =
        rustls_pemfile::certs(&mut Cursor::new(credential.trust_bundle_pem().as_bytes()))
            .collect::<std::result::Result<Vec<CertificateDer<'static>>, _>>()
            .map_err(|_| CoreError::Tls)?;
    if trust_bundle.is_empty() || roots.add_parsable_certificates(trust_bundle).1 != 0 {
        return Err(CoreError::Tls);
    }
    let provider = Arc::new(rustls::crypto::ring::default_provider());
    let mut config = ClientConfig::builder_with_provider(provider)
        .with_protocol_versions(&[&rustls::version::TLS13])
        .map_err(|_| CoreError::Tls)?
        .with_root_certificates(roots)
        .with_client_cert_resolver(Arc::new(ClientIdentity(Arc::new(CertifiedKey::new(
            certificate_chain,
            Arc::new(CallbackSigningKey::new(callback)),
        )))));
    config.alpn_protocols = vec![b"versevdi-gateway-v1".to_vec()];
    config.enable_early_data = false;
    Ok(config)
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::{AtomicBool, Ordering};

    use super::*;

    struct DropMarker(Arc<AtomicBool>);

    impl Drop for DropMarker {
        fn drop(&mut self) {
            self.0.store(true, Ordering::SeqCst);
        }
    }

    #[test]
    fn tls_callback_is_consumed_and_released_after_one_signature() {
        let dropped = Arc::new(AtomicBool::new(false));
        let marker = DropMarker(Arc::clone(&dropped));
        let callback: Arc<SignCallback> = Arc::new(move |_| {
            let _ = &marker;
            Ok([7; 64])
        });
        let key = CallbackSigningKey::new(callback);
        let signer = key
            .choose_scheme(&[SignatureScheme::ED25519])
            .expect("ED25519 signer");
        let second_signer = key
            .choose_scheme(&[SignatureScheme::ED25519])
            .expect("second ED25519 signer");
        drop(key);

        assert_eq!(
            signer.sign(b"handshake").expect("first signature"),
            vec![7; 64]
        );
        assert!(dropped.load(Ordering::SeqCst), "callback remained retained");
        assert!(
            second_signer.sign(b"second request").is_err(),
            "signer was reusable"
        );
    }
}