From 636390ec38be250c04738a96de442d87ea889181 Mon Sep 17 00:00:00 2001 From: KKRainbow <443152178@qq.com> Date: Sat, 15 Aug 2026 00:21:07 +0800 Subject: [PATCH] feat(peer): echo liveness probes on data traffic (#2497) * feat(peer): echo liveness probes on data traffic Advertise a liveness-echo capability during classic and Noise handshakes. After a ping failure, tag outgoing peer packets with a short probe token and accept only the matching echoed token as round-trip proof. Keep one ping request outstanding and coalesce scheduler triggers so high traffic cannot reorder timeout results. Preserve one-way failure detection because unrelated ingress never clears the loss counter. * test(three_node): relax disconnect wait for sequential pingpong proxy_three_node_disconnect_test assumed the old pingpong timing, where overlapping pings failed fast and the connection closed well inside the 11s wait (see the old [4, 9)s comment). The liveness-echo change keeps one ping outstanding: each failure now takes a full 2s timeout, so the fifth consecutive failure and the connection close land at ~11s. Both proto variants timed out at the 11s bound in CI. Widen the wait to 15s and update the timing comment. --- easytier-core/src/packet/mod.rs | 42 ++- easytier-core/src/peers/conn/mod.rs | 1 + easytier-core/src/peers/conn/peer_conn.rs | 23 +- .../src/peers/conn/peer_conn_liveness.rs | 310 ++++++++++++++++++ .../src/peers/conn/peer_conn_ping.rs | 229 +++++++------ easytier-core/src/peers/test_support.rs | 6 + easytier-core/src/peers/tests.rs | 49 +++ easytier-proto/proto/peer_rpc.proto | 2 + easytier/src/tests/three_node.rs | 6 +- 9 files changed, 560 insertions(+), 108 deletions(-) create mode 100644 easytier-core/src/peers/conn/peer_conn_liveness.rs diff --git a/easytier-core/src/packet/mod.rs b/easytier-core/src/packet/mod.rs index 87340314..705ada86 100644 --- a/easytier-core/src/packet/mod.rs +++ b/easytier-core/src/packet/mod.rs @@ -114,9 +114,9 @@ bitflags::bitflags! { const NO_PROXY = 0b0000_1000; const COMPRESSED = 0b0001_0000; // deprecated flags, can be reused. - // const KCP_SRC_MODIFIED = 0b0010_0000; - // const QUIC_SRC_MODIFIED = 0b1000_0000; + const LIVENESS_PROBE = 0b0010_0000; const NOT_SEND_TO_TUN = 0b0100_0000; + const LIVENESS_ECHO = 0b1000_0000; const _ = !0; } @@ -220,6 +220,44 @@ impl PeerManagerHeader { self } + pub(crate) fn liveness_probe_token(&self) -> Option { + PeerManagerHeaderFlags::from_bits(self.flags) + .unwrap() + .contains(PeerManagerHeaderFlags::LIVENESS_PROBE) + .then_some(self.reserved) + } + + pub(crate) fn liveness_echo_token(&self) -> Option { + PeerManagerHeaderFlags::from_bits(self.flags) + .unwrap() + .contains(PeerManagerHeaderFlags::LIVENESS_ECHO) + .then_some(self.reserved) + } + + pub(crate) fn set_liveness_probe(&mut self, token: u8) { + self.set_liveness_marker(PeerManagerHeaderFlags::LIVENESS_PROBE, token); + } + + pub(crate) fn set_liveness_echo(&mut self, token: u8) { + self.set_liveness_marker(PeerManagerHeaderFlags::LIVENESS_ECHO, token); + } + + pub(crate) fn clear_liveness_marker(&mut self) { + let mut flags = PeerManagerHeaderFlags::from_bits(self.flags).unwrap(); + flags + .remove(PeerManagerHeaderFlags::LIVENESS_PROBE | PeerManagerHeaderFlags::LIVENESS_ECHO); + self.flags = flags.bits(); + self.reserved = 0; + } + + fn set_liveness_marker(&mut self, marker: PeerManagerHeaderFlags, token: u8) { + self.clear_liveness_marker(); + let mut flags = PeerManagerHeaderFlags::from_bits(self.flags).unwrap(); + flags.insert(marker); + self.flags = flags.bits(); + self.reserved = token; + } + pub fn mark_kcp_src_modified(&mut self) -> &mut Self { assert_eq!(self.packet_type, PacketType::Data as u8); self.packet_type = PacketType::DataWithKcpSrcModified as u8; diff --git a/easytier-core/src/peers/conn/mod.rs b/easytier-core/src/peers/conn/mod.rs index 5a407ba9..ee202d40 100644 --- a/easytier-core/src/peers/conn/mod.rs +++ b/easytier-core/src/peers/conn/mod.rs @@ -3,6 +3,7 @@ pub(crate) mod peer; pub(crate) mod peer_conn; +pub(crate) mod peer_conn_liveness; pub(crate) mod peer_conn_ping; pub(crate) mod peer_map; pub(crate) mod peer_session; diff --git a/easytier-core/src/peers/conn/peer_conn.rs b/easytier-core/src/peers/conn/peer_conn.rs index 15f3f8fb..c226b4a2 100644 --- a/easytier-core/src/peers/conn/peer_conn.rs +++ b/easytier-core/src/peers/conn/peer_conn.rs @@ -29,6 +29,7 @@ use snow::{HandshakeState, params::NoiseParams}; use crate::foundation::time::{Duration, timeout}; use super::{ + peer_conn_liveness::{FEATURE as LIVENESS_ECHO_FEATURE, PeerConnLiveness}, peer_conn_ping::PeerConnPinger, peer_session::{PeerSession, PeerSessionAction}, }; @@ -88,6 +89,7 @@ struct NoiseHandshakeResult { // foreign network manager use this to verify peer. // the challenge will be sent to authorized peer and compare the proof against it. client_secret_proof: Option, + remote_features: Vec, } #[derive(Clone)] @@ -300,6 +302,7 @@ pub struct PeerConn { latency_stats: Arc, throughput: Arc, loss_rate_stats: Arc, + liveness: PeerConnLiveness, peer_session_store: Arc, my_encrypt_algo: String, @@ -347,7 +350,9 @@ impl PeerConn { let peer_conn_tunnel_filter = StatsRecorderTunnelFilter::new(); let throughput = peer_conn_tunnel_filter.filter_output(); - let filter_chain = TunnelFilterChain::new(session_filter.clone(), peer_conn_tunnel_filter); + let liveness = PeerConnLiveness::new(); + let filter_chain = TunnelFilterChain::new(session_filter.clone(), peer_conn_tunnel_filter) + .chain(liveness.clone()); let peer_conn_tunnel = TunnelWithFilter::new(tunnel, filter_chain); let mut mpsc_tunnel = MpscTunnel::new(peer_conn_tunnel, Some(Duration::from_secs(7))); @@ -389,6 +394,7 @@ impl PeerConn { latency_stats: Arc::new(WindowLatency::new(15)), throughput, loss_rate_stats: Arc::new(AtomicU32::new(0)), + liveness, peer_session_store, my_encrypt_algo, @@ -515,7 +521,7 @@ impl PeerConn { magic: MAGIC, my_peer_id: self.my_peer_id, version: VERSION, - features: Vec::new(), + features: vec![LIVENESS_ECHO_FEATURE.to_owned()], network_name: network.network_name.clone(), ..Default::default() }; @@ -796,6 +802,7 @@ impl PeerConn { a_session_generation, a_conn_id: Some(a_conn_id.into()), client_encryption_algorithm: self.my_encrypt_algo.clone(), + features: vec![LIVENESS_ECHO_FEATURE.to_owned()], }; let mut hs = builder @@ -941,6 +948,7 @@ impl PeerConn { // we have authorized the peer with noise handshake, so just set secret digest same as us even remote is a shared node. secret_digest, client_secret_proof: None, + remote_features: msg2_pb.features, }) } @@ -1067,6 +1075,7 @@ impl PeerConn { a_conn_id_echo: msg1_pb.a_conn_id, secret_proof_32, server_encryption_algorithm: algo, + features: vec![LIVENESS_ECHO_FEATURE.to_owned()], }; self.send_noise_msg( msg2_pb, @@ -1151,6 +1160,7 @@ impl PeerConn { challenge: handshake_hash_for_proof, proof: p.clone(), }), + remote_features: msg1_pb.features, }) } @@ -1162,7 +1172,7 @@ impl PeerConn { version: VERSION, network_name: noise.remote_network_name.clone(), - features: Vec::new(), + features: noise.remote_features.clone(), network_secret_digest: noise.secret_digest.clone(), } } @@ -1217,6 +1227,9 @@ impl PeerConn { ))); } + self.liveness + .set_remote_features(&self.info.as_ref().unwrap().features); + if self.get_peer_id() == self.my_peer_id { Err(Error::WaitRespError("peer id conflict".to_owned())) } else { @@ -1245,6 +1258,9 @@ impl PeerConn { self.is_client = Some(true); } + self.liveness + .set_remote_features(&self.info.as_ref().unwrap().features); + if self.get_peer_id() == self.my_peer_id { Err(Error::WaitRespError( "peer id conflict, are you connecting to yourself?".to_owned(), @@ -1348,6 +1364,7 @@ impl PeerConn { self.throughput.clone(), self.context.clone(), self.get_conn_info().network_name, + self.liveness.clone(), ); let close_event_notifier = self.close_event_notifier.clone(); diff --git a/easytier-core/src/peers/conn/peer_conn_liveness.rs b/easytier-core/src/peers/conn/peer_conn_liveness.rs new file mode 100644 index 00000000..939cbe6e --- /dev/null +++ b/easytier-core/src/peers/conn/peer_conn_liveness.rs @@ -0,0 +1,310 @@ +use std::sync::{ + Arc, + atomic::{AtomicBool, AtomicU8, AtomicU16, Ordering}, +}; + +use tokio::sync::Notify; + +use crate::tunnel::{SinkItem, StreamItem, filter::TunnelFilter}; + +pub(super) const FEATURE: &str = "liveness-echo-v1"; + +#[derive(Clone, Default)] +pub(super) struct PeerConnLiveness { + inner: Arc, +} + +#[derive(Default)] +struct PeerConnLivenessInner { + enabled: AtomicBool, + next_token: AtomicU8, + active_probe: AtomicU16, + acknowledged_probe: AtomicU16, + pending_echo: AtomicU16, + echo_received: Notify, +} + +impl PeerConnLiveness { + pub(super) fn new() -> Self { + Self::default() + } + + pub(super) fn set_enabled(&self, enabled: bool) { + self.inner.enabled.store(enabled, Ordering::Release); + if !enabled { + self.inner.active_probe.store(0, Ordering::Release); + self.inner.acknowledged_probe.store(0, Ordering::Release); + self.inner.pending_echo.store(0, Ordering::Release); + } + } + + pub(super) fn set_remote_features(&self, features: &[String]) { + self.set_enabled(features.iter().any(|feature| feature == FEATURE)); + } + + pub(super) fn start_probe(&self) -> Option { + if !self.inner.enabled.load(Ordering::Acquire) { + return None; + } + + let active = self.inner.active_probe.load(Ordering::Acquire); + if active != 0 { + return Some(Self::decode_probe(active)); + } + + let token = self + .inner + .next_token + .fetch_add(1, Ordering::Relaxed) + .wrapping_add(1); + self.inner.acknowledged_probe.store(0, Ordering::Release); + self.inner + .active_probe + .store(Self::encode_probe(token), Ordering::Release); + Some(token) + } + + pub(super) async fn wait_for_echo(&self, token: u8) { + let expected = Self::encode_probe(token); + loop { + let notified = self.inner.echo_received.notified(); + if self.inner.acknowledged_probe.load(Ordering::Acquire) == expected { + return; + } + notified.await; + } + } + + pub(super) fn finish_probe(&self, token: u8) { + let probe = Self::encode_probe(token); + let _ = + self.inner + .active_probe + .compare_exchange(probe, 0, Ordering::AcqRel, Ordering::Acquire); + let _ = self.inner.acknowledged_probe.compare_exchange( + probe, + 0, + Ordering::AcqRel, + Ordering::Acquire, + ); + } + + fn encode_probe(token: u8) -> u16 { + u16::from(token) + 1 + } + + fn decode_probe(probe: u16) -> u8 { + (probe - 1) as u8 + } + + fn active_probe(&self) -> Option { + let probe = self.inner.active_probe.load(Ordering::Acquire); + (probe != 0).then(|| Self::decode_probe(probe)) + } + + fn queue_echo(&self, token: u8) { + self.inner + .pending_echo + .store(Self::encode_probe(token), Ordering::Release); + } + + fn take_echo(&self) -> Option { + let echo = self.inner.pending_echo.swap(0, Ordering::AcqRel); + (echo != 0).then(|| Self::decode_probe(echo)) + } + + fn acknowledge(&self, token: u8) { + let probe = Self::encode_probe(token); + if self + .inner + .active_probe + .compare_exchange(probe, 0, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + { + self.inner + .acknowledged_probe + .store(probe, Ordering::Release); + self.inner.echo_received.notify_one(); + } + } +} + +impl TunnelFilter for PeerConnLiveness { + type FilterOutput = (); + + fn before_send(&self, mut data: SinkItem) -> Option { + if !self.inner.enabled.load(Ordering::Acquire) { + return Some(data); + } + + let Some(header) = data.mut_peer_manager_header() else { + return Some(data); + }; + header.clear_liveness_marker(); + if let Some(token) = self.take_echo() { + header.set_liveness_echo(token); + } else if let Some(token) = self.active_probe() { + header.set_liveness_probe(token); + } + Some(data) + } + + fn after_received(&self, data: StreamItem) -> Option { + let Ok(mut packet) = data else { + return Some(data); + }; + if !self.inner.enabled.load(Ordering::Acquire) { + return Some(Ok(packet)); + } + + let Some(header) = packet.mut_peer_manager_header() else { + return Some(Ok(packet)); + }; + let probe = header.liveness_probe_token(); + let echo = header.liveness_echo_token(); + header.clear_liveness_marker(); + + match (probe, echo) { + (Some(token), None) => self.queue_echo(token), + (None, Some(token)) => self.acknowledge(token), + _ => {} + } + Some(Ok(packet)) + } + + fn filter_output(&self) {} +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use crate::{ + packet::{PacketType, ZCPacket}, + tunnel::filter::TunnelFilter, + }; + + use super::{FEATURE, PeerConnLiveness}; + + fn data_packet(from: u32, to: u32) -> ZCPacket { + let mut packet = ZCPacket::new_with_payload(b"payload"); + packet.fill_peer_manager_hdr(from, to, PacketType::Data as u8); + packet + } + + #[tokio::test] + async fn echoed_business_packet_proves_round_trip_liveness() { + let local = PeerConnLiveness::new(); + let remote = PeerConnLiveness::new(); + local.set_enabled(true); + remote.set_enabled(true); + + let token = local.start_probe().expect("feature enabled"); + let outbound = local.before_send(data_packet(1, 2)).unwrap(); + let received = remote.after_received(Ok(outbound)).unwrap().unwrap(); + assert_eq!(received.payload(), b"payload"); + let header = received.peer_manager_header().unwrap(); + assert_eq!(header.liveness_probe_token(), None); + assert_eq!(header.liveness_echo_token(), None); + + let reply = remote.before_send(data_packet(2, 1)).unwrap(); + let received = local.after_received(Ok(reply)).unwrap().unwrap(); + assert_eq!(received.payload(), b"payload"); + let header = received.peer_manager_header().unwrap(); + assert_eq!(header.liveness_probe_token(), None); + assert_eq!(header.liveness_echo_token(), None); + + tokio::time::timeout(Duration::from_millis(50), local.wait_for_echo(token)) + .await + .expect("matching echo was not observed"); + } + + #[test] + fn old_peer_without_feature_keeps_liveness_markers_disabled() { + let liveness = PeerConnLiveness::new(); + liveness.set_remote_features(&[]); + assert_eq!(liveness.start_probe(), None); + + liveness.set_remote_features(&[FEATURE.to_owned()]); + assert!(liveness.start_probe().is_some()); + } + + #[tokio::test] + async fn unrelated_ingress_does_not_acknowledge_an_active_probe() { + let liveness = PeerConnLiveness::new(); + liveness.set_enabled(true); + let token = liveness.start_probe().unwrap(); + + for _ in 0..8 { + liveness + .after_received(Ok(data_packet(2, 1))) + .unwrap() + .unwrap(); + } + + assert!( + tokio::time::timeout(Duration::from_millis(20), liveness.wait_for_echo(token)) + .await + .is_err(), + "unrelated one-way ingress acknowledged the local probe" + ); + } + + #[tokio::test] + async fn stale_echo_does_not_acknowledge_a_new_probe() { + let liveness = PeerConnLiveness::new(); + liveness.set_enabled(true); + + let old_token = liveness.start_probe().unwrap(); + let mut old_echo = data_packet(2, 1); + old_echo + .mut_peer_manager_header() + .unwrap() + .set_liveness_echo(old_token); + liveness.after_received(Ok(old_echo)).unwrap().unwrap(); + liveness.wait_for_echo(old_token).await; + + let new_token = liveness.start_probe().unwrap(); + assert_ne!(old_token, new_token); + let mut stale_echo = data_packet(2, 1); + stale_echo + .mut_peer_manager_header() + .unwrap() + .set_liveness_echo(old_token); + liveness.after_received(Ok(stale_echo)).unwrap().unwrap(); + + assert!( + tokio::time::timeout(Duration::from_millis(20), liveness.wait_for_echo(new_token)) + .await + .is_err(), + "stale echo acknowledged a newer probe" + ); + } + + #[tokio::test] + async fn simultaneous_probes_are_echoed_in_both_directions() { + let left = PeerConnLiveness::new(); + let right = PeerConnLiveness::new(); + left.set_enabled(true); + right.set_enabled(true); + let left_token = left.start_probe().unwrap(); + let right_token = right.start_probe().unwrap(); + + let left_probe = left.before_send(data_packet(1, 2)).unwrap(); + let right_probe = right.before_send(data_packet(2, 1)).unwrap(); + right.after_received(Ok(left_probe)).unwrap().unwrap(); + left.after_received(Ok(right_probe)).unwrap().unwrap(); + + let left_echo = left.before_send(data_packet(1, 2)).unwrap(); + let right_echo = right.before_send(data_packet(2, 1)).unwrap(); + right.after_received(Ok(left_echo)).unwrap().unwrap(); + left.after_received(Ok(right_echo)).unwrap().unwrap(); + + tokio::time::timeout(Duration::from_millis(50), left.wait_for_echo(left_token)) + .await + .unwrap(); + tokio::time::timeout(Duration::from_millis(50), right.wait_for_echo(right_token)) + .await + .unwrap(); + } +} diff --git a/easytier-core/src/peers/conn/peer_conn_ping.rs b/easytier-core/src/peers/conn/peer_conn_ping.rs index 1826f3ff..77ddecc4 100644 --- a/easytier-core/src/peers/conn/peer_conn_ping.rs +++ b/easytier-core/src/peers/conn/peer_conn_ping.rs @@ -7,21 +7,28 @@ use std::{ }; use rand::{Rng, thread_rng}; -use tokio::{sync::broadcast, task::JoinSet}; -use tracing::Instrument; +use tokio::{ + sync::{broadcast, mpsc::error::TrySendError}, + task::JoinSet, +}; use crate::{ config::PeerId, foundation::time::{Interval, interval, timeout}, packet::{PacketType, ZCPacket}, - peers::{context::ArcPeerContext, error::Error}, + peers::{conn::peer_conn_liveness::PeerConnLiveness, context::ArcPeerContext, error::Error}, tunnel::{ - TunnelError, mpsc::MpscTunnelSender, stats::{Throughput, WindowLatency}, }, }; +#[derive(Debug)] +enum PingResponse { + Pong(u128), + LivenessEcho, +} + struct PingIntervalController { throughput: Arc, loss_counter: Arc, @@ -118,7 +125,7 @@ pub struct PeerConnPinger { throughput_stats: Arc, context: ArcPeerContext, network_name: String, - tasks: JoinSet>, + liveness: PeerConnLiveness, } impl std::fmt::Debug for PeerConnPinger { @@ -132,7 +139,7 @@ impl std::fmt::Debug for PeerConnPinger { impl PeerConnPinger { #[allow(clippy::too_many_arguments)] - pub(crate) fn new( + pub(super) fn new( my_peer_id: PeerId, peer_id: PeerId, sink: MpscTunnelSender, @@ -142,18 +149,19 @@ impl PeerConnPinger { throughput_stats: Arc, context: ArcPeerContext, network_name: String, + liveness: PeerConnLiveness, ) -> Self { Self { my_peer_id, peer_id, sink, - tasks: JoinSet::new(), latency_stats, ctrl_sender, loss_rate_stats, throughput_stats, context, network_name, + liveness, } } @@ -164,23 +172,19 @@ impl PeerConnPinger { } async fn do_pingpong_once( - my_node_id: PeerId, - peer_id: PeerId, - sink: &MpscTunnelSender, - context: &ArcPeerContext, - network_name: &str, + &self, receiver: &mut broadcast::Receiver, seq: u32, - ) -> Result { + liveness_token: Option, + ) -> Result { // should add seq here. so latency can be calculated more accurately - let req = Self::new_ping_packet(my_node_id, peer_id, seq); + let req = Self::new_ping_packet(self.my_peer_id, self.peer_id, seq); let req_len = req.buf_len() as u64; - sink.send(req).await?; - context.record_control_tx(network_name, req_len); + self.sink.send(req).await?; + self.context.record_control_tx(&self.network_name, req_len); let now = Instant::now(); - // wait until we get a pong packet in ctrl_resp_receiver - let resp = timeout(Duration::from_secs(2), async { + let wait_for_pong = async { loop { match receiver.recv().await { Ok(p) => { @@ -203,6 +207,18 @@ impl PeerConnPinger { } } Ok(()) + }; + let resp = timeout(Duration::from_secs(2), async { + if let Some(token) = liveness_token { + tokio::select! { + ret = wait_for_pong => ret.map(|()| PingResponse::Pong(now.elapsed().as_micros())), + () = self.liveness.wait_for_echo(token) => Ok(PingResponse::LivenessEcho), + } + } else { + wait_for_pong + .await + .map(|()| PingResponse::Pong(now.elapsed().as_micros())) + } }) .await; @@ -214,98 +230,59 @@ impl PeerConnPinger { )); } - if resp.as_ref().unwrap().is_err() { - return Err(resp.unwrap().err().unwrap()); - } - - Ok(now.elapsed().as_micros()) + resp.unwrap() } pub async fn pingpong(&mut self) { - let sink = self.sink.clone(); - let context = self.context.clone(); - let network_name = self.network_name.clone(); let my_node_id = self.my_peer_id; - let peer_id = self.peer_id; - let latency_stats = self.latency_stats.clone(); - - let (ping_res_sender, mut ping_res_receiver) = tokio::sync::mpsc::channel(100); // one with 1% precision let loss_rate_stats_1 = WindowLatency::new(100); // disconnect the connection if lost 5 pingpong consecutively let loss_counter = Arc::new(AtomicU32::new(0)); - let stopped = Arc::new(AtomicU32::new(0)); - - // generate a pingpong task every 200ms - let mut pingpong_tasks = JoinSet::new(); - let ctrl_resp_sender = self.ctrl_sender.clone(); - let stopped_clone = stopped.clone(); - let mut controller = - PingIntervalController::new(self.throughput_stats.clone(), loss_counter.clone()); - self.tasks.spawn( - async move { - let mut req_seq = 0; - loop { - controller.tick().await; - - if stopped_clone.load(Ordering::Relaxed) != 0 { - return Ok(()); - } - - while pingpong_tasks.len() > 5 { - pingpong_tasks.join_next().await; - } - - if !controller.should_send_ping() { - continue; - } - - tracing::debug!( - "pingpong controller send pingpong task, seq: {}, node_id: {}, controller: {:?}", - req_seq, - my_node_id, - controller, - ); - - let sink = sink.clone(); - let context = context.clone(); - let network_name = network_name.clone(); - let receiver = ctrl_resp_sender.subscribe(); - let ping_res_sender = ping_res_sender.clone(); - pingpong_tasks.spawn(async move { - let mut receiver = receiver.resubscribe(); - let pingpong_once_ret = Self::do_pingpong_once( - my_node_id, - peer_id, - &sink, - &context, - &network_name, - &mut receiver, - req_seq, - ) - .await; - - if let Err(e) = ping_res_sender.send(pingpong_once_ret).await { - tracing::info!(?e, "pingpong task send result error, exit.."); - }; - }); - - req_seq = req_seq.wrapping_add(1); + let (trigger_sender, mut trigger_receiver) = tokio::sync::mpsc::channel(1); + let mut controller_tasks = JoinSet::new(); + let throughput = self.throughput_stats.clone(); + let controller_loss_counter = loss_counter.clone(); + controller_tasks.spawn(async move { + let mut controller = PingIntervalController::new(throughput, controller_loss_counter); + loop { + controller.tick().await; + if !controller.should_send_ping() { + continue; + } + match trigger_sender.try_send(()) { + Ok(()) | Err(TrySendError::Full(())) => {} + Err(TrySendError::Closed(())) => break, } } - .instrument(tracing::info_span!( - "pingpong_controller", - ?my_node_id, - ?peer_id - )), - ); + }); - while let Some(ret) = ping_res_receiver.recv().await { - if let Ok(lat) = ret { - latency_stats.record_latency(lat as u32); + let mut req_seq = 0u32; + while trigger_receiver.recv().await.is_some() { + tracing::debug!( + "pingpong controller send pingpong task, seq: {}, node_id: {}", + req_seq, + my_node_id, + ); + let liveness_token = (loss_counter.load(Ordering::Relaxed) > 0) + .then(|| self.liveness.start_probe()) + .flatten(); + let mut receiver = self.ctrl_sender.subscribe(); + let ret = self + .do_pingpong_once(&mut receiver, req_seq, liveness_token) + .await; + req_seq = req_seq.wrapping_add(1); + + if let Ok(response) = &ret { + if let PingResponse::Pong(lat) = response { + self.latency_stats.record_latency(*lat as u32); + } + if let Some(token) = liveness_token { + self.liveness.finish_probe(token); + } loss_rate_stats_1.record_latency(0); loss_counter.store(0, Ordering::Relaxed); } else { @@ -343,18 +320,21 @@ impl PeerConnPinger { self.loss_rate_stats .store((loss_rate_1 * 100.0) as u32, Ordering::Relaxed); } - - stopped.store(1, Ordering::Relaxed); - ping_res_receiver.close(); } } #[cfg(test)] mod tests { + use futures::{SinkExt, StreamExt}; + use super::*; use crate::{ + packet::PacketType, + peers::conn::peer_conn_liveness::PeerConnLiveness, peers::test_support::NoopPeerContext, - tunnel::{mpsc::MpscTunnel, ring::create_ring_tunnel_pair}, + tunnel::{ + Tunnel, filter::TunnelWithFilter, mpsc::MpscTunnel, ring::create_ring_tunnel_pair, + }, }; #[tokio::test(flavor = "current_thread")] @@ -373,6 +353,7 @@ mod tests { throughput.clone(), Arc::new(NoopPeerContext::default()), "test".to_owned(), + PeerConnLiveness::new(), ); let ingress = tokio::spawn(async move { @@ -390,4 +371,52 @@ mod tests { "unrelated ingress traffic kept a failed round-trip alive" ); } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn echoed_business_traffic_keeps_connection_alive_when_pongs_are_lost() { + let local_liveness = PeerConnLiveness::new(); + let remote_liveness = PeerConnLiveness::new(); + local_liveness.set_enabled(true); + remote_liveness.set_enabled(true); + + let (local_transport, remote_transport) = create_ring_tunnel_pair(); + let local_transport = TunnelWithFilter::new(local_transport, local_liveness.clone()); + let mut local_tunnel = MpscTunnel::new(local_transport, None); + let mut local_stream = local_tunnel.get_stream(); + let local_reader = + tokio::spawn(async move { while local_stream.next().await.is_some() {} }); + + let remote_transport = TunnelWithFilter::new(remote_transport, remote_liveness); + let (mut remote_stream, mut remote_sink) = remote_transport.split(); + let remote = tokio::spawn(async move { + while let Some(Ok(_packet)) = remote_stream.next().await { + let mut reply = ZCPacket::new_with_payload(b"business traffic"); + reply.fill_peer_manager_hdr(2, 1, PacketType::Data as u8); + remote_sink.send(reply).await.unwrap(); + } + }); + + let (ctrl_sender, _) = broadcast::channel(16); + let mut pinger = PeerConnPinger::new( + 1, + 2, + local_tunnel.get_sink(), + ctrl_sender, + Arc::new(WindowLatency::new(15)), + Arc::new(AtomicU32::new(0)), + Arc::new(Throughput::new()), + Arc::new(NoopPeerContext::default()), + "test".to_owned(), + local_liveness, + ); + + let result = timeout(Duration::from_secs(12), pinger.pingpong()).await; + remote.abort(); + local_reader.abort(); + + assert!( + result.is_err(), + "matching business-packet echoes did not keep the connection alive" + ); + } } diff --git a/easytier-core/src/peers/test_support.rs b/easytier-core/src/peers/test_support.rs index 7da13bb1..c7f15ff2 100644 --- a/easytier-core/src/peers/test_support.rs +++ b/easytier-core/src/peers/test_support.rs @@ -70,6 +70,12 @@ impl NoopPeerContext { secure_mode: None, } } + + pub(crate) fn with_secure_mode(mut self, secure_mode: SecureModeConfig) -> Self { + self.flags.encryption_algorithm = "aes-gcm".to_owned(); + self.secure_mode = Some(secure_mode); + self + } } impl Default for NoopPeerContext { diff --git a/easytier-core/src/peers/tests.rs b/easytier-core/src/peers/tests.rs index a2a93c3b..d2d6256d 100644 --- a/easytier-core/src/peers/tests.rs +++ b/easytier-core/src/peers/tests.rs @@ -1,5 +1,8 @@ use std::sync::Arc; +use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64_STANDARD}; +use x25519_dalek::{PublicKey, StaticSecret}; + use crate::foundation::time::{Duration, timeout}; use crate::{ @@ -44,6 +47,52 @@ async fn peer_conn_handshake_over_memory_tunnel() { server_ret.unwrap(); assert_eq!(client.get_peer_id(), 2); assert_eq!(server.get_peer_id(), 1); + assert_eq!(client.get_conn_info().features, ["liveness-echo-v1"]); + assert_eq!(server.get_conn_info().features, ["liveness-echo-v1"]); +} + +#[tokio::test] +async fn peer_conn_noise_handshake_advertises_liveness_echo() { + fn context(peer_key: u8) -> Arc { + let private = StaticSecret::from([peer_key; 32]); + let public = PublicKey::from(&private); + Arc::new( + NoopPeerContext::new(NetworkIdentity { + network_name: "net".to_owned(), + network_secret: Some("secret".to_owned()), + network_secret_digest: None, + }) + .with_secure_mode(crate::proto::common::SecureModeConfig { + enabled: true, + local_private_key: Some(BASE64_STANDARD.encode(private.as_bytes())), + local_public_key: Some(BASE64_STANDARD.encode(public.as_bytes())), + }), + ) + } + + let (client_tunnel, server_tunnel) = create_ring_tunnel_pair(); + let mut client = PeerConn::new( + 1, + context(1), + client_tunnel, + Arc::new(PeerSessionStore::new()), + ); + let mut server = PeerConn::new( + 2, + context(2), + server_tunnel, + Arc::new(PeerSessionStore::new()), + ); + + let (client_ret, server_ret) = tokio::join!( + client.do_handshake_as_client(), + server.do_handshake_as_server() + ); + + client_ret.unwrap(); + server_ret.unwrap(); + assert_eq!(client.get_conn_info().features, ["liveness-echo-v1"]); + assert_eq!(server.get_conn_info().features, ["liveness-echo-v1"]); } #[tokio::test] diff --git a/easytier-proto/proto/peer_rpc.proto b/easytier-proto/proto/peer_rpc.proto index 7c8956d9..69ea0224 100644 --- a/easytier-proto/proto/peer_rpc.proto +++ b/easytier-proto/proto/peer_rpc.proto @@ -348,6 +348,7 @@ message PeerConnNoiseMsg1Pb { optional uint32 a_session_generation = 3; common.UUID a_conn_id = 4; string client_encryption_algorithm = 5; + repeated string features = 6; } message PeerConnNoiseMsg2Pb { @@ -361,6 +362,7 @@ message PeerConnNoiseMsg2Pb { common.UUID a_conn_id_echo = 8; optional bytes secret_proof_32 = 9; string server_encryption_algorithm = 10; + repeated string features = 11; } message RelayNoiseMsg1Pb { diff --git a/easytier/src/tests/three_node.rs b/easytier/src/tests/three_node.rs index 15ffa902..8955b3de 100644 --- a/easytier/src/tests/three_node.rs +++ b/easytier/src/tests/three_node.rs @@ -1541,9 +1541,9 @@ pub async fn proxy_three_node_disconnect_test(#[values("tcp", "wg")] proto: &str .any(|r| *r == inst4.peer_id()) }, // 0 down, assume last packet is recv in -0.01 - // [2, 7) send ping - // [4, 9) ping fail and close connection - Duration::from_secs(11), + // one ping outstanding at a time, each waits up to 2s: + // 5 consecutive failures close the connection at ~[4, 11) + Duration::from_secs(15), ) .await;