Fix credential ospf logic, fix udp subnet proxy loop protection (#2315)

easytier/src/gateway/udp_proxy.rs

@@ -40,6 +40,16 @@ use crate::{
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 struct UdpNatKey {
     src_socket: SocketAddr,
+    dst_socket: SocketAddr,
+}
+
+impl UdpNatKey {
+    fn new(src_socket: SocketAddr, dst_socket: SocketAddr) -> Self {
+        Self {
+            src_socket,
+            dst_socket,
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -204,13 +214,23 @@ impl UdpNatEntry {
 
                 self_clone.mark_active();
 
-                if src_v4.ip().is_loopback() {
-                    src_v4.set_ip(virtual_ipv4);
+                let has_mapped_dst = real_ipv4 != mapped_ipv4;
+                let mut reply_src_ip = *src_v4.ip();
+
+                // Preserve the existing priority for proxy rules that expose a
+                // real loopback address as a mapped address. Other loopback
+                // replies come from local delivery to 127.0.0.1 for the local
+                // virtual IP and may need the mapped rewrite below.
+                if has_mapped_dst && reply_src_ip == real_ipv4 {
+                    reply_src_ip = mapped_ipv4;
+                } else if reply_src_ip.is_loopback() {
+                    reply_src_ip = virtual_ipv4;
                 }
 
-                if *src_v4.ip() == real_ipv4 {
-                    src_v4.set_ip(mapped_ipv4);
+                if has_mapped_dst && reply_src_ip == real_ipv4 {
+                    reply_src_ip = mapped_ipv4;
                 }
+                src_v4.set_ip(reply_src_ip);
 
                 let Ok(_) = Self::compose_ipv4_packet(
                     &self_clone,
@@ -321,9 +341,10 @@ impl UdpProxy {
             "udp nat packet request received"
         );
 
-        let nat_key = UdpNatKey {
-            src_socket: SocketAddr::new(ipv4.get_source().into(), udp_packet.get_source()),
-        };
+        let nat_key = UdpNatKey::new(
+            SocketAddr::new(ipv4.get_source().into(), udp_packet.get_source()),
+            SocketAddr::new(ipv4.get_destination().into(), udp_packet.get_destination()),
+        );
         let nat_entry = self
             .nat_table
             .entry(nat_key)
@@ -487,3 +508,309 @@ impl Drop for UdpProxy {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::{
+        net::{Ipv4Addr, SocketAddr},
+        sync::Arc,
+        time::Duration,
+    };
+
+    use pnet::packet::{
+        MutablePacket, Packet,
+        ip::IpNextHeaderProtocols,
+        ipv4::{self, Ipv4Packet, MutableIpv4Packet},
+        udp::{self, MutableUdpPacket, UdpPacket},
+    };
+    use tokio::{net::UdpSocket, sync::mpsc::Receiver, time::timeout};
+
+    use crate::{
+        common::{config::ConfigLoader, global_ctx::tests::get_mock_global_ctx},
+        peers::{
+            create_packet_recv_chan,
+            peer_manager::{PeerManager, RouteAlgoType},
+        },
+        tunnel::packet_def::{PacketType, ZCPacket},
+    };
+
+    use super::UdpProxy;
+
+    fn build_udp_proxy_packet(
+        src_ip: Ipv4Addr,
+        src_port: u16,
+        dst_socket: SocketAddr,
+        payload: &[u8],
+    ) -> ZCPacket {
+        let SocketAddr::V4(dst_socket) = dst_socket else {
+            panic!("test only builds IPv4 UDP packets");
+        };
+        let dst_ip = *dst_socket.ip();
+        let mut packet = vec![0; 20 + 8 + payload.len()];
+        let packet_len = packet.len() as u16;
+
+        {
+            let mut ipv4_packet = MutableIpv4Packet::new(&mut packet).unwrap();
+            ipv4_packet.set_version(4);
+            ipv4_packet.set_header_length(5);
+            ipv4_packet.set_total_length(packet_len);
+            ipv4_packet.set_ttl(64);
+            ipv4_packet.set_next_level_protocol(IpNextHeaderProtocols::Udp);
+            ipv4_packet.set_source(src_ip);
+            ipv4_packet.set_destination(dst_ip);
+        }
+
+        {
+            let mut udp_packet = MutableUdpPacket::new(&mut packet[20..]).unwrap();
+            udp_packet.set_source(src_port);
+            udp_packet.set_destination(dst_socket.port());
+            udp_packet.set_length((8 + payload.len()) as u16);
+            udp_packet.payload_mut().copy_from_slice(payload);
+            udp_packet.set_checksum(udp::ipv4_checksum(
+                &udp_packet.to_immutable(),
+                &src_ip,
+                &dst_ip,
+            ));
+        }
+
+        {
+            let mut ipv4_packet = MutableIpv4Packet::new(&mut packet).unwrap();
+            ipv4_packet.set_checksum(ipv4::checksum(&ipv4_packet.to_immutable()));
+        }
+
+        let mut packet = ZCPacket::new_with_payload(&packet);
+        packet.fill_peer_manager_hdr(1009867077, 3831440917, PacketType::Data as u8);
+        packet
+    }
+
+    async fn wait_proxy_cidr_loaded(proxy: &UdpProxy) {
+        timeout(Duration::from_secs(1), async {
+            while proxy.cidr_set.is_empty() {
+                tokio::time::sleep(Duration::from_millis(10)).await;
+            }
+        })
+        .await
+        .unwrap();
+    }
+
+    async fn recv_payload(socket: &UdpSocket) -> (Vec<u8>, SocketAddr) {
+        let mut buf = [0; 64];
+        let (len, addr) = timeout(Duration::from_secs(1), socket.recv_from(&mut buf))
+            .await
+            .unwrap()
+            .unwrap();
+        (buf[..len].to_vec(), addr)
+    }
+
+    async fn recv_response_packet(receiver: &mut Receiver<ZCPacket>) -> ZCPacket {
+        timeout(Duration::from_secs(1), receiver.recv())
+            .await
+            .unwrap()
+            .unwrap()
+    }
+
+    fn assert_udp_response(
+        packet: ZCPacket,
+        src_socket: SocketAddr,
+        dst_ip: Ipv4Addr,
+        dst_port: u16,
+        payload: &[u8],
+    ) {
+        let SocketAddr::V4(src_socket) = src_socket else {
+            panic!("test only checks IPv4 UDP packets");
+        };
+        let ipv4_packet = Ipv4Packet::new(packet.payload()).unwrap();
+        assert_eq!(ipv4_packet.get_source(), *src_socket.ip());
+        assert_eq!(ipv4_packet.get_destination(), dst_ip);
+
+        let udp_packet = UdpPacket::new(ipv4_packet.payload()).unwrap();
+        assert_eq!(udp_packet.get_source(), src_socket.port());
+        assert_eq!(udp_packet.get_destination(), dst_port);
+        assert_eq!(udp_packet.payload(), payload);
+    }
+
+    async fn stop_nat_entries(proxy: &UdpProxy) {
+        let nat_socket_addrs = proxy
+            .nat_table
+            .iter()
+            .filter_map(|entry| {
+                entry
+                    .socket
+                    .as_ref()
+                    .and_then(|socket| socket.local_addr().ok())
+                    .map(|addr| SocketAddr::from((Ipv4Addr::LOCALHOST, addr.port())))
+            })
+            .collect::<Vec<_>>();
+
+        for entry in proxy.nat_table.iter() {
+            entry.stop();
+        }
+
+        let wake_socket = UdpSocket::bind((Ipv4Addr::LOCALHOST, 0)).await.unwrap();
+        for addr in nat_socket_addrs {
+            let _ = wake_socket.send_to(b"wake", addr).await;
+        }
+    }
+
+    #[tokio::test]
+    async fn udp_proxy_rewrites_unmapped_loopback_reply_to_virtual_ip() {
+        let global_ctx = get_mock_global_ctx();
+        global_ctx.set_ipv4(Some("10.144.144.204/24".parse().unwrap()));
+        global_ctx
+            .config
+            .add_proxy_cidr("127.0.0.1/32".parse().unwrap(), None)
+            .unwrap();
+
+        let (packet_sender, _packet_receiver) = create_packet_recv_chan();
+        let peer_manager = Arc::new(PeerManager::new(
+            RouteAlgoType::Ospf,
+            global_ctx.clone(),
+            packet_sender,
+        ));
+        let proxy = UdpProxy::new(global_ctx, peer_manager).unwrap();
+        wait_proxy_cidr_loaded(&proxy).await;
+        let mut response_receiver = proxy.receiver.lock().await.take().unwrap();
+
+        let real_dst = UdpSocket::bind((Ipv4Addr::LOCALHOST, 0)).await.unwrap();
+        let real_dst_port = real_dst.local_addr().unwrap().port();
+        let dst_socket = SocketAddr::from((Ipv4Addr::LOCALHOST, real_dst_port));
+        let src_ip = Ipv4Addr::new(10, 144, 144, 206);
+        let src_port = 53864;
+
+        let packet = build_udp_proxy_packet(src_ip, src_port, dst_socket, b"request");
+        assert!(proxy.try_handle_packet(&packet).await.is_some());
+        let (payload, nat_socket) = recv_payload(&real_dst).await;
+        assert_eq!(payload, b"request");
+
+        real_dst.send_to(b"reply", nat_socket).await.unwrap();
+        assert_udp_response(
+            recv_response_packet(&mut response_receiver).await,
+            SocketAddr::from((Ipv4Addr::new(10, 144, 144, 204), real_dst_port)),
+            src_ip,
+            src_port,
+            b"reply",
+        );
+
+        stop_nat_entries(&proxy).await;
+    }
+
+    #[tokio::test]
+    async fn udp_proxy_maps_local_virtual_destination_reply_to_mapped_source() {
+        let global_ctx = get_mock_global_ctx();
+        global_ctx.set_ipv4(Some("10.144.144.204/24".parse().unwrap()));
+        global_ctx
+            .config
+            .add_proxy_cidr(
+                "10.144.144.204/32".parse().unwrap(),
+                Some("10.10.10.3/32".parse().unwrap()),
+            )
+            .unwrap();
+
+        let (packet_sender, _packet_receiver) = create_packet_recv_chan();
+        let peer_manager = Arc::new(PeerManager::new(
+            RouteAlgoType::Ospf,
+            global_ctx.clone(),
+            packet_sender,
+        ));
+        let proxy = UdpProxy::new(global_ctx, peer_manager).unwrap();
+        wait_proxy_cidr_loaded(&proxy).await;
+        let mut response_receiver = proxy.receiver.lock().await.take().unwrap();
+
+        let real_dst = UdpSocket::bind((Ipv4Addr::LOCALHOST, 0)).await.unwrap();
+        let real_dst_port = real_dst.local_addr().unwrap().port();
+        let mapped_dst = SocketAddr::from((Ipv4Addr::new(10, 10, 10, 3), real_dst_port));
+        let src_ip = Ipv4Addr::new(10, 144, 144, 206);
+        let src_port = 53864;
+
+        let packet = build_udp_proxy_packet(src_ip, src_port, mapped_dst, b"request");
+        assert!(proxy.try_handle_packet(&packet).await.is_some());
+        let (payload, nat_socket) = recv_payload(&real_dst).await;
+        assert_eq!(payload, b"request");
+
+        real_dst.send_to(b"reply", nat_socket).await.unwrap();
+        assert_udp_response(
+            recv_response_packet(&mut response_receiver).await,
+            mapped_dst,
+            src_ip,
+            src_port,
+            b"reply",
+        );
+
+        stop_nat_entries(&proxy).await;
+    }
+
+    #[tokio::test]
+    async fn udp_proxy_separates_same_source_port_to_multiple_mapped_destinations() {
+        let global_ctx = get_mock_global_ctx();
+        global_ctx.set_ipv4(Some("10.144.144.204/24".parse().unwrap()));
+        global_ctx
+            .config
+            .add_proxy_cidr(
+                "127.0.0.1/32".parse().unwrap(),
+                Some("10.10.10.1/32".parse().unwrap()),
+            )
+            .unwrap();
+        global_ctx
+            .config
+            .add_proxy_cidr(
+                "127.0.0.1/32".parse().unwrap(),
+                Some("10.10.10.2/32".parse().unwrap()),
+            )
+            .unwrap();
+
+        let (packet_sender, _packet_receiver) = create_packet_recv_chan();
+        let peer_manager = Arc::new(PeerManager::new(
+            RouteAlgoType::Ospf,
+            global_ctx.clone(),
+            packet_sender,
+        ));
+        let proxy = UdpProxy::new(global_ctx, peer_manager).unwrap();
+        wait_proxy_cidr_loaded(&proxy).await;
+        let mut response_receiver = proxy.receiver.lock().await.take().unwrap();
+
+        let real_dst = UdpSocket::bind((Ipv4Addr::LOCALHOST, 0)).await.unwrap();
+        let real_dst_port = real_dst.local_addr().unwrap().port();
+        let first_mapped_dst = SocketAddr::from((Ipv4Addr::new(10, 10, 10, 1), real_dst_port));
+        let second_mapped_dst = SocketAddr::from((Ipv4Addr::new(10, 10, 10, 2), real_dst_port));
+        let src_ip = Ipv4Addr::new(10, 144, 144, 206);
+        let src_port = 53864;
+
+        let first_packet = build_udp_proxy_packet(src_ip, src_port, first_mapped_dst, b"first");
+        assert!(proxy.try_handle_packet(&first_packet).await.is_some());
+        let (payload, first_nat_socket) = recv_payload(&real_dst).await;
+        assert_eq!(payload, b"first");
+
+        let second_packet = build_udp_proxy_packet(src_ip, src_port, second_mapped_dst, b"second");
+        assert!(proxy.try_handle_packet(&second_packet).await.is_some());
+        let (payload, second_nat_socket) = recv_payload(&real_dst).await;
+        assert_eq!(payload, b"second");
+
+        assert_eq!(proxy.nat_table.len(), 2);
+
+        real_dst
+            .send_to(b"first-reply", first_nat_socket)
+            .await
+            .unwrap();
+        assert_udp_response(
+            recv_response_packet(&mut response_receiver).await,
+            first_mapped_dst,
+            src_ip,
+            src_port,
+            b"first-reply",
+        );
+
+        real_dst
+            .send_to(b"second-reply", second_nat_socket)
+            .await
+            .unwrap();
+        assert_udp_response(
+            recv_response_packet(&mut response_receiver).await,
+            second_mapped_dst,
+            src_ip,
+            src_port,
+            b"second-reply",
+        );
+
+        stop_nat_entries(&proxy).await;
+    }
+}