diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index db1a76a3..17cb99b1 100644 --- a/.github/workflows/rust.yml +++ b/.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 diff --git a/pixie-server/Cargo.lock b/pixie-server/Cargo.lock index 04caf99d..e3015805 100644 --- a/pixie-server/Cargo.lock +++ b/pixie-server/Cargo.lock @@ -1003,6 +1003,7 @@ dependencies = [ "ipnet", "macaddr", "serde", + "thiserror 2.0.12", ] [[package]] diff --git a/pixie-server/src/ping.rs b/pixie-server/src/ping.rs index 02e6fb60..d222b4b0 100644 --- a/pixie-server/src/ping.rs +++ b/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) -> 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?, diff --git a/pixie-server/src/udp.rs b/pixie-server/src/udp.rs index bb27bcc1..41042755 100644 --- a/pixie-server/src/udp.rs +++ b/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, ) -> Result<()> { let mut queue = BTreeSet::::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,7 +67,7 @@ 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 { @@ -75,40 +75,16 @@ async fn broadcast_chunks( 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?, diff --git a/pixie-shared/Cargo.toml b/pixie-shared/Cargo.toml index d0d6db79..75168286 100644 --- a/pixie-shared/Cargo.toml +++ b/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"] diff --git a/pixie-shared/src/chunk_codec.rs b/pixie-shared/src/chunk_codec.rs new file mode 100644 index 00000000..6e8e26c6 --- /dev/null +++ b/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, + missing_packet: Vec, + 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::>() + .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> { + 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, + groups: usize, + idx: usize, +} + +impl Encoder { + pub fn new(data: Vec) -> 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 { + 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); + } +} diff --git a/pixie-shared/src/lib.rs b/pixie-shared/src/lib.rs index 5be3d1a3..ee1c41da 100644 --- a/pixie-shared/src/lib.rs +++ b/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")] diff --git a/pixie-uefi/Cargo.lock b/pixie-uefi/Cargo.lock index 2042a500..a69fd636 100644 --- a/pixie-uefi/Cargo.lock +++ b/pixie-uefi/Cargo.lock @@ -319,6 +319,7 @@ version = "0.1.0" dependencies = [ "blake3", "serde", + "thiserror", ] [[package]] @@ -507,18 +508,18 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.4" +version = "2.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f49a1853cf82743e3b7950f77e0f4d622ca36cf4317cba00c767838bac8d490" +checksum = "567b8a2dae586314f7be2a752ec7474332959c6460e02bde30d702a66d488708" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "2.0.4" +version = "2.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8381894bb3efe0c4acac3ded651301ceee58a15d47c2e34885ed1908ad667061" +checksum = "7f7cf42b4507d8ea322120659672cf1b9dbb93f8f2d4ecfd6e51350ff5b17a1d" dependencies = [ "proc-macro2", "quote", diff --git a/pixie-uefi/src/flash.rs b/pixie-uefi/src/flash.rs index b2a7488e..1d2e84ed 100644 --- a/pixie-uefi/src/flash.rs +++ b/pixie-uefi/src/flash.rs @@ -11,37 +11,10 @@ use futures::future::{select, Either}; use log::info; use lz4_flex::decompress; use pixie_shared::{ - ChunkHash, Image, TcpRequest, UdpRequest, BODY_LEN, CHUNKS_PORT, HEADER_LEN, MAX_CHUNK_SIZE, + chunk_codec::Decoder, ChunkHash, Image, TcpRequest, UdpRequest, CHUNKS_PORT, MAX_CHUNK_SIZE, }; use uefi::proto::console::text::Color; -struct PartialChunk { - data: Vec, - missing_first: Vec, - missing_second: [u16; 32], - missing_third: u16, -} - -impl PartialChunk { - fn new(csize: usize) -> Self { - let num_packets = csize.div_ceil(BODY_LEN); - let data = vec![0; 32 * BODY_LEN + csize]; - let missing_first = vec![true; 32 + num_packets]; - let missing_second: [u16; 32] = (0..32) - .map(|i| ((num_packets + 31 - i) / 32) as u16) - .collect::>() - .try_into() - .unwrap(); - let missing_third = missing_second.iter().map(|&x| (x != 0) as u16).sum(); - PartialChunk { - data, - missing_first, - missing_second, - missing_third, - } - } -} - async fn fetch_image(stream: &TcpStream) -> Result { let req = TcpRequest::GetImage; let mut buf = postcard::to_allocvec(&req)?; @@ -65,76 +38,32 @@ struct Stats { fn handle_packet( buf: &[u8], chunks_info: &mut BTreeMap)>, - received: &mut BTreeMap, + received: &mut BTreeMap, last_seen: &mut Vec, ) -> Result, Vec)>> { let hash: ChunkHash = buf[..32].try_into().unwrap(); - let index = u16::from_le_bytes(buf[32..34].try_into().unwrap()) as usize; let csize = match chunks_info.get(&hash) { Some(&(_, csize, _)) => csize, _ => return Ok(None), }; - let pchunk = received - .entry(hash) - .or_insert_with(|| PartialChunk::new(csize)); + let decoder = received.entry(hash).or_insert_with(|| Decoder::new(csize)); last_seen.retain(|x| x != &hash); last_seen.push(hash); - let rot_index = (index as u16).wrapping_add(32) as usize; - match &mut pchunk.missing_first[rot_index] { - false => return Ok(None), - x @ true => *x = false, - } - - let start = rot_index * BODY_LEN; - pchunk.data[start..start + buf.len() - HEADER_LEN].clone_from_slice(&buf[HEADER_LEN..]); - - let group = index & 31; - match &mut pchunk.missing_second[group] { - 0 => return Ok(None), - x @ 1 => *x = 0, - x @ 2.. => { - *x -= 1; - return Ok(None); - } - } - - match &mut pchunk.missing_third { - 0 => unreachable!(), - x @ 1 => *x = 0, - x @ 2.. => { - *x -= 1; - return Ok(None); - } + if let Err(e) = decoder.add_packet(&buf[32..]) { + log::warn!("Received invalid packet for chunk {hash:02x?}: {e}"); + return Ok(None); } + let Some(cdata) = decoder.finish() else { + return Ok(None); + }; let (size, _, pos) = chunks_info.remove(&hash).unwrap(); - let mut pchunk = received.remove(&hash).unwrap(); + received.remove(&hash).unwrap(); last_seen.retain(|x| x != &hash); - let mut xor = [[0; BODY_LEN]; 32]; - for packet in 0..pchunk.missing_first.len() { - if !pchunk.missing_first[packet] { - let group = packet & 31; - pchunk.data[BODY_LEN * packet..] - .iter() - .zip(xor[group].iter_mut()) - .for_each(|(a, b)| *b ^= a); - } - } - for packet in 0..pchunk.missing_first.len() { - if pchunk.missing_first[packet] { - let group = packet & 31; - pchunk.data[BODY_LEN * packet..] - .iter_mut() - .zip(xor[group].iter()) - .for_each(|(a, b)| *a = *b); - } - } - - let data = decompress(&pchunk.data[32 * BODY_LEN..], size) - .map_err(|e| Error::Generic(e.to_string()))?; + let data = decompress(&cdata, size).map_err(|e| Error::Generic(e.to_string()))?; assert_eq!(data.len(), size); Ok(Some((pos, data))) diff --git a/pixie-web/Cargo.lock b/pixie-web/Cargo.lock index c34f50c0..08712bbc 100644 --- a/pixie-web/Cargo.lock +++ b/pixie-web/Cargo.lock @@ -241,7 +241,7 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d3ad3122b0001c7f140cf4d605ef9a9e2c24d96ab0b4fb4347b76de2425f445" dependencies = [ - "thiserror", + "thiserror 1.0.69", ] [[package]] @@ -594,7 +594,7 @@ dependencies = [ "pin-project", "serde", "serde_json", - "thiserror", + "thiserror 1.0.69", "wasm-bindgen", "wasm-bindgen-futures", "web-sys", @@ -1030,7 +1030,7 @@ dependencies = [ "lazy_static", "leptos", "paste", - "thiserror", + "thiserror 1.0.69", "unic-langid", "wasm-bindgen", "wasm-bindgen-futures", @@ -1046,7 +1046,7 @@ dependencies = [ "config", "regex", "serde", - "thiserror", + "thiserror 1.0.69", "typed-builder", ] @@ -1141,7 +1141,7 @@ dependencies = [ "serde-wasm-bindgen", "serde_json", "slotmap", - "thiserror", + "thiserror 1.0.69", "tracing", "wasm-bindgen", "wasm-bindgen-futures", @@ -1160,7 +1160,7 @@ dependencies = [ "leptos_reactive", "serde", "server_fn", - "thiserror", + "thiserror 1.0.69", "tracing", ] @@ -1297,7 +1297,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c51ebcefb2f0b9a5e0bea115532c8ae4215d1b01eff176d0f4ba4192895c2708" dependencies = [ "serde", - "thiserror", + "thiserror 1.0.69", ] [[package]] @@ -1459,6 +1459,7 @@ dependencies = [ "ipnet", "macaddr", "serde", + "thiserror 2.0.12", ] [[package]] @@ -1712,7 +1713,7 @@ dependencies = [ "quote", "syn", "syn_derive", - "thiserror", + "thiserror 1.0.69", ] [[package]] @@ -1820,7 +1821,7 @@ checksum = "0431a35568651e363364210c91983c1da5eb29404d9f0928b67d4ebcfa7d330c" dependencies = [ "percent-encoding", "serde", - "thiserror", + "thiserror 1.0.69", ] [[package]] @@ -1864,7 +1865,7 @@ dependencies = [ "serde_json", "serde_qs", "server_fn_macro_default", - "thiserror", + "thiserror 1.0.69", "url", "wasm-bindgen", "wasm-bindgen-futures", @@ -2051,7 +2052,16 @@ version = "1.0.69" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" dependencies = [ - "thiserror-impl", + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "567b8a2dae586314f7be2a752ec7474332959c6460e02bde30d702a66d488708" +dependencies = [ + "thiserror-impl 2.0.12", ] [[package]] @@ -2065,6 +2075,17 @@ dependencies = [ "syn", ] +[[package]] +name = "thiserror-impl" +version = "2.0.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f7cf42b4507d8ea322120659672cf1b9dbb93f8f2d4ecfd6e51350ff5b17a1d" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "time" version = "0.3.37" diff --git a/pixie-web/Cargo.toml b/pixie-web/Cargo.toml index 23953cc7..ec2a5ac2 100644 --- a/pixie-web/Cargo.toml +++ b/pixie-web/Cargo.toml @@ -3,6 +3,11 @@ name = "pixie-web" version = "0.1.0" edition = "2021" +[[bin]] +name = "pixie-web" +path = "src/main.rs" +test = false + [dependencies] bytes = "1.7.2" console_error_panic_hook = "0.1.7"