Runtime Ownership hardened

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>

src/tls_front/cache.rs

@@ -1,13 +1,14 @@
-use std::collections::HashMap;
-use std::collections::hash_map::DefaultHasher;
-use std::hash::{Hash, Hasher};
+use std::collections::{HashMap, HashSet};
+use std::collections::hash_map::RandomState;
+use std::hash::{BuildHasher, Hash, Hasher};
 use std::net::IpAddr;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
-use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
 use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
 
 use tokio::sync::RwLock;
+use tokio::io::AsyncReadExt;
 use tokio::time::sleep;
 use tracing::{debug, info, warn};
 
@@ -16,21 +17,160 @@ use crate::tls_front::types::{
     TlsProfileSource,
 };
 
-const FULL_CERT_SENT_SWEEP_INTERVAL_SECS: u64 = 30;
-const FULL_CERT_SENT_MAX_IPS: usize = 65_536;
+const FULL_CERT_SENT_SWEEP_INTERVAL_SECS: u64 = 1;
+const FULL_CERT_SENT_MAX_ENTRIES: usize = 65_536;
 const FULL_CERT_SENT_SHARDS: usize = 64;
+const FULL_CERT_SENT_MAX_ENTRIES_PER_SHARD: usize =
+    (FULL_CERT_SENT_MAX_ENTRIES / FULL_CERT_SENT_SHARDS) * 2;
+const TLS_FRONT_DISK_ENTRY_MAX_BYTES: u64 = 1024 * 1024;
 
