use s0_const::{ACK_BLOCKS, ACK_SCAN_BLOCKS, DEFAULT_PKT, MAX_BLOCK, MAX_CNT, TUNNEL_PKT};
use s0_wire::StreamDataHeader;
use crate::{ArgError, UdpError};
const MAX_SLOTS: usize = 65_536;
const MAX_POOL: usize = 1024 * 1024 * 1024;
const MIN_PKT: usize = 256;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Config {
pub slots: usize,
pub block: usize,
pub pkt: usize,
pub accept_zstd: bool,
}
impl Config {
pub fn new(slots: usize, block: usize, pkt: usize) -> Result<Self, UdpError> {
if slots == 0 || !slots.is_power_of_two() || slots > MAX_SLOTS {
return Err(ArgError::Slots.into());
}
if block == 0 || block > MAX_BLOCK {
return Err(ArgError::BlockSize.into());
}
if pkt < MIN_PKT || pkt > block {
return Err(ArgError::PktSize.into());
}
if block.div_ceil(pkt) > MAX_CNT {
return Err(ArgError::SegCnt.into());
}
slots
.checked_mul(block)
.filter(|&m| m <= MAX_POOL)
.ok_or(ArgError::PoolMem)?;
Ok(Self {
slots,
block,
pkt,
accept_zstd: false,
})
}
const fn pool_bytes(&self) -> usize {
self.slots * self.block
}
pub const fn with_zstd(mut self) -> Self {
self.accept_zstd = true;
self
}
const fn slots_mask(&self) -> usize {
self.slots - 1
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
enum State {
#[default]
Empty,
Filling,
Ready,
Done,
}
#[derive(Clone, Copy, Debug, Default)]
struct Slot {
blk: u32,
have: u128,
cnt: u8,
zstd: bool,
pkt: u16,
last: u32,
len: u32,
state: State,
}
impl Slot {
fn reset(&mut self, blk: u32, cnt: u8, zstd: bool, pkt: u16) {
*self = Self {
blk,
cnt,
zstd,
pkt,
state: State::Filling,
..Self::default()
};
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Outcome {
Stored,
BlockReady { blk: u32, len: usize },
Duplicate,
Stale,
NoCredit,
}
pub struct Receiver {
cfg: Config,
buf: Box<[u8]>,
meta: Box<[Slot]>,
done_lo: u32,
cur_pkt: usize,
}
impl Receiver {
pub fn new(cfg: Config) -> Result<Self, UdpError> {
let cfg = Config {
accept_zstd: cfg.accept_zstd,
..Config::new(cfg.slots, cfg.block, cfg.pkt)?
};
let buf = vec![0u8; cfg.pool_bytes()].into_boxed_slice();
let meta = vec![Slot::default(); cfg.slots].into_boxed_slice();
Ok(Self {
cfg,
buf,
meta,
done_lo: 0,
cur_pkt: cfg.pkt,
})
}
pub const fn done_lo(&self) -> u32 {
self.done_lo
}
pub const fn credit(&self) -> u32 {
self.done_lo.saturating_add(self.cfg.slots as u32)
}
pub fn on_data(&mut self, h: &StreamDataHeader, payload: &[u8]) -> Result<Outcome, UdpError> {
if h.is_zstd() && !self.cfg.accept_zstd {
return Err(ArgError::UnexpectedZstd.into());
}
if h.cnt == 0 || usize::from(h.cnt) > MAX_CNT || h.idx >= h.cnt {
return Err(ArgError::BadDataHdr.into());
}
let blk = h.blk;
if blk < self.done_lo {
return Ok(Outcome::Stale);
}
if blk >= self.credit() {
return Ok(Outcome::NoCredit);
}
if payload.is_empty() {
return Err(UdpError::Corrupt);
}
let tail = h.idx + 1 == h.cnt;
let allow_tunnel = self.cfg.pkt == DEFAULT_PKT;
let max_default_cnt = self.cfg.block.div_ceil(self.cfg.pkt);
let i = blk as usize & self.cfg.slots_mask();
let start = i * self.cfg.block;
let pkt_step = if !tail {
if payload.len() != self.cfg.pkt && !(allow_tunnel && payload.len() == TUNNEL_PKT) {
return Err(UdpError::Corrupt);
}
self.cur_pkt = payload.len();
payload.len()
} else if allow_tunnel && usize::from(h.cnt) > max_default_cnt {
self.cur_pkt = TUNNEL_PKT;
TUNNEL_PKT
} else {
let existing = &self.meta[i];
if matches!(existing.state, State::Filling | State::Ready) && existing.blk == blk {
if existing.cnt == h.cnt && existing.pkt > 0 {
usize::from(existing.pkt)
} else if allow_tunnel && h.cnt > existing.cnt {
self.cur_pkt = TUNNEL_PKT;
TUNNEL_PKT
} else {
self.cur_pkt
}
} else {
self.cur_pkt
}
};
if payload.len() > pkt_step.max(if h.cnt == 1 { self.cfg.pkt } else { pkt_step }) {
return Err(UdpError::Corrupt);
}
let slot = &mut self.meta[i];
match slot.state {
State::Filling | State::Ready => {
if slot.blk != blk || slot.zstd != h.is_zstd() {
return Err(UdpError::Corrupt);
}
if slot.cnt != h.cnt {
if allow_tunnel && (pkt_step == TUNNEL_PKT || usize::from(h.cnt) > max_default_cnt) {
slot.reset(blk, h.cnt, h.is_zstd(), TUNNEL_PKT as u16);
} else if allow_tunnel && slot.pkt == TUNNEL_PKT as u16 && h.cnt < slot.cnt {
return Ok(Outcome::Duplicate);
} else {
return Err(UdpError::Corrupt);
}
} else if !tail && slot.pkt != pkt_step as u16 {
let tail_bit = 1u128 << (h.cnt - 1);
if allow_tunnel && slot.have == tail_bit {
let old_off = start + usize::from(h.cnt - 1) * usize::from(slot.pkt);
let new_off = start + usize::from(h.cnt - 1) * pkt_step;
let last_len = slot.last as usize;
if new_off + last_len > start + self.cfg.block {
return Err(UdpError::Corrupt);
}
self.buf.copy_within(old_off..old_off + last_len, new_off);
slot.pkt = pkt_step as u16;
} else if allow_tunnel && slot.pkt == TUNNEL_PKT as u16 && pkt_step > TUNNEL_PKT {
return Ok(Outcome::Duplicate);
} else {
return Err(UdpError::Corrupt);
}
}
if slot.have >> h.idx & 1 == 1 {
return Ok(Outcome::Duplicate);
}
}
State::Done if slot.blk == blk => {
return Ok(Outcome::Duplicate);
}
_ => slot.reset(blk, h.cnt, h.is_zstd(), pkt_step as u16),
}
let eff_pkt = usize::from(slot.pkt);
let idx = usize::from(h.idx);
let pkt_off = idx * eff_pkt;
let end_off = pkt_off + payload.len();
if end_off > self.cfg.block {
return Err(UdpError::Corrupt);
}
if tail {
slot.last = payload.len() as u32;
}
let base = start + pkt_off;
self.buf[base..base + payload.len()].copy_from_slice(payload);
slot.have |= 1u128 << h.idx;
if slot.have.count_ones() != u32::from(h.cnt) {
return Ok(Outcome::Stored);
}
let len = (usize::from(h.cnt) - 1) * eff_pkt + slot.last as usize;
slot.len = len as u32;
slot.state = State::Ready;
Ok(Outcome::BlockReady { blk, len })
}
pub fn data(&self, blk: u32) -> Option<&[u8]> {
let i = blk as usize & self.cfg.slots_mask();
let slot = self.meta[i];
if slot.blk != blk || slot.state != State::Ready {
return None;
}
let base = i * self.cfg.block;
Some(&self.buf[base..base + slot.len as usize])
}
pub fn partial(&self, blk: u32) -> Option<u128> {
let slot = self.meta[blk as usize & self.cfg.slots_mask()];
(slot.blk == blk && slot.state == State::Filling).then_some(slot.have)
}
pub fn collect_ack_blocks(&self) -> ([(u32, u128); ACK_BLOCKS], usize) {
let mut blocks = [(0u32, 0u128); ACK_BLOCKS];
let mut nb = 0usize;
let lo = self.done_lo;
let hi = lo.saturating_add(ACK_SCAN_BLOCKS).min(self.credit());
for (dst, item) in blocks
.iter_mut()
.zip((lo..hi).filter_map(|b| self.partial(b).map(|bits| (b, bits))))
{
*dst = item;
nb += 1;
}
(blocks, nb)
}
pub fn release(&mut self, blk: u32) -> Result<(), UdpError> {
if blk < self.done_lo || blk >= self.credit() {
return Err(UdpError::Corrupt);
}
let mask = self.cfg.slots_mask();
let slot = &mut self.meta[blk as usize & mask];
if slot.blk != blk || slot.state != State::Ready {
return Err(UdpError::Corrupt);
}
slot.state = State::Done;
while self.done_lo < u32::MAX {
let s = self.meta[self.done_lo as usize & mask];
if s.state != State::Done || s.blk != self.done_lo {
break;
}
self.done_lo += 1;
}
Ok(())
}
}