product-os-proxy 0.0.19

//! Certificate authority implementation using the `rcgen` crate.
//!
//! This module provides [`RcgenAuthority`], which generates TLS certificates
//! on-the-fly for MITM proxy interception using the `rcgen` library.
//! Generated certificates are cached using `moka` for performance.

use crate::mitm::certificate_authority::{CertificateAuthority, CACHE_TTL};
use moka::future::Cache;
use product_os_http::uri::Authority;
use product_os_security::certificates::CaProfile;
use product_os_utilities::ProductOSError;
use rand::thread_rng;
use rand::Rng;
use rcgen::{Certificate, CertificateParams, KeyPair};
// CertificateParams used in fallback cert generation
use std::sync::Arc;
use tokio_rustls::rustls::{
    pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer},
    ServerConfig,
};
use tracing::debug;

/// Certificate authority backed by `rcgen` with an in-memory certificate cache.
///
/// This implementation generates TLS certificates signed by a provided CA certificate
/// and private key. Generated server configs are cached using `moka` to avoid
/// regenerating certificates for frequently-accessed hosts. Certificate generation
/// runs in a blocking task to avoid blocking the async runtime.
///
/// # Certificate Properties
///
/// Generated certificates have:
/// - A random serial number
/// - A validity period handled by [`CaProfile::server_leaf_params`] (includes clock-skew backdate)
/// - The hostname as both Common Name (CN) and Subject Alternative Name (SAN)
/// - ALPN protocols: `h2` and `http/1.1`
///
/// # Cache Behavior
///
/// Certificates are cached by authority (hostname:port) with a TTL of [`CACHE_TTL`]
/// (half the certificate lifetime). The cache has a configurable maximum capacity.
pub struct RcgenAuthority {
    /// The CA key pair used for signing generated certificates (Arc for `spawn_blocking`)
    key_pair: Arc<KeyPair>,
    /// The CA certificate used as the issuer for generated certificates (Arc for `spawn_blocking`)
    ca_cert: Arc<Certificate>,
    /// LRU cache mapping authorities to their generated server configs
    cache: Cache<Authority, Arc<ServerConfig>>,
}

impl RcgenAuthority {
    /// Creates a new `RcgenAuthority` with the given CA key pair and certificate.
    ///
    /// # Arguments
    ///
    /// * `key_pair` - The CA key pair for signing certificates
    /// * `ca_cert` - The CA certificate (issuer)
    /// * `cache_size` - Maximum number of certificates to cache
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let authority = RcgenAuthority::new(key_pair, ca_cert, 1000);
    /// ```
    pub fn new(key_pair: KeyPair, ca_cert: Certificate, cache_size: u64) -> Self {
        Self {
            key_pair: Arc::new(key_pair),
            ca_cert: Arc::new(ca_cert),
            cache: Cache::builder()
                .max_capacity(cache_size)
                .time_to_live(std::time::Duration::from_secs(CACHE_TTL))
                .build(),
        }
    }

    /// Generates a certificate for the given authority (hostname or IP).
    /// Runs on the current thread; used from within `spawn_blocking`.
    fn gen_cert_sync(
        ca_key_pair: &KeyPair,
        ca_cert: &Certificate,
        authority: &Authority,
    ) -> Result<(CertificateDer<'static>, PrivateKeyDer<'static>), rcgen::Error> {
        let mut params = CaProfile::server_leaf_params(authority.host())?;
        params.serial_number = Some(thread_rng().gen::<u64>().into());

        let leaf_key = KeyPair::generate()?;
        let cert = params.signed_by(&leaf_key, ca_cert, ca_key_pair)?;
        let leaf_private = PrivateKeyDer::from(PrivatePkcs8KeyDer::from(leaf_key.serialize_der()));
        Ok((cert.into(), leaf_private))
    }

    /// Generates a certificate for the given authority (used by tests and sync fallbacks).
    #[allow(dead_code)]
    pub(crate) fn gen_cert(
        &self,
        authority: &Authority,
    ) -> Result<CertificateDer<'static>, rcgen::Error> {
        Ok(Self::gen_cert_sync(self.key_pair.as_ref(), self.ca_cert.as_ref(), authority)?.0)
    }
}

