chore: refactors crl watcher

Signed-off-by: nilekh <[email protected]>

Author: nilekhc

src/proxy.rs

@@ -173,6 +173,7 @@ pub struct Proxy {
     outbound: Outbound,
     socks5: Option<Socks5>,
     policy_watcher: PolicyWatcher,
+    crl_watcher: Option<crate::tls::crl_watcher::CrlWatcher>,
 }
 
 pub struct LocalWorkloadInformation {
@@ -302,10 +303,6 @@ impl Proxy {
         // We setup all the listeners first so we can capture any errors that should block startup
         let inbound = Inbound::new(pi.clone(), drain.clone()).await?;
 
-        if let Some(ref crl_mgr) = pi.crl_manager {
-            crl_mgr.register_connection_manager(pi.connection_manager.clone());
-        }
-
         // This exists for `direct` integ tests, no other reason
         #[cfg(any(test, feature = "testing"))]
         if pi.cfg.fake_self_inbound {
@@ -326,15 +323,29 @@ impl Proxy {
         } else {
             None
         };
-        let policy_watcher =
-            PolicyWatcher::new(pi.state.clone(), drain, pi.connection_manager.clone());
+        let policy_watcher = PolicyWatcher::new(
+            pi.state.clone(),
+            drain.clone(),
+            pi.connection_manager.clone(),
+        );
+
+        let crl_watcher = if let Some(ref crl_mgr) = pi.crl_manager {
+            Some(crate::tls::crl_watcher::CrlWatcher::new(
+                crl_mgr.clone(),
+                drain,
+                pi.connection_manager.clone(),
+            ))
+        } else {
+            None
+        };
 
         Ok(Proxy {
             inbound,
             inbound_passthrough,
             outbound,
             socks5,
             policy_watcher,
+            crl_watcher,
         })
     }
 
@@ -350,6 +361,10 @@ impl Proxy {
             tasks.push(tokio::spawn(socks5.run().in_current_span()));
         };
 
+        if let Some(crl_watcher) = self.crl_watcher {
+            tasks.push(tokio::spawn(crl_watcher.run().in_current_span()));
+        };
+
         futures::future::join_all(tasks).await;
     }
 

src/proxy/connection_manager.rs

@@ -59,7 +59,6 @@ impl ConnectionDrain {
 pub struct ConnectionManager {
     drains: Arc<RwLock<HashMap<InboundConnection, ConnectionDrain>>>,
     outbound_connections: Arc<RwLock<HashSet<OutboundConnection>>>,
-    close_tx: Option<tokio::sync::mpsc::UnboundedSender<InboundConnection>>,
 }
 
 impl std::fmt::Debug for ConnectionManager {
@@ -73,7 +72,6 @@ impl Default for ConnectionManager {
         ConnectionManager {
             drains: Arc::new(RwLock::new(HashMap::new())),
             outbound_connections: Arc::new(RwLock::new(HashSet::new())),
-            close_tx: None,
         }
     }
 }
@@ -159,38 +157,6 @@ pub struct InboundConnection {
 }
 
 impl ConnectionManager {
-    // Initialize the connection manager with CRL support
-    // This spawns an async task that handles connection closes from the CRL watcher
-    pub fn new_with_crl_support() -> Self {
-        let (close_tx, mut close_rx) = tokio::sync::mpsc::unbounded_channel::<InboundConnection>();
-
-        let cm = ConnectionManager {
-            drains: Arc::new(RwLock::new(HashMap::new())),
-            outbound_connections: Arc::new(RwLock::new(HashSet::new())),
-            close_tx: Some(close_tx),
-        };
-
-        // Spawn async task to handle close requests
-        let cm_clone = ConnectionManager {
-            drains: cm.drains.clone(),
-            outbound_connections: cm.outbound_connections.clone(),
-            close_tx: None,
-        };
-
-        tokio::spawn(async move {
-            while let Some(conn) = close_rx.recv().await {
-                tracing::debug!(
-                    "Processing close request for connection from {}",
-                    conn.ctx.conn.src
-                );
-                cm_clone.close(&conn).await;
-            }
-            tracing::debug!("Connection close handler terminated");
-        });
-
-        cm
-    }
-
     pub fn track_outbound(
         &self,
         src: SocketAddr,
@@ -306,7 +272,7 @@ impl ConnectionManager {
     }
 
     // signal all connections listening to this channel to take action (typically terminate traffic)
-    async fn close(&self, c: &InboundConnection) {
+    pub async fn close(&self, c: &InboundConnection) {
         let drain = { self.drains.write().expect("mutex").remove(c) };
         if let Some(cd) = drain {
             cd.drain().await;
@@ -321,120 +287,6 @@ impl ConnectionManager {
         // potentially large copy under read lock, could require optimization
         self.drains.read().expect("mutex").keys().cloned().collect()
     }
-
-    /// Close only inbound connections whose client certificates have been revoked
-    /// This is called when CRL updates with new revocations
-    pub fn close_revoked_connections(&self, revoked_serials: &std::collections::HashSet<Vec<u8>>) {
-        let conns = self.connections();
-
-        tracing::debug!(
-            "checking {} active inbound connections against {} revoked serial(s)",
-            conns.len(),
-            revoked_serials.len()
-        );
-
-        // Log all revoked serials we're checking against
-        for (idx, serial) in revoked_serials.iter().enumerate() {
-            tracing::debug!(
-                "revoked serial {}: {} bytes (hex: {})",
-                idx + 1,
-                serial.len(),
-                serial
-                    .iter()
-                    .map(|b| format!("{:02x}", b))
-                    .collect::<String>()
-            );
-        }
-
-        let mut closed_count = 0;
-        let mut no_serial_count = 0;
-
-        for (idx, conn) in conns.iter().enumerate() {
-            tracing::debug!(
-                "connection {}: src={}, has_serials={}",
-                idx + 1,
-                conn.ctx.conn.src,
-                conn.client_cert_serials.is_some()
-            );
-
-            if let Some(ref serials) = conn.client_cert_serials {
-                tracing::debug!(
-                    "  connection {} has {} certificate(s) in chain",
-                    idx + 1,
-                    serials.len()
-                );
-
-                // Check if ANY certificate in the chain is revoked
-                let mut is_revoked = false;
-                for (cert_idx, serial) in serials.iter().enumerate() {
-                    let serial_hex = serial
-                        .iter()
-                        .map(|b| format!("{:02x}", b))
-                        .collect::<String>();
-                    tracing::debug!(
-                        "    cert {} serial: {} bytes (hex: {})",
-                        cert_idx,
-                        serial.len(),
-                        serial_hex
-                    );
-
-                    if revoked_serials.contains(serial) {
-                        is_revoked = true;
-                        tracing::warn!("    cert {} serial MATCHES revoked serial!", cert_idx);
-                        break;
-                    }
-                }
-
-                tracing::debug!(
-                    "  connection {} has revoked certificate: {}",
-                    idx + 1,
-                    is_revoked
-                );
-
-                if is_revoked {
-                    tracing::warn!(
-                        "closing inbound connection {} from {} due to revoked certificate in chain",
-                        idx + 1,
-                        conn.ctx.conn.src
-                    );
-
-                    // Send close request through channel (works from blocking context)
-                    if let Some(ref tx) = self.close_tx {
-                        if let Err(e) = tx.send(conn.clone()) {
-                            tracing::error!("failed to send close request: {}", e);
-                        } else {
-                            closed_count += 1;
-                        }
-                    } else {
-                        tracing::warn!("CRL support not initialized - cannot close connection");
-                    }
-                }
-            } else {
-                no_serial_count += 1;
-                tracing::debug!(
-                    "  connection {} has no client certificate serials (skipping)",
-                    idx + 1
-                );
-            }
-        }
-
-        tracing::info!(
-            "connection closure summary: {} total connections checked, {} with serials, {} without serials, {} closed",
-            conns.len(),
-            conns.len() - no_serial_count,
-            no_serial_count,
-            closed_count
-        );
-
-        if closed_count > 0 {
-            tracing::info!(
-                "closed {} inbound connection(s) with revoked certificates",
-                closed_count
-            );
-        } else {
-            tracing::debug!("no active connections with newly revoked certificates");
-        }
-    }
 }
 
 #[derive(serde::Serialize)]

src/proxy/h2/server.rs

@@ -25,7 +25,7 @@ use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, Ordering};
 use tokio::net::TcpStream;
 use tokio::sync::{oneshot, watch};
-use tracing::{Instrument, debug};
+use tracing::{Instrument, debug, info};
 
 pub struct H2Request {
     request: Parts,
@@ -99,6 +99,7 @@ pub async fn serve_connection<F, Fut>(
     s: tokio_rustls::server::TlsStream<TcpStream>,
     drain: DrainWatcher,
     mut force_shutdown: watch::Receiver<()>,
+    tls_drain: Option<DrainWatcher>,
     handler: F,
 ) -> Result<(), Error>
 where
@@ -169,6 +170,16 @@ where
                 conn.graceful_shutdown();
                 break;
             }
+            _tls_shutdown = async {
+                match &tls_drain {
+                    Some(d) => d.clone().wait_for_drain().await,
+                    None => std::future::pending().await,
+                }
+            } => {
+                info!("starting graceful drain (TLS connection certificate revoked)");
+                conn.graceful_shutdown();
+                break;
+            }
         }
     }
     // Signal to the ping_pong it should also stop.