fix: preserve secure relay sessions during cleanup (#2393)

* fix: skip secure relay packets in peer conn filter
* fix: keep active secure relay sessions during GC

easytier/src/peers/peer_conn.rs

@@ -1,3 +1,4 @@
+use arc_swap::ArcSwapOption;
 use crossbeam::atomic::AtomicCell;
 use futures::{StreamExt, TryFutureExt};
 use std::{
@@ -10,13 +11,9 @@ use std::{
     },
 };
 
-#[cfg(feature = "hotpath")]
-use hotpath::wrap::std::sync::Mutex as StdMutex;
 #[cfg(feature = "hotpath")]
 use hotpath::wrap::tokio::sync::Mutex;
 #[cfg(not(feature = "hotpath"))]
-use std::sync::Mutex as StdMutex;
-#[cfg(not(feature = "hotpath"))]
 use tokio::sync::Mutex;
 
 use base64::Engine as _;
@@ -107,7 +104,7 @@ struct PeerSessionTunnelFilter {
     enabled: bool,
     my_peer_id: Arc<AtomicCell<PeerId>>,
     peer_id: Arc<AtomicCell<Option<PeerId>>>,
-    session: Arc<StdMutex<Option<Arc<PeerSession>>>>,
+    session: Arc<ArcSwapOption<PeerSession>>,
 }
 
 impl PeerSessionTunnelFilter {
@@ -116,7 +113,7 @@ impl PeerSessionTunnelFilter {
             enabled,
             my_peer_id: Arc::new(AtomicCell::new(PeerId::default())),
             peer_id: Arc::new(AtomicCell::new(None)),
-            session: Arc::new(hotpath::mutex!(std::sync::Mutex::new(None))),
+            session: Arc::new(ArcSwapOption::empty()),
         }
     }
 
@@ -125,7 +122,7 @@ impl PeerSessionTunnelFilter {
             enabled,
             my_peer_id: Arc::new(AtomicCell::new(my_peer_id)),
             peer_id: Arc::new(AtomicCell::new(None)),
-            session: Arc::new(hotpath::mutex!(std::sync::Mutex::new(None))),
+            session: Arc::new(ArcSwapOption::empty()),
         }
     }
 
