TlsFrontCache + X509 Parser + GREASE Tolerance

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

src/main.rs

@@ -260,67 +260,6 @@ async fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
         info!("IP limits configured for {} users", config.access.user_max_unique_ips.len());
     }
 
-    // TLS front cache (optional emulation)
-    let mut tls_domains = Vec::with_capacity(1 + config.censorship.tls_domains.len());
-    tls_domains.push(config.censorship.tls_domain.clone());
-    for d in &config.censorship.tls_domains {
-        if !tls_domains.contains(d) {
-            tls_domains.push(d.clone());
-        }
-    }
-
-    let tls_cache: Option<Arc<TlsFrontCache>> = if config.censorship.tls_emulation {
-        let cache = Arc::new(TlsFrontCache::new(
-            &tls_domains,
-            config.censorship.fake_cert_len,
-            &config.censorship.tls_front_dir,
-        ));
-
-        let port = config.censorship.mask_port;
-        // Initial synchronous fetch to warm cache before serving clients.
-        for domain in tls_domains.clone() {
-            match crate::tls_front::fetcher::fetch_real_tls(
-                &domain,
-                port,
-                &domain,
-                Duration::from_secs(5),
-            )
-            .await
-            {
-                Ok(res) => cache.update_from_fetch(&domain, res).await,
-                Err(e) => warn!(domain = %domain, error = %e, "TLS emulation fetch failed"),
-            }
-        }
-
-        // Periodic refresh with jitter.
-        let cache_clone = cache.clone();
-        let domains = tls_domains.clone();
-        tokio::spawn(async move {
-            loop {
-                let base_secs = rand::rng().random_range(4 * 3600..=6 * 3600);
-                let jitter_secs = rand::rng().random_range(0..=7200);
-                tokio::time::sleep(Duration::from_secs(base_secs + jitter_secs)).await;
-                for domain in &domains {
-                    match crate::tls_front::fetcher::fetch_real_tls(
-                        domain,
-                        port,
-                        domain,
-                        Duration::from_secs(5),
-                    )
-                    .await
-                    {
-                        Ok(res) => cache_clone.update_from_fetch(domain, res).await,
-                        Err(e) => warn!(domain = %domain, error = %e, "TLS emulation refresh failed"),
-                    }
-                }
-            }
-        });
-
-        Some(cache)
-    } else {
-        None
-    };
-
     // Connection concurrency limit
     let _max_connections = Arc::new(Semaphore::new(10_000));
 
@@ -499,6 +438,72 @@ match crate::transport::middle_proxy::fetch_proxy_secret(proxy_secret_path).awai
     let upstream_manager = Arc::new(UpstreamManager::new(config.upstreams.clone()));
     let buffer_pool = Arc::new(BufferPool::with_config(16 * 1024, 4096));
 
