server: refactor reject packet sender (#14)

Signed-off-by: Xiaobo Liu <[email protected]>
cppcoffee · Jan 17, 2024 · 91994f8 · 91994f8
1 parent 32be8c1
commit 91994f8
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 53 deletions.
diff --git a/server/src/reject/message.rs b/server/src/reject/message.rs
@@ -0,0 +1,20 @@
+use std::net::IpAddr;
+
+use pnet::packet::icmp::MutableIcmpPacket;
+use pnet::packet::icmpv6::MutableIcmpv6Packet;
+use pnet::packet::tcp::MutableTcpPacket;
+
+pub enum Message {
+    Icmp {
+        destination: IpAddr,
+        icmp_packet: MutableIcmpPacket<'static>,
+    },
+    Icmpv6 {
+        destination: IpAddr,
+        icmp_packet: MutableIcmpv6Packet<'static>,
+    },
+    Tcp {
+        destination: IpAddr,
+        tcp_packet: MutableTcpPacket<'static>,
+    },
+}
diff --git a/server/src/reject/mod.rs b/server/src/reject/mod.rs
@@ -0,0 +1,5 @@
+pub mod message;
+pub mod sender;
+
+pub use message::Message;
+pub use sender::RejectPacketSender;
diff --git a/server/src/reject.rs → server/src/reject/sender.rs b/server/src/reject.rs → server/src/reject/sender.rs
@@ -3,7 +3,7 @@ use std::net::IpAddr;
 use std::sync::mpsc;
 use std::{ptr, thread};
 
-use anyhow::{anyhow, bail, Result};
+use anyhow::{anyhow, bail, Context, Result};
 use log::error;
 use pnet::packet::icmp::MutableIcmpPacket;
 use pnet::packet::icmpv6::MutableIcmpv6Packet;
@@ -13,6 +13,7 @@ use pnet::transport::{
     transport_channel, TransportChannelType, TransportProtocol, TransportSender,
 };
 
+use crate::reject::Message;
 use crate::util;
 
 const ICMP_UNREACHABLE_HEADER_SIZE: usize = 8;
@@ -97,7 +98,7 @@ impl RejectPacketSender {
     ) -> Result<()> {
         let tcp_min_size = TcpPacket::minimum_packet_size();
 
-        let buffer = Box::new(MaybeUninit::<[u8; BUFFER_SIZE]>::uninit());
+        let buffer = MaybeUninit::<[u8; BUFFER_SIZE]>::uninit();
         let buffer = unsafe { ptr::read(buffer.as_ptr() as *const [u8; BUFFER_SIZE]) };
 
         let mut tcp_packet = MutableTcpPacket::owned(buffer[..tcp_min_size].to_vec())
@@ -114,21 +115,6 @@ impl RejectPacketSender {
     }
 }
 
-enum Message {
-    Icmp {
-        destination: IpAddr,
-        icmp_packet: MutableIcmpPacket<'static>,
-    },
-    Icmpv6 {
-        destination: IpAddr,
-        icmp_packet: MutableIcmpv6Packet<'static>,
-    },
-    Tcp {
-        destination: IpAddr,
-        tcp_packet: MutableTcpPacket<'static>,
-    },
-}
-
 struct Sender {
     icmp: TransportSender,
     tcp: TransportSender,
@@ -171,42 +157,50 @@ impl Sender {
 
         thread::spawn(move || {
             for msg in rx {
-                match msg {
-                    Message::Icmp {
-                        destination,
-                        icmp_packet,
-                    } => {
-                        if let Err(e) = self.icmp.send_to(icmp_packet, destination) {
-                            error!("Failed to send ICMP packet: {:?}", e);
-                        }
-                    }
-                    Message::Icmpv6 {
-                        destination,
-                        icmp_packet,
-                    } => {
-                        if let Err(e) = self.icmpv6.send_to(icmp_packet, destination) {
-                            error!("Failed to send ICMPv6 packet: {:?}", e);
-                        }
-                    }
-                    Message::Tcp {
-                        destination,
-                        tcp_packet,
-                    } => match destination {
-                        IpAddr::V4(_) => {
-                            if let Err(e) = self.tcp.send_to(tcp_packet, destination) {
-                                error!("Failed to send TCP packet: {:?}", e);
-                            }
-                        }
-                        IpAddr::V6(_) => {
-                            if let Err(e) = self.tcp6.send_to(tcp_packet, destination) {
-                                error!("Failed to send TCP packet: {:?}", e);
-                            }
-                        }
-                    },
+                if let Err(e) = self.message_handler(msg) {
+                    error!("Failed to handle message: {:?}", e);
                 }
             }
         });
 
         tx
     }
+
+    fn message_handler(&mut self, msg: Message) -> Result<()> {
+        match msg {
+            Message::Icmp {
+                destination,
+                icmp_packet,
+            } => {
+                self.icmp
+                    .send_to(icmp_packet, destination)
+                    .context("send ICMP packet")?;
+            }
+            Message::Icmpv6 {
+                destination,
+                icmp_packet,
+            } => {
+                self.icmpv6
+                    .send_to(icmp_packet, destination)
+                    .context("send ICMPv6 packet")?;
+            }
+            Message::Tcp {
+                destination,
+                tcp_packet,
+            } => match destination {
+                IpAddr::V4(_) => {
+                    self.tcp
+                        .send_to(tcp_packet, destination)
+                        .context("send TCP packet")?;
+                }
+                IpAddr::V6(_) => {
+                    self.tcp6
+                        .send_to(tcp_packet, destination)
+                        .context("send TCPv6 packet")?;
+                }
+            },
+        }
+
+        Ok(())
+    }
 }
diff --git a/server/src/util.rs b/server/src/util.rs
@@ -111,12 +111,11 @@ pub fn set_rlimit_nofile(n: u64) -> Result<u64> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use pnet::packet::udp::UdpPacket;
     use std::net::{IpAddr, Ipv4Addr};
 
     #[test]
     fn test_build_icmpv4_unreachable() {
-        let mut data = vec![0u8; 128];
+        let data = vec![0u8; 128];
         let mut icmp_packet = MutableIcmpPacket::owned(data.clone()).unwrap();
 
         build_icmpv4_unreachable(&mut icmp_packet);
@@ -136,7 +135,7 @@ mod tests {
 
     #[test]
     fn test_build_icmpv6_unreachable() {
-        let mut data = vec![0u8; 128];
+        let data = vec![0u8; 128];
         let src = Ipv6Addr::new(1, 1, 1, 1, 1, 1, 1, 1);
         let dest = Ipv6Addr::new(2, 2, 2, 2, 2, 2, 2, 2);
 
@@ -157,7 +156,7 @@ mod tests {
         let tcp_min_size = TcpPacket::minimum_packet_size();
         let incoming = TcpPacket::new(&[0u8; 40]).unwrap();
 
-        let mut buffer = vec![0u8; 128];
+        let buffer = vec![0u8; 128];
         let mut tcp_reset = MutableTcpPacket::owned(buffer[..tcp_min_size].to_owned()).unwrap();
 
         build_tcp_reset(