rustls-mitm 0.1.0

use std::num::NonZeroUsize;
use std::sync::{Arc, Mutex};

use lru::LruCache;
use rustls::ServerConfig;
use rustls::crypto::CryptoProvider;
use rustls::server::{ClientHello, ResolvesServerCert};
use rustls::sign::CertifiedKey;

use crate::CertificateAuthority;

const DEFAULT_CACHE_CAPACITY: usize = 1024;

/// A [`ResolvesServerCert`] implementation that generates per-host TLS
/// certificates on the fly, signed by a [`CertificateAuthority`].
///
/// When a client connects with SNI, the resolver generates (or returns a
/// cached) leaf certificate for that hostname. Certificates are cached in an
/// LRU cache to avoid regenerating them on every connection.
pub struct MitmCertResolver {
    ca: CertificateAuthority,
    cache: Mutex<LruCache<String, Arc<CertifiedKey>>>,
}

impl std::fmt::Debug for MitmCertResolver {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MitmCertResolver").finish_non_exhaustive()
    }
}

impl MitmCertResolver {
    /// Create a resolver with the default LRU cache capacity (1024).
    pub fn new(ca: CertificateAuthority) -> Self {
        Self::with_cache_capacity(ca, DEFAULT_CACHE_CAPACITY)
    }

    /// Create a resolver with a custom LRU cache capacity.
    pub fn with_cache_capacity(ca: CertificateAuthority, capacity: usize) -> Self {
        Self {
            ca,
            cache: Mutex::new(LruCache::new(NonZeroUsize::new(capacity).unwrap())),
        }
    }

    /// Build a [`ServerConfig`] that uses this resolver for certificate selection.
    ///
    /// The returned config has no client authentication and no ALPN protocols.
    /// Set `config.alpn_protocols` if your use case requires protocol
    /// negotiation (e.g. `[b"h2", b"http/1.1"]` for HTTP).
    pub fn into_server_config(self) -> ServerConfig {
        ServerConfig::builder()
            .with_no_client_auth()
            .with_cert_resolver(Arc::new(self))
    }
}

impl ResolvesServerCert for MitmCertResolver {
    fn resolve(&self, client_hello: ClientHello<'_>) -> Option<Arc<CertifiedKey>> {
        let hostname = client_hello.server_name()?;

        {
            let mut cache = self.cache.lock().unwrap();
            if let Some(key) = cache.get(hostname) {
                return Some(key.clone());
            }
        }

        let (cert_der, key_der) = self.ca.generate_cert(hostname).ok()?;
        let provider = CryptoProvider::get_default()?;
        let signing_key = provider.key_provider.load_private_key(key_der).ok()?;
        let certified = Arc::new(CertifiedKey::new(vec![cert_der], signing_key));

        let mut cache = self.cache.lock().unwrap();
        if let Some(existing) = cache.get(hostname) {
            return Some(existing.clone());
        }
        cache.put(hostname.to_string(), certified.clone());
        Some(certified)
    }
}