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

use std::collections::HashMap;
use serde::{Deserialize, Serialize};

/// Certificate metadata stored alongside certs.
#[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,
}

/// In-memory certificate store.
///
/// All persistence is owned by the consumer (TypeScript side).
/// This struct is a thin HashMap wrapper used as a runtime cache.
pub struct CertStore {
    cache: HashMap<String, CertBundle>,
}

impl CertStore {
    /// Create a new empty cert store.
    pub fn new() -> Self {
        Self {
            cache: HashMap::new(),
        }
    }

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

    /// Store a certificate in the cache.
    pub fn store(&mut self, domain: String, bundle: CertBundle) {
        self.cache.insert(domain, bundle);
    }

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

    /// 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 the cache.
    pub fn remove(&mut self, domain: &str) -> bool {
        self.cache.remove(domain).is_some()
    }
}

impl Default for CertStore {
    fn default() -> Self {
        Self::new()
    }
}

#[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_get() {
        let mut store = CertStore::new();

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

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

    #[test]
    fn test_store_with_ca_cert() {
        let mut store = CertStore::new();

        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);

        let loaded = store.get("secure.com").unwrap();
        assert!(loaded.ca_pem.is_some());
    }

    #[test]
    fn test_multiple_certs() {
        let mut store = CertStore::new();

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

        assert_eq!(store.count(), 3);
        assert!(store.has("a.com"));
        assert!(store.has("b.com"));
        assert!(store.has("c.com"));
    }

    #[test]
    fn test_remove_cert() {
        let mut store = CertStore::new();

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

        let removed = store.remove("remove-me.com");
        assert!(removed);
        assert!(!store.has("remove-me.com"));
    }

    #[test]
    fn test_remove_nonexistent() {
        let mut store = CertStore::new();
        assert!(!store.remove("nonexistent.com"));
    }

    #[test]
    fn test_wildcard_domain() {
        let mut store = CertStore::new();

        store.store("*.example.com".to_string(), make_test_bundle("*.example.com"));
        assert!(store.has("*.example.com"));

        let loaded = store.get("*.example.com").unwrap();
        assert_eq!(loaded.metadata.domain, "*.example.com");
    }
}