Multi IP packet codec (#4379)

* Codec implementation

* rustfmt

* Extract out magic numbers

Cargo.lock

@@ -6666,6 +6666,8 @@ dependencies = [
  "rand 0.8.5",
  "serde",
  "thiserror",
+ "tokio",
+ "tokio-util",
 ]
 
 [[package]]

common/ip-packet-requests/Cargo.toml

@@ -8,8 +8,6 @@ documentation.workspace = true
 edition.workspace = true
 license.workspace = true
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
-
 [dependencies]
 bincode = "1.3.3"
 bytes = "1.5.0"
@@ -17,3 +15,5 @@ nym-sphinx = { path = "../nymsphinx" }
 rand = "0.8.5"
 serde = { workspace = true, features = ["derive"] }
 thiserror = { workspace = true }
+tokio = { workspace = true, features = ["time"] }
+tokio-util = { workspace = true, features = ["codec"] }

common/ip-packet-requests/src/codec.rs

@@ -0,0 +1,123 @@
+use std::time::Duration;
+
+use bytes::{Buf, Bytes, BytesMut};
+use tokio_util::codec::{Decoder, Encoder};
+
+#[derive(thiserror::Error, Debug)]
+pub enum Error {
+    #[error("{0}")]
+    IO(#[from] std::io::Error),
+}
+
+pub const BUFFER_TIMEOUT: Duration = Duration::from_millis(20);
+
+// TODO: increase this to make max out effective sphinx payload size. Sphinx packets also carry the
+// MixAck so that's why we can't just use 2kb.
+pub const MAX_PACKET_SIZE: usize = 1500;
+
+// Each IP packet is prefixed by a 2 byte length prefix
+const LENGTH_PREFIX_SIZE: usize = 2;
+
+// 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. We need a timer so that we don't wait too
+// long for the buffer to fill up, since this kills latency.
+pub struct MultiIpPacketCodec {
+    buffer: BytesMut,
+    buffer_timeout: tokio::time::Interval,
+}
+
+impl MultiIpPacketCodec {
+    pub fn new(buffer_timeout: Duration) -> Self {
+        MultiIpPacketCodec {
+            buffer: BytesMut::new(),
+            buffer_timeout: tokio::time::interval(buffer_timeout),
+        }
+    }
+
+    // Append a packet to the buffer and return the buffer if it's full
+    pub fn append_packet(&mut self, packet: Bytes) -> Option<Bytes> {
+        let mut bundled_packets = BytesMut::new();
+        self.encode(packet, &mut bundled_packets).unwrap();
+        if bundled_packets.is_empty() {
+            None
+        } else {
+            // log::info!("Sphinx packet utilization: {:.2}", self.buffer.len() as f64 / MAX_PACKET_SIZE as f64);
+            Some(bundled_packets.freeze())
+        }
+    }
+
+    // Flush the current buffer and return it.
+    fn flush_current_buffer(&mut self) -> Bytes {
+        let mut output_buffer = BytesMut::new();
+        std::mem::swap(&mut output_buffer, &mut self.buffer);
+        output_buffer.freeze()
+    }
+
+    // Wait for the buffer_timeout to tick and then flush the buffer.
+    // This is useful when we want to send the buffer even if it's not full.
+    pub async fn buffer_timeout(&mut self) -> Option<Bytes> {
+        // Wait for buffer_timeout to tick
+        let _ = self.buffer_timeout.tick().await;
+
+        // Flush the buffer and return it
+        let packets = self.flush_current_buffer();
+        if packets.is_empty() {
+            None
+        } else {
+            Some(packets)
+        }
+    }
+}
+
+impl Encoder<Bytes> for MultiIpPacketCodec {
+    type Error = Error;
+
+    fn encode(&mut self, packet: Bytes, dst: &mut BytesMut) -> Result<(), Self::Error> {
+        if self.buffer.is_empty() {
+            self.buffer_timeout.reset();
+        }
+        let packet_size = packet.len();
+
+        if self.buffer.len() + packet_size + LENGTH_PREFIX_SIZE > MAX_PACKET_SIZE {
+            // If the packet doesn't fit in the buffer, send the buffer and then add it to the buffer
+            dst.extend_from_slice(&self.buffer);
+            self.buffer = BytesMut::new();
+            self.buffer_timeout.reset();
+        }
+
+        // Add the packet size
+        self.buffer
+            .extend_from_slice(&(packet_size as u16).to_be_bytes());
+        // Add the packet to the buffer
+        self.buffer.extend_from_slice(&packet);
+
+        Ok(())
+    }
+}
+
+impl Decoder for MultiIpPacketCodec {
+    type Item = Bytes;
+    type Error = Error;
+
+    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
+        if src.len() < LENGTH_PREFIX_SIZE {
+            // Not enough bytes to read the length prefix
+            return Ok(None);
+        }
+
+        let packet_size = u16::from_be_bytes([src[0], src[1]]) as usize;
+
+        if src.len() < packet_size + LENGTH_PREFIX_SIZE {
+            // Not enough bytes to read the packet
+            return Ok(None);
+        }
+
+        // Remove the length prefix
+        src.advance(LENGTH_PREFIX_SIZE);
+
+        // Read the packet
+        let packet = src.split_to(packet_size);
+
+        Ok(Some(packet.freeze()))
+    }
+}

common/ip-packet-requests/src/lib.rs

@@ -1,3 +1,4 @@
+pub mod codec;
 pub mod request;
 pub mod response;