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

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use serde::{Deserialize, Serialize};
use thiserror::Error;

#[derive(Debug, Error)]
pub enum CertStoreError {
    #[error("Certificate not found for domain: {0}")]
    NotFound(String),
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),
    #[error("Invalid certificate: {0}")]
    Invalid(String),
    #[error("JSON error: {0}")]
    Json(#[from] serde_json::Error),
}

/// Certificate metadata stored alongside certs on disk.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CertMetadata {
    pub domain: String,
    pub source: CertSource,
    pub issued_at: u64,
    pub expires_at: u64,
    pub renewed_at: Option<u64>,
}

/// How a certificate was obtained.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum CertSource {
    Acme,
    Static,
    Custom,
    SelfSigned,
}

/// An in-memory certificate bundle.
#[derive(Debug, Clone)]
pub struct CertBundle {
    pub key_pem: String,
    pub cert_pem: String,
    pub ca_pem: Option<String>,
    pub metadata: CertMetadata,
}

/// Filesystem-backed certificate store.
///
/// File layout per domain:
/// ```text
/// {base_dir}/{domain}/
///   key.pem
///   cert.pem
///   ca.pem       (optional)
///   metadata.json
/// ```
pub struct CertStore {
    base_dir: PathBuf,
    /// In-memory cache of loaded certs
    cache: HashMap<String, CertBundle>,
}

impl CertStore {
    /// Create a new cert store at the given directory.
    pub fn new(base_dir: impl AsRef<Path>) -> Self {
        Self {
            base_dir: base_dir.as_ref().to_path_buf(),
            cache: HashMap::new(),
        }
    }

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

    /// Store a certificate to both cache and filesystem.
    pub fn store(&mut self, domain: String, bundle: CertBundle) -> Result<(), CertStoreError> {
        // Sanitize domain for directory name (replace wildcards)
        let dir_name = domain.replace('*', "_wildcard_");
        let cert_dir = self.base_dir.join(&dir_name);

        // Create directory
        std::fs::create_dir_all(&cert_dir)?;

        // Write key
        std::fs::write(cert_dir.join("key.pem"), &bundle.key_pem)?;

        // Write cert
        std::fs::write(cert_dir.join("cert.pem"), &bundle.cert_pem)?;

        // Write CA cert if present
        if let Some(ref ca) = bundle.ca_pem {
            std::fs::write(cert_dir.join("ca.pem"), ca)?;
        }

        // Write metadata
        let metadata_json = serde_json::to_string_pretty(&bundle.metadata)?;
        std::fs::write(cert_dir.join("metadata.json"), metadata_json)?;

        // Update cache
        self.cache.insert(domain, bundle);
        Ok(())
    }

    /// Check if a certificate exists for a domain.
    pub fn has(&self, domain: &str) -> bool {
        self.cache.contains_key(domain)
    }

    /// Load all certificates from the base directory.
    pub fn load_all(&mut self) -> Result<usize, CertStoreError> {
        if !self.base_dir.exists() {
            return Ok(0);
        }

        let entries = std::fs::read_dir(&self.base_dir)?;
        let mut loaded = 0;

        for entry in entries {
            let entry = entry?;
            let path = entry.path();

            if !path.is_dir() {
                continue;
            }

            let metadata_path = path.join("metadata.json");
            let key_path = path.join("key.pem");
            let cert_path = path.join("cert.pem");

            // All three files must exist
            if !metadata_path.exists() || !key_path.exists() || !cert_path.exists() {
                continue;
            }

            // Load metadata
            let metadata_str = std::fs::read_to_string(&metadata_path)?;
            let metadata: CertMetadata = serde_json::from_str(&metadata_str)?;

            // Load key and cert
            let key_pem = std::fs::read_to_string(&key_path)?;
            let cert_pem = std::fs::read_to_string(&cert_path)?;

            // Load CA cert if present
            let ca_path = path.join("ca.pem");
            let ca_pem = if ca_path.exists() {
                Some(std::fs::read_to_string(&ca_path)?)
            } else {
                None
            };

            let domain = metadata.domain.clone();
            let bundle = CertBundle {
                key_pem,
                cert_pem,
                ca_pem,
                metadata,
            };

            self.cache.insert(domain, bundle);
            loaded += 1;
        }

        Ok(loaded)
    }

    /// Get the base directory.
    pub fn base_dir(&self) -> &Path {
        &self.base_dir
    }

    /// Get the number of cached certificates.
    pub fn count(&self) -> usize {
        self.cache.len()
    }

    /// Iterate over all cached certificates.
    pub fn iter(&self) -> impl Iterator<Item = (&String, &CertBundle)> {
        self.cache.iter()
    }

