server: impl NicPacketFilter

easytier/src/dns/server.rs

@@ -1,13 +1,18 @@
+use crate::common::PeerId;
 use crate::dns::node_mgr::DnsNodeMgr;
 use crate::dns::utils::addr::NameServerAddr;
+use crate::peer_center::instance::PeerCenterPeerManagerTrait;
 use crate::peers::peer_manager::PeerManager;
+use crate::peers::NicPacketFilter;
 use crate::proto::dns::DnsNodeMgrRpcServer;
 use crate::tunnel::common::bind_socket;
+use crate::tunnel::packet_def::ZCPacket;
 use derivative::Derivative;
 use derive_more::{Deref, DerefMut, From, Into};
 use hickory_proto::rr::Record;
-use hickory_proto::serialize::binary::BinEncoder;
+use hickory_proto::serialize::binary::{BinDecodable, BinEncoder};
 use hickory_proto::xfer::Protocol;
+use hickory_server::authority::MessageRequest;
 use hickory_server::{
     authority::{Catalog, MessageResponse},
     server::{Request, RequestHandler, ResponseHandler, ResponseInfo},
@@ -15,8 +20,14 @@ use hickory_server::{
 };
 use itertools::Itertools;
 use parking_lot::Mutex;
+use pnet::packet::icmp::{IcmpTypes, MutableIcmpPacket};
+use pnet::packet::ip::IpNextHeaderProtocols;
+use pnet::packet::ipv4::{Ipv4Packet, MutableIpv4Packet};
+use pnet::packet::udp::{MutableUdpPacket, UdpPacket};
+use pnet::packet::{icmp, ipv4, udp, MutablePacket, Packet};
 use std::collections::HashSet;
 use std::io;
+use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4};
 use std::{sync::Arc, time::Duration};
 use tokio::{sync::RwLock, task::JoinHandle};
 use tokio_util::sync::CancellationToken;
@@ -131,6 +142,9 @@ pub struct DnsServer {
 
     #[derivative(Debug = "ignore")]
     catalog: DynamicCatalog,
+
+    my_peer_id: PeerId,
+    addresses: Arc<RwLock<HashSet<NameServerAddr>>>,
 }
 
 const DNS_SERVER_LISTENER_TCP_TIMEOUT: Duration = Duration::from_secs(5);
@@ -150,23 +164,22 @@ impl DnsServer {
         Self {
             mgr,
             catalog: DynamicCatalog::new(),
+            my_peer_id: peer_mgr.my_peer_id(),
+            addresses: Arc::new(Default::default()),
         }
     }
 
-    async fn reload_addresses(
+    async fn reload_addresses(&self, addresses: impl IntoIterator<Item = NameServerAddr>) {
-        &self,
-        addresses: impl IntoIterator<Item = NameServerAddr>,
-        current: &mut HashSet<NameServerAddr>,
-    ) {
         let addresses = addresses.into_iter().collect::<HashSet<_>>();
+        let mut current = self.addresses.write().await; // TODO: read?
 
-        let added = addresses.difference(current).cloned().collect_vec();
+        if *current == addresses {
-        let removed = current.difference(&addresses).cloned().collect_vec();
-
-        if added.is_empty() && removed.is_empty() {
             return;
         }
 
+        let added = addresses.difference(&*current).cloned().collect_vec();
+        let removed = current.difference(&addresses).cloned().collect_vec();
+
         *current = addresses;
 
         // TODO
@@ -213,12 +226,10 @@ impl DnsServer {
                 }
             },
             async {
-                let mut addresses = HashSet::new();
                 loop {
                     dirty.addresses.notified().await;
                     if dirty.addresses.reset() {
-                        self.reload_addresses(self.mgr.iter_addresses(), &mut addresses)
+                        self.reload_addresses(self.mgr.iter_addresses()).await;
-                            .await;
                     }
                     tokio::time::sleep(Duration::from_secs(1)).await;
                 }
@@ -242,3 +253,173 @@ impl DnsServer {
         );
     }
 }
+
+const NIC_PIPELINE_NAME: &str = "magic_dns_server";
+
+#[async_trait::async_trait]
+impl NicPacketFilter for DnsServer {
+    async fn try_process_packet_from_nic(&self, zc_packet: &mut ZCPacket) -> bool {
+        self.handle_ip_packet(zc_packet).await.is_some()
+    }
+
+    fn id(&self) -> String {
+        NIC_PIPELINE_NAME.to_string()
+    }
+}
+
+impl DnsServer {
+    async fn is_hijacked_ip(&self, ip: &IpAddr) -> bool {
+        self.addresses
+            .read()
+            .await
+            .iter()
+            .any(|a| a.addr.ip() == *ip)
+    }
+
+    async fn is_hijacked_addr(&self, addr: &NameServerAddr) -> bool {
+        self.addresses.read().await.contains(addr)
+    }
+
+    /// 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 !self.is_hijacked_ip(&dst_ip.into()).await {
+            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;
+            }
+        }
+
+        // Swap source and destination IP addresses for the reply.
+        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()));
+
+        // Route the response back to ourselves so it goes through the tun device.
+        zc_packet.mut_peer_manager_header().unwrap().to_peer_id = self.my_peer_id.into();
+
+        Some(())
+    }
+
+    /// Extract the DNS request message from a UDP packet and send it to the catalog.
+    /// Replace the content of the UDP packet with the response message.
+    async fn handle_udp_packet(
+        &self,
+        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();
+
+            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)),
+                    Protocol::Udp,
+                ),
+                request_payload.len(),
+            )
+        };
+
+        if !self
+            .is_hijacked_addr(&SocketAddr::new(dst_ip.into(), dst_port).into())
+            .await
+        {
+            return None;
+        }
+
+        let response_payload = {
+            let response = Response::new(512);
+
+            self.catalog
+                .handle_request(&request, response.clone())
+                .await;
+
+            response.into_inner()?
+        };
+
+        let response_length = response_payload.len();
+        let delta_length = response_length as isize - request_length as isize;
+
+        // Resize the packet buffer to accommodate the response.
+        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(())
+    }
+
+    /// Handle ICMP echo request by turning it into an echo reply.
+    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(())
+    }
+}