diff --git a/Cargo.toml b/Cargo.toml index 5afe95d..e8b1f61 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,6 +3,9 @@ name = "flood-rs" version = "0.1.0" edition = "2024" +[profile.release] +debug = true + [dependencies] png = "0.17" libc = "0.2" diff --git a/flake.lock b/flake.lock new file mode 100644 index 0000000..7dac337 --- /dev/null +++ b/flake.lock @@ -0,0 +1,27 @@ +{ + "nodes": { + "nixpkgs": { + "locked": { + "lastModified": 1759733170, + "narHash": "sha256-TXnlsVb5Z8HXZ6mZoeOAIwxmvGHp1g4Dw89eLvIwKVI=", + "owner": "nixos", + "repo": "nixpkgs", + "rev": "8913c168d1c56dc49a7718685968f38752171c3b", + "type": "github" + }, + "original": { + "owner": "nixos", + "ref": "nixos-unstable", + "repo": "nixpkgs", + "type": "github" + } + }, + "root": { + "inputs": { + "nixpkgs": "nixpkgs" + } + } + }, + "root": "root", + "version": 7 +} diff --git a/flake.nix b/flake.nix new file mode 100644 index 0000000..004750b --- /dev/null +++ b/flake.nix @@ -0,0 +1,29 @@ +{ + description = "A very basic flake"; + + inputs = { + nixpkgs.url = "github:nixos/nixpkgs?ref=nixos-unstable"; + }; + + outputs = { self, nixpkgs }: { + devShells.x86_64-linux.default = + let + insecure-package-overlay = final: prev: { + nixpkgs.config.permittedInsecurePackages = [ + # Add the full name and version of the package here + "openssl-1.1.1w" + # "another-insecure-package-2.0" + ]; + }; + pkgs = import nixpkgs { + system = "x86_64-linux"; + overlays = []; + }; + in + pkgs.mkShell { + buildInputs = [ + pkgs.cargo + ]; + }; + }; +} diff --git a/src/bouncing_image.rs b/src/bouncing_image.rs new file mode 100644 index 0000000..bd4544c --- /dev/null +++ b/src/bouncing_image.rs @@ -0,0 +1,85 @@ +use crate::drawable::Drawable; +use crate::pixel_buf::PixelBuf; +use crate::png_data::PngData; +use crate::{DISPLAY_HEIGHT, DISPLAY_WIDTH}; +use crate::pixel::Pixel; +/// Represents a bouncing image on the screen. +pub struct BouncingImage { + img: PngData, + x: f32, + y: f32, + x1: f32, + y1: f32, + x2: f32, + y2: f32, + /// Horizontal velocity in pixels per second. + move_x: f32, + /// Vertical velocity in pixels per second. + move_y: f32, +} + +impl BouncingImage { + /// Initializes a new BouncingImage. + /// + /// - `move_x`, `move_y`: velocity in pixels per second. + /// - `start_x`, `start_y`: initial position (-1 to center on that axis). + pub fn new(img_file: &str, move_x: f32, move_y: f32, start_x: i32, start_y: i32) -> Self { + let img = PngData::open(img_file).expect("Could not load image"); + let x2 = DISPLAY_WIDTH as f32 - img.width as f32; + let y2 = DISPLAY_HEIGHT as f32 - img.height as f32; + let x1: f32 = 0.0; + let y1: f32 = 0.0; + + let x = if start_x == -1 { + (x1 + x2) / 2.0 + } else { + start_x as f32 + }; + let y = if start_y == -1 { + (y1 + y2) / 2.0 + } else { + start_y as f32 + }; + + BouncingImage { + img, + x, + y, + x1, + y1, + x2, + y2, + move_x, + move_y, + } + } +} +fn offset_pixels(pixels: &Vec, dx: u16, dy: u16) -> Vec { + let mut out = pixels.clone(); // memcpy, since Pixel is Copy + for p in out.iter_mut() { + p.x = p.x.saturating_add(dx); + p.y = p.y.saturating_add(dy); + } + out +} +impl Drawable for BouncingImage { + fn draw(&self, buf: &mut PixelBuf, _elapsed: f32) { + + let pixels = offset_pixels(&self.img.packed_pixels,self.x as u16,self.y as u16); + buf.set_pixels(pixels); + } + + fn update(&mut self, dt: f32, _elapsed: f32) { + self.x += self.move_x * dt; + self.y += self.move_y * dt; + + if self.x < self.x1 || self.x > self.x2 { + self.move_x *= -1.0; + self.x = self.x.clamp(self.x1, self.x2); + } + if self.y < self.y1 || self.y > self.y2 { + self.move_y *= -1.0; + self.y = self.y.clamp(self.y1, self.y2); + } + } +} diff --git a/src/circle.rs b/src/circle.rs new file mode 100644 index 0000000..c975228 --- /dev/null +++ b/src/circle.rs @@ -0,0 +1,109 @@ +use std::sync::{Arc, Mutex}; + +use crate::color; +use crate::pixel::Pixel; +use crate::drawable::Drawable; +use crate::pixel_buf::PixelBuf; + +/// Speed at which the radius shrinks, in pixels per second. +const RADIUS_SHRINK_RATE: f32 = 30.0; + +/// Speed at which the hue cycles, in degrees per second. +const HUE_CYCLE_RATE: f32 = 180.0; + +/// Represents a circle drawn at coordinates received via shared state. +pub struct Circle { + x: Arc>, + y: Arc>, + set: Arc>, + radius: f32, + /// Cached draw coordinates (read from mutex in update, used in draw). + draw_x: u32, + draw_y: u32, + color: color::Rgb +} + +impl Circle { + pub fn new(x: Arc>, y: Arc>, set: Arc>) -> Self { + Circle { + x, + y, + set, + radius: 0.0, + draw_x: 0, + draw_y: 0, + color: color::Rgb { r:0,g:0,b:0} + } + } + + /// Draws a circle using the Midpoint Circle Algorithm, writing 8 octant + /// points per step. + fn draw_circle( + buf: &mut PixelBuf, + cx: i32, + cy: i32, + radius: u32, + r: u8, + g: u8, + b: u8, + ) { + let rad = radius as i32; + let mut x: i32 = 0; + let mut y: i32 = rad; + let mut d: i32 = 3 - 2 * rad; + + while y >= x { + // 8 octant points. + buf.set_pixel(Pixel {x:(cx + x) as u16, y:(cy + y) as u16, r, g, b}); + buf.set_pixel(Pixel {x:(cx - x) as u16, y:(cy + y) as u16, r, g, b}); + buf.set_pixel(Pixel {x:(cx + x) as u16, y:(cy - y) as u16, r, g, b}); + buf.set_pixel(Pixel {x:(cx - x) as u16, y:(cy - y) as u16, r, g, b}); + buf.set_pixel(Pixel {x:(cx + y) as u16, y:(cy + x) as u16, r, g, b}); + buf.set_pixel(Pixel {x:(cx - y) as u16, y:(cy + x) as u16, r, g, b}); + buf.set_pixel(Pixel {x:(cx + y) as u16, y:(cy - x) as u16, r, g, b}); + buf.set_pixel(Pixel {x:(cx - y) as u16, y:(cy - x) as u16, r, g, b}); + + x += 1; + if d > 0 { + y -= 1; + d = d + 4 * (x - y) + 10; + } else { + d = d + 4 * x + 6; + } + } + } +} + +impl Drawable for Circle { + fn draw(&self, buf: &mut PixelBuf, _elapsed: f32) { + + + let cx = self.draw_y as i32; + let cy = self.draw_x as i32; + let radius = self.radius as u32; + + for i in 0..10 { + if radius > i { + Self::draw_circle(buf, cx, cy, radius - i, self.color.r, self.color.g, self.color.b); + } + } + } + + fn update(&mut self, dt: f32, elapsed: f32) { + self.radius = 150.0; + self.draw_y = (*self.x.lock().unwrap() + (elapsed as u32))%1080; + self.draw_x = (*self.y.lock().unwrap() + (elapsed as u32))%1920; + + self.radius -= RADIUS_SHRINK_RATE * dt; + if self.radius < 0.0 { + self.radius = 0.0; + } + let hue = ((elapsed * HUE_CYCLE_RATE) % 360.0) as u16; + let hsv_color = color::Hsv { + h: hue, + s: 1.0, + v: 1.0, + }; + self.color = hsv_color.into(); + } +} diff --git a/src/display.rs b/src/display.rs new file mode 100644 index 0000000..c0e178e --- /dev/null +++ b/src/display.rs @@ -0,0 +1,304 @@ +use std::collections::VecDeque; +use std::net::{ToSocketAddrs, UdpSocket}; +use std::sync::{Arc, Condvar, Mutex}; +use std::thread; +use std::time::Instant; + +use crate::raw_socket::RawSender; + +const QUEUE_LEN: usize = 1000; +/// Max pixels per packet. Constrained by MTU 1500: +/// 1500 (IP payload max) - 20 (IP hdr) - 8 (UDP hdr) - 2 (pixel hdr) = 1470 +/// 1470 / 7 bytes per pixel = 210 pixels. +const PIXELS_PER_PACKET: usize = 210; +const MSG_PAYLOAD_SIZE: usize = 7 * PIXELS_PER_PACKET; +const MSGSIZE: usize = 2 + MSG_PAYLOAD_SIZE; + +/// Maximum number of packets to batch into a single sendmmsg call. +const SEND_BATCH_SIZE: usize = 26; + +/// A filled buffer ready to be sent, carrying its valid data length. +struct Packet { + buf: Box<[u8; MSGSIZE]>, + len: usize, +} + +/// Shared state between all pixel writers and the sender thread. +struct SharedState { + /// Filled packets waiting to be sent. + pending: Mutex>, + /// Wakes the sender when packets are available. + condvar: Condvar, + /// Pool of empty, reusable buffers. + pool: Mutex>>, +} + +/// Allocates a fresh buffer with the protocol header bytes set. +fn alloc_buf() -> Box<[u8; MSGSIZE]> { + let mut buf = Box::new([0u8; MSGSIZE]); + buf[0] = 0x00; + buf[1] = 0x01; + buf +} + +/// A pixel writer that fills buffers and submits them for sending. +/// +/// Multiple writers can exist concurrently, each with its own buffer, +/// all sharing the same pending queue and buffer pool. This enables +/// multi-threaded drawing. +/// +/// Writers never block. If the pool is empty, they steal the oldest +/// unsent packet or allocate a fresh buffer. +pub struct Display { + shared: Arc, + /// The buffer currently being filled with pixel data. + current_buf: Box<[u8; MSGSIZE]>, + /// How many pixels have been written into `current_buf`. + pos_in_buf: usize, +} + +impl Display { + /// Creates a new Display and spawns the sender thread. + /// + /// The returned Display can create additional writers via + /// `create_writer()` for multi-threaded drawing. + pub fn new(host: &str, port: u16) -> Self { + let shared = Arc::new(SharedState { + pending: Mutex::new(VecDeque::with_capacity(QUEUE_LEN)), + condvar: Condvar::new(), + pool: Mutex::new(Vec::with_capacity(QUEUE_LEN)), + }); + + // Pre-allocate the buffer pool. Keep one out as the initial current_buf. + { + let mut pool = shared.pool.lock().unwrap(); + for _ in 1..QUEUE_LEN { + pool.push(alloc_buf()); + } + } + + for _ in [1,2,3,4,5] { + // Spawn the sender thread. + let shared_sender = Arc::clone(&shared); + match RawSender::new(host, port, 0) { + Ok(raw_sender) => { + eprintln!("[display] Using raw AF_PACKET sender"); + thread::Builder::new() + .name("sender-raw".into()) + .spawn(move || { + sender_loop_raw(raw_sender, shared_sender); + }).unwrap(); + } + Err(e) => { + eprintln!("[display] Raw AF_PACKET unavailable ({}), falling back to UDP", e); + let remote_addr = (host, port) + .to_socket_addrs() + .expect("Invalid remote address") + .next() + .expect("Could not resolve host"); + let socket = UdpSocket::bind("0.0.0.0:0").expect("Could not bind"); + socket.connect(remote_addr).expect("Could not connect"); + thread::Builder::new() + .name("sender-raw".into()) + .spawn(move || { + sender_loop_udp(socket, shared_sender); + }).unwrap(); + } + }; + } + + + Display { + shared, + current_buf: alloc_buf(), + pos_in_buf: 0, + } + } + + /// Grabs an empty buffer: try pool first, steal oldest pending, or allocate. + fn grab_buf(&self) -> Box<[u8; MSGSIZE]> { + // 1. Try the pool (fast path). + { + let mut pool = self.shared.pool.lock().unwrap(); + if let Some(buf) = pool.pop() { + return buf; + } + } + // 2. Pool empty -- steal the oldest unsent packet. + { + let mut pending = self.shared.pending.lock().unwrap(); + if let Some(oldest) = pending.pop_front() { + return oldest.buf; + } + } + // 3. Everything is in-flight. Allocate a fresh buffer. + alloc_buf() + } + + /// Flushes the current buffer if it contains pixel data. + /// + /// This method **never blocks**. + pub fn flush_frame(&mut self) { + if self.pos_in_buf == 0 { + return; + } + + let len = 2 + self.pos_in_buf * 7; + let new_buf = self.grab_buf(); + let filled_buf = std::mem::replace(&mut self.current_buf, new_buf); + + { + let mut pending = self.shared.pending.lock().unwrap(); + pending.push_back(Packet { + buf: filled_buf, + len, + }); + } + self.shared.condvar.notify_one(); + self.pos_in_buf = 0; + } + + #[inline] + pub fn write_raw_pixels(&mut self, pixels: &[[u8; 7]]) { + let mut remaining = pixels; + while !remaining.is_empty() { + let space = PIXELS_PER_PACKET - self.pos_in_buf; + let n = remaining.len().min(space); + let dst_offset = 2 + self.pos_in_buf * 7; + + // SAFETY: [u8; 7] has align 1 and no padding, so &[[u8; 7]] + // is a contiguous block of bytes we can copy in one shot. + let src = unsafe { + std::slice::from_raw_parts(remaining.as_ptr() as *const u8, n * 7) + }; + self.current_buf[dst_offset..dst_offset + n * 7].copy_from_slice(src); + + self.pos_in_buf += n; + remaining = &remaining[n..]; + + if self.pos_in_buf == PIXELS_PER_PACKET { + self.flush_frame(); + } + } + } + +} + +// --------------------------------------------------------------------------- +// Sender threads +// --------------------------------------------------------------------------- + +/// Raw AF_PACKET sender loop. +fn sender_loop_raw(raw_sender: RawSender, shared: Arc) { + let mut batch: Vec = Vec::with_capacity(SEND_BATCH_SIZE); + let mut frame_buf: Vec = Vec::with_capacity(SEND_BATCH_SIZE * (MSGSIZE + 42)); + + let mut stats_sent: u64 = 0; + let mut stats_bytes: u64 = 0; + let mut stats_errors: u64 = 0; + let mut stats_send_us: u128 = 0; + let mut stats_send_calls: u64 = 0; + let mut stats_last_report = Instant::now(); + + loop { + { + let mut pending = shared.pending.lock().unwrap(); + while pending.is_empty() { + pending = shared.condvar.wait(pending).unwrap(); + } + let n = pending.len().min(SEND_BATCH_SIZE); + batch.extend(pending.drain(..n)); + } + + let payload_ptrs: Vec<&[u8]> = batch.iter().map(|p| &p.buf[..p.len]).collect(); + + let send_start = Instant::now(); + match raw_sender.send_batch(&payload_ptrs, &mut frame_buf) { + Ok(n) => { + stats_sent += n as u64; + for packet in batch.iter().take(n) { + stats_bytes += (packet.len + raw_sender.header_size()) as u64; + } + } + Err(e) => { + if stats_errors == 0 { + eprintln!("sendmmsg (raw) failed: {}", e); + } + stats_errors += batch.len() as u64; + } + } + stats_send_us += send_start.elapsed().as_micros(); + stats_send_calls += 1; + + // Return buffers to the pool. + { + let mut pool = shared.pool.lock().unwrap(); + for packet in batch.drain(..) { + pool.push(packet.buf); + } + } + + let now = Instant::now(); + let elapsed = (now - stats_last_report).as_secs_f64(); + if elapsed >= 5.0 { + let pps = stats_sent as f64 / elapsed; + let mbps = stats_bytes as f64 * 8.0 / elapsed / 1_000_000.0; + eprintln!( + "[sender] {:.0} pkt/s, {:.1} Mbit/s sent | errors: {} | send: {:.0}us avg over {} calls", + pps, mbps, stats_errors, + stats_send_us as f64 / stats_send_calls.max(1) as f64, + stats_send_calls, + ); + stats_sent = 0; + stats_bytes = 0; + stats_errors = 0; + stats_send_us = 0; + stats_send_calls = 0; + stats_last_report = now; + } + } +} + +/// Fallback UDP sender loop. +fn sender_loop_udp(socket: UdpSocket, shared: Arc) { + let mut batch: Vec = Vec::with_capacity(SEND_BATCH_SIZE); + let mut stats_sent: u64 = 0; + let mut stats_bytes: u64 = 0; + let mut stats_last_report = Instant::now(); + + loop { + { + let mut pending = shared.pending.lock().unwrap(); + while pending.is_empty() { + pending = shared.condvar.wait(pending).unwrap(); + } + let n = pending.len().min(SEND_BATCH_SIZE); + batch.extend(pending.drain(..n)); + } + + for packet in &batch { + if let Ok(n) = socket.send(&packet.buf[..packet.len]) { + stats_sent += 1; + stats_bytes += n as u64; + } + } + + { + let mut pool = shared.pool.lock().unwrap(); + for packet in batch.drain(..) { + pool.push(packet.buf); + } + } + + let now = Instant::now(); + let elapsed = (now - stats_last_report).as_secs_f64(); + if elapsed >= 5.0 { + let pps = stats_sent as f64 / elapsed; + let mbps = stats_bytes as f64 * 8.0 / elapsed / 1_000_000.0; + eprintln!("[sender] {:.0} pkt/s, {:.1} Mbit/s (UDP fallback)", pps, mbps); + stats_sent = 0; + stats_bytes = 0; + stats_last_report = now; + } + } +} diff --git a/src/drawable.rs b/src/drawable.rs new file mode 100644 index 0000000..3850860 --- /dev/null +++ b/src/drawable.rs @@ -0,0 +1,8 @@ +use crate::pixel_buf::PixelBuf; + +pub trait Drawable: Send + Sync { + fn draw(&self, buf: &mut PixelBuf, elapsed: f32); + + + fn update(&mut self, dt: f32, elapsed: f32); +} diff --git a/src/main.rs b/src/main.rs index 8585c0b..ad679ca 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,396 +1,187 @@ -use std::fs::File; -use std::io::BufReader; -use std::net::{ToSocketAddrs, UdpSocket}; -use std::time::Duration; +use std::net::{SocketAddr, UdpSocket}; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::sync::{Arc, Barrier, Mutex, RwLock}; use std::thread; -use std::sync::{Mutex,Arc}; +use std::time::{Duration, Instant}; + use socket2::{Domain, Socket, Type}; -use std::net::{Ipv4Addr, SocketAddr}; + +mod bouncing_image; +mod circle; mod color; -// Constants from the C code -const QUEUE_LEN: usize = 1000; -const MSG_PAYLOAD_SIZE: usize = 7 * 211; -const MSGSIZE: usize = 2 + MSG_PAYLOAD_SIZE; +mod display; +mod drawable; +mod pixel_buf; +mod png_data; +mod raw_socket; +mod pixel; -const DISPLAY_HOST: &str = "100.65.0.2"; -const DISPLAY_PORT: u16 = 5005; -const DISPLAY_WIDTH: i32 = 1920; -const DISPLAY_HEIGHT: i32 = 1080; - -/// Represents the data decoded from a PNG file. -struct PngData { - width: u32, - height: u32, - pixels: Vec, -} - -impl PngData { - /// Loads and decodes a PNG image from the given path. - fn open(path: &str) -> Result { - let file = File::open(path).expect("Failed to open PNG file"); - let decoder = png::Decoder::new(BufReader::new(file)); - let mut reader = decoder.read_info()?; - let mut buf = vec![0; reader.output_buffer_size()]; - let info = reader.next_frame(&mut buf)?; - - Ok(PngData { - width: info.width, - height: info.height, - pixels: buf, - }) - } -} - -/// Represents a bouncing image on the screen. -struct BouncingImage { - img: PngData, - x: i32, - y: i32, - x1: i32, - y1: i32, - x2: i32, - y2: i32, - move_x: i32, - move_y: i32, - rate: u32, -} - -trait Drawable { - // Associated function signature; `Self` refers to the implementor type. - fn rate(&self) -> u32; - fn draw_and_move(&mut self, display: &mut Display,tick:u32); - -} -impl BouncingImage { - /// Initializes a new BouncingImage. - fn new(img_file: &str, move_x: i32, move_y: i32, rate: u32, start_x: i32, start_y: i32) -> Self { - let img = PngData::open(img_file).expect("Could not load image"); - let mut bb = BouncingImage { - x2: DISPLAY_WIDTH - img.width as i32, - y2: DISPLAY_HEIGHT - img.height as i32, - img, - x: start_x, - y: start_y, - x1: 0, - y1: 0, - move_x, - move_y, - rate, - }; - if bb.x == -1 { - bb.x = (bb.x1 + bb.x2) / 2; - } - if bb.y == -1 { - bb.y = (bb.y1 + bb.y2) / 2; - } - bb - } - - - /// Draws a PNG image at the given coordinates. - fn draw_png(&mut self, display: &mut Display, x: i32, y: i32) { - - for sy in 0..self.img.height { - for sx in 0..self.img.width { - let index = (sy * self.img.width + sx) as usize * 4; - let rgba = &self.img.pixels[index..index + 4]; - if rgba[3] > 0 { // Check alpha channel - //display.set_pixel((x + sx as i32) as u16, (y + sy as i32) as u16, rgba[0], rgba[1], rgba[2]); - } - } - } - } - -} -impl Drawable for BouncingImage { - fn rate(&self) -> u32 { - return self.rate; - } - - /// Draws the image and updates its position. - fn draw_and_move(&mut self, display: &mut Display,_: u32) { - self.draw_png(display, self.x, self.y); - - self.x += self.move_x; - self.y += self.move_y; - - if self.x < self.x1 || self.x > self.x2 { - self.move_x *= -1; - } - if self.y < self.y1 || self.y > self.y2 { - self.move_y *= -1; - } - } -} - -struct Circle { - x: Arc>, - y: Arc>, -} -impl Circle { - - fn new(x:Arc>, y: Arc>) -> Self { - Circle { - x, - y - } - } - - /// This exploits the eight-way symmetry of a circle. - fn draw_circle_octants(&mut self, display: &mut Display, cx: i32, cy: i32, x: i32, y: i32, r: u8, g: u8, b: u8) { - display.set_pixel(cx + x, cy + y, r, g, b); - display.set_pixel(cx - x, cy + y, r, g, b); - display.set_pixel(cx + x, cy - y, r, g, b); - display.set_pixel(cx - x, cy - y, r, g, b); - display.set_pixel(cx + y, cy + x, r, g, b); - display.set_pixel(cx - y, cy + x, r, g, b); - display.set_pixel(cx + y, cy - x, r, g, b); - display.set_pixel(cx - y, cy - x, r, g, b); - } - - /// Draws a circle using the Midpoint Circle Algorithm. - /// - /// # Arguments - /// * `center_x`: The x-coordinate of the circle's center. - /// * `center_y`: The y-coordinate of the circle's center. - /// * `radius`: The radius of the circle. Must be non-negative. - /// * `r`, `g`, `b`: The RGB color components for the circle. - pub fn draw_circle(&mut self, display: &mut Display, center_x: u32, center_y: u32, radius: u32, r: u8, g: u8, b: u8) { - if radius < 0 { - // Or return an error: Err("Radius cannot be negative".into()) - return; - } - let radius_i32:i32 = radius.try_into().unwrap(); - let mut x:i32 = 0; - let mut y:i32 = radius_i32; - // Initial decision parameter - let mut d:i32 = 3 - 2 * radius_i32; - - // Iterate through the first octant and draw points in all 8 octants - while y >= x { - self.draw_circle_octants(display,center_x.try_into().unwrap(), center_y.try_into().unwrap(), x, y, r, g, b); - - x += 1; - - // Update the decision parameter - if d > 0 { - y -= 1; - d = d + 4 * (x - y) + 10; - } else { - d = d + 4 * x + 6; - } - } - } - -} -impl Drawable for Circle { - - fn rate(&self) -> u32 { - 1 - } - - /// Helper method to draw the 8 symmetric points for a given (x, y) offset. - fn draw_and_move(&mut self, display: &mut Display,tick:u32) { - let hsv_color = color::Hsv { - h: ((tick/200)%360).try_into().unwrap(), - s: 1.0, - v: 1.0, - }; - let rgb: color::Rgb = hsv_color.into(); - let draw_y =*self.x.lock().unwrap(); - let draw_x = *self.y.lock().unwrap(); - let radius = (tick/30) % (300/2); - self.draw_circle(display,draw_y,draw_x,radius.try_into().unwrap(),rgb.r,rgb.g,rgb.b); - } -} - -/// Manages the connection and data sent to the display. -struct Display { - socket: UdpSocket, - bufs: Vec<[u8; MSGSIZE]>, - next_buf: usize, - pos_in_buf: usize, -} - -impl Display { - /// Creates a new Display and connects to the specified host and port. - fn new(host: &str, port: u16) -> Self { - let remote_addr = (host, port) - .to_socket_addrs() - .expect("Invalid remote address") - .next() - .expect("Could not resolve host"); - - let socket = UdpSocket::bind("0.0.0.0:0").expect("Could not bind to local port"); - socket.connect(remote_addr).expect("Could not connect to remote"); - - let mut bufs = vec![[0; MSGSIZE]; QUEUE_LEN]; - for buf in bufs.iter_mut() { - buf[0] = 0x00; - buf[1] = 0x01; - } - - Display { - socket, - bufs, - next_buf: 0, - pos_in_buf: 0, - } - } - - /// Flushes the current buffer if it contains pixel data. - fn flush_frame(&mut self) { - if self.pos_in_buf > 0 { - let len = 2 + self.pos_in_buf * 7; - let buf_to_send = &self.bufs[self.next_buf][..len]; - self.socket.send(buf_to_send).expect("Failed to send data"); - self.next_buf = (self.next_buf + 1) % QUEUE_LEN; - self.pos_in_buf = 0; - } - } - - /// Sets a pixel color at a specific coordinate. - fn set_pixel(&mut self, x: i32, y: i32, r: u8, g: u8, b: u8) { - if let (Ok(output_x),Ok(output_y)) = (u16::try_from(x), u16::try_from(y)) { - let offset = 2 + self.pos_in_buf * 7; - let buf = &mut self.bufs[self.next_buf][offset..offset + 7]; - buf[0] = output_x as u8; - buf[1] = (output_x >> 8) as u8; - buf[2] = output_y as u8; - buf[3] = (output_y >> 8) as u8; - buf[4] = r; - buf[5] = g; - buf[6] = b; - - self.pos_in_buf += 1; - if self.pos_in_buf == 211 { - self.flush_frame(); - } - } - - - } - - - /// Clears the entire screen to black. - #[allow(dead_code)] - fn blank_screen(&mut self) { - for x in 0..DISPLAY_WIDTH { - for y in 0..DISPLAY_HEIGHT { - //self.set_pixel(x as u16, y as u16, 0, 0, 0); - } - } - self.flush_frame(); - } -} +use bouncing_image::BouncingImage; +use circle::Circle; +use display::Display; +use drawable::Drawable; +use pixel_buf::PixelBuf; +// Display configuration constants +pub const DISPLAY_HOST: &str = "100.65.0.2"; +pub const DISPLAY_PORT: u16 = 5005; +pub const DISPLAY_WIDTH: i32 = 1920; +pub const DISPLAY_HEIGHT: i32 = 1080; /// Unpacks a 4-byte slice into two u16 values (little-endian). fn unpack_coordinates(buffer: &[u8]) -> Option<(u16, u16)> { if buffer.len() != 4 { return None; } - // Try to convert the first 2 bytes to a u16 for x. let x_bytes: [u8; 2] = buffer[0..2].try_into().ok()?; - // Try to convert the next 2 bytes to a u16 for y. let y_bytes: [u8; 2] = buffer[2..4].try_into().ok()?; - - // Reconstruct the u16 values from their little-endian byte representation. let x = u16::from_le_bytes(x_bytes); let y = u16::from_le_bytes(y_bytes); - Some((x, y)) } fn main() { - - - let x:Arc> = Arc::new(Mutex::new(0)); + let x: Arc> = Arc::new(Mutex::new(0)); let x_thread = x.clone(); - - let y:Arc>= Arc::new(Mutex::new(0)); + + let y: Arc> = Arc::new(Mutex::new(0)); let y_thread = y.clone(); - let circle = Box::new(Circle::new(x,y)); - let mut images:Vec> = vec![ -// Box::new(BouncingImage::new("images/unicorn_cc.png", 13, -10, 1, -1, -1)), -// Box::new(BouncingImage::new("images/windows_logo.png", -8, 3, 2, -1, -1)), -// Box::new(BouncingImage::new("images/spade.png", 32, -12, 1, 0, 0)), -// Box::new(BouncingImage::new("images/dvdvideo.png", 20, 6, 5, 1000, 800)), -// Box::new(BouncingImage::new("images/hackaday.png", 40, 18, 3, 500, 800)), - circle + let set: Arc> = Arc::new(Mutex::new(false)); + let set_thread = set.clone(); + + let circle = Box::new(Circle::new(x, y, set)); + let mut images: Vec> = vec![ + // Velocities are in pixels per second. + Box::new(BouncingImage::new("images/unicorn_cc.png", 30.0, -30.0, -1, -1)), + Box::new(BouncingImage::new("images/windows_logo.png", -20.0, 20.0, -1, -1)), + Box::new(BouncingImage::new("images/spade.png", 90.0, -60.0, 0, 0)), + Box::new(BouncingImage::new("images/dvdvideo.png", 60.0, 18.0, 1000, 800)), + Box::new(BouncingImage::new("images/hackaday.png", 40.0, 50.0, 500, 800)), + circle, ]; let mut display = Display::new(DISPLAY_HOST, DISPLAY_PORT); - let mut frame_counter: u32 = 0; + + // Spawn a UDP listener thread for receiving coordinates thread::spawn(move || { - let bind_address = format!("0.0.0.0:12345"); - let socket = Socket::new(Domain::IPV4, Type::DGRAM, None).unwrap(); - socket.set_reuse_address(true).unwrap(); - //socket.set_nonblocking(true).unwrap(); - socket.join_multicast_v4(&Ipv4Addr::new(239, 1, 1, 1), &Ipv4Addr::new(0, 0, 0, 0)).unwrap(); - socket.bind(&"0.0.0.0:1234".parse::().unwrap().into()).unwrap(); - // Bind the UDP socket to the specified address and port. - let socket: UdpSocket = socket.into(); + let bind_address = "0.0.0.0:12345"; + let socket = Socket::new(Domain::IPV4, Type::DGRAM, None).unwrap(); + socket.set_reuse_address(true).unwrap(); + socket + .bind(&"0.0.0.0:1234".parse::().unwrap().into()) + .unwrap(); + let socket: UdpSocket = socket.into(); - println!("Listening for UDP packets on {}", bind_address); + println!("Listening for UDP packets on {}", bind_address); - // Create a buffer to hold incoming data. 4 bytes for two u16 values. - let mut buf = [0u8; 4]; + let mut buf = [0u8; 4]; - loop { - // Wait for a packet to arrive. - match socket.recv_from(&mut buf) { - Ok((number_of_bytes, src_addr)) => { - println!("\nReceived {} bytes from {}", number_of_bytes, src_addr); + loop { + match socket.recv_from(&mut buf) { + Ok((number_of_bytes, src_addr)) => { + println!("\nReceived {} bytes from {}", number_of_bytes, src_addr); - // Ensure we received the correct number of bytes. - if number_of_bytes == 4 { - // Unpack the buffer into coordinates. - if let Some((x_rev, y_rev)) = unpack_coordinates(&buf) { - - println!("Received Coordinates: X = {}, Y = {}", x_rev, y_rev); let x_32:u32 = x_rev.into(); - let y_32:u32 = y_rev.into(); - *x_thread.lock().unwrap() = x_32; - *y_thread.lock().unwrap() = y_32; - - + if number_of_bytes == 4 { + if let Some((x_rev, y_rev)) = unpack_coordinates(&buf) { + println!("Received Coordinates: X = {}, Y = {}", x_rev, y_rev); + *x_thread.lock().unwrap() = x_rev.into(); + *y_thread.lock().unwrap() = y_rev.into(); + *set_thread.lock().unwrap() = true; + } else { + eprintln!("Error: Failed to unpack coordinate data."); + } } else { - // This case should ideally not be reached if number_of_bytes is 4. - eprintln!("Error: Failed to unpack coordinate data."); + eprintln!( + "Warning: Received packet with incorrect size ({} bytes). Expected 4.", + number_of_bytes + ); } - } else { - eprintln!( - "Warning: Received packet with incorrect size ({} bytes). Expected 4.", - number_of_bytes - ); + } + Err(e) => { + eprintln!("Error receiving data: {}", e); } } - Err(e) => { - eprintln!("Error receiving data: {}", e); - // Decide if you want to break the loop on an error. - // For a continuous server, you might just log and continue. - } } - } - }); - // display.blank_screen(); - let mut tick:u32 = 0; - loop { - for (i, bb) in images.iter_mut().enumerate() { - if bb.rate() > 0 && frame_counter % bb.rate() != 0 { - continue; + let num_drawables = images.len(); + + // Wrap drawables in Arc> for shared access: + // - Workers take read locks for draw(&self) -- parallel, no contention + // - Main takes write locks for update(&mut self) -- exclusive, 30fps only + let images: Vec>>> = images + .into_iter() + .map(|b| Arc::new(RwLock::new(b))) + .collect(); + + // Shared elapsed time for draw threads (as u32 bits of f32). + let elapsed_bits = Arc::new(AtomicU32::new(0)); + + // Barriers: main signals "go", workers draw, then signal "done". + let go_barrier = Arc::new(Barrier::new(num_drawables + 1)); + let done_barrier = Arc::new(Barrier::new(num_drawables + 1)); + + // Each worker gets its own PixelBuf wrapped in Mutex so the main + // thread can read the results after the done barrier. + let pixel_bufs: Vec>> = (0..num_drawables) + .map(|_| Arc::new(Mutex::new(PixelBuf::new(60_000)))) + .collect(); + + // Spawn persistent worker threads. + for i in 0..num_drawables { + let img = Arc::clone(&images[i]); + let buf = Arc::clone(&pixel_bufs[i]); + let elapsed = Arc::clone(&elapsed_bits); + let go = Arc::clone(&go_barrier); + let done = Arc::clone(&done_barrier); + + thread::spawn(move || { + loop { + go.wait(); + + let e = f32::from_bits(elapsed.load(Ordering::Relaxed)); + let drawable = img.read().unwrap(); + let mut pb = buf.lock().unwrap(); + pb.clear(); + drawable.draw(&mut pb, e); + drop(pb); + drop(drawable); + + done.wait(); } - bb.draw_and_move(&mut display,tick); - } + }); + } + + let start_time = Instant::now(); + let mut last_update = Instant::now(); + let update_interval = Duration::from_secs_f32(1.0 / 30.0); + + loop { + let now = Instant::now(); + let elapsed = (now - start_time).as_secs_f32(); + + // Update positions at 30 fps (main thread only, workers are idle). + if now - last_update >= update_interval { + let dt = (now - last_update).as_secs_f32(); + for img in &images { + img.write().unwrap().update(dt, elapsed); + } + last_update = now; + } + + // Publish elapsed time and signal workers to draw. + elapsed_bits.store(elapsed.to_bits(), Ordering::Relaxed); + go_barrier.wait(); + + // Wait for all workers to finish drawing. + done_barrier.wait(); + + // Merge all pixel bufs into the single Display for efficient packing. + for pb in &pixel_bufs { + let buf = pb.lock().unwrap(); + display.write_raw_pixels(&buf.pixels); + } display.flush_frame(); - tick+=1; - frame_counter += 1; - - // A small delay to control the frame rate - //std::thread::sleep(Duration::from_millis(16)); } } - diff --git a/src/pixel.rs b/src/pixel.rs new file mode 100644 index 0000000..ab3a3af --- /dev/null +++ b/src/pixel.rs @@ -0,0 +1,22 @@ +#[derive(Copy, Clone)] +pub struct Pixel { + pub x: u16, + pub y: u16, + pub r: u8, + pub g: u8, + pub b: u8 +} +impl Pixel { + pub fn format0(&self) -> [u8; 7] { + + let buf:[u8; 7] = [ self.x as u8 + , (self.x >> 8) as u8 + ,self.y as u8 + ,(self.y >> 8) as u8 + , self.r + , self.g + , self.b]; + return buf; + } + +} \ No newline at end of file diff --git a/src/pixel_buf.rs b/src/pixel_buf.rs new file mode 100644 index 0000000..91606ec --- /dev/null +++ b/src/pixel_buf.rs @@ -0,0 +1,37 @@ +/// A lightweight per-thread pixel accumulator. +/// +/// Stores already-adjusted 7-byte pixel entries `[x_lo, x_hi, y_lo, y_hi, r, g, b]` +/// ready to be merged into the main Display buffer. No packets, no flushing, +/// no queue -- just a growable buffer that each draw thread fills independently. +use crate::pixel::Pixel; +pub struct PixelBuf { + pub pixels: Vec<[u8;7]>, +} + +impl PixelBuf { + /// Creates a new empty PixelBuf with the given initial capacity. + pub fn new(capacity: usize) -> Self { + PixelBuf { + pixels: Vec::with_capacity(capacity), + } + } + + /// Clears the buffer for reuse without deallocating. + #[inline] + pub fn clear(&mut self) { + self.pixels.clear(); + } + + #[inline] + pub fn set_pixel(&mut self, pixel: Pixel) { + self.pixels.push(pixel.format0()); + } + + #[inline] + pub fn set_pixels(&mut self, pixels: Vec) { + self.pixels.reserve(pixels.len()); + for pixel in pixels { + self.pixels.push(pixel.format0()) + } + } +} diff --git a/src/png_data.rs b/src/png_data.rs new file mode 100644 index 0000000..aa7398f --- /dev/null +++ b/src/png_data.rs @@ -0,0 +1,55 @@ +use std::fs::File; +use std::io::BufReader; +use crate::pixel::Pixel; +/// Represents the data decoded from a PNG file, with pre-computed pixel data. +pub struct PngData { + pub width: u32, + pub height: u32, + pub packed_pixels: Vec, +} +impl PngData { + + pub fn open(path: &str) -> Result { + let file = File::open(path).expect("Failed to open PNG file"); + let decoder = png::Decoder::new(BufReader::new(file)); + let mut reader = decoder.read_info()?; + let mut buf = vec![0; reader.output_buffer_size()]; + let info = reader.next_frame(&mut buf)?; + + let width = info.width; + let height = info.height; + + let mut packed_pixels = Vec::new(); + for sy in 0..height { + for sx in 0..width { + let index = (sy * width + sx) as usize * 4; + let a = buf[index + 3]; + if a > 0 { + let r = buf[index]; + let g = buf[index + 1]; + let b = buf[index + 2]; + let x = sx as u16; + let y = sy as u16; + packed_pixels.push(Pixel { + x, + y, + r, + g, + b + }) + } + } + } + + eprintln!( + "[png] {}: {}x{}, {} opaque pixels pre-computed", + path, width, height, packed_pixels.len() + ); + + Ok(PngData { + width, + height, + packed_pixels, + }) + } +} diff --git a/src/raw_socket.rs b/src/raw_socket.rs new file mode 100644 index 0000000..39d0dab --- /dev/null +++ b/src/raw_socket.rs @@ -0,0 +1,426 @@ +//! Raw AF_PACKET sender that constructs complete Ethernet + IP + UDP frames +//! in userspace and sends them directly to the NIC, bypassing the kernel's +//! UDP/IP stack entirely. + +use std::fs; +use std::io::{self, BufRead, BufReader}; +use std::net::{Ipv4Addr, ToSocketAddrs, UdpSocket}; +use std::os::unix::io::AsRawFd; + +/// Size of Ethernet + IPv4 + UDP headers combined. +const HEADER_SIZE: usize = 14 + 20 + 8; // 42 bytes + +/// PACKET_QDISC_BYPASS socket option (skip the qdisc layer). +const PACKET_QDISC_BYPASS: libc::c_int = 20; + +/// A raw AF_PACKET sender that constructs full Ethernet frames. +pub struct RawSender { + fd: libc::c_int, + ifindex: i32, + /// Pre-built header template (Ethernet + IP + UDP). + /// Only the IP total_length, IP checksum, and UDP length fields + /// need to be patched per packet. + header_template: [u8; HEADER_SIZE], +} + +impl RawSender { + /// Creates a new RawSender. + /// + /// Auto-discovers the network interface, MAC addresses, and source IP + /// from the routing table and ARP cache. + /// + /// Requires `CAP_NET_RAW` or root. + pub fn new(dst_host: &str, dst_port: u16, src_port: u16) -> io::Result { + // Resolve destination IP. + let dst_ip: Ipv4Addr = dst_host + .parse() + .or_else(|_| { + (dst_host, dst_port) + .to_socket_addrs()? + .find_map(|a| match a { + std::net::SocketAddr::V4(v4) => Some(*v4.ip()), + _ => None, + }) + .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "Could not resolve host")) + })?; + + // Discover source IP and interface via a temporary connected UDP socket. + let probe = UdpSocket::bind("0.0.0.0:0")?; + probe.connect((dst_ip, dst_port))?; + let src_ip: Ipv4Addr = match probe.local_addr()? { + std::net::SocketAddr::V4(a) => *a.ip(), + _ => return Err(io::Error::new(io::ErrorKind::Other, "Not IPv4")), + }; + + // Discover interface name and index from routing. + let (ifname, ifindex) = discover_interface(dst_ip)?; + eprintln!("[raw] interface: {} (index {})", ifname, ifindex); + + // Read source MAC from sysfs. + let src_mac = read_mac_from_sysfs(&ifname)?; + eprintln!("[raw] src MAC: {:02x}:{:02x}:{:02x}:{:02x}:{:02x}:{:02x}", + src_mac[0], src_mac[1], src_mac[2], src_mac[3], src_mac[4], src_mac[5]); + + // Look up destination MAC from ARP cache. + let dst_mac = lookup_arp(dst_ip, &ifname)?; + eprintln!("[raw] dst MAC: {:02x}:{:02x}:{:02x}:{:02x}:{:02x}:{:02x}", + dst_mac[0], dst_mac[1], dst_mac[2], dst_mac[3], dst_mac[4], dst_mac[5]); + + eprintln!("[raw] {}:{} -> {}:{}", src_ip, src_port, dst_ip, dst_port); + + // Create AF_PACKET raw socket. + let fd = unsafe { + libc::socket( + libc::AF_PACKET, + libc::SOCK_RAW, + (libc::ETH_P_IP as u16).to_be() as i32, + ) + }; + if fd < 0 { + return Err(io::Error::last_os_error()); + } + + // Bind to the interface. + let mut sll: libc::sockaddr_ll = unsafe { std::mem::zeroed() }; + sll.sll_family = libc::AF_PACKET as u16; + sll.sll_protocol = (libc::ETH_P_IP as u16).to_be(); + sll.sll_ifindex = ifindex; + let ret = unsafe { + libc::bind( + fd, + &sll as *const libc::sockaddr_ll as *const libc::sockaddr, + std::mem::size_of::() as libc::socklen_t, + ) + }; + if ret < 0 { + unsafe { libc::close(fd); } + return Err(io::Error::last_os_error()); + } + + // Skip qdisc for lower latency and higher throughput. + let bypass: libc::c_int = 1; + unsafe { + libc::setsockopt( + fd, + libc::SOL_PACKET, + PACKET_QDISC_BYPASS, + &bypass as *const libc::c_int as *const libc::c_void, + std::mem::size_of::() as libc::socklen_t, + ); + } + + // Increase send buffer. + let sndbuf: libc::c_int = 25 * 1024 * 1024; + unsafe { + libc::setsockopt( + fd, + libc::SOL_SOCKET, + libc::SO_SNDBUF, + &sndbuf as *const libc::c_int as *const libc::c_void, + std::mem::size_of::() as libc::socklen_t, + ); + } + + // Build the header template. + let header_template = build_header_template( + &src_mac, &dst_mac, src_ip, dst_ip, src_port, dst_port, + ); + + Ok(RawSender { + fd, + ifindex, + header_template, + }) + } + + /// Sends a batch of payloads as raw Ethernet frames using sendmmsg. + /// + /// `frames` is a pre-allocated buffer of frame data. + /// Each frame is HEADER_SIZE + payload bytes. + /// Returns the number of frames successfully sent. + pub fn send_batch(&self, payloads: &[&[u8]], frame_buf: &mut Vec) -> io::Result { + let count = payloads.len(); + if count == 0 { + return Ok(0); + } + + // Build all frames into the contiguous buffer. + frame_buf.clear(); + let mut offsets: Vec<(usize, usize)> = Vec::with_capacity(count); + + for payload in payloads { + let frame_start = frame_buf.len(); + let total_len = (20 + 8 + payload.len()) as u16; // IP total length + let udp_len = (8 + payload.len()) as u16; + + // Copy header template. + frame_buf.extend_from_slice(&self.header_template); + + // Patch IP total length (bytes 16-17 of the frame = offset 2-3 in IP header). + let ip_start = frame_start + 14; + frame_buf[ip_start + 2] = (total_len >> 8) as u8; + frame_buf[ip_start + 3] = total_len as u8; + + // Patch IP header checksum (bytes 10-11 of IP header). + // Zero the checksum field first, then compute. + frame_buf[ip_start + 10] = 0; + frame_buf[ip_start + 11] = 0; + let cksum = ip_checksum(&frame_buf[ip_start..ip_start + 20]); + frame_buf[ip_start + 10] = (cksum >> 8) as u8; + frame_buf[ip_start + 11] = cksum as u8; + + // Patch UDP length (bytes 4-5 of UDP header). + let udp_start = frame_start + 14 + 20; + frame_buf[udp_start + 4] = (udp_len >> 8) as u8; + frame_buf[udp_start + 5] = udp_len as u8; + + // UDP checksum = 0 (optional for IPv4). + frame_buf[udp_start + 6] = 0; + frame_buf[udp_start + 7] = 0; + + // Append payload. + frame_buf.extend_from_slice(payload); + + let frame_end = frame_buf.len(); + offsets.push((frame_start, frame_end)); + } + + // Build iovec and mmsghdr arrays for sendmmsg. + let mut iovecs: Vec = Vec::with_capacity(count); + let mut msghdrs: Vec = Vec::with_capacity(count); + + // Destination sockaddr_ll. + let mut sll: libc::sockaddr_ll = unsafe { std::mem::zeroed() }; + sll.sll_family = libc::AF_PACKET as u16; + sll.sll_protocol = (libc::ETH_P_IP as u16).to_be(); + sll.sll_ifindex = self.ifindex; + sll.sll_halen = 6; + // dst MAC is in the first 6 bytes of each frame, but sockaddr_ll + // also needs it. Copy from the template. + sll.sll_addr[..6].copy_from_slice(&self.header_template[..6]); + + for &(start, end) in &offsets { + iovecs.push(libc::iovec { + iov_base: frame_buf[start..end].as_ptr() as *mut libc::c_void, + iov_len: end - start, + }); + } + + for iov in &iovecs { + let mut mhdr: libc::mmsghdr = unsafe { std::mem::zeroed() }; + mhdr.msg_hdr.msg_iov = iov as *const libc::iovec as *mut libc::iovec; + mhdr.msg_hdr.msg_iovlen = 1; + mhdr.msg_hdr.msg_name = &sll as *const libc::sockaddr_ll as *mut libc::c_void; + mhdr.msg_hdr.msg_namelen = std::mem::size_of::() as u32; + msghdrs.push(mhdr); + } + + let ret = unsafe { + libc::sendmmsg( + self.fd, + msghdrs.as_mut_ptr(), + count as libc::c_uint, + 0, + ) + }; + + if ret < 0 { + Err(io::Error::last_os_error()) + } else { + Ok(ret as usize) + } + } + + /// Returns the header size so callers can compute total frame sizes. + pub fn header_size(&self) -> usize { + HEADER_SIZE + } +} + +impl Drop for RawSender { + fn drop(&mut self) { + unsafe { + libc::close(self.fd); + } + } +} + +/// Builds the 42-byte header template: Ethernet (14) + IP (20) + UDP (8). +/// +/// Fields that vary per packet (IP total_length, IP checksum, UDP length) +/// are filled with placeholder values and must be patched before sending. +fn build_header_template( + src_mac: &[u8; 6], + dst_mac: &[u8; 6], + src_ip: Ipv4Addr, + dst_ip: Ipv4Addr, + src_port: u16, + dst_port: u16, +) -> [u8; HEADER_SIZE] { + let mut h = [0u8; HEADER_SIZE]; + + // -- Ethernet header (14 bytes) -- + h[0..6].copy_from_slice(dst_mac); + h[6..12].copy_from_slice(src_mac); + h[12] = 0x08; // EtherType: IPv4 + h[13] = 0x00; + + // -- IPv4 header (20 bytes, no options) -- + let ip = &mut h[14..34]; + ip[0] = 0x45; // Version 4, IHL 5 (20 bytes) + ip[1] = 0x00; // DSCP/ECN + // ip[2..4] = total length (patched per packet) + ip[4] = 0x00; // Identification + ip[5] = 0x00; + ip[6] = 0x40; // Flags: Don't Fragment + ip[7] = 0x00; // Fragment offset + ip[8] = 64; // TTL + ip[9] = 17; // Protocol: UDP + // ip[10..12] = header checksum (patched per packet) + ip[12..16].copy_from_slice(&src_ip.octets()); + ip[16..20].copy_from_slice(&dst_ip.octets()); + + // -- UDP header (8 bytes) -- + let udp = &mut h[34..42]; + udp[0] = (src_port >> 8) as u8; + udp[1] = src_port as u8; + udp[2] = (dst_port >> 8) as u8; + udp[3] = dst_port as u8; + // udp[4..6] = length (patched per packet) + // udp[6..8] = checksum (set to 0, optional for IPv4) + + h +} + +/// Computes the IPv4 header checksum (one's complement of the one's +/// complement sum of all 16-bit words in the header). +fn ip_checksum(header: &[u8]) -> u16 { + let mut sum: u32 = 0; + for i in (0..header.len()).step_by(2) { + let word = if i + 1 < header.len() { + ((header[i] as u32) << 8) | (header[i + 1] as u32) + } else { + (header[i] as u32) << 8 + }; + sum += word; + } + // Fold carry bits. + while sum > 0xFFFF { + sum = (sum & 0xFFFF) + (sum >> 16); + } + !(sum as u16) +} + +/// Discovers the outbound interface name and index for a destination IP. +fn discover_interface(dst_ip: Ipv4Addr) -> io::Result<(String, i32)> { + // Use a temporary UDP socket to discover the interface. + let sock = UdpSocket::bind("0.0.0.0:0")?; + sock.connect((dst_ip, 80))?; + + // Get the interface index via SO_BINDTODEVICE or by reading /proc/net/route. + // Simpler: use the socket's bound address to find the matching interface. + let local_ip = match sock.local_addr()? { + std::net::SocketAddr::V4(a) => *a.ip(), + _ => return Err(io::Error::new(io::ErrorKind::Other, "Not IPv4")), + }; + + // Scan /proc/net/if_inet6 or /proc/net/fib_trie... actually, simplest + // is to iterate /sys/class/net/*/address and check IPs. + // But even simpler: parse `ip route get` output... that's ugly. + // + // Let's iterate interfaces and match by IP. + let entries = fs::read_dir("/sys/class/net/")?; + for entry in entries { + let entry = entry?; + let ifname = entry.file_name().to_string_lossy().into_owned(); + + // Get the interface index. + let ifindex_path = format!("/sys/class/net/{}/ifindex", ifname); + let ifindex_str = match fs::read_to_string(&ifindex_path) { + Ok(s) => s, + Err(_) => continue, + }; + let ifindex: i32 = match ifindex_str.trim().parse() { + Ok(i) => i, + Err(_) => continue, + }; + + // Check if this interface has the local IP using ioctl. + let fd = sock.as_raw_fd(); + let mut ifr: libc::ifreq = unsafe { std::mem::zeroed() }; + let name_bytes = ifname.as_bytes(); + let copy_len = name_bytes.len().min(libc::IFNAMSIZ - 1); + unsafe { + std::ptr::copy_nonoverlapping( + name_bytes.as_ptr(), + ifr.ifr_name.as_mut_ptr() as *mut u8, + copy_len, + ); + } + + let ret = unsafe { + libc::ioctl(fd, libc::SIOCGIFADDR, &mut ifr) + }; + if ret < 0 { + continue; + } + + let addr = unsafe { &*(&ifr.ifr_ifru as *const _ as *const libc::sockaddr_in) }; + let iface_ip = Ipv4Addr::from(u32::from_be(addr.sin_addr.s_addr)); + + if iface_ip == local_ip { + return Ok((ifname, ifindex)); + } + } + + Err(io::Error::new( + io::ErrorKind::NotFound, + format!("Could not find interface for IP {}", local_ip), + )) +} + +/// Reads the MAC address of a network interface from sysfs. +fn read_mac_from_sysfs(ifname: &str) -> io::Result<[u8; 6]> { + let path = format!("/sys/class/net/{}/address", ifname); + let mac_str = fs::read_to_string(path)?; + parse_mac(mac_str.trim()) +} + +/// Looks up the MAC address for an IP in the kernel ARP cache. +fn lookup_arp(ip: Ipv4Addr, ifname: &str) -> io::Result<[u8; 6]> { + let file = fs::File::open("/proc/net/arp")?; + let reader = BufReader::new(file); + + let ip_str = ip.to_string(); + + for line in reader.lines().skip(1) { + let line = line?; + let fields: Vec<&str> = line.split_whitespace().collect(); + if fields.len() >= 6 && fields[0] == ip_str && fields[5] == ifname { + return parse_mac(fields[3]); + } + } + + Err(io::Error::new( + io::ErrorKind::NotFound, + format!("No ARP entry for {} on {}", ip, ifname), + )) +} + +/// Parses a MAC address string like "50:6b:4b:c3:fa:9c" into 6 bytes. +fn parse_mac(s: &str) -> io::Result<[u8; 6]> { + let parts: Vec<&str> = s.split(':').collect(); + if parts.len() != 6 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Invalid MAC: {}", s), + )); + } + let mut mac = [0u8; 6]; + for (i, part) in parts.iter().enumerate() { + mac[i] = u8::from_str_radix(part, 16).map_err(|_| { + io::Error::new(io::ErrorKind::InvalidData, format!("Invalid MAC byte: {}", part)) + })?; + } + Ok(mac) +}