impl CertificateAuthority for RcgenAuthority {
    /// Generates a [`ServerConfig`] for the given authority, with caching.
    ///
    /// If a cached config exists for this authority, it is returned immediately.
    /// Otherwise, certificate generation runs in `spawn_blocking` to avoid blocking
    /// the async runtime, then the config is cached and returned.
    ///
    /// The generated `ServerConfig` supports both HTTP/2 (`h2`) and HTTP/1.1 ALPN.
    async fn gen_server_config(
        &self,
        authority: &Authority,
    ) -> Result<Arc<ServerConfig>, ProductOSError> {
        if let Some(server_cfg) = self.cache.get(authority).await {
            debug!("Using cached server config");
            return Ok(server_cfg);
        }
        debug!("Generating server config");

        let key_pair = Arc::clone(&self.key_pair);
        let ca_cert = Arc::clone(&self.ca_cert);
        let authority_clone = authority.clone();

        let (cert, private_key) = match tokio::task::spawn_blocking(move || {
            Self::gen_cert_sync(&key_pair, &ca_cert, &authority_clone)
        })
        .await
        {
            Ok(Ok(pair)) => pair,
            Ok(Err(e)) => {
                tracing::error!(
                    "Failed to generate certificate for {}: {:?}; falling back to generic cert",
                    authority,
                    e
                );
                let key_pair = Arc::clone(&self.key_pair);
                let ca_cert = Arc::clone(&self.ca_cert);
                let fallback_cert = tokio::task::spawn_blocking(move || -> Result<
                    (CertificateDer<'static>, PrivateKeyDer<'static>),
                    rcgen::Error,
                > {
                    let mut fallback_params = CertificateParams::default();
                    fallback_params.serial_number = Some(thread_rng().gen::<u64>().into());
                    let leaf_key = KeyPair::generate()?;
                    match fallback_params.signed_by(&leaf_key, ca_cert.as_ref(), key_pair.as_ref()) {
                        Ok(c) => Ok((
                            c.into(),
                            PrivateKeyDer::from(PrivatePkcs8KeyDer::from(leaf_key.serialize_der())),
                        )),
                        Err(fallback_err) => {
                            tracing::error!("Fallback certificate generation also failed: {:?}", fallback_err);
                            let cert = CertificateParams::default().self_signed(&leaf_key)?;
                            Ok((
                                cert.into(),
                                PrivateKeyDer::from(PrivatePkcs8KeyDer::from(leaf_key.serialize_der())),
                            ))
                        }
                    }
                })
                .await;
                match fallback_cert {
                    Ok(Ok(pair)) => pair,
                    Ok(Err(_)) => {
                        tracing::error!("Emergency self-signed cert generation failed");
                        return Err(ProductOSError::GenericError(
                            "Cannot generate any certificate; CA key pair may be invalid".into(),
                        ));
                    }
                    Err(e) => {
                        tracing::error!("spawn_blocking failed for fallback cert: {:?}", e);
                        return Err(ProductOSError::GenericError(format!(
                            "Certificate generation task failed: {e:?}"
                        )));
                    }
                }
            }
            Err(e) => {
                tracing::error!("spawn_blocking failed for cert generation: {:?}", e);
                return Err(ProductOSError::GenericError(format!(
                    "Certificate generation task failed: {e:?}"
                )));
            }
        };
        let certs = vec![cert];

        let mut server_cfg = ServerConfig::builder()
            .with_no_client_auth()
            .with_single_cert(certs, private_key)
            .map_err(|e| {
                ProductOSError::GenericError(format!("Failed to build ServerConfig: {e:?}"))
            })?;

        server_cfg.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];

        let server_cfg = Arc::new(server_cfg);

        self.cache
            .insert(authority.clone(), Arc::clone(&server_cfg))
            .await;

        Ok(server_cfg)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio_rustls::rustls;

    /// Helper to create a test `RcgenAuthority` from generated CA certificates
    fn create_test_authority() -> RcgenAuthority {
        let _ = rustls::crypto::ring::default_provider().install_default();
        let key_pair = rcgen::KeyPair::generate().expect("generate test CA key pair");
        let mut cert_params = rcgen::CertificateParams::default();
        cert_params
            .distinguished_name
            .push(rcgen::DnType::CommonName, "test-ca");
        cert_params.is_ca = rcgen::IsCa::Ca(rcgen::BasicConstraints::Unconstrained);
        let ca_cert = cert_params
            .self_signed(&key_pair)
            .expect("self-sign test CA cert");

        RcgenAuthority::new(key_pair, ca_cert, 100)
    }

    #[test]
    fn test_rcgen_authority_creation() {
        let _authority = create_test_authority();
    }

    #[test]
    fn test_rcgen_authority_gen_cert() {
        let authority = create_test_authority();
        let host_authority: Authority = "example.com".parse().unwrap();
        let cert_der = authority
            .gen_cert(&host_authority)
            .expect("Failed to generate cert");

        // Certificate should be non-empty DER bytes
        assert!(
            !cert_der.is_empty(),
            "Generated certificate should not be empty"
        );
    }

    #[test]
    fn test_rcgen_authority_gen_cert_ip_address() {
        let authority = create_test_authority();
        let ip_authority: Authority = "127.0.0.1".parse().unwrap();
        let cert_der = authority
            .gen_cert(&ip_authority)
            .expect("Failed to generate cert for IP");

        // Certificate should be non-empty DER bytes
        assert!(
            !cert_der.is_empty(),
            "Generated certificate for IP should not be empty"
        );
    }

    #[tokio::test]
    async fn test_rcgen_authority_gen_server_config() {
        let authority = create_test_authority();
        let host_authority: Authority = "example.com".parse().unwrap();

        let config = authority
            .gen_server_config(&host_authority)
            .await
            .expect("gen_server_config should succeed");

        // Config should have ALPN protocols set
        assert_eq!(config.alpn_protocols.len(), 2);
        assert_eq!(config.alpn_protocols[0], b"h2");
        assert_eq!(config.alpn_protocols[1], b"http/1.1");
    }

    #[tokio::test]
    async fn test_rcgen_authority_cache_returns_same_config() {
        let authority = create_test_authority();
        let host_authority: Authority = "cached.example.com".parse().unwrap();

        let config1 = authority
            .gen_server_config(&host_authority)
            .await
            .expect("gen_server_config should succeed");
        let config2 = authority
            .gen_server_config(&host_authority)
            .await
            .expect("gen_server_config should succeed");

        // Both calls should return the same Arc (cached)
        assert!(
            Arc::ptr_eq(&config1, &config2),
            "Second call should return cached config"
        );
    }

    #[tokio::test]
    async fn test_rcgen_authority_different_hosts_different_configs() {
        let authority = create_test_authority();
        let host1: Authority = "host1.example.com".parse().unwrap();
        let host2: Authority = "host2.example.com".parse().unwrap();

        let config1 = authority
            .gen_server_config(&host1)
            .await
            .expect("gen_server_config should succeed");
        let config2 = authority
            .gen_server_config(&host2)
            .await
            .expect("gen_server_config should succeed");

        // Different hosts should get different configs
        assert!(
            !Arc::ptr_eq(&config1, &config2),
            "Different hosts should have different configs"
        );
    }
}