+    // TLS front cache (optional emulation)
+    let mut tls_domains = Vec::with_capacity(1 + config.censorship.tls_domains.len());
+    tls_domains.push(config.censorship.tls_domain.clone());
+    for d in &config.censorship.tls_domains {
+        if !tls_domains.contains(d) {
+            tls_domains.push(d.clone());
+        }
+    }
+
+    let tls_cache: Option<Arc<TlsFrontCache>> = if config.censorship.tls_emulation {
+        let cache = Arc::new(TlsFrontCache::new(
+            &tls_domains,
+            config.censorship.fake_cert_len,
+            &config.censorship.tls_front_dir,
+        ));
+
+        cache.load_from_disk().await;
+
+        let port = config.censorship.mask_port;
+        // Initial synchronous fetch to warm cache before serving clients.
+        for domain in tls_domains.clone() {
+            match crate::tls_front::fetcher::fetch_real_tls(
+                &domain,
+                port,
+                &domain,
+                Duration::from_secs(5),
+                Some(upstream_manager.clone()),
+            )
+            .await
+            {
+                Ok(res) => cache.update_from_fetch(&domain, res).await,
+                Err(e) => warn!(domain = %domain, error = %e, "TLS emulation fetch failed"),
+            }
+        }
+
+        // Periodic refresh with jitter.
+        let cache_clone = cache.clone();
+        let domains = tls_domains.clone();
+        let upstream_for_task = upstream_manager.clone();
+        tokio::spawn(async move {
+            loop {
+                let base_secs = rand::rng().random_range(4 * 3600..=6 * 3600);
+                let jitter_secs = rand::rng().random_range(0..=7200);
+                tokio::time::sleep(Duration::from_secs(base_secs + jitter_secs)).await;
+                for domain in &domains {
+                    match crate::tls_front::fetcher::fetch_real_tls(
+                        domain,
+                        port,
+                        domain,
+                        Duration::from_secs(5),
+                        Some(upstream_for_task.clone()),
+                    )
+                    .await
+                    {
+                        Ok(res) => cache_clone.update_from_fetch(domain, res).await,
+                        Err(e) => warn!(domain = %domain, error = %e, "TLS emulation refresh failed"),
+                    }
+                }
+            }
+        });
+
+        Some(cache)
+    } else {
+        None
+    };
+
     // Middle-End ping before DC connectivity
     if let Some(ref pool) = me_pool {
         let me_results = run_me_ping(pool, &rng).await;

src/tls_front/cache.rs

@@ -59,6 +59,45 @@ impl TlsFrontCache {
         guard.insert(domain.to_string(), Arc::new(data));
     }
 
+    pub async fn load_from_disk(&self) {
+        let path = self.disk_path.clone();
+        if tokio::fs::create_dir_all(&path).await.is_err() {
+            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 {
+                        if let Ok(cached) = serde_json::from_slice::<CachedTlsData>(&data) {
+                            let domain = cached.domain.clone();
+                            self.set(&domain, cached).await;
+                            loaded += 1;
+                        }
+                    }
+                }
+            }
+        }
+        if loaded > 0 {
+            info!(count = loaded, "Loaded TLS cache entries from disk");
+        }
+    }
+
+    async fn persist(&self, domain: &str, data: &CachedTlsData) {
+        if tokio::fs::create_dir_all(&self.disk_path).await.is_err() {
+            return;
+        }
+        let fname = format!("{}.json", domain.replace(['/', '\\'], "_"));
+        let path = self.disk_path.join(fname);
+        if let Ok(json) = serde_json::to_vec_pretty(data) {
+            // best-effort write
+            let _ = tokio::fs::write(path, json).await;
+        }
+    }
+
     /// Spawn background updater that periodically refreshes cached domains using provided fetcher.
     pub fn spawn_updater<F>(
         self: Arc<Self>,
@@ -82,14 +121,15 @@ impl TlsFrontCache {
     pub async fn update_from_fetch(&self, domain: &str, fetched: TlsFetchResult) {
         let data = CachedTlsData {
             server_hello_template: fetched.server_hello_parsed,
-            cert_info: None,
+            cert_info: fetched.cert_info,
             app_data_records_sizes: fetched.app_data_records_sizes.clone(),
             total_app_data_len: fetched.total_app_data_len,
             fetched_at: SystemTime::now(),
             domain: domain.to_string(),
         };
 
-        self.set(domain, data).await;
+        self.set(domain, data.clone()).await;
+        self.persist(domain, &data).await;
         debug!(domain = %domain, len = fetched.total_app_data_len, "TLS cache updated");
     }
 

src/tls_front/fetcher.rs

@@ -14,9 +14,12 @@ use rustls::client::ClientConfig;
 use rustls::pki_types::{CertificateDer, ServerName, UnixTime};
 use rustls::{DigitallySignedStruct, Error as RustlsError};
 
+use x509_parser::prelude::FromDer;
+use x509_parser::certificate::X509Certificate;
+
 use crate::crypto::SecureRandom;
 use crate::protocol::constants::{TLS_RECORD_APPLICATION, TLS_RECORD_HANDSHAKE};
-use crate::tls_front::types::{ParsedServerHello, TlsExtension, TlsFetchResult};
+use crate::tls_front::types::{ParsedServerHello, TlsExtension, TlsFetchResult, ParsedCertificateInfo};
 
 /// No-op verifier: accept any certificate (we only need lengths and metadata).
 #[derive(Debug)]
@@ -266,6 +269,52 @@ fn parse_server_hello(body: &[u8]) -> Option<ParsedServerHello> {
     })
 }
 
