refactor: change magicdns to internal redirect (#1428)

To resolve issue #1419, DNS request packets are read directly and responses are sent back internally instead of being forwarded to the listening port.

The DNS service on fake_ip (100.100.100.101) no longer supports DNS-over-TCP.

Co-authored-by: Sijie.Sun <sunsijie@buaa.edu.cn>

easytier/src/instance/dns_server/server_instance.rs

@@ -1,4 +1,4 @@
-// single-instance server in one machine, every easytier instance that has ip address and tun device will try create a server instance.
+// single-instance server in one machine, every easytier instance that has ip address and tun device will try to create a server instance.
 
 // magic dns client will connect to this server to update the dns records.
 // magic dns server will add the dns server ip address to the tun device, and forward the dns request to the dns server
@@ -6,22 +6,12 @@
 // magic dns client will establish a long live tcp connection to the magic dns server, and when the server stops or crashes,
 // all the clients will exit and let the easytier instance to launch a new server instance.
 
-use std::{collections::BTreeMap, net::Ipv4Addr, str::FromStr, sync::Arc, time::Duration};
-
-use anyhow::Context;
-use cidr::Ipv4Inet;
-use dashmap::DashMap;
-use hickory_proto::rr::LowerName;
-use multimap::MultiMap;
-use pnet::packet::{
-    icmp::{self, IcmpTypes, MutableIcmpPacket},
-    ip::IpNextHeaderProtocols,
-    ipv4::{self, MutableIpv4Packet},
-    tcp::{self, MutableTcpPacket},
-    udp::{self, MutableUdpPacket},
-    MutablePacket,
+use super::{
+    config::{GeneralConfigBuilder, RunConfigBuilder},
+    server::Server,
+    system_config::{OSConfig, SystemConfig},
+    MAGIC_DNS_INSTANCE_ADDR,
 };
-
 use crate::{
     common::{
         ifcfg::{IfConfiger, IfConfiguerTrait},
@@ -46,13 +36,27 @@ use crate::{
     },
     tunnel::{packet_def::ZCPacket, tcp::TcpTunnelListener},
 };
-
-use super::{
-    config::{GeneralConfigBuilder, RunConfigBuilder},
-    server::Server,
-    system_config::{OSConfig, SystemConfig},
-    MAGIC_DNS_INSTANCE_ADDR,
+use anyhow::Context;
+use cidr::Ipv4Inet;
+use dashmap::DashMap;
+use hickory_proto::rr::LowerName;
+use hickory_proto::serialize::binary::{BinDecodable, BinEncoder};
+use hickory_server::authority::{MessageRequest, MessageResponse};
+use hickory_server::server::{Request, RequestHandler, ResponseHandler, ResponseInfo};
+use multimap::MultiMap;
+use pnet::packet::icmp::{IcmpTypes, MutableIcmpPacket};
+use pnet::packet::ipv4::Ipv4Packet;
+use pnet::packet::udp::UdpPacket;
+use pnet::packet::{
+    icmp,
+    ip::IpNextHeaderProtocols,
+    ipv4::{self, MutableIpv4Packet},
+    udp::{self, MutableUdpPacket},
+    MutablePacket, Packet,
 };
+use std::net::{SocketAddr, SocketAddrV4};
+use std::sync::Mutex;
+use std::{collections::BTreeMap, io, net::Ipv4Addr, str::FromStr, sync::Arc, time::Duration};
 
 static NIC_PIPELINE_NAME: &str = "magic_dns_server";
 
@@ -117,7 +121,7 @@ impl MagicDnsServerInstanceData {
         self.dns_server
             .upsert(
                 LowerName::from_str(zone)
-                    .with_context(|| "Invalid zone name, expect fomat like \"et.net.\"")?,
+                    .with_context(|| "Invalid zone name, expect format like \"et.net.\"")?,
                 Arc::new(authority),
             )
             .await;
@@ -160,6 +164,14 @@ impl MagicDnsServerRpc for MagicDnsServerInstanceData {
         Ok(Default::default())
     }
 
+    async fn heartbeat(
+        &self,
+        _ctrl: Self::Controller,
+        _input: Void,
+    ) -> crate::proto::rpc_types::error::Result<Void> {
+        Ok(Default::default())
+    }
+
     async fn update_dns_record(
         &self,
         ctrl: Self::Controller,
@@ -204,97 +216,203 @@ impl MagicDnsServerRpc for MagicDnsServerInstanceData {
         }
         Ok(GetDnsRecordResponse { records: ret })
     }
+}
 
-    async fn heartbeat(
+// This should only be used for UDP response.
+// For other protocols, the variable `max_size` in `send_response` should be u16::MAX.
+#[derive(Clone)]
+struct ResponseWrapper {
+    response: Arc<Mutex<Vec<u8>>>,
+}
+
+trait RecordIter<'a>: Iterator<Item = &'a hickory_proto::rr::Record> + Send + 'a {}
+impl<'a, T> RecordIter<'a> for T where T: Iterator<Item = &'a hickory_proto::rr::Record> + Send + 'a {}
+
+#[async_trait::async_trait]
+impl ResponseHandler for ResponseWrapper {
+    async fn send_response<'a>(
+        &mut self,
+        response: MessageResponse<
+            '_,
+            'a,
+            impl RecordIter<'a>,
+            impl RecordIter<'a>,
+            impl RecordIter<'a>,
+            impl RecordIter<'a>,
+        >,
+    ) -> io::Result<ResponseInfo> {
+        let mut buffer = self
+            .response
+            .lock()
+            .map_err(|_| io::Error::other("lock poisoned"))?;
+
+        let mut encoder = BinEncoder::new(&mut buffer);
+
+        // `max_size` should be u16::MAX for protocol other than UDP.
+        let max_size = if let Some(edns) = response.get_edns() {
+            edns.max_payload()
+        } else {
+            hickory_proto::udp::MAX_RECEIVE_BUFFER_SIZE as u16
+        };
+
+        encoder.set_max_size(max_size);
+        response
+            .destructive_emit(&mut encoder)
+            .map_err(io::Error::other)
+    }
+}
+
+impl MagicDnsServerInstanceData {
+    /// Replace content of incoming UDP DNS request and ICMP echo request packet with reply data,
+    /// and swap source and destination IP addresses to send it back.
+    async fn handle_ip_packet(&self, zc_packet: &mut ZCPacket) -> Option<()> {
+        let (ip_header_length, ip_protocol, src_ip, dst_ip) = {
+            let ip_packet = Ipv4Packet::new(zc_packet.payload())?;
+
+            if ip_packet.get_version() != 4 {
+                return None;
+            }
+
+            (
+                ip_packet.get_header_length() as usize * 4,
+                ip_packet.get_next_level_protocol(),
+                ip_packet.get_source(),
+                ip_packet.get_destination(),
+            )
+        };
+
+        if dst_ip != self.fake_ip {
+            return None;
+        }
+
+        match ip_protocol {
+            IpNextHeaderProtocols::Udp => {
+                self.handle_udp_packet(zc_packet, ip_header_length, src_ip, dst_ip)
+                    .await?;
+            }
+            IpNextHeaderProtocols::Icmp => {
+                self.handle_icmp_packet(zc_packet, ip_header_length)?;
+            }
+            _ => {
+                return None;
+            }
+        }
+
+        let mut ip_packet = MutableIpv4Packet::new(zc_packet.mut_payload())?;
+        ip_packet.set_source(dst_ip);
+        ip_packet.set_destination(src_ip);
+
+        ip_packet.set_checksum(ipv4::checksum(&ip_packet.to_immutable()));
+
+        zc_packet.mut_peer_manager_header().unwrap().to_peer_id = self.my_peer_id.into();
+
+        Some(())
+    }
+
+    /// Extract the DNS request message and send it to the hickory-dns server instance.
+    /// Replace the content of the UDP packet with the response message.
+    async fn handle_udp_packet(
         &self,
-        _ctrl: Self::Controller,
-        _input: Void,
-    ) -> crate::proto::rpc_types::error::Result<Void> {
-        Ok(Default::default())
+        zc_packet: &mut ZCPacket,
+        ip_header_length: usize,
+        src_ip: Ipv4Addr,
+        dst_ip: Ipv4Addr,
+    ) -> Option<()> {
+        let (src_port, dst_port, request, request_length) = {
+            let udp_packet = UdpPacket::new(&zc_packet.payload()[ip_header_length..])?;
+
+            let src_port = udp_packet.get_source();
+            let dst_port = udp_packet.get_destination();
+
+            // Remove this to support any UDP port
+            if dst_port != 53 {
+                return None;
+            }
+
+            let request_payload = udp_packet.payload();
+
+            (
+                src_port,
+                dst_port,
+                Request::new(
+                    MessageRequest::from_bytes(request_payload).ok()?,
+                    SocketAddr::from(SocketAddrV4::new(src_ip, src_port)),
+                    hickory_proto::xfer::Protocol::Udp,
+                ),
+                request_payload.len(),
+            )
+        };
+
+        let response_payload = {
+            let response_payload_arc = Arc::new(Mutex::new(Vec::with_capacity(512)));
+
+            self.dns_server
+                .read_catalog()
+                .await
+                .handle_request(
+                    &request,
+                    ResponseWrapper {
+                        response: response_payload_arc.clone(),
+                    },
+                )
+                .await;
+
+            Arc::into_inner(response_payload_arc)?.into_inner().ok()?
+        };
+
+        let response_length = response_payload.len();
+        let delta_length = response_length as isize - request_length as isize;
+
+        let inner_length = (zc_packet.buf_len() as isize + delta_length) as usize;
+        if zc_packet.mut_inner().capacity() < inner_length {
+            let header_length = inner_length - response_length;
+            zc_packet.mut_inner().truncate(header_length);
+        }
+        zc_packet.mut_inner().resize(inner_length, 0);
+
+        let mut ip_packet = MutableIpv4Packet::new(zc_packet.mut_payload())?;
+
+        let ip_length = (ip_packet.get_total_length() as isize + delta_length) as u16;
+        ip_packet.set_total_length(ip_length);
+
+        let mut udp_packet = MutableUdpPacket::new(ip_packet.payload_mut())?;
+
+        let udp_length = (udp_packet.get_length() as isize + delta_length) as u16;
+        udp_packet.set_length(udp_length);
+
+        udp_packet.set_source(dst_port);
+        udp_packet.set_destination(src_port);
+
+        udp_packet.payload_mut().copy_from_slice(&response_payload);
+
+        udp_packet.set_checksum(udp::ipv4_checksum(
+            &udp_packet.to_immutable(),
+            &dst_ip,
+            &src_ip,
+        ));
+
+        Some(())
+    }
+
+    fn handle_icmp_packet(&self, zc_packet: &mut ZCPacket, ip_header_length: usize) -> Option<()> {
+        let mut icmp_packet =
+            MutableIcmpPacket::new(&mut zc_packet.mut_payload()[ip_header_length..])?;
+
+        if icmp_packet.get_icmp_type() != IcmpTypes::EchoRequest {
+            return None;
+        }
+
+        icmp_packet.set_icmp_type(IcmpTypes::EchoReply);
+        icmp_packet.set_checksum(icmp::checksum(&icmp_packet.to_immutable()));
+
+        Some(())
     }
 }
 
 #[async_trait::async_trait]
 impl NicPacketFilter for MagicDnsServerInstanceData {
     async fn try_process_packet_from_nic(&self, zc_packet: &mut ZCPacket) -> bool {
-        let data = zc_packet.mut_payload();
-        let mut ip_packet = MutableIpv4Packet::new(data).unwrap();
-        if ip_packet.get_version() != 4 || ip_packet.get_destination() != self.fake_ip {
-            return false;
-        }
-
-        match ip_packet.get_next_level_protocol() {
-            IpNextHeaderProtocols::Udp => {
-                let Some(dns_udp_addr) = self.dns_server.udp_local_addr() else {
-                    return false;
-                };
-
-                let Some(mut udp_packet) = MutableUdpPacket::new(ip_packet.payload_mut()) else {
-                    return false;
-                };
-                if udp_packet.get_destination() == 53 {
-                    // for dns request
-                    udp_packet.set_destination(dns_udp_addr.port());
-                } else if udp_packet.get_source() == dns_udp_addr.port() {
-                    // for dns response
-                    udp_packet.set_source(53);
-                } else {
-                    return false;
-                }
-                udp_packet.set_checksum(udp::ipv4_checksum(
-                    &udp_packet.to_immutable(),
-                    &self.fake_ip,
-                    &self.tun_ip,
-                ));
-            }
-
-            IpNextHeaderProtocols::Tcp => {
-                let Some(dns_tcp_addr) = self.dns_server.tcp_local_addr() else {
-                    return false;
-                };
-
-                let Some(mut tcp_packet) = MutableTcpPacket::new(ip_packet.payload_mut()) else {
-                    return false;
-                };
-                if tcp_packet.get_destination() == 53 {
-                    // for dns request
-                    tcp_packet.set_destination(dns_tcp_addr.port());
-                } else if tcp_packet.get_source() == dns_tcp_addr.port() {
-                    // for dns response
-                    tcp_packet.set_source(53);
-                } else {
-                    return false;
-                }
-                tcp_packet.set_checksum(tcp::ipv4_checksum(
-                    &tcp_packet.to_immutable(),
-                    &self.fake_ip,
-                    &self.tun_ip,
-                ));
-            }
-
-            IpNextHeaderProtocols::Icmp => {
-                let Some(mut icmp_packet) = MutableIcmpPacket::new(ip_packet.payload_mut()) else {
-                    return false;
-                };
-                if icmp_packet.get_icmp_type() != IcmpTypes::EchoRequest {
-                    return false;
-                }
-                icmp_packet.set_icmp_type(IcmpTypes::EchoReply);
-                icmp_packet.set_checksum(icmp::checksum(&icmp_packet.to_immutable()));
-            }
-
-            _ => {
-                return false;
-            }
-        }
-
-        ip_packet.set_source(self.fake_ip);
-        ip_packet.set_destination(self.tun_ip);
-
-        ip_packet.set_checksum(ipv4::checksum(&ip_packet.to_immutable()));
-        zc_packet.mut_peer_manager_header().unwrap().to_peer_id = self.my_peer_id.into();
-
-        true
+        self.handle_ip_packet(zc_packet).await.is_some()
     }
 
     fn id(&self) -> String {
@@ -363,23 +481,14 @@ impl MagicDnsServerInstance {
         tun_inet: Ipv4Inet,
         fake_ip: Ipv4Addr,
     ) -> Result<Self, anyhow::Error> {
-        let tcp_listener = TcpTunnelListener::new(MAGIC_DNS_INSTANCE_ADDR.parse().unwrap());
+        let tcp_listener = TcpTunnelListener::new(MAGIC_DNS_INSTANCE_ADDR.parse()?);
         let mut rpc_server = StandAloneServer::new(tcp_listener);
         rpc_server.serve().await?;
 
-        let bind_addr = tun_inet.address();
-
         let dns_config = RunConfigBuilder::default()
-            .general(
-                GeneralConfigBuilder::default()
-                    .listen_udp(format!("{}:0", bind_addr))
-                    .listen_tcp(format!("{}:0", bind_addr))
-                    .build()
-                    .unwrap(),
-            )
+            .general(GeneralConfigBuilder::default().build()?)
             .excluded_forward_nameservers(vec![fake_ip.into()])
-            .build()
-            .unwrap();
+            .build()?;
         let mut dns_server = Server::new(dns_config);
         dns_server.run().await?;