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

//! ACME (Let's Encrypt) integration using instant-acme.
//!
//! This module handles HTTP-01 challenge creation and certificate provisioning.
//! Supports persisting ACME account credentials to disk for reuse across restarts.

use std::path::{Path, PathBuf};
use instant_acme::{
    Account, NewAccount, NewOrder, Identifier, ChallengeType, OrderStatus,
    AccountCredentials,
};
use rcgen::{CertificateParams, KeyPair};
use thiserror::Error;
use tracing::{debug, info, warn};

#[derive(Debug, Error)]
pub enum AcmeError {
    #[error("ACME account creation failed: {0}")]
    AccountCreation(String),
    #[error("ACME order failed: {0}")]
    OrderFailed(String),
    #[error("Challenge failed: {0}")]
    ChallengeFailed(String),
    #[error("Certificate finalization failed: {0}")]
    FinalizationFailed(String),
    #[error("No HTTP-01 challenge found")]
    NoHttp01Challenge,
    #[error("Timeout waiting for order: {0}")]
    Timeout(String),
    #[error("Account persistence error: {0}")]
    Persistence(String),
}

/// Pending HTTP-01 challenge that needs to be served.
pub struct PendingChallenge {
    pub token: String,
    pub key_authorization: String,
    pub domain: String,
}

/// ACME client wrapper around instant-acme.
pub struct AcmeClient {
    use_production: bool,
    email: String,
    /// Optional directory where account.json is persisted.
    account_dir: Option<PathBuf>,
}

impl AcmeClient {
    pub fn new(email: String, use_production: bool) -> Self {
        Self {
            use_production,
            email,
            account_dir: None,
        }
    }

    /// Create a new client with account persistence at the given directory.
    pub fn with_persistence(email: String, use_production: bool, account_dir: impl AsRef<Path>) -> Self {
        Self {
            use_production,
            email,
            account_dir: Some(account_dir.as_ref().to_path_buf()),
        }
    }

    /// Get or create an ACME account, persisting credentials if account_dir is set.
    async fn get_or_create_account(&self) -> Result<Account, AcmeError> {
        let directory_url = self.directory_url();

        // Try to restore from persisted credentials
        if let Some(ref dir) = self.account_dir {
            let account_file = dir.join("account.json");
            if account_file.exists() {
                match std::fs::read_to_string(&account_file) {
                    Ok(json) => {
                        match serde_json::from_str::<AccountCredentials>(&json) {
                            Ok(credentials) => {
                                match Account::from_credentials(credentials).await {
                                    Ok(account) => {
                                        debug!("Restored ACME account from {}", account_file.display());
                                        return Ok(account);
                                    }
                                    Err(e) => {
                                        warn!("Failed to restore ACME account, creating new: {}", e);
                                    }
                                }
                            }
                            Err(e) => {
                                warn!("Invalid account.json, creating new account: {}", e);
                            }
                        }
                    }
                    Err(e) => {
                        warn!("Could not read account.json: {}", e);
                    }
                }
            }
        }

        // Create a new account
        let contact = format!("mailto:{}", self.email);
        let (account, credentials) = Account::create(
            &NewAccount {
                contact: &[&contact],
                terms_of_service_agreed: true,
                only_return_existing: false,
            },
            directory_url,
            None,
        )
        .await
        .map_err(|e| AcmeError::AccountCreation(e.to_string()))?;

        debug!("ACME account created");

        // Persist credentials if we have a directory
        if let Some(ref dir) = self.account_dir {
            if let Err(e) = std::fs::create_dir_all(dir) {
                warn!("Failed to create account directory {}: {}", dir.display(), e);
            } else {
                let account_file = dir.join("account.json");
                match serde_json::to_string_pretty(&credentials) {
                    Ok(json) => {
                        if let Err(e) = std::fs::write(&account_file, &json) {
                            warn!("Failed to persist ACME account to {}: {}", account_file.display(), e);
                        } else {
                            info!("ACME account credentials persisted to {}", account_file.display());
                        }
                    }
                    Err(e) => {
                        warn!("Failed to serialize account credentials: {}", e);
                    }
                }
            }
        }

        Ok(account)
    }

