diff --git a/src/sniffing/fakedns/server.rs b/src/sniffing/fakedns/server.rs index c620f66..eb1c819 100644 --- a/src/sniffing/fakedns/server.rs +++ b/src/sniffing/fakedns/server.rs @@ -1,24 +1,555 @@ -use crate::sniffing::fakedns::resolver::FakeIPResolver; +use crate::sniffing::fakedns::resolver::{FakeIPOptions, FakeIPResolver, IPStrategies}; use hickory_proto::{ - op::Message, - op::MessageType, - serialize::{binary, txt}, + op::{Edns, Message, MessageType, OpCode, ResponseCode}, + rr::{RData, Record, RecordType, rdata::A, rdata::AAAA}, }; -use std::net::SocketAddr; +use std::net::{IpAddr, SocketAddr}; +use std::sync::Arc; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, UdpSocket}; +use tokio::sync::RwLock; +/// Default TTL for DNS responses generated by FakeDNS server (in seconds). +pub const DEFAULT_DNS_TTL: u32 = 60; + +/// FakeDNS server that captures DNS queries and resolves them via `FakeIPResolver`. +#[derive(Clone)] pub struct FakeIPServer { bind: SocketAddr, - resolver: FakeIPResolver, + resolver: Arc>, + ttl: u32, } impl FakeIPServer { + /// Creates a new `FakeIPServer` binding to the given socket address with the provided resolver. #[must_use] pub fn new(bind: SocketAddr, resolver: FakeIPResolver) -> Self { - Self { bind, resolver } + Self { + bind, + resolver: Arc::new(RwLock::new(resolver)), + ttl: DEFAULT_DNS_TTL, + } } - pub async fn serve() -> Result<(), std::io::Error> { todo!() } - async fn serve_tcp() -> Result<(), std::io::Error> { todo!() } - async fn serve_udp() -> Result<(), std::io::Error> { todo!() } - fn set_as_default_resolver() -> Result<(), std::io::Error> { todo!() } + /// Creates a new `FakeIPServer` with a shared `Arc>`. + #[must_use] + pub fn with_shared_resolver(bind: SocketAddr, resolver: Arc>) -> Self { + Self { + bind, + resolver, + ttl: DEFAULT_DNS_TTL, + } + } + + /// Configures the TTL for DNS responses. + #[must_use] + pub fn with_ttl(mut self, ttl: u32) -> Self { + self.ttl = ttl; + self + } + + /// Returns the bound socket address. + #[must_use] + pub fn bind(&self) -> SocketAddr { + self.bind + } + + /// Returns a shared reference to the inner resolver. + #[must_use] + pub fn resolver(&self) -> Arc> { + self.resolver.clone() + } + + /// Returns the configured response TTL. + #[must_use] + pub fn ttl(&self) -> u32 { + self.ttl + } + + /// Starts serving both UDP and TCP DNS requests concurrently. + pub async fn serve(&self) -> Result<(), std::io::Error> { + tokio::try_join!(self.serve_udp(), self.serve_tcp()).map(|_| ()) + } + + /// Starts the UDP DNS server loop. + pub async fn serve_udp(&self) -> Result<(), std::io::Error> { + let socket = Arc::new(UdpSocket::bind(self.bind).await?); + let mut buf = [0u8; 4096]; + + loop { + let (amt, src) = socket.recv_from(&mut buf).await?; + let packet = buf[..amt].to_vec(); + let socket_clone = socket.clone(); + let server_clone = self.clone(); + + tokio::spawn(async move { + if let Some(resp_bytes) = server_clone.handle_packet(&packet).await { + let _ = socket_clone.send_to(&resp_bytes, src).await; + } + }); + } + } + + /// Starts the TCP DNS server loop. + pub async fn serve_tcp(&self) -> Result<(), std::io::Error> { + let listener = TcpListener::bind(self.bind).await?; + + loop { + let (mut stream, _src) = listener.accept().await?; + let server_clone = self.clone(); + + tokio::spawn(async move { + let mut len_buf = [0u8; 2]; + loop { + match stream.read_exact(&mut len_buf).await { + Ok(_) => {} + Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => break, + Err(_) => break, + } + + let msg_len = u16::from_be_bytes(len_buf) as usize; + if msg_len == 0 { + continue; + } + + let mut msg_buf = vec![0u8; msg_len]; + if stream.read_exact(&mut msg_buf).await.is_err() { + break; + } + + if let Some(resp_bytes) = server_clone.handle_packet(&msg_buf).await { + let Ok(resp_len) = u16::try_from(resp_bytes.len()) else { + break; + }; + if stream.write_all(&resp_len.to_be_bytes()).await.is_err() + || stream.write_all(&resp_bytes).await.is_err() + || stream.flush().await.is_err() + { + break; + } + } + } + }); + } + } + + /// Decodes a raw DNS packet, handles the query, and encodes the DNS response. + pub async fn handle_packet(&self, packet: &[u8]) -> Option> { + let request = Message::from_vec(packet).ok()?; + let response = self.handle_message(&request).await; + response.to_vec().ok() + } + + /// Processes a DNS request `Message` and generates a DNS response `Message`. + pub async fn handle_message(&self, request: &Message) -> Message { + // If not a query message type, return format error + if request.message_type != MessageType::Query { + return Message::error_msg(request.id, request.op_code, ResponseCode::FormErr); + } + + // If not standard query opcode, return Not Implemented + if request.op_code != OpCode::Query { + return Message::error_msg(request.id, request.op_code, ResponseCode::NotImp); + } + + let mut response = Message::response(request.id, request.op_code); + response.metadata.authoritative = true; + response.metadata.recursion_desired = request.recursion_desired; + response.metadata.recursion_available = true; + response.metadata.response_code = ResponseCode::NoError; + + // Echo client queries + for query in &request.queries { + response.add_query(query.clone()); + } + + // Echo EDNS if provided by client + if let Some(ref client_edns) = request.edns { + let mut edns = Edns::new(); + edns.set_max_payload(client_edns.max_payload().max(512)); + edns.set_version(client_edns.version()); + response.set_edns(edns); + } + + for query in &request.queries { + let qname = &query.name; + let domain_str = qname.to_ascii(); + let clean_domain = domain_str.trim_end_matches('.').to_lowercase(); + + if clean_domain.is_empty() { + continue; + } + + match query.query_type { + RecordType::A => { + let maybe_ip = { + let res = self.resolver.read().await; + res.get(&clean_domain) + }; + + let resolved_ip = match maybe_ip { + Some(ip) => Some(ip), + None => { + let mut res = self.resolver.write().await; + if let Some(ip) = res.get(&clean_domain) { + Some(ip) + } else if res.resolver.is_some() { + match res + .add_real_ip(clean_domain.clone(), IPStrategies::V4Only) + .await + { + Ok(ip) => Some(ip), + Err(_) => Some(IpAddr::V4(res.add(clean_domain.clone()))), + } + } else { + Some(IpAddr::V4(res.add(clean_domain.clone()))) + } + } + }; + + if let Some(IpAddr::V4(ipv4)) = resolved_ip { + let rdata = RData::A(A(ipv4)); + let record = Record::from_rdata(qname.clone(), self.ttl, rdata); + response.add_answer(record); + } + } + RecordType::AAAA => { + let maybe_ip = { + let res = self.resolver.read().await; + res.get(&clean_domain) + }; + + let resolved_ip = match maybe_ip { + Some(ip) => Some(ip), + None => { + let mut res = self.resolver.write().await; + if let Some(ip) = res.get(&clean_domain) { + Some(ip) + } else if res.resolver.is_some() { + res.add_real_ip(clean_domain.clone(), IPStrategies::V6Only) + .await + .ok() + } else { + None + } + } + }; + + if let Some(IpAddr::V6(ipv6)) = resolved_ip { + let rdata = RData::AAAA(AAAA(ipv6)); + let record = Record::from_rdata(qname.clone(), self.ttl, rdata); + response.add_answer(record); + } + } + RecordType::ANY => { + let maybe_ip = { + let res = self.resolver.read().await; + res.get(&clean_domain) + }; + + let resolved_ip = match maybe_ip { + Some(ip) => Some(ip), + None => { + let mut res = self.resolver.write().await; + if let Some(ip) = res.get(&clean_domain) { + Some(ip) + } else if res.resolver.is_some() { + match res + .add_real_ip(clean_domain.clone(), IPStrategies::Any) + .await + { + Ok(ip) => Some(ip), + Err(_) => Some(IpAddr::V4(res.add(clean_domain.clone()))), + } + } else { + Some(IpAddr::V4(res.add(clean_domain.clone()))) + } + } + }; + + match resolved_ip { + Some(IpAddr::V4(ipv4)) => { + let rdata = RData::A(A(ipv4)); + let record = Record::from_rdata(qname.clone(), self.ttl, rdata); + response.add_answer(record); + } + Some(IpAddr::V6(ipv6)) => { + let rdata = RData::AAAA(AAAA(ipv6)); + let record = Record::from_rdata(qname.clone(), self.ttl, rdata); + response.add_answer(record); + } + None => {} + } + } + _ => { + // Other record types (HTTPS, MX, TXT, SOA) return NOERROR with no answers + } + } + } + + response + } + + /// Sets this server's address as the system default resolver in `/etc/resolv.conf`. + #[allow(dead_code)] + pub fn set_as_default_resolver(&self) -> Result<(), std::io::Error> { + Self::set_system_resolver(self.bind.ip()) + } + + /// Sets the given IP address as the system default resolver in `/etc/resolv.conf`. + #[allow(dead_code)] + pub fn set_system_resolver(ip: IpAddr) -> Result<(), std::io::Error> { + let content = format!("nameserver {ip}\n"); + std::fs::write("/etc/resolv.conf", content) + } } + +#[cfg(test)] +mod tests { + use super::*; + use hickory_proto::op::Query; + use hickory_proto::rr::Name; + + #[tokio::test] + async fn test_fake_dns_a_query() { + let resolver = FakeIPResolver::new(); + let server = FakeIPServer::new("127.0.0.1:0".parse().unwrap(), resolver); + + let mut query_msg = Message::new(42, MessageType::Query, OpCode::Query); + let qname = Name::from_ascii("example.com.").unwrap(); + let query = Query::query(qname.clone(), RecordType::A); + query_msg.add_query(query); + + let query_bytes = query_msg.to_vec().unwrap(); + let resp_bytes = server + .handle_packet(&query_bytes) + .await + .expect("should return response bytes"); + + let resp_msg = Message::from_vec(&resp_bytes).expect("should parse response message"); + assert_eq!(resp_msg.id, 42); + assert_eq!(resp_msg.metadata.message_type, MessageType::Response); + assert_eq!(resp_msg.metadata.response_code, ResponseCode::NoError); + assert_eq!(resp_msg.answers.len(), 1); + + let answer = &resp_msg.answers[0]; + assert_eq!(&answer.name, &qname); + assert_eq!(answer.record_type(), RecordType::A); + + let ip = match &answer.data { + RData::A(A(ip)) => *ip, + other => panic!("expected A record, got {other:?}"), + }; + + // RFC 2544 benchmarking space: 198.18.0.0/15 -> 198.18.0.0 - 198.19.255.255 + let octets = ip.octets(); + assert_eq!(octets[0], 198); + assert_eq!(octets[1] & 0xFE, 18); + + // Ensure consecutive queries for the same domain return the exact same IP + let resp_bytes2 = server.handle_packet(&query_bytes).await.unwrap(); + let resp_msg2 = Message::from_vec(&resp_bytes2).unwrap(); + let ip2 = match &resp_msg2.answers[0].data { + RData::A(A(ip)) => *ip, + other => panic!("expected A record, got {other:?}"), + }; + assert_eq!(ip, ip2); + + // Check backend resolver state + let res_lock = server.resolver.read().await; + assert_eq!(res_lock.get("example.com"), Some(IpAddr::V4(ip))); + } + + #[tokio::test] + async fn test_fake_dns_aaaa_query() { + let resolver = FakeIPResolver::new(); + let server = FakeIPServer::new("127.0.0.1:0".parse().unwrap(), resolver); + + let mut query_msg = Message::new(100, MessageType::Query, OpCode::Query); + let qname = Name::from_ascii("nosystemd.org.").unwrap(); + let query = Query::query(qname.clone(), RecordType::AAAA); + query_msg.add_query(query); + + let query_bytes = query_msg.to_vec().unwrap(); + let resp_bytes = server + .handle_packet(&query_bytes) + .await + .expect("should return response bytes"); + + let resp_msg = Message::from_vec(&resp_bytes).expect("should parse response message"); + assert_eq!(resp_msg.id, 100); + assert_eq!(resp_msg.metadata.response_code, ResponseCode::NoError); + assert_eq!(resp_msg.answers.len(), 0); // NODATA + } + + #[tokio::test] + async fn test_fake_dns_unsupported_opcode_and_type() { + let resolver = FakeIPResolver::new(); + let server = FakeIPServer::new("127.0.0.1:0".parse().unwrap(), resolver); + + // Test non-Query opcode + let status_msg = Message::new(101, MessageType::Query, OpCode::Status); + let resp = server.handle_message(&status_msg).await; + assert_eq!(resp.metadata.response_code, ResponseCode::NotImp); + + // Test TXT query returns NOERROR with 0 answers + let mut txt_msg = Message::new(102, MessageType::Query, OpCode::Query); + let qname = Name::from_ascii("example.org.").unwrap(); + txt_msg.add_query(Query::query(qname, RecordType::TXT)); + let resp_txt = server.handle_message(&txt_msg).await; + assert_eq!(resp_txt.metadata.response_code, ResponseCode::NoError); + assert_eq!(resp_txt.answers.len(), 0); + } + + #[tokio::test] + async fn test_fake_dns_udp_serve() { + let probe = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let addr = probe.local_addr().unwrap(); + drop(probe); + + let resolver = FakeIPResolver::new(); + let server = FakeIPServer::new(addr, resolver); + + let server_clone = server.clone(); + tokio::spawn(async move { + let _ = server_clone.serve_udp().await; + }); + + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + + let client = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + + let mut query_msg = Message::new(55, MessageType::Query, OpCode::Query); + let qname = Name::from_ascii("udp-test.org.").unwrap(); + query_msg.add_query(Query::query(qname.clone(), RecordType::A)); + let qbytes = query_msg.to_vec().unwrap(); + + client.send_to(&qbytes, addr).await.unwrap(); + + let mut buf = [0u8; 1024]; + let (amt, _) = tokio::time::timeout(tokio::time::Duration::from_secs(2), client.recv_from(&mut buf)) + .await + .expect("UDP receive timed out") + .expect("UDP receive failed"); + + let resp_msg = Message::from_vec(&buf[..amt]).unwrap(); + assert_eq!(resp_msg.id, 55); + assert_eq!(resp_msg.answers.len(), 1); + assert_eq!(&resp_msg.answers[0].name, &qname); + } + + #[tokio::test] + async fn test_fake_dns_tcp_serve() { + let probe = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = probe.local_addr().unwrap(); + drop(probe); + + let resolver = FakeIPResolver::new(); + let server = FakeIPServer::new(addr, resolver); + + let server_clone = server.clone(); + tokio::spawn(async move { + let _ = server_clone.serve_tcp().await; + }); + + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + + let mut stream = tokio::time::timeout( + tokio::time::Duration::from_secs(2), + tokio::net::TcpStream::connect(addr), + ) + .await + .expect("TCP connect timed out") + .expect("TCP connect failed"); + + let mut query_msg = Message::new(77, MessageType::Query, OpCode::Query); + let qname = Name::from_ascii("tcp-test.org.").unwrap(); + query_msg.add_query(Query::query(qname.clone(), RecordType::A)); + let qbytes = query_msg.to_vec().unwrap(); + + let qlen = (qbytes.len() as u16).to_be_bytes(); + stream.write_all(&qlen).await.unwrap(); + stream.write_all(&qbytes).await.unwrap(); + stream.flush().await.unwrap(); + + let mut rlen_buf = [0u8; 2]; + stream.read_exact(&mut rlen_buf).await.unwrap(); + let rlen = u16::from_be_bytes(rlen_buf) as usize; + + let mut rbuf = vec![0u8; rlen]; + stream.read_exact(&mut rbuf).await.unwrap(); + + let resp_msg = Message::from_vec(&rbuf).unwrap(); + assert_eq!(resp_msg.id, 77); + assert_eq!(resp_msg.answers.len(), 1); + assert_eq!(&resp_msg.answers[0].name, &qname); + } + + #[tokio::test] + async fn test_fake_dns_serve_combined() { + let probe = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let addr = probe.local_addr().unwrap(); + drop(probe); + + let resolver = Arc::new(RwLock::new(FakeIPResolver::new())); + let server = FakeIPServer::with_shared_resolver(addr, resolver.clone()) + .with_ttl(30); + + assert_eq!(server.bind(), addr); + assert_eq!(server.ttl(), 30); + assert!(Arc::ptr_eq(&server.resolver(), &resolver)); + + let server_task = tokio::spawn(async move { + let _ = server.serve().await; + }); + + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + + // Test UDP through serve() + let client_udp = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let mut query_msg = Message::new(10, MessageType::Query, OpCode::Query); + let qname = Name::from_ascii("combined-test.com.").unwrap(); + query_msg.add_query(Query::query(qname.clone(), RecordType::A)); + let qbytes = query_msg.to_vec().unwrap(); + + client_udp.send_to(&qbytes, addr).await.unwrap(); + let mut buf = [0u8; 1024]; + let (amt, _) = tokio::time::timeout(tokio::time::Duration::from_secs(2), client_udp.recv_from(&mut buf)) + .await + .expect("UDP timeout") + .expect("UDP recv error"); + let resp_udp = Message::from_vec(&buf[..amt]).unwrap(); + assert_eq!(resp_udp.id, 10); + assert_eq!(resp_udp.answers.len(), 1); + + // Test TCP through serve() + let mut client_tcp = tokio::time::timeout( + tokio::time::Duration::from_secs(2), + tokio::net::TcpStream::connect(addr), + ) + .await + .expect("TCP timeout") + .expect("TCP connect error"); + + let mut query_tcp = Message::new(20, MessageType::Query, OpCode::Query); + query_tcp.add_query(Query::query(qname.clone(), RecordType::A)); + let qbytes_tcp = query_tcp.to_vec().unwrap(); + let qlen = (qbytes_tcp.len() as u16).to_be_bytes(); + client_tcp.write_all(&qlen).await.unwrap(); + client_tcp.write_all(&qbytes_tcp).await.unwrap(); + client_tcp.flush().await.unwrap(); + + let mut rlen_buf = [0u8; 2]; + client_tcp.read_exact(&mut rlen_buf).await.unwrap(); + let rlen = u16::from_be_bytes(rlen_buf) as usize; + let mut rbuf = vec![0u8; rlen]; + client_tcp.read_exact(&mut rbuf).await.unwrap(); + let resp_tcp = Message::from_vec(&rbuf).unwrap(); + assert_eq!(resp_tcp.id, 20); + assert_eq!(resp_tcp.answers.len(), 1); + + server_task.abort(); + } +} + +