From 60c9aefddaeea20f1c464edf7df1589013a2df47 Mon Sep 17 00:00:00 2001 From: Alex Zenla Date: Mon, 12 Feb 2024 17:24:40 +0000 Subject: [PATCH] network: optimize proxynat packet reconstruction --- network/src/proxynat/icmp.rs | 18 +++++++++++------- network/src/proxynat/tcp.rs | 18 +++++++++--------- network/src/proxynat/udp.rs | 10 ++++++---- 3 files changed, 26 insertions(+), 20 deletions(-) diff --git a/network/src/proxynat/icmp.rs b/network/src/proxynat/icmp.rs index f5e5765..fa973af 100644 --- a/network/src/proxynat/icmp.rs +++ b/network/src/proxynat/icmp.rs @@ -5,7 +5,7 @@ use std::{ use anyhow::{anyhow, Result}; use async_trait::async_trait; -use bytes::BytesMut; +use bytes::{BufMut, BytesMut}; use etherparse::{ IcmpEchoHeader, Icmpv4Header, Icmpv4Type, Icmpv6Header, Icmpv6Type, IpNumber, Ipv4Slice, Ipv6Slice, NetSlice, PacketBuilder, SlicedPacket, @@ -221,9 +221,11 @@ impl ProxyIcmpHandler { } }; let packet = packet.icmpv4_echo_reply(echo.id, echo.seq); - let mut buffer: Vec = Vec::new(); - packet.write(&mut buffer, &payload)?; - if let Err(error) = context.try_send(buffer.as_slice().into()) { + let buffer = BytesMut::with_capacity(packet.size(payload.len())); + let mut writer = buffer.writer(); + packet.write(&mut writer, &payload)?; + let buffer = writer.into_inner(); + if let Err(error) = context.try_send(buffer) { debug!("failed to transmit icmp packet: {}", error); } Ok(()) @@ -264,9 +266,11 @@ impl ProxyIcmpHandler { } }; let packet = packet.icmpv6_echo_reply(echo.id, echo.seq); - let mut buffer: Vec = Vec::new(); - packet.write(&mut buffer, &payload)?; - if let Err(error) = context.try_send(buffer.as_slice().into()) { + let buffer = BytesMut::with_capacity(packet.size(payload.len())); + let mut writer = buffer.writer(); + packet.write(&mut writer, &payload)?; + let buffer = writer.into_inner(); + if let Err(error) = context.try_send(buffer) { debug!("failed to transmit icmp packet: {}", error); } Ok(()) diff --git a/network/src/proxynat/tcp.rs b/network/src/proxynat/tcp.rs index 4fbeafd..659aacd 100644 --- a/network/src/proxynat/tcp.rs +++ b/network/src/proxynat/tcp.rs @@ -203,7 +203,7 @@ impl ProxyTcpHandler { } ProxyTcpAcceptSelect::TxIpPacket(payload) => { - let mut buffer: Vec = Vec::new(); + let mut buffer = BytesMut::with_capacity(Ethernet2Header::LEN + payload.len()); let header = Ethernet2Header { source: context.key.local_mac.0, destination: context.key.client_mac.0, @@ -212,9 +212,9 @@ impl ProxyTcpHandler { IpAddress::Ipv6(_) => EtherType::IPV6, }, }; - header.write(&mut buffer)?; + buffer.extend_from_slice(&header.to_bytes()); buffer.extend_from_slice(&payload); - if let Err(error) = context.try_send(buffer.as_slice().into()) { + if let Err(error) = context.try_send(buffer) { debug!("failed to transmit tcp packet: {}", error); } } @@ -376,7 +376,7 @@ impl ProxyTcpHandler { } ProxyTcpDataSelect::TxIpPacket(payload) => { - let mut buffer: Vec = Vec::new(); + let mut buffer = BytesMut::with_capacity(Ethernet2Header::LEN + payload.len()); let header = Ethernet2Header { source: context.key.local_mac.0, destination: context.key.client_mac.0, @@ -385,9 +385,9 @@ impl ProxyTcpHandler { IpAddress::Ipv6(_) => EtherType::IPV6, }, }; - header.write(&mut buffer)?; + buffer.extend_from_slice(&header.to_bytes()); buffer.extend_from_slice(&payload); - if let Err(error) = context.try_send(buffer.as_slice().into()) { + if let Err(error) = context.try_send(buffer) { debug!("failed to transmit tcp packet: {}", error); } } @@ -436,7 +436,7 @@ impl ProxyTcpHandler { } ProxyTcpFinishSelect::TxIpPacket(payload) => { - let mut buffer: Vec = Vec::new(); + let mut buffer = BytesMut::with_capacity(Ethernet2Header::LEN + payload.len()); let header = Ethernet2Header { source: context.key.local_mac.0, destination: context.key.client_mac.0, @@ -445,9 +445,9 @@ impl ProxyTcpHandler { IpAddress::Ipv6(_) => EtherType::IPV6, }, }; - header.write(&mut buffer)?; + buffer.extend_from_slice(&header.to_bytes()); buffer.extend_from_slice(&payload); - if let Err(error) = context.try_send(buffer.as_slice().into()) { + if let Err(error) = context.try_send(buffer) { debug!("failed to transmit tcp packet: {}", error); } } diff --git a/network/src/proxynat/udp.rs b/network/src/proxynat/udp.rs index fb1ecd7..d902c25 100644 --- a/network/src/proxynat/udp.rs +++ b/network/src/proxynat/udp.rs @@ -5,7 +5,7 @@ use std::{ use anyhow::{anyhow, Result}; use async_trait::async_trait; -use bytes::BytesMut; +use bytes::{BufMut, BytesMut}; use etherparse::{PacketBuilder, SlicedPacket, UdpSlice}; use log::{debug, warn}; use smoltcp::wire::IpAddress; @@ -107,9 +107,11 @@ impl ProxyUdpHandler { }; let packet = packet.udp(context.key.external_ip.port, context.key.client_ip.port); - let mut buffer: Vec = Vec::new(); - packet.write(&mut buffer, data)?; - if let Err(error) = context.try_send(buffer.as_slice().into()) { + let buffer = BytesMut::with_capacity(packet.size(data.len())); + let mut writer = buffer.writer(); + packet.write(&mut writer, data)?; + let buffer = writer.into_inner(); + if let Err(error) = context.try_send(buffer) { debug!("failed to transmit udp packet: {}", error); } }