use core::{net::SocketAddr, time::Duration};
use std::{
collections::{HashMap, HashSet},
sync::{
Arc, RwLock,
atomic::{AtomicBool, Ordering},
},
};
use kameo::{
actor::ActorRef,
message::{Context, Message},
};
use tokio::{sync::Mutex, task::JoinSet};
use ts_keys::{DiscoPublicKey, NodePublicKey};
use ts_magicsock::{BindingVerifier, DirectTransport, MagicSock, SelfEndpoint};
use ts_transport::{
BatchRecvIter, PeerId, PeerLookup, UnderlayTransport, UnderlayTransportExt, UnderlayTransportId,
};
use crate::{
Env, Error,
dataplane::{
DatapathActivity, DataplaneActor, NewUnderlayTransport, UnderlayFromDataplane,
UnderlayToDataplane,
},
multiderp::{self, Multiderp},
peer_tracker::{DiscoKeyMatch, PeerDb, PeerState},
};
#[derive(Debug, Clone, Copy)]
pub struct DiscoKeyObserved {
pub peer: PeerId,
pub key: DiscoPublicKey,
}
type DiscoKeyObserver = tokio::sync::mpsc::UnboundedSender<DiscoKeyObserved>;
async fn run_disco_key_observer(
mut observed: tokio::sync::mpsc::UnboundedReceiver<DiscoKeyObserved>,
env: Env,
mut shutdown: tokio::sync::watch::Receiver<bool>,
) {
while !*shutdown.borrow() {
let msg = tokio::select! {
_ = shutdown.changed() => break,
msg = observed.recv() => match msg {
Some(msg) => msg,
None => break,
},
};
if let Err(e) = env.publish(msg).await {
tracing::error!(error = %e, "publishing an observed peer disco key");
}
}
}
const PING_INTERVAL: Duration = Duration::from_secs(2);
const STUN_PROBE_INTERVAL_MIN: Duration = Duration::from_secs(20);
const STUN_PROBE_INTERVAL_MAX: Duration = Duration::from_secs(26);
fn stun_probe_delay() -> Duration {
rand::random_range(STUN_PROBE_INTERVAL_MIN..STUN_PROBE_INTERVAL_MAX)
}
const ADVERTISE_INTERVAL: Duration = Duration::from_secs(5);
#[derive(Clone)]
pub struct EndpointAdvertisement {
pub endpoints: Arc<Vec<SelfEndpoint>>,
}
#[cfg(test)]
const BIND_ADDR: &str = "0.0.0.0:0";
async fn bind_unspecified_with_fallback(
ip: core::net::IpAddr,
listen_port: u16,
our_disco: &ts_keys::DiscoPrivateKey,
our_node_key: NodePublicKey,
) -> Result<MagicSock, ts_magicsock::Error> {
let pinned = SocketAddr::new(ip, listen_port);
match MagicSock::bind(pinned, our_disco.clone(), our_node_key).await {
Ok(sock) => Ok(sock),
Err(_) if listen_port != 0 => {
tracing::warn!(
%pinned,
"underlay bind on pinned port failed; falling back to an ephemeral port",
);
MagicSock::bind(SocketAddr::new(ip, 0), our_disco.clone(), our_node_key).await
}
Err(e) => Err(e),
}
}
async fn bind_underlay_addr(
enable_ipv6: bool,
listen_port: u16,
our_disco: ts_keys::DiscoPrivateKey,
our_node_key: NodePublicKey,
) -> Result<MagicSock, ts_magicsock::Error> {
use core::net::{Ipv4Addr, Ipv6Addr};
if !enable_ipv6 {
return bind_unspecified_with_fallback(
Ipv4Addr::UNSPECIFIED.into(),
listen_port,
&our_disco,
our_node_key,
)
.await;
}
match bind_unspecified_with_fallback(
Ipv6Addr::UNSPECIFIED.into(),
listen_port,
&our_disco,
our_node_key,
)
.await
{
Ok(sock) => Ok(sock),
Err(e) => {
tracing::warn!(
error = %e,
"dual-stack underlay bind failed (host IPv6 disabled?); falling back to IPv4-only",
);
bind_unspecified_with_fallback(
Ipv4Addr::UNSPECIFIED.into(),
listen_port,
&our_disco,
our_node_key,
)
.await
}
}
}
pub struct DirectManager {
sock: Option<Arc<MagicSock>>,
transport_id: Option<UnderlayTransportId>,
peer_db: Arc<RwLock<Option<Arc<PeerDb>>>>,
multiderp: ActorRef<Multiderp>,
stun_cursor: Arc<Mutex<usize>>,
force_background_stun: Arc<AtomicBool>,
#[allow(dead_code)]
tasks: JoinSet<()>,
}
#[kameo::messages]
impl DirectManager {
#[message]
pub fn direct_transport_id(&self) -> Option<UnderlayTransportId> {
self.transport_id
}
#[message]
pub fn best_addrs(&self, ids: Vec<PeerId>) -> HashMap<PeerId, SocketAddr> {
let mut addrs = HashMap::new();
let Some(sock) = self.sock.as_ref() else {
return addrs;
};
let db = poisoned_read(&self.peer_db);
let Some(db) = db.as_ref() else {
return addrs;
};
for id in ids {
let Some((_, node)) = db.get(&id) else {
continue;
};
let Some(disco) = node.disco_key else {
continue;
};
if let Some(addr) = sock.best_addr(&disco) {
addrs.insert(id, addr);
}
}
addrs
}
#[message]
pub fn direct_path_latency(&self, disco: DiscoPublicKey) -> Option<(SocketAddr, Duration)> {
self.sock.as_ref()?.best_addr_and_latency(&disco)
}
#[message]
pub fn sock_handle(&self) -> Option<Arc<MagicSock>> {
self.sock.clone()
}
#[message]
pub fn peers_with_direct_path(&self, ids: Vec<PeerId>) -> HashSet<PeerId> {
self.best_addrs(ids).into_keys().collect()
}
#[message]
pub async fn rebind(&self) -> Result<(), ts_magicsock::Error> {
match self.sock.as_ref() {
Some(sock) => sock.rebind().await,
None => Ok(()),
}
}
#[message]
pub async fn rebind_and_reprobe(&self) -> Result<(), ts_magicsock::Error> {
let Some(sock) = self.sock.as_ref() else {
return Ok(());
};
sock.rebind().await?;
if let Err(e) = sock.send_pings().await {
tracing::trace!(error = %e, "rebind-and-reprobe: re-ping after rebind");
}
self.stun_sweep_once(sock).await;
Ok(())
}
#[message]
pub async fn re_stun(&self) -> Result<(), ts_magicsock::Error> {
let Some(sock) = self.sock.as_ref() else {
return Ok(());
};
self.stun_sweep_once(sock).await;
Ok(())
}
async fn stun_sweep_once(&self, sock: &MagicSock) {
if stun_probe_should_run(&self.peer_db) {
match self.multiderp.ask(multiderp::StunServersV4).await {
Ok((servers,)) => probe_stun_servers_once(sock, &servers, &self.stun_cursor).await,
Err(e) => {
tracing::trace!(error = %e, "stun sweep: querying stun servers");
}
}
}
}
}
fn verify_binding(
peer_db: &RwLock<Option<Arc<PeerDb>>>,
observed: &DiscoKeyObserver,
disco: &DiscoPublicKey,
claimed_node_key: Option<&NodePublicKey>,
) -> bool {
let resolved = {
let db = poisoned_read(peer_db);
let db = db.as_ref().and_then(|db| db.peer_by_known_disco_key(disco));
db.map(|(peer, node, matched)| (peer, node.node_key, matched))
};
let Some((peer, node_key, matched)) = resolved else {
return false;
};
let bound = match claimed_node_key {
Some(claimed) => node_key == *claimed,
None => true,
};
if !bound {
return false;
}
if matched == DiscoKeyMatch::Inactive {
let _ = observed.send(DiscoKeyObserved { peer, key: *disco });
}
true
}
fn poisoned_read(
lock: &RwLock<Option<Arc<PeerDb>>>,
) -> std::sync::RwLockReadGuard<'_, Option<Arc<PeerDb>>> {
lock.read().unwrap_or_else(|poisoned| poisoned.into_inner())
}
fn poisoned_write(
lock: &RwLock<Option<Arc<PeerDb>>>,
) -> std::sync::RwLockWriteGuard<'_, Option<Arc<PeerDb>>> {
lock.write()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
struct DiscoPeerLookup(Arc<RwLock<Option<Arc<PeerDb>>>>);
impl PeerLookup<PeerId, DiscoPublicKey> for DiscoPeerLookup {
fn lookup_key(&self, id: PeerId) -> Option<DiscoPublicKey> {
let db = poisoned_read(&self.0);
let db = db.as_ref()?;
let (_, node) = db.get(&id)?;
node.disco_key
}
}
impl PeerLookup<DiscoPublicKey, PeerId> for DiscoPeerLookup {
fn lookup_key(&self, key: DiscoPublicKey) -> Option<PeerId> {
let db = poisoned_read(&self.0);
let db = db.as_ref()?;
let (id, _, _) = db.peer_by_known_disco_key(&key)?;
Some(id)
}
}
async fn run_direct(
transport: impl UnderlayTransport<PeerKey = PeerId, Error = ts_magicsock::Error>,
mut from_dataplane: UnderlayFromDataplane,
to_dataplane: UnderlayToDataplane,
mut shutdown: tokio::sync::watch::Receiver<bool>,
) {
while !*shutdown.borrow() {
tokio::select! {
_ = shutdown.changed() => break,
from_direct = transport.recv() => {
for ret in from_direct.batch_iter() {
match ret {
Ok((peer_id, pkts)) => {
let pkts = pkts.into_iter().collect::<Vec<_>>();
if to_dataplane.send((peer_id, pkts)).is_err() {
tracing::error!("underlay receive channel closed");
return;
}
}
Err(e) => {
tracing::trace!(error = %e, "ignoring undecodable direct packet");
}
}
}
}
from_net = from_dataplane.recv() => {
let Some(from_net) = from_net else {
tracing::warn!("direct underlay queue closed");
break;
};
if let Err(e) = transport.send([from_net]).await {
tracing::trace!(error = %e, "sending direct packet");
}
}
}
}
}
async fn run_pinger(sock: Arc<MagicSock>, mut shutdown: tokio::sync::watch::Receiver<bool>) {
let mut interval = tokio::time::interval(PING_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
while !*shutdown.borrow() {
tokio::select! {
_ = shutdown.changed() => break,
_ = interval.tick() => {
if let Err(e) = sock.send_pings().await {
tracing::trace!(error = %e, "sending disco pings");
}
}
}
}
}
async fn run_stun_prober(
sock: Arc<MagicSock>,
peer_db: Arc<RwLock<Option<Arc<PeerDb>>>>,
multiderp: ActorRef<Multiderp>,
stun_cursor: Arc<Mutex<usize>>,
activity: DatapathActivity,
force_background_stun: Arc<AtomicBool>,
mut shutdown: tokio::sync::watch::Receiver<bool>,
) {
while !*shutdown.borrow() {
tokio::select! {
_ = shutdown.changed() => break,
_ = tokio::time::sleep(stun_probe_delay()) => {
if !should_do_periodic_restun(&peer_db, &activity, &force_background_stun) {
continue;
}
let servers = match multiderp.ask(multiderp::StunServersV4).await {
Ok((servers,)) => servers,
Err(e) => {
tracing::trace!(error = %e, "querying stun servers from multiderp");
continue;
}
};
probe_stun_servers_once(&sock, &servers, &stun_cursor).await;
}
}
}
}
const SESSION_ACTIVE_TIMEOUT: Duration = Duration::from_secs(45);
fn stun_probe_should_run(peer_db: &RwLock<Option<Arc<PeerDb>>>) -> bool {
let db = poisoned_read(peer_db);
db.as_ref().is_some_and(|db| !db.peers().is_empty())
}
fn should_do_periodic_restun(
peer_db: &RwLock<Option<Arc<PeerDb>>>,
activity: &DatapathActivity,
force_background_stun: &AtomicBool,
) -> bool {
if !stun_probe_should_run(peer_db) {
return false;
}
if activity.idle_duration() > SESSION_ACTIVE_TIMEOUT {
return force_background_stun.load(Ordering::Relaxed);
}
true
}
async fn probe_stun_servers_once(sock: &MagicSock, servers: &[SocketAddr], cursor: &Mutex<usize>) {
if sock.only_tcp_443() {
return;
}
let mut cursor = cursor.lock().await;
let (window, next) = stun_probe_window(servers, *cursor, sock.stun_in_flight_remaining());
for s in window {
if let Err(e) = sock.send_stun_request(s).await {
tracing::trace!(error = %e, server = %s, "sending stun binding request");
}
}
*cursor = next;
}
fn stun_probe_window(
servers: &[SocketAddr],
cursor: usize,
budget: usize,
) -> (Vec<SocketAddr>, usize) {
if servers.is_empty() {
return (Vec::new(), 0);
}
let start = cursor % servers.len();
let take = budget.min(servers.len());
let window = (0..take)
.map(|i| servers[(start + i) % servers.len()])
.collect();
(window, (start + take) % servers.len())
}
async fn run_advertiser(
sock: Arc<MagicSock>,
env: Env,
mut shutdown: tokio::sync::watch::Receiver<bool>,
) {
let mut interval = tokio::time::interval(ADVERTISE_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
let mut last: Vec<SelfEndpoint> = Vec::new();
while !*shutdown.borrow() {
tokio::select! {
_ = shutdown.changed() => break,
_ = interval.tick() => {
let mut eps = sock.self_endpoints();
eps.sort_by_key(|e| (e.addr, e.ty as u8));
if eps == last {
continue;
}
last = eps.clone();
if let Err(e) = env
.publish(EndpointAdvertisement {
endpoints: Arc::new(eps),
})
.await
{
tracing::error!(error = %e, "publishing endpoint advertisement");
}
}
}
}
}
async fn run_call_me_maybe(
sock: Arc<MagicSock>,
peer_db: Arc<RwLock<Option<Arc<PeerDb>>>>,
multiderp: ActorRef<Multiderp>,
mut shutdown: tokio::sync::watch::Receiver<bool>,
) {
let mut interval = tokio::time::interval(ADVERTISE_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
while !*shutdown.borrow() {
tokio::select! {
_ = shutdown.changed() => break,
_ = interval.tick() => {
let have_reflexive = sock
.self_endpoints()
.iter()
.any(|e| e.ty == ts_magicsock::SelfEndpointType::Stun);
if !have_reflexive {
continue;
}
let targets: Vec<(ts_keys::NodePublicKey, DiscoPublicKey, Option<ts_derp::RegionId>)> = {
let db = poisoned_read(&peer_db);
let Some(db) = db.as_ref() else { continue; };
db.peers()
.values()
.filter_map(|node| {
let disco = node.disco_key?;
if sock.best_addr(&disco).is_some() {
return None;
}
Some((node.node_key, disco, node.derp_region))
})
.collect()
};
for (node_key, disco, netmap_region) in targets {
let region = match netmap_region {
Some(region) => Some(region),
None => {
match multiderp.ask(multiderp::RegionForNode { node: node_key }).await {
Ok(region) => region,
Err(e) => {
tracing::trace!(error = %e, "inferring call-me-maybe relay region");
None
}
}
}
};
let Some(region) = region else {
continue;
};
let frame = match sock.seal_call_me_maybe(&disco) {
Ok(frame) => frame,
Err(e) => {
tracing::trace!(error = %e, "sealing call-me-maybe");
continue;
}
};
if let Err(e) = multiderp
.tell(multiderp::SendDisco {
peer: node_key,
region,
frame,
})
.await
{
tracing::trace!(error = %e, "relaying call-me-maybe to multiderp");
}
}
}
}
}
}
impl kameo::Actor for DirectManager {
type Args = (Env, ActorRef<DataplaneActor>, ActorRef<Multiderp>);
type Error = Error;
async fn on_start(
(env, dataplane, multiderp): Self::Args,
slf: ActorRef<Self>,
) -> Result<Self, Self::Error> {
env.subscribe::<Arc<PeerState>>(&slf).await?;
env.subscribe::<Arc<ts_control::StateUpdate>>(&slf).await?;
let peer_db: Arc<RwLock<Option<Arc<PeerDb>>>> = Default::default();
let force_background_stun = Arc::new(AtomicBool::new(false));
let stun_cursor: Arc<Mutex<usize>> = Default::default();
let mut tasks = JoinSet::new();
let (observed_tx, observed_rx) = tokio::sync::mpsc::unbounded_channel();
tasks.spawn(run_disco_key_observer(
observed_rx,
env.clone(),
env.shutdown.clone(),
));
let verifier_db = peer_db.clone();
let binding_verifier: BindingVerifier = Arc::new(move |disco, claimed_node_key| {
verify_binding(&verifier_db, &observed_tx, disco, claimed_node_key)
});
let sock = match bind_underlay_addr(
env.enable_ipv6,
env.wireguard_listen_port.unwrap_or(0),
env.keys.disco_keys.private.clone(),
env.keys.node_keys.public,
)
.await
{
Ok(sock) => Arc::new(
sock.with_enable_ipv6(env.enable_ipv6)
.with_binding_verifier(binding_verifier),
),
Err(e) => {
tracing::error!(
error = %e,
enable_ipv6 = env.enable_ipv6,
"direct underlay udp bind failed; direct manager inert, staying DERP-only",
);
return Ok(Self {
sock: None,
transport_id: None,
peer_db,
multiderp,
stun_cursor,
force_background_stun,
tasks,
});
}
};
let (transport_id, from_dataplane, to_dataplane) =
dataplane.ask(NewUnderlayTransport).await?;
let transport =
DirectTransport::new(sock.clone()).with_key_lookup(DiscoPeerLookup(peer_db.clone()));
tasks.spawn(run_direct(
transport,
from_dataplane,
to_dataplane,
env.shutdown.clone(),
));
tasks.spawn(run_pinger(sock.clone(), env.shutdown.clone()));
tasks.spawn(run_advertiser(
sock.clone(),
env.clone(),
env.shutdown.clone(),
));
tasks.spawn(run_stun_prober(
sock.clone(),
peer_db.clone(),
multiderp.clone(),
stun_cursor.clone(),
env.datapath_activity.clone(),
force_background_stun.clone(),
env.shutdown.clone(),
));
if let Err(e) = multiderp
.tell(multiderp::SetDirectSock { sock: sock.clone() })
.await
{
tracing::warn!(error = %e, "could not install direct socket on multiderp");
}
let multiderp_for_field = multiderp.clone();
tasks.spawn(run_call_me_maybe(
sock.clone(),
peer_db.clone(),
multiderp,
env.shutdown.clone(),
));
Ok(Self {
sock: Some(sock),
transport_id: Some(transport_id),
peer_db,
multiderp: multiderp_for_field,
stun_cursor,
force_background_stun,
tasks,
})
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct DiscoKeyRotation {
node_key: NodePublicKey,
previous: DiscoPublicKey,
current: DiscoPublicKey,
}
fn disco_key_rotations(previous: Option<&PeerDb>, current: &PeerDb) -> Vec<DiscoKeyRotation> {
let Some(previous) = previous else {
return Vec::new();
};
current
.peers()
.values()
.filter_map(|node| {
let current_key = node.disco_key?;
let (_, was) = previous.get(&node.node_key)?;
let previous_key = was.disco_key?;
(previous_key != current_key).then_some(DiscoKeyRotation {
node_key: node.node_key,
previous: previous_key,
current: current_key,
})
})
.collect()
}
impl Message<Arc<ts_control::StateUpdate>> for DirectManager {
type Reply = ();
async fn handle(
&mut self,
msg: Arc<ts_control::StateUpdate>,
_ctx: &mut Context<Self, Self::Reply>,
) {
apply_self_node_knobs(
msg.node.as_ref(),
&self.force_background_stun,
self.sock.as_deref(),
);
}
}
fn apply_self_node_knobs(
self_node: Option<&ts_control::Node>,
force_background_stun: &AtomicBool,
sock: Option<&MagicSock>,
) {
let Some(self_node) = self_node else {
return;
};
let force = self_node.force_background_stun();
if force_background_stun.swap(force, Ordering::Relaxed) != force {
tracing::info!(
force_background_stun = force,
"control changed the background-STUN override; the periodic sweep's idle stop \
follows it",
);
}
if let Some(sock) = sock {
sock.set_only_tcp_443(self_node.only_tcp_443());
}
}
async fn reprobe_rotated_peers(sock: &MagicSock, rotated: &[DiscoPublicKey]) -> usize {
let mut sent = 0;
for peer in rotated {
match sock.send_pings_to_peer_now(peer).await {
Ok(n) => sent += n,
Err(e) => {
tracing::warn!(
error = %e,
"re-probing a peer whose disco key changed; it stays on DERP until the \
periodic pinger re-confirms a path",
);
}
}
}
sent
}
impl Message<Arc<PeerState>> for DirectManager {
type Reply = ();
async fn handle(&mut self, msg: Arc<PeerState>, _ctx: &mut Context<Self, Self::Reply>) {
if let Some(sock) = self.sock.as_ref() {
sock.set_silent_disco(msg.silent_disco);
let rotations = {
let previous = poisoned_read(&self.peer_db);
disco_key_rotations(previous.as_deref(), &msg.peers)
};
let mut invalidated = Vec::new();
for rotation in rotations {
tracing::info!(
node_key = %rotation.node_key,
previous = %rotation.previous,
current = %rotation.current,
"peer disco key changed; invalidating its trusted direct path",
);
if sock.changed_active_disco(&rotation.previous, &rotation.current) {
invalidated.push(rotation.current);
}
}
let mut live = HashSet::new();
for node in msg.peers.peers().values() {
let Some(disco) = node.disco_key else {
continue;
};
live.insert(disco);
sock.set_netmap_endpoints(disco, node.underlay_addresses.iter().copied());
}
sock.retain_peers(&live);
let pings = reprobe_rotated_peers(sock, &invalidated).await;
if pings > 0 {
tracing::debug!(
pings,
peers = invalidated.len(),
"re-probing rotated peers now rather than on the next pinger tick",
);
}
}
let mut db = poisoned_write(&self.peer_db);
*db = Some(msg.peers.clone());
}
}
#[cfg(test)]
mod tests {
use core::net::Ipv4Addr;
use ts_control::{Node, StableNodeId, TailnetAddress};
use ts_keys::{DiscoPrivateKey, NodePrivateKey};
use super::*;
use crate::peer_tracker::PeerDb;
fn node_with_keys(disco: DiscoPublicKey, node_key: NodePublicKey, stable: &str) -> Node {
Node {
id: 1,
stable_id: StableNodeId(stable.to_string()),
hostname: "peer".to_string(),
user_id: 0,
tailnet: Some("ts.net".to_string()),
tags: vec![],
addresses: vec![
"100.64.0.9/32".parse().unwrap(),
"fd7a::9/128".parse().unwrap(),
],
tailnet_address: TailnetAddress {
ipv4: "100.64.0.9/32".parse().unwrap(),
ipv6: "fd7a::9/128".parse().unwrap(),
},
node_key,
node_key_expiry: None,
expired: false,
online: None,
last_seen: None,
key_signature: vec![],
machine_key: None,
disco_key: Some(disco),
accepted_routes: vec![],
underlay_addresses: vec![],
derp_region: None,
cap: Default::default(),
cap_map: Default::default(),
peerapi_port: None,
peerapi_dns_proxy: false,
is_wireguard_only: false,
exit_node_dns_resolvers: vec![],
peer_relay: false,
ssh_host_keys: vec![],
service_vips: Default::default(),
unsigned_peer_api_only: false,
}
}
fn db_with(node: Node) -> Arc<RwLock<Option<Arc<PeerDb>>>> {
let mut db = PeerDb::default();
db.upsert(&node);
Arc::new(RwLock::new(Some(Arc::new(db))))
}
fn db_with_inactive_key(
node: Node,
inactive: DiscoPublicKey,
) -> Arc<RwLock<Option<Arc<PeerDb>>>> {
let mut db = PeerDb::default();
let id = db.upsert(&node);
db.set_inactive_disco_key(id, Some(inactive));
Arc::new(RwLock::new(Some(Arc::new(db))))
}
fn observer() -> (
DiscoKeyObserver,
tokio::sync::mpsc::UnboundedReceiver<DiscoKeyObserved>,
) {
tokio::sync::mpsc::unbounded_channel()
}
#[test]
fn verify_binding_ping_requires_exact_node_key() {
let disco = DiscoPrivateKey::random().public_key();
let node_key = NodePrivateKey::random().public_key();
let other_key = NodePrivateKey::random().public_key();
let db = db_with(node_with_keys(disco, node_key, "n1"));
let (tx, mut rx) = observer();
assert!(
verify_binding(&db, &tx, &disco, Some(&node_key)),
"correct disco<->node-key binding must be accepted"
);
assert!(
!verify_binding(&db, &tx, &disco, Some(&other_key)),
"a claimed node key that is not the bound one must be rejected"
);
let unknown_disco = DiscoPrivateKey::random().public_key();
assert!(
!verify_binding(&db, &tx, &unknown_disco, Some(&node_key)),
"a disco key not in the netmap must be rejected"
);
let empty: Arc<RwLock<Option<Arc<PeerDb>>>> = Default::default();
assert!(
!verify_binding(&empty, &tx, &disco, Some(&node_key)),
"with no netmap loaded the verifier fails closed"
);
assert!(
rx.try_recv().is_err(),
"a frame under the peer's ACTIVE key reports no active-key switch"
);
}
#[test]
fn verify_binding_call_me_maybe_is_membership_only() {
let disco = DiscoPrivateKey::random().public_key();
let node_key = NodePrivateKey::random().public_key();
let db = db_with(node_with_keys(disco, node_key, "n1"));
let (tx, _rx) = observer();
assert!(
verify_binding(&db, &tx, &disco, None),
"a netmap-member disco key must be accepted for a CallMeMaybe"
);
let stranger = DiscoPrivateKey::random().public_key();
assert!(
!verify_binding(&db, &tx, &stranger, None),
"a non-member disco key must be rejected for a CallMeMaybe"
);
}
#[test]
fn verify_binding_accepts_the_peers_inactive_disco_key_and_reports_it() {
let active = DiscoPrivateKey::random().public_key();
let inactive = DiscoPrivateKey::random().public_key();
let node_key = NodePrivateKey::random().public_key();
let db = db_with_inactive_key(node_with_keys(active, node_key, "n1"), inactive);
let peer = db
.read()
.unwrap()
.as_ref()
.and_then(|db| db.has(&node_key))
.expect("the peer is in the db");
let (tx, mut rx) = observer();
assert!(
verify_binding(&db, &tx, &inactive, Some(&node_key)),
"a ping under the peer's other known disco key must be accepted"
);
let seen = rx.try_recv().expect("the sighting is reported");
assert_eq!(seen.peer, peer);
assert_eq!(seen.key, inactive);
assert!(
verify_binding(&db, &tx, &inactive, None),
"a CallMeMaybe under the peer's other known disco key must be accepted"
);
assert_eq!(rx.try_recv().expect("reported too").key, inactive);
}
#[test]
fn verify_binding_refuses_a_key_in_neither_slot() {
let active = DiscoPrivateKey::random().public_key();
let inactive = DiscoPrivateKey::random().public_key();
let third = DiscoPrivateKey::random().public_key();
let node_key = NodePrivateKey::random().public_key();
let other_node_key = NodePrivateKey::random().public_key();
let db = db_with_inactive_key(node_with_keys(active, node_key, "n1"), inactive);
let (tx, mut rx) = observer();
assert!(
!verify_binding(&db, &tx, &third, Some(&node_key)),
"a disco key in neither slot must be refused even for a known node key"
);
assert!(
!verify_binding(&db, &tx, &third, None),
"and for a CallMeMaybe, which has no node key to check at all"
);
assert!(
!verify_binding(&db, &tx, &inactive, Some(&other_node_key)),
"the inactive key is still bound to its own node key"
);
assert!(
rx.try_recv().is_err(),
"a refused frame must never report an active-key switch"
);
}
#[tokio::test]
async fn probe_stun_servers_once_sends_binding_request() {
let sock = Arc::new(
MagicSock::bind(
BIND_ADDR.parse().unwrap(),
DiscoPrivateKey::random(),
NodePrivateKey::random().public_key(),
)
.await
.unwrap(),
);
let sink = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
let server: SocketAddr = sink.local_addr().unwrap();
let cursor = Mutex::new(0);
probe_stun_servers_once(&sock, &[server], &cursor).await;
let mut buf = [0u8; 64];
let (n, _from) = tokio::time::timeout(Duration::from_secs(2), sink.recv_from(&mut buf))
.await
.expect("a STUN binding request must arrive at the v4 server")
.unwrap();
assert_eq!(
n, 40,
"a STUN Binding Request is the 40-byte SOFTWARE+FINGERPRINT form"
);
assert_eq!(
&buf[0..2],
&0x0001u16.to_be_bytes(),
"message type must be Binding Request (0x0001)"
);
assert_eq!(
&buf[2..4],
&0x0014u16.to_be_bytes(),
"message length must be 0x0014 (the 20 trailing attribute bytes)"
);
assert_eq!(
&buf[4..8],
&0x2112_A442u32.to_be_bytes(),
"the STUN magic cookie must be present at bytes[4..8]"
);
assert_eq!(
&buf[20..22],
&0x8022u16.to_be_bytes(),
"SOFTWARE attribute type"
);
assert_eq!(&buf[24..32], b"tailnode", "SOFTWARE value must be tailnode");
assert_eq!(
&buf[32..34],
&0x8028u16.to_be_bytes(),
"FINGERPRINT attribute type"
);
}
#[tokio::test]
async fn probe_stun_servers_once_stops_at_the_in_flight_budget() {
let sock = Arc::new(
MagicSock::bind(
BIND_ADDR.parse().unwrap(),
DiscoPrivateKey::random(),
NodePrivateKey::random().public_key(),
)
.await
.unwrap(),
);
let budget = sock.stun_in_flight_remaining();
assert!(budget > 0, "a fresh socket must admit some requests");
let mut sinks = Vec::new();
for _ in 0..(budget + 4) {
sinks.push(tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap());
}
let servers: Vec<SocketAddr> = sinks.iter().map(|s| s.local_addr().unwrap()).collect();
assert!(
servers.len() > budget,
"the test needs a server list longer than the in-flight budget"
);
let cursor = Mutex::new(0);
probe_stun_servers_once(&sock, &servers, &cursor).await;
assert_eq!(
sock.stun_in_flight_remaining(),
0,
"the round must use the whole in-flight budget"
);
assert_eq!(
*cursor.lock().await,
budget,
"the cursor must resume at the first server this round could not reach"
);
let (next_round, _) = stun_probe_window(&servers, *cursor.lock().await, budget);
for starved in &servers[budget..] {
assert!(
next_round.contains(starved),
"{starved} was starved this round and must be probed in the next"
);
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn concurrent_stun_sweeps_advance_one_rotation_and_never_rewind() {
const ROUNDS: usize = 25;
let budget = MagicSock::bind(
BIND_ADDR.parse().unwrap(),
DiscoPrivateKey::random(),
NodePrivateKey::random().public_key(),
)
.await
.unwrap()
.stun_in_flight_remaining();
assert!(budget > 0, "a fresh socket must admit some requests");
let mut sinks = Vec::new();
for _ in 0..(budget + 2) {
sinks.push(tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap());
}
let servers: Arc<Vec<SocketAddr>> =
Arc::new(sinks.iter().map(|s| s.local_addr().unwrap()).collect());
for round in 0..ROUNDS {
let sock = Arc::new(
MagicSock::bind(
BIND_ADDR.parse().unwrap(),
DiscoPrivateKey::random(),
NodePrivateKey::random().public_key(),
)
.await
.unwrap(),
);
let cursor: Arc<Mutex<usize>> = Arc::new(Mutex::new(0));
let sweeps: Vec<_> = (0..2)
.map(|_| {
let (sock, servers, cursor) = (sock.clone(), servers.clone(), cursor.clone());
tokio::spawn(async move {
probe_stun_servers_once(&sock, &servers, &cursor).await;
})
})
.collect();
for sweep in sweeps {
sweep.await.unwrap();
}
let parked = *cursor.lock().await;
assert!(
parked >= budget,
"round {round}: the rotation rewound by {} — the late sweep wrote a `next` sized \
off the cursor the first sweep had already moved past",
budget - parked
);
assert_eq!(
parked, budget,
"round {round}: the two sweeps must advance the rotation exactly one round; \
anything past the budget is a window nobody actually probed"
);
assert_eq!(
sock.stun_in_flight_remaining(),
0,
"round {round}: the pair must spend one round's budget between them, not size two \
rounds off the same budget"
);
let (next_round, _) = stun_probe_window(&servers, parked, budget);
for starved in &servers[budget..] {
assert!(
next_round.contains(starved),
"round {round}: {starved} was starved by the overlapping pair and must be \
probed next round"
);
}
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn stun_sweep_waits_for_the_round_ahead_of_it() {
let sock = Arc::new(
MagicSock::bind(
BIND_ADDR.parse().unwrap(),
DiscoPrivateKey::random(),
NodePrivateKey::random().public_key(),
)
.await
.unwrap(),
);
let budget = sock.stun_in_flight_remaining();
let sink = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
let server: SocketAddr = sink.local_addr().unwrap();
let cursor: Arc<Mutex<usize>> = Arc::new(Mutex::new(0));
let round_in_flight = cursor.clone().lock_owned().await;
let sweep = tokio::spawn({
let (sock, cursor) = (sock.clone(), cursor.clone());
async move {
probe_stun_servers_once(&sock, &[server], &cursor).await;
}
});
tokio::time::sleep(Duration::from_millis(300)).await;
assert!(
!sweep.is_finished(),
"a sweep must wait for the round in front of it rather than run beside it"
);
assert_eq!(
sock.stun_in_flight_remaining(),
budget,
"a waiting sweep must not have spent any of the budget the round ahead of it is sizing \
against"
);
drop(round_in_flight);
tokio::time::timeout(Duration::from_secs(2), sweep)
.await
.expect("the sweep must proceed once the cursor is released")
.unwrap();
assert_eq!(
sock.stun_in_flight_remaining(),
budget - 1,
"the released sweep spends exactly its one request"
);
}
#[test]
fn stun_probe_window_rotates_across_rounds() {
let servers: Vec<SocketAddr> = (1..=20u8)
.map(|i| SocketAddr::from((Ipv4Addr::new(192, 0, 2, i), 3478)))
.collect();
let budget = 16;
let (round1, next) = stun_probe_window(&servers, 0, budget);
assert_eq!(round1, servers[..16], "round 1 probes the head of the list");
assert_eq!(
next, 16,
"round 2 resumes at the first server round 1 skipped"
);
let (round2, next) = stun_probe_window(&servers, next, budget);
let expected2: Vec<SocketAddr> = servers[16..]
.iter()
.chain(servers[..12].iter())
.copied()
.collect();
assert_eq!(round2, expected2, "round 2 wraps from the cursor");
assert_eq!(next, 12);
for s in &servers[16..] {
assert!(
round2.contains(s),
"{s} was starved in round 1 and must be probed in round 2"
);
}
for s in &servers {
assert!(
round1.contains(s) || round2.contains(s),
"{s} must be probed within two rounds"
);
}
}
#[test]
fn stun_probe_window_handles_short_lists() {
let (window, next) = stun_probe_window(&[], 7, 16);
assert!(window.is_empty(), "an empty derp map probes nothing");
assert_eq!(next, 0, "an empty list resets the cursor");
let servers: Vec<SocketAddr> = (1..=3u8)
.map(|i| SocketAddr::from((Ipv4Addr::new(192, 0, 2, i), 3478)))
.collect();
let (window, next) = stun_probe_window(&servers, 2, 16);
assert_eq!(
window.len(),
servers.len(),
"a budget wider than the list probes each server exactly once"
);
assert_eq!(window[0], servers[2], "the round starts at the cursor");
for s in &servers {
assert_eq!(
window.iter().filter(|w| *w == s).count(),
1,
"{s} must be probed exactly once per round"
);
}
assert_eq!(next, 2, "a full pass leaves the cursor where it started");
let (window, _) = stun_probe_window(&servers, 99, 1);
assert_eq!(window, vec![servers[99 % 3]]);
}
#[tokio::test]
async fn bind_underlay_addr_v4_default_is_unchanged() {
let sock = bind_underlay_addr(
false,
0,
DiscoPrivateKey::random(),
NodePrivateKey::random().public_key(),
)
.await
.expect("the IPv4 underlay bind must succeed");
let local = sock.local_addr().expect("a bound socket has a local addr");
assert!(
local.is_ipv4(),
"with enable_ipv6 == false the underlay must bind the v4 family, got {local}"
);
assert_eq!(
local.ip(),
"0.0.0.0".parse::<core::net::IpAddr>().unwrap(),
"the v4 default binds the unspecified v4 address"
);
}
#[tokio::test]
async fn bind_underlay_addr_pins_requested_port_when_free() {
let probe = tokio::net::UdpSocket::bind("0.0.0.0:0").await.unwrap();
let want = probe.local_addr().unwrap().port();
drop(probe);
let sock = bind_underlay_addr(
false,
want,
DiscoPrivateKey::random(),
NodePrivateKey::random().public_key(),
)
.await
.expect("pinned-port underlay bind must succeed");
let local = sock.local_addr().expect("a bound socket has a local addr");
assert!(local.is_ipv4(), "still the v4 family, got {local}");
assert_eq!(
local.port(),
want,
"a free pinned port must be bound exactly (got {local})"
);
}
#[tokio::test]
async fn bind_underlay_addr_falls_back_to_ephemeral_when_port_taken() {
let occupier = tokio::net::UdpSocket::bind("0.0.0.0:0").await.unwrap();
let taken = occupier.local_addr().unwrap().port();
let sock = bind_underlay_addr(
false,
taken,
DiscoPrivateKey::random(),
NodePrivateKey::random().public_key(),
)
.await
.expect("a taken pinned port must fall back to ephemeral, never error");
let local = sock.local_addr().expect("a bound socket has a local addr");
assert!(local.is_ipv4(), "still the v4 family, got {local}");
assert_ne!(
local.port(),
taken,
"the occupied port must not be bound; an ephemeral port is used instead"
);
assert_ne!(local.port(), 0, "a real ephemeral port must be assigned");
drop(occupier);
}
#[tokio::test]
async fn bind_underlay_addr_v6_attempts_dual_stack_or_falls_back() {
let sock = bind_underlay_addr(
true,
0,
DiscoPrivateKey::random(),
NodePrivateKey::random().public_key(),
)
.await
.expect("bind must succeed (dual-stack, else inert IPv4 fallback) and never error");
let local = sock.local_addr().expect("a bound socket has a local addr");
match tokio::net::UdpSocket::bind("[::]:0").await {
Ok(_) => assert!(
local.is_ipv6(),
"on a v6-capable host enable_ipv6 == true must bind the v6 (dual-stack) family, \
got {local}"
),
Err(_) => assert!(
local.is_ipv4(),
"on a host that cannot bind v6 the inert fallback must yield a v4 socket, got \
{local}"
),
}
}
#[tokio::test]
async fn probe_stun_servers_once_empty_list_is_noop() {
let sock = Arc::new(
MagicSock::bind(
BIND_ADDR.parse().unwrap(),
DiscoPrivateKey::random(),
NodePrivateKey::random().public_key(),
)
.await
.unwrap(),
);
let cursor = Mutex::new(0);
probe_stun_servers_once(&sock, &[], &cursor).await;
assert_eq!(*cursor.lock().await, 0);
}
#[tokio::test]
async fn only_tcp_443_skips_the_whole_stun_sweep() {
let sock = Arc::new(
MagicSock::bind(
BIND_ADDR.parse().unwrap(),
DiscoPrivateKey::random(),
NodePrivateKey::random().public_key(),
)
.await
.unwrap(),
);
let mut sinks = Vec::new();
for _ in 0..3 {
sinks.push(tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap());
}
let servers: Vec<SocketAddr> = sinks.iter().map(|s| s.local_addr().unwrap()).collect();
let budget = sock.stun_in_flight_remaining();
assert!(
budget > servers.len(),
"the test needs a budget the whole server list fits inside"
);
sock.set_only_tcp_443(true);
let cursor = Mutex::new(0);
probe_stun_servers_once(&sock, &servers, &cursor).await;
assert_eq!(
sock.stun_in_flight_remaining(),
budget,
"under only-tcp-443 no Binding Request may be sent, so no transaction is recorded"
);
assert_eq!(
*cursor.lock().await,
0,
"a skipped round must not advance the rotation — the next real round starts here"
);
sock.set_only_tcp_443(false);
probe_stun_servers_once(&sock, &servers, &cursor).await;
assert_eq!(
sock.stun_in_flight_remaining(),
budget - servers.len(),
"withdrawing the attribute restores STUN with no restart"
);
assert_eq!(
*cursor.lock().await,
0,
"the resumed round wrapped the whole (short) list, parking the cursor back at the head"
);
}
#[tokio::test]
async fn only_tcp_443_rides_the_self_node_to_the_underlay_socket() {
let sock = MagicSock::bind(
BIND_ADDR.parse().unwrap(),
DiscoPrivateKey::random(),
NodePrivateKey::random().public_key(),
)
.await
.unwrap();
let force = AtomicBool::new(false);
let mut self_node = node_with_keys(
DiscoPrivateKey::random().public_key(),
NodePrivateKey::random().public_key(),
"self",
);
apply_self_node_knobs(Some(&self_node), &force, Some(&sock));
assert!(
!sock.only_tcp_443(),
"a self node without the attribute leaves UDP alone"
);
self_node
.cap_map
.insert("only-tcp-443".to_string(), Vec::new());
apply_self_node_knobs(Some(&self_node), &force, Some(&sock));
assert!(
sock.only_tcp_443(),
"the only-tcp-443 cap-map key is what silences the underlay socket"
);
apply_self_node_knobs(None, &force, Some(&sock));
assert!(
sock.only_tcp_443(),
"a netmap with no self node must leave the flag exactly as it was"
);
self_node.cap_map.remove("only-tcp-443");
apply_self_node_knobs(Some(&self_node), &force, Some(&sock));
assert!(
!sock.only_tcp_443(),
"control restating the self node without the attribute restores UDP, with no restart"
);
}
#[test]
fn applying_the_self_node_knobs_without_a_socket_is_a_noop() {
let force = AtomicBool::new(false);
let mut self_node = node_with_keys(
DiscoPrivateKey::random().public_key(),
NodePrivateKey::random().public_key(),
"self",
);
self_node
.cap_map
.insert("only-tcp-443".to_string(), Vec::new());
self_node
.cap_map
.insert("debug-always-stun".to_string(), Vec::new());
apply_self_node_knobs(Some(&self_node), &force, None);
assert!(
force.load(Ordering::Relaxed),
"the background-STUN knob needs no socket and must still be applied"
);
}
#[test]
fn stun_probe_gated_on_peer_presence() {
let empty: Arc<RwLock<Option<Arc<PeerDb>>>> = Default::default();
assert!(
!stun_probe_should_run(&empty),
"with no netmap loaded the prober must not STUN"
);
let no_peers: Arc<RwLock<Option<Arc<PeerDb>>>> =
Arc::new(RwLock::new(Some(Arc::new(PeerDb::default()))));
assert!(
!stun_probe_should_run(&no_peers),
"an empty peer set must not STUN (Go's len(peerSet)==0 stop)"
);
let disco = DiscoPrivateKey::random().public_key();
let node_key = NodePrivateKey::random().public_key();
let with_peer = db_with(node_with_keys(disco, node_key, "n1"));
assert!(
stun_probe_should_run(&with_peer),
"with a peer present the prober resumes STUN"
);
}
#[tokio::test]
async fn re_stun_sweep_gate_matches_periodic_prober() {
let no_peers: Arc<RwLock<Option<Arc<PeerDb>>>> = Default::default();
assert!(
!stun_probe_should_run(&no_peers),
"re_stun must skip the sweep with no peers, like the periodic prober"
);
let disco = DiscoPrivateKey::random().public_key();
let node_key = NodePrivateKey::random().public_key();
let with_peer = db_with(node_with_keys(disco, node_key, "n1"));
assert!(
stun_probe_should_run(&with_peer),
"re_stun sweeps once a peer is present"
);
let sock = Arc::new(
MagicSock::bind(
BIND_ADDR.parse().unwrap(),
DiscoPrivateKey::random(),
NodePrivateKey::random().public_key(),
)
.await
.unwrap(),
);
probe_stun_servers_once(&sock, &[], &Mutex::new(0)).await;
}
#[test]
fn periodic_stun_stops_when_the_datapath_has_been_idle_past_the_session_timeout() {
let disco = DiscoPrivateKey::random().public_key();
let node_key = NodePrivateKey::random().public_key();
let with_peer = db_with(node_with_keys(disco, node_key, "n1"));
let no_override = AtomicBool::new(false);
let busy = DatapathActivity::new();
assert!(
should_do_periodic_restun(&with_peer, &busy, &no_override),
"an active datapath with a peer must keep STUNning"
);
let nearly_idle =
DatapathActivity::idle_for(SESSION_ACTIVE_TIMEOUT - Duration::from_secs(5));
assert!(
should_do_periodic_restun(&with_peer, &nearly_idle, &no_override),
"idle for under sessionActiveTimeout must not stop the sweep"
);
let idle = DatapathActivity::idle_for(SESSION_ACTIVE_TIMEOUT + Duration::from_secs(1));
assert!(
!should_do_periodic_restun(&with_peer, &idle, &no_override),
"a datapath idle past sessionActiveTimeout must stop the periodic sweep"
);
assert!(
stun_probe_should_run(&with_peer),
"the on-demand sweep gate must ignore idleness, as Go's Conn.ReSTUN does"
);
}
#[test]
fn debug_always_stun_overrides_the_idle_stop_and_nothing_else() {
let disco = DiscoPrivateKey::random().public_key();
let node_key = NodePrivateKey::random().public_key();
let with_peer = db_with(node_with_keys(disco, node_key, "n1"));
let idle = DatapathActivity::idle_for(SESSION_ACTIVE_TIMEOUT + Duration::from_secs(1));
let forced = AtomicBool::new(true);
assert!(
should_do_periodic_restun(&with_peer, &idle, &forced),
"debug-always-stun must keep the sweep running on an idle datapath"
);
let no_peers: Arc<RwLock<Option<Arc<PeerDb>>>> = Default::default();
assert!(
!should_do_periodic_restun(&no_peers, &idle, &forced),
"debug-always-stun must not make a peerless node STUN"
);
}
#[test]
fn the_override_is_read_off_the_self_nodes_cap_map() {
let mut self_node = node_with_keys(
DiscoPrivateKey::random().public_key(),
NodePrivateKey::random().public_key(),
"self",
);
assert!(
!self_node.force_background_stun(),
"a self node without the attribute leaves the idle stop in force"
);
self_node
.cap_map
.insert("debug-always-stun".to_string(), vec![]);
assert!(
self_node.force_background_stun(),
"the debug-always-stun cap-map key is what turns the override on"
);
}
#[test]
fn traffic_resuming_restarts_the_sweep_without_a_netmap_change() {
let disco = DiscoPrivateKey::random().public_key();
let node_key = NodePrivateKey::random().public_key();
let with_peer = db_with(node_with_keys(disco, node_key, "n1"));
let no_override = AtomicBool::new(false);
let activity = DatapathActivity::idle_for(SESSION_ACTIVE_TIMEOUT + Duration::from_secs(1));
assert!(
!should_do_periodic_restun(&with_peer, &activity, &no_override),
"precondition: the idle datapath has stopped the sweep"
);
let (up, _down) = crate::dataplane::OverlayToDataplane::for_test(activity.clone());
up.send(vec![ts_packet::PacketMut::from(vec![0u8; 20])])
.expect("the dataplane receiver is still alive");
assert!(
should_do_periodic_restun(&with_peer, &activity, &no_override),
"the sweep must resume on the next tick after traffic returns, with the same netmap"
);
}
#[test]
fn stun_probe_delay_is_jittered_within_go_bounds() {
let mut seen = std::collections::HashSet::new();
for _ in 0..1000 {
let d = stun_probe_delay();
assert!(
d >= STUN_PROBE_INTERVAL_MIN && d < STUN_PROBE_INTERVAL_MAX,
"delay {d:?} out of [20s, 26s)"
);
assert!(
d < Duration::from_secs(30),
"delay {d:?} must stay under the 30s UDP-NAT-timeout ceiling"
);
seen.insert(d);
}
assert!(
seen.len() > 100,
"expected jittered delays, got only {} distinct value(s)",
seen.len()
);
}
fn db_of(nodes: impl IntoIterator<Item = Node>) -> PeerDb {
let mut db = PeerDb::default();
for node in nodes {
db.upsert(&node);
}
db
}
#[test]
fn disco_key_rotation_is_detected_across_snapshots() {
let node_key = NodePrivateKey::random().public_key();
let old = DiscoPrivateKey::random().public_key();
let new = DiscoPrivateKey::random().public_key();
let before = db_of([node_with_keys(old, node_key, "n1")]);
let after = db_of([node_with_keys(new, node_key, "n1")]);
assert_eq!(
disco_key_rotations(Some(&before), &after),
vec![DiscoKeyRotation {
node_key,
previous: old,
current: new,
}],
);
assert!(disco_key_rotations(Some(&after), &after).is_empty());
}
#[test]
fn disco_key_rotations_ignores_non_rotations() {
let node_key = NodePrivateKey::random().public_key();
let disco = DiscoPrivateKey::random().public_key();
let with_key = node_with_keys(disco, node_key, "n1");
let without_key = Node {
disco_key: None,
..with_key.clone()
};
assert!(
disco_key_rotations(None, &db_of([with_key.clone()])).is_empty(),
"the first snapshot has nothing to diff against"
);
assert!(
disco_key_rotations(
Some(&db_of([without_key.clone()])),
&db_of([with_key.clone()])
)
.is_empty(),
"acquiring a first disco key is not a rotation"
);
assert!(
disco_key_rotations(Some(&db_of([with_key.clone()])), &db_of([without_key])).is_empty(),
"losing the disco key leaves nothing to carry over"
);
let rekeyed = Node {
node_key: NodePrivateKey::random().public_key(),
disco_key: Some(DiscoPrivateKey::random().public_key()),
..with_key.clone()
};
assert!(
disco_key_rotations(Some(&db_of([with_key])), &db_of([rekeyed])).is_empty(),
"a node-key rotation is a new endpoint, not a disco-key change on the old one"
);
}
#[test]
fn disco_key_rotations_reports_only_the_peer_that_rotated() {
let rotator = NodePrivateKey::random().public_key();
let steady = NodePrivateKey::random().public_key();
let old = DiscoPrivateKey::random().public_key();
let new = DiscoPrivateKey::random().public_key();
let steady_disco = DiscoPrivateKey::random().public_key();
let before = db_of([
node_with_keys(old, rotator, "n1"),
node_with_keys(steady_disco, steady, "n2"),
]);
let after = db_of([
node_with_keys(new, rotator, "n1"),
node_with_keys(steady_disco, steady, "n2"),
]);
assert_eq!(
disco_key_rotations(Some(&before), &after),
vec![DiscoKeyRotation {
node_key: rotator,
previous: old,
current: new,
}],
);
}
#[tokio::test]
async fn a_rotated_peer_is_reprobed_immediately() {
let sock = MagicSock::bind(
"127.0.0.1:0".parse().unwrap(),
DiscoPrivateKey::random(),
NodePrivateKey::random().public_key(),
)
.await
.expect("the test underlay socket must bind");
let sink = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
let advertised = sink.local_addr().unwrap();
let second = SocketAddr::new(advertised.ip(), advertised.port().wrapping_add(1).max(1));
let old = DiscoPrivateKey::random().public_key();
let new = DiscoPrivateKey::random().public_key();
sock.set_netmap_endpoints(old, [advertised, second]);
assert!(
sock.changed_active_disco(&old, &new),
"precondition: the rotation carried the peer's path state onto the new key"
);
let pinged = reprobe_rotated_peers(&sock, &[new]).await;
assert_eq!(
pinged, 2,
"every candidate the rotation carried over is re-probed straight away, so the peer \
re-confirms a direct path in one round trip instead of waiting out a pinger tick"
);
}
#[tokio::test]
async fn reprobing_nothing_sends_nothing() {
let sock = MagicSock::bind(
"127.0.0.1:0".parse().unwrap(),
DiscoPrivateKey::random(),
NodePrivateKey::random().public_key(),
)
.await
.expect("the test underlay socket must bind");
assert_eq!(reprobe_rotated_peers(&sock, &[]).await, 0);
assert_eq!(
reprobe_rotated_peers(&sock, &[DiscoPrivateKey::random().public_key()]).await,
0,
"a peer with no path state is skipped rather than failing the sweep"
);
}
}