use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap, VecDeque};
use std::fmt;
pub const FLAG_FIN: u8 = 0b0000_0001;
pub const FLAG_RST: u8 = 0b0000_0010;
pub const FLAG_SYN: u8 = 0b0000_0100;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub struct FrameFlags(pub u8);
impl FrameFlags {
pub const SYN: u8 = 0x01;
pub const FIN: u8 = 0x02;
pub const RST: u8 = 0x04;
pub const ACK: u8 = 0x08;
pub const DATA: u8 = 0x10;
#[inline]
pub fn new(raw: u8) -> Self {
Self(raw)
}
#[inline]
pub fn is_syn(self) -> bool {
self.0 & Self::SYN != 0
}
#[inline]
pub fn is_fin(self) -> bool {
self.0 & Self::FIN != 0
}
#[inline]
pub fn is_rst(self) -> bool {
self.0 & Self::RST != 0
}
#[inline]
pub fn is_ack(self) -> bool {
self.0 & Self::ACK != 0
}
#[inline]
pub fn is_data(self) -> bool {
self.0 & Self::DATA != 0
}
#[inline]
pub fn with(self, flag: u8) -> Self {
Self(self.0 | flag)
}
}
impl fmt::Display for FrameFlags {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut parts: Vec<&str> = Vec::new();
if self.is_syn() {
parts.push("SYN");
}
if self.is_fin() {
parts.push("FIN");
}
if self.is_rst() {
parts.push("RST");
}
if self.is_ack() {
parts.push("ACK");
}
if self.is_data() {
parts.push("DATA");
}
if parts.is_empty() {
write!(f, "NONE")
} else {
write!(f, "{}", parts.join("|"))
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct StreamId(pub u32);
impl fmt::Display for StreamId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "stream:{}", self.0)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum StreamPriority {
Background = 0,
Low = 1,
Normal = 2,
High = 4,
Critical = 8,
}
impl StreamPriority {
#[inline]
pub fn weight(self) -> u32 {
match self {
Self::Critical => 8,
Self::High => 4,
Self::Normal => 2,
Self::Low => 1,
Self::Background => 0,
}
}
}
pub fn priority_from_u8(p: u8) -> StreamPriority {
match p {
0 => StreamPriority::Background,
1..=63 => StreamPriority::Low,
64..=127 => StreamPriority::Normal,
128..=191 => StreamPriority::High,
192..=255 => StreamPriority::Critical,
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StreamState {
Opening,
Open,
HalfClosed,
Closed,
Reset,
}
const FRAME_HEADER_LEN: usize = 25;
#[derive(Debug, Clone)]
pub struct StreamFrame {
pub stream_id: StreamId,
pub sequence: u64,
pub data: Vec<u8>,
pub flags: u8,
pub timestamp: u64,
}
impl StreamFrame {
#[inline]
pub fn is_fin(&self) -> bool {
self.flags & FLAG_FIN != 0
}
#[inline]
pub fn is_rst(&self) -> bool {
self.flags & FLAG_RST != 0
}
#[inline]
pub fn is_syn(&self) -> bool {
self.flags & FLAG_SYN != 0
}
#[inline]
pub fn is_control(&self) -> bool {
self.flags != 0
}
pub fn encode(&self) -> Vec<u8> {
let payload_len = self.data.len();
let mut buf = Vec::with_capacity(FRAME_HEADER_LEN + payload_len);
buf.extend_from_slice(&self.stream_id.0.to_le_bytes());
buf.extend_from_slice(&self.sequence.to_le_bytes());
buf.extend_from_slice(&(payload_len as u32).to_le_bytes());
buf.push(self.flags);
buf.extend_from_slice(&self.timestamp.to_le_bytes());
buf.extend_from_slice(&self.data);
buf
}
pub fn decode(data: &[u8]) -> Result<Self, MuxError> {
if data.len() < FRAME_HEADER_LEN {
return Err(MuxError::FrameTooLarge(data.len()));
}
let stream_id = u32::from_le_bytes(
data[0..4]
.try_into()
.map_err(|_| MuxError::FrameTooLarge(data.len()))?,
);
let sequence = u64::from_le_bytes(
data[4..12]
.try_into()
.map_err(|_| MuxError::FrameTooLarge(data.len()))?,
);
let payload_len = u32::from_le_bytes(
data[12..16]
.try_into()
.map_err(|_| MuxError::FrameTooLarge(data.len()))?,
) as usize;
let flags = data[16];
let timestamp = u64::from_le_bytes(
data[17..25]
.try_into()
.map_err(|_| MuxError::FrameTooLarge(data.len()))?,
);
let expected_total = FRAME_HEADER_LEN + payload_len;
if data.len() < expected_total {
return Err(MuxError::FrameTooLarge(data.len()));
}
let payload = data[FRAME_HEADER_LEN..FRAME_HEADER_LEN + payload_len].to_vec();
Ok(Self {
stream_id: StreamId(stream_id),
sequence,
data: payload,
flags,
timestamp,
})
}
}
#[derive(Debug, Clone)]
pub struct StreamInfo {
pub id: StreamId,
pub state: StreamState,
pub bytes_sent: u64,
pub bytes_received: u64,
pub frames_sent: u64,
pub frames_received: u64,
pub opened_at: u64,
pub last_activity: u64,
pub priority: u8,
}
#[derive(Debug, Clone)]
pub struct LogicalStream {
pub id: StreamId,
pub priority: StreamPriority,
pub state: StreamState,
pub send_window: u32,
pub recv_window: u32,
pub send_seq: u64,
pub recv_seq: u64,
pub bytes_sent: u64,
pub bytes_received: u64,
pub frames_sent: u64,
pub frames_received: u64,
pub opened_at: u64,
pub last_activity: u64,
pub priority_raw: u8,
recv_buffer: VecDeque<Vec<u8>>,
recv_buffer_max: usize,
}
#[derive(Debug, Clone)]
pub struct MultiplexerConfig {
pub max_streams: usize,
pub default_window_size: u32,
pub max_frame_size: usize,
pub enable_priority: bool,
pub recv_buffer_capacity: usize,
pub send_buffer_capacity: usize,
pub idle_timeout_us: u64,
}
impl Default for MultiplexerConfig {
fn default() -> Self {
Self {
max_streams: 256,
default_window_size: 65_536,
max_frame_size: 16_384,
enable_priority: true,
recv_buffer_capacity: 256,
send_buffer_capacity: 256,
idle_timeout_us: 30_000_000, }
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MultiplexerStats {
pub active_streams: usize,
pub total_streams_opened: u64,
pub total_frames_sent: u64,
pub total_frames_received: u64,
pub total_bytes_sent: u64,
pub total_bytes_received: u64,
pub dropped_frames: u64,
}
#[derive(Debug)]
struct PrioritizedFrame {
priority_weight: u32,
sequence: u64,
frame: StreamFrame,
}
impl PartialEq for PrioritizedFrame {
fn eq(&self, other: &Self) -> bool {
self.priority_weight == other.priority_weight && self.sequence == other.sequence
}
}
impl Eq for PrioritizedFrame {}
impl PartialOrd for PrioritizedFrame {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl Ord for PrioritizedFrame {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.priority_weight
.cmp(&other.priority_weight)
.then_with(|| Reverse(self.sequence).cmp(&Reverse(other.sequence)))
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MuxEvent {
StreamOpened(StreamId),
StreamClosed(StreamId),
StreamReset(StreamId),
FrameReceived(StreamId),
SendBufferFull(StreamId),
IdleStreamExpired(StreamId),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MuxError {
MaxStreamsReached,
StreamNotFound(u32),
StreamNotOpen(u32),
WindowExhausted(u32),
SequenceError {
expected: u64,
got: u64,
},
FrameTooLarge(usize),
BufferOverflow(u32),
StreamAlreadyOpen(u32),
}
impl fmt::Display for MuxError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::MaxStreamsReached => write!(f, "maximum number of streams reached"),
Self::StreamNotFound(id) => write!(f, "stream {id} not found"),
Self::StreamNotOpen(id) => write!(f, "stream {id} is not open"),
Self::WindowExhausted(id) => write!(f, "send window exhausted on stream {id}"),
Self::SequenceError { expected, got } => {
write!(f, "sequence error: expected {expected}, got {got}")
}
Self::FrameTooLarge(len) => write!(f, "frame data too short to decode ({len} bytes)"),
Self::BufferOverflow(id) => write!(f, "buffer overflow on stream {id}"),
Self::StreamAlreadyOpen(id) => write!(f, "stream {id} is already open"),
}
}
}
impl std::error::Error for MuxError {}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MuxStats {
pub total_streams: usize,
pub open_streams: usize,
pub queued_frames: usize,
pub total_frames_sent: u64,
pub total_frames_received: u64,
pub total_bytes_sent: u64,
}
pub struct StreamMultiplexer {
pub config: MultiplexerConfig,
streams: HashMap<StreamId, LogicalStream>,
send_queue: BinaryHeap<PrioritizedFrame>,
next_stream_id: u32,
enqueue_seq: u64,
total_frames_sent: u64,
total_frames_received: u64,
total_streams_opened: u64,
dropped_frames: u64,
}
impl StreamMultiplexer {
pub fn new(config: MultiplexerConfig) -> Self {
Self {
streams: HashMap::with_capacity(config.max_streams.min(256)),
send_queue: BinaryHeap::new(),
next_stream_id: 1,
enqueue_seq: 0,
total_frames_sent: 0,
total_frames_received: 0,
total_streams_opened: 0,
dropped_frames: 0,
config,
}
}
pub fn open_stream(&mut self, priority: StreamPriority) -> Result<StreamId, MuxError> {
if self.streams.len() >= self.config.max_streams {
return Err(MuxError::MaxStreamsReached);
}
let id = StreamId(self.next_stream_id);
self.next_stream_id = self.next_stream_id.wrapping_add(1);
let recv_cap = self.config.recv_buffer_capacity;
let stream = LogicalStream {
id,
priority,
state: StreamState::Open,
send_window: self.config.default_window_size,
recv_window: self.config.default_window_size,
send_seq: 1, recv_seq: 0,
bytes_sent: 0,
bytes_received: 0,
frames_sent: 1, frames_received: 0,
opened_at: 0,
last_activity: 0,
priority_raw: priority.weight() as u8,
recv_buffer: VecDeque::with_capacity(recv_cap.min(64)),
recv_buffer_max: recv_cap,
};
self.streams.insert(id, stream);
self.total_streams_opened = self.total_streams_opened.wrapping_add(1);
self.enqueue_control(id, FLAG_SYN, 0, priority, 0);
Ok(id)
}
pub fn open_stream_with_priority(&mut self, priority_raw: u8) -> Result<StreamId, MuxError> {
if self.streams.len() >= self.config.max_streams {
return Err(MuxError::MaxStreamsReached);
}
let priority = priority_from_u8(priority_raw);
let id = StreamId(self.next_stream_id);
self.next_stream_id = self.next_stream_id.wrapping_add(1);
let recv_cap = self.config.recv_buffer_capacity;
let stream = LogicalStream {
id,
priority,
state: StreamState::Open,
send_window: self.config.default_window_size,
recv_window: self.config.default_window_size,
send_seq: 1,
recv_seq: 0,
bytes_sent: 0,
bytes_received: 0,
frames_sent: 1,
frames_received: 0,
opened_at: 0,
last_activity: 0,
priority_raw,
recv_buffer: VecDeque::with_capacity(recv_cap.min(64)),
recv_buffer_max: recv_cap,
};
self.streams.insert(id, stream);
self.total_streams_opened = self.total_streams_opened.wrapping_add(1);
self.enqueue_control(id, FLAG_SYN, 0, priority, 0);
Ok(id)
}
pub fn close_stream(&mut self, stream_id: StreamId) -> Result<(), MuxError> {
let stream = self
.streams
.get_mut(&stream_id)
.ok_or(MuxError::StreamNotFound(stream_id.0))?;
match stream.state {
StreamState::Closed | StreamState::Reset => {
return Err(MuxError::StreamNotOpen(stream_id.0));
}
_ => {}
}
let seq = stream.send_seq;
stream.send_seq = stream.send_seq.wrapping_add(1);
stream.frames_sent = stream.frames_sent.wrapping_add(1);
let priority = stream.priority;
stream.state = StreamState::HalfClosed;
self.enqueue_control(stream_id, FLAG_FIN, seq, priority, 0);
Ok(())
}
pub fn reset_stream(&mut self, stream_id: StreamId) -> Result<(), MuxError> {
let stream = self
.streams
.get_mut(&stream_id)
.ok_or(MuxError::StreamNotFound(stream_id.0))?;
let seq = stream.send_seq;
stream.send_seq = stream.send_seq.wrapping_add(1);
stream.frames_sent = stream.frames_sent.wrapping_add(1);
let priority = stream.priority;
stream.state = StreamState::Reset;
self.enqueue_control(stream_id, FLAG_RST, seq, priority, 0);
Ok(())
}
pub fn send(
&mut self,
stream_id: StreamId,
data: Vec<u8>,
now: u64,
) -> Result<usize, MuxError> {
{
let stream = self
.streams
.get(&stream_id)
.ok_or(MuxError::StreamNotFound(stream_id.0))?;
if stream.state != StreamState::Open {
return Err(MuxError::StreamNotOpen(stream_id.0));
}
if data.len() as u64 > stream.send_window as u64 {
return Err(MuxError::WindowExhausted(stream_id.0));
}
}
let total = data.len();
let chunk_size = self.config.max_frame_size.max(1);
let mut offset = 0usize;
while offset < data.len() {
let end = (offset + chunk_size).min(data.len());
let chunk = data[offset..end].to_vec();
let chunk_len = chunk.len();
let stream = self
.streams
.get_mut(&stream_id)
.ok_or(MuxError::StreamNotFound(stream_id.0))?;
let seq = stream.send_seq;
stream.send_seq = stream.send_seq.wrapping_add(1);
stream.bytes_sent = stream.bytes_sent.saturating_add(chunk_len as u64);
stream.frames_sent = stream.frames_sent.wrapping_add(1);
stream.send_window = stream.send_window.saturating_sub(chunk_len as u32);
stream.last_activity = now;
let priority = stream.priority;
let weight = if self.config.enable_priority {
priority.weight()
} else {
0
};
let enqueue_seq = self.enqueue_seq;
self.enqueue_seq = self.enqueue_seq.wrapping_add(1);
self.send_queue.push(PrioritizedFrame {
priority_weight: weight,
sequence: enqueue_seq,
frame: StreamFrame {
stream_id,
sequence: seq,
data: chunk,
flags: 0,
timestamp: now,
},
});
offset = end;
}
Ok(total)
}
pub fn receive_events(
&mut self,
frame: StreamFrame,
now: u64,
) -> Result<Vec<MuxEvent>, MuxError> {
let mut events = Vec::new();
if frame.is_syn() {
if self.streams.contains_key(&frame.stream_id) {
return Err(MuxError::StreamAlreadyOpen(frame.stream_id.0));
}
if self.streams.len() >= self.config.max_streams {
return Err(MuxError::MaxStreamsReached);
}
let recv_cap = self.config.recv_buffer_capacity;
let stream = LogicalStream {
id: frame.stream_id,
priority: StreamPriority::Normal,
state: StreamState::Open,
send_window: self.config.default_window_size,
recv_window: self.config.default_window_size,
send_seq: 0,
recv_seq: 1, bytes_sent: 0,
bytes_received: 0,
frames_sent: 0,
frames_received: 1,
opened_at: now,
last_activity: now,
priority_raw: StreamPriority::Normal.weight() as u8,
recv_buffer: VecDeque::with_capacity(recv_cap.min(64)),
recv_buffer_max: recv_cap,
};
self.streams.insert(frame.stream_id, stream);
self.total_streams_opened = self.total_streams_opened.wrapping_add(1);
self.total_frames_received = self.total_frames_received.wrapping_add(1);
events.push(MuxEvent::StreamOpened(frame.stream_id));
return Ok(events);
}
let stream = self
.streams
.get_mut(&frame.stream_id)
.ok_or(MuxError::StreamNotFound(frame.stream_id.0))?;
if frame.is_rst() {
stream.state = StreamState::Reset;
stream.last_activity = now;
stream.frames_received = stream.frames_received.wrapping_add(1);
self.total_frames_received = self.total_frames_received.wrapping_add(1);
events.push(MuxEvent::StreamReset(frame.stream_id));
return Ok(events);
}
if frame.sequence != stream.recv_seq {
return Err(MuxError::SequenceError {
expected: stream.recv_seq,
got: frame.sequence,
});
}
stream.recv_seq = stream.recv_seq.wrapping_add(1);
let data_len = frame.data.len();
stream.frames_received = stream.frames_received.wrapping_add(1);
stream.last_activity = now;
self.total_frames_received = self.total_frames_received.wrapping_add(1);
if frame.is_fin() {
stream.bytes_received = stream.bytes_received.saturating_add(data_len as u64);
stream.state = if stream.state == StreamState::HalfClosed {
StreamState::Closed
} else {
StreamState::HalfClosed
};
events.push(MuxEvent::StreamClosed(frame.stream_id));
} else if !frame.data.is_empty() {
stream.bytes_received = stream.bytes_received.saturating_add(data_len as u64);
if stream.recv_buffer.len() >= stream.recv_buffer_max {
self.dropped_frames = self.dropped_frames.wrapping_add(1);
events.push(MuxEvent::SendBufferFull(frame.stream_id));
} else {
stream.recv_buffer.push_back(frame.data);
events.push(MuxEvent::FrameReceived(frame.stream_id));
}
} else {
events.push(MuxEvent::FrameReceived(frame.stream_id));
}
Ok(events)
}
pub fn receive(&mut self, frame: StreamFrame, _now: u64) -> Result<Vec<u8>, MuxError> {
let stream = self
.streams
.get_mut(&frame.stream_id)
.ok_or(MuxError::StreamNotFound(frame.stream_id.0))?;
if frame.is_rst() {
stream.state = StreamState::Reset;
self.total_frames_received = self.total_frames_received.wrapping_add(1);
return Ok(Vec::new());
}
if !frame.is_syn() && frame.sequence != stream.recv_seq {
return Err(MuxError::SequenceError {
expected: stream.recv_seq,
got: frame.sequence,
});
}
stream.recv_seq = stream.recv_seq.wrapping_add(1);
let data_len = frame.data.len();
stream.bytes_received = stream.bytes_received.saturating_add(data_len as u64);
if frame.is_fin() {
stream.state = StreamState::HalfClosed;
}
self.total_frames_received = self.total_frames_received.wrapping_add(1);
Ok(frame.data)
}
pub fn drain_recv_buffer(&mut self, stream_id: StreamId) -> Result<Vec<Vec<u8>>, MuxError> {
let stream = self
.streams
.get_mut(&stream_id)
.ok_or(MuxError::StreamNotFound(stream_id.0))?;
let payloads: Vec<Vec<u8>> = stream.recv_buffer.drain(..).collect();
Ok(payloads)
}
pub fn dequeue_frame(&mut self) -> Option<StreamFrame> {
let pf = self.send_queue.pop()?;
self.total_frames_sent = self.total_frames_sent.wrapping_add(1);
Some(pf.frame)
}
pub fn dequeue_frames(&mut self, n: usize) -> Vec<StreamFrame> {
let mut result = Vec::with_capacity(n);
for _ in 0..n {
match self.dequeue_frame() {
Some(f) => result.push(f),
None => break,
}
}
result
}
pub fn drain_send_queue(&mut self) -> Vec<StreamFrame> {
let mut result = Vec::with_capacity(self.send_queue.len());
while let Some(pf) = self.send_queue.pop() {
self.total_frames_sent = self.total_frames_sent.wrapping_add(1);
result.push(pf.frame);
}
result
}
pub fn expire_idle(&mut self, current_ts: u64) -> Vec<MuxEvent> {
let timeout = self.config.idle_timeout_us;
let mut expired = Vec::new();
for (id, stream) in self.streams.iter_mut() {
if matches!(
stream.state,
StreamState::Open | StreamState::Opening | StreamState::HalfClosed
) && stream.last_activity.saturating_add(timeout) < current_ts
{
stream.state = StreamState::Closed;
expired.push(*id);
}
}
expired
.into_iter()
.map(MuxEvent::IdleStreamExpired)
.collect()
}
pub fn update_window(&mut self, stream_id: StreamId, increment: u32) -> bool {
if let Some(stream) = self.streams.get_mut(&stream_id) {
stream.send_window = stream.send_window.saturating_add(increment);
true
} else {
false
}
}
pub fn stream_state(&self, stream_id: StreamId) -> Option<&LogicalStream> {
self.streams.get(&stream_id)
}
pub fn stream_info(&self, stream_id: StreamId) -> Result<StreamInfo, MuxError> {
let s = self
.streams
.get(&stream_id)
.ok_or(MuxError::StreamNotFound(stream_id.0))?;
Ok(StreamInfo {
id: s.id,
state: s.state,
bytes_sent: s.bytes_sent,
bytes_received: s.bytes_received,
frames_sent: s.frames_sent,
frames_received: s.frames_received,
opened_at: s.opened_at,
last_activity: s.last_activity,
priority: s.priority_raw,
})
}
pub fn active_streams(&self) -> Vec<StreamId> {
self.streams
.values()
.filter(|s| matches!(s.state, StreamState::Open | StreamState::Opening))
.map(|s| s.id)
.collect()
}
pub fn open_stream_count(&self) -> usize {
self.streams
.values()
.filter(|s| s.state == StreamState::Open)
.count()
}
pub fn stats(&self) -> MuxStats {
let total_bytes_sent = self
.streams
.values()
.map(|s| s.bytes_sent)
.fold(0u64, |acc, b| acc.saturating_add(b));
MuxStats {
total_streams: self.streams.len(),
open_streams: self.open_stream_count(),
queued_frames: self.send_queue.len(),
total_frames_sent: self.total_frames_sent,
total_frames_received: self.total_frames_received,
total_bytes_sent,
}
}
pub fn multiplexer_stats(&self) -> MultiplexerStats {
let (total_bytes_sent, total_bytes_received) =
self.streams.values().fold((0u64, 0u64), |(s, r), stream| {
(
s.saturating_add(stream.bytes_sent),
r.saturating_add(stream.bytes_received),
)
});
let active = self
.streams
.values()
.filter(|s| matches!(s.state, StreamState::Open | StreamState::Opening))
.count();
MultiplexerStats {
active_streams: active,
total_streams_opened: self.total_streams_opened,
total_frames_sent: self.total_frames_sent,
total_frames_received: self.total_frames_received,
total_bytes_sent,
total_bytes_received,
dropped_frames: self.dropped_frames,
}
}
fn enqueue_control(
&mut self,
stream_id: StreamId,
flags: u8,
sequence: u64,
priority: StreamPriority,
now: u64,
) {
let weight = if self.config.enable_priority {
priority.weight()
} else {
0
};
let enqueue_seq = self.enqueue_seq;
self.enqueue_seq = self.enqueue_seq.wrapping_add(1);
self.send_queue.push(PrioritizedFrame {
priority_weight: weight,
sequence: enqueue_seq,
frame: StreamFrame {
stream_id,
sequence,
data: Vec::new(),
flags,
timestamp: now,
},
});
}
}
pub fn xorshift64(state: &mut u64) -> u64 {
let mut x = *state;
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
*state = x;
x
}
#[cfg(test)]
#[path = "stream_multiplexer_tests.rs"]
mod tests;