diff options
| -rw-r--r-- | src/main.rs | 12 | ||||
| -rw-r--r-- | src/platform.rs | 2 | ||||
| -rw-r--r-- | src/platform/linux.rs | 4 | ||||
| -rw-r--r-- | src/platform/windows.rs | 7 |
4 files changed, 16 insertions, 9 deletions
diff --git a/src/main.rs b/src/main.rs index dd3b471..12bd303 100644 --- a/src/main.rs +++ b/src/main.rs @@ -84,7 +84,7 @@ fn split_packet(pkt: &pkt::PktView, start: u32, end: Option<u32>, } /// Return Ok(true) if packet is handled -fn handle_packet(pkt: &[u8]) -> Result<bool> { +fn handle_packet(pkt: &[u8], buf: &mut Vec::<u8>) -> Result<bool> { use platform::send_to_raw; #[cfg(target_os = "linux")] @@ -102,14 +102,12 @@ fn handle_packet(pkt: &[u8]) -> Result<bool> { // TODO: if clienthello packet has been (unlikely) fragmented, // we should find the second part and drop, reassemble it here. - let mut buf = Vec::<u8>::with_capacity(2048); - - split_packet(&view, 0, Some(1), &mut buf)?; + split_packet(&view, 0, Some(1), buf)?; send_to_raw(&buf)?; std::thread::sleep(std::time::Duration::from_millis(delay_ms())); - split_packet(&view, 1, None, &mut buf)?; + split_packet(&view, 1, None, buf)?; send_to_raw(&buf)?; #[cfg(debug_assertions)] @@ -120,8 +118,8 @@ fn handle_packet(pkt: &[u8]) -> Result<bool> { #[macro_export] macro_rules! handle_packet { - ($bytes:expr, handled => $on_handled:expr, rejected => $on_rejected:expr $(,)?) => {{ - match handle_packet($bytes) { + ($bytes:expr, $buf:expr, handled => $on_handled:expr, rejected => $on_rejected:expr $(,)?) => {{ + match handle_packet($bytes, $buf) { Ok(true) => { $on_handled } Ok(false) => { $on_rejected } Err(e) => { diff --git a/src/platform.rs b/src/platform.rs index a618237..5b4b907 100644 --- a/src/platform.rs +++ b/src/platform.rs @@ -15,6 +15,8 @@ // You should have received a copy of the GNU General Public License // along with DPIBreak. If not, see <https://www.gnu.org/licenses/>. +pub const PACKET_SIZE_CAP: usize = 2048; + #[cfg(windows)] pub mod windows; diff --git a/src/platform/linux.rs b/src/platform/linux.rs index f19adc3..0faef31 100644 --- a/src/platform/linux.rs +++ b/src/platform/linux.rs @@ -213,6 +213,7 @@ pub fn run() -> Result<()> { }; use nfq::Queue; use crate::handle_packet; + use super::PACKET_SIZE_CAP; let mut q = Queue::open()?; q.bind(queue_num())?; @@ -227,6 +228,8 @@ pub fn run() -> Result<()> { splash!("{MESSAGE_AT_RUN}"); + let mut buf = Vec::<u8>::with_capacity(PACKET_SIZE_CAP); + while crate::RUNNING.load(Ordering::SeqCst) { { let fd = q.as_fd(); @@ -246,6 +249,7 @@ pub fn run() -> Result<()> { while let Ok(mut msg) = q.recv() { let verdict = handle_packet!( &msg.get_payload(), + &mut buf, handled => nfq::Verdict::Drop, rejected => nfq::Verdict::Accept, ); diff --git a/src/platform/windows.rs b/src/platform/windows.rs index 35ce956..f3ee2c0 100644 --- a/src/platform/windows.rs +++ b/src/platform/windows.rs @@ -62,16 +62,19 @@ pub fn send_to_raw(pkt: &[u8]) -> Result<()> { pub fn run() -> Result<()> { use crate::{handle_packet, RUNNING, MESSAGE_AT_RUN}; + use super::PACKET_SIZE_CAP; - let mut buf = vec![0u8; 65536]; + let mut windivert_buf = vec![0u8; 65536]; + let mut buf = Vec::<u8>::with_capacity(PACKET_SIZE_CAP); splash!("{MESSAGE_AT_RUN}"); while RUNNING.load(Ordering::SeqCst) { - let pkt = WINDIVERT_HANDLE.recv(Some(&mut buf))?; + let pkt = WINDIVERT_HANDLE.recv(Some(&mut windivert_buf))?; handle_packet!( &pkt.data, + &mut buf, handled => {}, rejected => { WINDIVERT_HANDLE.send(&pkt)?; } ); |