    /// Request a certificate for a domain using the HTTP-01 challenge.
    ///
    /// Returns (cert_chain_pem, private_key_pem) on success.
    ///
    /// The caller must serve the HTTP-01 challenge at:
    ///   `http://<domain>/.well-known/acme-challenge/<token>`
    ///
    /// The `challenge_handler` closure is called with a `PendingChallenge`
    /// and must arrange for the challenge response to be served. It should
    /// return once the challenge is ready to be validated.
    pub async fn provision<F, Fut>(
        &self,
        domain: &str,
        challenge_handler: F,
    ) -> Result<(String, String), AcmeError>
    where
        F: FnOnce(PendingChallenge) -> Fut,
        Fut: std::future::Future<Output = Result<(), AcmeError>>,
    {
        info!("Starting ACME provisioning for {} via {}", domain, self.directory_url());

        // 1. Get or create ACME account (with persistence)
        let account = self.get_or_create_account().await?;

        // 2. Create order
        let identifier = Identifier::Dns(domain.to_string());
        let mut order = account
            .new_order(&NewOrder {
                identifiers: &[identifier],
            })
            .await
            .map_err(|e| AcmeError::OrderFailed(e.to_string()))?;

        debug!("ACME order created");

        // 3. Get authorizations and find HTTP-01 challenge
        let authorizations = order
            .authorizations()
            .await
            .map_err(|e| AcmeError::OrderFailed(e.to_string()))?;

        // Find the HTTP-01 challenge
        let (challenge_token, challenge_url) = authorizations
            .iter()
            .flat_map(|auth| auth.challenges.iter())
            .find(|c| c.r#type == ChallengeType::Http01)
            .map(|c| {
                let key_auth = order.key_authorization(c);
                (
                    PendingChallenge {
                        token: c.token.clone(),
                        key_authorization: key_auth.as_str().to_string(),
                        domain: domain.to_string(),
                    },
                    c.url.clone(),
                )
            })
            .ok_or(AcmeError::NoHttp01Challenge)?;

        // Call the handler to set up challenge serving
        challenge_handler(challenge_token).await?;

        // 4. Notify ACME server that challenge is ready
        order
            .set_challenge_ready(&challenge_url)
            .await
            .map_err(|e| AcmeError::ChallengeFailed(e.to_string()))?;

        debug!("Challenge marked as ready, waiting for validation...");

        // 5. Poll for order to become ready
        let mut attempts = 0;
        let state = loop {
            tokio::time::sleep(std::time::Duration::from_secs(2)).await;
            let state = order
                .refresh()
                .await
                .map_err(|e| AcmeError::OrderFailed(e.to_string()))?;

            match state.status {
                OrderStatus::Ready | OrderStatus::Valid => break state.status,
                OrderStatus::Invalid => {
                    return Err(AcmeError::ChallengeFailed(
                        "Order became invalid (challenge failed)".to_string(),
                    ));
                }
                _ => {
                    attempts += 1;
                    if attempts > 30 {
                        return Err(AcmeError::Timeout(
                            "Order did not become ready within 60 seconds".to_string(),
                        ));
                    }
                }
            }
        };

        debug!("Order ready, finalizing...");

        // 6. Generate CSR and finalize
        let key_pair = KeyPair::generate().map_err(|e| {
            AcmeError::FinalizationFailed(format!("Key generation failed: {}", e))
        })?;

        let mut params = CertificateParams::new(vec![domain.to_string()]).map_err(|e| {
            AcmeError::FinalizationFailed(format!("CSR params failed: {}", e))
        })?;
        params.distinguished_name.push(rcgen::DnType::CommonName, domain);

        let csr = params.serialize_request(&key_pair).map_err(|e| {
            AcmeError::FinalizationFailed(format!("CSR serialization failed: {}", e))
        })?;

        if state == OrderStatus::Ready {
            order
                .finalize(csr.der())
                .await
                .map_err(|e| AcmeError::FinalizationFailed(e.to_string()))?;
        }

        // 7. Wait for certificate to be issued
        let mut attempts = 0;
        loop {
            let state = order
                .refresh()
                .await
                .map_err(|e| AcmeError::OrderFailed(e.to_string()))?;
            if state.status == OrderStatus::Valid {
                break;
            }
            if state.status == OrderStatus::Invalid {
                return Err(AcmeError::FinalizationFailed(
                    "Order became invalid during finalization".to_string(),
                ));
            }
            attempts += 1;
            if attempts > 15 {
                return Err(AcmeError::Timeout(
                    "Certificate not issued within 30 seconds".to_string(),
                ));
            }
            tokio::time::sleep(std::time::Duration::from_secs(2)).await;
        }

        // 8. Download certificate
        let cert_chain_pem = order
            .certificate()
            .await
            .map_err(|e| AcmeError::FinalizationFailed(e.to_string()))?
            .ok_or_else(|| {
                AcmeError::FinalizationFailed("No certificate returned".to_string())
            })?;

        let private_key_pem = key_pair.serialize_pem();

        info!("Certificate provisioned successfully for {}", domain);

        Ok((cert_chain_pem, private_key_pem))
    }

    /// Restore an ACME account from stored credentials.
    pub async fn restore_account(
        &self,
        credentials: AccountCredentials,
    ) -> Result<Account, AcmeError> {
        Account::from_credentials(credentials)
            .await
            .map_err(|e| AcmeError::AccountCreation(e.to_string()))
    }

    /// Get the ACME directory URL based on production/staging.
    pub fn directory_url(&self) -> &str {
        if self.use_production {
            "https://acme-v02.api.letsencrypt.org/directory"
        } else {
            "https://acme-staging-v02.api.letsencrypt.org/directory"
        }
    }

    /// Whether this client is configured for production.
    pub fn is_production(&self) -> bool {
        self.use_production
    }
}

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

    #[test]
    fn test_directory_url_staging() {
        let client = AcmeClient::new("test@example.com".to_string(), false);
        assert!(client.directory_url().contains("staging"));
        assert!(!client.is_production());
    }

    #[test]
    fn test_directory_url_production() {
        let client = AcmeClient::new("test@example.com".to_string(), true);
        assert!(!client.directory_url().contains("staging"));
        assert!(client.is_production());
    }

    #[test]
    fn test_with_persistence_sets_account_dir() {
        let tmp = tempfile::tempdir().unwrap();
        let client = AcmeClient::with_persistence(
            "test@example.com".to_string(),
            false,
            tmp.path(),
        );
        assert!(client.account_dir.is_some());
        assert_eq!(client.account_dir.unwrap(), tmp.path());
    }

    #[test]
    fn test_without_persistence_no_account_dir() {
        let client = AcmeClient::new("test@example.com".to_string(), false);
        assert!(client.account_dir.is_none());
    }
}