Refactoring.

`find_mac` takes an `IpAddr` instead of `Ipv4Addr`.
Deduplicate `BytesFmt` and move it to `pixie_shared::utils`.
Using `BytesFmt` more often in `pixie-uefi`.
Derive `From` for `pixie_uefi::os::error::Error`.
Use `gloo-net` instead of `reqwest` on `pixie-web`.
Fix images table alignment.
Add swap partition to the test.

Author: Virv12

pixie-server/src/main.rs

@@ -25,7 +25,7 @@ use tokio::{signal::unix::SignalKind, task::JoinHandle};
 ///
 /// This function searches the address in the arp cache, if it is not available it tries to
 /// populate it by pinging the peer.
-fn find_mac(ip: Ipv4Addr) -> Result<MacAddr6> {
+fn find_mac(ip: IpAddr) -> Result<MacAddr6> {
     struct Zombie {
         inner: Child,
     }
@@ -41,8 +41,6 @@ fn find_mac(ip: Ipv4Addr) -> Result<MacAddr6> {
         bail!("localhost not supported");
     }
 
-    let s = ip.to_string();
-
     // Repeat twice, sending a ping if looking at ip neigh the first time fails.
     for _ in 0..2 {
         let mut child = Zombie {
@@ -60,7 +58,7 @@ fn find_mac(ip: Ipv4Addr) -> Result<MacAddr6> {
             let line = line?;
             let mut parts = line.split(' ');
 
-            if parts.next() == Some(&s) {
+            if parts.next().and_then(|s| s.parse().ok()) == Some(ip) {
                 let mac = parts.nth(3).unwrap();
                 if let Ok(mac) = mac.parse() {
                     return Ok(mac);
@@ -69,7 +67,7 @@ fn find_mac(ip: Ipv4Addr) -> Result<MacAddr6> {
         }
 
         let _ = Command::new("ping")
-            .args([&s, "-c", "1", "-W", "0.1"])
+            .args([&ip.to_string(), "-c", "1", "-W", "0.1"])
             .stdout(Stdio::null())
             .spawn()?
             .wait();

pixie-server/src/ping.rs

@@ -4,13 +4,9 @@ use crate::{
     find_mac,
     state::{State, UnitSelector},
 };
-use anyhow::{bail, Result};
+use anyhow::Result;
 use pixie_shared::{PING_PORT, UDP_BODY_LEN};
-use std::{
-    net::{IpAddr, Ipv4Addr},
-    sync::Arc,
-    time::SystemTime,
-};
+use std::{net::Ipv4Addr, sync::Arc, time::SystemTime};
 use tokio::net::UdpSocket;
 
 pub async fn main(state: Arc<State>) -> Result<()> {
@@ -23,10 +19,7 @@ pub async fn main(state: Arc<State>) -> Result<()> {
             x = socket.recv_from(&mut buf) => x?,
             _ = state.cancel_token.cancelled() => break,
         };
-        let IpAddr::V4(peer_ip) = peer_addr.ip() else {
-            bail!("IPv6 is not supported")
-        };
-        let peer_mac = match find_mac(peer_ip) {
+        let peer_mac = match find_mac(peer_addr.ip()) {
             Ok(peer_mac) => peer_mac,
             Err(err) => {
                 log::error!("Error handling ping packet: {err}");

pixie-server/src/tcp.rs

@@ -4,12 +4,12 @@ use crate::{
     find_mac,
     state::{State, UnitSelector},
 };
-use anyhow::{bail, Context, Result};
+use anyhow::{Context, Result};
 use macaddr::MacAddr6;
 use pixie_shared::{TcpRequest, ACTION_PORT};
 use std::{
     io::ErrorKind,
-    net::{IpAddr, Ipv4Addr, SocketAddr},
+    net::{Ipv4Addr, SocketAddr},
     sync::Arc,
 };
 use tokio::{
@@ -57,10 +57,7 @@ async fn handle_connection(
     mut stream: TcpStream,
     peer_addr: SocketAddr,
 ) -> Result<()> {
-    let IpAddr::V4(peer_ip) = peer_addr.ip() else {
-        bail!("IPv6 is not supported")
-    };
-    let peer_mac = match find_mac(peer_ip) {
+    let peer_mac = match find_mac(peer_addr.ip()) {
         Ok(peer_mac) => peer_mac,
         Err(err) => {
             log::error!("Error handling tcp connection: {err}");

pixie-server/src/udp.rs

@@ -4,14 +4,14 @@ use crate::{
     find_mac, find_network,
     state::{State, UnitSelector},
 };
-use anyhow::{bail, ensure, Context, Result};
+use anyhow::{ensure, Context, Result};
 use pixie_shared::{
     chunk_codec::Encoder, ChunkHash, HintPacket, RegistrationInfo, UdpRequest, ACTION_PORT,
     CHUNKS_PORT, HINT_PORT, UDP_BODY_LEN,
 };
 use std::{
     collections::BTreeSet,
-    net::{IpAddr, Ipv4Addr, SocketAddrV4},
+    net::{Ipv4Addr, SocketAddrV4},
     ops::Bound,
     sync::Arc,
 };
@@ -168,20 +168,17 @@ 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; UDP_BODY_LEN];
     loop {
-        let (len, addr) = tokio::select! {
+        let (len, peer_addr) = tokio::select! {
             x = socket.recv_from(&mut buf) => x?,
             _ = state.cancel_token.cancelled() => break,
         };
         let req: postcard::Result<UdpRequest> = postcard::from_bytes(&buf[..len]);
         match req {
             Ok(UdpRequest::Discover) => {
-                socket.send_to(&[], addr).await?;
+                socket.send_to(&[], peer_addr).await?;
             }
             Ok(UdpRequest::ActionProgress(frac, tot)) => {
-                let IpAddr::V4(peer_ip) = addr.ip() else {
-                    bail!("IPv6 is not supported")
-                };
-                match find_mac(peer_ip) {
+                match find_mac(peer_addr.ip()) {
                     Ok(peer_mac) => {
                         state.set_unit_progress(UnitSelector::MacAddr(peer_mac), Some((frac, tot)));
                     }
@@ -196,7 +193,7 @@ async fn handle_requests(state: &State, socket: &UdpSocket, tx: Sender<[u8; 32]>
                 }
             }
             Err(e) => {
-                log::warn!("Invalid request from {addr}: {e}");
+                log::warn!("Invalid request from {peer_addr}: {e}");
             }
         }
     }

pixie-shared/src/lib.rs

@@ -6,6 +6,7 @@ pub mod bijection;
 pub mod chunk_codec;
 #[cfg(feature = "std")]
 pub mod config;
+pub mod util;
 
 use alloc::{collections::BTreeMap, string::String, vec::Vec};
 use blake3::OUT_LEN;

pixie-shared/src/util.rs

@@ -0,0 +1,17 @@
+pub struct BytesFmt(pub u64);
+
+impl core::fmt::Display for BytesFmt {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        if self.0 < 1 << 10 {
+            write!(f, "{} B", self.0)
+        } else if self.0 < 1 << 20 {
+            write!(f, "{:.2} KiB", self.0 as f64 / (1i64 << 10) as f64)
+        } else if self.0 < 1 << 30 {
+            write!(f, "{:.2} MiB", self.0 as f64 / (1i64 << 20) as f64)
+        } else if self.0 < 1 << 40 {
+            write!(f, "{:.2} GiB", self.0 as f64 / (1i64 << 30) as f64)
+        } else {
+            write!(f, "{:.2} TiB", self.0 as f64 / (1i64 << 40) as f64)
+        }
+    }
+}

pixie-uefi/Cargo.lock

@@ -16,6 +16,7 @@ coverage = ["dep:minicov"]
 anstyle = { version = "1.0.11", default-features = false }
 blake3 = { version = "1.8.2", default-features = false, features = ["prefer_intrinsics", "no_avx512", "pure"] }
 core_detect = "1.0.0"
+derive_more = { version = "2.0.1", features = ["from"], default-features = false }
 futures = { version = "0.3.29", default-features = false, features = ["alloc", "async-await"] }
 gpt_disk_io = "0.16.2"
 log = "0.4.27"

pixie-uefi/Cargo.toml

@@ -11,7 +11,8 @@ use futures::future::{select, Either};
 use log::info;
 use lz4_flex::decompress;
 use pixie_shared::{
-    chunk_codec::Decoder, ChunkHash, Image, TcpRequest, UdpRequest, CHUNKS_PORT, MAX_CHUNK_SIZE,
+    chunk_codec::Decoder, util::BytesFmt, ChunkHash, Image, TcpRequest, UdpRequest, CHUNKS_PORT,
+    MAX_CHUNK_SIZE,
 };
 use uefi::proto::console::text::Color;
 
@@ -168,7 +169,10 @@ pub async fn flash(os: UefiOS, server_addr: SocketAddrV4) -> Result<()> {
         let mut last_seen = Vec::new();
         let total_mem = os.get_total_mem();
         let max_chunks = (total_mem.saturating_sub(MIN_MEMORY) as usize / MAX_CHUNK_SIZE).max(128);
-        log::debug!("Total memory: {total_mem}. Max chunks in memory: {max_chunks}");
+        log::debug!(
+            "Total memory: {}. Max chunks in memory: {max_chunks}",
+            BytesFmt(total_mem)
+        );
         while !chunks_info.is_empty() {
             let recv = Box::pin(socket.recv(&mut buf));
             let sleep = Box::pin(os.sleep_us(100_000));