411 lines
14 KiB
Rust
411 lines
14 KiB
Rust
// WIP mesh/transport protocol — suppress dead code warnings
|
|
#![allow(dead_code)]
|
|
//! Chunked message protocol with Reed-Solomon FEC for LoRa transport.
|
|
//!
|
|
//! Splits payloads larger than a single LoRa frame (160 bytes) into
|
|
//! numbered chunks with forward error correction, enabling reliable
|
|
//! transfer over lossy radio links.
|
|
//!
|
|
//! Chunk wire format (8 bytes header + payload):
|
|
//! ```text
|
|
//! [0x01: type] [msg_id: u32 LE] [chunk_idx: u8] [total: u8] [is_parity: u8] [payload...]
|
|
//! ```
|
|
|
|
use anyhow::{Context, Result};
|
|
use reed_solomon_erasure::galois_8::ReedSolomon;
|
|
use std::collections::HashMap;
|
|
use std::time::Instant;
|
|
|
|
/// Header size for each chunk frame.
|
|
const CHUNK_HEADER_SIZE: usize = 8;
|
|
|
|
/// Maximum payload per chunk after header.
|
|
/// 132 bytes available after ChaCha20-Poly1305 encryption overhead (12 nonce + 16 tag),
|
|
/// minus 8 byte chunk header = 124 bytes of user data per chunk.
|
|
pub const MAX_CHUNK_PAYLOAD: usize = 124;
|
|
|
|
/// Chunk type marker in the wire format.
|
|
const CHUNK_TYPE_MARKER: u8 = 0x01;
|
|
|
|
/// FEC redundancy ratio: 25% parity shards.
|
|
const FEC_RATIO_DENOMINATOR: usize = 4;
|
|
|
|
/// Maximum age of pending reassembly entries before garbage collection.
|
|
const REASSEMBLY_TIMEOUT_SECS: u64 = 60;
|
|
|
|
/// Maximum practical chunks for LoRa (airtime budget).
|
|
pub const MAX_PRACTICAL_CHUNKS: usize = 20;
|
|
|
|
/// A single chunk ready for transmission.
|
|
#[derive(Debug, Clone)]
|
|
pub struct Chunk {
|
|
pub message_id: u32,
|
|
pub chunk_index: u8,
|
|
pub total_chunks: u8,
|
|
pub is_parity: bool,
|
|
pub payload: Vec<u8>,
|
|
}
|
|
|
|
impl Chunk {
|
|
/// Serialize chunk to wire format.
|
|
pub fn to_bytes(&self) -> Vec<u8> {
|
|
let mut buf = Vec::with_capacity(CHUNK_HEADER_SIZE + self.payload.len());
|
|
buf.push(CHUNK_TYPE_MARKER);
|
|
buf.extend_from_slice(&self.message_id.to_le_bytes());
|
|
buf.push(self.chunk_index);
|
|
buf.push(self.total_chunks);
|
|
buf.push(if self.is_parity { 1 } else { 0 });
|
|
buf.extend_from_slice(&self.payload);
|
|
buf
|
|
}
|
|
|
|
/// Parse chunk from wire format.
|
|
pub fn from_bytes(data: &[u8]) -> Result<Self> {
|
|
if data.len() < CHUNK_HEADER_SIZE {
|
|
anyhow::bail!("Chunk too small: {} bytes", data.len());
|
|
}
|
|
if data[0] != CHUNK_TYPE_MARKER {
|
|
anyhow::bail!("Not a chunked message (marker: 0x{:02x})", data[0]);
|
|
}
|
|
let message_id = u32::from_le_bytes([data[1], data[2], data[3], data[4]]);
|
|
let chunk_index = data[5];
|
|
let total_chunks = data[6];
|
|
let is_parity = data[7] != 0;
|
|
let payload = data[CHUNK_HEADER_SIZE..].to_vec();
|
|
|
|
Ok(Self {
|
|
message_id,
|
|
chunk_index,
|
|
total_chunks,
|
|
is_parity,
|
|
payload,
|
|
})
|
|
}
|
|
|
|
/// Check if a raw byte slice starts with the chunk type marker.
|
|
pub fn is_chunked_message(data: &[u8]) -> bool {
|
|
!data.is_empty() && data[0] == CHUNK_TYPE_MARKER
|
|
}
|
|
}
|
|
|
|
/// Encode a payload into chunks with Reed-Solomon FEC parity.
|
|
///
|
|
/// Returns a vector of chunks ready for sequential transmission.
|
|
/// Each chunk's payload is exactly `shard_size` bytes (padded if needed).
|
|
pub fn encode_chunked(data: &[u8]) -> Result<Vec<Chunk>> {
|
|
if data.is_empty() {
|
|
anyhow::bail!("Cannot chunk empty data");
|
|
}
|
|
|
|
let shard_size = MAX_CHUNK_PAYLOAD;
|
|
let data_shard_count = (data.len() + shard_size - 1) / shard_size;
|
|
|
|
if data_shard_count > MAX_PRACTICAL_CHUNKS {
|
|
anyhow::bail!(
|
|
"Payload too large for LoRa chunking: {} bytes ({} chunks, max {})",
|
|
data.len(),
|
|
data_shard_count,
|
|
MAX_PRACTICAL_CHUNKS
|
|
);
|
|
}
|
|
|
|
let parity_shard_count = (data_shard_count + FEC_RATIO_DENOMINATOR - 1) / FEC_RATIO_DENOMINATOR;
|
|
let total_shards = data_shard_count + parity_shard_count;
|
|
|
|
if total_shards > 255 {
|
|
anyhow::bail!("Too many shards: {}", total_shards);
|
|
}
|
|
|
|
// Split data into equal-size shards
|
|
let mut shards: Vec<Vec<u8>> = Vec::with_capacity(total_shards);
|
|
for i in 0..data_shard_count {
|
|
let start = i * shard_size;
|
|
let end = (start + shard_size).min(data.len());
|
|
let mut shard = vec![0u8; shard_size];
|
|
shard[..end - start].copy_from_slice(&data[start..end]);
|
|
shards.push(shard);
|
|
}
|
|
|
|
// Add empty parity shards
|
|
for _ in 0..parity_shard_count {
|
|
shards.push(vec![0u8; shard_size]);
|
|
}
|
|
|
|
// Generate parity
|
|
let rs = ReedSolomon::new(data_shard_count, parity_shard_count)
|
|
.context("Failed to create Reed-Solomon codec")?;
|
|
rs.encode(&mut shards)
|
|
.context("Reed-Solomon encoding failed")?;
|
|
|
|
// Build chunk frames
|
|
let message_id: u32 = rand::random();
|
|
let total = total_shards as u8;
|
|
let mut chunks = Vec::with_capacity(total_shards);
|
|
|
|
for (i, shard) in shards.into_iter().enumerate() {
|
|
chunks.push(Chunk {
|
|
message_id,
|
|
chunk_index: i as u8,
|
|
total_chunks: total,
|
|
is_parity: i >= data_shard_count,
|
|
payload: shard,
|
|
});
|
|
}
|
|
|
|
// Encode the original data length in the first chunk's first 4 bytes
|
|
// so the receiver can trim padding after reconstruction.
|
|
let data_len = data.len() as u32;
|
|
chunks[0].payload[..4].copy_from_slice(&data_len.to_le_bytes());
|
|
// Re-encode FEC to reflect the length header change
|
|
let mut shard_data: Vec<Vec<u8>> = chunks.iter().map(|c| c.payload.clone()).collect();
|
|
rs.encode(&mut shard_data)
|
|
.context("Reed-Solomon re-encoding failed")?;
|
|
for (i, shard) in shard_data.into_iter().enumerate() {
|
|
chunks[i].payload = shard;
|
|
}
|
|
|
|
Ok(chunks)
|
|
}
|
|
|
|
/// In-progress reassembly of a chunked message.
|
|
struct PendingMessage {
|
|
shards: Vec<Option<Vec<u8>>>,
|
|
data_shard_count: usize,
|
|
parity_shard_count: usize,
|
|
received_count: usize,
|
|
created_at: Instant,
|
|
}
|
|
|
|
/// Reassembles chunked messages from incoming chunks.
|
|
pub struct ChunkReassembler {
|
|
pending: HashMap<u32, PendingMessage>,
|
|
}
|
|
|
|
impl ChunkReassembler {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
pending: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
/// Feed a chunk into the reassembler.
|
|
/// Returns `Some(data)` if the message is fully reconstructed.
|
|
pub fn feed(&mut self, chunk: &Chunk) -> Result<Option<Vec<u8>>> {
|
|
// Garbage collect stale entries
|
|
self.pending.retain(|_, pm| {
|
|
pm.created_at.elapsed().as_secs() < REASSEMBLY_TIMEOUT_SECS
|
|
});
|
|
|
|
let total = chunk.total_chunks as usize;
|
|
let entry = self.pending.entry(chunk.message_id).or_insert_with(|| {
|
|
// Infer data vs parity count from chunks we've seen
|
|
// The first non-parity chunk tells us the split point
|
|
let data_count = if chunk.is_parity {
|
|
// Best guess: 80% data, 20% parity
|
|
(total * FEC_RATIO_DENOMINATOR) / (FEC_RATIO_DENOMINATOR + 1)
|
|
} else {
|
|
// We know this index is data — parity starts after all data
|
|
// Exact split point: smallest i where chunk_index >= data_count AND is_parity
|
|
total - (total + FEC_RATIO_DENOMINATOR) / (FEC_RATIO_DENOMINATOR + 1)
|
|
};
|
|
let parity_count = total - data_count;
|
|
|
|
PendingMessage {
|
|
shards: vec![None; total],
|
|
data_shard_count: data_count,
|
|
parity_shard_count: parity_count,
|
|
received_count: 0,
|
|
created_at: Instant::now(),
|
|
}
|
|
});
|
|
|
|
let idx = chunk.chunk_index as usize;
|
|
if idx >= total {
|
|
anyhow::bail!("Chunk index {} out of range (total {})", idx, total);
|
|
}
|
|
|
|
if entry.shards[idx].is_none() {
|
|
entry.shards[idx] = Some(chunk.payload.clone());
|
|
entry.received_count += 1;
|
|
}
|
|
|
|
// Need at least data_shard_count shards to reconstruct
|
|
if entry.received_count >= entry.data_shard_count {
|
|
self.try_reconstruct(chunk.message_id)
|
|
} else {
|
|
Ok(None)
|
|
}
|
|
}
|
|
|
|
fn try_reconstruct(&mut self, message_id: u32) -> Result<Option<Vec<u8>>> {
|
|
let entry = match self.pending.get_mut(&message_id) {
|
|
Some(e) => e,
|
|
None => return Ok(None),
|
|
};
|
|
|
|
let rs = ReedSolomon::new(entry.data_shard_count, entry.parity_shard_count)
|
|
.context("Failed to create Reed-Solomon codec for reconstruction")?;
|
|
|
|
let mut shards: Vec<Option<Vec<u8>>> = entry.shards.clone();
|
|
|
|
match rs.reconstruct(&mut shards) {
|
|
Ok(()) => {
|
|
// Concatenate data shards (not parity)
|
|
let mut result = Vec::new();
|
|
for shard in shards.iter().take(entry.data_shard_count) {
|
|
if let Some(data) = shard {
|
|
result.extend_from_slice(data);
|
|
}
|
|
}
|
|
|
|
// Extract original length from first 4 bytes
|
|
if result.len() < 4 {
|
|
anyhow::bail!("Reconstructed data too small for length header");
|
|
}
|
|
let original_len =
|
|
u32::from_le_bytes([result[0], result[1], result[2], result[3]]) as usize;
|
|
|
|
// The actual data starts at byte 4 of the first shard
|
|
// But wait — the length is embedded in shard 0 bytes 0..4, and the
|
|
// actual payload starts at byte 4 of shard 0, then continues in subsequent shards.
|
|
// Actually, encode_chunked puts the length in the first 4 bytes of shard 0,
|
|
// and the rest of shard 0 + all other shards contain the original data.
|
|
// So we need to skip 4 bytes from the beginning.
|
|
if 4 + original_len > result.len() {
|
|
anyhow::bail!(
|
|
"Original length {} exceeds reconstructed data ({})",
|
|
original_len,
|
|
result.len() - 4
|
|
);
|
|
}
|
|
|
|
let data = result[4..4 + original_len].to_vec();
|
|
self.pending.remove(&message_id);
|
|
Ok(Some(data))
|
|
}
|
|
Err(_) => {
|
|
// Not enough shards yet
|
|
Ok(None)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Default for ChunkReassembler {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_chunk_roundtrip_small() {
|
|
// Small payload fits in 1 data chunk + 1 parity chunk
|
|
let data = b"Hello, mesh network!";
|
|
let chunks = encode_chunked(data).unwrap();
|
|
|
|
// 1 data + 1 parity = 2 chunks
|
|
assert_eq!(chunks.len(), 2);
|
|
assert!(!chunks[0].is_parity);
|
|
assert!(chunks[1].is_parity);
|
|
|
|
let mut reassembler = ChunkReassembler::new();
|
|
// Feed data chunk — should reconstruct immediately (1 data shard needed)
|
|
let result = reassembler.feed(&chunks[0]).unwrap();
|
|
assert!(result.is_some());
|
|
assert_eq!(result.unwrap(), data);
|
|
}
|
|
|
|
#[test]
|
|
fn test_chunk_roundtrip_medium() {
|
|
// ~500 bytes: 4 data chunks + 1 parity
|
|
let data: Vec<u8> = (0..500).map(|i| (i % 256) as u8).collect();
|
|
let chunks = encode_chunked(&data).unwrap();
|
|
|
|
let data_chunks: Vec<_> = chunks.iter().filter(|c| !c.is_parity).collect();
|
|
let parity_chunks: Vec<_> = chunks.iter().filter(|c| c.is_parity).collect();
|
|
assert_eq!(data_chunks.len(), 4); // ceil(500/124) = 5... wait
|
|
// Actually: ceil(500/124) = ceil(4.03) = 5 data shards
|
|
// But the first shard has 4 bytes of length header embedded, so
|
|
// the actual data capacity is 124 * N - 0 (length is IN the shard data).
|
|
// Let's just check it roundtrips.
|
|
|
|
let mut reassembler = ChunkReassembler::new();
|
|
let mut result = None;
|
|
for chunk in &chunks {
|
|
if let Some(data) = reassembler.feed(chunk).unwrap() {
|
|
result = Some(data);
|
|
break;
|
|
}
|
|
}
|
|
assert!(result.is_some());
|
|
assert_eq!(result.unwrap(), data);
|
|
}
|
|
|
|
#[test]
|
|
fn test_chunk_wire_format() {
|
|
let chunk = Chunk {
|
|
message_id: 0x12345678,
|
|
chunk_index: 2,
|
|
total_chunks: 5,
|
|
is_parity: false,
|
|
payload: vec![0xAA, 0xBB],
|
|
};
|
|
let bytes = chunk.to_bytes();
|
|
assert_eq!(bytes[0], CHUNK_TYPE_MARKER);
|
|
assert_eq!(&bytes[1..5], &0x12345678u32.to_le_bytes());
|
|
assert_eq!(bytes[5], 2);
|
|
assert_eq!(bytes[6], 5);
|
|
assert_eq!(bytes[7], 0);
|
|
assert_eq!(&bytes[8..], &[0xAA, 0xBB]);
|
|
|
|
let parsed = Chunk::from_bytes(&bytes).unwrap();
|
|
assert_eq!(parsed.message_id, 0x12345678);
|
|
assert_eq!(parsed.chunk_index, 2);
|
|
assert_eq!(parsed.total_chunks, 5);
|
|
assert!(!parsed.is_parity);
|
|
assert_eq!(parsed.payload, vec![0xAA, 0xBB]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_chunk_is_chunked_message() {
|
|
assert!(Chunk::is_chunked_message(&[0x01, 0x00]));
|
|
assert!(!Chunk::is_chunked_message(&[0x02, 0x00]));
|
|
assert!(!Chunk::is_chunked_message(&[]));
|
|
}
|
|
|
|
#[test]
|
|
fn test_chunk_with_missing_chunk() {
|
|
// Verify FEC can recover from a missing data chunk
|
|
let data: Vec<u8> = (0..300).map(|i| (i % 256) as u8).collect();
|
|
let chunks = encode_chunked(&data).unwrap();
|
|
|
|
let mut reassembler = ChunkReassembler::new();
|
|
// Skip chunk index 1 (simulate loss)
|
|
for chunk in &chunks {
|
|
if chunk.chunk_index == 1 {
|
|
continue;
|
|
}
|
|
if let Some(recovered) = reassembler.feed(chunk).unwrap() {
|
|
assert_eq!(recovered, data);
|
|
return;
|
|
}
|
|
}
|
|
panic!("Failed to reconstruct with one missing chunk");
|
|
}
|
|
|
|
#[test]
|
|
fn test_empty_data_rejected() {
|
|
assert!(encode_chunked(&[]).is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_too_large_rejected() {
|
|
let data = vec![0u8; MAX_CHUNK_PAYLOAD * (MAX_PRACTICAL_CHUNKS + 1)];
|
|
assert!(encode_chunked(&data).is_err());
|
|
}
|
|
}
|