-static FULL_CERT_SENT_IPS_GAUGE: AtomicU64 = AtomicU64::new(0);
-static FULL_CERT_SENT_CAP_DROPS: AtomicU64 = AtomicU64::new(0);
-
-/// Current number of IPs tracked by the TLS full-cert budget gate.
-pub(crate) fn full_cert_sent_ips_for_metrics() -> u64 {
-    FULL_CERT_SENT_IPS_GAUGE.load(Ordering::Relaxed)
+#[derive(Debug, Clone, Hash, PartialEq, Eq)]
+struct FullCertBudgetKey {
+    domain: Arc<str>,
+    client_ip: IpAddr,
 }
 
-/// Number of new IPs denied a full-cert budget slot because the cap was reached.
-pub(crate) fn full_cert_sent_cap_drops_for_metrics() -> u64 {
-    FULL_CERT_SENT_CAP_DROPS.load(Ordering::Relaxed)
+#[derive(Debug)]
+struct FullCertBudgetEntry {
+    expires_at: Option<Instant>,
+}
+
+/// Process-owned, bounded TLS full-certificate admission history.
+#[derive(Debug)]
+pub(crate) struct TlsFullCertBudget {
+    shards: Vec<RwLock<HashMap<FullCertBudgetKey, FullCertBudgetEntry>>>,
+    hash_builder: RandomState,
+    entries: AtomicUsize,
+    cap_drops: AtomicU64,
+    last_sweep_epoch_secs: AtomicU64,
+    sweep_cursor: AtomicUsize,
+}
+
+impl TlsFullCertBudget {
+    /// Creates an empty process-owned full-certificate budget.
+    pub(crate) fn new() -> Self {
+        Self {
+            shards: (0..FULL_CERT_SENT_SHARDS)
+                .map(|_| RwLock::new(HashMap::new()))
+                .collect(),
+            hash_builder: RandomState::new(),
+            entries: AtomicUsize::new(0),
+            cap_drops: AtomicU64::new(0),
+            last_sweep_epoch_secs: AtomicU64::new(0),
+            sweep_cursor: AtomicUsize::new(0),
+        }
+    }
+
+    fn shard_index(&self, key: &FullCertBudgetKey) -> usize {
+        let mut hasher = self.hash_builder.build_hasher();
+        key.hash(&mut hasher);
+        (hasher.finish() as usize) % FULL_CERT_SENT_SHARDS
+    }
+
+    fn decrement_entries(&self, amount: usize) {
+        if amount == 0 {
+            return;
+        }
+        let _ = self
+            .entries
+            .fetch_update(Ordering::AcqRel, Ordering::Relaxed, |current| {
+                Some(current.saturating_sub(amount))
+            });
+    }
+
+    fn try_reserve_entry(&self) -> bool {
+        let mut current = self.entries.load(Ordering::Relaxed);
+        loop {
+            if current >= FULL_CERT_SENT_MAX_ENTRIES {
+                return false;
+            }
+            match self.entries.compare_exchange_weak(
+                current,
+                current.saturating_add(1),
+                Ordering::AcqRel,
+                Ordering::Relaxed,
+            ) {
+                Ok(_) => return true,
+                Err(actual) => current = actual,
+            }
+        }
+    }
+
+    async fn sweep_one_shard(&self, now: Instant) {
+        let shard_index = self.sweep_cursor.fetch_add(1, Ordering::Relaxed)
+            % FULL_CERT_SENT_SHARDS;
+        let mut guard = self.shards[shard_index].write().await;
+        let before = guard.len();
+        guard.retain(|_, entry| entry.expires_at.map_or(true, |expires_at| expires_at > now));
+        self.decrement_entries(before.saturating_sub(guard.len()));
+    }
+
+    async fn maybe_sweep(&self, now: Instant) {
+        let now_epoch_secs = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap_or_default()
+            .as_secs();
+        let should_sweep = self
+            .last_sweep_epoch_secs
+            .fetch_update(Ordering::AcqRel, Ordering::Relaxed, |last_sweep| {
+                if now_epoch_secs.saturating_sub(last_sweep)
+                    >= FULL_CERT_SENT_SWEEP_INTERVAL_SECS
+                {
+                    Some(now_epoch_secs)
+                } else {
+                    None
+                }
+            })
+            .is_ok();
+        if should_sweep {
+            self.sweep_one_shard(now).await;
+        }
+    }
+
+    async fn take(&self, domain: Arc<str>, client_ip: IpAddr, ttl: Duration) -> bool {
+        if ttl.is_zero() {
+            return true;
+        }
+
+        let now = Instant::now();
+        self.maybe_sweep(now).await;
+        let expires_at = now.checked_add(ttl);
+        let key = FullCertBudgetKey { domain, client_ip };
+        let shard_index = self.shard_index(&key);
+        let mut guard = self.shards[shard_index].write().await;
+
+        if let Some(entry) = guard.get_mut(&key) {
+            if entry.expires_at.is_some_and(|expires_at| expires_at <= now) {
+                entry.expires_at = expires_at;
+                return true;
+            }
+            return false;
+        }
+
+        if guard.len() >= FULL_CERT_SENT_MAX_ENTRIES_PER_SHARD {
+            let before = guard.len();
+            guard.retain(|_, entry| {
+                entry.expires_at.map_or(true, |expires_at| expires_at > now)
+            });
+            self.decrement_entries(before.saturating_sub(guard.len()));
+        }
+        if guard.len() >= FULL_CERT_SENT_MAX_ENTRIES_PER_SHARD || !self.try_reserve_entry() {
+            self.cap_drops.fetch_add(1, Ordering::Relaxed);
+            return false;
+        }
+        guard.insert(key, FullCertBudgetEntry { expires_at });
+        true
+    }
+
+    /// Returns the current number of retained domain and client IP entries.
+    pub(crate) fn entries_for_metrics(&self) -> u64 {
+        self.entries.load(Ordering::Relaxed) as u64
+    }
+
+    /// Returns the cumulative number of entries rejected by hard bounds.
+    pub(crate) fn cap_drops_for_metrics(&self) -> u64 {
+        self.cap_drops.load(Ordering::Relaxed)
+    }
 }
 
 /// Lightweight in-memory + optional on-disk cache for TLS fronting data.
@@ -38,8 +178,9 @@ pub(crate) fn full_cert_sent_cap_drops_for_metrics() -> u64 {
 pub struct TlsFrontCache {
     memory: RwLock<HashMap<String, Arc<CachedTlsData>>>,
     default: Arc<CachedTlsData>,
-    full_cert_sent_shards: Vec<RwLock<HashMap<IpAddr, Instant>>>,
-    full_cert_sent_last_sweep_epoch_secs: AtomicU64,
+    full_cert_budget: Arc<TlsFullCertBudget>,
+    full_cert_domain_keys: HashMap<String, Arc<str>>,
+    disk_entry_names: HashSet<String>,
     disk_path: PathBuf,
 }
 
@@ -91,6 +232,21 @@ fn key_share_group_label(group: Option<u16>) -> &'static str {
 #[allow(dead_code)]
 impl TlsFrontCache {
     pub fn new(domains: &[String], default_len: usize, disk_path: impl AsRef<Path>) -> Self {
+        Self::new_with_full_cert_budget(
+            domains,
+            default_len,
+            disk_path,
+            Arc::new(TlsFullCertBudget::new()),
+        )
+    }
+
+    /// Creates a generation-local cache backed by the process-owned full-cert budget.
+    pub(crate) fn new_with_full_cert_budget(
+        domains: &[String],
+        default_len: usize,
+        disk_path: impl AsRef<Path>,
+        full_cert_budget: Arc<TlsFullCertBudget>,
+    ) -> Self {
         let default_template = ParsedServerHello {
             version: [0x03, 0x03],
             random: [0u8; 32],
@@ -112,17 +268,24 @@ impl TlsFrontCache {
         });
 
         let mut map = HashMap::new();
+        let mut full_cert_domain_keys = HashMap::new();
+        let mut disk_entry_names = HashSet::new();
         for d in domains {
             map.insert(d.clone(), default.clone());
+            disk_entry_names.insert(format!("{}.json", d.replace(['/', '\\'], "_")));
+            let canonical: Arc<str> = Arc::from(normalize_dns_name(d));
+            full_cert_domain_keys.insert(d.clone(), canonical.clone());
+            full_cert_domain_keys
+                .entry(canonical.to_string())
+                .or_insert(canonical);
         }
 
         Self {
             memory: RwLock::new(map),
             default,
-            full_cert_sent_shards: (0..FULL_CERT_SENT_SHARDS)
-                .map(|_| RwLock::new(HashMap::new()))
-                .collect(),
-            full_cert_sent_last_sweep_epoch_secs: AtomicU64::new(0),
+            full_cert_budget,
+            full_cert_domain_keys,
+            disk_entry_names,
             disk_path: disk_path.as_ref().to_path_buf(),
         }
     }
@@ -202,114 +365,66 @@ impl TlsFrontCache {
             .collect()
     }
 
-    fn full_cert_sent_shard_index(client_ip: IpAddr) -> usize {
-        let mut hasher = DefaultHasher::new();
-        client_ip.hash(&mut hasher);
-        (hasher.finish() as usize) % FULL_CERT_SENT_SHARDS
+    fn full_cert_domain_key(&self, domain: &str) -> Arc<str> {
+        self.full_cert_domain_keys
+            .get(domain)
+            .cloned()
+            .unwrap_or_else(|| Arc::from(normalize_dns_name(domain)))
     }
 
-    fn full_cert_sent_shard(&self, client_ip: IpAddr) -> &RwLock<HashMap<IpAddr, Instant>> {
-        &self.full_cert_sent_shards[Self::full_cert_sent_shard_index(client_ip)]
+    /// Returns true when the selected domain and client IP may receive a full cert payload.
+    pub async fn take_full_cert_budget_for_ip(
+        &self,
+        domain: &str,
+        client_ip: IpAddr,
+        ttl: Duration,
+    ) -> bool {
+        self.full_cert_budget
+            .take(self.full_cert_domain_key(domain), client_ip, ttl)
+            .await
     }
 
-    fn decrement_full_cert_sent_entries(amount: usize) {
-        if amount == 0 {
-            return;
-        }
-        let amount = amount as u64;
-        let _ =
-            FULL_CERT_SENT_IPS_GAUGE.fetch_update(Ordering::AcqRel, Ordering::Relaxed, |current| {
-                Some(current.saturating_sub(amount))
-            });
+    /// Returns the current process-owned full-cert budget entry count.
+    pub(crate) fn full_cert_budget_entries_for_metrics(&self) -> u64 {
+        self.full_cert_budget.entries_for_metrics()
     }
 
-    fn try_reserve_full_cert_sent_entry() -> bool {
-        let mut current = FULL_CERT_SENT_IPS_GAUGE.load(Ordering::Relaxed);
-        loop {
-            if current >= FULL_CERT_SENT_MAX_IPS as u64 {
-                return false;
-            }
-            match FULL_CERT_SENT_IPS_GAUGE.compare_exchange_weak(
-                current,
-                current.saturating_add(1),
-                Ordering::AcqRel,
-                Ordering::Relaxed,
-            ) {
-                Ok(_) => return true,
-                Err(actual) => current = actual,
-            }
-        }
-    }
-
-    async fn sweep_full_cert_sent_shards(&self, now: Instant, ttl: Duration) {
-        for shard in &self.full_cert_sent_shards {
-            let mut guard = shard.write().await;
-            let before = guard.len();
-            guard.retain(|_, seen_at| now.duration_since(*seen_at) < ttl);
-            Self::decrement_full_cert_sent_entries(before.saturating_sub(guard.len()));
-        }
-    }
-
-    /// Returns true when full cert payload should be sent for client_ip
-    /// according to TTL policy.
-    pub async fn take_full_cert_budget_for_ip(&self, client_ip: IpAddr, ttl: Duration) -> bool {
-        if ttl.is_zero() {
-            return true;
-        }
-
-        let now = Instant::now();
-        let now_epoch_secs = SystemTime::now()
-            .duration_since(UNIX_EPOCH)
-            .unwrap_or_default()
-            .as_secs();
-        let should_sweep = self
-            .full_cert_sent_last_sweep_epoch_secs
-            .fetch_update(Ordering::AcqRel, Ordering::Relaxed, |last_sweep| {
-                if now_epoch_secs.saturating_sub(last_sweep) >= FULL_CERT_SENT_SWEEP_INTERVAL_SECS {
-                    Some(now_epoch_secs)
-                } else {
-                    None
-                }
-            })
-            .is_ok();
-
-        if should_sweep {
-            self.sweep_full_cert_sent_shards(now, ttl).await;
-        }
-
-        let mut guard = self.full_cert_sent_shard(client_ip).write().await;
-        let allowed = match guard.get_mut(&client_ip) {
-            Some(seen_at) => {
-                if now.duration_since(*seen_at) >= ttl {
-                    *seen_at = now;
-                    true
-                } else {
-                    false
-                }
-            }
-            None => {
-                if !Self::try_reserve_full_cert_sent_entry() {
-                    FULL_CERT_SENT_CAP_DROPS.fetch_add(1, Ordering::Relaxed);
-                    return false;
-                }
-                guard.insert(client_ip, now);
-                true
-            }
-        };
-        allowed
+    /// Returns the cumulative process-owned full-cert budget cap drops.
+    pub(crate) fn full_cert_budget_cap_drops_for_metrics(&self) -> u64 {
+        self.full_cert_budget.cap_drops_for_metrics()
     }
 
     #[cfg(test)]
-    async fn insert_full_cert_sent_for_tests(&self, client_ip: IpAddr, seen_at: Instant) {
-        let mut guard = self.full_cert_sent_shard(client_ip).write().await;
-        if guard.insert(client_ip, seen_at).is_none() {
-            FULL_CERT_SENT_IPS_GAUGE.fetch_add(1, Ordering::Relaxed);
+    async fn insert_full_cert_sent_for_tests(
+        &self,
+        domain: &str,
+        client_ip: IpAddr,
+        expires_at: Instant,
+    ) {
+        let key = FullCertBudgetKey {
+            domain: self.full_cert_domain_key(domain),
+            client_ip,
+        };
+        let shard_index = self.full_cert_budget.shard_index(&key);
+        let mut guard = self.full_cert_budget.shards[shard_index].write().await;
+        if guard
+            .insert(
+                key,
+                FullCertBudgetEntry {
+                    expires_at: Some(expires_at),
+                },
+            )
+            .is_none()
+        {
+            self.full_cert_budget
+                .entries
+                .fetch_add(1, Ordering::Relaxed);
         }
     }
 
     #[cfg(test)]
     async fn full_cert_sent_is_empty_for_tests(&self) -> bool {
-        for shard in &self.full_cert_sent_shards {
+        for shard in &self.full_cert_budget.shards {
             if !shard.read().await.is_empty() {
                 return false;
             }
@@ -318,11 +433,16 @@ impl TlsFrontCache {
     }
 
     #[cfg(test)]
-    async fn full_cert_sent_contains_for_tests(&self, client_ip: IpAddr) -> bool {
-        self.full_cert_sent_shard(client_ip)
+    async fn full_cert_sent_contains_for_tests(&self, domain: &str, client_ip: IpAddr) -> bool {
+        let key = FullCertBudgetKey {
+            domain: self.full_cert_domain_key(domain),
+            client_ip,
+        };
+        let shard_index = self.full_cert_budget.shard_index(&key);
+        self.full_cert_budget.shards[shard_index]
             .read()
             .await
-            .contains_key(&client_ip)
+            .contains_key(&key)
     }
 
     pub async fn set(&self, domain: &str, data: CachedTlsData) {
@@ -336,54 +456,60 @@ impl TlsFrontCache {
             return;
         }
         let mut loaded = 0usize;
-        if let Ok(mut dir) = tokio::fs::read_dir(&path).await {
-            while let Ok(Some(entry)) = dir.next_entry().await {
-                if let Ok(name) = entry.file_name().into_string() {
-                    if !name.ends_with(".json") {
-                        continue;
-                    }
-                    if let Ok(data) = tokio::fs::read(entry.path()).await
-                        && let Ok(mut cached) = serde_json::from_slice::<CachedTlsData>(&data)
-                    {
-                        if cached.domain.is_empty()
-                            || cached.domain.len() > 255
-                            || !cached
-                                .domain
-                                .chars()
-                                .all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '-')
-                        {
-                            warn!(file = %name, "Skipping TLS cache entry with invalid domain");
-                            continue;
-                        }
-                        if !cert_info_matches_domain(&cached) {
-                            warn!(
-                                file = %name,
-                                domain = %cached.domain,
-                                "Skipping TLS cache entry with mismatched certificate metadata"
-                            );
-                            continue;
-                        }
-                        // fetched_at is skipped during deserialization; approximate with file mtime if available.
-                        if let Ok(meta) = entry.metadata().await
-                            && let Ok(modified) = meta.modified()
-                        {
-                            cached.fetched_at = modified;
-                        }
-                        // Drop entries older than 72h
-                        if let Ok(age) = cached.fetched_at.elapsed()
-                            && age > Duration::from_secs(72 * 3600)
-                        {
-                            warn!(domain = %cached.domain, "Skipping stale TLS cache entry (>72h)");
-                            continue;
-                        }
-                        cached
-                            .behavior_profile
-                            .refresh_server_hello_summary(&cached.server_hello_template);
-                        let domain = cached.domain.clone();
-                        self.set(&domain, cached).await;
-                        loaded += 1;
-                    }
+        for name in &self.disk_entry_names {
+            let entry_path = path.join(name);
+            let Ok(metadata) = tokio::fs::symlink_metadata(&entry_path).await else {
+                continue;
+            };
+            if !metadata.file_type().is_file() {
+                continue;
+            }
+            if let Ok(data) = read_disk_entry_bounded(&entry_path).await
+                && let Ok(mut cached) = serde_json::from_slice::<CachedTlsData>(&data)
+            {
+                if cached.domain.is_empty()
+                    || cached.domain.len() > 255
+                    || !cached
+                        .domain
+                        .chars()
+                        .all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '-')
+                {
+                    warn!(file = %name, "Skipping TLS cache entry with invalid domain");
+                    continue;
                 }
+                if !self.full_cert_domain_keys.contains_key(&cached.domain) {
+                    warn!(
+                        file = %name,
+                        domain = %cached.domain,
+                        "Skipping TLS cache entry outside configured domains"
+                    );
+                    continue;
+                }
+                if !cert_info_matches_domain(&cached) {
+                    warn!(
+                        file = %name,
+                        domain = %cached.domain,
+                        "Skipping TLS cache entry with mismatched certificate metadata"
+                    );
+                    continue;
+                }
+                // fetched_at is skipped during deserialization; approximate with file mtime if available.
+                if let Ok(modified) = metadata.modified() {
+                    cached.fetched_at = modified;
+                }
+                // Drop entries older than 72h
+                if let Ok(age) = cached.fetched_at.elapsed()
+                    && age > Duration::from_secs(72 * 3600)
+                {
+                    warn!(domain = %cached.domain, "Skipping stale TLS cache entry (>72h)");
+                    continue;
+                }
+                cached
+                    .behavior_profile
+                    .refresh_server_hello_summary(&cached.server_hello_template);
+                let domain = cached.domain.clone();
+                self.set(&domain, cached).await;
+                loaded += 1;
             }
         }
         if loaded > 0 {
@@ -398,6 +524,14 @@ impl TlsFrontCache {
         let fname = format!("{}.json", domain.replace(['/', '\\'], "_"));
         let path = self.disk_path.join(fname);
         if let Ok(json) = serde_json::to_vec_pretty(data) {
+            if json.len() as u64 > TLS_FRONT_DISK_ENTRY_MAX_BYTES {
+                warn!(
+                    domain,
+                    bytes = json.len(),
+                    "Skipping oversized TLS cache persistence"
+                );
+                return;
+            }
             // best-effort write
             let _ = tokio::fs::write(path, json).await;
         }
@@ -464,6 +598,27 @@ impl TlsFrontCache {
     }
 }
 
+async fn read_disk_entry_bounded(path: &Path) -> std::io::Result<Vec<u8>> {
+    let file = tokio::fs::File::open(path).await?;
+    if file.metadata().await?.len() > TLS_FRONT_DISK_ENTRY_MAX_BYTES {
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidData,
+            "TLS cache entry exceeds the 1 MiB limit",
+        ));
+    }
+    let mut bytes = Vec::new();
+    file.take(TLS_FRONT_DISK_ENTRY_MAX_BYTES.saturating_add(1))
+        .read_to_end(&mut bytes)
+        .await?;
+    if bytes.len() as u64 > TLS_FRONT_DISK_ENTRY_MAX_BYTES {
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidData,
+            "TLS cache entry grew beyond the 1 MiB limit while reading",
+        ));
+    }
+    Ok(bytes)
+}
+
 fn cert_info_matches_domain(cached: &CachedTlsData) -> bool {
     let Some(cert_info) = cached.cert_info.as_ref() else {
         return true;
@@ -581,12 +736,24 @@ mod tests {
         let ip: IpAddr = "127.0.0.1".parse().expect("ip");
         let ttl = Duration::from_millis(80);
 
-        assert!(cache.take_full_cert_budget_for_ip(ip, ttl).await);
-        assert!(!cache.take_full_cert_budget_for_ip(ip, ttl).await);
+        assert!(
+            cache
+                .take_full_cert_budget_for_ip("example.com", ip, ttl)
+                .await
+        );
+        assert!(
+            !cache
+                .take_full_cert_budget_for_ip("example.com", ip, ttl)
+                .await
+        );
 
         tokio::time::sleep(Duration::from_millis(90)).await;
 
-        assert!(cache.take_full_cert_budget_for_ip(ip, ttl).await);
+        assert!(
+            cache
+                .take_full_cert_budget_for_ip("example.com", ip, ttl)
+                .await
+        );
     }
 
     #[tokio::test]
@@ -601,7 +768,11 @@ mod tests {
                 ((idx >> 8) & 0xff) as u8,
                 (idx & 0xff) as u8,
             ));
-            assert!(cache.take_full_cert_budget_for_ip(ip, ttl).await);
+            assert!(
+                cache
+                    .take_full_cert_budget_for_ip("example.com", ip, ttl)
+                    .await
+            );
         }
 
         assert!(cache.full_cert_sent_is_empty_for_tests().await);
@@ -613,21 +784,42 @@ mod tests {
         let stale_ip: IpAddr = "127.0.0.1".parse().expect("ip");
         let new_ip: IpAddr = "127.0.0.2".parse().expect("ip");
         let ttl = Duration::from_secs(1);
-        let stale_seen_at = Instant::now()
-            .checked_sub(Duration::from_secs(10))
+        let stale_expires_at = Instant::now()
+            .checked_sub(Duration::from_secs(1))
             .unwrap_or_else(Instant::now);
 
         cache
-            .insert_full_cert_sent_for_tests(stale_ip, stale_seen_at)
+            .insert_full_cert_sent_for_tests("example.com", stale_ip, stale_expires_at)
             .await;
+        let stale_key = FullCertBudgetKey {
+            domain: cache.full_cert_domain_key("example.com"),
+            client_ip: stale_ip,
+        };
+        cache.full_cert_budget.sweep_cursor.store(
+            cache.full_cert_budget.shard_index(&stale_key),
+            Ordering::Relaxed,
+        );
         cache
-            .full_cert_sent_last_sweep_epoch_secs
+            .full_cert_budget
+            .last_sweep_epoch_secs
             .store(0, Ordering::Relaxed);
 
-        assert!(cache.take_full_cert_budget_for_ip(new_ip, ttl).await);
+        assert!(
+            cache
+                .take_full_cert_budget_for_ip("example.com", new_ip, ttl)
+                .await
+        );
 
-        assert!(!cache.full_cert_sent_contains_for_tests(stale_ip).await);
-        assert!(cache.full_cert_sent_contains_for_tests(new_ip).await);
+        assert!(
+            !cache
+                .full_cert_sent_contains_for_tests("example.com", stale_ip)
+                .await
+        );
+        assert!(
+            cache
+                .full_cert_sent_contains_for_tests("example.com", new_ip)
+                .await
+        );
     }
 
     #[tokio::test]
@@ -636,8 +828,8 @@ mod tests {
         let stale_ip: IpAddr = "127.0.0.1".parse().expect("ip");
         let new_ip: IpAddr = "127.0.0.2".parse().expect("ip");
         let ttl = Duration::from_secs(1);
-        let stale_seen_at = Instant::now()
-            .checked_sub(Duration::from_secs(10))
+        let stale_expires_at = Instant::now()
+            .checked_sub(Duration::from_secs(1))
             .unwrap_or_else(Instant::now);
         let now_epoch_secs = SystemTime::now()
             .duration_since(UNIX_EPOCH)
@@ -645,15 +837,126 @@ mod tests {
             .as_secs();
 
         cache
-            .insert_full_cert_sent_for_tests(stale_ip, stale_seen_at)
+            .insert_full_cert_sent_for_tests("example.com", stale_ip, stale_expires_at)
             .await;
         cache
-            .full_cert_sent_last_sweep_epoch_secs
+            .full_cert_budget
+            .last_sweep_epoch_secs
             .store(now_epoch_secs, Ordering::Relaxed);
 
-        assert!(cache.take_full_cert_budget_for_ip(new_ip, ttl).await);
+        assert!(
+            cache
+                .take_full_cert_budget_for_ip("example.com", new_ip, ttl)
+                .await
+        );
 
-        assert!(cache.full_cert_sent_contains_for_tests(stale_ip).await);
-        assert!(cache.full_cert_sent_contains_for_tests(new_ip).await);
+        assert!(
+            cache
+                .full_cert_sent_contains_for_tests("example.com", stale_ip)
+                .await
+        );
+        assert!(
+            cache
+                .full_cert_sent_contains_for_tests("example.com", new_ip)
+                .await
+        );
+    }
+
+    #[tokio::test]
+    async fn full_cert_budget_is_shared_across_cache_generations_and_scoped_by_domain() {
+        let budget = Arc::new(TlsFullCertBudget::new());
+        let domains = ["one.example".to_string(), "two.example".to_string()];
+        let first = TlsFrontCache::new_with_full_cert_budget(
+            &domains,
+            1024,
+            "tlsfront-test-cache",
+            budget.clone(),
+        );
+        let second = TlsFrontCache::new_with_full_cert_budget(
+            &domains,
+            1024,
+            "tlsfront-test-cache",
+            budget,
+        );
+        let ip: IpAddr = "127.0.0.1".parse().expect("ip");
+        let ttl = Duration::from_secs(60);
+
+        assert!(
+            first
+                .take_full_cert_budget_for_ip("one.example", ip, ttl)
+                .await
+        );
+        assert!(
+            !second
+                .take_full_cert_budget_for_ip("one.example", ip, ttl)
+                .await
+        );
+        assert!(
+            second
+                .take_full_cert_budget_for_ip("two.example", ip, ttl)
+                .await
+        );
+        assert_eq!(second.full_cert_budget_entries_for_metrics(), 2);
+    }
+
+    #[tokio::test]
+    async fn existing_full_cert_entry_keeps_its_own_expiry_after_ttl_change() {
+        let cache = TlsFrontCache::new(&["example.com".to_string()], 1024, "tlsfront-test-cache");
+        let ip: IpAddr = "127.0.0.1".parse().expect("ip");
+
+        assert!(
+            cache
+                .take_full_cert_budget_for_ip("example.com", ip, Duration::from_millis(80))
+                .await
+        );
+        tokio::time::sleep(Duration::from_millis(20)).await;
+        assert!(
+            !cache
+                .take_full_cert_budget_for_ip("example.com", ip, Duration::from_millis(1))
+                .await
+        );
+        tokio::time::sleep(Duration::from_millis(70)).await;
+        assert!(
+            cache
+                .take_full_cert_budget_for_ip("example.com", ip, Duration::from_millis(1))
+                .await
+        );
+    }
+
+    #[tokio::test]
+    async fn disk_reader_rejects_an_entry_above_the_hard_limit() {
+        let directory = tempfile::tempdir().unwrap();
+        let path = directory.path().join("oversized.json");
+        tokio::fs::write(
+            &path,
+            vec![0u8; TLS_FRONT_DISK_ENTRY_MAX_BYTES as usize + 1],
+        )
+        .await
+        .unwrap();
+
+        let error = read_disk_entry_bounded(&path).await.unwrap_err();
+
+        assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
+    }
+
+    #[cfg(unix)]
+    #[tokio::test]
+    async fn disk_loader_does_not_follow_a_configured_name_symlink() {
+        let directory = tempfile::tempdir().unwrap();
+        let target = directory.path().join("outside.json");
+        let cached = cached_with_cert_info("example.com", None, Vec::new());
+        tokio::fs::write(&target, serde_json::to_vec(&cached).unwrap())
+            .await
+            .unwrap();
+        std::os::unix::fs::symlink(&target, directory.path().join("example.com.json")).unwrap();
+        let cache = TlsFrontCache::new(
+            &["example.com".to_string()],
+            1024,
+            directory.path(),
+        );
+
+        cache.load_from_disk().await;
+
+        assert_eq!(cache.get("example.com").await.domain, "default");
     }
 }

src/tls_front/fetcher.rs

@@ -30,7 +30,6 @@ use x509_parser::prelude::FromDer;
 
 use crate::config::TlsFetchProfile;
 use crate::crypto::{SecureRandom, sha256};
-use crate::network::dns_overrides::resolve_socket_addr;
 use crate::protocol::constants::{
     TLS_RECORD_APPLICATION, TLS_RECORD_CHANGE_CIPHER, TLS_RECORD_HANDSHAKE,
 };
@@ -901,9 +900,6 @@ async fn connect_with_dns_override(
     port: u16,
     connect_timeout: Duration,
 ) -> Result<TcpStream> {
-    if let Some(addr) = resolve_socket_addr(host, port) {
-        return Ok(timeout(connect_timeout, TcpStream::connect(addr)).await??);
-    }
     Ok(timeout(connect_timeout, TcpStream::connect((host, port))).await??)
 }