From deef1092a82521233a6a31a7bfae17956340cfb2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= <jon.haggblad@gmail.com>
Date: Thu, 1 Feb 2024 18:30:50 +0100
Subject: [PATCH] wip: bundle packets from tun

---
 service-providers/ip-packet-router/src/lib.rs |  2 +-
 .../ip-packet-router/src/mixnet_listener.rs   |  4 ++--
 .../ip-packet-router/src/tun_listener.rs      | 20 ++++++++++++++++---
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/service-providers/ip-packet-router/src/lib.rs b/service-providers/ip-packet-router/src/lib.rs
index f8aa663be8..5ddb2c4641 100644
--- a/service-providers/ip-packet-router/src/lib.rs
+++ b/service-providers/ip-packet-router/src/lib.rs
@@ -9,7 +9,7 @@ mod constants;
 pub mod error;
 mod ip_packet_router;
 mod mixnet_client;
-mod mixnet_listener;
+pub mod mixnet_listener;
 pub mod request_filter;
 mod tun_listener;
 mod util;
diff --git a/service-providers/ip-packet-router/src/mixnet_listener.rs b/service-providers/ip-packet-router/src/mixnet_listener.rs
index 993bc0db8f..2e6e427bd5 100644
--- a/service-providers/ip-packet-router/src/mixnet_listener.rs
+++ b/service-providers/ip-packet-router/src/mixnet_listener.rs
@@ -30,12 +30,12 @@ use crate::{
 
 // Tokio codec for bundling multiple IP packets into one buffer that is at most 1500 bytes long.
 // These packets are separated by a 2 byte length prefix.
-struct BundledIpPacketCodec {
+pub struct BundledIpPacketCodec {
     buffer: BytesMut,
 }
 
 impl BundledIpPacketCodec {
-    fn new() -> Self {
+    pub fn new() -> Self {
         BundledIpPacketCodec {
             buffer: BytesMut::new(),
         }
diff --git a/service-providers/ip-packet-router/src/tun_listener.rs b/service-providers/ip-packet-router/src/tun_listener.rs
index 1f36328863..ab3a235f46 100644
--- a/service-providers/ip-packet-router/src/tun_listener.rs
+++ b/service-providers/ip-packet-router/src/tun_listener.rs
@@ -23,7 +23,7 @@ pub(crate) struct TunListener {
 
 #[cfg(target_os = "linux")]
 impl TunListener {
-    async fn handle_packet(&mut self, buf: &[u8], len: usize) -> Result<()> {
+    async fn handle_packet(&mut self, buf: &[u8], len: usize, bundled_packet_codec: &mut mixnet_listener::BundledIpPacketCodec) -> Result<()> {
         let Some(dst_addr) = parse_dst_addr(&buf[..len]) else {
             log::warn!("Failed to parse packet");
             return Ok(());
@@ -36,7 +36,18 @@ impl TunListener {
         }) = self.connected_clients.get(&dst_addr)
         {
             let packet = buf[..len].to_vec();
-            let response_packet = IpPacketResponse::new_ip_packet(packet.into())
+
+            // Bunch together
+            let packet_bytes = Bytes::from(packet);
+            let mut bundled_packets = BytesMut::new();
+            bundled_packet_codec.encode(packet_bytes, &mut bundled_packets).unwrap();
+            if bundled_packets.is_empty() {
+                return Ok(());
+            }
+            let bundled_packets = bundled_packets.freeze();
+
+            // let response_packet = IpPacketResponse::new_ip_packet(packet.into())
+            let response_packet = IpPacketResponse::new_ip_packet(bundled_packets)
                 .to_bytes()
                 .map_err(|err| IpPacketRouterError::FailedToSerializeResponsePacket {
                     source: err,
@@ -56,6 +67,9 @@ impl TunListener {
 
     async fn run(mut self) -> Result<()> {
         let mut buf = [0u8; 65535];
+
+        let mut bundled_packet_codec = mixnet_listener::BundledIpPacketCodec::new();
+
         while !self.task_client.is_shutdown() {
             tokio::select! {
                 _ = self.task_client.recv() => {
@@ -71,7 +85,7 @@ impl TunListener {
                 },
                 len = self.tun_reader.read(&mut buf) => match len {
                     Ok(len) => {
-                        if let Err(err) = self.handle_packet(&buf, len).await {
+                        if let Err(err) = self.handle_packet(&buf, len, &mut bundled_packet_codec).await {
                             log::error!("tun: failed to handle packet: {err}");
                         }
                     },