impl(fakedns): fakedns server

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::{Edns, Message, MessageType, OpCode, ResponseCode},
-    op::MessageType,
+    rr::{RData, Record, RecordType, rdata::A, rdata::AAAA},
-    serialize::{binary, txt},
 };
-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<RwLock<FakeIPResolver>>,
+    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<RwLock<FakeIPResolver>>`.
+    #[must_use]
+    pub fn with_shared_resolver(bind: SocketAddr, resolver: Arc<RwLock<FakeIPResolver>>) -> 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<RwLock<FakeIPResolver>> {
+        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<Vec<u8>> {
+        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();
+    }
+}
+
+