From 8c5851e2eefa40f9fbeb50be941e5b93f80c29c5 Mon Sep 17 00:00:00 2001 From: Luna Yao <40349250+ZnqbuZ@users.noreply.github.com> Date: Mon, 6 Apr 2026 12:04:39 +0200 Subject: [PATCH] server: unit tests --- easytier/src/dns/server.rs | 311 +++++++++++++++++++++++++++++++++++++ 1 file changed, 311 insertions(+) diff --git a/easytier/src/dns/server.rs b/easytier/src/dns/server.rs index 71c8c889..dd157d6d 100644 --- a/easytier/src/dns/server.rs +++ b/easytier/src/dns/server.rs @@ -479,6 +479,9 @@ impl DnsServer { mod tests { use super::*; use crate::peers::tests::create_mock_peer_manager; + use crate::proto::common::Url; + use crate::proto::dns::{DnsNodeMgrRpc, DnsSnapshot, HeartbeatRequest, ZoneData}; + use crate::proto::rpc_types::controller::BaseController; use hickory_client::client::{Client, ClientHandle}; use hickory_proto::op::{Message, MessageType, OpCode, Query}; use hickory_proto::rr::{rdata, DNSClass, Name, RData, Record, RecordType}; @@ -496,6 +499,9 @@ mod tests { use std::net::{Ipv4Addr, SocketAddr}; use std::str::FromStr; use std::time::Duration; + use tokio::net::UdpSocket; + use tokio::time::{sleep, timeout}; + use uuid::Uuid; /// Build a `Catalog` containing a single A record: `test.example.com -> 1.2.3.4`. fn build_test_catalog() -> Catalog { @@ -612,6 +618,49 @@ mod tests { buf } + fn valid_zone_data(origin: &str, record: &str) -> ZoneData { + ZoneData { + id: Some(Uuid::new_v4().into()), + origin: origin.to_string(), + ttl: 60, + records: vec![format!("@ IN A {record}")], + forwarders: vec![], + } + } + + 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(), + } + } + + fn heartbeat_with_snapshot(id: Uuid, snapshot: DnsSnapshot) -> HeartbeatRequest { + let mut hb = HeartbeatRequest { + id: Some(id.into()), + ..Default::default() + }; + hb.update(snapshot); + hb + } + + async fn wait_until(mut f: impl FnMut() -> bool) { + for _ in 0..80 { + if f() { + return; + } + sleep(Duration::from_millis(50)).await; + } + panic!("condition not met in time"); + } + // ─── Tests ─────────────────────────────────────────────────────────── #[tokio::test] @@ -676,6 +725,10 @@ mod tests { // Verify IP addresses are swapped. assert_eq!(ip.get_source(), dst_ip); assert_eq!(ip.get_destination(), src_ip); + + // Verify route-to-self rewrite in peer manager header. + let hdr = zc.peer_manager_header().unwrap(); + assert_eq!(hdr.to_peer_id.get(), server.peer_mgr.my_peer_id() as u32); } #[tokio::test] @@ -771,6 +824,23 @@ mod tests { let udp_reply = UdpPacket::new(ip.payload()).unwrap(); assert_eq!(udp_reply.get_source(), dns_port); assert_eq!(udp_reply.get_destination(), client_port); + assert_eq!( + udp_reply.get_length() as usize, + 8 + udp_reply.payload().len(), + "UDP length should match payload size" + ); + + assert_eq!( + udp_reply.get_checksum(), + udp::ipv4_checksum(&udp_reply, &dst_ip, &src_ip), + "UDP checksum should be recomputed for swapped src/dst IP" + ); + + assert_eq!( + ip.get_total_length() as usize, + 20 + udp_reply.packet().len(), + "IP total length should match rewritten packet" + ); let dns_reply = Message::from_vec(udp_reply.payload()).unwrap(); assert_eq!(dns_reply.id(), 0x1234); @@ -844,4 +914,245 @@ mod tests { // Shutdown the server. shutdown_token.cancel(); } + + #[tokio::test] + async fn test_try_process_packet_from_nic_and_pipeline_id() { + let server = create_test_server().await; + let src_ip: Ipv4Addr = "10.0.0.1".parse().unwrap(); + let dst_ip: Ipv4Addr = "10.0.0.53".parse().unwrap(); + + server + .addresses + .write() + .insert(SocketAddr::new(dst_ip.into(), 53).into()); + + let icmp_payload = build_icmp_echo_request(); + let ip_bytes = build_ipv4_packet(src_ip, dst_ip, IpNextHeaderProtocols::Icmp, &icmp_payload); + let mut zc = ZCPacket::new_with_payload(&ip_bytes); + zc.fill_peer_manager_hdr(1, 2, crate::tunnel::packet_def::PacketType::Data as u8); + + assert!(server.try_process_packet_from_nic(&mut zc).await); + assert_eq!(server.id(), NIC_PIPELINE_NAME); + } + + #[tokio::test] + async fn test_handle_ip_packet_non_ipv4_version_ignored() { + let server = create_test_server().await; + let src_ip: Ipv4Addr = "10.0.0.1".parse().unwrap(); + let dst_ip: Ipv4Addr = "10.0.0.53".parse().unwrap(); + + server + .addresses + .write() + .insert(SocketAddr::new(dst_ip.into(), 53).into()); + + let icmp_payload = build_icmp_echo_request(); + let mut ip_bytes = build_ipv4_packet(src_ip, dst_ip, IpNextHeaderProtocols::Icmp, &icmp_payload); + ip_bytes[0] = (6 << 4) | 5; // fake IPv6 version in IPv4 header + + let mut zc = ZCPacket::new_with_payload(&ip_bytes); + zc.fill_peer_manager_hdr(1, 2, crate::tunnel::packet_def::PacketType::Data as u8); + + assert!(server.handle_ip_packet(&mut zc).await.is_none()); + } + + #[tokio::test] + async fn test_handle_ip_packet_unsupported_protocol_ignored() { + let server = create_test_server().await; + let src_ip: Ipv4Addr = "10.0.0.1".parse().unwrap(); + let dst_ip: Ipv4Addr = "10.0.0.53".parse().unwrap(); + + server + .addresses + .write() + .insert(SocketAddr::new(dst_ip.into(), 53).into()); + + let tcp_like_payload = vec![0u8; 20]; + let ip_bytes = build_ipv4_packet( + src_ip, + dst_ip, + IpNextHeaderProtocols::Tcp, + &tcp_like_payload, + ); + + let mut zc = ZCPacket::new_with_payload(&ip_bytes); + zc.fill_peer_manager_hdr(1, 2, crate::tunnel::packet_def::PacketType::Data as u8); + + assert!(server.handle_ip_packet(&mut zc).await.is_none()); + } + + #[tokio::test] + async fn test_handle_udp_packet_wrong_port_ignored() { + let server = create_test_server().await; + let dst_ip: Ipv4Addr = "10.0.0.53".parse().unwrap(); + let src_ip: Ipv4Addr = "10.0.0.1".parse().unwrap(); + + server + .addresses + .write() + .insert(SocketAddr::new(dst_ip.into(), 53).into()); + server.catalog.replace(build_test_catalog()).await; + + let dns_bytes = build_dns_query_bytes("test.example.com."); + let udp_bytes = build_udp_packet(12345, 5353, &dns_bytes, src_ip, dst_ip); + let ip_bytes = build_ipv4_packet(src_ip, dst_ip, IpNextHeaderProtocols::Udp, &udp_bytes); + + let mut zc = ZCPacket::new_with_payload(&ip_bytes); + zc.fill_peer_manager_hdr(1, 2, crate::tunnel::packet_def::PacketType::Data as u8); + + assert!(server.handle_ip_packet(&mut zc).await.is_none()); + } + + #[tokio::test] + async fn test_handle_udp_packet_invalid_dns_payload_ignored() { + let server = create_test_server().await; + let dst_ip: Ipv4Addr = "10.0.0.53".parse().unwrap(); + let src_ip: Ipv4Addr = "10.0.0.1".parse().unwrap(); + + server + .addresses + .write() + .insert(SocketAddr::new(dst_ip.into(), 53).into()); + + let invalid_dns = vec![0xde, 0xad, 0xbe]; + let udp_bytes = build_udp_packet(12345, 53, &invalid_dns, src_ip, dst_ip); + let ip_bytes = build_ipv4_packet(src_ip, dst_ip, IpNextHeaderProtocols::Udp, &udp_bytes); + + let mut zc = ZCPacket::new_with_payload(&ip_bytes); + zc.fill_peer_manager_hdr(1, 2, crate::tunnel::packet_def::PacketType::Data as u8); + + assert!(server.handle_ip_packet(&mut zc).await.is_none()); + } + + #[tokio::test] + async fn test_reload_addresses_updates_public_addresses() { + let server = create_test_server().await; + + let addrs = vec![ + "udp://10.10.10.53:53".parse::().unwrap(), + "tcp://10.10.10.54:5353".parse::().unwrap(), + ]; + server.reload_addresses(addrs.clone()).await.unwrap(); + + let as_socket = server.addresses(); + assert_eq!(as_socket.len(), 2); + assert!(as_socket.contains(&addrs[0].addr)); + assert!(as_socket.contains(&addrs[1].addr)); + + // No-op reload should keep the same content. + server.reload_addresses(addrs).await.unwrap(); + assert_eq!(server.addresses().len(), 2); + } + + #[tokio::test] + async fn test_reload_listeners_partial_bind_failure_still_serves_working_listener() { + let server = create_test_server().await; + server.catalog.replace(build_test_catalog()).await; + + let occupied = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let occupied_addr = occupied.local_addr().unwrap(); + + let probe = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let good_addr = probe.local_addr().unwrap(); + drop(probe); + + let listeners = vec![ + NameServerAddr { + protocol: Protocol::Udp, + addr: occupied_addr, + }, + NameServerAddr { + protocol: Protocol::Udp, + addr: good_addr, + }, + ]; + + let mut runtime = None; + server.reload_listeners(listeners, &mut runtime).await.unwrap(); + + let conn = UdpClientStream::builder(good_addr, TokioRuntimeProvider::default()).build(); + let (mut client, bg) = timeout(Duration::from_secs(2), Client::connect(conn)) + .await + .expect("client connect timeout") + .expect("client connect failed"); + tokio::spawn(async move { + let _ = bg.await; + }); + + let response = timeout( + Duration::from_secs(2), + client.query( + Name::from_str("test.example.com.").unwrap(), + DNSClass::IN, + RecordType::A, + ), + ) + .await + .expect("query timeout") + .expect("query failed"); + + assert!(!response.answers().is_empty()); + + if let Some(runtime) = runtime.take() { + let _ = runtime.stop().await; + } + } + + #[tokio::test] + async fn test_run_applies_snapshot_updates_and_clears_on_shutdown() { + let server = create_test_server().await; + let token = CancellationToken::new(); + let run_server = server.clone(); + let run_token = token.clone(); + let run_task = tokio::spawn(async move { + run_server.run(run_token).await; + }); + + let node_id = Uuid::new_v4(); + let snapshot = snapshot_with( + vec![valid_zone_data("run-loop.test", "7.7.7.7")], + vec!["udp://10.0.0.53:53"], + vec![], + ); + + DnsNodeMgrRpc::heartbeat( + &*server.mgr, + BaseController::default(), + heartbeat_with_snapshot(node_id, snapshot), + ) + .await + .unwrap(); + + wait_until(|| { + server + .addresses() + .contains(&"10.0.0.53:53".parse::().unwrap()) + }) + .await; + + // Verify catalog hot-reload by issuing a hijacked DNS packet query. + let src_ip: Ipv4Addr = "10.0.0.1".parse().unwrap(); + let dst_ip: Ipv4Addr = "10.0.0.53".parse().unwrap(); + let dns_bytes = build_dns_query_bytes("run-loop.test."); + let udp_bytes = build_udp_packet(12000, 53, &dns_bytes, src_ip, dst_ip); + let ip_bytes = build_ipv4_packet(src_ip, dst_ip, IpNextHeaderProtocols::Udp, &udp_bytes); + let mut zc = ZCPacket::new_with_payload(&ip_bytes); + zc.fill_peer_manager_hdr(1, 2, crate::tunnel::packet_def::PacketType::Data as u8); + + let mut handled = false; + for _ in 0..80 { + if server.handle_ip_packet(&mut zc).await.is_some() { + handled = true; + break; + } + sleep(Duration::from_millis(50)).await; + } + assert!(handled, "catalog should be hot-reloaded before timeout"); + + token.cancel(); + let _ = run_task.await; + + assert!(server.addresses().is_empty(), "run() should clear addresses on exit"); + assert!(server.listeners.read().is_empty(), "run() should clear listeners on exit"); + } }