network: implement icmp nat support

network/src/icmp.rs

@@ -0,0 +1,245 @@
+use crate::raw_socket::{RawSocketHandle, RawSocketProtocol};
+use anyhow::{anyhow, Result};
+use etherparse::{
+    IcmpEchoHeader, Icmpv4Header, Icmpv4Slice, Icmpv4Type, Icmpv6Header, Icmpv6Slice, Icmpv6Type,
+    IpNumber, NetSlice, SlicedPacket,
+};
+use log::warn;
+use std::{
+    collections::HashMap,
+    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
+    os::fd::{FromRawFd, IntoRawFd},
+    sync::Arc,
+    time::Duration,
+};
+use tokio::{
+    net::UdpSocket,
+    sync::{oneshot, Mutex},
+    task::JoinHandle,
+    time::timeout,
+};
+
+#[derive(Debug)]
+pub enum IcmpProtocol {
+    Icmp4,
+    Icmp6,
+}
+
+impl IcmpProtocol {
+    pub fn to_socket_protocol(&self) -> RawSocketProtocol {
+        match self {
+            IcmpProtocol::Icmp4 => RawSocketProtocol::Icmpv4,
+            IcmpProtocol::Icmp6 => RawSocketProtocol::Icmpv6,
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct IcmpHandlerToken(IpAddr, Option<u16>, u16);
+
+#[derive(Debug)]
+pub enum IcmpReply {
+    Icmp4 {
+        header: Icmpv4Header,
+        echo: IcmpEchoHeader,
+        payload: Vec<u8>,
+    },
+
+    Icmp6 {
+        header: Icmpv6Header,
+        echo: IcmpEchoHeader,
+        payload: Vec<u8>,
+    },
+}
+
+type IcmpHandlerMap = Arc<Mutex<HashMap<IcmpHandlerToken, oneshot::Sender<IcmpReply>>>>;
+pub struct IcmpClient {
+    socket: Arc<UdpSocket>,
+    handlers: IcmpHandlerMap,
+    task: Arc<JoinHandle<Result<()>>>,
+}
+
+impl IcmpClient {
+    pub fn new(protocol: IcmpProtocol) -> Result<IcmpClient> {
+        let handle = RawSocketHandle::new(protocol.to_socket_protocol())?;
+        let socket = unsafe { std::net::UdpSocket::from_raw_fd(handle.into_raw_fd()) };
+        let socket: Arc<UdpSocket> = Arc::new(socket.try_into()?);
+        let handlers = Arc::new(Mutex::new(HashMap::new()));
+        let task = Arc::new(tokio::task::spawn(IcmpClient::process(
+            protocol,
+            socket.clone(),
+            handlers.clone(),
+        )));
+        Ok(IcmpClient {
+            socket,
+            handlers,
+            task,
+        })
+    }
+
+    async fn process(
+        protocol: IcmpProtocol,
+        socket: Arc<UdpSocket>,
+        handlers: IcmpHandlerMap,
+    ) -> Result<()> {
+        let mut buffer = vec![0u8; 2048];
+        loop {
+            let (size, addr) = socket.recv_from(&mut buffer).await?;
+            let packet = &buffer[0..size];
+
+            let (token, reply) = match protocol {
+                IcmpProtocol::Icmp4 => {
+                    let sliced = match SlicedPacket::from_ip(packet) {
+                        Ok(sliced) => sliced,
+                        Err(error) => {
+                            warn!("received icmp packet but failed to parse it: {}", error);
+                            continue;
+                        }
+                    };
+
+                    let Some(NetSlice::Ipv4(ipv4)) = sliced.net else {
+                        continue;
+                    };
+
+                    if ipv4.header().protocol() != IpNumber::ICMP {
+                        continue;
+                    }
+
+                    let Ok(icmpv4) = Icmpv4Slice::from_slice(ipv4.payload().payload) else {
+                        continue;
+                    };
+
+                    let Icmpv4Type::EchoReply(echo) = icmpv4.header().icmp_type else {
+                        continue;
+                    };
+
+                    let token = IcmpHandlerToken(
+                        IpAddr::V4(ipv4.header().source_addr()),
+                        Some(echo.id),
+                        echo.seq,
+                    );
+                    let reply = IcmpReply::Icmp4 {
+                        header: icmpv4.header(),
+                        echo,
+                        payload: icmpv4.payload().to_vec(),
+                    };
+                    (token, reply)
+                }
+
+                IcmpProtocol::Icmp6 => {
+                    let Ok(icmpv6) = Icmpv6Slice::from_slice(packet) else {
+                        continue;
+                    };
+
+                    let Icmpv6Type::EchoReply(echo) = icmpv6.header().icmp_type else {
+                        continue;
+                    };
+
+                    let SocketAddr::V6(addr) = addr else {
+                        continue;
+                    };
+
+                    let token = IcmpHandlerToken(IpAddr::V6(*addr.ip()), Some(echo.id), echo.seq);
+
+                    let reply = IcmpReply::Icmp6 {
+                        header: icmpv6.header(),
+                        echo,
+                        payload: icmpv6.payload().to_vec(),
+                    };
+                    (token, reply)
+                }
+            };
+
+            if let Some(sender) = handlers.lock().await.remove(&token) {
+                let _ = sender.send(reply);
+            }
+        }
+    }
+
+    async fn add_handler(&self, token: IcmpHandlerToken) -> Result<oneshot::Receiver<IcmpReply>> {
+        let (tx, rx) = oneshot::channel();
+        if self
+            .handlers
+            .lock()
+            .await
+            .insert(token.clone(), tx)
+            .is_some()
+        {
+            return Err(anyhow!("duplicate icmp request: {:?}", token));
+        }
+        Ok(rx)
+    }
+
+    async fn remove_handler(&self, token: IcmpHandlerToken) -> Result<()> {
+        self.handlers.lock().await.remove(&token);
+        Ok(())
+    }
+
+    pub async fn ping4(
+        &self,
+        addr: Ipv4Addr,
+        id: u16,
+        seq: u16,
+        payload: &[u8],
+        deadline: Duration,
+    ) -> Result<Option<IcmpReply>> {
+        let token = IcmpHandlerToken(IpAddr::V4(addr), Some(id), seq);
+        let rx = self.add_handler(token.clone()).await?;
+
+        let echo = IcmpEchoHeader { id, seq };
+        let header = Icmpv4Header::new(Icmpv4Type::EchoRequest(echo));
+        let mut buffer: Vec<u8> = Vec::new();
+        header.write(&mut buffer)?;
+        buffer.extend_from_slice(payload);
+        self.socket
+            .send_to(&buffer, SocketAddr::V4(SocketAddrV4::new(addr, 0)))
+            .await?;
+
+        let result = timeout(deadline, rx).await;
+        self.remove_handler(token).await?;
+        let reply = match result {
+            Ok(Ok(packet)) => Some(packet),
+            Ok(Err(err)) => return Err(anyhow!("failed to wait for icmp packet: {}", err)),
+            Err(_) => None,
+        };
+        Ok(reply)
+    }
+
+    pub async fn ping6(
+        &self,
+        addr: Ipv6Addr,
+        id: u16,
+        seq: u16,
+        payload: &[u8],
+        deadline: Duration,
+    ) -> Result<Option<IcmpReply>> {
+        let token = IcmpHandlerToken(IpAddr::V6(addr), Some(id), seq);
+        let rx = self.add_handler(token.clone()).await?;
+
+        let echo = IcmpEchoHeader { id, seq };
+        let header = Icmpv6Header::new(Icmpv6Type::EchoRequest(echo));
+        let mut buffer: Vec<u8> = Vec::new();
+        header.write(&mut buffer)?;
+        buffer.extend_from_slice(payload);
+        self.socket
+            .send_to(&buffer, SocketAddr::V6(SocketAddrV6::new(addr, 0, 0, 0)))
+            .await?;
+
+        let result = timeout(deadline, rx).await;
+        self.remove_handler(token).await?;
+        let reply = match result {
+            Ok(Ok(packet)) => Some(packet),
+            Ok(Err(err)) => return Err(anyhow!("failed to wait for icmp packet: {}", err)),
+            Err(_) => None,
+        };
+        Ok(reply)
+    }
+}
+
+impl Drop for IcmpClient {
+    fn drop(&mut self) {
+        if Arc::strong_count(&self.task) <= 1 {
+            self.task.abort();
+        }
+    }
+}