109 lines
3.0 KiB
Rust
109 lines
3.0 KiB
Rust
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub struct ReassemblyState {
|
|
pub active: bool,
|
|
pub batch_id: u16,
|
|
pub next_index: u8,
|
|
pub expected_chunks: u8,
|
|
pub total_len: u16,
|
|
pub received_len: u16,
|
|
pub last_rx_ms: u32,
|
|
pub timeout_ms: u32,
|
|
}
|
|
|
|
impl Default for ReassemblyState {
|
|
fn default() -> Self {
|
|
Self {
|
|
active: false,
|
|
batch_id: 0,
|
|
next_index: 0,
|
|
expected_chunks: 0,
|
|
total_len: 0,
|
|
received_len: 0,
|
|
last_rx_ms: 0,
|
|
timeout_ms: 0,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum ReassemblyStatus {
|
|
InProgress,
|
|
Complete { complete_len: u16 },
|
|
ErrorReset,
|
|
}
|
|
|
|
pub fn reset_reassembly(state: &mut ReassemblyState) {
|
|
*state = ReassemblyState::default();
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn push_chunk(
|
|
state: &mut ReassemblyState,
|
|
batch_id: u16,
|
|
chunk_index: u8,
|
|
chunk_count: u8,
|
|
total_len: u16,
|
|
chunk_data: &[u8],
|
|
now_ms: u32,
|
|
timeout_ms_for_new_batch: u32,
|
|
max_total_len: u16,
|
|
buffer: &mut [u8],
|
|
) -> ReassemblyStatus {
|
|
if chunk_data.len() > 0 && total_len == 0 {
|
|
reset_reassembly(state);
|
|
return ReassemblyStatus::ErrorReset;
|
|
}
|
|
|
|
let expired = state.timeout_ms > 0
|
|
&& now_ms.wrapping_sub(state.last_rx_ms) > state.timeout_ms;
|
|
|
|
if !state.active || batch_id != state.batch_id || expired {
|
|
if chunk_index != 0 {
|
|
reset_reassembly(state);
|
|
return ReassemblyStatus::ErrorReset;
|
|
}
|
|
if total_len == 0 || total_len > max_total_len || chunk_count == 0 {
|
|
reset_reassembly(state);
|
|
return ReassemblyStatus::ErrorReset;
|
|
}
|
|
|
|
state.active = true;
|
|
state.batch_id = batch_id;
|
|
state.expected_chunks = chunk_count;
|
|
state.total_len = total_len;
|
|
state.received_len = 0;
|
|
state.next_index = 0;
|
|
state.last_rx_ms = now_ms;
|
|
state.timeout_ms = timeout_ms_for_new_batch;
|
|
}
|
|
|
|
if !state.active || chunk_index != state.next_index || chunk_count != state.expected_chunks {
|
|
reset_reassembly(state);
|
|
return ReassemblyStatus::ErrorReset;
|
|
}
|
|
|
|
let next_received = state.received_len as usize + chunk_data.len();
|
|
if next_received > state.total_len as usize
|
|
|| next_received > max_total_len as usize
|
|
|| next_received > buffer.len()
|
|
{
|
|
reset_reassembly(state);
|
|
return ReassemblyStatus::ErrorReset;
|
|
}
|
|
|
|
let start = state.received_len as usize;
|
|
let end = start + chunk_data.len();
|
|
buffer[start..end].copy_from_slice(chunk_data);
|
|
|
|
state.received_len = next_received as u16;
|
|
state.next_index = state.next_index.wrapping_add(1);
|
|
state.last_rx_ms = now_ms;
|
|
|
|
if state.next_index == state.expected_chunks && state.received_len == state.total_len {
|
|
let complete_len = state.received_len;
|
|
reset_reassembly(state);
|
|
return ReassemblyStatus::Complete { complete_len };
|
|
}
|
|
|
|
ReassemblyStatus::InProgress
|
|
} |