diff --git a/Cargo.lock b/Cargo.lock index 5fcfffb..6f35d93 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1432,7 +1432,7 @@ dependencies = [ [[package]] name = "mqp2p" -version = "0.1.0" +version = "0.2.0" dependencies = [ "clap", "flume", diff --git a/crates/mqp2p/Cargo.toml b/crates/mqp2p/Cargo.toml index 856e098..1a3900a 100644 --- a/crates/mqp2p/Cargo.toml +++ b/crates/mqp2p/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "mqp2p" -version = "0.1.0" +version = "0.2.0" edition = "2024" rust-version = "1.88" license = "Apache-2.0" @@ -28,3 +28,4 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] } clap = { version = "4", features = ["derive"] } tempfile = "3" rand = "0.10.0" +mqtt5 = "0.38" diff --git a/crates/mqp2p/src/peer.rs b/crates/mqp2p/src/peer.rs index 34ce99f..e042ed8 100644 --- a/crates/mqp2p/src/peer.rs +++ b/crates/mqp2p/src/peer.rs @@ -1,7 +1,8 @@ use crate::error::{Error, Result}; use crate::quic::{self, CertIdentity, QuicEndpoint}; use crate::signaling::{ - Candidate, CandidateKind, ConnectionOffer, ConnectionSync, PeerInfo, SignalingClient, + Candidate, CandidateKind, ConnectionOffer, ConnectionSync, MIN_LEASE_SECS, PEER_LEASE_SECS, + PeerInfo, SignalingClient, heartbeat_interval_secs, }; use crate::stun; use crate::transfer::{self, FileOffer, TransferProgress, TransferResult}; @@ -18,6 +19,7 @@ pub struct PeerConfig { pub bind_addr: SocketAddr, pub stun_server: Option, pub credentials: Option<(String, String)>, + pub lease_secs: u64, } impl PeerConfig { @@ -28,9 +30,15 @@ impl PeerConfig { bind_addr: "0.0.0.0:0".parse().unwrap_or_else(|_| unreachable!()), stun_server: Some(stun::DEFAULT_STUN_SERVER.into()), credentials: None, + lease_secs: PEER_LEASE_SECS, } } + pub fn with_lease_secs(mut self, lease_secs: u64) -> Self { + self.lease_secs = lease_secs; + self + } + pub fn with_bind_addr(mut self, addr: SocketAddr) -> Self { self.bind_addr = addr; self @@ -63,10 +71,13 @@ pub struct Peer { identity: CertIdentity, signaling: SignalingClient, quic: QuicEndpoint, + record_id: String, peer_id: Option, std_socket: std::net::UdpSocket, host_addr: SocketAddr, srflx_addr: Option, + lease_secs: u64, + heartbeat: Option>, } impl Peer { @@ -131,26 +142,67 @@ impl Peer { identity, signaling, quic, + record_id: uuid::Uuid::new_v4().to_string(), peer_id: None, std_socket, host_addr, srflx_addr, + lease_secs: config.lease_secs.max(MIN_LEASE_SECS), + heartbeat: None, }) } pub async fn register(&mut self) -> Result { let quic_port = self.quic.local_addr()?.port(); - - let peer_id = self - .signaling - .register_peer(&self.name, quic_port, &self.identity.fingerprint) + let peer_id = self.record_id.clone(); + let public_addr = self.srflx_addr.map(|a| a.to_string()); + + self.signaling + .register_peer( + &peer_id, + &self.name, + quic_port, + &self.identity.fingerprint, + public_addr.as_deref(), + self.lease_secs, + ) .await?; - if let Some(addr) = self.srflx_addr { - self.signaling.update_peer_addr(&peer_id, addr).await?; + self.peer_id = Some(peer_id.clone()); + + if let Some(handle) = self.heartbeat.take() { + handle.abort(); } - self.peer_id = Some(peer_id.clone()); + let signaling = self.signaling.clone(); + let renew_id = peer_id.clone(); + let renew_name = self.name.clone(); + let renew_fingerprint = self.identity.fingerprint.clone(); + let renew_public_addr = public_addr; + let lease_secs = self.lease_secs; + self.heartbeat = Some(tokio::spawn(async move { + let mut ticker = tokio::time::interval(std::time::Duration::from_secs( + heartbeat_interval_secs(lease_secs), + )); + ticker.tick().await; + loop { + ticker.tick().await; + if let Err(e) = signaling + .ensure_registered( + &renew_id, + &renew_name, + quic_port, + &renew_fingerprint, + renew_public_addr.as_deref(), + lease_secs, + ) + .await + { + warn!(peer_id = renew_id, error = %e, "peer presence renewal failed"); + } + } + })); + info!(peer_id, name = self.name, "peer registered"); Ok(peer_id) } @@ -371,6 +423,7 @@ impl Peer { quic_port: remote_addr.port(), cert_fingerprint: offer.cert_fingerprint, public_addr: Some(remote_addr.to_string()), + expires_at: None, }; info!( @@ -385,7 +438,10 @@ impl Peer { }) } - pub async fn shutdown(self) -> Result<()> { + pub async fn shutdown(mut self) -> Result<()> { + if let Some(handle) = self.heartbeat.take() { + handle.abort(); + } if let Some(peer_id) = &self.peer_id { let _ = self.signaling.deregister_peer(peer_id).await; } @@ -395,6 +451,14 @@ impl Peer { } } +impl Drop for Peer { + fn drop(&mut self) { + if let Some(handle) = self.heartbeat.take() { + handle.abort(); + } + } +} + pub struct P2pConnection { connection: Connection, remote_peer: PeerInfo, diff --git a/crates/mqp2p/src/signaling.rs b/crates/mqp2p/src/signaling.rs index 5c55510..ff061a8 100644 --- a/crates/mqp2p/src/signaling.rs +++ b/crates/mqp2p/src/signaling.rs @@ -4,11 +4,30 @@ use mqtt5::types::{PublishOptions, PublishProperties}; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::net::SocketAddr; -use std::time::Duration; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tracing::{debug, error, trace}; const DEFAULT_TIMEOUT_MS: u64 = 5000; +pub const PEER_LEASE_SECS: u64 = 15; + +pub const MIN_LEASE_SECS: u64 = 2; + +pub fn heartbeat_interval_secs(lease_secs: u64) -> u64 { + (lease_secs / 3).max(1) +} + +fn unix_now_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) +} + +fn peer_is_live(expires_at: Option, now: u64) -> bool { + expires_at.map(|e| e > now).unwrap_or(true) +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PeerInfo { pub id: String, @@ -18,6 +37,8 @@ pub struct PeerInfo { pub cert_fingerprint: String, #[serde(default)] pub public_addr: Option, + #[serde(rename = "_expires_at", default)] + pub expires_at: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -49,6 +70,7 @@ pub enum CandidateKind { Srflx, } +#[derive(Clone)] pub struct SignalingClient { client: MqttClient, } @@ -73,31 +95,73 @@ impl SignalingClient { pub async fn register_peer( &self, + id: &str, name: &str, quic_port: u16, cert_fingerprint: &str, - ) -> Result { - let payload = serde_json::json!({ + public_addr: Option<&str>, + lease_secs: u64, + ) -> Result<()> { + let mut payload = serde_json::json!({ + "id": id, "name": name, "status": "online", "quic_port": quic_port, "cert_fingerprint": cert_fingerprint, + "_expires_at": unix_now_secs() + lease_secs, }); + if let Some(addr) = public_addr { + payload["public_addr"] = Value::String(addr.to_string()); + } let response = self .publish_and_wait("$DB/peers/create", &serde_json::to_vec(&payload)?) .await?; check_response(&response)?; + debug!(peer_id = id, name, "peer registered"); + Ok(()) + } - let id = response - .pointer("/data/id") - .and_then(|v| v.as_str()) - .ok_or_else(|| Error::Signaling("no peer ID in create response".into()))? - .to_string(); + pub async fn renew_lease(&self, peer_id: &str, lease_secs: u64) -> Result<()> { + let topic = format!("$DB/peers/{peer_id}/update"); + let payload = serde_json::json!({ + "_expires_at": unix_now_secs() + lease_secs, + }); - debug!(peer_id = id, name, "peer registered"); - Ok(id) + let response = self + .publish_and_wait(&topic, &serde_json::to_vec(&payload)?) + .await?; + + check_response(&response)?; + trace!(peer_id, "peer lease renewed"); + Ok(()) + } + + pub async fn ensure_registered( + &self, + id: &str, + name: &str, + quic_port: u16, + cert_fingerprint: &str, + public_addr: Option<&str>, + lease_secs: u64, + ) -> Result<()> { + match self.renew_lease(id, lease_secs).await { + Ok(()) => Ok(()), + Err(e) => { + debug!(peer_id = id, error = %e, "lease renewal failed, re-creating record"); + self.register_peer( + id, + name, + quic_port, + cert_fingerprint, + public_addr, + lease_secs, + ) + .await + } + } } pub async fn update_peer_addr(&self, peer_id: &str, public_addr: SocketAddr) -> Result<()> { @@ -131,9 +195,11 @@ impl SignalingClient { other => vec![other], }; + let now = unix_now_secs(); let peers: Vec = items .into_iter() - .filter_map(|v| serde_json::from_value(v).ok()) + .filter_map(|v| serde_json::from_value::(v).ok()) + .filter(|p| peer_is_live(p.expires_at, now)) .collect(); debug!(count = peers.len(), "peers discovered"); @@ -306,3 +372,66 @@ fn check_response(response: &Value) -> Result<()> { Err(Error::Signaling(format!("DB error: {msg}"))) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn leaseless_peer_is_live() { + assert!(peer_is_live(None, 1000)); + } + + #[test] + fn future_lease_is_live() { + assert!(peer_is_live(Some(1001), 1000)); + } + + #[test] + fn lease_at_current_second_is_stale() { + assert!(!peer_is_live(Some(1000), 1000)); + } + + #[test] + fn past_lease_is_stale() { + assert!(!peer_is_live(Some(999), 1000)); + } + + #[test] + fn heartbeat_interval_is_a_third_of_lease_min_one() { + assert_eq!(heartbeat_interval_secs(15), 5); + assert_eq!(heartbeat_interval_secs(3), 1); + assert_eq!(heartbeat_interval_secs(2), 1); + assert_eq!(heartbeat_interval_secs(1), 1); + } + + #[test] + fn expires_at_round_trips_as_mqdb_field() { + let info = PeerInfo { + id: "p1".into(), + name: "edge".into(), + status: "online".into(), + quic_port: 4433, + cert_fingerprint: "ab".into(), + public_addr: None, + expires_at: Some(1234), + }; + let json = serde_json::to_value(&info).expect("serialize"); + assert_eq!(json.get("_expires_at").and_then(Value::as_u64), Some(1234)); + let back: PeerInfo = serde_json::from_value(json).expect("deserialize"); + assert_eq!(back.expires_at, Some(1234)); + } + + #[test] + fn missing_expires_at_deserializes_to_none() { + let json = serde_json::json!({ + "id": "p1", + "name": "edge", + "status": "online", + "quic_port": 4433, + "cert_fingerprint": "ab", + }); + let info: PeerInfo = serde_json::from_value(json).expect("deserialize"); + assert_eq!(info.expires_at, None); + } +} diff --git a/crates/mqp2p/tests/integration.rs b/crates/mqp2p/tests/integration.rs index cbf0dd4..716d4d1 100644 --- a/crates/mqp2p/tests/integration.rs +++ b/crates/mqp2p/tests/integration.rs @@ -53,6 +53,60 @@ fn wait_for_port(port: u16) { panic!("broker did not start on port {port} within 10s"); } +fn delete_peer_record(port: u16, id: &str) { + let broker = format!("127.0.0.1:{port}"); + let status = Command::new("mqdb") + .args([ + "delete", + "peers", + id, + "--broker", + &broker, + "--user", + "testuser", + "--pass", + "testpass", + "--timeout", + "5", + ]) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status() + .expect("failed to run mqdb delete"); + assert!(status.success(), "out-of-band delete of peer {id} failed"); +} + +fn wait_for_broker_ready(port: u16) { + let broker = format!("127.0.0.1:{port}"); + for _ in 0..100 { + let ok = Command::new("mqdb") + .args([ + "list", + "peers", + "--broker", + &broker, + "--user", + "testuser", + "--pass", + "testpass", + "--timeout", + "2", + "--limit", + "1", + ]) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status() + .map(|s| s.success()) + .unwrap_or(false); + if ok { + return; + } + std::thread::sleep(Duration::from_millis(200)); + } + panic!("broker on port {port} did not answer a $DB round-trip within 20s"); +} + fn create_passwd_file(dir: &std::path::Path) -> std::path::PathBuf { let passwd_path = dir.join("passwd.txt"); let output = Command::new("mqdb") @@ -86,6 +140,7 @@ async fn test_full_file_transfer() { let broker_port = pick_free_port(); let _broker = start_broker(broker_port, &db_dir, &passwd_file); wait_for_port(broker_port); + wait_for_broker_ready(broker_port); let broker_addr = format!("127.0.0.1:{broker_port}"); let bind_addr: SocketAddr = "127.0.0.1:0".parse().unwrap(); @@ -180,3 +235,208 @@ async fn test_full_file_transfer() { peer_a.shutdown().await.expect("peer A shutdown failed"); } + +#[tokio::test] +async fn presence_lease_expires_without_heartbeat() { + let _ = tracing_subscriber::fmt() + .with_env_filter("warn") + .with_test_writer() + .try_init(); + + let tmpdir = tempfile::tempdir().expect("failed to create tmpdir"); + let db_dir = tmpdir.path().join("db"); + std::fs::create_dir_all(&db_dir).expect("failed to create db dir"); + let passwd_file = create_passwd_file(tmpdir.path()); + + let broker_port = pick_free_port(); + let _broker = start_broker(broker_port, &db_dir, &passwd_file); + wait_for_port(broker_port); + wait_for_broker_ready(broker_port); + + let broker_addr = format!("127.0.0.1:{broker_port}"); + let bind_addr: SocketAddr = "127.0.0.1:0".parse().unwrap(); + let lease_secs = 2; + + let mut observer = Peer::new( + PeerConfig::new("observer", &broker_addr) + .with_bind_addr(bind_addr) + .without_stun() + .with_credentials("testuser", "testpass") + .with_lease_secs(lease_secs), + ) + .await + .expect("observer creation failed"); + observer + .register() + .await + .expect("observer registration failed"); + + let ephemeral_id = { + let mut ephemeral = Peer::new( + PeerConfig::new("ephemeral", &broker_addr) + .with_bind_addr(bind_addr) + .without_stun() + .with_credentials("testuser", "testpass") + .with_lease_secs(lease_secs), + ) + .await + .expect("ephemeral creation failed"); + let id = ephemeral + .register() + .await + .expect("ephemeral registration failed"); + + let seen = observer + .discover_peers() + .await + .expect("discovery failed") + .into_iter() + .any(|p| p.id == id); + assert!( + seen, + "ephemeral peer should be discoverable while registered" + ); + + id + }; + + let deadline = tokio::time::Instant::now() + Duration::from_secs(8); + loop { + let peers = observer.discover_peers().await.expect("discovery failed"); + let ephemeral_present = peers.iter().any(|p| p.id == ephemeral_id); + let observer_present = peers + .iter() + .any(|p| Some(p.id.as_str()) == observer.peer_id()); + + assert!( + observer_present, + "observer must stay present while its heartbeat renews the lease" + ); + + if !ephemeral_present { + break; + } + + if tokio::time::Instant::now() >= deadline { + panic!("ephemeral peer did not expire from discovery after its lease lapsed"); + } + tokio::time::sleep(Duration::from_millis(250)).await; + } + + observer.shutdown().await.expect("observer shutdown failed"); +} + +#[tokio::test] +async fn presence_recreates_after_record_reaped() { + let _ = tracing_subscriber::fmt() + .with_env_filter("warn") + .with_test_writer() + .try_init(); + + let tmpdir = tempfile::tempdir().expect("failed to create tmpdir"); + let db_dir = tmpdir.path().join("db"); + std::fs::create_dir_all(&db_dir).expect("failed to create db dir"); + let passwd_file = create_passwd_file(tmpdir.path()); + + let broker_port = pick_free_port(); + let _broker = start_broker(broker_port, &db_dir, &passwd_file); + wait_for_port(broker_port); + wait_for_broker_ready(broker_port); + + let broker_addr = format!("127.0.0.1:{broker_port}"); + let bind_addr: SocketAddr = "127.0.0.1:0".parse().unwrap(); + + let mut resilient = Peer::new( + PeerConfig::new("resilient", &broker_addr) + .with_bind_addr(bind_addr) + .without_stun() + .with_credentials("testuser", "testpass") + .with_lease_secs(2), + ) + .await + .expect("peer creation failed"); + let id = resilient.register().await.expect("registration failed"); + + let present = |peers: &[mqp2p::PeerInfo]| peers.iter().any(|p| p.id == id); + assert!( + present(&resilient.discover_peers().await.expect("discovery failed")), + "peer should be present right after registration" + ); + + delete_peer_record(broker_port, &id); + + let deadline = tokio::time::Instant::now() + Duration::from_secs(8); + loop { + let peers = resilient.discover_peers().await.expect("discovery failed"); + if present(&peers) { + break; + } + if tokio::time::Instant::now() >= deadline { + panic!("peer did not re-register after its record was reaped"); + } + tokio::time::sleep(Duration::from_millis(250)).await; + } + + resilient.shutdown().await.expect("shutdown failed"); +} + +#[tokio::test] +async fn reregister_preserves_public_addr() { + let _ = tracing_subscriber::fmt() + .with_env_filter("warn") + .with_test_writer() + .try_init(); + + let tmpdir = tempfile::tempdir().expect("failed to create tmpdir"); + let db_dir = tmpdir.path().join("db"); + std::fs::create_dir_all(&db_dir).expect("failed to create db dir"); + let passwd_file = create_passwd_file(tmpdir.path()); + + let broker_port = pick_free_port(); + let _broker = start_broker(broker_port, &db_dir, &passwd_file); + wait_for_port(broker_port); + wait_for_broker_ready(broker_port); + + let broker_addr = format!("127.0.0.1:{broker_port}"); + let opts = mqtt5::ConnectOptions::new("sig-test".to_string()) + .with_credentials("testuser", b"testpass"); + let client = mqtt5::client::MqttClient::with_options(opts); + let sig = mqp2p::signaling::SignalingClient::new(client); + sig.connect(&broker_addr) + .await + .expect("signaling connect failed"); + + let id = "peer-with-public-addr"; + let public_addr = "203.0.113.5:9000"; + sig.register_peer(id, "edge", 4433, "fp-abc", Some(public_addr), 2) + .await + .expect("register failed"); + + let found = sig + .list_peers() + .await + .expect("list failed") + .into_iter() + .find(|p| p.id == id) + .expect("peer should be present after registration"); + assert_eq!(found.public_addr.as_deref(), Some(public_addr)); + + delete_peer_record(broker_port, id); + + sig.ensure_registered(id, "edge", 4433, "fp-abc", Some(public_addr), 2) + .await + .expect("ensure_registered failed"); + + let found = sig + .list_peers() + .await + .expect("list failed") + .into_iter() + .find(|p| p.id == id) + .expect("peer should be present after re-create"); + assert_eq!( + found.public_addr.as_deref(), + Some(public_addr), + "public_addr must survive re-registration" + ); +} diff --git a/crates/stitch-p2p/Cargo.toml b/crates/stitch-p2p/Cargo.toml index 6b532eb..906802d 100644 --- a/crates/stitch-p2p/Cargo.toml +++ b/crates/stitch-p2p/Cargo.toml @@ -13,7 +13,7 @@ readme = "README.md" [dependencies] ed25519-dalek = { version = "3", default-features = false, features = ["fast"], optional = true } fjall = { version = "3.0", optional = true } -mqp2p = { version = "0.1.0", path = "../mqp2p", optional = true } +mqp2p = { version = "0.2.0", path = "../mqp2p", optional = true } serde_json = { version = "1.0.150", optional = true } thiserror = "2.0.18" tokio = { version = "1.52.3", features = ["sync", "io-util", "macros", "rt", "rt-multi-thread", "time", "net"] } diff --git a/crates/stitch-tasks/Cargo.toml b/crates/stitch-tasks/Cargo.toml index 8098e60..df6b914 100644 --- a/crates/stitch-tasks/Cargo.toml +++ b/crates/stitch-tasks/Cargo.toml @@ -6,7 +6,7 @@ license = "Apache-2.0" publish = false [dependencies] -mqp2p = { version = "0.1.0", path = "../mqp2p" } +mqp2p = { version = "0.2.0", path = "../mqp2p" } rand = "0.10.1" serde = { version = "1.0.228", features = ["derive"] } serde_json = "1.0.150"