+fn parse_cert_info(certs: &[CertificateDer<'static>]) -> Option<ParsedCertificateInfo> {
+    let first = certs.first()?;
+    let (_rem, cert) = X509Certificate::from_der(first.as_ref()).ok()?;
+
+    let not_before = Some(cert.validity().not_before.to_datetime().unix_timestamp());
+    let not_after = Some(cert.validity().not_after.to_datetime().unix_timestamp());
+
+    let issuer_cn = cert
+        .issuer()
+        .iter_common_name()
+        .next()
+        .and_then(|cn| cn.as_str().ok())
+        .map(|s| s.to_string());
+
+    let subject_cn = cert
+        .subject()
+        .iter_common_name()
+        .next()
+        .and_then(|cn| cn.as_str().ok())
+        .map(|s| s.to_string());
+
+    let san_names = cert
+        .subject_alternative_name()
+        .ok()
+        .flatten()
+        .map(|san| {
+            san.value
+                .general_names
+                .iter()
+                .filter_map(|gn| match gn {
+                    x509_parser::extensions::GeneralName::DNSName(n) => Some(n.to_string()),
+                    _ => None,
+                })
+                .collect::<Vec<_>>()
+        })
+        .unwrap_or_default();
+
+    Some(ParsedCertificateInfo {
+        not_after_unix: not_after,
+        not_before_unix: not_before,
+        issuer_cn,
+        subject_cn,
+        san_names,
+    })
+}
+
 async fn fetch_via_raw_tls(
     host: &str,
     port: u16,
@@ -318,6 +367,7 @@ async fn fetch_via_raw_tls(
             app_sizes
         },
         total_app_data_len,
+        cert_info: None,
     })
 }
 
@@ -327,6 +377,7 @@ pub async fn fetch_real_tls(
     port: u16,
     sni: &str,
     connect_timeout: Duration,
+    upstream: Option<std::sync::Arc<crate::transport::UpstreamManager>>,
 ) -> Result<TlsFetchResult> {
     // Preferred path: raw TLS probe for accurate record sizing
     match fetch_via_raw_tls(host, port, sni, connect_timeout).await {
@@ -337,8 +388,26 @@ pub async fn fetch_real_tls(
     }
 
     // Fallback: rustls handshake to at least get certificate sizes
-    let addr = format!("{host}:{port}");
+    let stream = if let Some(manager) = upstream {
-    let stream = timeout(connect_timeout, TcpStream::connect(addr)).await??;
+        // Resolve host to SocketAddr
+        if let Ok(mut addrs) = tokio::net::lookup_host((host, port)).await {
+            if let Some(addr) = addrs.find(|a| a.is_ipv4()) {
+                match manager.connect(addr, None, None).await {
+                    Ok(s) => s,
+                    Err(e) => {
+                        warn!(sni = %sni, error = %e, "Upstream connect failed, using direct connect");
+                        timeout(connect_timeout, TcpStream::connect((host, port))).await??
+                    }
+                }
+            } else {
+                timeout(connect_timeout, TcpStream::connect((host, port))).await??
+            }
+        } else {
+            timeout(connect_timeout, TcpStream::connect((host, port))).await??
+        }
+    } else {
+        timeout(connect_timeout, TcpStream::connect((host, port))).await??
+    };
 
     let config = build_client_config();
     let connector = TlsConnector::from(config);
@@ -362,6 +431,7 @@ pub async fn fetch_real_tls(
         .unwrap_or_default();
 
     let total_cert_len: usize = certs.iter().map(|c| c.len()).sum::<usize>().max(1024);
+    let cert_info = parse_cert_info(&certs);
 
     // Heuristic: split across two records if large to mimic real servers a bit.
     let app_data_records_sizes = if total_cert_len > 3000 {
@@ -390,5 +460,6 @@ pub async fn fetch_real_tls(
         server_hello_parsed: parsed,
         app_data_records_sizes: app_data_records_sizes.clone(),
         total_app_data_len: app_data_records_sizes.iter().sum(),
+        cert_info,
     })
 }

src/tls_front/types.rs

@@ -1,7 +1,8 @@
 use std::time::SystemTime;
+use serde::{Serialize, Deserialize};
 
 /// Parsed representation of an unencrypted TLS ServerHello.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ParsedServerHello {
     pub version: [u8; 2],
     pub random: [u8; 32],
@@ -12,14 +13,14 @@ pub struct ParsedServerHello {
 }
 
 /// Generic TLS extension container.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct TlsExtension {
     pub ext_type: u16,
     pub data: Vec<u8>,
 }
 
 /// Basic certificate metadata (optional, informative).
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ParsedCertificateInfo {
     pub not_after_unix: Option<i64>,
     pub not_before_unix: Option<i64>,
@@ -29,20 +30,26 @@ pub struct ParsedCertificateInfo {
 }
 
 /// Cached data per SNI used by the emulator.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct CachedTlsData {
     pub server_hello_template: ParsedServerHello,
     pub cert_info: Option<ParsedCertificateInfo>,
     pub app_data_records_sizes: Vec<usize>,
     pub total_app_data_len: usize,
+    #[serde(default = "now_system_time", skip_serializing, skip_deserializing)]
     pub fetched_at: SystemTime,
     pub domain: String,
 }
 
+fn now_system_time() -> SystemTime {
+    SystemTime::now()
+}
+
 /// Result of attempting to fetch real TLS artifacts.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct TlsFetchResult {
     pub server_hello_parsed: ParsedServerHello,
     pub app_data_records_sizes: Vec<usize>,
     pub total_app_data_len: usize,
+    pub cert_info: Option<ParsedCertificateInfo>,
 }

