diff --git a/crates/tsunagi-wireguard/src/device.rs b/crates/tsunagi-wireguard/src/device.rs index 7f5157f..2c4c12f 100644 --- a/crates/tsunagi-wireguard/src/device.rs +++ b/crates/tsunagi-wireguard/src/device.rs @@ -14,20 +14,28 @@ use bytes::Bytes; use iroh::EndpointId; use tokio::task::JoinHandle; -use tsunagi::dataplane::PacketSink; -use tsunagi::dataplane::PluginError; use tsunagi::dataplane::routing::{FlowId, flow::ip_flow}; use tsunagi::dataplane::transport::{SharedLink, TransportError}; +use tsunagi::dataplane::{PacketSink, PluginContext, PluginError}; use tsunagi::identity::NetworkId; use crate::keys::{WgPublicKey, WgSecretKey}; +use crate::transport::WireguardTransport; const TIMER_INTERVAL: Duration = Duration::from_millis(250); const SCRATCH: usize = 4096; +/// Maximum consecutive rekey timeout retransmissions before declaring a tunnel dead. +pub const MAX_REKEY_TIMEOUTS: u32 = 3; + +/// Overall maximum duration trying to establish or rekey a handshake before declaring it dead. +pub const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(15); + struct FlowTunnel { tunn: Tunn, queued: VecDeque, + handshake_started_at: Option, + rekey_timeouts: u32, } fn is_data(result: &TunnResult<'_>) -> bool { @@ -45,6 +53,12 @@ impl FlowTunnel { if is_data(&result) { return (result, flow); } + if matches!(result, TunnResult::WriteToNetwork(ref b) if !b.is_empty() && b[0] == 1) + && self.handshake_started_at.is_none() + { + self.handshake_started_at = Some(std::time::Instant::now()); + self.rekey_timeouts = 0; + } if self.queued.len() < 256 { self.queued.push_back(flow); } @@ -57,6 +71,10 @@ impl FlowTunnel { scratch: &'a mut [u8], ) -> (TunnResult<'a>, FlowId) { let result = self.tunn.decapsulate(None, packet, scratch); + if self.tunn.time_since_last_handshake().is_some() { + self.handshake_started_at = None; + self.rekey_timeouts = 0; + } let mut flow = 0; if packet.is_empty() { if is_data(&result) { @@ -68,15 +86,37 @@ impl FlowTunnel { (result, flow) } - fn update_timers<'a>(&mut self, scratch: &'a mut [u8]) -> TunnResult<'a> { + fn update_timers<'a>(&mut self, scratch: &'a mut [u8]) -> (TunnResult<'a>, bool) { let result = self.tunn.update_timers(scratch); if matches!( result, TunnResult::Err(boringtun::noise::errors::WireGuardError::ConnectionExpired) - ) { + ) || self.tunn.is_expired() + { self.queued.clear(); + self.handshake_started_at = None; + self.rekey_timeouts = 0; + return (result, true); } - result + + if let TunnResult::WriteToNetwork(ref bytes) = result + && !bytes.is_empty() + && bytes[0] == 1 + && self.handshake_started_at.is_none() + { + self.handshake_started_at = Some(std::time::Instant::now()); + } + + if let Some(started) = self.handshake_started_at + && started.elapsed() >= HANDSHAKE_TIMEOUT + { + self.queued.clear(); + self.handshake_started_at = None; + self.rekey_timeouts = 0; + return (result, true); + } + + (result, false) } } @@ -205,9 +245,23 @@ struct Inner { private_key: WgSecretKey, peers: RwLock>>, sink: Arc, + transport: Arc, + context: PluginContext, next_index: AtomicU32, } +impl Inner { + fn on_tunnel_failed(&self, peer: EndpointId) { + self.transport.close_link(self.network, peer); + self.context.report_peer_protocol_failure( + self.network, + peer, + crate::announcement::WG_PROTOCOL, + "WireGuard rekey timeout", + ); + } +} + impl std::fmt::Debug for Inner { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("Inner") @@ -225,12 +279,20 @@ pub struct WireguardDevice { impl WireguardDevice { /// Starts the WireGuard tunnels for one network. - pub fn start(network: NetworkId, private_key: WgSecretKey, sink: Arc) -> Self { + pub fn start( + network: NetworkId, + private_key: WgSecretKey, + sink: Arc, + transport: Arc, + context: PluginContext, + ) -> Self { let inner = Arc::new(Inner { network, private_key, peers: RwLock::new(HashMap::new()), sink, + transport, + context, next_index: AtomicU32::new(1), }); @@ -312,6 +374,8 @@ impl WireguardDevice { tunn: Mutex::new(FlowTunnel { tunn, queued: VecDeque::new(), + handshake_started_at: None, + rekey_timeouts: 0, }), link, counters: Arc::new(PeerCounters::default()), @@ -405,7 +469,11 @@ fn kick_handshake(peer: &Peer) { Err(poisoned) => poisoned.into_inner(), }; match tunn.tunn.format_handshake_initiation(&mut scratch, false) { - TunnResult::WriteToNetwork(out) => Some(out.len()), + TunnResult::WriteToNetwork(out) => { + tunn.handshake_started_at = Some(std::time::Instant::now()); + tunn.rekey_timeouts = 0; + Some(out.len()) + } _ => None, } }; @@ -505,23 +573,92 @@ async fn drive_timers(inner: Arc) { let peers: Vec> = read_lock(&inner.peers).values().cloned().collect(); for peer in peers { let mut scratch = vec![0u8; SCRATCH]; - let len = { + let (len, failed) = { let mut tunn = match peer.tunn.lock() { Ok(guard) => guard, Err(poisoned) => poisoned.into_inner(), }; - match tunn.update_timers(&mut scratch) { + let (result, failed) = tunn.update_timers(&mut scratch); + let len = match result { TunnResult::WriteToNetwork(out) => Some(out.len()), TunnResult::Err(err) => { tracing::trace!(?err, "WireGuard timer error"); None } _ => None, - } + }; + (len, failed) }; if let Some(len) = len { send_to_peer(&peer, &scratch[..len]); } + if failed { + tracing::warn!( + peer = %peer.endpoint_id.fmt_short(), + "WireGuard tunnel failed due to rekey timeout, triggering fallback" + ); + inner.on_tunnel_failed(peer.endpoint_id); + } } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_flow_tunnel_rekey_timeout_triggers_failure() { + let my_secret = WgSecretKey::generate(); + let their_secret = WgSecretKey::generate(); + let their_public = their_secret.public(); + let tunn = Tunn::new( + my_secret.to_static_secret(), + their_public.into_x25519(), + None, + None, + 1, + None, + ); + + let mut tunnel = FlowTunnel { + tunn, + queued: VecDeque::new(), + handshake_started_at: Some(std::time::Instant::now() - Duration::from_secs(16)), + rekey_timeouts: 0, + }; + + let mut scratch = vec![0u8; SCRATCH]; + // When handshake timeout exceeds HANDSHAKE_TIMEOUT (15s), update_timers returns true (failed) + let (_, failed) = tunnel.update_timers(&mut scratch); + assert!(failed, "handshake timeout should report failure"); + assert!(tunnel.handshake_started_at.is_none()); + } + + #[test] + fn test_flow_tunnel_reset_on_handshake() { + let my_secret = WgSecretKey::generate(); + let their_secret = WgSecretKey::generate(); + let their_public = their_secret.public(); + let tunn = Tunn::new( + my_secret.to_static_secret(), + their_public.into_x25519(), + None, + None, + 1, + None, + ); + + let mut tunnel = FlowTunnel { + tunn, + queued: VecDeque::new(), + handshake_started_at: Some(std::time::Instant::now()), + rekey_timeouts: 2, + }; + + let mut scratch = vec![0u8; SCRATCH]; + let _ = tunnel.decapsulate(&[], &mut scratch); + // decapsulate without handshake response doesn't establish session + assert_eq!(tunnel.rekey_timeouts, 2); + } +} diff --git a/crates/tsunagi-wireguard/src/plugin.rs b/crates/tsunagi-wireguard/src/plugin.rs index ccc20df..91b54be 100644 --- a/crates/tsunagi-wireguard/src/plugin.rs +++ b/crates/tsunagi-wireguard/src/plugin.rs @@ -576,7 +576,18 @@ impl Worker { Some(sink) => Arc::clone(sink), None => Arc::new(tsunagi::dataplane::DiscardPackets) as Arc, }; - let device = Arc::new(WireguardDevice::start(network, key, sink)); + let context = self + .context + .get() + .cloned() + .unwrap_or_else(PluginContext::detached); + let device = Arc::new(WireguardDevice::start( + network, + key, + sink, + Arc::clone(&self.transport), + context, + )); let mut shared = self.lock_shared(); if let Some(state) = shared.networks.get_mut(&network) { diff --git a/crates/tsunagi-wireguard/src/transport.rs b/crates/tsunagi-wireguard/src/transport.rs index e49ee02..7413c37 100644 --- a/crates/tsunagi-wireguard/src/transport.rs +++ b/crates/tsunagi-wireguard/src/transport.rs @@ -222,6 +222,13 @@ impl WireguardLink { } } + /// Explicitly closes this link, notifying any pending readers or waiters. + pub fn close(&self) { + if !self.is_closed.swap(true, Ordering::AcqRel) { + self.closed_notify.notify_waiters(); + } + } + /// Returns the current target address. pub fn target_addr(&self) -> SocketAddr { match self.target_addr.read() { @@ -261,8 +268,17 @@ impl PacketLink for WireguardLink { fn recv(&self) -> BoxFuture<'_, Option> { Box::pin(async move { - let mut guard = self.rx.lock().await; - guard.recv().await + if self.is_closed.load(Ordering::Acquire) { + return None; + } + tokio::select! { + biased; + _ = self.closed_notify.notified() => None, + msg = async { + let mut guard = self.rx.lock().await; + guard.recv().await + } => msg, + } }) } @@ -388,6 +404,20 @@ impl WireguardTransport { guard.get(&peer).cloned().unwrap_or_default() } + /// Explicitly closes any active direct link for the given network and peer. + pub fn close_link(&self, network: NetworkId, peer: EndpointId) { + let link = { + let guard = match self.links.read() { + Ok(g) => g, + Err(p) => p.into_inner(), + }; + guard.get(&(network, peer)).cloned() + }; + if let Some(link) = link { + link.close(); + } + } + /// Notifies the transport that a peer's capability was announced. /// /// When `we_accept` is true, sends a hole-punch probe to candidate addresses diff --git a/crates/tsunagi-wireguard/tests/transport_tests.rs b/crates/tsunagi-wireguard/tests/transport_tests.rs index 0d1c1fc..d1cf929 100644 --- a/crates/tsunagi-wireguard/tests/transport_tests.rs +++ b/crates/tsunagi-wireguard/tests/transport_tests.rs @@ -144,3 +144,77 @@ async fn test_wireguard_transport_unreachable_probe_times_out() { other => panic!("expected Unreachable error, got {other:?}"), } } + +#[tokio::test] +async fn test_wireguard_link_close_and_transport_close_link() { + let identity_a = DeviceIdentity::generate(); + let peer_a = identity_a.endpoint_id(); + let identity_b = DeviceIdentity::generate(); + let peer_b = identity_b.endpoint_id(); + let network = NetworkId::from_bytes([55u8; 32]); + + let transport_b = Arc::new( + WireguardTransport::bind(&identity_b, Some(0)) + .await + .unwrap(), + ); + let port_b = transport_b.bound_port(); + + let transport_a = Arc::new( + WireguardTransport::bind(&identity_a, Some(0)) + .await + .unwrap(), + ); + let port_a = transport_a.bound_port(); + + let addr_b: SocketAddr = format!("127.0.0.1:{port_b}").parse().unwrap(); + transport_a.set_peer_addr(peer_b, addr_b); + transport_a.set_peer_network(peer_b, network); + + let addr_a: SocketAddr = format!("127.0.0.1:{port_a}").parse().unwrap(); + transport_b.set_peer_addr(peer_a, addr_a); + transport_b.set_peer_network(peer_a, network); + + let transport_a_clone = Arc::clone(&transport_a); + tokio::spawn(async move { + transport_a_clone.accept_loop(|_| {}).await; + }); + + let transport_b_clone = Arc::clone(&transport_b); + tokio::spawn(async move { + transport_b_clone.accept_loop(|_| {}).await; + }); + + let link_a = transport_a + .open(network, peer_b, "wg") + .await + .expect("open link A"); + assert!(!link_a.is_closed()); + + // Spawn a reader task waiting on recv() + let link_a_clone = Arc::clone(&link_a); + let recv_task = tokio::spawn(async move { link_a_clone.recv().await }); + + // Explicitly closing via transport.close_link() + transport_a.close_link(network, peer_b); + assert!(link_a.is_closed()); + + // Wait on closed() notification + tokio::time::timeout(Duration::from_secs(1), link_a.closed()) + .await + .expect("link.closed() should resolve promptly"); + + // recv() must have returned None + let recv_res = tokio::time::timeout(Duration::from_secs(1), recv_task) + .await + .expect("recv task should complete") + .expect("task join"); + assert_eq!(recv_res, None); + + // Send after close must return Closed error + let send_res = link_a.send(Bytes::from_static(b"after close")); + assert!(matches!( + send_res, + Err(tsunagi::dataplane::transport::TransportError::Closed) + )); +} diff --git a/crates/tsunagi/src/agent/mod.rs b/crates/tsunagi/src/agent/mod.rs index c1fe6bb..ad300ca 100644 --- a/crates/tsunagi/src/agent/mod.rs +++ b/crates/tsunagi/src/agent/mod.rs @@ -1118,6 +1118,26 @@ async fn plugin_request_loop(weak: Weak, mut requests: mpsc::Receiver { + let sender = { + let networks = inner.networks.read().await; + networks.get(&network).map(|handle| handle.commands.clone()) + }; + if let Some(sender) = sender { + let _ = sender + .send(NetCommand::ProtocolFailure { + peer, + protocol, + reason, + }) + .await; + } + } } } } diff --git a/crates/tsunagi/src/agent/network.rs b/crates/tsunagi/src/agent/network.rs index 76c7298..d8304bb 100644 --- a/crates/tsunagi/src/agent/network.rs +++ b/crates/tsunagi/src/agent/network.rs @@ -86,6 +86,12 @@ pub(crate) enum NetCommand { /// Human readable reason, free of secrets. reason: String, }, + /// An IP plugin reported that data plane connectivity or a tunnel for a peer failed. + ProtocolFailure { + peer: EndpointId, + protocol: String, + reason: String, + }, } impl std::fmt::Debug for NetCommand { @@ -104,6 +110,9 @@ impl std::fmt::Debug for NetCommand { NetCommand::SetHostname(_) => f.write_str("SetHostname"), NetCommand::Release { .. } => f.write_str("Release"), NetCommand::PluginError { protocol, .. } => write!(f, "PluginError({protocol})"), + NetCommand::ProtocolFailure { peer, protocol, .. } => { + write!(f, "ProtocolFailure({}, {protocol})", peer.fmt_short()) + } } } } @@ -522,6 +531,26 @@ impl Runtime { reason, }); } + NetCommand::ProtocolFailure { + peer, + protocol, + reason, + } => { + if self.links.remove(&(peer, protocol.clone())).is_some() { + self.hub.clear_direct(peer, &protocol); + } + self.metrics.data_link_failures += 1; + self.record_protocol_failure(peer, &protocol); + self.emit(Event::DataLinkDown { + network: self.network_id, + peer, + protocol, + reason, + }); + self.ensure_links(); + self.update_paths(); + self.announce_reach(false); + } } } @@ -2162,4 +2191,70 @@ mod tests { assert_eq!(data_link_backoff(6), Duration::from_secs(60)); assert_eq!(data_link_backoff(10), Duration::from_secs(60)); } + + #[test] + fn test_protocol_dial_state_cooldown_and_fallback_selection() { + let peer = crate::identity::DeviceIdentity::generate().endpoint_id(); + let supported = ["wg", "tcp-tls", "wg-quic"]; + let mut states: HashMap<(EndpointId, String), ProtocolDialState> = HashMap::new(); + + let now = std::time::Instant::now(); + + // 1. Without cooldown, the highest priority protocol "wg" is chosen + let chosen = + supported + .iter() + .find(|proto| match states.get(&(peer, (*proto).to_string())) { + Some(state) => state.cooldown_until <= now, + None => true, + }); + assert_eq!(chosen.copied(), Some("wg")); + + // 2. Record failure for "wg" putting it on cooldown + states.insert( + (peer, "wg".to_string()), + ProtocolDialState { + failures: 1, + cooldown_until: now + Duration::from_secs(3), + }, + ); + + // With "wg" on cooldown, selection must fall back to next priority "tcp-tls" + let chosen_after_failure = + supported + .iter() + .find(|proto| match states.get(&(peer, (*proto).to_string())) { + Some(state) => state.cooldown_until <= now, + None => true, + }); + assert_eq!(chosen_after_failure.copied(), Some("tcp-tls")); + + // 3. If "tcp-tls" also fails, selection falls back to "wg-quic" + states.insert( + (peer, "tcp-tls".to_string()), + ProtocolDialState { + failures: 1, + cooldown_until: now + Duration::from_secs(3), + }, + ); + let chosen_quic = + supported + .iter() + .find(|proto| match states.get(&(peer, (*proto).to_string())) { + Some(state) => state.cooldown_until <= now, + None => true, + }); + assert_eq!(chosen_quic.copied(), Some("wg-quic")); + + // 4. Once success is recorded (entry removed), "wg" is preferred again + states.remove(&(peer, "wg".to_string())); + let chosen_restored = + supported + .iter() + .find(|proto| match states.get(&(peer, (*proto).to_string())) { + Some(state) => state.cooldown_until <= now, + None => true, + }); + assert_eq!(chosen_restored.copied(), Some("wg")); + } } diff --git a/crates/tsunagi/src/dataplane/mod.rs b/crates/tsunagi/src/dataplane/mod.rs index ba3adf2..7ed05ce 100644 --- a/crates/tsunagi/src/dataplane/mod.rs +++ b/crates/tsunagi/src/dataplane/mod.rs @@ -119,7 +119,7 @@ pub enum PluginError { /// A request a plugin makes of the agent that owns it. #[derive(Debug)] -pub(crate) enum PluginRequest { +pub enum PluginRequest { /// Re-send this agent's announcement to every peer of a network. Reannounce(NetworkId), /// Surface a plugin error on the agent's event stream. @@ -131,6 +131,17 @@ pub(crate) enum PluginRequest { /// Human readable reason, free of secrets. reason: String, }, + /// A plugin reported that a data link or tunnel for a peer failed. + ProtocolFailure { + /// Network the failure occurred in. + network: NetworkId, + /// Peer whose protocol failed. + peer: EndpointId, + /// Protocol that failed. + protocol: String, + /// Human readable reason. + reason: String, + }, } /// The agent-side handle a plugin is given when it is attached. @@ -192,6 +203,20 @@ impl PluginContext { } } + /// A context configured with a channel for plugin requests and custom sink, for tests. + #[cfg(feature = "testing")] + pub fn for_test( + sender: mpsc::Sender, + local: EndpointId, + sink: Option>, + ) -> Self { + Self { + sender: Some(sender), + local: Some(local), + sink, + } + } + /// This agent's own endpoint id, when the context is attached. /// /// A plugin needs it to find itself in the agreed allocation. @@ -233,6 +258,22 @@ impl PluginContext { reason: reason.into(), }); } + + /// Reports that data plane connectivity or a tunnel for a peer has failed. + pub fn report_peer_protocol_failure( + &self, + network: NetworkId, + peer: EndpointId, + protocol: impl Into, + reason: impl Into, + ) { + self.send(PluginRequest::ProtocolFailure { + network, + peer, + protocol: protocol.into(), + reason: reason.into(), + }); + } } impl std::fmt::Debug for PluginContext {