diff --git a/easytier/src/peers/peer_conn.rs b/easytier/src/peers/peer_conn.rs index 2e640547..789a6e1e 100644 --- a/easytier/src/peers/peer_conn.rs +++ b/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>, peer_id: Arc>>, - session: Arc>>>, + session: Arc>, } 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) { - *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 Some(session) = guard.as_mut() else { + let session_guard = self.session.load(); + 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()); diff --git a/easytier/src/peers/peer_session.rs b/easytier/src/peers/peer_session.rs index e15a1710..0a9ad175 100644 --- a/easytier/src/peers/peer_session.rs +++ b/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, pub action: PeerSessionAction, @@ -44,7 +49,25 @@ impl SessionKey { #[derive(Clone)] pub struct PeerSessionStore { - sessions: Arc>>, + sessions: Arc>, +} + +struct PeerSessionEntry { + session: Arc, + last_used_at: AtomicCell, +} + +impl PeerSessionEntry { + fn new(session: Arc) -> 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> { - 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) { - self.sessions.insert(key, session); + self.sessions.insert(key, PeerSessionEntry::new(session)); } pub fn evict_unused_sessions(&self) { - self.sessions - .retain(|_key, session| Arc::strong_count(session) > 1); + self.evict_unused_sessions_idle(SESSION_IDLE_TIMEOUT); + } + + 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 { - 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 - .sessions - .entry(key.clone()) - .or_insert_with(|| { - Arc::new(PeerSession::new( + let session = { + let entry = self.sessions.entry(key.clone()).or_insert_with(|| { + PeerSessionEntry::new(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" + ); + } } diff --git a/easytier/src/tests/three_node.rs b/easytier/src/tests/three_node.rs index 32f262ae..7bcea386 100644 --- a/easytier/src/tests/three_node.rs +++ b/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) },