src/transport/proxy_protocol.rs

@@ -283,6 +283,58 @@ impl Default for ProxyProtocolV1Builder {
     }
 }
 
+/// Builder for PROXY protocol v2 header
+pub struct ProxyProtocolV2Builder {
+    src: Option<SocketAddr>,
+    dst: Option<SocketAddr>,
+}
+
+impl ProxyProtocolV2Builder {
+    pub fn new() -> Self {
+        Self { src: None, dst: None }
+    }
+
+    pub fn with_addrs(mut self, src: SocketAddr, dst: SocketAddr) -> Self {
+        self.src = Some(src);
+        self.dst = Some(dst);
+        self
+    }
+
+    pub fn build(&self) -> Vec<u8> {
+        let mut header = Vec::new();
+        header.extend_from_slice(PROXY_V2_SIGNATURE);
+        // version 2, PROXY command
+        header.push(0x21);
+
+        match (self.src, self.dst) {
+            (Some(SocketAddr::V4(src)), Some(SocketAddr::V4(dst))) => {
+                header.push(0x11); // INET + STREAM
+                header.extend_from_slice(&(12u16).to_be_bytes());
+                header.extend_from_slice(&src.ip().octets());
+                header.extend_from_slice(&dst.ip().octets());
+                header.extend_from_slice(&src.port().to_be_bytes());
+                header.extend_from_slice(&dst.port().to_be_bytes());
+            }
+            (Some(SocketAddr::V6(src)), Some(SocketAddr::V6(dst))) => {
+                header.push(0x21); // INET6 + STREAM
+                header.extend_from_slice(&(36u16).to_be_bytes());
+                header.extend_from_slice(&src.ip().octets());
+                header.extend_from_slice(&dst.ip().octets());
+                header.extend_from_slice(&src.port().to_be_bytes());
+                header.extend_from_slice(&dst.port().to_be_bytes());
+            }
+            _ => {
+                // LOCAL/UNSPEC: no address information
+                header[12] = 0x20; // version 2, LOCAL command
+                header.push(0x00);
+                header.extend_from_slice(&0u16.to_be_bytes());
+            }
+        }
+
+        header
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;