node_mgr: unit tests

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<ZoneData>,
+        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<u8>) -> 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<Message> {
+        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<ZoneData> = 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<NameServerAddr> = 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<ZoneData> = 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<NameServerAddr> = 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);
+    }
 }