Jon Häggblad
131 lines
4.2 KiB
Rust
131 lines
4.2 KiB
Rust
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),
|
|
}
|
|
}
|
|
|
|
pub fn bundle_one_packet(packet: Bytes) -> Bytes {
|
|
let mut bundled_packets = BytesMut::new();
|
|
bundled_packets.extend_from_slice(&(packet.len() as u16).to_be_bytes());
|
|
bundled_packets.extend_from_slice(&packet);
|
|
bundled_packets.freeze()
|
|
}
|
|
|
|
// 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.
|
|
pub 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()))
|
|
}
|
|
}
|