#![forbid(unsafe_code)]
use std::collections::HashMap;
use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::io;
use std::net::{SocketAddr, ToSocketAddrs, UdpSocket};
use std::sync::{Arc, Mutex};
use sectorsync_core::prelude::{ClientId, StationId, Tick};
pub const DEFAULT_UDP_RECV_BUFFER_SIZE: usize = 16 * 1024;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct OutboundPacket {
pub client_id: ClientId,
pub bytes: Vec<u8>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct InboundPacket {
pub client_id: Option<ClientId>,
pub remote_addr: SocketAddr,
pub bytes: Vec<u8>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StationOutboundPacket {
pub source_station: StationId,
pub target_station: StationId,
pub bytes: Vec<u8>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StationInboundPacket {
pub source_station: StationId,
pub target_station: StationId,
pub bytes: Vec<u8>,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct StationPacketBatch {
pub packets: Vec<StationOutboundPacket>,
}
impl StationPacketBatch {
pub const fn new() -> Self {
Self {
packets: Vec::new(),
}
}
pub fn push(&mut self, packet: StationOutboundPacket) {
self.packets.push(packet);
}
pub fn len(&self) -> usize {
self.packets.len()
}
pub fn is_empty(&self) -> bool {
self.packets.is_empty()
}
pub fn bytes_len(&self) -> usize {
self.packets.iter().map(|packet| packet.bytes.len()).sum()
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct PacketBatch {
pub packets: Vec<OutboundPacket>,
}
impl PacketBatch {
pub const fn new() -> Self {
Self {
packets: Vec::new(),
}
}
pub fn push(&mut self, packet: OutboundPacket) {
self.packets.push(packet);
}
pub fn len(&self) -> usize {
self.packets.len()
}
pub fn is_empty(&self) -> bool {
self.packets.is_empty()
}
pub fn bytes_len(&self) -> usize {
self.packets.iter().map(|packet| packet.bytes.len()).sum()
}
}
pub trait TransportSink {
type Error;
fn send(&mut self, packet: OutboundPacket) -> Result<(), Self::Error>;
fn send_batch(&mut self, batch: PacketBatch) -> Result<(), Self::Error> {
for packet in batch.packets {
self.send(packet)?;
}
Ok(())
}
}
pub trait TransportReceiver {
type Error;
fn try_recv(&mut self) -> Result<Option<InboundPacket>, Self::Error>;
}
const PACKET_SECURITY_MAGIC: [u8; 4] = *b"SSEC";
pub const PACKET_SECURITY_HEADER_BYTES: usize = 22;
pub const DEFAULT_PACKET_SECURITY_MAX_PAYLOAD_BYTES: usize =
(16 * 1024) - PACKET_SECURITY_HEADER_BYTES;
pub const DEFAULT_PACKET_SECURITY_MAX_TAG_BYTES: usize = 128;
pub const DEFAULT_PACKET_SECURITY_REPLAY_HISTORY: usize = 4096;
pub const DEFAULT_PACKET_KEY_RING_MAX_KEYS: usize = 32;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct PacketSecurityConfig {
pub max_payload_bytes: usize,
pub max_tag_bytes: usize,
pub max_replay_history: usize,
}
impl Default for PacketSecurityConfig {
fn default() -> Self {
Self {
max_payload_bytes: DEFAULT_PACKET_SECURITY_MAX_PAYLOAD_BYTES,
max_tag_bytes: DEFAULT_PACKET_SECURITY_MAX_TAG_BYTES,
max_replay_history: DEFAULT_PACKET_SECURITY_REPLAY_HISTORY,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PacketKeyState {
Active,
Retiring,
Revoked,
}
impl PacketKeyState {
fn can_send(self) -> bool {
self == Self::Active
}
fn can_accept(self) -> bool {
matches!(self, Self::Active | Self::Retiring)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct PacketKeyDescriptor {
pub key_id: u32,
pub state: PacketKeyState,
pub created_at: Tick,
pub activated_at: Tick,
pub retires_at: Option<Tick>,
pub expires_at: Option<Tick>,
pub send_priority: u32,
}
impl PacketKeyDescriptor {
pub const fn active(key_id: u32, now: Tick, send_priority: u32) -> Self {
Self {
key_id,
state: PacketKeyState::Active,
created_at: now,
activated_at: now,
retires_at: None,
expires_at: None,
send_priority,
}
}
#[must_use]
pub const fn with_expiry(mut self, expires_at: Tick) -> Self {
self.expires_at = Some(expires_at);
self
}
pub fn is_send_eligible(self, now: Tick) -> bool {
self.state.can_send() && self.is_activated_at(now) && !self.is_expired_at(now)
}
pub fn is_accept_eligible(self, now: Tick) -> bool {
self.state.can_accept() && self.is_activated_at(now) && !self.is_expired_at(now)
}
fn is_activated_at(self, now: Tick) -> bool {
now >= self.activated_at
}
fn is_expired_at(self, now: Tick) -> bool {
self.expires_at.is_some_and(|expires_at| now >= expires_at)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct PacketKeyRingConfig {
pub max_keys: usize,
}
impl Default for PacketKeyRingConfig {
fn default() -> Self {
Self {
max_keys: DEFAULT_PACKET_KEY_RING_MAX_KEYS,
}
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct PacketKeyRingStats {
pub keys_inserted: usize,
pub keys_activated: usize,
pub keys_retired: usize,
pub keys_revoked: usize,
pub keys_expired_removed: usize,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PacketKeyRingError {
CapacityFull {
capacity: usize,
},
DuplicateKey(u32),
MissingKey(u32),
NoSendKey,
KeyNotSendable {
key_id: u32,
state: PacketKeyState,
},
KeyNotAccepted {
key_id: u32,
state: PacketKeyState,
},
}
impl core::fmt::Display for PacketKeyRingError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::CapacityFull { capacity } => {
write!(f, "packet key ring capacity full: capacity {capacity}")
}
Self::DuplicateKey(key_id) => write!(f, "packet key {key_id} already exists"),
Self::MissingKey(key_id) => write!(f, "packet key {key_id} is missing"),
Self::NoSendKey => f.write_str("no packet key is eligible for send"),
Self::KeyNotSendable { key_id, state } => {
write!(f, "packet key {key_id} is not sendable in state {state:?}")
}
Self::KeyNotAccepted { key_id, state } => {
write!(f, "packet key {key_id} is not accepted in state {state:?}")
}
}
}
}
impl std::error::Error for PacketKeyRingError {}
#[derive(Clone, Debug)]
pub struct PacketKeyRing {
config: PacketKeyRingConfig,
keys: BTreeMap<u32, PacketKeyDescriptor>,
stats: PacketKeyRingStats,
}
impl PacketKeyRing {
pub fn new(config: PacketKeyRingConfig) -> Self {
Self {
config,
keys: BTreeMap::new(),
stats: PacketKeyRingStats::default(),
}
}
pub fn with_defaults() -> Self {
Self::new(PacketKeyRingConfig::default())
}
pub const fn config(&self) -> PacketKeyRingConfig {
self.config
}
pub fn len(&self) -> usize {
self.keys.len()
}
pub fn is_empty(&self) -> bool {
self.keys.is_empty()
}
pub const fn stats(&self) -> PacketKeyRingStats {
self.stats
}
pub fn get(&self, key_id: u32) -> Option<PacketKeyDescriptor> {
self.keys.get(&key_id).copied()
}
pub fn iter(&self) -> impl Iterator<Item = &PacketKeyDescriptor> {
self.keys.values()
}
pub fn insert(&mut self, descriptor: PacketKeyDescriptor) -> Result<(), PacketKeyRingError> {
if self.keys.contains_key(&descriptor.key_id) {
return Err(PacketKeyRingError::DuplicateKey(descriptor.key_id));
}
if self.keys.len() >= self.config.max_keys {
return Err(PacketKeyRingError::CapacityFull {
capacity: self.config.max_keys,
});
}
self.keys.insert(descriptor.key_id, descriptor);
self.stats.keys_inserted = self.stats.keys_inserted.saturating_add(1);
Ok(())
}
pub fn insert_active(
&mut self,
key_id: u32,
now: Tick,
send_priority: u32,
) -> Result<(), PacketKeyRingError> {
self.insert(PacketKeyDescriptor::active(key_id, now, send_priority))
}
pub fn activate(&mut self, key_id: u32, activated_at: Tick) -> Result<(), PacketKeyRingError> {
let descriptor = self
.keys
.get_mut(&key_id)
.ok_or(PacketKeyRingError::MissingKey(key_id))?;
descriptor.state = PacketKeyState::Active;
descriptor.activated_at = activated_at;
descriptor.retires_at = None;
self.stats.keys_activated = self.stats.keys_activated.saturating_add(1);
Ok(())
}
pub fn retire(&mut self, key_id: u32, retires_at: Tick) -> Result<(), PacketKeyRingError> {
let descriptor = self
.keys
.get_mut(&key_id)
.ok_or(PacketKeyRingError::MissingKey(key_id))?;
descriptor.state = PacketKeyState::Retiring;
descriptor.retires_at = Some(retires_at);
self.stats.keys_retired = self.stats.keys_retired.saturating_add(1);
Ok(())
}
pub fn revoke(&mut self, key_id: u32) -> Result<(), PacketKeyRingError> {
let descriptor = self
.keys
.get_mut(&key_id)
.ok_or(PacketKeyRingError::MissingKey(key_id))?;
descriptor.state = PacketKeyState::Revoked;
self.stats.keys_revoked = self.stats.keys_revoked.saturating_add(1);
Ok(())
}
pub fn set_expiry(
&mut self,
key_id: u32,
expires_at: Option<Tick>,
) -> Result<(), PacketKeyRingError> {
let descriptor = self
.keys
.get_mut(&key_id)
.ok_or(PacketKeyRingError::MissingKey(key_id))?;
descriptor.expires_at = expires_at;
Ok(())
}
pub fn remove_expired(&mut self, now: Tick) -> usize {
let before = self.keys.len();
self.keys
.retain(|_, descriptor| !descriptor.is_expired_at(now));
let removed = before.saturating_sub(self.keys.len());
self.stats.keys_expired_removed = self.stats.keys_expired_removed.saturating_add(removed);
removed
}
pub fn select_send_key(&self, now: Tick) -> Result<PacketKeyDescriptor, PacketKeyRingError> {
self.keys
.values()
.copied()
.filter(|descriptor| descriptor.is_send_eligible(now))
.max_by_key(|descriptor| {
(
descriptor.send_priority,
descriptor.activated_at,
descriptor.key_id,
)
})
.ok_or(PacketKeyRingError::NoSendKey)
}
pub fn accept_key(
&self,
key_id: u32,
now: Tick,
) -> Result<PacketKeyDescriptor, PacketKeyRingError> {
let descriptor = self
.keys
.get(&key_id)
.copied()
.ok_or(PacketKeyRingError::MissingKey(key_id))?;
if descriptor.is_accept_eligible(now) {
Ok(descriptor)
} else {
Err(PacketKeyRingError::KeyNotAccepted {
key_id,
state: descriptor.state,
})
}
}
}
impl Default for PacketKeyRing {
fn default() -> Self {
Self::with_defaults()
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PacketSecurityEnvelope {
pub key_id: u32,
pub nonce: u64,
pub payload: Vec<u8>,
pub tag: Vec<u8>,
}
impl PacketSecurityEnvelope {
pub fn encode(
&self,
config: PacketSecurityConfig,
out: &mut Vec<u8>,
) -> Result<(), PacketSecurityEncodeError> {
if self.payload.len() > config.max_payload_bytes {
return Err(PacketSecurityEncodeError::PayloadTooLarge {
budget: config.max_payload_bytes,
actual: self.payload.len(),
});
}
if self.tag.len() > config.max_tag_bytes {
return Err(PacketSecurityEncodeError::TagTooLarge {
budget: config.max_tag_bytes,
actual: self.tag.len(),
});
}
out.extend_from_slice(&PACKET_SECURITY_MAGIC);
out.extend_from_slice(&self.key_id.to_le_bytes());
out.extend_from_slice(&self.nonce.to_le_bytes());
let payload_len = u32::try_from(self.payload.len()).map_err(|_| {
PacketSecurityEncodeError::PayloadTooLarge {
budget: config.max_payload_bytes,
actual: self.payload.len(),
}
})?;
let tag_len =
u16::try_from(self.tag.len()).map_err(|_| PacketSecurityEncodeError::TagTooLarge {
budget: config.max_tag_bytes,
actual: self.tag.len(),
})?;
out.extend_from_slice(&payload_len.to_le_bytes());
out.extend_from_slice(&tag_len.to_le_bytes());
out.extend_from_slice(&self.payload);
out.extend_from_slice(&self.tag);
Ok(())
}
pub fn decode(
config: PacketSecurityConfig,
input: &[u8],
) -> Result<Self, PacketSecurityDecodeError> {
let mut cursor = SecurityCursor::new(input);
let magic = cursor.read_array::<4>()?;
if magic != PACKET_SECURITY_MAGIC {
return Err(PacketSecurityDecodeError::BadMagic);
}
let key_id = cursor.read_u32()?;
let nonce = cursor.read_u64()?;
let payload_len = cursor.read_u32()? as usize;
let tag_len = cursor.read_u16()? as usize;
if payload_len > config.max_payload_bytes {
return Err(PacketSecurityDecodeError::PayloadTooLarge {
budget: config.max_payload_bytes,
actual: payload_len,
});
}
if tag_len > config.max_tag_bytes {
return Err(PacketSecurityDecodeError::TagTooLarge {
budget: config.max_tag_bytes,
actual: tag_len,
});
}
let payload = cursor.read_bytes(payload_len)?;
let tag = cursor.read_bytes(tag_len)?;
cursor.finish()?;
Ok(Self {
key_id,
nonce,
payload,
tag,
})
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PacketSecurityEncodeError {
PayloadTooLarge {
budget: usize,
actual: usize,
},
TagTooLarge {
budget: usize,
actual: usize,
},
}
impl core::fmt::Display for PacketSecurityEncodeError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::PayloadTooLarge { budget, actual } => write!(
f,
"packet security payload exceeded byte budget: budget {budget}, actual {actual}"
),
Self::TagTooLarge { budget, actual } => write!(
f,
"packet security tag exceeded byte budget: budget {budget}, actual {actual}"
),
}
}
}
impl std::error::Error for PacketSecurityEncodeError {}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum PacketSecurityDecodeError {
BadMagic,
Truncated {
needed: usize,
available: usize,
},
PayloadTooLarge {
budget: usize,
actual: usize,
},
TagTooLarge {
budget: usize,
actual: usize,
},
TrailingBytes(usize),
}
impl core::fmt::Display for PacketSecurityDecodeError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::BadMagic => f.write_str("bad packet security envelope magic"),
Self::Truncated { needed, available } => {
write!(
f,
"truncated packet security envelope: needed {needed}, available {available}"
)
}
Self::PayloadTooLarge { budget, actual } => write!(
f,
"packet security payload exceeded byte budget: budget {budget}, actual {actual}"
),
Self::TagTooLarge { budget, actual } => write!(
f,
"packet security tag exceeded byte budget: budget {budget}, actual {actual}"
),
Self::TrailingBytes(bytes) => {
write!(f, "packet security envelope has {bytes} trailing bytes")
}
}
}
}
impl std::error::Error for PacketSecurityDecodeError {}
pub trait PacketAuthenticator {
type Error;
fn sign(
&mut self,
key_id: u32,
nonce: u64,
payload: &[u8],
out: &mut Vec<u8>,
) -> Result<(), Self::Error>;
fn verify(
&mut self,
key_id: u32,
nonce: u64,
payload: &[u8],
tag: &[u8],
) -> Result<bool, Self::Error>;
}
pub trait PacketCipher {
type Error;
fn seal(&mut self, key_id: u32, nonce: u64, payload: &mut Vec<u8>) -> Result<(), Self::Error>;
fn open(&mut self, key_id: u32, nonce: u64, payload: &mut Vec<u8>) -> Result<(), Self::Error>;
}
#[derive(Clone, Copy, Debug, Default)]
pub struct PlaintextPacketCipher;
impl PacketCipher for PlaintextPacketCipher {
type Error = core::convert::Infallible;
fn seal(
&mut self,
_key_id: u32,
_nonce: u64,
_payload: &mut Vec<u8>,
) -> Result<(), Self::Error> {
Ok(())
}
fn open(
&mut self,
_key_id: u32,
_nonce: u64,
_payload: &mut Vec<u8>,
) -> Result<(), Self::Error> {
Ok(())
}
}
#[derive(Clone, Debug)]
pub struct PacketReplayWindow {
max_seen: usize,
seen: BTreeSet<(u32, u64)>,
order: VecDeque<(u32, u64)>,
}
impl PacketReplayWindow {
pub fn new(max_seen: usize) -> Self {
Self {
max_seen,
seen: BTreeSet::new(),
order: VecDeque::new(),
}
}
pub fn len(&self) -> usize {
self.seen.len()
}
pub fn is_empty(&self) -> bool {
self.seen.is_empty()
}
pub fn contains(&self, key_id: u32, nonce: u64) -> bool {
self.seen.contains(&(key_id, nonce))
}
pub fn accept(&mut self, key_id: u32, nonce: u64) -> bool {
if self.max_seen == 0 {
return true;
}
let key = (key_id, nonce);
if self.seen.contains(&key) {
return false;
}
self.seen.insert(key);
self.order.push_back(key);
while self.order.len() > self.max_seen {
if let Some(old) = self.order.pop_front() {
self.seen.remove(&old);
}
}
true
}
}
impl Default for PacketReplayWindow {
fn default() -> Self {
Self::new(DEFAULT_PACKET_SECURITY_REPLAY_HISTORY)
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct PacketSecurityStats {
pub sealed: usize,
pub opened: usize,
pub key_rejected: usize,
pub auth_failed: usize,
pub replay_rejected: usize,
}
#[derive(Debug)]
pub enum PacketSecurityError<A, C> {
Encode(PacketSecurityEncodeError),
Decode(PacketSecurityDecodeError),
Authenticator(A),
Cipher(C),
Key(PacketKeyRingError),
AuthenticationFailed {
key_id: u32,
nonce: u64,
},
Replay {
key_id: u32,
nonce: u64,
},
}
impl<A: core::fmt::Display, C: core::fmt::Display> core::fmt::Display
for PacketSecurityError<A, C>
{
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::Encode(error) => write!(f, "{error}"),
Self::Decode(error) => write!(f, "{error}"),
Self::Authenticator(error) => write!(f, "{error}"),
Self::Cipher(error) => write!(f, "{error}"),
Self::Key(error) => write!(f, "{error}"),
Self::AuthenticationFailed { key_id, nonce } => write!(
f,
"packet authentication failed for key {key_id} nonce {nonce}"
),
Self::Replay { key_id, nonce } => {
write!(f, "packet replay rejected for key {key_id} nonce {nonce}")
}
}
}
}
impl<A, C> std::error::Error for PacketSecurityError<A, C>
where
A: std::error::Error + 'static,
C: std::error::Error + 'static,
{
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Encode(error) => Some(error),
Self::Decode(error) => Some(error),
Self::Authenticator(error) => Some(error),
Self::Cipher(error) => Some(error),
Self::Key(error) => Some(error),
Self::AuthenticationFailed { .. } | Self::Replay { .. } => None,
}
}
}
#[derive(Clone, Debug)]
pub struct PacketSecurityBox<A, C> {
config: PacketSecurityConfig,
authenticator: A,
cipher: C,
replay: PacketReplayWindow,
next_nonce: BTreeMap<u32, u64>,
stats: PacketSecurityStats,
}
impl<A, C> PacketSecurityBox<A, C> {
pub fn new(config: PacketSecurityConfig, authenticator: A, cipher: C) -> Self {
Self {
config,
authenticator,
cipher,
replay: PacketReplayWindow::new(config.max_replay_history),
next_nonce: BTreeMap::new(),
stats: PacketSecurityStats::default(),
}
}
pub const fn config(&self) -> PacketSecurityConfig {
self.config
}
pub const fn stats(&self) -> PacketSecurityStats {
self.stats
}
pub const fn replay(&self) -> &PacketReplayWindow {
&self.replay
}
pub fn into_inner(self) -> (A, C, PacketReplayWindow) {
(self.authenticator, self.cipher, self.replay)
}
}
impl<A, C> PacketSecurityBox<A, C>
where
A: PacketAuthenticator,
C: PacketCipher,
{
pub fn seal(
&mut self,
key_id: u32,
payload: &[u8],
) -> Result<Vec<u8>, PacketSecurityError<A::Error, C::Error>> {
let nonce = self.allocate_nonce(key_id);
self.seal_with_nonce(key_id, nonce, payload)
}
pub fn seal_with_key_ring(
&mut self,
key_ring: &PacketKeyRing,
payload: &[u8],
now: Tick,
) -> Result<Vec<u8>, PacketSecurityError<A::Error, C::Error>> {
let descriptor = key_ring
.select_send_key(now)
.map_err(PacketSecurityError::Key)?;
self.seal(descriptor.key_id, payload)
}
pub fn seal_with_nonce(
&mut self,
key_id: u32,
nonce: u64,
payload: &[u8],
) -> Result<Vec<u8>, PacketSecurityError<A::Error, C::Error>> {
if payload.len() > self.config.max_payload_bytes {
return Err(PacketSecurityError::Encode(
PacketSecurityEncodeError::PayloadTooLarge {
budget: self.config.max_payload_bytes,
actual: payload.len(),
},
));
}
let mut sealed_payload = payload.to_vec();
self.cipher
.seal(key_id, nonce, &mut sealed_payload)
.map_err(PacketSecurityError::Cipher)?;
let mut tag = Vec::new();
self.authenticator
.sign(key_id, nonce, &sealed_payload, &mut tag)
.map_err(PacketSecurityError::Authenticator)?;
let envelope = PacketSecurityEnvelope {
key_id,
nonce,
payload: sealed_payload,
tag,
};
let mut out = Vec::with_capacity(
PACKET_SECURITY_HEADER_BYTES
.saturating_add(envelope.payload.len())
.saturating_add(envelope.tag.len()),
);
envelope
.encode(self.config, &mut out)
.map_err(PacketSecurityError::Encode)?;
self.stats.sealed = self.stats.sealed.saturating_add(1);
Ok(out)
}
pub fn open(
&mut self,
input: &[u8],
) -> Result<Vec<u8>, PacketSecurityError<A::Error, C::Error>> {
let envelope = PacketSecurityEnvelope::decode(self.config, input)
.map_err(PacketSecurityError::Decode)?;
self.open_decoded(envelope)
}
pub fn open_with_key_ring(
&mut self,
key_ring: &PacketKeyRing,
input: &[u8],
now: Tick,
) -> Result<Vec<u8>, PacketSecurityError<A::Error, C::Error>> {
let envelope = PacketSecurityEnvelope::decode(self.config, input)
.map_err(PacketSecurityError::Decode)?;
if let Err(error) = key_ring.accept_key(envelope.key_id, now) {
self.stats.key_rejected = self.stats.key_rejected.saturating_add(1);
return Err(PacketSecurityError::Key(error));
}
self.open_decoded(envelope)
}
fn open_decoded(
&mut self,
envelope: PacketSecurityEnvelope,
) -> Result<Vec<u8>, PacketSecurityError<A::Error, C::Error>> {
let verified = self
.authenticator
.verify(
envelope.key_id,
envelope.nonce,
&envelope.payload,
&envelope.tag,
)
.map_err(PacketSecurityError::Authenticator)?;
if !verified {
self.stats.auth_failed = self.stats.auth_failed.saturating_add(1);
return Err(PacketSecurityError::AuthenticationFailed {
key_id: envelope.key_id,
nonce: envelope.nonce,
});
}
if !self.replay.accept(envelope.key_id, envelope.nonce) {
self.stats.replay_rejected = self.stats.replay_rejected.saturating_add(1);
return Err(PacketSecurityError::Replay {
key_id: envelope.key_id,
nonce: envelope.nonce,
});
}
let mut payload = envelope.payload;
self.cipher
.open(envelope.key_id, envelope.nonce, &mut payload)
.map_err(PacketSecurityError::Cipher)?;
self.stats.opened = self.stats.opened.saturating_add(1);
Ok(payload)
}
fn allocate_nonce(&mut self, key_id: u32) -> u64 {
let next = self.next_nonce.entry(key_id).or_insert(1);
let nonce = *next;
*next = next.saturating_add(1);
nonce
}
}
struct SecurityCursor<'a> {
input: &'a [u8],
offset: usize,
}
impl<'a> SecurityCursor<'a> {
const fn new(input: &'a [u8]) -> Self {
Self { input, offset: 0 }
}
fn read_u16(&mut self) -> Result<u16, PacketSecurityDecodeError> {
let bytes = self.read_array::<2>()?;
Ok(u16::from_le_bytes(bytes))
}
fn read_u32(&mut self) -> Result<u32, PacketSecurityDecodeError> {
let bytes = self.read_array::<4>()?;
Ok(u32::from_le_bytes(bytes))
}
fn read_u64(&mut self) -> Result<u64, PacketSecurityDecodeError> {
let bytes = self.read_array::<8>()?;
Ok(u64::from_le_bytes(bytes))
}
fn read_array<const N: usize>(&mut self) -> Result<[u8; N], PacketSecurityDecodeError> {
self.require(N)?;
let mut out = [0_u8; N];
out.copy_from_slice(&self.input[self.offset..self.offset + N]);
self.offset += N;
Ok(out)
}
fn read_bytes(&mut self, len: usize) -> Result<Vec<u8>, PacketSecurityDecodeError> {
self.require(len)?;
let bytes = self.input[self.offset..self.offset + len].to_vec();
self.offset += len;
Ok(bytes)
}
fn require(&self, count: usize) -> Result<(), PacketSecurityDecodeError> {
let needed = self.offset.saturating_add(count);
if needed > self.input.len() {
Err(PacketSecurityDecodeError::Truncated {
needed,
available: self.input.len(),
})
} else {
Ok(())
}
}
fn finish(&self) -> Result<(), PacketSecurityDecodeError> {
if self.offset == self.input.len() {
Ok(())
} else {
Err(PacketSecurityDecodeError::TrailingBytes(
self.input.len().saturating_sub(self.offset),
))
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ClientTransportLimits {
pub max_queued_packets_per_client: usize,
pub max_packet_bytes: usize,
}
impl Default for ClientTransportLimits {
fn default() -> Self {
Self {
max_queued_packets_per_client: 4096,
max_packet_bytes: 16 * 1024,
}
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct InMemoryTransportStats {
pub packets_sent: usize,
pub packets_received: usize,
pub bytes_sent: usize,
pub bytes_received: usize,
pub packets_rejected_full: usize,
pub packets_rejected_bytes: usize,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum InMemoryTransportError {
MissingLocal(ClientId),
MissingTarget(ClientId),
QueueFull {
client_id: ClientId,
capacity: usize,
},
PacketTooLarge {
budget: usize,
actual: usize,
},
Poisoned,
}
impl core::fmt::Display for InMemoryTransportError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::MissingLocal(client_id) => {
write!(
f,
"in-memory transport local client {} is missing",
client_id.get()
)
}
Self::MissingTarget(client_id) => {
write!(
f,
"in-memory transport target client {} is missing",
client_id.get()
)
}
Self::QueueFull {
client_id,
capacity,
} => write!(
f,
"in-memory transport target client {} queue is full at capacity {capacity}",
client_id.get()
),
Self::PacketTooLarge { budget, actual } => {
write!(
f,
"in-memory transport packet exceeded byte budget: budget {budget}, actual {actual}"
)
}
Self::Poisoned => f.write_str("in-memory transport state is poisoned"),
}
}
}
impl std::error::Error for InMemoryTransportError {}
#[derive(Clone, Debug)]
struct InMemoryTransportClient {
remote_addr: SocketAddr,
queue: VecDeque<InboundPacket>,
}
#[derive(Debug)]
struct InMemoryTransportInner {
limits: ClientTransportLimits,
clients: BTreeMap<ClientId, InMemoryTransportClient>,
addr_to_client: HashMap<SocketAddr, ClientId>,
stats: InMemoryTransportStats,
}
#[derive(Clone, Debug)]
pub struct InMemoryTransportHub {
inner: Arc<Mutex<InMemoryTransportInner>>,
}
impl InMemoryTransportHub {
pub fn new(limits: ClientTransportLimits) -> Self {
Self {
inner: Arc::new(Mutex::new(InMemoryTransportInner {
limits,
clients: BTreeMap::new(),
addr_to_client: HashMap::new(),
stats: InMemoryTransportStats::default(),
})),
}
}
pub fn register_client(
&self,
client_id: ClientId,
remote_addr: SocketAddr,
) -> Result<Option<SocketAddr>, InMemoryTransportError> {
let mut inner = self.lock_inner()?;
let queue_capacity = inner.limits.max_queued_packets_per_client;
let old_addr = inner.clients.insert(
client_id,
InMemoryTransportClient {
remote_addr,
queue: VecDeque::with_capacity(queue_capacity),
},
);
let previous_addr = old_addr.map(|client| client.remote_addr);
if let Some(previous_addr) = previous_addr {
inner.addr_to_client.remove(&previous_addr);
}
if let Some(old_client) = inner.addr_to_client.insert(remote_addr, client_id)
&& old_client != client_id
{
inner.clients.remove(&old_client);
}
Ok(previous_addr)
}
pub fn endpoint(
&self,
client_id: ClientId,
remote_addr: SocketAddr,
) -> Result<InMemoryTransportEndpoint, InMemoryTransportError> {
self.register_client(client_id, remote_addr)?;
Ok(self.endpoint_for_registered(client_id))
}
pub fn endpoint_for_registered(&self, client_id: ClientId) -> InMemoryTransportEndpoint {
InMemoryTransportEndpoint {
local_client_id: client_id,
hub: self.clone(),
}
}
pub fn queued_len(&self, client_id: ClientId) -> Result<Option<usize>, InMemoryTransportError> {
let inner = self.lock_inner()?;
Ok(inner
.clients
.get(&client_id)
.map(|client| client.queue.len()))
}
pub fn limits(&self) -> Result<ClientTransportLimits, InMemoryTransportError> {
let inner = self.lock_inner()?;
Ok(inner.limits)
}
pub fn stats(&self) -> Result<InMemoryTransportStats, InMemoryTransportError> {
let inner = self.lock_inner()?;
Ok(inner.stats)
}
fn lock_inner(
&self,
) -> Result<std::sync::MutexGuard<'_, InMemoryTransportInner>, InMemoryTransportError> {
self.inner
.lock()
.map_err(|_| InMemoryTransportError::Poisoned)
}
}
impl Default for InMemoryTransportHub {
fn default() -> Self {
Self::new(ClientTransportLimits::default())
}
}
#[derive(Clone, Debug)]
pub struct InMemoryTransportEndpoint {
local_client_id: ClientId,
hub: InMemoryTransportHub,
}
impl InMemoryTransportEndpoint {
pub const fn local_client_id(&self) -> ClientId {
self.local_client_id
}
pub fn local_addr(&self) -> Result<Option<SocketAddr>, InMemoryTransportError> {
let inner = self.hub.lock_inner()?;
Ok(inner
.clients
.get(&self.local_client_id)
.map(|client| client.remote_addr))
}
}
impl TransportSink for InMemoryTransportEndpoint {
type Error = InMemoryTransportError;
fn send(&mut self, packet: OutboundPacket) -> Result<(), Self::Error> {
let actual = packet.bytes.len();
let mut inner = self.hub.lock_inner()?;
let limits = inner.limits;
if actual > limits.max_packet_bytes {
inner.stats.packets_rejected_bytes =
inner.stats.packets_rejected_bytes.saturating_add(1);
return Err(InMemoryTransportError::PacketTooLarge {
budget: limits.max_packet_bytes,
actual,
});
}
let source_addr = inner
.clients
.get(&self.local_client_id)
.ok_or(InMemoryTransportError::MissingLocal(self.local_client_id))?
.remote_addr;
let queue_len = inner
.clients
.get(&packet.client_id)
.ok_or(InMemoryTransportError::MissingTarget(packet.client_id))?
.queue
.len();
if queue_len >= limits.max_queued_packets_per_client {
inner.stats.packets_rejected_full = inner.stats.packets_rejected_full.saturating_add(1);
return Err(InMemoryTransportError::QueueFull {
client_id: packet.client_id,
capacity: limits.max_queued_packets_per_client,
});
}
inner.stats.packets_sent = inner.stats.packets_sent.saturating_add(1);
inner.stats.bytes_sent = inner.stats.bytes_sent.saturating_add(actual);
let target = inner
.clients
.get_mut(&packet.client_id)
.ok_or(InMemoryTransportError::MissingTarget(packet.client_id))?;
target.queue.push_back(InboundPacket {
client_id: Some(self.local_client_id),
remote_addr: source_addr,
bytes: packet.bytes,
});
Ok(())
}
}
impl TransportReceiver for InMemoryTransportEndpoint {
type Error = InMemoryTransportError;
fn try_recv(&mut self) -> Result<Option<InboundPacket>, Self::Error> {
let mut inner = self.hub.lock_inner()?;
let local = inner
.clients
.get_mut(&self.local_client_id)
.ok_or(InMemoryTransportError::MissingLocal(self.local_client_id))?;
let Some(packet) = local.queue.pop_front() else {
return Ok(None);
};
inner.stats.packets_received = inner.stats.packets_received.saturating_add(1);
inner.stats.bytes_received = inner
.stats
.bytes_received
.saturating_add(packet.bytes.len());
Ok(Some(packet))
}
}
const RELIABLE_CLIENT_MAGIC: [u8; 4] = *b"SSCR";
const RELIABLE_CLIENT_KIND_DATA: u8 = 0;
const RELIABLE_CLIENT_KIND_ACK: u8 = 1;
pub const RELIABLE_CLIENT_DATA_HEADER_BYTES: usize = 17;
pub const RELIABLE_CLIENT_ACK_BYTES: usize = 13;
pub const DEFAULT_RELIABLE_CLIENT_MAX_PAYLOAD_BYTES: usize =
(16 * 1024) - RELIABLE_CLIENT_DATA_HEADER_BYTES;
pub const DEFAULT_RELIABLE_CLIENT_DELIVERED_HISTORY: usize = 4096;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ReliableClientConfig {
pub max_in_flight_per_peer: usize,
pub retry_after_ticks: u64,
pub max_attempts: u8,
pub max_payload_bytes: usize,
pub max_delivered_history: usize,
}
impl Default for ReliableClientConfig {
fn default() -> Self {
Self {
max_in_flight_per_peer: 1024,
retry_after_ticks: 2,
max_attempts: 4,
max_payload_bytes: DEFAULT_RELIABLE_CLIENT_MAX_PAYLOAD_BYTES,
max_delivered_history: DEFAULT_RELIABLE_CLIENT_DELIVERED_HISTORY,
}
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ReliableClientStats {
pub data_sent: usize,
pub retries_sent: usize,
pub acks_sent: usize,
pub acks_received: usize,
pub data_delivered: usize,
pub duplicates_suppressed: usize,
pub timed_out: usize,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ReliableClientFrame {
Data {
sequence: u64,
payload: Vec<u8>,
},
Ack {
sequence: u64,
},
}
impl ReliableClientFrame {
pub fn encode(&self, out: &mut Vec<u8>) -> Result<(), ReliableClientEncodeError> {
out.extend_from_slice(&RELIABLE_CLIENT_MAGIC);
match self {
Self::Data { sequence, payload } => {
out.push(RELIABLE_CLIENT_KIND_DATA);
out.extend_from_slice(&sequence.to_le_bytes());
let len = u32::try_from(payload.len()).map_err(|_| {
ReliableClientEncodeError::PayloadTooLarge {
actual: payload.len(),
}
})?;
out.extend_from_slice(&len.to_le_bytes());
out.extend_from_slice(payload);
}
Self::Ack { sequence } => {
out.push(RELIABLE_CLIENT_KIND_ACK);
out.extend_from_slice(&sequence.to_le_bytes());
}
}
Ok(())
}
pub fn decode(input: &[u8]) -> Result<Self, ReliableClientDecodeError> {
let mut cursor = ReliableCursor::new(input);
let magic = cursor
.read_array::<4>()
.map_err(ReliableClientDecodeError::from_station_decode)?;
if magic != RELIABLE_CLIENT_MAGIC {
return Err(ReliableClientDecodeError::BadMagic);
}
let kind = cursor
.read_u8()
.map_err(ReliableClientDecodeError::from_station_decode)?;
let sequence = cursor
.read_u64()
.map_err(ReliableClientDecodeError::from_station_decode)?;
let frame = match kind {
RELIABLE_CLIENT_KIND_DATA => {
let len = cursor
.read_u32()
.map_err(ReliableClientDecodeError::from_station_decode)?
as usize;
let payload = cursor
.read_bytes(len)
.map_err(ReliableClientDecodeError::from_station_decode)?;
Self::Data { sequence, payload }
}
RELIABLE_CLIENT_KIND_ACK => Self::Ack { sequence },
other => return Err(ReliableClientDecodeError::UnknownKind(other)),
};
cursor
.finish()
.map_err(ReliableClientDecodeError::from_station_decode)?;
Ok(frame)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ReliableClientEncodeError {
PayloadTooLarge {
actual: usize,
},
}
impl core::fmt::Display for ReliableClientEncodeError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::PayloadTooLarge { actual } => {
write!(f, "reliable client payload too large: {actual} bytes")
}
}
}
}
impl std::error::Error for ReliableClientEncodeError {}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ReliableClientDecodeError {
BadMagic,
UnknownKind(u8),
Truncated {
needed: usize,
available: usize,
},
TrailingBytes(usize),
}
impl ReliableClientDecodeError {
fn from_station_decode(error: ReliableStationDecodeError) -> Self {
match error {
ReliableStationDecodeError::BadMagic => Self::BadMagic,
ReliableStationDecodeError::UnknownKind(kind) => Self::UnknownKind(kind),
ReliableStationDecodeError::Truncated { needed, available } => {
Self::Truncated { needed, available }
}
ReliableStationDecodeError::TrailingBytes(bytes) => Self::TrailingBytes(bytes),
}
}
}
impl core::fmt::Display for ReliableClientDecodeError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::BadMagic => f.write_str("bad reliable client frame magic"),
Self::UnknownKind(kind) => write!(f, "unknown reliable client frame kind {kind}"),
Self::Truncated { needed, available } => {
write!(
f,
"truncated reliable client frame: needed {needed}, available {available}"
)
}
Self::TrailingBytes(bytes) => {
write!(f, "reliable client frame has {bytes} trailing bytes")
}
}
}
}
impl std::error::Error for ReliableClientDecodeError {}
#[derive(Debug)]
pub enum ReliableClientError<E> {
Transport(E),
MissingSourceClient,
PayloadTooLarge {
budget: usize,
actual: usize,
},
WindowFull {
peer_client: ClientId,
capacity: usize,
},
Encode(ReliableClientEncodeError),
Decode(ReliableClientDecodeError),
}
impl<E: core::fmt::Display> core::fmt::Display for ReliableClientError<E> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::Transport(error) => write!(f, "{error}"),
Self::MissingSourceClient => f.write_str("reliable client packet source is unknown"),
Self::PayloadTooLarge { budget, actual } => {
write!(
f,
"reliable client payload exceeded byte budget: budget {budget}, actual {actual}"
)
}
Self::WindowFull {
peer_client,
capacity,
} => write!(
f,
"reliable client peer {} window is full at capacity {capacity}",
peer_client.get()
),
Self::Encode(error) => write!(f, "{error}"),
Self::Decode(error) => write!(f, "{error}"),
}
}
}
impl<E> std::error::Error for ReliableClientError<E>
where
E: std::error::Error + 'static,
{
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Transport(error) => Some(error),
Self::Encode(error) => Some(error),
Self::Decode(error) => Some(error),
Self::MissingSourceClient | Self::PayloadTooLarge { .. } | Self::WindowFull { .. } => {
None
}
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
struct InFlightReliableClientPacket {
peer_client: ClientId,
sequence: u64,
payload: Vec<u8>,
first_sent_tick: u64,
last_sent_tick: u64,
attempts: u8,
}
#[derive(Clone, Debug)]
pub struct ReliableClientSender {
config: ReliableClientConfig,
next_sequence: BTreeMap<ClientId, u64>,
in_flight: BTreeMap<(ClientId, u64), InFlightReliableClientPacket>,
stats: ReliableClientStats,
}
impl ReliableClientSender {
pub fn new(config: ReliableClientConfig) -> Self {
Self {
config,
next_sequence: BTreeMap::new(),
in_flight: BTreeMap::new(),
stats: ReliableClientStats::default(),
}
}
pub const fn config(&self) -> ReliableClientConfig {
self.config
}
pub const fn stats(&self) -> ReliableClientStats {
self.stats
}
pub fn in_flight_len(&self) -> usize {
self.in_flight.len()
}
pub fn in_flight_for(&self, peer_client: ClientId) -> usize {
self.in_flight
.keys()
.filter(|(client_id, _)| *client_id == peer_client)
.count()
}
pub fn send<T: TransportSink>(
&mut self,
transport: &mut T,
packet: OutboundPacket,
now_tick: u64,
) -> Result<u64, ReliableClientError<T::Error>> {
self.validate_payload(packet.bytes.len())?;
if self.in_flight_for(packet.client_id) >= self.config.max_in_flight_per_peer {
return Err(ReliableClientError::WindowFull {
peer_client: packet.client_id,
capacity: self.config.max_in_flight_per_peer,
});
}
let sequence = self.allocate_sequence(packet.client_id);
Self::send_data_frame(transport, packet.client_id, sequence, &packet.bytes)?;
self.in_flight.insert(
(packet.client_id, sequence),
InFlightReliableClientPacket {
peer_client: packet.client_id,
sequence,
payload: packet.bytes,
first_sent_tick: now_tick,
last_sent_tick: now_tick,
attempts: 1,
},
);
self.stats.data_sent = self.stats.data_sent.saturating_add(1);
Ok(sequence)
}
pub fn acknowledge(&mut self, ack_source_client: ClientId, sequence: u64) -> bool {
let removed = self
.in_flight
.remove(&(ack_source_client, sequence))
.is_some();
if removed {
self.stats.acks_received = self.stats.acks_received.saturating_add(1);
}
removed
}
pub fn retry_due<T: TransportSink>(
&mut self,
transport: &mut T,
now_tick: u64,
) -> Result<ReliableRetryReport, ReliableClientError<T::Error>> {
let keys = self
.in_flight
.iter()
.filter_map(|(key, packet)| {
let due =
now_tick.saturating_sub(packet.last_sent_tick) >= self.config.retry_after_ticks;
due.then_some(*key)
})
.collect::<Vec<_>>();
let mut report = ReliableRetryReport::default();
for key in keys {
let Some(packet) = self.in_flight.get(&key).cloned() else {
continue;
};
if packet.attempts >= self.config.max_attempts {
self.in_flight.remove(&key);
self.stats.timed_out = self.stats.timed_out.saturating_add(1);
report.timed_out = report.timed_out.saturating_add(1);
continue;
}
Self::send_data_frame(
transport,
packet.peer_client,
packet.sequence,
&packet.payload,
)?;
if let Some(stored) = self.in_flight.get_mut(&key) {
stored.last_sent_tick = now_tick;
stored.attempts = stored.attempts.saturating_add(1);
}
self.stats.retries_sent = self.stats.retries_sent.saturating_add(1);
report.retried = report.retried.saturating_add(1);
}
Ok(report)
}
fn validate_payload<E>(&self, bytes: usize) -> Result<(), ReliableClientError<E>> {
if bytes > self.config.max_payload_bytes {
Err(ReliableClientError::PayloadTooLarge {
budget: self.config.max_payload_bytes,
actual: bytes,
})
} else {
Ok(())
}
}
fn allocate_sequence(&mut self, peer_client: ClientId) -> u64 {
let next = self.next_sequence.entry(peer_client).or_insert(1);
let sequence = *next;
*next = next.saturating_add(1);
sequence
}
fn send_data_frame<T: TransportSink>(
transport: &mut T,
peer_client: ClientId,
sequence: u64,
payload: &[u8],
) -> Result<(), ReliableClientError<T::Error>> {
let mut bytes = Vec::with_capacity(
payload
.len()
.saturating_add(RELIABLE_CLIENT_DATA_HEADER_BYTES),
);
ReliableClientFrame::Data {
sequence,
payload: payload.to_vec(),
}
.encode(&mut bytes)
.map_err(ReliableClientError::Encode)?;
transport
.send(OutboundPacket {
client_id: peer_client,
bytes,
})
.map_err(ReliableClientError::Transport)
}
}
impl Default for ReliableClientSender {
fn default() -> Self {
Self::new(ReliableClientConfig::default())
}
}
#[derive(Clone, Debug)]
pub struct ReliableClientReceiver {
config: ReliableClientConfig,
delivered: BTreeSet<(ClientId, u64)>,
delivered_order: VecDeque<(ClientId, u64)>,
stats: ReliableClientStats,
}
impl ReliableClientReceiver {
pub fn new(config: ReliableClientConfig) -> Self {
Self {
config,
delivered: BTreeSet::new(),
delivered_order: VecDeque::new(),
stats: ReliableClientStats::default(),
}
}
pub const fn config(&self) -> ReliableClientConfig {
self.config
}
pub const fn stats(&self) -> ReliableClientStats {
self.stats
}
pub fn handle_data<T: TransportSink>(
&mut self,
transport: &mut T,
packet: InboundPacket,
source_client: ClientId,
sequence: u64,
payload: Vec<u8>,
) -> Result<Option<InboundPacket>, ReliableClientError<T::Error>> {
let InboundPacket {
remote_addr,
bytes: wire_bytes,
..
} = packet;
drop(wire_bytes);
self.send_ack(transport, source_client, sequence)?;
if !self.record_unique(source_client, sequence) {
self.stats.duplicates_suppressed = self.stats.duplicates_suppressed.saturating_add(1);
return Ok(None);
}
self.stats.data_delivered = self.stats.data_delivered.saturating_add(1);
Ok(Some(InboundPacket {
client_id: Some(source_client),
remote_addr,
bytes: payload,
}))
}
fn send_ack<T: TransportSink>(
&mut self,
transport: &mut T,
target_client: ClientId,
sequence: u64,
) -> Result<(), ReliableClientError<T::Error>> {
let mut bytes = Vec::with_capacity(RELIABLE_CLIENT_ACK_BYTES);
ReliableClientFrame::Ack { sequence }
.encode(&mut bytes)
.map_err(ReliableClientError::Encode)?;
transport
.send(OutboundPacket {
client_id: target_client,
bytes,
})
.map_err(ReliableClientError::Transport)?;
self.stats.acks_sent = self.stats.acks_sent.saturating_add(1);
Ok(())
}
fn record_unique(&mut self, source_client: ClientId, sequence: u64) -> bool {
if self.config.max_delivered_history == 0 {
return true;
}
let key = (source_client, sequence);
if self.delivered.contains(&key) {
return false;
}
self.delivered.insert(key);
self.delivered_order.push_back(key);
while self.delivered_order.len() > self.config.max_delivered_history {
if let Some(old) = self.delivered_order.pop_front() {
self.delivered.remove(&old);
}
}
true
}
}
impl Default for ReliableClientReceiver {
fn default() -> Self {
Self::new(ReliableClientConfig::default())
}
}
#[derive(Clone, Debug)]
pub struct ReliableClientEndpoint {
pub sender: ReliableClientSender,
pub receiver: ReliableClientReceiver,
}
impl ReliableClientEndpoint {
pub fn new(config: ReliableClientConfig) -> Self {
Self {
sender: ReliableClientSender::new(config),
receiver: ReliableClientReceiver::new(config),
}
}
pub fn send<T: TransportSink>(
&mut self,
transport: &mut T,
packet: OutboundPacket,
now_tick: u64,
) -> Result<u64, ReliableClientError<T::Error>> {
self.sender.send(transport, packet, now_tick)
}
pub fn retry_due<T: TransportSink>(
&mut self,
transport: &mut T,
now_tick: u64,
) -> Result<ReliableRetryReport, ReliableClientError<T::Error>> {
self.sender.retry_due(transport, now_tick)
}
pub fn handle_inbound<T: TransportSink>(
&mut self,
transport: &mut T,
packet: InboundPacket,
) -> Result<Option<InboundPacket>, ReliableClientError<T::Error>> {
let source_client = packet
.client_id
.ok_or(ReliableClientError::MissingSourceClient)?;
match ReliableClientFrame::decode(&packet.bytes).map_err(ReliableClientError::Decode)? {
ReliableClientFrame::Data { sequence, payload } => {
self.receiver
.handle_data(transport, packet, source_client, sequence, payload)
}
ReliableClientFrame::Ack { sequence } => {
self.sender.acknowledge(source_client, sequence);
Ok(None)
}
}
}
}
impl Default for ReliableClientEndpoint {
fn default() -> Self {
Self::new(ReliableClientConfig::default())
}
}
pub trait StationTransportSink {
type Error;
fn send_station(&mut self, packet: StationOutboundPacket) -> Result<(), Self::Error>;
fn send_station_batch(&mut self, batch: StationPacketBatch) -> Result<(), Self::Error> {
for packet in batch.packets {
self.send_station(packet)?;
}
Ok(())
}
}
pub trait StationTransportReceiver {
type Error;
fn try_recv_station(
&mut self,
target_station: StationId,
) -> Result<Option<StationInboundPacket>, Self::Error>;
}
const RELIABLE_STATION_MAGIC: [u8; 4] = *b"SSRP";
const RELIABLE_KIND_DATA: u8 = 0;
const RELIABLE_KIND_ACK: u8 = 1;
pub const RELIABLE_STATION_DATA_HEADER_BYTES: usize = 17;
pub const RELIABLE_STATION_ACK_BYTES: usize = 13;
pub const DEFAULT_RELIABLE_STATION_MAX_PAYLOAD_BYTES: usize =
(16 * 1024) - RELIABLE_STATION_DATA_HEADER_BYTES;
pub const DEFAULT_RELIABLE_STATION_DELIVERED_HISTORY: usize = 4096;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ReliableStationConfig {
pub max_in_flight_per_target: usize,
pub retry_after_ticks: u64,
pub max_attempts: u8,
pub max_payload_bytes: usize,
pub max_delivered_history: usize,
}
impl Default for ReliableStationConfig {
fn default() -> Self {
Self {
max_in_flight_per_target: 1024,
retry_after_ticks: 2,
max_attempts: 4,
max_payload_bytes: DEFAULT_RELIABLE_STATION_MAX_PAYLOAD_BYTES,
max_delivered_history: DEFAULT_RELIABLE_STATION_DELIVERED_HISTORY,
}
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ReliableStationStats {
pub data_sent: usize,
pub retries_sent: usize,
pub acks_sent: usize,
pub acks_received: usize,
pub data_delivered: usize,
pub duplicates_suppressed: usize,
pub timed_out: usize,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ReliableStationFrame {
Data {
sequence: u64,
payload: Vec<u8>,
},
Ack {
sequence: u64,
},
}
impl ReliableStationFrame {
pub fn encode(&self, out: &mut Vec<u8>) -> Result<(), ReliableStationEncodeError> {
out.extend_from_slice(&RELIABLE_STATION_MAGIC);
match self {
Self::Data { sequence, payload } => {
out.push(RELIABLE_KIND_DATA);
out.extend_from_slice(&sequence.to_le_bytes());
let len = u32::try_from(payload.len()).map_err(|_| {
ReliableStationEncodeError::PayloadTooLarge {
actual: payload.len(),
}
})?;
out.extend_from_slice(&len.to_le_bytes());
out.extend_from_slice(payload);
}
Self::Ack { sequence } => {
out.push(RELIABLE_KIND_ACK);
out.extend_from_slice(&sequence.to_le_bytes());
}
}
Ok(())
}
pub fn decode(input: &[u8]) -> Result<Self, ReliableStationDecodeError> {
let mut cursor = ReliableCursor::new(input);
let magic = cursor.read_array::<4>()?;
if magic != RELIABLE_STATION_MAGIC {
return Err(ReliableStationDecodeError::BadMagic);
}
let kind = cursor.read_u8()?;
let sequence = cursor.read_u64()?;
let frame = match kind {
RELIABLE_KIND_DATA => {
let len = cursor.read_u32()? as usize;
let payload = cursor.read_bytes(len)?;
Self::Data { sequence, payload }
}
RELIABLE_KIND_ACK => Self::Ack { sequence },
other => return Err(ReliableStationDecodeError::UnknownKind(other)),
};
cursor.finish()?;
Ok(frame)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ReliableStationEncodeError {
PayloadTooLarge {
actual: usize,
},
}
impl core::fmt::Display for ReliableStationEncodeError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::PayloadTooLarge { actual } => {
write!(f, "reliable station payload too large: {actual} bytes")
}
}
}
}
impl std::error::Error for ReliableStationEncodeError {}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ReliableStationDecodeError {
BadMagic,
UnknownKind(u8),
Truncated {
needed: usize,
available: usize,
},
TrailingBytes(usize),
}
impl core::fmt::Display for ReliableStationDecodeError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::BadMagic => f.write_str("bad reliable station frame magic"),
Self::UnknownKind(kind) => write!(f, "unknown reliable station frame kind {kind}"),
Self::Truncated { needed, available } => {
write!(
f,
"truncated reliable station frame: needed {needed}, available {available}"
)
}
Self::TrailingBytes(bytes) => {
write!(f, "reliable station frame has {bytes} trailing bytes")
}
}
}
}
impl std::error::Error for ReliableStationDecodeError {}
#[derive(Debug)]
pub enum ReliableStationError<E> {
Transport(E),
PayloadTooLarge {
budget: usize,
actual: usize,
},
WindowFull {
target_station: StationId,
capacity: usize,
},
Encode(ReliableStationEncodeError),
Decode(ReliableStationDecodeError),
}
impl<E: core::fmt::Display> core::fmt::Display for ReliableStationError<E> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::Transport(error) => write!(f, "{error}"),
Self::PayloadTooLarge { budget, actual } => {
write!(
f,
"reliable station payload exceeded byte budget: budget {budget}, actual {actual}"
)
}
Self::WindowFull {
target_station,
capacity,
} => write!(
f,
"reliable station target {} window is full at capacity {capacity}",
target_station.get()
),
Self::Encode(error) => write!(f, "{error}"),
Self::Decode(error) => write!(f, "{error}"),
}
}
}
impl<E> std::error::Error for ReliableStationError<E>
where
E: std::error::Error + 'static,
{
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Transport(error) => Some(error),
Self::Encode(error) => Some(error),
Self::Decode(error) => Some(error),
Self::PayloadTooLarge { .. } | Self::WindowFull { .. } => None,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
struct InFlightReliableStationPacket {
source_station: StationId,
target_station: StationId,
sequence: u64,
payload: Vec<u8>,
first_sent_tick: u64,
last_sent_tick: u64,
attempts: u8,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ReliableRetryReport {
pub retried: usize,
pub timed_out: usize,
}
#[derive(Clone, Debug)]
pub struct ReliableStationSender {
config: ReliableStationConfig,
next_sequence: BTreeMap<StationId, u64>,
in_flight: BTreeMap<(StationId, u64), InFlightReliableStationPacket>,
stats: ReliableStationStats,
}
impl ReliableStationSender {
pub fn new(config: ReliableStationConfig) -> Self {
Self {
config,
next_sequence: BTreeMap::new(),
in_flight: BTreeMap::new(),
stats: ReliableStationStats::default(),
}
}
pub const fn config(&self) -> ReliableStationConfig {
self.config
}
pub const fn stats(&self) -> ReliableStationStats {
self.stats
}
pub fn in_flight_len(&self) -> usize {
self.in_flight.len()
}
pub fn in_flight_for(&self, target_station: StationId) -> usize {
self.in_flight
.keys()
.filter(|(station_id, _)| *station_id == target_station)
.count()
}
pub fn send<T: StationTransportSink>(
&mut self,
transport: &mut T,
packet: StationOutboundPacket,
now_tick: u64,
) -> Result<u64, ReliableStationError<T::Error>> {
self.validate_payload(packet.bytes.len())?;
if self.in_flight_for(packet.target_station) >= self.config.max_in_flight_per_target {
return Err(ReliableStationError::WindowFull {
target_station: packet.target_station,
capacity: self.config.max_in_flight_per_target,
});
}
let sequence = self.allocate_sequence(packet.target_station);
Self::send_data_frame(
transport,
packet.source_station,
packet.target_station,
sequence,
&packet.bytes,
)?;
self.in_flight.insert(
(packet.target_station, sequence),
InFlightReliableStationPacket {
source_station: packet.source_station,
target_station: packet.target_station,
sequence,
payload: packet.bytes,
first_sent_tick: now_tick,
last_sent_tick: now_tick,
attempts: 1,
},
);
self.stats.data_sent = self.stats.data_sent.saturating_add(1);
Ok(sequence)
}
pub fn acknowledge(&mut self, ack_source_station: StationId, sequence: u64) -> bool {
let removed = self
.in_flight
.remove(&(ack_source_station, sequence))
.is_some();
if removed {
self.stats.acks_received = self.stats.acks_received.saturating_add(1);
}
removed
}
pub fn retry_due<T: StationTransportSink>(
&mut self,
transport: &mut T,
now_tick: u64,
) -> Result<ReliableRetryReport, ReliableStationError<T::Error>> {
let keys = self
.in_flight
.iter()
.filter_map(|(key, packet)| {
let due =
now_tick.saturating_sub(packet.last_sent_tick) >= self.config.retry_after_ticks;
due.then_some(*key)
})
.collect::<Vec<_>>();
let mut report = ReliableRetryReport::default();
for key in keys {
let Some(packet) = self.in_flight.get(&key).cloned() else {
continue;
};
if packet.attempts >= self.config.max_attempts {
self.in_flight.remove(&key);
self.stats.timed_out = self.stats.timed_out.saturating_add(1);
report.timed_out = report.timed_out.saturating_add(1);
continue;
}
Self::send_data_frame(
transport,
packet.source_station,
packet.target_station,
packet.sequence,
&packet.payload,
)?;
if let Some(stored) = self.in_flight.get_mut(&key) {
stored.last_sent_tick = now_tick;
stored.attempts = stored.attempts.saturating_add(1);
}
self.stats.retries_sent = self.stats.retries_sent.saturating_add(1);
report.retried = report.retried.saturating_add(1);
}
Ok(report)
}
fn validate_payload<E>(&self, bytes: usize) -> Result<(), ReliableStationError<E>> {
if bytes > self.config.max_payload_bytes {
Err(ReliableStationError::PayloadTooLarge {
budget: self.config.max_payload_bytes,
actual: bytes,
})
} else {
Ok(())
}
}
fn allocate_sequence(&mut self, target_station: StationId) -> u64 {
let next = self.next_sequence.entry(target_station).or_insert(1);
let sequence = *next;
*next = next.saturating_add(1);
sequence
}
fn send_data_frame<T: StationTransportSink>(
transport: &mut T,
source_station: StationId,
target_station: StationId,
sequence: u64,
payload: &[u8],
) -> Result<(), ReliableStationError<T::Error>> {
let mut bytes = Vec::with_capacity(
payload
.len()
.saturating_add(RELIABLE_STATION_DATA_HEADER_BYTES),
);
ReliableStationFrame::Data {
sequence,
payload: payload.to_vec(),
}
.encode(&mut bytes)
.map_err(ReliableStationError::Encode)?;
transport
.send_station(StationOutboundPacket {
source_station,
target_station,
bytes,
})
.map_err(ReliableStationError::Transport)
}
}
impl Default for ReliableStationSender {
fn default() -> Self {
Self::new(ReliableStationConfig::default())
}
}
#[derive(Clone, Debug)]
pub struct ReliableStationReceiver {
config: ReliableStationConfig,
delivered: BTreeSet<(StationId, u64)>,
delivered_order: VecDeque<(StationId, u64)>,
stats: ReliableStationStats,
}
impl ReliableStationReceiver {
pub fn new(config: ReliableStationConfig) -> Self {
Self {
config,
delivered: BTreeSet::new(),
delivered_order: VecDeque::new(),
stats: ReliableStationStats::default(),
}
}
pub const fn config(&self) -> ReliableStationConfig {
self.config
}
pub const fn stats(&self) -> ReliableStationStats {
self.stats
}
pub fn handle_data<T: StationTransportSink>(
&mut self,
transport: &mut T,
packet: StationInboundPacket,
sequence: u64,
payload: Vec<u8>,
) -> Result<Option<StationInboundPacket>, ReliableStationError<T::Error>> {
let StationInboundPacket {
source_station,
target_station,
bytes: wire_bytes,
} = packet;
drop(wire_bytes);
self.send_ack(transport, target_station, source_station, sequence)?;
if !self.record_unique(source_station, sequence) {
self.stats.duplicates_suppressed = self.stats.duplicates_suppressed.saturating_add(1);
return Ok(None);
}
self.stats.data_delivered = self.stats.data_delivered.saturating_add(1);
Ok(Some(StationInboundPacket {
source_station,
target_station,
bytes: payload,
}))
}
fn send_ack<T: StationTransportSink>(
&mut self,
transport: &mut T,
source_station: StationId,
target_station: StationId,
sequence: u64,
) -> Result<(), ReliableStationError<T::Error>> {
let mut bytes = Vec::with_capacity(RELIABLE_STATION_ACK_BYTES);
ReliableStationFrame::Ack { sequence }
.encode(&mut bytes)
.map_err(ReliableStationError::Encode)?;
transport
.send_station(StationOutboundPacket {
source_station,
target_station,
bytes,
})
.map_err(ReliableStationError::Transport)?;
self.stats.acks_sent = self.stats.acks_sent.saturating_add(1);
Ok(())
}
fn record_unique(&mut self, source_station: StationId, sequence: u64) -> bool {
if self.config.max_delivered_history == 0 {
return true;
}
let key = (source_station, sequence);
if self.delivered.contains(&key) {
return false;
}
self.delivered.insert(key);
self.delivered_order.push_back(key);
while self.delivered_order.len() > self.config.max_delivered_history {
if let Some(old) = self.delivered_order.pop_front() {
self.delivered.remove(&old);
}
}
true
}
}
impl Default for ReliableStationReceiver {
fn default() -> Self {
Self::new(ReliableStationConfig::default())
}
}
#[derive(Clone, Debug)]
pub struct ReliableStationEndpoint {
pub sender: ReliableStationSender,
pub receiver: ReliableStationReceiver,
}
impl ReliableStationEndpoint {
pub fn new(config: ReliableStationConfig) -> Self {
Self {
sender: ReliableStationSender::new(config),
receiver: ReliableStationReceiver::new(config),
}
}
pub fn send<T: StationTransportSink>(
&mut self,
transport: &mut T,
packet: StationOutboundPacket,
now_tick: u64,
) -> Result<u64, ReliableStationError<T::Error>> {
self.sender.send(transport, packet, now_tick)
}
pub fn retry_due<T: StationTransportSink>(
&mut self,
transport: &mut T,
now_tick: u64,
) -> Result<ReliableRetryReport, ReliableStationError<T::Error>> {
self.sender.retry_due(transport, now_tick)
}
pub fn handle_inbound<T: StationTransportSink>(
&mut self,
transport: &mut T,
packet: StationInboundPacket,
) -> Result<Option<StationInboundPacket>, ReliableStationError<T::Error>> {
match ReliableStationFrame::decode(&packet.bytes).map_err(ReliableStationError::Decode)? {
ReliableStationFrame::Data { sequence, payload } => self
.receiver
.handle_data(transport, packet, sequence, payload),
ReliableStationFrame::Ack { sequence } => {
self.sender.acknowledge(packet.source_station, sequence);
Ok(None)
}
}
}
}
impl Default for ReliableStationEndpoint {
fn default() -> Self {
Self::new(ReliableStationConfig::default())
}
}
struct ReliableCursor<'a> {
input: &'a [u8],
offset: usize,
}
impl<'a> ReliableCursor<'a> {
const fn new(input: &'a [u8]) -> Self {
Self { input, offset: 0 }
}
fn read_u8(&mut self) -> Result<u8, ReliableStationDecodeError> {
self.require(1)?;
let value = self.input[self.offset];
self.offset += 1;
Ok(value)
}
fn read_u32(&mut self) -> Result<u32, ReliableStationDecodeError> {
let bytes = self.read_array::<4>()?;
Ok(u32::from_le_bytes(bytes))
}
fn read_u64(&mut self) -> Result<u64, ReliableStationDecodeError> {
let bytes = self.read_array::<8>()?;
Ok(u64::from_le_bytes(bytes))
}
fn read_array<const N: usize>(&mut self) -> Result<[u8; N], ReliableStationDecodeError> {
self.require(N)?;
let mut out = [0_u8; N];
out.copy_from_slice(&self.input[self.offset..self.offset + N]);
self.offset += N;
Ok(out)
}
fn read_bytes(&mut self, len: usize) -> Result<Vec<u8>, ReliableStationDecodeError> {
self.require(len)?;
let bytes = self.input[self.offset..self.offset + len].to_vec();
self.offset += len;
Ok(bytes)
}
fn require(&self, count: usize) -> Result<(), ReliableStationDecodeError> {
let needed = self.offset.saturating_add(count);
if needed > self.input.len() {
Err(ReliableStationDecodeError::Truncated {
needed,
available: self.input.len(),
})
} else {
Ok(())
}
}
fn finish(&self) -> Result<(), ReliableStationDecodeError> {
if self.offset == self.input.len() {
Ok(())
} else {
Err(ReliableStationDecodeError::TrailingBytes(
self.input.len().saturating_sub(self.offset),
))
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct StationTransportLimits {
pub max_queued_packets_per_station: usize,
pub max_packet_bytes: usize,
}
impl Default for StationTransportLimits {
fn default() -> Self {
Self {
max_queued_packets_per_station: 4096,
max_packet_bytes: 16 * 1024,
}
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct InMemoryStationTransportStats {
pub packets_sent: usize,
pub packets_received: usize,
pub bytes_sent: usize,
pub bytes_received: usize,
pub packets_rejected_full: usize,
pub packets_rejected_bytes: usize,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct UdpStationTransportStats {
pub packets_sent: usize,
pub packets_received: usize,
pub bytes_sent: usize,
pub bytes_received: usize,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum StationTransportError {
MissingTarget(StationId),
QueueFull {
station_id: StationId,
capacity: usize,
},
PacketTooLarge {
budget: usize,
actual: usize,
},
}
impl core::fmt::Display for StationTransportError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::MissingTarget(station_id) => {
write!(
f,
"station transport target {} is missing",
station_id.get()
)
}
Self::QueueFull {
station_id,
capacity,
} => write!(
f,
"station transport target {} queue is full at capacity {capacity}",
station_id.get()
),
Self::PacketTooLarge { budget, actual } => {
write!(
f,
"station transport packet exceeded byte budget: budget {budget}, actual {actual}"
)
}
}
}
}
impl std::error::Error for StationTransportError {}
#[derive(Clone, Debug)]
pub struct InMemoryStationTransport {
limits: StationTransportLimits,
queues: BTreeMap<StationId, VecDeque<StationInboundPacket>>,
stats: InMemoryStationTransportStats,
}
impl InMemoryStationTransport {
pub fn new(limits: StationTransportLimits) -> Self {
Self {
limits,
queues: BTreeMap::new(),
stats: InMemoryStationTransportStats::default(),
}
}
pub fn register_station(&mut self, station_id: StationId) {
self.queues
.entry(station_id)
.or_insert_with(|| VecDeque::with_capacity(self.limits.max_queued_packets_per_station));
}
pub fn queued_len(&self, station_id: StationId) -> Option<usize> {
self.queues.get(&station_id).map(VecDeque::len)
}
pub const fn limits(&self) -> StationTransportLimits {
self.limits
}
pub const fn stats(&self) -> InMemoryStationTransportStats {
self.stats
}
}
impl Default for InMemoryStationTransport {
fn default() -> Self {
Self::new(StationTransportLimits::default())
}
}
impl StationTransportSink for InMemoryStationTransport {
type Error = StationTransportError;
fn send_station(&mut self, packet: StationOutboundPacket) -> Result<(), Self::Error> {
let actual = packet.bytes.len();
if actual > self.limits.max_packet_bytes {
self.stats.packets_rejected_bytes = self.stats.packets_rejected_bytes.saturating_add(1);
return Err(StationTransportError::PacketTooLarge {
budget: self.limits.max_packet_bytes,
actual,
});
}
let queue = self
.queues
.get_mut(&packet.target_station)
.ok_or(StationTransportError::MissingTarget(packet.target_station))?;
if queue.len() >= self.limits.max_queued_packets_per_station {
self.stats.packets_rejected_full = self.stats.packets_rejected_full.saturating_add(1);
return Err(StationTransportError::QueueFull {
station_id: packet.target_station,
capacity: self.limits.max_queued_packets_per_station,
});
}
self.stats.packets_sent = self.stats.packets_sent.saturating_add(1);
self.stats.bytes_sent = self.stats.bytes_sent.saturating_add(actual);
queue.push_back(StationInboundPacket {
source_station: packet.source_station,
target_station: packet.target_station,
bytes: packet.bytes,
});
Ok(())
}
}
impl StationTransportReceiver for InMemoryStationTransport {
type Error = StationTransportError;
fn try_recv_station(
&mut self,
target_station: StationId,
) -> Result<Option<StationInboundPacket>, Self::Error> {
let queue = self
.queues
.get_mut(&target_station)
.ok_or(StationTransportError::MissingTarget(target_station))?;
let Some(packet) = queue.pop_front() else {
return Ok(None);
};
self.stats.packets_received = self.stats.packets_received.saturating_add(1);
self.stats.bytes_received = self.stats.bytes_received.saturating_add(packet.bytes.len());
Ok(Some(packet))
}
}
#[derive(Debug)]
pub enum UdpStationTransportError {
UnknownStation(StationId),
UnknownRemote(SocketAddr),
LocalStationMismatch {
local_station: StationId,
packet_source: StationId,
},
TargetStationMismatch {
local_station: StationId,
requested_target: StationId,
},
Io(io::Error),
}
impl core::fmt::Display for UdpStationTransportError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::UnknownStation(station_id) => {
write!(
f,
"udp station target {} is not registered",
station_id.get()
)
}
Self::UnknownRemote(addr) => {
write!(f, "udp station remote address {addr} is not registered")
}
Self::LocalStationMismatch {
local_station,
packet_source,
} => write!(
f,
"udp station local source mismatch: local {}, packet source {}",
local_station.get(),
packet_source.get()
),
Self::TargetStationMismatch {
local_station,
requested_target,
} => write!(
f,
"udp station receive target mismatch: local {}, requested {}",
local_station.get(),
requested_target.get()
),
Self::Io(error) => write!(f, "{error}"),
}
}
}
impl std::error::Error for UdpStationTransportError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Io(error) => Some(error),
Self::UnknownStation(_)
| Self::UnknownRemote(_)
| Self::LocalStationMismatch { .. }
| Self::TargetStationMismatch { .. } => None,
}
}
}
impl From<io::Error> for UdpStationTransportError {
fn from(error: io::Error) -> Self {
Self::Io(error)
}
}
#[derive(Debug)]
pub struct UdpStationTransport {
local_station: StationId,
socket: UdpSocket,
stations: HashMap<StationId, SocketAddr>,
addr_to_station: HashMap<SocketAddr, StationId>,
recv_buffer: Vec<u8>,
stats: UdpStationTransportStats,
}
impl UdpStationTransport {
pub fn bind<A: ToSocketAddrs>(local_station: StationId, addr: A) -> io::Result<Self> {
let socket = UdpSocket::bind(addr)?;
Self::from_socket(local_station, socket)
}
pub fn from_socket(local_station: StationId, socket: UdpSocket) -> io::Result<Self> {
socket.set_nonblocking(true)?;
Ok(Self {
local_station,
socket,
stations: HashMap::new(),
addr_to_station: HashMap::new(),
recv_buffer: vec![0; DEFAULT_UDP_RECV_BUFFER_SIZE],
stats: UdpStationTransportStats::default(),
})
}
pub const fn local_station(&self) -> StationId {
self.local_station
}
pub fn local_addr(&self) -> io::Result<SocketAddr> {
self.socket.local_addr()
}
pub const fn socket(&self) -> &UdpSocket {
&self.socket
}
pub fn socket_mut(&mut self) -> &mut UdpSocket {
&mut self.socket
}
pub fn register_station(
&mut self,
station_id: StationId,
addr: SocketAddr,
) -> Option<SocketAddr> {
let old_addr = self.stations.insert(station_id, addr);
if let Some(old_addr) = old_addr {
self.addr_to_station.remove(&old_addr);
}
if let Some(old_station) = self.addr_to_station.insert(addr, station_id)
&& old_station != station_id
{
self.stations.remove(&old_station);
}
old_addr
}
pub fn unregister_station(&mut self, station_id: StationId) -> Option<SocketAddr> {
let addr = self.stations.remove(&station_id)?;
self.addr_to_station.remove(&addr);
Some(addr)
}
pub fn station_addr(&self, station_id: StationId) -> Option<SocketAddr> {
self.stations.get(&station_id).copied()
}
pub fn station_for_addr(&self, addr: SocketAddr) -> Option<StationId> {
self.addr_to_station.get(&addr).copied()
}
pub fn set_recv_buffer_size(&mut self, bytes: usize) {
self.recv_buffer.resize(bytes.max(1), 0);
}
pub fn recv_buffer_size(&self) -> usize {
self.recv_buffer.len()
}
pub const fn stats(&self) -> UdpStationTransportStats {
self.stats
}
}
impl StationTransportSink for UdpStationTransport {
type Error = UdpStationTransportError;
fn send_station(&mut self, packet: StationOutboundPacket) -> Result<(), Self::Error> {
if packet.source_station != self.local_station {
return Err(UdpStationTransportError::LocalStationMismatch {
local_station: self.local_station,
packet_source: packet.source_station,
});
}
let addr = self.stations.get(&packet.target_station).copied().ok_or(
UdpStationTransportError::UnknownStation(packet.target_station),
)?;
let sent = self.socket.send_to(&packet.bytes, addr)?;
if sent != packet.bytes.len() {
return Err(io::Error::new(
io::ErrorKind::WriteZero,
"udp station socket reported a partial datagram send",
)
.into());
}
self.stats.packets_sent = self.stats.packets_sent.saturating_add(1);
self.stats.bytes_sent = self.stats.bytes_sent.saturating_add(packet.bytes.len());
Ok(())
}
}
impl StationTransportReceiver for UdpStationTransport {
type Error = UdpStationTransportError;
fn try_recv_station(
&mut self,
target_station: StationId,
) -> Result<Option<StationInboundPacket>, Self::Error> {
if target_station != self.local_station {
return Err(UdpStationTransportError::TargetStationMismatch {
local_station: self.local_station,
requested_target: target_station,
});
}
match self.socket.recv_from(&mut self.recv_buffer) {
Ok((len, remote_addr)) => {
let source_station = self
.addr_to_station
.get(&remote_addr)
.copied()
.ok_or(UdpStationTransportError::UnknownRemote(remote_addr))?;
self.stats.packets_received = self.stats.packets_received.saturating_add(1);
self.stats.bytes_received = self.stats.bytes_received.saturating_add(len);
Ok(Some(StationInboundPacket {
source_station,
target_station: self.local_station,
bytes: self.recv_buffer[..len].to_vec(),
}))
}
Err(error) if error.kind() == io::ErrorKind::WouldBlock => Ok(None),
Err(error) => Err(error.into()),
}
}
}
#[derive(Debug)]
pub enum UdpTransportError {
UnknownClient(ClientId),
Io(io::Error),
}
impl core::fmt::Display for UdpTransportError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::UnknownClient(client_id) => {
write!(f, "udp target client {} is not registered", client_id.get())
}
Self::Io(error) => write!(f, "{error}"),
}
}
}
impl std::error::Error for UdpTransportError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::UnknownClient(_) => None,
Self::Io(error) => Some(error),
}
}
}
impl From<io::Error> for UdpTransportError {
fn from(error: io::Error) -> Self {
Self::Io(error)
}
}
#[derive(Debug)]
pub struct UdpTransport {
socket: UdpSocket,
clients: HashMap<ClientId, SocketAddr>,
addr_to_client: HashMap<SocketAddr, ClientId>,
recv_buffer: Vec<u8>,
}
impl UdpTransport {
pub fn bind<A: ToSocketAddrs>(addr: A) -> io::Result<Self> {
let socket = UdpSocket::bind(addr)?;
Self::from_socket(socket)
}
pub fn from_socket(socket: UdpSocket) -> io::Result<Self> {
socket.set_nonblocking(true)?;
Ok(Self {
socket,
clients: HashMap::new(),
addr_to_client: HashMap::new(),
recv_buffer: vec![0; DEFAULT_UDP_RECV_BUFFER_SIZE],
})
}
pub fn local_addr(&self) -> io::Result<SocketAddr> {
self.socket.local_addr()
}
pub const fn socket(&self) -> &UdpSocket {
&self.socket
}
pub fn socket_mut(&mut self) -> &mut UdpSocket {
&mut self.socket
}
pub fn register_client(&mut self, client_id: ClientId, addr: SocketAddr) -> Option<SocketAddr> {
let old_addr = self.clients.insert(client_id, addr);
if let Some(old_addr) = old_addr {
self.addr_to_client.remove(&old_addr);
}
if let Some(old_client) = self.addr_to_client.insert(addr, client_id)
&& old_client != client_id
{
self.clients.remove(&old_client);
}
old_addr
}
pub fn unregister_client(&mut self, client_id: ClientId) -> Option<SocketAddr> {
let addr = self.clients.remove(&client_id)?;
self.addr_to_client.remove(&addr);
Some(addr)
}
pub fn client_addr(&self, client_id: ClientId) -> Option<SocketAddr> {
self.clients.get(&client_id).copied()
}
pub fn client_for_addr(&self, addr: SocketAddr) -> Option<ClientId> {
self.addr_to_client.get(&addr).copied()
}
pub fn set_recv_buffer_size(&mut self, bytes: usize) {
self.recv_buffer.resize(bytes.max(1), 0);
}
pub fn recv_buffer_size(&self) -> usize {
self.recv_buffer.len()
}
}
impl TransportSink for UdpTransport {
type Error = UdpTransportError;
fn send(&mut self, packet: OutboundPacket) -> Result<(), Self::Error> {
let addr = self
.clients
.get(&packet.client_id)
.copied()
.ok_or(UdpTransportError::UnknownClient(packet.client_id))?;
let sent = self.socket.send_to(&packet.bytes, addr)?;
if sent == packet.bytes.len() {
Ok(())
} else {
Err(io::Error::new(
io::ErrorKind::WriteZero,
"udp socket reported a partial datagram send",
)
.into())
}
}
}
impl TransportReceiver for UdpTransport {
type Error = UdpTransportError;
fn try_recv(&mut self) -> Result<Option<InboundPacket>, Self::Error> {
match self.socket.recv_from(&mut self.recv_buffer) {
Ok((len, remote_addr)) => Ok(Some(InboundPacket {
client_id: self.addr_to_client.get(&remote_addr).copied(),
remote_addr,
bytes: self.recv_buffer[..len].to_vec(),
})),
Err(error) if error.kind() == io::ErrorKind::WouldBlock => Ok(None),
Err(error) => Err(error.into()),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TransportError<E> {
Inner(E),
ByteBudgetExceeded {
budget: usize,
actual: usize,
},
}
impl<E: core::fmt::Display> core::fmt::Display for TransportError<E> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::Inner(error) => write!(f, "{error}"),
Self::ByteBudgetExceeded { budget, actual } => {
write!(
f,
"transport byte budget exceeded: budget {budget}, actual {actual}"
)
}
}
}
}
impl<E> std::error::Error for TransportError<E> where E: std::error::Error {}
#[derive(Clone, Debug)]
pub struct BudgetedTransport<T> {
inner: T,
max_packet_bytes: usize,
max_batch_bytes: usize,
}
impl<T> BudgetedTransport<T> {
pub const fn new(inner: T, max_packet_bytes: usize, max_batch_bytes: usize) -> Self {
Self {
inner,
max_packet_bytes,
max_batch_bytes,
}
}
pub fn into_inner(self) -> T {
self.inner
}
pub const fn inner(&self) -> &T {
&self.inner
}
}
impl<T: TransportSink> TransportSink for BudgetedTransport<T> {
type Error = TransportError<T::Error>;
fn send(&mut self, packet: OutboundPacket) -> Result<(), Self::Error> {
let bytes = packet.bytes.len();
if bytes > self.max_packet_bytes {
return Err(TransportError::ByteBudgetExceeded {
budget: self.max_packet_bytes,
actual: bytes,
});
}
self.inner.send(packet).map_err(TransportError::Inner)
}
fn send_batch(&mut self, batch: PacketBatch) -> Result<(), Self::Error> {
let bytes = batch.bytes_len();
if bytes > self.max_batch_bytes {
return Err(TransportError::ByteBudgetExceeded {
budget: self.max_batch_bytes,
actual: bytes,
});
}
for packet in &batch.packets {
if packet.bytes.len() > self.max_packet_bytes {
return Err(TransportError::ByteBudgetExceeded {
budget: self.max_packet_bytes,
actual: packet.bytes.len(),
});
}
}
self.inner.send_batch(batch).map_err(TransportError::Inner)
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct FakeTransport {
packets: usize,
bytes: usize,
batches: usize,
}
impl FakeTransport {
pub const fn packets_sent(&self) -> usize {
self.packets
}
pub const fn bytes_sent(&self) -> usize {
self.bytes
}
pub const fn batches_sent(&self) -> usize {
self.batches
}
}
impl TransportSink for FakeTransport {
type Error = core::convert::Infallible;
fn send(&mut self, packet: OutboundPacket) -> Result<(), Self::Error> {
self.packets += 1;
self.bytes += packet.bytes.len();
Ok(())
}
fn send_batch(&mut self, batch: PacketBatch) -> Result<(), Self::Error> {
self.batches += 1;
self.packets += batch.packets.len();
self.bytes += batch.bytes_len();
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::thread;
use std::time::Duration;
fn packet(bytes: usize) -> OutboundPacket {
OutboundPacket {
client_id: ClientId::new(1),
bytes: vec![0; bytes],
}
}
fn station_packet(bytes: usize) -> StationOutboundPacket {
StationOutboundPacket {
source_station: StationId::new(1),
target_station: StationId::new(2),
bytes: vec![1; bytes],
}
}
fn memory_addr(port: u16) -> SocketAddr {
SocketAddr::from(([127, 0, 0, 1], port))
}
fn recv_with_retry(transport: &mut UdpTransport) -> InboundPacket {
for _ in 0..50 {
if let Some(packet) = transport.try_recv().expect("udp receive should work") {
return packet;
}
thread::sleep(Duration::from_millis(2));
}
panic!("udp packet was not received");
}
fn recv_station_with_retry(
transport: &mut UdpStationTransport,
station_id: StationId,
) -> StationInboundPacket {
for _ in 0..50 {
if let Some(packet) = transport
.try_recv_station(station_id)
.expect("udp station receive should work")
{
return packet;
}
thread::sleep(Duration::from_millis(2));
}
panic!("udp station packet was not received");
}
#[derive(Clone, Debug, Default)]
struct TestAuthenticator;
impl PacketAuthenticator for TestAuthenticator {
type Error = core::convert::Infallible;
fn sign(
&mut self,
key_id: u32,
nonce: u64,
payload: &[u8],
out: &mut Vec<u8>,
) -> Result<(), Self::Error> {
out.extend_from_slice(&test_tag(key_id, nonce, payload));
Ok(())
}
fn verify(
&mut self,
key_id: u32,
nonce: u64,
payload: &[u8],
tag: &[u8],
) -> Result<bool, Self::Error> {
Ok(tag == test_tag(key_id, nonce, payload))
}
}
fn test_tag(key_id: u32, nonce: u64, payload: &[u8]) -> [u8; 8] {
let mut acc = u64::from(key_id)
.wrapping_mul(0x9E37_79B9_7F4A_7C15)
.wrapping_add(nonce.rotate_left(17));
for (index, byte) in payload.iter().copied().enumerate() {
acc = acc.rotate_left(5) ^ (u64::from(byte) << ((index % 8) * 8));
acc = acc.wrapping_mul(0x1000_0000_01B3);
}
acc.to_le_bytes()
}
#[test]
fn fake_transport_counts_batches_without_storing_packets() {
let mut batch = PacketBatch::new();
batch.push(packet(3));
batch.push(packet(5));
let mut transport = FakeTransport::default();
transport
.send_batch(batch)
.expect("fake transport is infallible");
assert_eq!(transport.batches_sent(), 1);
assert_eq!(transport.packets_sent(), 2);
assert_eq!(transport.bytes_sent(), 8);
}
#[test]
fn budgeted_transport_rejects_large_batch() {
let mut batch = PacketBatch::new();
batch.push(packet(8));
batch.push(packet(8));
let mut transport = BudgetedTransport::new(FakeTransport::default(), 16, 12);
let error = transport
.send_batch(batch)
.expect_err("batch should exceed budget");
assert_eq!(
error,
TransportError::ByteBudgetExceeded {
budget: 12,
actual: 16
}
);
assert_eq!(transport.inner().packets_sent(), 0);
}
#[test]
fn packet_security_envelope_roundtrips_and_enforces_limits() {
let config = PacketSecurityConfig {
max_payload_bytes: 4,
max_tag_bytes: 8,
max_replay_history: 4,
};
let envelope = PacketSecurityEnvelope {
key_id: 9,
nonce: 42,
payload: b"move".to_vec(),
tag: vec![1, 2, 3, 4],
};
let mut bytes = Vec::new();
envelope
.encode(config, &mut bytes)
.expect("envelope should encode");
assert_eq!(
PacketSecurityEnvelope::decode(config, &bytes).expect("envelope should decode"),
envelope
);
let too_large = PacketSecurityEnvelope {
key_id: 9,
nonce: 43,
payload: b"large".to_vec(),
tag: Vec::new(),
}
.encode(config, &mut Vec::new())
.expect_err("payload should exceed configured budget");
assert_eq!(
too_large,
PacketSecurityEncodeError::PayloadTooLarge {
budget: 4,
actual: 5
}
);
let mut bad = bytes;
bad[16..20].copy_from_slice(&5_u32.to_le_bytes());
assert_eq!(
PacketSecurityEnvelope::decode(config, &bad)
.expect_err("decoded payload length should exceed budget"),
PacketSecurityDecodeError::PayloadTooLarge {
budget: 4,
actual: 5
}
);
}
#[test]
fn packet_key_ring_selects_active_key_and_accepts_retiring_key() {
let now = Tick::new(10);
let mut ring = PacketKeyRing::new(PacketKeyRingConfig { max_keys: 4 });
ring.insert_active(1, now, 1)
.expect("first key should insert");
ring.insert_active(2, Tick::new(11), 10)
.expect("second key should insert");
assert_eq!(
ring.select_send_key(now).expect("key 1 should send").key_id,
1
);
assert_eq!(
ring.select_send_key(Tick::new(11))
.expect("key 2 should send")
.key_id,
2
);
ring.retire(1, Tick::new(12)).expect("key 1 should retire");
assert_eq!(
ring.accept_key(1, Tick::new(12))
.expect("retiring key should still receive")
.state,
PacketKeyState::Retiring
);
assert_eq!(
ring.select_send_key(Tick::new(12))
.expect("active key should win")
.key_id,
2
);
assert_eq!(ring.stats().keys_inserted, 2);
assert_eq!(ring.stats().keys_retired, 1);
}
#[test]
fn packet_key_ring_rejects_revoked_expired_and_over_capacity_keys() {
let mut ring = PacketKeyRing::new(PacketKeyRingConfig { max_keys: 2 });
ring.insert(PacketKeyDescriptor::active(1, Tick::new(1), 1).with_expiry(Tick::new(5)))
.expect("expiring key should insert");
ring.insert_active(2, Tick::new(1), 2)
.expect("second key should insert");
assert_eq!(
ring.insert_active(3, Tick::new(1), 3)
.expect_err("ring should be full"),
PacketKeyRingError::CapacityFull { capacity: 2 }
);
assert_eq!(
ring.insert_active(2, Tick::new(1), 2)
.expect_err("duplicate should reject"),
PacketKeyRingError::DuplicateKey(2)
);
assert!(ring.accept_key(1, Tick::new(4)).is_ok());
assert_eq!(
ring.accept_key(1, Tick::new(5))
.expect_err("expired key should reject"),
PacketKeyRingError::KeyNotAccepted {
key_id: 1,
state: PacketKeyState::Active
}
);
ring.revoke(2).expect("key should revoke");
assert_eq!(
ring.accept_key(2, Tick::new(4))
.expect_err("revoked key should reject"),
PacketKeyRingError::KeyNotAccepted {
key_id: 2,
state: PacketKeyState::Revoked
}
);
assert_eq!(ring.remove_expired(Tick::new(5)), 1);
assert_eq!(ring.len(), 1);
assert_eq!(ring.stats().keys_revoked, 1);
assert_eq!(ring.stats().keys_expired_removed, 1);
}
#[test]
fn packet_security_box_seals_opens_and_rejects_replay() {
let mut sender = PacketSecurityBox::new(
PacketSecurityConfig::default(),
TestAuthenticator,
PlaintextPacketCipher,
);
let mut receiver = PacketSecurityBox::new(
PacketSecurityConfig::default(),
TestAuthenticator,
PlaintextPacketCipher,
);
let sealed = sender.seal(7, b"command").expect("packet should seal");
assert_eq!(sender.stats().sealed, 1);
let opened = receiver.open(&sealed).expect("packet should open");
assert_eq!(opened, b"command");
assert_eq!(receiver.stats().opened, 1);
let replay = receiver.open(&sealed).expect_err("replay should reject");
match replay {
PacketSecurityError::Replay { key_id, nonce } => {
assert_eq!(key_id, 7);
assert_eq!(nonce, 1);
}
other => panic!("unexpected error: {other}"),
}
assert_eq!(receiver.stats().replay_rejected, 1);
}
#[test]
fn packet_security_box_uses_key_ring_for_rotation_policy() {
let mut sender_ring = PacketKeyRing::default();
let mut receiver_ring = PacketKeyRing::default();
sender_ring
.insert_active(7, Tick::new(10), 1)
.expect("sender key should insert");
receiver_ring
.insert_active(7, Tick::new(10), 1)
.expect("receiver key should insert");
let mut sender = PacketSecurityBox::new(
PacketSecurityConfig::default(),
TestAuthenticator,
PlaintextPacketCipher,
);
let mut receiver = PacketSecurityBox::new(
PacketSecurityConfig::default(),
TestAuthenticator,
PlaintextPacketCipher,
);
let sealed = sender
.seal_with_key_ring(&sender_ring, b"command", Tick::new(10))
.expect("packet should seal through selected key");
let envelope = PacketSecurityEnvelope::decode(PacketSecurityConfig::default(), &sealed)
.expect("packet should decode");
assert_eq!(envelope.key_id, 7);
assert_eq!(
receiver
.open_with_key_ring(&receiver_ring, &sealed, Tick::new(10))
.expect("packet should open through accepted key"),
b"command"
);
sender_ring
.insert_active(8, Tick::new(11), 10)
.expect("rotated sender key should insert");
receiver_ring
.insert_active(8, Tick::new(11), 10)
.expect("rotated receiver key should insert");
receiver_ring
.retire(7, Tick::new(11))
.expect("old key should retire");
let rotated = sender
.seal_with_key_ring(&sender_ring, b"ack", Tick::new(11))
.expect("rotated packet should seal");
let rotated_envelope =
PacketSecurityEnvelope::decode(PacketSecurityConfig::default(), &rotated)
.expect("rotated packet should decode");
assert_eq!(rotated_envelope.key_id, 8);
assert_eq!(
receiver
.open_with_key_ring(&receiver_ring, &rotated, Tick::new(11))
.expect("rotated packet should open"),
b"ack"
);
receiver_ring.revoke(7).expect("old key should revoke");
let stale = sender
.seal_with_nonce(7, 99, b"stale")
.expect("explicit stale-key packet should seal");
let error = receiver
.open_with_key_ring(&receiver_ring, &stale, Tick::new(12))
.expect_err("revoked key should reject before auth");
match error {
PacketSecurityError::Key(PacketKeyRingError::KeyNotAccepted { key_id, state }) => {
assert_eq!(key_id, 7);
assert_eq!(state, PacketKeyState::Revoked);
}
other => panic!("unexpected error: {other}"),
}
assert_eq!(receiver.stats().key_rejected, 1);
}
#[test]
fn packet_security_box_rejects_tampered_payload() {
let mut sender = PacketSecurityBox::new(
PacketSecurityConfig::default(),
TestAuthenticator,
PlaintextPacketCipher,
);
let mut receiver = PacketSecurityBox::new(
PacketSecurityConfig::default(),
TestAuthenticator,
PlaintextPacketCipher,
);
let mut sealed = sender
.seal_with_nonce(7, 10, b"command")
.expect("packet should seal");
let payload_offset = PACKET_SECURITY_HEADER_BYTES;
sealed[payload_offset] ^= 0x55;
let error = receiver
.open(&sealed)
.expect_err("tampered payload should reject");
match error {
PacketSecurityError::AuthenticationFailed { key_id, nonce } => {
assert_eq!(key_id, 7);
assert_eq!(nonce, 10);
}
other => panic!("unexpected error: {other}"),
}
assert_eq!(receiver.stats().auth_failed, 1);
assert_eq!(receiver.replay().len(), 0);
}
#[test]
fn packet_replay_window_bounds_history() {
let mut replay = PacketReplayWindow::new(2);
assert!(replay.accept(1, 1));
assert!(replay.accept(1, 2));
assert!(!replay.accept(1, 2));
assert!(replay.accept(1, 3));
assert_eq!(replay.len(), 2);
assert!(!replay.contains(1, 1));
assert!(replay.accept(1, 1));
}
#[test]
fn in_memory_transport_delivers_bounded_packets() {
let client_id = ClientId::new(7);
let server_id = ClientId::new(0);
let hub = InMemoryTransportHub::new(ClientTransportLimits {
max_queued_packets_per_client: 2,
max_packet_bytes: 8,
});
let mut client = hub
.endpoint(client_id, memory_addr(20007))
.expect("client should register");
let mut server = hub
.endpoint(server_id, memory_addr(20000))
.expect("server should register");
client
.send(OutboundPacket {
client_id: server_id,
bytes: b"command".to_vec(),
})
.expect("client packet should send");
assert_eq!(
hub.queued_len(server_id).expect("queue should exist"),
Some(1)
);
let inbound = server
.try_recv()
.expect("server receive should work")
.expect("packet should exist");
assert_eq!(inbound.client_id, Some(client_id));
assert_eq!(inbound.remote_addr, memory_addr(20007));
assert_eq!(inbound.bytes, b"command");
let stats = hub.stats().expect("stats should read");
assert_eq!(stats.packets_sent, 1);
assert_eq!(stats.packets_received, 1);
assert_eq!(stats.bytes_sent, 7);
assert_eq!(stats.bytes_received, 7);
}
#[test]
fn in_memory_transport_rejects_full_queue_and_large_packet() {
let client_id = ClientId::new(7);
let server_id = ClientId::new(0);
let hub = InMemoryTransportHub::new(ClientTransportLimits {
max_queued_packets_per_client: 1,
max_packet_bytes: 4,
});
let mut client = hub
.endpoint(client_id, memory_addr(20007))
.expect("client should register");
hub.endpoint(server_id, memory_addr(20000))
.expect("server should register");
client
.send(OutboundPacket {
client_id: server_id,
bytes: vec![0; 4],
})
.expect("first packet should send");
let full = client
.send(OutboundPacket {
client_id: server_id,
bytes: vec![0; 4],
})
.expect_err("queue should be full");
assert_eq!(
full,
InMemoryTransportError::QueueFull {
client_id: server_id,
capacity: 1
}
);
let large = client
.send(OutboundPacket {
client_id: server_id,
bytes: vec![0; 5],
})
.expect_err("packet should exceed budget");
assert_eq!(
large,
InMemoryTransportError::PacketTooLarge {
budget: 4,
actual: 5
}
);
let stats = hub.stats().expect("stats should read");
assert_eq!(stats.packets_rejected_full, 1);
assert_eq!(stats.packets_rejected_bytes, 1);
}
#[test]
fn reliable_client_frame_roundtrips_data_and_ack() {
let data = ReliableClientFrame::Data {
sequence: 42,
payload: b"command".to_vec(),
};
let mut bytes = Vec::new();
data.encode(&mut bytes).expect("data frame should encode");
assert_eq!(
ReliableClientFrame::decode(&bytes).expect("data frame should decode"),
data
);
let ack = ReliableClientFrame::Ack { sequence: 42 };
bytes.clear();
ack.encode(&mut bytes).expect("ack frame should encode");
assert_eq!(
ReliableClientFrame::decode(&bytes).expect("ack frame should decode"),
ack
);
}
#[test]
fn reliable_client_endpoint_delivers_payload_and_acknowledges() {
let client_id = ClientId::new(7);
let server_id = ClientId::new(0);
let hub = InMemoryTransportHub::default();
let mut client_transport = hub
.endpoint(client_id, memory_addr(20007))
.expect("client should register");
let mut server_transport = hub
.endpoint(server_id, memory_addr(20000))
.expect("server should register");
let mut client = ReliableClientEndpoint::default();
let mut server = ReliableClientEndpoint::default();
let sequence = client
.send(
&mut client_transport,
OutboundPacket {
client_id: server_id,
bytes: b"command".to_vec(),
},
0,
)
.expect("reliable command should send");
assert_eq!(sequence, 1);
assert_eq!(client.sender.in_flight_len(), 1);
let raw = server_transport
.try_recv()
.expect("server receive should work")
.expect("data packet should exist");
let delivered = server
.handle_inbound(&mut server_transport, raw)
.expect("data packet should handle")
.expect("first data packet should deliver");
assert_eq!(delivered.client_id, Some(client_id));
assert_eq!(delivered.remote_addr, memory_addr(20007));
assert_eq!(delivered.bytes, b"command");
assert_eq!(server.receiver.stats().data_delivered, 1);
assert_eq!(server.receiver.stats().acks_sent, 1);
let ack = client_transport
.try_recv()
.expect("client ACK receive should work")
.expect("ACK packet should exist");
assert_eq!(
client
.handle_inbound(&mut client_transport, ack)
.expect("ACK should handle"),
None
);
assert_eq!(client.sender.in_flight_len(), 0);
assert_eq!(client.sender.stats().acks_received, 1);
}
#[test]
fn reliable_client_endpoint_retries_and_suppresses_duplicate_delivery() {
let client_id = ClientId::new(7);
let server_id = ClientId::new(0);
let hub = InMemoryTransportHub::default();
let mut client_transport = hub
.endpoint(client_id, memory_addr(20007))
.expect("client should register");
let mut server_transport = hub
.endpoint(server_id, memory_addr(20000))
.expect("server should register");
let mut client = ReliableClientEndpoint::default();
let mut server = ReliableClientEndpoint::default();
client
.send(
&mut client_transport,
OutboundPacket {
client_id: server_id,
bytes: b"idempotent-command".to_vec(),
},
0,
)
.expect("reliable command should send");
let retry = client
.retry_due(&mut client_transport, 2)
.expect("retry should send");
assert_eq!(retry.retried, 1);
assert_eq!(retry.timed_out, 0);
assert_eq!(client.sender.stats().retries_sent, 1);
assert_eq!(
hub.queued_len(server_id).expect("queue should exist"),
Some(2)
);
let first_raw = server_transport
.try_recv()
.expect("server receive should work")
.expect("first data packet should exist");
let delivered = server
.handle_inbound(&mut server_transport, first_raw)
.expect("first data packet should handle")
.expect("first data packet should deliver");
assert_eq!(delivered.bytes, b"idempotent-command");
let duplicate_raw = server_transport
.try_recv()
.expect("server receive should work")
.expect("duplicate data packet should exist");
assert_eq!(
server
.handle_inbound(&mut server_transport, duplicate_raw)
.expect("duplicate data packet should handle"),
None
);
assert_eq!(server.receiver.stats().data_delivered, 1);
assert_eq!(server.receiver.stats().duplicates_suppressed, 1);
assert_eq!(server.receiver.stats().acks_sent, 2);
}
#[test]
fn reliable_client_receiver_bounds_duplicate_history() {
let client_id = ClientId::new(7);
let server_id = ClientId::new(0);
let hub = InMemoryTransportHub::default();
hub.endpoint(client_id, memory_addr(20007))
.expect("client should register");
let mut server_transport = hub
.endpoint(server_id, memory_addr(20000))
.expect("server should register");
let config = ReliableClientConfig {
max_in_flight_per_peer: 8,
retry_after_ticks: 2,
max_attempts: 4,
max_payload_bytes: DEFAULT_RELIABLE_CLIENT_MAX_PAYLOAD_BYTES,
max_delivered_history: 1,
};
let mut server = ReliableClientEndpoint::new(config);
let packet = |sequence: u64, payload: &[u8]| {
let mut bytes = Vec::new();
ReliableClientFrame::Data {
sequence,
payload: payload.to_vec(),
}
.encode(&mut bytes)
.expect("data frame should encode");
InboundPacket {
client_id: Some(client_id),
remote_addr: memory_addr(20007),
bytes,
}
};
assert!(
server
.handle_inbound(&mut server_transport, packet(1, b"first"))
.expect("first data packet should handle")
.is_some()
);
assert_eq!(
server
.handle_inbound(&mut server_transport, packet(1, b"first-duplicate"))
.expect("duplicate data packet should handle"),
None
);
assert!(
server
.handle_inbound(&mut server_transport, packet(2, b"second"))
.expect("second data packet should handle")
.is_some()
);
assert!(
server
.handle_inbound(&mut server_transport, packet(1, b"first-after-eviction"))
.expect("evicted data packet should handle")
.is_some()
);
assert_eq!(server.receiver.stats().data_delivered, 3);
assert_eq!(server.receiver.stats().duplicates_suppressed, 1);
}
#[test]
fn reliable_client_sender_enforces_payload_and_window_limits() {
let client_id = ClientId::new(7);
let server_id = ClientId::new(0);
let hub = InMemoryTransportHub::default();
let mut client_transport = hub
.endpoint(client_id, memory_addr(20007))
.expect("client should register");
hub.endpoint(server_id, memory_addr(20000))
.expect("server should register");
let config = ReliableClientConfig {
max_in_flight_per_peer: 1,
retry_after_ticks: 2,
max_attempts: 4,
max_payload_bytes: 4,
max_delivered_history: DEFAULT_RELIABLE_CLIENT_DELIVERED_HISTORY,
};
let mut client = ReliableClientEndpoint::new(config);
let too_large = client
.send(
&mut client_transport,
OutboundPacket {
client_id: server_id,
bytes: vec![0; 5],
},
0,
)
.expect_err("payload should exceed configured budget");
match too_large {
ReliableClientError::PayloadTooLarge { budget, actual } => {
assert_eq!(budget, 4);
assert_eq!(actual, 5);
}
other => panic!("unexpected error: {other}"),
}
client
.send(
&mut client_transport,
OutboundPacket {
client_id: server_id,
bytes: vec![0; 4],
},
0,
)
.expect("first packet should fit");
let full = client
.send(
&mut client_transport,
OutboundPacket {
client_id: server_id,
bytes: vec![1; 4],
},
0,
)
.expect_err("in-flight window should be full");
match full {
ReliableClientError::WindowFull {
peer_client,
capacity,
} => {
assert_eq!(peer_client, server_id);
assert_eq!(capacity, 1);
}
other => panic!("unexpected error: {other}"),
}
}
#[test]
fn reliable_client_endpoint_rejects_unknown_packet_source() {
let mut endpoint = ReliableClientEndpoint::default();
let mut transport = FakeTransport::default();
let mut bytes = Vec::new();
ReliableClientFrame::Ack { sequence: 1 }
.encode(&mut bytes)
.expect("ACK should encode");
let error = endpoint
.handle_inbound(
&mut transport,
InboundPacket {
client_id: None,
remote_addr: memory_addr(20007),
bytes,
},
)
.expect_err("unknown source should be rejected");
match error {
ReliableClientError::MissingSourceClient => {}
other => panic!("unexpected error: {other}"),
}
}
#[test]
fn in_memory_station_transport_delivers_bounded_packets() {
let mut transport = InMemoryStationTransport::new(StationTransportLimits {
max_queued_packets_per_station: 2,
max_packet_bytes: 8,
});
transport.register_station(StationId::new(2));
transport
.send_station(station_packet(4))
.expect("station packet should send");
assert_eq!(transport.queued_len(StationId::new(2)), Some(1));
let packet = transport
.try_recv_station(StationId::new(2))
.expect("receive should work")
.expect("packet should exist");
assert_eq!(packet.source_station, StationId::new(1));
assert_eq!(packet.target_station, StationId::new(2));
assert_eq!(packet.bytes, vec![1; 4]);
assert_eq!(transport.stats().packets_sent, 1);
assert_eq!(transport.stats().packets_received, 1);
}
#[test]
fn in_memory_station_transport_rejects_full_queue_and_large_packet() {
let mut transport = InMemoryStationTransport::new(StationTransportLimits {
max_queued_packets_per_station: 1,
max_packet_bytes: 4,
});
transport.register_station(StationId::new(2));
transport
.send_station(station_packet(4))
.expect("first packet should send");
let full = transport
.send_station(station_packet(4))
.expect_err("queue should be full");
assert_eq!(
full,
StationTransportError::QueueFull {
station_id: StationId::new(2),
capacity: 1
}
);
let large = transport
.send_station(station_packet(5))
.expect_err("packet should exceed budget");
assert_eq!(
large,
StationTransportError::PacketTooLarge {
budget: 4,
actual: 5
}
);
assert_eq!(transport.stats().packets_rejected_full, 1);
assert_eq!(transport.stats().packets_rejected_bytes, 1);
}
#[test]
fn reliable_station_frame_roundtrips_data_and_ack() {
let data = ReliableStationFrame::Data {
sequence: 42,
payload: b"station-event".to_vec(),
};
let mut bytes = Vec::new();
data.encode(&mut bytes).expect("data frame should encode");
assert_eq!(
ReliableStationFrame::decode(&bytes).expect("data frame should decode"),
data
);
let ack = ReliableStationFrame::Ack { sequence: 42 };
bytes.clear();
ack.encode(&mut bytes).expect("ack frame should encode");
assert_eq!(
ReliableStationFrame::decode(&bytes).expect("ack frame should decode"),
ack
);
}
#[test]
fn reliable_station_endpoint_delivers_payload_and_acknowledges() {
let station_one = StationId::new(1);
let station_two = StationId::new(2);
let mut transport = InMemoryStationTransport::default();
transport.register_station(station_one);
transport.register_station(station_two);
let mut first = ReliableStationEndpoint::default();
let mut second = ReliableStationEndpoint::default();
let sequence = first
.send(
&mut transport,
StationOutboundPacket {
source_station: station_one,
target_station: station_two,
bytes: b"event".to_vec(),
},
0,
)
.expect("reliable packet should send");
assert_eq!(sequence, 1);
assert_eq!(first.sender.in_flight_len(), 1);
let raw = transport
.try_recv_station(station_two)
.expect("receive should work")
.expect("data packet should exist");
let delivered = second
.handle_inbound(&mut transport, raw)
.expect("data packet should handle")
.expect("first data packet should deliver");
assert_eq!(delivered.source_station, station_one);
assert_eq!(delivered.target_station, station_two);
assert_eq!(delivered.bytes, b"event");
assert_eq!(second.receiver.stats().data_delivered, 1);
assert_eq!(second.receiver.stats().acks_sent, 1);
let ack = transport
.try_recv_station(station_one)
.expect("ack receive should work")
.expect("ack packet should exist");
assert_eq!(
first
.handle_inbound(&mut transport, ack)
.expect("ack should handle"),
None
);
assert_eq!(first.sender.in_flight_len(), 0);
assert_eq!(first.sender.stats().acks_received, 1);
}
#[test]
fn reliable_station_endpoint_retries_and_suppresses_duplicate_delivery() {
let station_one = StationId::new(1);
let station_two = StationId::new(2);
let mut transport = InMemoryStationTransport::default();
transport.register_station(station_one);
transport.register_station(station_two);
let mut first = ReliableStationEndpoint::default();
let mut second = ReliableStationEndpoint::default();
first
.send(
&mut transport,
StationOutboundPacket {
source_station: station_one,
target_station: station_two,
bytes: b"idempotent-event".to_vec(),
},
0,
)
.expect("reliable packet should send");
let retry = first
.retry_due(&mut transport, 2)
.expect("retry should send");
assert_eq!(retry.retried, 1);
assert_eq!(retry.timed_out, 0);
assert_eq!(first.sender.stats().retries_sent, 1);
assert_eq!(transport.queued_len(station_two), Some(2));
let first_raw = transport
.try_recv_station(station_two)
.expect("receive should work")
.expect("first data packet should exist");
let delivered = second
.handle_inbound(&mut transport, first_raw)
.expect("first data packet should handle")
.expect("first data packet should deliver");
assert_eq!(delivered.bytes, b"idempotent-event");
let duplicate_raw = transport
.try_recv_station(station_two)
.expect("receive should work")
.expect("duplicate data packet should exist");
assert_eq!(
second
.handle_inbound(&mut transport, duplicate_raw)
.expect("duplicate data packet should handle"),
None
);
assert_eq!(second.receiver.stats().data_delivered, 1);
assert_eq!(second.receiver.stats().duplicates_suppressed, 1);
assert_eq!(second.receiver.stats().acks_sent, 2);
}
#[test]
fn reliable_station_receiver_bounds_duplicate_history() {
let station_one = StationId::new(1);
let station_two = StationId::new(2);
let mut transport = InMemoryStationTransport::default();
transport.register_station(station_one);
transport.register_station(station_two);
let config = ReliableStationConfig {
max_in_flight_per_target: 8,
retry_after_ticks: 2,
max_attempts: 4,
max_payload_bytes: DEFAULT_RELIABLE_STATION_MAX_PAYLOAD_BYTES,
max_delivered_history: 1,
};
let mut endpoint = ReliableStationEndpoint::new(config);
let packet = |sequence: u64, payload: &[u8]| {
let mut bytes = Vec::new();
ReliableStationFrame::Data {
sequence,
payload: payload.to_vec(),
}
.encode(&mut bytes)
.expect("data frame should encode");
StationInboundPacket {
source_station: station_one,
target_station: station_two,
bytes,
}
};
assert!(
endpoint
.handle_inbound(&mut transport, packet(1, b"first"))
.expect("first data packet should handle")
.is_some()
);
assert_eq!(
endpoint
.handle_inbound(&mut transport, packet(1, b"first-duplicate"))
.expect("duplicate data packet should handle"),
None
);
assert!(
endpoint
.handle_inbound(&mut transport, packet(2, b"second"))
.expect("second data packet should handle")
.is_some()
);
assert!(
endpoint
.handle_inbound(&mut transport, packet(1, b"first-after-eviction"))
.expect("evicted data packet should handle")
.is_some()
);
assert_eq!(endpoint.receiver.stats().data_delivered, 3);
assert_eq!(endpoint.receiver.stats().duplicates_suppressed, 1);
}
#[test]
fn reliable_station_sender_enforces_payload_and_window_limits() {
let station_one = StationId::new(1);
let station_two = StationId::new(2);
let mut transport = InMemoryStationTransport::default();
transport.register_station(station_one);
transport.register_station(station_two);
let config = ReliableStationConfig {
max_in_flight_per_target: 1,
retry_after_ticks: 2,
max_attempts: 4,
max_payload_bytes: 4,
max_delivered_history: DEFAULT_RELIABLE_STATION_DELIVERED_HISTORY,
};
let mut endpoint = ReliableStationEndpoint::new(config);
let too_large = endpoint
.send(
&mut transport,
StationOutboundPacket {
source_station: station_one,
target_station: station_two,
bytes: vec![0; 5],
},
0,
)
.expect_err("payload should exceed configured budget");
match too_large {
ReliableStationError::PayloadTooLarge { budget, actual } => {
assert_eq!(budget, 4);
assert_eq!(actual, 5);
}
other => panic!("unexpected error: {other}"),
}
endpoint
.send(
&mut transport,
StationOutboundPacket {
source_station: station_one,
target_station: station_two,
bytes: vec![0; 4],
},
0,
)
.expect("first packet should fit");
let full = endpoint
.send(
&mut transport,
StationOutboundPacket {
source_station: station_one,
target_station: station_two,
bytes: vec![1; 4],
},
0,
)
.expect_err("in-flight window should be full");
match full {
ReliableStationError::WindowFull {
target_station,
capacity,
} => {
assert_eq!(target_station, station_two);
assert_eq!(capacity, 1);
}
other => panic!("unexpected error: {other}"),
}
}
#[test]
fn udp_transport_sends_and_receives_registered_client() {
let client_id = ClientId::new(7);
let server_id = ClientId::new(0);
let mut server = UdpTransport::bind("127.0.0.1:0").expect("server should bind");
let mut client = UdpTransport::bind("127.0.0.1:0").expect("client should bind");
let server_addr = server.local_addr().expect("server addr should exist");
let client_addr = client.local_addr().expect("client addr should exist");
server.register_client(client_id, client_addr);
client.register_client(server_id, server_addr);
client
.send(OutboundPacket {
client_id: server_id,
bytes: b"command".to_vec(),
})
.expect("client should send");
let inbound = recv_with_retry(&mut server);
assert_eq!(inbound.client_id, Some(client_id));
assert_eq!(inbound.remote_addr, client_addr);
assert_eq!(inbound.bytes, b"command");
server
.send(OutboundPacket {
client_id,
bytes: b"replication".to_vec(),
})
.expect("server should send");
let inbound = recv_with_retry(&mut client);
assert_eq!(inbound.client_id, Some(server_id));
assert_eq!(inbound.remote_addr, server_addr);
assert_eq!(inbound.bytes, b"replication");
}
#[test]
fn udp_transport_rejects_unknown_client() {
let mut transport = UdpTransport::bind("127.0.0.1:0").expect("transport should bind");
let error = transport
.send(OutboundPacket {
client_id: ClientId::new(99),
bytes: Vec::new(),
})
.expect_err("unknown client should fail");
match error {
UdpTransportError::UnknownClient(client_id) => {
assert_eq!(client_id, ClientId::new(99));
}
UdpTransportError::Io(error) => panic!("unexpected io error: {error}"),
}
}
#[test]
fn udp_station_transport_sends_and_receives_registered_stations() {
let station_one = StationId::new(1);
let station_two = StationId::new(2);
let mut first =
UdpStationTransport::bind(station_one, "127.0.0.1:0").expect("first should bind");
let mut second =
UdpStationTransport::bind(station_two, "127.0.0.1:0").expect("second should bind");
let first_addr = first.local_addr().expect("first addr should exist");
let second_addr = second.local_addr().expect("second addr should exist");
first.register_station(station_two, second_addr);
second.register_station(station_one, first_addr);
first
.send_station(StationOutboundPacket {
source_station: station_one,
target_station: station_two,
bytes: b"handoff-prepare".to_vec(),
})
.expect("first station should send");
let inbound = recv_station_with_retry(&mut second, station_two);
assert_eq!(inbound.source_station, station_one);
assert_eq!(inbound.target_station, station_two);
assert_eq!(inbound.bytes, b"handoff-prepare");
second
.send_station(StationOutboundPacket {
source_station: station_two,
target_station: station_one,
bytes: b"handoff-commit".to_vec(),
})
.expect("second station should send");
let inbound = recv_station_with_retry(&mut first, station_one);
assert_eq!(inbound.source_station, station_two);
assert_eq!(inbound.target_station, station_one);
assert_eq!(inbound.bytes, b"handoff-commit");
assert_eq!(first.stats().packets_sent, 1);
assert_eq!(first.stats().packets_received, 1);
assert_eq!(second.stats().packets_sent, 1);
assert_eq!(second.stats().packets_received, 1);
}
#[test]
fn udp_station_transport_rejects_invalid_station_endpoints() {
let local = StationId::new(1);
let mut transport =
UdpStationTransport::bind(local, "127.0.0.1:0").expect("transport should bind");
let source_mismatch = transport
.send_station(StationOutboundPacket {
source_station: StationId::new(9),
target_station: StationId::new(2),
bytes: Vec::new(),
})
.expect_err("source should match local station");
match source_mismatch {
UdpStationTransportError::LocalStationMismatch {
local_station,
packet_source,
} => {
assert_eq!(local_station, local);
assert_eq!(packet_source, StationId::new(9));
}
other => panic!("unexpected error: {other}"),
}
let unknown = transport
.send_station(StationOutboundPacket {
source_station: local,
target_station: StationId::new(2),
bytes: Vec::new(),
})
.expect_err("target station should be registered");
match unknown {
UdpStationTransportError::UnknownStation(station_id) => {
assert_eq!(station_id, StationId::new(2));
}
other => panic!("unexpected error: {other}"),
}
let target_mismatch = transport
.try_recv_station(StationId::new(99))
.expect_err("receive target should match local station");
match target_mismatch {
UdpStationTransportError::TargetStationMismatch {
local_station,
requested_target,
} => {
assert_eq!(local_station, local);
assert_eq!(requested_target, StationId::new(99));
}
other => panic!("unexpected error: {other}"),
}
}
}