From 91994f875fe1be6a70ee689c4e4b30ec787f9f79 Mon Sep 17 00:00:00 2001 From: Xiaobo Liu Date: Wed, 17 Jan 2024 22:37:16 +0800 Subject: [PATCH] server: refactor reject packet sender (#14) Signed-off-by: Xiaobo Liu --- server/src/reject/message.rs | 20 +++++ server/src/reject/mod.rs | 5 ++ server/src/{reject.rs => reject/sender.rs} | 92 ++++++++++------------ server/src/util.rs | 7 +- 4 files changed, 71 insertions(+), 53 deletions(-) create mode 100644 server/src/reject/message.rs create mode 100644 server/src/reject/mod.rs rename server/src/{reject.rs => reject/sender.rs} (71%) diff --git a/server/src/reject/message.rs b/server/src/reject/message.rs new file mode 100644 index 0000000..0f1dc8e --- /dev/null +++ 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 new file mode 100644 index 0000000..4f7d94b --- /dev/null +++ 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 b/server/src/reject/sender.rs similarity index 71% rename from server/src/reject.rs rename to server/src/reject/sender.rs index f7177ec..2bd0fb1 100644 --- a/server/src/reject.rs +++ b/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 index 43b93cb..cd70d40 100644 --- a/server/src/util.rs +++ b/server/src/util.rs @@ -111,12 +111,11 @@ pub fn set_rlimit_nofile(n: u64) -> Result { #[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(