From b422e603c85b3ffdc610219cbaf1b2f8637067a1 Mon Sep 17 00:00:00 2001 From: Luna Yao <40349250+ZnqbuZ@users.noreply.github.com> Date: Mon, 6 Apr 2026 04:56:22 +0200 Subject: [PATCH] node_mgr: unit tests --- easytier/src/dns/node_mgr.rs | 452 +++++++++++++++++++++++++++++++---- 1 file changed, 404 insertions(+), 48 deletions(-) diff --git a/easytier/src/dns/node_mgr.rs b/easytier/src/dns/node_mgr.rs index 4b9eafe1..995c3dc3 100644 --- a/easytier/src/dns/node_mgr.rs +++ b/easytier/src/dns/node_mgr.rs @@ -161,68 +161,106 @@ impl DnsNodeMgrRpc for DnsNodeMgr { mod tests { use super::*; use crate::common::log; + use crate::dns::tests::new_request; use crate::dns::utils::response::ResponseHandle; - use hickory_proto::op::{Message, MessageType, OpCode, Query, ResponseCode}; - use hickory_proto::rr::{rdata, Name, RData, RecordType}; - use hickory_proto::serialize::binary::{BinDecodable, BinEncodable, BinEncoder}; - use hickory_proto::xfer::Protocol; - use hickory_server::authority::MessageRequest; - use hickory_server::server::Request; - use std::net::{Ipv4Addr, SocketAddrV4}; + use crate::proto::common::Url; + use crate::proto::dns::ZoneData; + use hickory_proto::op::{Message, ResponseCode}; + use hickory_proto::rr::{rdata, RData, RecordType}; + use std::net::Ipv4Addr; + use std::str::FromStr; + use tokio::time::{sleep, Duration}; - #[tokio::test] - async fn test_force_insert_node_info_then_catalog_lookup() -> anyhow::Result<()> { - log::tests::init(); - - let mgr = DnsNodeMgr::new(); - - let zone = Zone::try_from(&crate::proto::dns::ZoneData { + fn valid_zone_data(origin: &str, record: &str, forwarders: Vec<&str>) -> ZoneData { + ZoneData { id: Some(Uuid::new_v4().into()), - origin: "catalog.test".to_string(), + origin: origin.to_string(), ttl: 60, - records: vec!["@ IN A 10.20.30.40".to_string()], - forwarders: vec![], - })?; + records: vec![format!("@ IN A {record}")], + forwarders: forwarders + .into_iter() + .map(|f| Url::from_str(f).expect("invalid forwarder")) + .collect(), + } + } - mgr.nodes - .insert( - Uuid::new_v4(), - DnsNodeInfo { - digest: vec![], - zones: vec![zone].into(), - addresses: Default::default(), - listeners: Default::default(), - }, - ) - .await; + fn snapshot_with( + zones: Vec, + addresses: Vec<&str>, + listeners: Vec<&str>, + ) -> DnsSnapshot { + DnsSnapshot { + zones, + addresses: addresses + .into_iter() + .map(|a| Url::from_str(a).expect("invalid address")) + .collect(), + listeners: listeners + .into_iter() + .map(|l| Url::from_str(l).expect("invalid listener")) + .collect(), + } + } - let mut query = Message::new(); - query.set_id(0x1234); - query.set_message_type(MessageType::Query); - query.set_op_code(OpCode::Query); - query.set_recursion_desired(true); - query.add_query(Query::query( - Name::from_ascii("catalog.test.")?, - RecordType::A, - )); + fn heartbeat_with_snapshot(id: Uuid, snapshot: DnsSnapshot) -> HeartbeatRequest { + let mut hb = HeartbeatRequest { + id: Some(id.into()), + ..Default::default() + }; + hb.update(snapshot); + hb + } - let mut request = Vec::new(); - let mut encoder = BinEncoder::new(&mut request); - query.emit(&mut encoder)?; + fn heartbeat_digest_only(id: Uuid, digest: Vec) -> HeartbeatRequest { + HeartbeatRequest { + id: Some(id.into()), + digest, + snapshot: None, + } + } - let request = Request::new( - MessageRequest::from_bytes(&request)?, - SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0).into(), - Protocol::Udp, - ); + fn reset_all_dirty(mgr: &DnsNodeMgr) { + let _ = mgr.dirty.catalog.reset(); + let _ = mgr.dirty.addresses.reset(); + let _ = mgr.dirty.listeners.reset(); + } + async fn send_heartbeat(mgr: &DnsNodeMgr, input: HeartbeatRequest) -> HeartbeatResponse { + DnsNodeMgrRpc::heartbeat(mgr, BaseController::default(), input) + .await + .expect("heartbeat should succeed") + } + + fn ns(s: &str) -> NameServerAddr { + s.parse().expect("invalid nameserver") + } + + async fn lookup_a_record(mgr: &DnsNodeMgr, name: &str) -> anyhow::Result { + let request = new_request(name, RecordType::A)?; let response = ResponseHandle::new(512); let info = mgr.catalog().lookup(&request, None, response.clone()).await; assert_eq!(info.response_code(), ResponseCode::NoError); - let response = response.into_inner().unwrap(); - let message = Message::from_vec(&response)?; + let response = response.into_inner().expect("response should exist"); + Message::from_vec(&response).map_err(Into::into) + } + + #[tokio::test] + async fn catalog_lookup_returns_record_after_snapshot_heartbeat() -> anyhow::Result<()> { + log::tests::init(); + + let mgr = DnsNodeMgr::new(); + let id = Uuid::new_v4(); + let snapshot = snapshot_with( + vec![valid_zone_data("catalog.test", "10.20.30.40", vec![])], + vec![], + vec![], + ); + + let _ = send_heartbeat(&mgr, heartbeat_with_snapshot(id, snapshot)).await; + + let message = lookup_a_record(&mgr, "catalog.test.").await?; assert!(message.answers().iter().any(|record| { matches!( record.data(), @@ -232,4 +270,322 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn catalog_lookup_aggregates_records_from_multiple_nodes() -> anyhow::Result<()> { + let mgr = DnsNodeMgr::new(); + + let snap_a = snapshot_with( + vec![valid_zone_data("node-a.test", "10.11.12.13", vec![])], + vec!["udp://10.0.1.1:53"], + vec![], + ); + let snap_b = snapshot_with( + vec![valid_zone_data("node-b.test", "10.21.22.23", vec![])], + vec!["udp://10.0.2.1:53"], + vec![], + ); + + let _ = send_heartbeat(&mgr, heartbeat_with_snapshot(Uuid::new_v4(), snap_a)).await; + let _ = send_heartbeat(&mgr, heartbeat_with_snapshot(Uuid::new_v4(), snap_b)).await; + + let message_a = lookup_a_record(&mgr, "node-a.test.").await?; + let message_b = lookup_a_record(&mgr, "node-b.test.").await?; + + assert!(message_a.answers().iter().any(|record| { + matches!( + record.data(), + RData::A(addr) if *addr == rdata::a::A(Ipv4Addr::new(10, 11, 12, 13)) + ) + })); + assert!(message_b.answers().iter().any(|record| { + matches!( + record.data(), + RData::A(addr) if *addr == rdata::a::A(Ipv4Addr::new(10, 21, 22, 23)) + ) + })); + + Ok(()) + } + + #[tokio::test] + async fn heartbeat_digest_only_resync_behavior() { + let mgr = DnsNodeMgr::new(); + let id = Uuid::new_v4(); + + let first = send_heartbeat(&mgr, heartbeat_digest_only(id, vec![1, 2, 3])).await; + assert!(first.resync); + + let snapshot = snapshot_with( + vec![valid_zone_data("resync.test", "10.0.0.10", vec![])], + vec!["udp://10.0.0.1:53"], + vec!["udp://10.0.0.2:53"], + ); + let digest = snapshot.digest(); + let full = send_heartbeat(&mgr, heartbeat_with_snapshot(id, snapshot)).await; + assert!(!full.resync); + + let same = send_heartbeat(&mgr, heartbeat_digest_only(id, digest.clone())).await; + assert!(!same.resync); + + let different = send_heartbeat(&mgr, heartbeat_digest_only(id, vec![9, 9, 9])).await; + assert!(different.resync); + } + + #[tokio::test] + async fn heartbeat_with_snapshot_marks_dirty_flags_by_field_changes() { + let mgr = DnsNodeMgr::new(); + let id = Uuid::new_v4(); + + reset_all_dirty(&mgr); + + let first = snapshot_with( + vec![valid_zone_data("dirty.test", "10.0.0.1", vec![])], + vec!["udp://10.10.10.1:53"], + vec!["udp://10.10.10.2:53"], + ); + let _ = send_heartbeat(&mgr, heartbeat_with_snapshot(id, first)).await; + assert!(mgr.dirty.catalog.peek()); + assert!(mgr.dirty.addresses.peek()); + assert!(mgr.dirty.listeners.peek()); + + reset_all_dirty(&mgr); + + let record_changed = snapshot_with( + vec![valid_zone_data("dirty.test", "10.0.0.2", vec![])], + vec!["udp://10.10.10.1:53"], + vec!["udp://10.10.10.2:53"], + ); + let _ = send_heartbeat(&mgr, heartbeat_with_snapshot(id, record_changed)).await; + assert!(mgr.dirty.catalog.peek()); + assert!(!mgr.dirty.addresses.peek()); + assert!(!mgr.dirty.listeners.peek()); + + reset_all_dirty(&mgr); + + let addr_listener_changed = snapshot_with( + vec![valid_zone_data("dirty.test", "10.0.0.2", vec![])], + vec!["udp://10.10.10.10:53"], + vec!["udp://10.10.10.20:53"], + ); + let _ = send_heartbeat(&mgr, heartbeat_with_snapshot(id, addr_listener_changed)).await; + assert!(mgr.dirty.catalog.peek()); + assert!(mgr.dirty.addresses.peek()); + assert!(mgr.dirty.listeners.peek()); + } + + #[tokio::test] + async fn heartbeat_with_same_snapshot_digest_is_noop_for_dirty() { + let mgr = DnsNodeMgr::new(); + let id = Uuid::new_v4(); + + let snapshot = snapshot_with( + vec![valid_zone_data("stable.test", "10.30.40.50", vec![])], + vec!["udp://10.3.0.1:53"], + vec!["udp://10.3.0.2:53"], + ); + let _ = send_heartbeat(&mgr, heartbeat_with_snapshot(id, snapshot.clone())).await; + + reset_all_dirty(&mgr); + + let _ = send_heartbeat(&mgr, heartbeat_with_snapshot(id, snapshot)).await; + assert!(!mgr.dirty.catalog.peek()); + assert!(!mgr.dirty.addresses.peek()); + assert!(!mgr.dirty.listeners.peek()); + } + + #[tokio::test] + async fn heartbeat_missing_id_returns_error() { + let mgr = DnsNodeMgr::new(); + let err = + DnsNodeMgrRpc::heartbeat(&mgr, BaseController::default(), HeartbeatRequest::default()) + .await + .expect_err("missing id should error"); + assert!(err.to_string().contains("missing id")); + } + + #[tokio::test] + async fn iter_addresses_and_listeners_deduplicate_across_multiple_nodes() -> anyhow::Result<()> + { + let mgr = DnsNodeMgr::new(); + let zone_a = Zone::try_from(&valid_zone_data("iter-a.test", "10.1.1.1", vec![]))?; + let zone_b = Zone::try_from(&valid_zone_data("iter-b.test", "10.1.1.2", vec![]))?; + + mgr.nodes + .insert( + Uuid::new_v4(), + DnsNodeInfo { + digest: vec![1], + zones: vec![zone_a].into(), + addresses: [ns("udp://10.100.0.1:53"), ns("udp://10.100.0.2:53")] + .into_iter() + .collect(), + listeners: [ns("udp://10.200.0.1:53")].into_iter().collect(), + }, + ) + .await; + mgr.nodes + .insert( + Uuid::new_v4(), + DnsNodeInfo { + digest: vec![2], + zones: vec![zone_b].into(), + addresses: [ns("udp://10.100.0.2:53"), ns("udp://10.100.0.3:53")] + .into_iter() + .collect(), + listeners: [ns("udp://10.200.0.1:53"), ns("udp://10.200.0.2:53")] + .into_iter() + .collect(), + }, + ) + .await; + + let addresses: HashSet<_> = mgr.iter_addresses().collect(); + let listeners: HashSet<_> = mgr.iter_listeners().collect(); + + assert_eq!(addresses.len(), 3); + assert!(addresses.contains(&ns("udp://10.100.0.1:53"))); + assert!(addresses.contains(&ns("udp://10.100.0.2:53"))); + assert!(addresses.contains(&ns("udp://10.100.0.3:53"))); + + assert_eq!(listeners.len(), 2); + assert!(listeners.contains(&ns("udp://10.200.0.1:53"))); + assert!(listeners.contains(&ns("udp://10.200.0.2:53"))); + + Ok(()) + } + + #[tokio::test] + async fn collect_zones_filters_out_local_forwarders() -> anyhow::Result<()> { + let mgr = DnsNodeMgr::new(); + let zone = Zone::try_from(&valid_zone_data( + "filter-loop.test", + "10.2.3.4", + vec![ + "udp://10.0.0.10:53", + "tcp://10.0.0.11:53", + "udp://1.1.1.1:53", + ], + ))?; + + mgr.nodes + .insert( + Uuid::new_v4(), + DnsNodeInfo { + digest: vec![1], + zones: vec![zone].into(), + addresses: [ns("udp://10.0.0.10:53")].into_iter().collect(), + listeners: [ns("tcp://10.0.0.11:53")].into_iter().collect(), + }, + ) + .await; + + let zones: Vec = mgr.collect_zones().into_iter().map(Into::into).collect(); + let loop_zone = zones + .into_iter() + .find(|z| z.origin.trim_end_matches('.') == "filter-loop.test") + .expect("test zone should exist"); + + let forwarders: HashSet = loop_zone + .forwarders + .iter() + .map(|u| NameServerAddr::try_from(u).expect("forwarder should be valid")) + .collect(); + + assert_eq!(forwarders.len(), 1); + assert!(forwarders.contains(&ns("udp://1.1.1.1:53"))); + + Ok(()) + } + + #[tokio::test] + async fn collect_zones_filters_cross_node_local_forwarders() -> anyhow::Result<()> { + let mgr = DnsNodeMgr::new(); + + let node_a = snapshot_with( + vec![valid_zone_data( + "cross-node-filter.test", + "10.8.8.8", + vec![ + "udp://10.50.0.1:53", + "udp://10.50.0.2:53", + "udp://8.8.8.8:53", + ], + )], + vec!["udp://10.50.0.1:53"], + vec![], + ); + let node_b = snapshot_with( + vec![valid_zone_data( + "cross-node-helper.test", + "10.9.9.9", + vec![], + )], + vec![], + vec!["udp://10.50.0.2:53"], + ); + + let _ = send_heartbeat(&mgr, heartbeat_with_snapshot(Uuid::new_v4(), node_a)).await; + let _ = send_heartbeat(&mgr, heartbeat_with_snapshot(Uuid::new_v4(), node_b)).await; + + let zones: Vec = mgr.collect_zones().into_iter().map(Into::into).collect(); + let zone = zones + .into_iter() + .find(|z| z.origin.trim_end_matches('.') == "cross-node-filter.test") + .expect("test zone should exist"); + + let forwarders: HashSet = zone + .forwarders + .iter() + .map(|u| NameServerAddr::try_from(u).expect("forwarder should be valid")) + .collect(); + + assert_eq!(forwarders.len(), 1); + assert!(forwarders.contains(&ns("udp://8.8.8.8:53"))); + + Ok(()) + } + + #[tokio::test] + async fn heartbeat_digest_resync_is_node_scoped() { + let mgr = DnsNodeMgr::new(); + let node_a = Uuid::new_v4(); + let node_b = Uuid::new_v4(); + + let snap_a = snapshot_with( + vec![valid_zone_data("scope-a.test", "10.60.0.1", vec![])], + vec!["udp://10.60.0.2:53"], + vec![], + ); + let digest_a = snap_a.digest(); + + let _ = send_heartbeat(&mgr, heartbeat_with_snapshot(node_a, snap_a)).await; + + let a_same = send_heartbeat(&mgr, heartbeat_digest_only(node_a, digest_a)).await; + assert!(!a_same.resync); + + let b_unknown = send_heartbeat(&mgr, heartbeat_digest_only(node_b, vec![1, 2, 3])).await; + assert!(b_unknown.resync); + } + + #[tokio::test] + async fn heartbeat_resync_after_node_idle_ttl_expiry() { + let mgr = DnsNodeMgr::new(); + let id = Uuid::new_v4(); + let snapshot = snapshot_with( + vec![valid_zone_data("ttl.test", "10.9.9.9", vec![])], + vec!["udp://10.9.0.1:53"], + vec![], + ); + let digest = snapshot.digest(); + + let _ = send_heartbeat(&mgr, heartbeat_with_snapshot(id, snapshot)).await; + let before_expiry = send_heartbeat(&mgr, heartbeat_digest_only(id, digest.clone())).await; + assert!(!before_expiry.resync); + + sleep(DNS_NODE_TTI + Duration::from_millis(300)).await; + + let after_expiry = send_heartbeat(&mgr, heartbeat_digest_only(id, digest)).await; + assert!(after_expiry.resync); + } }