    /// Remove a certificate from cache and filesystem.
    pub fn remove(&mut self, domain: &str) -> Result<bool, CertStoreError> {
        let removed = self.cache.remove(domain).is_some();
        if removed {
            let dir_name = domain.replace('*', "_wildcard_");
            let cert_dir = self.base_dir.join(&dir_name);
            if cert_dir.exists() {
                std::fs::remove_dir_all(&cert_dir)?;
            }
        }
        Ok(removed)
    }
}

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

    fn make_test_bundle(domain: &str) -> CertBundle {
        CertBundle {
            key_pem: "-----BEGIN PRIVATE KEY-----\ntest-key\n-----END PRIVATE KEY-----\n".to_string(),
            cert_pem: "-----BEGIN CERTIFICATE-----\ntest-cert\n-----END CERTIFICATE-----\n".to_string(),
            ca_pem: None,
            metadata: CertMetadata {
                domain: domain.to_string(),
                source: CertSource::Static,
                issued_at: 1700000000,
                expires_at: 1700000000 + 90 * 86400,
                renewed_at: None,
            },
        }
    }

    #[test]
    fn test_store_and_load_roundtrip() {
        let tmp = tempfile::tempdir().unwrap();
        let mut store = CertStore::new(tmp.path());

        let bundle = make_test_bundle("example.com");
        store.store("example.com".to_string(), bundle.clone()).unwrap();

        // Verify files exist
        let cert_dir = tmp.path().join("example.com");
        assert!(cert_dir.join("key.pem").exists());
        assert!(cert_dir.join("cert.pem").exists());
        assert!(cert_dir.join("metadata.json").exists());
        assert!(!cert_dir.join("ca.pem").exists()); // No CA cert

        // Load into a fresh store
        let mut store2 = CertStore::new(tmp.path());
        let loaded = store2.load_all().unwrap();
        assert_eq!(loaded, 1);

        let loaded_bundle = store2.get("example.com").unwrap();
        assert_eq!(loaded_bundle.key_pem, bundle.key_pem);
        assert_eq!(loaded_bundle.cert_pem, bundle.cert_pem);
        assert_eq!(loaded_bundle.metadata.domain, "example.com");
        assert_eq!(loaded_bundle.metadata.source, CertSource::Static);
    }

    #[test]
    fn test_store_with_ca_cert() {
        let tmp = tempfile::tempdir().unwrap();
        let mut store = CertStore::new(tmp.path());

        let mut bundle = make_test_bundle("secure.com");
        bundle.ca_pem = Some("-----BEGIN CERTIFICATE-----\nca-cert\n-----END CERTIFICATE-----\n".to_string());
        store.store("secure.com".to_string(), bundle).unwrap();

        let cert_dir = tmp.path().join("secure.com");
        assert!(cert_dir.join("ca.pem").exists());

        let mut store2 = CertStore::new(tmp.path());
        store2.load_all().unwrap();
        let loaded = store2.get("secure.com").unwrap();
        assert!(loaded.ca_pem.is_some());
    }

    #[test]
    fn test_load_all_multiple_certs() {
        let tmp = tempfile::tempdir().unwrap();
        let mut store = CertStore::new(tmp.path());

        store.store("a.com".to_string(), make_test_bundle("a.com")).unwrap();
        store.store("b.com".to_string(), make_test_bundle("b.com")).unwrap();
        store.store("c.com".to_string(), make_test_bundle("c.com")).unwrap();

        let mut store2 = CertStore::new(tmp.path());
        let loaded = store2.load_all().unwrap();
        assert_eq!(loaded, 3);
        assert!(store2.has("a.com"));
        assert!(store2.has("b.com"));
        assert!(store2.has("c.com"));
    }

    #[test]
    fn test_load_all_missing_directory() {
        let mut store = CertStore::new("/nonexistent/path/to/certs");
        let loaded = store.load_all().unwrap();
        assert_eq!(loaded, 0);
    }

    #[test]
    fn test_remove_cert() {
        let tmp = tempfile::tempdir().unwrap();
        let mut store = CertStore::new(tmp.path());

        store.store("remove-me.com".to_string(), make_test_bundle("remove-me.com")).unwrap();
        assert!(store.has("remove-me.com"));

        let removed = store.remove("remove-me.com").unwrap();
        assert!(removed);
        assert!(!store.has("remove-me.com"));
        assert!(!tmp.path().join("remove-me.com").exists());
    }

    #[test]
    fn test_wildcard_domain_storage() {
        let tmp = tempfile::tempdir().unwrap();
        let mut store = CertStore::new(tmp.path());

        store.store("*.example.com".to_string(), make_test_bundle("*.example.com")).unwrap();

        // Directory should use sanitized name
        assert!(tmp.path().join("_wildcard_.example.com").exists());

        let mut store2 = CertStore::new(tmp.path());
        store2.load_all().unwrap();
        assert!(store2.has("*.example.com"));
    }
}