smartproxy/rust/crates/rustproxy-tls/src/cert_manager.rs

use std::time::{SystemTime, UNIX_EPOCH};
use thiserror::Error;
use tracing::info;

use crate::cert_store::{CertStore, CertBundle, CertMetadata, CertSource};
use crate::acme::AcmeClient;

#[derive(Debug, Error)]
pub enum CertManagerError {
    #[error("ACME provisioning failed for {domain}: {message}")]
    AcmeFailure { domain: String, message: String },
    #[error("No ACME email configured")]
    NoEmail,
}

/// Certificate lifecycle manager.
/// Handles ACME provisioning, static cert loading, and renewal.
pub struct CertManager {
    store: CertStore,
    acme_email: Option<String>,
    use_production: bool,
    renew_before_days: u32,
}

impl CertManager {
    pub fn new(
        store: CertStore,
        acme_email: Option<String>,
        use_production: bool,
        renew_before_days: u32,
    ) -> Self {
        Self {
            store,
            acme_email,
            use_production,
            renew_before_days,
        }
    }

    /// Get a certificate for a domain (from cache).
    pub fn get_cert(&self, domain: &str) -> Option<&CertBundle> {
        self.store.get(domain)
    }

    /// Create an ACME client using this manager's configuration.
    /// Returns None if no ACME email is configured.
    pub fn acme_client(&self) -> Option<AcmeClient> {
        self.acme_email.as_ref().map(|email| {
            AcmeClient::new(email.clone(), self.use_production)
        })
    }

    /// Load a static certificate into the store (infallible — pure cache insert).
    pub fn load_static(
        &mut self,
        domain: String,
        bundle: CertBundle,
    ) {
        self.store.store(domain, bundle);
    }

    /// Check and return domains that need certificate renewal.
    ///
    /// A certificate needs renewal if it expires within `renew_before_days`.
    /// Returns a list of domain names needing renewal.
    pub fn check_renewals(&self) -> Vec<String> {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        let renewal_threshold = self.renew_before_days as u64 * 86400;
        let mut needs_renewal = Vec::new();

        for (domain, bundle) in self.store.iter() {
            // Only auto-renew ACME certs
            if bundle.metadata.source != CertSource::Acme {
                continue;
            }

            let time_until_expiry = bundle.metadata.expires_at.saturating_sub(now);
            if time_until_expiry < renewal_threshold {
                info!(
                    "Certificate for {} needs renewal (expires in {} days)",
                    domain,
                    time_until_expiry / 86400
                );
                needs_renewal.push(domain.clone());
            }
        }

        needs_renewal
    }

    /// Renew a certificate for a domain.
    ///
    /// Performs the full ACME provision+store flow. The `challenge_setup` closure
    /// is called to arrange for the HTTP-01 challenge to be served. It receives
    /// (token, key_authorization) and must make the challenge response available.
    ///
    /// Returns the new CertBundle on success.
    pub async fn renew_domain<F, Fut>(
        &mut self,
        domain: &str,
        challenge_setup: F,
    ) -> Result<CertBundle, CertManagerError>
    where
        F: FnOnce(String, String) -> Fut,
        Fut: std::future::Future<Output = ()>,
    {
        let acme_client = self.acme_client()
            .ok_or(CertManagerError::NoEmail)?;

        info!("Renewing certificate for {}", domain);

        let domain_owned = domain.to_string();
        let result = acme_client.provision(&domain_owned, |pending| {
            let token = pending.token.clone();
            let key_auth = pending.key_authorization.clone();
            async move {
                challenge_setup(token, key_auth).await;
                Ok(())
            }
        }).await.map_err(|e| CertManagerError::AcmeFailure {
            domain: domain.to_string(),
            message: e.to_string(),
        })?;

        let (cert_pem, key_pem) = result;
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        let bundle = CertBundle {
            cert_pem,
            key_pem,
            ca_pem: None,
            metadata: CertMetadata {
                domain: domain.to_string(),
                source: CertSource::Acme,
                issued_at: now,
                expires_at: now + 90 * 86400,
                renewed_at: Some(now),
            },
        };

        self.store.store(domain.to_string(), bundle.clone());
        info!("Certificate renewed and stored for {}", domain);

        Ok(bundle)
    }

    /// Whether this manager has an ACME email configured.
    pub fn has_acme(&self) -> bool {
        self.acme_email.is_some()
    }

    /// Get reference to the underlying store.
    pub fn store(&self) -> &CertStore {
        &self.store
    }

    /// Get mutable reference to the underlying store.
    pub fn store_mut(&mut self) -> &mut CertStore {
        &mut self.store
    }
}