@@ -138,7 +135,7 @@ impl PeerSessionTunnelFilter {
     }
 
     fn set_session(&self, session: Arc<PeerSession>) {
-        *self.session.lock().unwrap() = Some(session);
+        self.session.store(Some(session));
     }
 
     fn should_skip_encrypt(&self, hdr: &crate::tunnel::packet_def::PeerManagerHeader) -> bool {
@@ -172,16 +169,15 @@ impl TunnelFilter for PeerSessionTunnelFilter {
             return Some(data);
         };
 
-        let mut guard = self.session.lock().unwrap();
-        let Some(session) = guard.as_mut() else {
-            return Some(data);
-        };
-
         let my_peer_id = self.my_peer_id.load();
-        if my_peer_id != hdr.from_peer_id.get() {
+        if my_peer_id != hdr.from_peer_id.get() || hdr.to_peer_id.get() != peer_id {
             return Some(data);
         }
 
+        let session_guard = self.session.load();
+        let Some(session) = session_guard.as_deref() else {
+            return Some(data);
+        };
         if let Err(e) = session.encrypt_payload(my_peer_id, peer_id, &mut data) {
             tracing::warn!(
                 ?my_peer_id,
@@ -226,8 +222,8 @@ impl TunnelFilter for PeerSessionTunnelFilter {
             return Some(Ok(data));
         }
 
-        let mut guard = self.session.lock().unwrap();
+        let session_guard = self.session.load();
-        let Some(session) = guard.as_mut() else {
+        let Some(session) = session_guard.as_deref() else {
             return Some(Ok(data));
         };
 
@@ -1653,6 +1649,45 @@ pub mod tests {
             .unwrap_or(0)
     }
 
+    #[test]
+    fn peer_session_filter_skips_relay_packet_for_next_hop() {
+        let my_peer_id = 10;
+        let next_hop_peer_id = 20;
+        let dst_peer_id = 30;
+        let filter = PeerSessionTunnelFilter::new_with_peer(my_peer_id, true);
+        filter.set_peer_id(next_hop_peer_id);
+
+        let session = Arc::new(PeerSession::new(
+            next_hop_peer_id,
+            PeerSession::new_root_key(),
+            1,
+            0,
+            "aes-gcm".to_string(),
+            "aes-gcm".to_string(),
+            None,
+        ));
+        session.invalidate();
+        filter.set_session(session);
+
+        let mut packet = ZCPacket::new_with_payload(b"relay payload");
+        packet.fill_peer_manager_hdr(my_peer_id, dst_peer_id, PacketType::Data as u8);
+        packet
+            .mut_peer_manager_header()
+            .unwrap()
+            .set_encrypted(true);
+        let original_len = packet.buf_len();
+
+        let packet = filter
+            .before_send(packet)
+            .expect("relay packet should bypass next-hop session");
+
+        let hdr = packet.peer_manager_header().unwrap();
+        assert_eq!(hdr.from_peer_id.get(), my_peer_id);
+        assert_eq!(hdr.to_peer_id.get(), dst_peer_id);
+        assert!(hdr.is_encrypted());
+        assert_eq!(packet.buf_len(), original_len);
+    }
+
     #[tokio::test]
     async fn peer_conn_handshake_same_id() {
         let ps = Arc::new(PeerSessionStore::new());

easytier/src/peers/peer_session.rs

@@ -2,9 +2,12 @@ use std::sync::{
     Arc, RwLock,
     atomic::{AtomicBool, Ordering},
 };
+use std::time::Duration;
 
 use anyhow::anyhow;
+use crossbeam::atomic::AtomicCell;
 use dashmap::DashMap;
+use hotpath::instant::Instant;
 
 use super::secure_datagram::{SecureDatagramDirection, SecureDatagramSession};
 use crate::{
@@ -12,6 +15,8 @@ use crate::{
     tunnel::packet_def::ZCPacket,
 };
 
+const SESSION_IDLE_TIMEOUT: Duration = Duration::from_secs(60);
+
 pub struct UpsertResponderSessionReturn {
     pub session: Arc<PeerSession>,
     pub action: PeerSessionAction,
@@ -44,7 +49,25 @@ impl SessionKey {
 
 #[derive(Clone)]
 pub struct PeerSessionStore {
-    sessions: Arc<DashMap<SessionKey, Arc<PeerSession>>>,
+    sessions: Arc<DashMap<SessionKey, PeerSessionEntry>>,
+}
+
+struct PeerSessionEntry {
+    session: Arc<PeerSession>,
+    last_used_at: AtomicCell<Instant>,
+}
+
+impl PeerSessionEntry {
+    fn new(session: Arc<PeerSession>) -> Self {
+        Self {
+            session,
+            last_used_at: AtomicCell::new(Instant::now()),
+        }
+    }
+
+    fn touch(&self) {
+        self.last_used_at.store(Instant::now());
+    }
 }
 
 impl Default for PeerSessionStore {
@@ -61,7 +84,11 @@ impl PeerSessionStore {
     }
 
     pub fn get(&self, key: &SessionKey) -> Option<Arc<PeerSession>> {
-        let session = self.sessions.get(key)?.clone();
+        let session = {
+            let entry = self.sessions.get(key)?;
+            entry.touch();
+            entry.session.clone()
+        };
         if session.is_valid() {
             Some(session)
         } else {
@@ -75,12 +102,20 @@ impl PeerSessionStore {
     }
 
     pub fn insert_session(&self, key: SessionKey, session: Arc<PeerSession>) {
-        self.sessions.insert(key, session);
+        self.sessions.insert(key, PeerSessionEntry::new(session));
     }
 
     pub fn evict_unused_sessions(&self) {
-        self.sessions
+        self.evict_unused_sessions_idle(SESSION_IDLE_TIMEOUT);
-            .retain(|_key, session| Arc::strong_count(session) > 1);
+    }
+
+    pub fn evict_unused_sessions_idle(&self, idle: Duration) {
+        let now = Instant::now();
+        self.sessions.retain(|_key, entry| {
+            entry.session.is_valid()
+                && (Arc::strong_count(&entry.session) > 1
+                    || now.saturating_duration_since(entry.last_used_at.load()) < idle)
+        });
         shrink_dashmap(&self.sessions, None);
     }
 
@@ -93,11 +128,14 @@ impl PeerSessionStore {
         recv_algorithm: String,
         peer_static_pubkey: Option<[u8; 32]>,
     ) -> Result<UpsertResponderSessionReturn, anyhow::Error> {
-        tracing::event!(tracing::Level::INFO, "upsert_responder_session {:?}", key);
+        tracing::event!(tracing::Level::INFO, ?key, "upsert_responder_session");
         let existing = self
             .sessions
             .get(key)
-            .map(|v| v.clone())
+            .map(|v| {
+                v.touch();
+                v.session.clone()
+            })
             .filter(|s| s.is_valid());
         match existing {
             None => {
@@ -113,7 +151,8 @@ impl PeerSessionStore {
                     recv_algorithm,
                     peer_static_pubkey,
                 ));
-                self.sessions.insert(key.clone(), session.clone());
+                self.sessions
+                    .insert(key.clone(), PeerSessionEntry::new(session.clone()));
                 Ok(UpsertResponderSessionReturn {
                     session,
                     action: PeerSessionAction::Create,
@@ -178,16 +217,14 @@ impl PeerSessionStore {
             PeerSessionAction::Sync | PeerSessionAction::Create => {
                 let root_key = root_key_32.ok_or_else(|| anyhow!("missing root_key"))?;
                 if let Some(existing) = self.sessions.get(key)
-                    && !existing.is_valid()
+                    && !existing.session.is_valid()
                 {
                     drop(existing);
                     self.sessions.remove(key);
                 }
-                let session = self
+                let session = {
-                    .sessions
+                    let entry = self.sessions.entry(key.clone()).or_insert_with(|| {
-                    .entry(key.clone())
+                        PeerSessionEntry::new(Arc::new(PeerSession::new(
-                    .or_insert_with(|| {
-                        Arc::new(PeerSession::new(
                             key.peer_id,
                             root_key,
                             b_session_generation,
@@ -195,9 +232,11 @@ impl PeerSessionStore {
                             send_algorithm.clone(),
                             recv_algorithm.clone(),
                             peer_static_pubkey,
-                        ))
+                        )))
-                    })
+                    });
-                    .clone();
+                    entry.touch();
+                    entry.session.clone()
+                };
                 session.check_encrypt_algo_same(&send_algorithm, &recv_algorithm)?;
                 session.check_or_set_peer_static_pubkey(peer_static_pubkey)?;
                 session.sync_root_key(
@@ -421,4 +460,77 @@ mod tests {
             SecureDatagramSession::SYNC_RX_GRACE_AFTER_MS
         );
     }
+
+    #[test]
+    fn peer_session_store_keeps_recent_session_without_external_refs() {
+        let store = PeerSessionStore::new();
+        let key = SessionKey::new("net".to_string(), 20);
+        let session = Arc::new(PeerSession::new(
+            20,
+            PeerSession::new_root_key(),
+            1,
+            0,
+            "aes-gcm".to_string(),
+            "aes-gcm".to_string(),
+            None,
+        ));
+        store.insert_session(key.clone(), session);
+
+        assert!(store.get(&key).is_some());
+        store.evict_unused_sessions();
+
+        assert!(
+            store.get(&key).is_some(),
+            "recent relay sessions should survive the periodic GC"
+        );
+    }
+
+    #[test]
+    fn peer_session_store_evicts_idle_session_without_external_refs() {
+        let store = PeerSessionStore::new();
+        let key = SessionKey::new("net".to_string(), 20);
+        let session = Arc::new(PeerSession::new(
+            20,
+            PeerSession::new_root_key(),
+            1,
+            0,
+            "aes-gcm".to_string(),
+            "aes-gcm".to_string(),
+            None,
+        ));
+        store.insert_session(key.clone(), session);
+
+        store.evict_unused_sessions_idle(Duration::from_millis(0));
+
+        assert!(
+            store.get(&key).is_none(),
+            "idle sessions without external users should still be collected"
+        );
+    }
+
+    #[test]
+    fn peer_session_store_evicts_invalid_recent_session() {
+        let store = PeerSessionStore::new();
+        let key = SessionKey::new("net".to_string(), 20);
+        let session = Arc::new(PeerSession::new(
+            20,
+            PeerSession::new_root_key(),
+            1,
+            0,
+            "aes-gcm".to_string(),
+            "aes-gcm".to_string(),
+            None,
+        ));
+        store.insert_session(key.clone(), session);
+
+        let session = store.get(&key).unwrap();
+        session.invalidate();
+        drop(session);
+        store.evict_unused_sessions();
+
+        assert!(
+            !store.sessions.contains_key(&key),
+            "invalid sessions should not be kept by recent activity"
+        );
+    }
 }

easytier/src/tests/three_node.rs

@@ -4340,7 +4340,7 @@ pub async fn relay_peer_session_cleanup() {
     insts[0]
         .get_peer_manager()
         .get_peer_session_store()
-        .evict_unused_sessions();
+        .evict_unused_sessions_idle(Duration::from_millis(0));
 
     wait_for_condition(
         || async { !relay_map_1.has_session(inst3_peer_id) },