From fc4be559433a0ac417ab6e08670876e49304c5a9 Mon Sep 17 00:00:00 2001 From: SajjadPourali Date: Mon, 2 Nov 2026 21:11:14 -0700 Subject: [PATCH] Optimize outgoing packet serialization --- src/lib.rs | 63 ++++++++++++++++++++++++++++++++++++++++++++------- src/packet.rs | 30 +++++++++++++++++++----- 2 files changed, 79 insertions(+), 14 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 8f09a5f..d67d7b2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -314,6 +314,7 @@ fn run( let pi = config.packet_information; let offset = if pi && cfg!(unix) { 4 } else { 0 }; let mut buffer = vec![0_u8; config.mtu as usize + offset]; + let mut outgoing_buffer = Vec::with_capacity(config.mtu as usize + offset); let (up_pkt_sender, mut up_pkt_receiver) = mpsc::unbounded_channel::(); tokio::spawn(async move { @@ -348,7 +349,7 @@ fn run( remove_closed_session(&mut sessions, network_tuple); } Some(packet) = up_pkt_receiver.recv() => { - process_upstream_recv(packet, &mut device, #[cfg(unix)]pi).await?; + process_upstream_recv(packet, &mut device, &mut outgoing_buffer, #[cfg(unix)]pi).await?; } } } @@ -477,22 +478,33 @@ fn create_stream( async fn process_upstream_recv( up_packet: NetworkPacket, device: &mut Device, + packet_bytes: &mut Vec, #[cfg(unix)] packet_information: bool, ) -> Result<()> { - #[allow(unused_mut)] - let Ok(mut packet_bytes) = up_packet.to_bytes() else { - log::warn!("to_bytes error"); - return Ok(()); + #[cfg(unix)] + let metadata_len = if packet_information { + TUN_FLAGS.len() + TUN_PROTO_IP4.len() + } else { + 0 }; + #[cfg(not(unix))] + let metadata_len = 0; + packet_bytes.clear(); + packet_bytes.reserve(metadata_len + up_packet.serialized_len()); #[cfg(unix)] if packet_information { + packet_bytes.extend_from_slice(&TUN_FLAGS); if up_packet.src_addr().is_ipv4() { - packet_bytes.splice(0..0, [TUN_FLAGS, TUN_PROTO_IP4].concat()); + packet_bytes.extend_from_slice(&TUN_PROTO_IP4); } else { - packet_bytes.splice(0..0, [TUN_FLAGS, TUN_PROTO_IP6].concat()); + packet_bytes.extend_from_slice(&TUN_PROTO_IP6); } } - device.write_all(&packet_bytes).await?; + if up_packet.write_to(packet_bytes).is_err() { + log::warn!("to_bytes error"); + return Ok(()); + } + device.write_all(packet_bytes).await?; // device.flush().await?; Ok(()) @@ -528,6 +540,41 @@ mod tests { .unwrap() } + #[cfg(unix)] + #[tokio::test] + async fn upstream_packet_information_precedes_serialized_ipv4_and_ipv6_packets() { + use tokio::io::AsyncReadExt; + + let mut buffer = Vec::new(); + for ipv6 in [false, true] { + let mut raw = Vec::new(); + if ipv6 { + etherparse::PacketBuilder::ipv6([0; 16], [1; 16], 64) + .udp(1234, 5678) + .write(&mut raw, b"payload") + .unwrap(); + } else { + raw = udp_packet(b"payload"); + } + for packet_information in [false, true] { + let packet = NetworkPacket::parse(&raw).unwrap(); + let mut expected = Vec::new(); + if packet_information { + expected.extend_from_slice(&TUN_FLAGS); + expected.extend_from_slice(if ipv6 { &TUN_PROTO_IP6 } else { &TUN_PROTO_IP4 }); + } + expected.extend_from_slice(&packet.to_bytes().unwrap()); + let (mut device, mut peer) = duplex(256); + process_upstream_recv(packet, &mut device, &mut buffer, packet_information) + .await + .unwrap(); + let mut actual = vec![0; expected.len()]; + peer.read_exact(&mut actual).await.unwrap(); + assert_eq!(actual, expected); + } + } + } + /// Exercise dispatch before the old removal notification is consumed, then pass that /// notification through the same handler used by the event loop. #[tokio::test] diff --git a/src/packet.rs b/src/packet.rs index c2d89fc..88f9707 100644 --- a/src/packet.rs +++ b/src/packet.rs @@ -137,21 +137,39 @@ impl NetworkPacket { } } pub fn to_bytes(&self) -> Result, IpStackError> { - let mut buf = Vec::new(); + let mut buf = Vec::with_capacity(self.serialized_len()); + self.write_to(&mut buf)?; + Ok(buf) + } + + pub(crate) fn serialized_len(&self) -> usize { + let ip_len = match &self.ip { + IpHeader::Ipv4(ip) => ip.header_len(), + IpHeader::Ipv6(ip) => ip.header_len(), + }; + let transport_len = match &self.transport { + TransportHeader::Tcp(tcp) => tcp.header_len(), + TransportHeader::Udp(_) => UdpHeader::LEN, + TransportHeader::Unknown => 0, + }; + ip_len + transport_len + self.payload.as_ref().map_or(0, Vec::len) + } + + pub(crate) fn write_to(&self, buf: &mut Vec) -> Result<(), IpStackError> { match self.ip { - IpHeader::Ipv4(ref ip) => ip.write(&mut buf)?, - IpHeader::Ipv6(ref ip) => ip.write(&mut buf)?, + IpHeader::Ipv4(ref ip) => ip.write(&mut *buf)?, + IpHeader::Ipv6(ref ip) => ip.write(&mut *buf)?, } match self.transport { - TransportHeader::Tcp(ref h) => h.write(&mut buf)?, - TransportHeader::Udp(ref h) => h.write(&mut buf)?, + TransportHeader::Tcp(ref h) => h.write(&mut *buf)?, + TransportHeader::Udp(ref h) => h.write(&mut *buf)?, _ => {} }; if let Some(payload) = &self.payload { buf.extend_from_slice(payload); } - Ok(buf) + Ok(()) } pub fn ttl(&self) -> u8 { match &self.ip {