Refactor chunk_codec

.github/workflows/rust.yml

@@ -8,7 +8,7 @@ on:
 
 jobs:
   clippy:
-    name: Format & Clippy
+    name: Cargo Format & Clippy & Test
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
@@ -31,6 +31,7 @@ jobs:
             pushd $t
             cargo fmt --check
             cargo clippy -- -D warnings
+            cargo test --no-fail-fast --all-features
             popd
           done
 

pixie-server/src/ping.rs

@@ -5,7 +5,7 @@ use crate::{
     state::{State, UnitSelector},
 };
 use anyhow::{bail, Result};
-use pixie_shared::{PACKET_LEN, PING_PORT};
+use pixie_shared::{PING_PORT, UDP_BODY_LEN};
 use std::{
     net::{IpAddr, Ipv4Addr},
     sync::Arc,
@@ -17,7 +17,7 @@ pub async fn main(state: Arc<State>) -> Result<()> {
     let socket = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, PING_PORT)).await?;
     log::info!("Listening on {}", socket.local_addr()?);
 
-    let mut buf = [0; PACKET_LEN];
+    let mut buf = [0; UDP_BODY_LEN];
     loop {
         let (len, peer_addr) = tokio::select! {
             x = socket.recv_from(&mut buf) => x?,

pixie-server/src/udp.rs

@@ -6,8 +6,8 @@ use crate::{
 };
 use anyhow::{bail, ensure, Context, Result};
 use pixie_shared::{
-    ChunkHash, HintPacket, RegistrationInfo, UdpRequest, ACTION_PORT, BODY_LEN, CHUNKS_PORT,
-    HINT_PORT, PACKET_LEN,
+    chunk_codec::Encoder, ChunkHash, HintPacket, RegistrationInfo, UdpRequest, ACTION_PORT,
+    CHUNKS_PORT, HINT_PORT, UDP_BODY_LEN,
 };
 use std::{
     collections::BTreeSet,
@@ -28,7 +28,7 @@ async fn broadcast_chunks(
     mut rx: Receiver<ChunkHash>,
 ) -> Result<()> {
     let mut queue = BTreeSet::<ChunkHash>::new();
-    let mut write_buf = [0; PACKET_LEN];
+    let mut write_buf = [0; UDP_BODY_LEN];
     let mut wait_for = Instant::now();
     let mut index = [0; 32];
 
@@ -67,48 +67,24 @@ async fn broadcast_chunks(
             _ = state.cancel_token.cancelled() => break,
         };
 
-        let Some(data) = state
+        let Some(cdata) = state
             .get_chunk_cdata(index)
             .with_context(|| format!("get chunk {}", hex::encode(index)))?
         else {
             log::warn!("Chunk {} not found", hex::encode(index));
             continue;
         };
 
-        let num_packets = data.len().div_ceil(BODY_LEN);
-        write_buf[..32].clone_from_slice(&index);
-
-        let mut xor = [[0; BODY_LEN]; 32];
-
         let hosts_cfg = &state.config.hosts;
         let chunks_addr = SocketAddrV4::new(ip, CHUNKS_PORT);
 
-        for index in 0..num_packets {
-            write_buf[32..34].clone_from_slice(&(index as u16).to_le_bytes());
-            let start = index * BODY_LEN;
-            let len = BODY_LEN.min(data.len() - start);
-            let body = &data[start..start + len];
-            let group = index & 31;
-            body.iter()
-                .zip(xor[group].iter_mut())
-                .for_each(|(a, b)| *b ^= a);
-            write_buf[34..34 + len].clone_from_slice(body);
-
+        let mut encoder = Encoder::new(cdata);
+        write_buf[..32].clone_from_slice(&index);
+        while let Some(len) = encoder.next_packet(&mut write_buf[32..]) {
             time::sleep_until(wait_for).await;
 
-            let sent_len = socket.send_to(&write_buf[..34 + len], chunks_addr).await?;
-            ensure!(sent_len == 34 + len, "Could not send packet");
-            wait_for += 8 * (sent_len as u32) * Duration::from_secs(1) / hosts_cfg.broadcast_speed;
-        }
-
-        for (index, body) in xor.iter().enumerate().take(num_packets) {
-            write_buf[32..34].clone_from_slice(&(index as u16).wrapping_sub(32).to_le_bytes());
-            let len = BODY_LEN;
-            write_buf[34..34 + len].clone_from_slice(body);
-
-            time::sleep_until(wait_for).await;
-            let sent_len = socket.send_to(&write_buf[..34 + len], chunks_addr).await?;
-            ensure!(sent_len == 34 + len, "Could not send packet");
+            let sent_len = socket.send_to(&write_buf[..32 + len], chunks_addr).await?;
+            ensure!(sent_len == 32 + len, "Could not send packet");
             wait_for += 8 * (sent_len as u32) * Duration::from_secs(1) / hosts_cfg.broadcast_speed;
         }
     }
@@ -190,7 +166,7 @@ async fn broadcast_hint(state: &State, socket: &UdpSocket, ip: Ipv4Addr) -> Resu
 }
 
 async fn handle_requests(state: &State, socket: &UdpSocket, tx: Sender<[u8; 32]>) -> Result<()> {
-    let mut buf = [0; PACKET_LEN];
+    let mut buf = [0; UDP_BODY_LEN];
     loop {
         let (len, addr) = tokio::select! {
             x = socket.recv_from(&mut buf) => x?,

pixie-shared/Cargo.toml

@@ -8,6 +8,7 @@ blake3 = { version = "1", default-features = false }
 ipnet = { version = "2.9.0", optional = true, features = ["serde"] }
 macaddr = { version = "1.0.1", optional = true, features = ["serde"] }
 serde =  { version = "1.0.193", default-features = false, features = ["derive", "alloc"] }
+thiserror = { version = "2.0.12", default-features = false }
 
 [features]
 std = ["ipnet", "macaddr", "serde/std"]

pixie-shared/src/chunk_codec.rs

@@ -0,0 +1,227 @@
+use crate::UDP_BODY_LEN;
+use alloc::{vec, vec::Vec};
+use thiserror::Error;
+
+const PACKET_LEN: usize = UDP_BODY_LEN - 32;
+const HEADER_LEN: usize = 2;
+const BODY_LEN: usize = PACKET_LEN - HEADER_LEN;
+
+const MIN_SIZE: usize = HEADER_LEN;
+const MAX_SIZE: usize = PACKET_LEN;
+
+#[derive(Error, Debug)]
+pub enum DecoderError {
+    #[error("Packet too small; got {0} bytes, expected at least {MIN_SIZE} bytes")]
+    PacketTooSmall(usize),
+    #[error("Packet too big; got {0} bytes, expected at most {MAX_SIZE} bytes")]
+    PacketTooBig(usize),
+    #[error("Invalid index: 0x{0:04x}")]
+    InvalidIndex(u16),
+}
+
+pub struct Decoder {
+    data: Vec<u8>,
+    missing_packet: Vec<bool>,
+    missing_packets_per_group: [u16; 32],
+    missing_groups: u16,
+}
+
+impl Decoder {
+    pub fn new(size: usize) -> Self {
+        let num_packets = size.div_ceil(BODY_LEN);
+        let data = vec![0; 32 * BODY_LEN + size];
+        let missing_packet = vec![true; 32 + num_packets];
+        let missing_packets_per_group: [u16; 32] = (0..32)
+            .map(|i| ((num_packets + 31 - i) / 32) as u16)
+            .collect::<Vec<_>>()
+            .try_into()
+            .unwrap();
+        let missing_groups = missing_packets_per_group
+            .iter()
+            .map(|&x| (x != 0) as u16)
+            .sum();
+        Decoder {
+            data,
+            missing_packet,
+            missing_packets_per_group,
+            missing_groups,
+        }
+    }
+
+    pub fn add_packet(&mut self, buf: &[u8]) -> Result<(), DecoderError> {
+        if buf.len() < MIN_SIZE {
+            return Err(DecoderError::PacketTooSmall(buf.len()));
+        }
+        if buf.len() > MAX_SIZE {
+            return Err(DecoderError::PacketTooBig(buf.len()));
+        }
+
+        let index = u16::from_le_bytes(buf[..2].try_into().unwrap());
+
+        let rot_index = index.wrapping_add(32) as usize;
+        let missing = self
+            .missing_packet
+            .get_mut(rot_index)
+            .ok_or(DecoderError::InvalidIndex(index))?;
+        match missing {
+            false => return Ok(()),
+            x @ true => *x = false,
+        }
+
+        let start = rot_index * BODY_LEN;
+        self.data[start..start + buf.len() - 2].clone_from_slice(&buf[2..]);
+
+        let group = index & 31;
+        match &mut self.missing_packets_per_group[group as usize] {
+            0 => return Ok(()),
+            x @ 1 => *x = 0,
+            x @ 2.. => {
+                *x -= 1;
+                return Ok(());
+            }
+        }
+
+        match &mut self.missing_groups {
+            0 => unreachable!(),
+            x @ 1.. => *x -= 1,
+        }
+
+        Ok(())
+    }
+
+    pub fn finish(&mut self) -> Option<Vec<u8>> {
+        if self.missing_groups != 0 {
+            return None;
+        }
+
+        let mut xor = [[0; BODY_LEN]; 32];
+        for packet in 0..self.missing_packet.len() {
+            if !self.missing_packet[packet] {
+                let group = packet & 31;
+                self.data[BODY_LEN * packet..]
+                    .iter()
+                    .zip(xor[group].iter_mut())
+                    .for_each(|(a, b)| *b ^= a);
+            }
+        }
+        for packet in 0..self.missing_packet.len() {
+            if self.missing_packet[packet] {
+                let group = packet & 31;
+                self.data[BODY_LEN * packet..]
+                    .iter_mut()
+                    .zip(xor[group].iter())
+                    .for_each(|(a, b)| *a = *b);
+            }
+        }
+        Some(self.data[32 * BODY_LEN..].to_vec())
+    }
+}
+
+pub struct Encoder {
+    data: Vec<u8>,
+    groups: usize,
+    idx: usize,
+}
+
+impl Encoder {
+    pub fn new(data: Vec<u8>) -> Self {
+        let idx = 0;
+        let groups = data.len().div_ceil(BODY_LEN).min(32);
+        Encoder { data, groups, idx }
+    }
+
+    pub fn next_packet(&mut self, out_buf: &mut [u8]) -> Option<usize> {
+        let start = self.idx * BODY_LEN;
+        if start < self.data.len() {
+            let end = self.data.len().min(start + BODY_LEN);
+            let len = end - start;
+            out_buf[0..2].copy_from_slice(&(self.idx as u16).to_le_bytes());
+            out_buf[2..2 + len].copy_from_slice(&self.data[start..end]);
+            self.idx += 1;
+            Some(2 + len)
+        } else if self.groups > 0 {
+            self.groups -= 1;
+            out_buf[0..2].copy_from_slice(&(self.groups as u16).wrapping_sub(32).to_le_bytes());
+            out_buf[2..].fill(0);
+            (0..)
+                .map(|x| (x * 32 + self.groups) * BODY_LEN)
+                .take_while(|x| *x < self.data.len())
+                .for_each(|start| {
+                    let end = self.data.len().min(start + BODY_LEN);
+                    out_buf[2..2 + end - start]
+                        .iter_mut()
+                        .zip(self.data[start..end].iter())
+                        .for_each(|(a, b)| *a ^= *b);
+                });
+            Some(2 + BODY_LEN.min(self.data.len() - self.groups * BODY_LEN))
+        } else {
+            None
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::UDP_BODY_LEN;
+
+    fn test_chunk_skip_packet(chunk: &[u8]) {
+        let mut encoder = Encoder::new(chunk.to_vec());
+        let mut packets = Vec::new();
+        let mut buf = [0u8; UDP_BODY_LEN];
+        while let Some(len) = encoder.next_packet(&mut buf) {
+            packets.push(buf[..len].to_vec());
+        }
+
+        packets.sort_by_key(|p| {
+            p.iter().take(6).fold(0u64, |acc, &x| {
+                acc.wrapping_mul(0x5DEECE66D).wrapping_add(x as u64)
+            })
+        });
+
+        for skip_idx in 0..packets.len() {
+            let mut decoder = Decoder::new(chunk.len());
+            for (idx, packet) in packets.iter().enumerate() {
+                if idx != skip_idx {
+                    decoder.add_packet(packet).expect("Failed to add packet");
+                }
+            }
+            let decoded = decoder.finish().expect("Failed to decode chunk");
+            assert_eq!(
+                decoded, chunk,
+                "Failed to decode chunk with skip index {skip_idx}"
+            );
+        }
+
+        let mut decoder = Decoder::new(chunk.len());
+        for (idx, packet) in packets.iter().enumerate() {
+            if !(idx.is_multiple_of(33) && idx / 33 < 32) {
+                decoder.add_packet(packet).expect("Failed to add packet");
+            }
+        }
+        let decoded = decoder.finish().expect("Failed to decode chunk");
+        assert_eq!(decoded, chunk, "Failed to decode chunk with multiple skips");
+    }
+
+    #[test]
+    fn test_small_chunk() {
+        let mut chunk = vec![0u8; 20];
+        let mut val = u64::MAX / 5;
+        for x in &mut chunk {
+            val = val.wrapping_mul(0x5DEECE66D).wrapping_add(0xB);
+            *x = val.to_be_bytes()[0];
+        }
+        test_chunk_skip_packet(&chunk);
+    }
+
+    #[test]
+    fn test_big_chunk() {
+        let mut chunk = vec![0u8; 200 << 10];
+        let mut val = u64::MAX / 5;
+        for x in &mut chunk {
+            val = val.wrapping_mul(0x5DEECE66D).wrapping_add(0xB);
+            *x = val.to_be_bytes()[0];
+        }
+        test_chunk_skip_packet(&chunk);
+    }
+}

pixie-shared/src/lib.rs

@@ -3,6 +3,7 @@
 extern crate alloc;
 
 pub mod bijection;
+pub mod chunk_codec;
 #[cfg(feature = "std")]
 pub mod config;
 
@@ -110,9 +111,7 @@ pub struct HintPacket {
 }
 
 /// The maximum number of bytes in a udp packet with mtu 1500
-pub const PACKET_LEN: usize = 1472;
-pub const HEADER_LEN: usize = 34;
-pub const BODY_LEN: usize = PACKET_LEN - HEADER_LEN;
+pub const UDP_BODY_LEN: usize = 1472;
 
 #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
 #[serde(rename_all = "lowercase")]