From 8b2f80b03c696cefb2deebcf6fd8e6c52424b594 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= Date: Thu, 8 Feb 2024 11:01:01 +0100 Subject: [PATCH] Multi IP packet codec (#4379) * Codec implementation * rustfmt * Extract out magic numbers --- Cargo.lock | 2 + common/ip-packet-requests/Cargo.toml | 4 +- common/ip-packet-requests/src/codec.rs | 123 +++++++++++++++++++++++++ common/ip-packet-requests/src/lib.rs | 1 + 4 files changed, 128 insertions(+), 2 deletions(-) create mode 100644 common/ip-packet-requests/src/codec.rs diff --git a/Cargo.lock b/Cargo.lock index 74f012a2c3..141f1a846a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6666,6 +6666,8 @@ dependencies = [ "rand 0.8.5", "serde", "thiserror", + "tokio", + "tokio-util", ] [[package]] diff --git a/common/ip-packet-requests/Cargo.toml b/common/ip-packet-requests/Cargo.toml index 2a92715f82..c6b5d918e1 100644 --- a/common/ip-packet-requests/Cargo.toml +++ b/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"] } diff --git a/common/ip-packet-requests/src/codec.rs b/common/ip-packet-requests/src/codec.rs new file mode 100644 index 0000000000..86df1ce8f7 --- /dev/null +++ b/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 { + 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 { + // 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 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, 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())) + } +} diff --git a/common/ip-packet-requests/src/lib.rs b/common/ip-packet-requests/src/lib.rs index 5301cf6dd4..616357dedd 100644 --- a/common/ip-packet-requests/src/lib.rs +++ b/common/ip-packet-requests/src/lib.rs @@ -1,3 +1,4 @@ +pub mod codec; pub mod request; pub mod response;