use std::sync::Arc;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::sync::{broadcast, mpsc};
use crate::network::NetworkProvider;
use crate::network::{
HealthInfo, IncomingConnection, NetworkError, NetworkPeer, NetworkPeerEvent,
NetworkTcpListener, NetworkUdpSocket, NodeIdentity, PeerAddr, PingResult,
};
use crate::transport::websocket::WebSocketTransport;
use crate::transport::WsConfig;
use super::reconnect::ReconnectBackoff;
use super::{
format_peer_ref, parse_peer_ref, PeerEvent, PeerIdentity, PeerRegistry, PeerRegistryOptions,
SessionError,
};
struct MockNetworkProvider {
identity: NodeIdentity,
local_addr: PeerAddr,
peer_event_tx: broadcast::Sender<NetworkPeerEvent>,
}
impl MockNetworkProvider {
#[allow(dead_code)]
fn new(id: &str) -> Self {
Self::new_with_app("test", id)
}
fn new_with_app(app_id: &str, id: &str) -> Self {
let (peer_event_tx, _) = broadcast::channel(64);
Self {
identity: NodeIdentity {
app_id: app_id.to_string(),
device_id: format!("dev-{id}"),
device_name: format!("Test Node {id}"),
tailscale_hostname: format!("truffle-{app_id}-{id}"),
tailscale_id: id.to_string(),
dns_name: None,
ip: Some("127.0.0.1".parse().unwrap()),
},
local_addr: PeerAddr {
ip: Some("127.0.0.1".parse().unwrap()),
hostname: format!("truffle-{app_id}-{id}"),
dns_name: None,
},
peer_event_tx,
}
}
fn event_sender(&self) -> broadcast::Sender<NetworkPeerEvent> {
self.peer_event_tx.clone()
}
}
impl NetworkProvider for MockNetworkProvider {
async fn start(&mut self) -> Result<(), NetworkError> {
Ok(())
}
async fn stop(&self) -> Result<(), NetworkError> {
Ok(())
}
fn local_identity(&self) -> NodeIdentity {
self.identity.clone()
}
fn local_addr(&self) -> PeerAddr {
self.local_addr.clone()
}
fn peer_events(&self) -> broadcast::Receiver<NetworkPeerEvent> {
self.peer_event_tx.subscribe()
}
async fn peers(&self) -> Vec<NetworkPeer> {
vec![]
}
async fn dial_tcp(&self, addr: &str, port: u16) -> Result<TcpStream, NetworkError> {
let target = format!("{addr}:{port}");
TcpStream::connect(&target)
.await
.map_err(|e| NetworkError::DialFailed(format!("mock dial {target}: {e}")))
}
async fn listen_tcp(&self, port: u16) -> Result<NetworkTcpListener, NetworkError> {
let listener = tokio::net::TcpListener::bind(format!("127.0.0.1:{port}"))
.await
.map_err(|e| NetworkError::ListenFailed(format!("mock listen :{port}: {e}")))?;
let actual_port = listener.local_addr().unwrap().port();
let (tx, rx) = mpsc::channel::<IncomingConnection>(64);
tokio::spawn(async move {
loop {
match listener.accept().await {
Ok((stream, addr)) => {
let conn = IncomingConnection {
stream,
remote_addr: addr.to_string(),
remote_identity: String::new(),
port: actual_port,
};
if tx.send(conn).await.is_err() {
break;
}
}
Err(e) => {
tracing::debug!("mock listener error: {e}");
break;
}
}
}
});
Ok(NetworkTcpListener {
port: actual_port,
incoming: rx,
})
}
async fn unlisten_tcp(&self, _port: u16) -> Result<(), NetworkError> {
Ok(())
}
async fn bind_udp(&self, _port: u16) -> Result<NetworkUdpSocket, NetworkError> {
Err(NetworkError::Internal("mock: UDP not supported".into()))
}
async fn ping(&self, _addr: &str) -> Result<PingResult, NetworkError> {
Ok(PingResult {
latency: Duration::from_millis(1),
connection: "direct".to_string(),
peer_addr: None,
})
}
async fn health(&self) -> HealthInfo {
HealthInfo {
state: "running".to_string(),
healthy: true,
..Default::default()
}
}
}
fn make_network_peer(id: &str, ip: &str) -> NetworkPeer {
NetworkPeer {
id: id.to_string(),
hostname: format!("host-{id}"),
ip: ip.parse().unwrap(),
online: true,
cur_addr: Some(format!("{ip}:41641")),
relay: None,
os: Some("linux".to_string()),
last_seen: Some("2026-03-25T12:00:00Z".to_string()),
key_expiry: None,
dns_name: None,
}
}
async fn random_port() -> u16 {
let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
l.local_addr().unwrap().port()
}
fn ws_config(port: u16) -> WsConfig {
WsConfig {
port,
ping_interval: Duration::from_secs(300), pong_timeout: Duration::from_secs(300),
..Default::default()
}
}
fn make_loopback_peer(id: &str) -> NetworkPeer {
NetworkPeer {
id: id.to_string(),
hostname: format!("truffle-test-{id}"),
ip: "127.0.0.1".parse().unwrap(),
online: true,
cur_addr: Some("127.0.0.1:41641".to_string()),
relay: None,
os: None,
last_seen: None,
key_expiry: None,
dns_name: None,
}
}
fn build_registry(
id: &str,
port: u16,
) -> (
PeerRegistry<MockNetworkProvider>,
broadcast::Sender<NetworkPeerEvent>,
) {
build_registry_with_app("test", id, port)
}
fn build_registry_with_app(
app_id: &str,
id: &str,
port: u16,
) -> (
PeerRegistry<MockNetworkProvider>,
broadcast::Sender<NetworkPeerEvent>,
) {
build_registry_with_options(app_id, id, port, PeerRegistryOptions::default())
}
fn build_registry_with_options(
app_id: &str,
id: &str,
port: u16,
options: PeerRegistryOptions,
) -> (
PeerRegistry<MockNetworkProvider>,
broadcast::Sender<NetworkPeerEvent>,
) {
let provider = MockNetworkProvider::new_with_app(app_id, id);
let event_sender = provider.event_sender();
let network = Arc::new(provider);
let ws_transport = Arc::new(WebSocketTransport::new(network.clone(), ws_config(port)));
let registry = PeerRegistry::with_options(network, ws_transport, options);
(registry, event_sender)
}
#[tokio::test]
async fn test_peers_from_network_events() {
let port = random_port().await;
let (registry, event_sender) = build_registry("node-a", port);
registry.start().await;
let peer1 = make_network_peer("peer-1", "100.64.0.1");
let peer2 = make_network_peer("peer-2", "100.64.0.2");
event_sender.send(NetworkPeerEvent::Joined(peer1)).unwrap();
event_sender.send(NetworkPeerEvent::Joined(peer2)).unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
let peers = registry.peers().await;
assert_eq!(peers.len(), 2, "should have 2 peers");
let ids: Vec<String> = peers.iter().map(|p| p.id.clone()).collect();
assert!(ids.contains(&"peer-1".to_string()));
assert!(ids.contains(&"peer-2".to_string()));
}
#[tokio::test]
async fn test_peer_left_removes() {
let port = random_port().await;
let (registry, event_sender) = build_registry("node-a", port);
registry.start().await;
let peer = make_network_peer("peer-1", "100.64.0.1");
event_sender.send(NetworkPeerEvent::Joined(peer)).unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(registry.peers().await.len(), 1);
event_sender
.send(NetworkPeerEvent::Left("peer-1".to_string()))
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(registry.peers().await.len(), 0, "peer should be removed");
}
#[tokio::test]
async fn test_peers_online_without_connection() {
let port = random_port().await;
let (registry, event_sender) = build_registry("node-a", port);
registry.start().await;
let peer = make_network_peer("peer-1", "100.64.0.1");
event_sender.send(NetworkPeerEvent::Joined(peer)).unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
let peers = registry.peers().await;
assert_eq!(peers.len(), 1);
let p = &peers[0];
assert_eq!(p.id, "peer-1");
assert!(p.online, "peer should be online");
assert!(
!p.ws_connected,
"peer should NOT be ws_connected (no WS yet)"
);
}
#[tokio::test]
async fn test_peer_updated_preserves_connected() {
let port = random_port().await;
let (registry, event_sender) = build_registry("node-a", port);
registry.start().await;
let peer = make_network_peer("peer-1", "100.64.0.1");
event_sender.send(NetworkPeerEvent::Joined(peer)).unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
let mut updated_peer = make_network_peer("peer-1", "100.64.0.1");
updated_peer.relay = Some("sfo".to_string());
updated_peer.cur_addr = None;
event_sender
.send(NetworkPeerEvent::Updated(updated_peer))
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
let peers = registry.peers().await;
assert_eq!(peers.len(), 1);
let p = &peers[0];
assert_eq!(p.connection_type, "relay:sfo");
assert!(!p.ws_connected, "ws_connected should be preserved as false");
}
#[tokio::test]
async fn test_peer_event_subscription() {
let port = random_port().await;
let (registry, event_sender) = build_registry("node-a", port);
let mut rx = registry.on_peer_change();
registry.start().await;
let peer = make_network_peer("peer-1", "100.64.0.1");
event_sender.send(NetworkPeerEvent::Joined(peer)).unwrap();
let event = tokio::time::timeout(Duration::from_millis(200), rx.recv())
.await
.expect("should receive event within timeout")
.expect("should not be a recv error");
match event {
PeerEvent::Joined(state) => {
assert_eq!(state.id, "peer-1");
assert!(state.online);
}
other => panic!("expected PeerEvent::Joined, got: {other:?}"),
}
}
#[tokio::test]
async fn test_send_unknown_peer_errors() {
let port = random_port().await;
let (registry, _event_sender) = build_registry("node-a", port);
registry.start().await;
let result = registry.send("nonexistent-peer", b"hello").await;
assert!(result.is_err(), "send to unknown peer should fail");
let err = result.unwrap_err();
assert!(
err.to_string().contains("unknown peer"),
"error should mention unknown peer: {err}"
);
}
#[tokio::test]
async fn test_send_offline_peer_errors() {
let port = random_port().await;
let (registry, event_sender) = build_registry("node-a", port);
registry.start().await;
let mut peer = make_network_peer("peer-1", "100.64.0.1");
peer.online = false;
event_sender.send(NetworkPeerEvent::Joined(peer)).unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
let result = registry.send("peer-1", b"hello").await;
assert!(result.is_err(), "send to offline peer should fail");
let err = result.unwrap_err();
assert!(
err.to_string().contains("offline"),
"error should mention offline: {err}"
);
}
#[tokio::test]
async fn test_send_lazy_connects() {
let server_port = random_port().await;
let (server_registry, _server_es) = build_registry("server", server_port);
let mut server_incoming = server_registry.subscribe();
server_registry.start().await;
let (client_registry, client_es) = build_registry("client", server_port);
client_registry.start().await;
client_es
.send(NetworkPeerEvent::Joined(make_loopback_peer("server")))
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
let peers = client_registry.peers().await;
assert_eq!(peers.len(), 1);
assert!(
!peers[0].ws_connected,
"should not be connected before send"
);
let msg = b"hello from lazy connect";
client_registry.send("server", msg).await.unwrap();
let peers = client_registry.peers().await;
assert!(peers[0].ws_connected, "should be connected after send");
let incoming = tokio::time::timeout(Duration::from_millis(500), server_incoming.recv())
.await
.expect("should receive message within timeout")
.expect("should not be a recv error");
assert_eq!(incoming.from, "client");
assert_eq!(incoming.data, msg);
}
#[tokio::test]
async fn test_send_reuses_connection() {
let server_port = random_port().await;
let (server_registry, _) = build_registry("server", server_port);
let mut server_incoming = server_registry.subscribe();
server_registry.start().await;
let (client_registry, client_es) = build_registry("client", server_port);
let mut client_events = client_registry.on_peer_change();
client_registry.start().await;
client_es
.send(NetworkPeerEvent::Joined(make_loopback_peer("server")))
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
client_registry.send("server", b"msg-1").await.unwrap();
let event = tokio::time::timeout(Duration::from_millis(200), async {
loop {
match client_events.recv().await {
Ok(PeerEvent::WsConnected(_)) => return,
_ => continue,
}
}
})
.await;
assert!(event.is_ok(), "should receive Connected event");
client_registry.send("server", b"msg-2").await.unwrap();
let msg1 = tokio::time::timeout(Duration::from_millis(200), server_incoming.recv())
.await
.unwrap()
.unwrap();
assert_eq!(msg1.data, b"msg-1");
let msg2 = tokio::time::timeout(Duration::from_millis(200), server_incoming.recv())
.await
.unwrap()
.unwrap();
assert_eq!(msg2.data, b"msg-2");
}
#[tokio::test]
async fn test_broadcast_sends_to_all() {
let bcast_port = random_port().await;
let (bcast_reg, _) = build_registry("broadcaster", bcast_port);
bcast_reg.start().await;
let (client1_reg, client1_es) = build_registry("client1", bcast_port);
let mut client1_incoming = client1_reg.subscribe();
client1_reg.start().await;
let (client2_reg, client2_es) = build_registry("client2", bcast_port);
let mut client2_incoming = client2_reg.subscribe();
client2_reg.start().await;
client1_es
.send(NetworkPeerEvent::Joined(make_loopback_peer("broadcaster")))
.unwrap();
client2_es
.send(NetworkPeerEvent::Joined(make_loopback_peer("broadcaster")))
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
client1_reg.send("broadcaster", b"hello1").await.unwrap();
client2_reg.send("broadcaster", b"hello2").await.unwrap();
tokio::time::sleep(Duration::from_millis(200)).await;
bcast_reg.broadcast(b"broadcast-msg").await;
let msg1 = tokio::time::timeout(Duration::from_millis(500), client1_incoming.recv())
.await
.expect("client1 should receive broadcast")
.expect("should not error");
assert_eq!(msg1.data, b"broadcast-msg");
let msg2 = tokio::time::timeout(Duration::from_millis(500), client2_incoming.recv())
.await
.expect("client2 should receive broadcast")
.expect("should not error");
assert_eq!(msg2.data, b"broadcast-msg");
}
#[tokio::test]
async fn test_incoming_message() {
let server_port = random_port().await;
let (server_registry, _) = build_registry("server", server_port);
let mut server_incoming = server_registry.subscribe();
server_registry.start().await;
let (client_registry, client_es) = build_registry("sender", server_port);
client_registry.start().await;
client_es
.send(NetworkPeerEvent::Joined(make_loopback_peer("server")))
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
let payload = b"test payload 12345";
client_registry.send("server", payload).await.unwrap();
let msg = tokio::time::timeout(Duration::from_millis(500), server_incoming.recv())
.await
.expect("should receive within timeout")
.expect("should not error");
assert_eq!(msg.from, "sender");
assert_eq!(msg.data, payload);
}
#[tokio::test]
async fn test_disconnect_reconnect() {
let server_port = random_port().await;
let (server_registry, _) = build_registry("server", server_port);
let mut server_incoming = server_registry.subscribe();
server_registry.start().await;
let (client_registry, client_es) = build_registry("client", server_port);
client_registry.start().await;
client_es
.send(NetworkPeerEvent::Joined(make_loopback_peer("server")))
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
client_registry.send("server", b"msg-1").await.unwrap();
let msg1 = tokio::time::timeout(Duration::from_millis(500), server_incoming.recv())
.await
.unwrap()
.unwrap();
assert_eq!(msg1.data, b"msg-1");
client_registry.disconnect("server").await;
tokio::time::sleep(Duration::from_millis(100)).await;
let peers = client_registry.peers().await;
let server_peer = peers.iter().find(|p| p.id == "server").unwrap();
assert!(
!server_peer.ws_connected,
"should be disconnected after disconnect()"
);
client_registry.send("server", b"msg-2").await.unwrap();
let peers = client_registry.peers().await;
let server_peer = peers.iter().find(|p| p.id == "server").unwrap();
assert!(
server_peer.ws_connected,
"should be reconnected after second send"
);
let msg2 = tokio::time::timeout(Duration::from_millis(500), server_incoming.recv())
.await
.unwrap()
.unwrap();
assert_eq!(msg2.data, b"msg-2");
}
#[tokio::test]
async fn test_peer_left_closes_ws_connection() {
let server_port = random_port().await;
let (server_registry, _) = build_registry("server", server_port);
server_registry.start().await;
let (client_registry, client_es) = build_registry("client", server_port);
let mut client_events = client_registry.on_peer_change();
client_registry.start().await;
client_es
.send(NetworkPeerEvent::Joined(make_loopback_peer("server")))
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
client_registry.send("server", b"hello").await.unwrap();
tokio::time::timeout(Duration::from_millis(200), async {
loop {
if let Ok(PeerEvent::WsConnected(id)) = client_events.recv().await {
if id == "server" {
return;
}
}
}
})
.await
.expect("should receive Connected event");
let peers = client_registry.peers().await;
assert!(peers[0].ws_connected, "should be connected before Left");
client_es
.send(NetworkPeerEvent::Left("server".to_string()))
.unwrap();
let mut got_disconnected = false;
let mut got_left = false;
let _ = tokio::time::timeout(Duration::from_millis(500), async {
loop {
match client_events.recv().await {
Ok(PeerEvent::WsDisconnected(id)) if id == "server" => {
got_disconnected = true;
}
Ok(PeerEvent::Left(state)) if state.id == "server" => {
assert!(!state.online && !state.ws_connected);
got_left = true;
return;
}
_ => continue,
}
}
})
.await;
assert!(got_disconnected, "should emit Disconnected before Left");
assert!(got_left, "should emit Left");
let peers = client_registry.peers().await;
assert!(
peers.iter().all(|p| p.id != "server"),
"peer should be removed after Left"
);
}
#[test]
fn test_reconnect_backoff_basic() {
let mut backoff = ReconnectBackoff::new();
assert!(backoff.should_retry().is_some());
backoff.failure();
assert!(backoff.should_retry().is_none());
let wait = backoff.retry_after();
assert!(wait > Duration::ZERO, "should have non-zero retry_after");
assert!(
wait <= Duration::from_millis(100),
"first backoff should be <= 100ms"
);
backoff.failure();
let wait2 = backoff.retry_after();
assert!(wait2 > Duration::ZERO);
}
#[test]
fn test_reconnect_backoff_resets_on_success() {
let mut backoff = ReconnectBackoff::new();
backoff.failure();
backoff.failure();
backoff.failure();
assert!(
backoff.should_retry().is_none(),
"should be in backoff after 3 failures"
);
backoff.success();
assert!(
backoff.should_retry().is_some(),
"should allow retry after success"
);
assert_eq!(
backoff.retry_after(),
Duration::ZERO,
"retry_after should be zero after success"
);
}
#[tokio::test]
async fn test_broadcast_with_zero_connections() {
let port = random_port().await;
let (registry, _) = build_registry("node-a", port);
registry.start().await;
registry.broadcast(b"hello nobody").await;
let peers = registry.peers().await;
assert!(peers.is_empty(), "should have no peers");
}
#[tokio::test]
async fn test_hello_exchange_populates_identity() {
let server_port = random_port().await;
let (server_registry, server_es) = build_registry("server", server_port);
let mut server_incoming = server_registry.subscribe();
server_registry.start().await;
let (client_registry, client_es) = build_registry("client", server_port);
client_registry.start().await;
server_es
.send(NetworkPeerEvent::Joined(make_loopback_peer("client")))
.unwrap();
client_es
.send(NetworkPeerEvent::Joined(make_loopback_peer("server")))
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
client_registry.send("server", b"hello").await.unwrap();
tokio::time::sleep(Duration::from_millis(150)).await;
let _msg = tokio::time::timeout(Duration::from_millis(500), server_incoming.recv())
.await
.expect("server should receive the inbound frame")
.expect("broadcast channel should not error");
let client_peers = client_registry.peers().await;
let server_view = client_peers
.iter()
.find(|p| p.id == "server")
.expect("client should know about server");
let server_identity = server_view
.identity
.as_ref()
.expect("hello should have landed on the client side");
assert_eq!(server_identity.app_id, "test");
assert_eq!(server_identity.device_id, "dev-server");
assert_eq!(server_identity.device_name, "Test Node server");
assert_eq!(server_identity.tailscale_id, "server");
assert!(!server_identity.os.is_empty(), "os should be populated");
let server_peers = server_registry.peers().await;
let client_view = server_peers
.iter()
.find(|p| p.id == "client")
.expect("server should know about client");
let client_identity = client_view
.identity
.as_ref()
.expect("hello should have landed on the server side");
assert_eq!(client_identity.app_id, "test");
assert_eq!(client_identity.device_id, "dev-client");
assert_eq!(client_identity.device_name, "Test Node client");
assert_eq!(client_identity.tailscale_id, "client");
}
#[tokio::test]
async fn test_hello_exchange_rejects_app_mismatch() {
let server_port = random_port().await;
let (server_registry, _server_es) = build_registry_with_app("chat", "server", server_port);
server_registry.start().await;
let (client_registry, client_es) = build_registry_with_app("playground", "client", server_port);
client_registry.start().await;
client_es
.send(NetworkPeerEvent::Joined(make_loopback_peer("server")))
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
let result = client_registry.send("server", b"unauthorised").await;
assert!(
result.is_err(),
"sending across an app mismatch should fail, got: {result:?}"
);
let err = result.unwrap_err().to_string();
assert!(
err.contains("app mismatch") || err.contains("connect failed"),
"error should indicate the hello failure, got: {err}"
);
let client_peers = client_registry.peers().await;
let server_entry = client_peers
.iter()
.find(|p| p.id == "server")
.expect("peer still known at Layer 3");
assert!(
server_entry.identity.is_none(),
"app mismatch must leave identity unpopulated"
);
}
#[tokio::test]
async fn test_hello_exchange_rejects_malformed_hello() {
use futures_util::{SinkExt, StreamExt};
use tokio_tungstenite::tungstenite::Message;
let server_port = random_port().await;
let (server_registry, _server_es) = build_registry("server", server_port);
server_registry.start().await;
tokio::time::sleep(Duration::from_millis(50)).await;
let (mut ws, _resp) =
tokio_tungstenite::connect_async(format!("ws://127.0.0.1:{server_port}/ws"))
.await
.expect("ws upgrade should succeed");
ws.send(Message::Text(
r#"{"kind":"not_hello","version":2,"identity":{"app_id":"test","device_id":"x","device_name":"y","os":"darwin","tailscale_id":"z"}}"#
.to_string()
.into(),
))
.await
.expect("send should succeed before server close");
let mut saw_close = false;
for _ in 0..10 {
match ws.next().await {
Some(Ok(Message::Close(Some(frame)))) => {
assert_eq!(u16::from(frame.code), super::CLOSE_HELLO_PROTOCOL);
saw_close = true;
break;
}
Some(Ok(_)) => continue,
Some(Err(_)) | None => break,
}
}
assert!(
saw_close,
"server should close with CLOSE_HELLO_PROTOCOL after bogus hello"
);
}
#[tokio::test]
async fn test_hello_exchange_hello_timeout() {
use futures_util::StreamExt;
use tokio_tungstenite::tungstenite::Message;
let server_port = random_port().await;
let (server_registry, _server_es) = build_registry("server", server_port);
server_registry.start().await;
tokio::time::sleep(Duration::from_millis(50)).await;
let (mut ws, _resp) =
tokio_tungstenite::connect_async(format!("ws://127.0.0.1:{server_port}/ws"))
.await
.expect("ws upgrade should succeed");
let mut saw_close = false;
let deadline = tokio::time::Instant::now() + Duration::from_secs(7);
while tokio::time::Instant::now() < deadline {
let remaining = deadline.duration_since(tokio::time::Instant::now());
match tokio::time::timeout(remaining, ws.next()).await {
Ok(Some(Ok(Message::Close(Some(frame))))) => {
assert_eq!(u16::from(frame.code), super::CLOSE_HELLO_PROTOCOL);
saw_close = true;
break;
}
Ok(Some(Ok(_))) => continue,
Ok(Some(Err(_))) | Ok(None) => break,
Err(_) => break,
}
}
assert!(
saw_close,
"server should close with CLOSE_HELLO_PROTOCOL after hello timeout"
);
}
#[tokio::test]
async fn test_rfc022_generation_bumps_on_rejoin() {
let port = random_port().await;
let (registry, es) = build_registry("local", port);
registry.start().await;
es.send(NetworkPeerEvent::Joined(make_loopback_peer("peer-a")))
.unwrap();
tokio::time::sleep(Duration::from_millis(40)).await;
let gen1 = registry
.peers()
.await
.into_iter()
.find(|p| p.id == "peer-a")
.unwrap()
.generation;
assert_eq!(gen1, 1);
es.send(NetworkPeerEvent::Left("peer-a".into())).unwrap();
tokio::time::sleep(Duration::from_millis(40)).await;
es.send(NetworkPeerEvent::Joined(make_loopback_peer("peer-a")))
.unwrap();
tokio::time::sleep(Duration::from_millis(40)).await;
let gen2 = registry
.peers()
.await
.into_iter()
.find(|p| p.id == "peer-a")
.unwrap()
.generation;
assert_eq!(gen2, 2);
assert_ne!(
format_peer_ref("peer-a", gen1),
format_peer_ref("peer-a", gen2)
);
}
#[tokio::test]
async fn test_rfc022_first_wins_duplicate_ulid() {
let port = random_port().await;
let (registry, es) = build_registry("local", port);
registry.start().await;
es.send(NetworkPeerEvent::Joined(make_loopback_peer("a")))
.unwrap();
es.send(NetworkPeerEvent::Joined(make_loopback_peer("b")))
.unwrap();
tokio::time::sleep(Duration::from_millis(40)).await;
let uid = "01J4K9M2Z8AB3RNYQPW6H5TC0X";
let ident = |ts: &str| PeerIdentity {
app_id: "test".into(),
device_id: uid.into(),
device_name: format!("Node {ts}"),
os: "linux".into(),
tailscale_id: ts.into(),
};
assert!(registry.test_stamp_identity("a", ident("a")).await);
assert!(registry.test_stamp_identity("b", ident("b")).await);
let peers = registry.peers().await;
let a = peers.iter().find(|p| p.id == "a").unwrap();
let b = peers.iter().find(|p| p.id == "b").unwrap();
assert!(!a.identity_suppressed);
assert_eq!(a.published_device_id(), Some(uid));
assert!(b.identity_suppressed);
assert!(b.published_device_id().is_none());
assert_eq!(registry.test_by_device(uid).await.as_deref(), Some("a"));
let pa: crate::node::Peer = a.clone().into();
let pb: crate::node::Peer = b.clone().into();
assert_eq!(pa.device_id.as_deref(), Some(uid));
assert!(pb.device_id.is_none());
}
#[tokio::test]
async fn test_rfc022_ghost_retire_on_offline_holder() {
let port = random_port().await;
let (registry, es) = build_registry("local", port);
registry.start().await;
es.send(NetworkPeerEvent::Joined(make_loopback_peer("old")))
.unwrap();
es.send(NetworkPeerEvent::Joined(make_loopback_peer("new")))
.unwrap();
tokio::time::sleep(Duration::from_millis(40)).await;
let uid = "01J4K9M2Z8AB3RNYQPW6H5TC0Y";
let ident = |ts: &str| PeerIdentity {
app_id: "test".into(),
device_id: uid.into(),
device_name: format!("Node {ts}"),
os: "linux".into(),
tailscale_id: ts.into(),
};
assert!(registry.test_stamp_identity("old", ident("old")).await);
{
let mut p = make_loopback_peer("old");
p.online = false;
es.send(NetworkPeerEvent::Updated(p)).unwrap();
}
tokio::time::sleep(Duration::from_millis(40)).await;
assert!(registry.test_stamp_identity("new", ident("new")).await);
let peers = registry.peers().await;
assert!(
peers.iter().all(|p| p.id != "old"),
"ghost should be retired"
);
let n = peers.iter().find(|p| p.id == "new").unwrap();
assert_eq!(n.published_device_id(), Some(uid));
assert_eq!(registry.test_by_device(uid).await.as_deref(), Some("new"));
}
#[tokio::test]
async fn test_rfc022_inbound_from_is_tailscale_id() {
let server_port = random_port().await;
let (server_registry, server_es) = build_registry("server", server_port);
let mut server_incoming = server_registry.subscribe();
server_registry.start().await;
let (client_registry, client_es) = build_registry("client", server_port);
client_registry.start().await;
server_es
.send(NetworkPeerEvent::Joined(make_loopback_peer("client")))
.unwrap();
client_es
.send(NetworkPeerEvent::Joined(make_loopback_peer("server")))
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
client_registry.send("server", b"ping").await.unwrap();
let msg = tokio::time::timeout(Duration::from_millis(800), server_incoming.recv())
.await
.expect("timeout")
.expect("recv");
assert_eq!(msg.from, "client");
}
#[tokio::test]
async fn test_rfc022_eager_identity_without_app_send() {
let server_port = random_port().await;
let (server_registry, server_es) = build_registry("server", server_port);
server_registry.start().await;
let (client_registry, client_es) = build_registry("client", server_port);
client_registry.start().await;
server_es
.send(NetworkPeerEvent::Joined(make_loopback_peer("client")))
.unwrap();
client_es
.send(NetworkPeerEvent::Joined(make_loopback_peer("server")))
.unwrap();
let deadline = tokio::time::Instant::now() + Duration::from_secs(3);
let mut saw = false;
while tokio::time::Instant::now() < deadline {
let peers = client_registry.peers().await;
if let Some(p) = peers.iter().find(|p| p.id == "server") {
if p.published_device_id().is_some() {
saw = true;
break;
}
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
assert!(
saw,
"eager identity should populate device_id without app send"
);
let peers = server_registry.peers().await;
let client = peers.iter().find(|p| p.id == "client");
if client.and_then(|p| p.published_device_id()).is_none() {
let deadline = tokio::time::Instant::now() + Duration::from_secs(2);
loop {
let peers = server_registry.peers().await;
if peers
.iter()
.find(|p| p.id == "client")
.and_then(|p| p.published_device_id())
.is_some()
{
break;
}
if tokio::time::Instant::now() >= deadline {
break;
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
}
let peers = server_registry.peers().await;
let client = peers.iter().find(|p| p.id == "client").unwrap();
assert!(
client.published_device_id().is_some(),
"server should see client identity after eager exchange"
);
}
#[tokio::test]
async fn test_rfc022_eager_identity_disabled_no_auto_hello() {
let server_port = random_port().await;
let opts = PeerRegistryOptions {
eager_identity: false,
..Default::default()
};
let (server_registry, _server_es) =
build_registry_with_options("test", "server", server_port, opts.clone());
server_registry.start().await;
let (client_registry, client_es) =
build_registry_with_options("test", "client", server_port, opts);
client_registry.start().await;
client_es
.send(NetworkPeerEvent::Joined(make_loopback_peer("server")))
.unwrap();
tokio::time::sleep(Duration::from_millis(400)).await;
let peers = client_registry.peers().await;
let server = peers.iter().find(|p| p.id == "server").unwrap();
assert!(
server.published_device_id().is_none(),
"with eager_identity=false, no auto hello"
);
assert!(!server.ws_connected);
}
#[test]
fn test_parse_peer_ref_forms() {
assert_eq!(parse_peer_ref("nABC123:3"), Some(("nABC123", 3)));
assert_eq!(parse_peer_ref("nABC123"), None); assert_eq!(parse_peer_ref("fd7a:115c::1"), None); assert_eq!(parse_peer_ref(":3"), None); assert_eq!(parse_peer_ref("nABC:"), None); assert_eq!(parse_peer_ref("nABC:3a"), None); }
#[tokio::test]
async fn test_rfc022_routing_ignores_suppressed_and_checks_generation() {
let port = random_port().await;
let (registry, es) = build_registry("local", port);
registry.start().await;
es.send(NetworkPeerEvent::Joined(make_loopback_peer("a")))
.unwrap();
es.send(NetworkPeerEvent::Joined(make_loopback_peer("b")))
.unwrap();
tokio::time::sleep(Duration::from_millis(40)).await;
let uid = "01J4K9M2Z8AB3RNYQPW6H5TC0X";
let ident = |ts: &str, name: &str| PeerIdentity {
app_id: "test".into(),
device_id: uid.into(),
device_name: name.into(),
os: "linux".into(),
tailscale_id: ts.into(),
};
assert!(
registry
.test_stamp_identity("a", ident("a", "Holder"))
.await
);
assert!(
registry
.test_stamp_identity("b", ident("b", "Claimant"))
.await
);
assert_eq!(registry.resolve_routing_key(uid).await.unwrap(), "a");
assert!(matches!(
registry.resolve_routing_key("Claimant").await,
Err(SessionError::UnknownPeer(_))
));
assert_eq!(registry.resolve_routing_key("a:1").await.unwrap(), "a");
assert!(matches!(
registry.resolve_routing_key("a:9").await,
Err(SessionError::PeerGone(_))
));
assert!(matches!(
registry.resolve_routing_key("gone:1").await,
Err(SessionError::PeerGone(_))
));
}
#[tokio::test]
async fn test_rfc022_rehello_emits_no_duplicate_identity() {
let port = random_port().await;
let (registry, es) = build_registry("local", port);
let mut events = registry.on_peer_change();
registry.start().await;
es.send(NetworkPeerEvent::Joined(make_loopback_peer("a")))
.unwrap();
tokio::time::sleep(Duration::from_millis(40)).await;
let ident = |name: &str| PeerIdentity {
app_id: "test".into(),
device_id: "01J4K9M2Z8AB3RNYQPW6H5TC0X".into(),
device_name: name.into(),
os: "linux".into(),
tailscale_id: "a".into(),
};
assert!(registry.test_stamp_identity("a", ident("Alice")).await);
assert!(registry.test_stamp_identity("a", ident("Alice")).await);
assert!(registry.test_stamp_identity("a", ident("Alice II")).await);
let mut identity_events = 0;
let mut updated_with_identity = 0;
while let Ok(ev) = events.try_recv() {
match ev {
PeerEvent::Identity(_) => identity_events += 1,
PeerEvent::Updated(s) if s.identity.is_some() => updated_with_identity += 1,
_ => {}
}
}
assert_eq!(identity_events, 1, "re-hello must not re-emit identity");
assert_eq!(updated_with_identity, 1, "rename surfaces as updated");
}
#[test]
fn test_eager_jitter_delay_bounded_and_stable() {
use super::eager_jitter_delay;
assert_eq!(eager_jitter_delay("server", 0), Duration::ZERO);
assert_eq!(eager_jitter_delay("anything", 0), Duration::ZERO);
let window = 250;
for id in ["server", "client", "peer-a", "peer-b", "01J4K9M2Z8AB"] {
let d = eager_jitter_delay(id, window);
assert!(
d < Duration::from_millis(window),
"delay {d:?} for {id} must be < {window}ms window"
);
assert_eq!(
d,
eager_jitter_delay(id, window),
"delay must be stable per peer id"
);
}
let distinct: std::collections::HashSet<Duration> = (0..32)
.map(|i| eager_jitter_delay(&format!("peer-{i}"), window))
.collect();
assert!(
distinct.len() > 1,
"jitter should spread peers across the window, got {} distinct value(s)",
distinct.len()
);
}
use proptest::prelude::*;
proptest! {
#[test]
fn parse_peer_ref_never_panics(s in ".{0,128}") {
let _ = super::parse_peer_ref(&s);
}
#[test]
fn peer_ref_roundtrip(id in "[A-Za-z0-9_.-]{1,32}", generation in 0u64..1_000_000) {
let formatted = super::format_peer_ref(&id, generation);
prop_assert_eq!(super::parse_peer_ref(&formatted), Some((id.as_str(), generation)));
}
}