From 13a4810cbf6e4c3f6500690c835e7e3832b9d868 Mon Sep 17 00:00:00 2001 From: Ditin2 Date: Sat, 27 Dec 2025 14:09:05 +0100 Subject: [PATCH] b --- Cargo.lock | 56 +++++++- Cargo.toml | 2 + src/args.rs | 30 +++- src/image_manager.rs | 326 ++++++++++++++++++++++++++++-------------- src/main.rs | 143 +++++++++++++++++- src/painter/handle.rs | 13 +- 6 files changed, 444 insertions(+), 126 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 031f036..d8fdf16 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -578,6 +578,15 @@ dependencies = [ "objc2 0.4.1", ] +[[package]] +name = "block2" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdeb9d870516001442e364c5220d3574d2da8dc765554b4a617230d33fa58ef5" +dependencies = [ + "objc2 0.6.3", +] + [[package]] name = "blocking" version = "1.6.2" @@ -716,6 +725,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fd16c4719339c4530435d38e511904438d07cce7950afa3718a84ac36c10e89e" +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + [[package]] name = "cgl" version = "0.3.2" @@ -989,6 +1004,17 @@ dependencies = [ "typenum", ] +[[package]] +name = "ctrlc" +version = "3.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73736a89c4aff73035ba2ed2e565061954da00d4970fc9ac25dcc85a2a20d790" +dependencies = [ + "dispatch2", + "nix 0.30.1", + "windows-sys 0.61.2", +] + [[package]] name = "cursor-icon" version = "1.2.0" @@ -1039,6 +1065,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "89a09f22a6c6069a18470eb92d2298acf25463f14256d24778e1230d789a2aec" dependencies = [ "bitflags 2.10.0", + "block2 0.6.2", + "libc", "objc2 0.6.3", ] @@ -1529,7 +1557,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "18fcd4ae4e86d991ad1300b8f57166e5be0c95ef1f63f3f5b827f8a164548746" dependencies = [ "bitflags 2.10.0", - "cfg_aliases", + "cfg_aliases 0.1.1", "cgl", "core-foundation", "dispatch", @@ -1552,7 +1580,7 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ebcdfba24f73b8412c5181e56f092b5eff16671c514ce896b258a0a64bd7735" dependencies = [ - "cfg_aliases", + "cfg_aliases 0.1.1", "glutin", "raw-window-handle 0.5.2", "winit", @@ -2207,6 +2235,18 @@ dependencies = [ "memoffset 0.7.1", ] +[[package]] +name = "nix" +version = "0.30.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6" +dependencies = [ + "bitflags 2.10.0", + "cfg-if", + "cfg_aliases 0.2.1", + "libc", +] + [[package]] name = "nohash-hasher" version = "0.2.0" @@ -2530,9 +2570,11 @@ version = "0.1.0" dependencies = [ "bufstream", "clap", + "ctrlc", "eframe", "glob", "image 0.23.14", + "libc", "num_cpus", "rayon", "regex", @@ -3726,7 +3768,7 @@ checksum = "cbd7311dbd2abcfebaabf1841a2824ed7c8be443a0f29166e5d3c6a53a762c01" dependencies = [ "arrayvec", "cfg-if", - "cfg_aliases", + "cfg_aliases 0.1.1", "js-sys", "log", "parking_lot", @@ -3751,7 +3793,7 @@ dependencies = [ "arrayvec", "bit-vec", "bitflags 2.10.0", - "cfg_aliases", + "cfg_aliases 0.1.1", "codespan-reporting", "indexmap", "log", @@ -3778,7 +3820,7 @@ dependencies = [ "arrayvec", "ash", "bitflags 2.10.0", - "cfg_aliases", + "cfg_aliases 0.1.1", "core-graphics-types", "glow", "glutin_wgl_sys", @@ -4150,7 +4192,7 @@ dependencies = [ "bitflags 2.10.0", "bytemuck", "calloop 0.12.4", - "cfg_aliases", + "cfg_aliases 0.1.1", "core-foundation", "core-graphics", "cursor-icon", @@ -4337,7 +4379,7 @@ dependencies = [ "futures-sink", "futures-util", "hex", - "nix", + "nix 0.26.4", "once_cell", "ordered-stream", "rand", diff --git a/Cargo.toml b/Cargo.toml index be07298..84aa9f2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,9 +16,11 @@ categories = [ [dependencies] bufstream = "0.1" clap = { version = "4.4", features = [ "derive" ] } +ctrlc = "3.4" eframe = "0.26" glob = "0.3" image = "0.23" +libc = "0.2" num_cpus = "1.13.1" regex = "1.5" rayon = "1.5.1" diff --git a/src/args.rs b/src/args.rs index 458313a..651b288 100644 --- a/src/args.rs +++ b/src/args.rs @@ -10,14 +10,15 @@ pub struct Arguments { help: Option, /// The host to pwn "host:port" - host: String, + #[arg(required_unless_present = "config")] + host: Option, /// Image path(s) #[arg( short, long, value_name = "PATH", - required = true, + required_unless_present = "config", alias = "images", num_args(1..) )] @@ -37,6 +38,14 @@ pub struct Arguments { #[arg(short, value_name = "PIXELS", default_value_t = 0)] y: u16, + /// Path to a JSON config file + #[arg(long, value_name = "PATH")] + config: Option, + + /// Path to write JSON status updates + #[arg(long, value_name = "PATH")] + status: Option, + /// Tile images across the screen, ignoring offsets #[arg(short, long)] tile: bool, @@ -72,7 +81,22 @@ impl ArgHandler { /// Get the host property. pub fn host(&self) -> &str { - self.data.host.as_str() + self.data.host.as_deref().expect("Host is required") + } + + /// Get the optional host property. + pub fn host_optional(&self) -> Option<&str> { + self.data.host.as_deref() + } + + /// Get the config path. + pub fn config_path(&self) -> Option<&str> { + self.data.config.as_deref() + } + + /// Get the status path. + pub fn status_path(&self) -> Option<&str> { + self.data.status.as_deref() } /// Get the thread count. diff --git a/src/image_manager.rs b/src/image_manager.rs index 6f35ec9..598acb6 100644 --- a/src/image_manager.rs +++ b/src/image_manager.rs @@ -1,29 +1,48 @@ -use image::codecs::gif::GifDecoder; -use rayon::prelude::*; use glob::glob; -use std::fs::{self, File}; +use rayon::prelude::*; +use std::fs; use std::path::{Path, PathBuf}; +use std::process; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::thread::sleep; use std::time::Duration; +use std::time::{SystemTime, UNIX_EPOCH}; use crate::pix::canvas::Canvas; use image::imageops::FilterType; -use image::{AnimationDecoder, DynamicImage, RgbaImage}; +use image::{DynamicImage, RgbaImage}; + +enum ImageMode { + Direct { size: (u16, u16) }, + Tiled { tile_size: (u16, u16), canvas_size: (u16, u16) }, +} /// A manager that manages all images to print. pub struct ImageManager { - /// Image frames and their preferred delay. - images: Vec<(DynamicImage, Option)>, + frames: Vec, + mode: ImageMode, + preprocessed: bool, + cache_dir: Option, + cached_single: Option, /// Define whether the first image has been drawn first: bool, - index: isize, + index: usize, } impl ImageManager { /// Intantiate the image manager. - pub fn from(images: Vec<(DynamicImage, Option)>) -> ImageManager { + fn from( + frames: Vec, + mode: ImageMode, + preprocessed: bool, + cache_dir: Option, + ) -> ImageManager { ImageManager { - images, + frames, + mode, + preprocessed, + cache_dir, + cached_single: None, first: false, index: 0, } @@ -35,19 +54,20 @@ impl ImageManager { println!("Load and process {} path(s)...", paths.len()); // Load the images from the paths - let image_manager = ImageManager::from( - paths - .par_iter() - .flat_map(|path| load_image(path, size)) - .collect(), - ); + let frames: Vec = paths + .par_iter() + .flat_map(|path| load_image_frames(path)) + .collect(); + let mode = ImageMode::Direct { size }; + let (frames, preprocessed, cache_dir) = preprocess_frames(&frames, &mode); + let image_manager = ImageManager::from(frames, mode, preprocessed, cache_dir); // TODO: process the image slices // We succeeded println!( "Loaded {} frame(s) from {} path(s).", - image_manager.images.len(), + image_manager.frames.len(), paths.len() ); @@ -63,16 +83,20 @@ impl ImageManager { // Show a status message println!("Load and process {} path(s)...", paths.len()); - let image_manager = ImageManager::from( - paths - .par_iter() - .flat_map(|path| load_image_tiled(path, tile_size, canvas_size)) - .collect(), - ); + let frames: Vec = paths + .par_iter() + .flat_map(|path| load_image_frames(path)) + .collect(); + let mode = ImageMode::Tiled { + tile_size, + canvas_size, + }; + let (frames, preprocessed, cache_dir) = preprocess_frames(&frames, &mode); + let image_manager = ImageManager::from(frames, mode, preprocessed, cache_dir); println!( "Loaded {} frame(s) from {} path(s).", - image_manager.images.len(), + image_manager.frames.len(), paths.len() ); @@ -84,7 +108,7 @@ impl ImageManager { /// Returns the desired duration for othis frame. pub fn tick(&mut self, canvas: &mut Canvas) -> Option { // Get the image index bound - let bound = self.images.len(); + let bound = self.frames.len(); // Just return if the bound is one, as nothing should be updated if self.first && bound == 1 { @@ -92,10 +116,10 @@ impl ImageManager { } // Get the image to use - let (image, duration) = &mut self.images[self.index as usize % bound]; + let mut image = self.load_frame(self.index % bound); // Update the image on the canvas - canvas.update_image(image); + canvas.update_image(&mut image); // Increase the index self.index += 1; @@ -103,7 +127,7 @@ impl ImageManager { // We have rendered the first image self.first = true; - *duration + None } /// Start working in the image manager. @@ -111,53 +135,172 @@ impl ImageManager { /// This will start walking through all image frames, /// and pushes each frame to all painters, /// with the specified frames per second. - pub fn work(&mut self, canvas: &mut Canvas, fps: u32) { - loop { - // Determine duration to wait, use frame direction and fall back to FPS - let frame_delay = self.tick(canvas); - let delay = frame_delay - .unwrap_or_else(|| Duration::from_millis((1000f32 / (fps as f32)) as u64)); + pub fn work( + &mut self, + canvas: &mut Canvas, + fps: u32, + shutdown: &AtomicBool, + mut on_frame: F, + ) + where + F: FnMut(usize, usize), + { + let fps = fps.max(1); + let frame_ms = 1000f64 / (fps as f64); + let start = std::time::Instant::now(); + let mut next_tick = 0f64; - // Sleep until we need to show the next image - sleep(delay); + loop { + if shutdown.load(Ordering::SeqCst) { + break; + } + + let elapsed_ms = start.elapsed().as_secs_f64() * 1000.0; + let target_index = (elapsed_ms / frame_ms).floor() as usize; + if target_index != self.index { + self.index = target_index; + self.tick(canvas); + let total = self.frames.len(); + let current = (self.index % total) + 1; + on_frame(current, total); + } + + next_tick += frame_ms; + let sleep_ms = (next_tick - elapsed_ms).max(1.0) as u64; + sleep(Duration::from_millis(sleep_ms)); + } + } + + pub fn frame_count(&self) -> usize { + self.frames.len() + } + + fn load_frame(&mut self, index: usize) -> DynamicImage { + if self.frames.len() == 1 { + if let Some(image) = self.cached_single.as_ref() { + return image.clone(); + } + } + + let path = &self.frames[index]; + let image = if self.preprocessed { + image::open(path).unwrap() + } else { + match self.mode { + ImageMode::Direct { size } => image::open(path) + .unwrap() + .resize_exact(size.0 as u32, size.1 as u32, FilterType::Gaussian), + ImageMode::Tiled { + tile_size, + canvas_size, + } => tile_image( + image::open(path).unwrap(), + (tile_size.0 as u32, tile_size.1 as u32), + (canvas_size.0 as u32, canvas_size.1 as u32), + ), + } + }; + + if self.frames.len() == 1 { + self.cached_single = Some(image.clone()); + } + + image + } +} + +static CACHE_COUNTER: AtomicUsize = AtomicUsize::new(0); + +fn preprocess_frames( + frames: &[PathBuf], + mode: &ImageMode, +) -> (Vec, bool, Option) { + if frames.len() <= 1 { + return (frames.to_vec(), false, None); + } + + let counter = CACHE_COUNTER.fetch_add(1, Ordering::SeqCst); + let stamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis(); + let cache_dir = std::env::temp_dir().join(format!( + "pixelpwnr_cache_{}_{}_{}", + process::id(), + stamp, + counter + )); + fs::create_dir_all(&cache_dir).expect("failed to create cache directory"); + match *mode { + ImageMode::Direct { size } => println!( + "Resizing {} frame(s) to {}x{} (cache: {})", + frames.len(), + size.0, + size.1, + cache_dir.display() + ), + ImageMode::Tiled { + tile_size, + canvas_size, + } => println!( + "Resizing {} frame(s) to tile {}x{} over {}x{} (cache: {})", + frames.len(), + tile_size.0, + tile_size.1, + canvas_size.0, + canvas_size.1, + cache_dir.display() + ), + } + + let total = frames.len(); + let progress = AtomicUsize::new(0); + let output_paths: Vec = frames + .par_iter() + .enumerate() + .map(|(index, path)| { + let output = cache_dir.join(format!("{:06}.png", index)); + let processed = match *mode { + ImageMode::Direct { size } => image::open(path) + .unwrap() + .resize_exact(size.0 as u32, size.1 as u32, FilterType::Gaussian), + ImageMode::Tiled { + tile_size, + canvas_size, + } => tile_image( + image::open(path).unwrap(), + (tile_size.0 as u32, tile_size.1 as u32), + (canvas_size.0 as u32, canvas_size.1 as u32), + ), + }; + processed + .save(&output) + .expect("failed to save cached frame"); + let done = progress.fetch_add(1, Ordering::SeqCst) + 1; + if done == total || done % 100 == 0 { + println!("Resized {done}/{total} frame(s)..."); + } + output + }) + .collect(); + + println!("Cached {} frame(s).", output_paths.len()); + + (output_paths, true, Some(cache_dir)) +} + +impl Drop for ImageManager { + fn drop(&mut self) { + if let Some(cache_dir) = &self.cache_dir { + match fs::remove_dir_all(cache_dir) { + Ok(()) => println!("Removed cache directory {}", cache_dir.display()), + Err(err) => eprintln!("Failed to remove cache directory: {err}"), + } } } } -/// Load the image at the given path, and size it correctly -fn load_image(path: &str, size: (u16, u16)) -> Vec<(DynamicImage, Option)> { - load_image_frames(path) - .into_iter() - .map(|(image, frame_delay)| { - ( - image.resize_exact(size.0 as u32, size.1 as u32, FilterType::Gaussian), - frame_delay, - ) - }) - .collect() -} - -fn load_image_tiled( - path: &str, - tile_size: (u16, u16), - canvas_size: (u16, u16), -) -> Vec<(DynamicImage, Option)> { - load_image_frames(path) - .into_iter() - .map(|(image, frame_delay)| { - ( - tile_image( - image, - (tile_size.0 as u32, tile_size.1 as u32), - (canvas_size.0 as u32, canvas_size.1 as u32), - ), - frame_delay, - ) - }) - .collect() -} - -fn load_image_frames(path: &str) -> Vec<(DynamicImage, Option)> { +fn load_image_frames(path: &str) -> Vec { if has_glob_pattern(path) { let mut entries: Vec = glob(path) .expect("failed to read frames glob") @@ -166,16 +309,11 @@ fn load_image_frames(path: &str) -> Vec<(DynamicImage, Option)> { .collect(); entries.sort_by(|a, b| a.file_name().cmp(&b.file_name())); - let mut frames = Vec::new(); - for entry in entries { - frames.extend(load_file_frames(&entry)); - } - - if frames.is_empty() { + if entries.is_empty() { panic!("The given frames glob matched no files"); } - return frames; + return entries; } let path = Path::new(&path); @@ -189,54 +327,24 @@ fn load_image_frames(path: &str) -> Vec<(DynamicImage, Option)> { .collect(); entries.sort_by(|a, b| a.file_name().cmp(&b.file_name())); - let mut frames = Vec::new(); - for entry in entries { - frames.extend(load_file_frames(&entry)); - } - - if frames.is_empty() { + if entries.is_empty() { panic!("The given frames directory is empty"); } - return frames; + return entries; } if !path.is_file() { panic!("The given path does not exist or is not a file"); } - load_file_frames(path) + vec![path.to_path_buf()] } fn has_glob_pattern(path: &str) -> bool { path.contains('*') || path.contains('?') || path.contains('[') } -fn load_file_frames(path: &Path) -> Vec<(DynamicImage, Option)> { - let extension = path - .extension() - .and_then(|e| e.to_str()) - .map(|e| e.to_lowercase()); - - match extension.as_deref() { - Some("gif") => GifDecoder::new(File::open(path).unwrap()) - .expect("failed to decode GIF file") - .into_frames() - .collect_frames() - .expect("failed to parse GIF frames") - .into_iter() - .map(|frame| { - let frame_delay = Duration::from(frame.delay()); - ( - DynamicImage::ImageRgba8(frame.into_buffer()), - Some(frame_delay), - ) - }) - .collect(), - _ => vec![(image::open(path).unwrap(), None)], - } -} - fn tile_image(image: DynamicImage, tile_size: (u32, u32), canvas_size: (u32, u32)) -> DynamicImage { let tile = image .resize_exact(tile_size.0, tile_size.1, FilterType::Gaussian) diff --git a/src/main.rs b/src/main.rs index 3caafd6..b3adf67 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,6 @@ mod args; mod color; +mod config; mod image_manager; mod painter; mod pix; @@ -8,9 +9,14 @@ mod rect; use std::io::Error; use args::ArgHandler; +use config::ProjectConfig; use image_manager::ImageManager; use pix::canvas::Canvas; use pix::client::Client; +use serde::Serialize; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; /// Main application entrypoint. fn main() { @@ -26,6 +32,13 @@ fn start(arg_handler: &ArgHandler) { // Start println!("Starting... (use CTRL+C to stop)"); + let shutdown = setup_shutdown_handler(); + + if let Some(config_path) = arg_handler.config_path() { + start_from_config(arg_handler, config_path, shutdown); + return; + } + // Gather facts about the host let screen_size = gather_host_facts(arg_handler).expect("Failed to gather facts about pixelflut server"); @@ -55,7 +68,135 @@ fn start(arg_handler: &ArgHandler) { }; // Start the work in the image manager, to walk through the frames - image_manager.work(&mut canvas, arg_handler.fps()); + image_manager.work(&mut canvas, arg_handler.fps(), &shutdown, |_current, _total| {}); +} + +#[derive(Clone, Serialize)] +struct StatusEntry { + name: String, + fps: f64, + current: usize, + total: usize, +} + +#[derive(Serialize)] +struct StatusFile { + entries: Vec, +} + +fn start_from_config(arg_handler: &ArgHandler, config_path: &str, shutdown: Arc) { + let config = load_config(config_path); + let host = arg_handler + .host_optional() + .map(|host| host.to_string()) + .unwrap_or_else(|| config.host.clone()); + if host.trim().is_empty() { + panic!("No host specified in CLI or config"); + } + + let default_threads = config.threads.unwrap_or_else(|| arg_handler.count()); + let binary = config.binary; + let flush = config.flush; + + let status_path = arg_handler.status_path().map(|path| path.to_string()); + let stats: Arc>> = Arc::new(std::sync::Mutex::new( + config + .entries + .iter() + .map(|entry| StatusEntry { + name: entry.name.clone(), + fps: 0.0, + current: 0, + total: 0, + }) + .collect(), + )); + if let Some(status_path) = status_path.clone() { + let stats = Arc::clone(&stats); + let shutdown = Arc::clone(&shutdown); + std::thread::spawn(move || { + while !shutdown.load(Ordering::SeqCst) { + if let Ok(locked) = stats.lock() { + let payload = StatusFile { + entries: locked.clone(), + }; + if let Ok(json) = serde_json::to_string_pretty(&payload) { + let _ = std::fs::write(&status_path, json); + } + } + std::thread::sleep(Duration::from_millis(500)); + } + }); + } + + let mut handles = Vec::new(); + for (index, entry) in config.entries.into_iter().enumerate() { + if !entry.enabled { + continue; + } + let host = host.clone(); + let threads = entry.threads.unwrap_or(default_threads); + let stats = Arc::clone(&stats); + let shutdown = Arc::clone(&shutdown); + handles.push(std::thread::spawn(move || { + let mut canvas = Canvas::new( + &host, + threads, + (entry.width, entry.height), + (entry.x, entry.y), + binary, + flush, + ); + let mut image_manager = ImageManager::load(&[entry.source.as_str()], (entry.width, entry.height)); + let total = image_manager.frame_count(); + if let Ok(mut locked) = stats.lock() { + if let Some(stat) = locked.get_mut(index) { + stat.total = total; + } + } + let mut frames = 0usize; + let mut last_report = Instant::now(); + image_manager.work(&mut canvas, entry.fps, &shutdown, |current, total| { + frames += 1; + if let Ok(mut locked) = stats.lock() { + if let Some(stat) = locked.get_mut(index) { + stat.current = current; + stat.total = total; + } + } + let elapsed = last_report.elapsed(); + if elapsed >= Duration::from_secs(1) { + let fps = frames as f64 / elapsed.as_secs_f64(); + if let Ok(mut locked) = stats.lock() { + if let Some(stat) = locked.get_mut(index) { + stat.fps = fps; + } + } + frames = 0; + last_report = Instant::now(); + } + }); + })); + } + + for handle in handles { + let _ = handle.join(); + } +} + +fn setup_shutdown_handler() -> Arc { + let shutdown = Arc::new(AtomicBool::new(false)); + let shutdown_handler = Arc::clone(&shutdown); + ctrlc::set_handler(move || { + shutdown_handler.store(true, Ordering::SeqCst); + }) + .expect("Failed to set CTRL+C handler"); + shutdown +} + +fn load_config(path: &str) -> ProjectConfig { + let data = std::fs::read_to_string(path).expect("Failed to read config file"); + serde_json::from_str::(&data).expect("Failed to parse config file") } /// Gather important facts about the host. diff --git a/src/painter/handle.rs b/src/painter/handle.rs index 1f345ae..3f65174 100644 --- a/src/painter/handle.rs +++ b/src/painter/handle.rs @@ -28,18 +28,19 @@ impl Handle { /// Push an image update. pub fn update_image(&self, full_image: &mut DynamicImage) { - // Crop the image to the area - let image = full_image.crop( + // Crop without mutating the source image so each painter sees full data. + let image = image::imageops::crop_imm( + full_image, self.area.x as u32, self.area.y as u32, self.area.w as u32, self.area.h as u32, - ); + ) + .to_image(); + let image = DynamicImage::ImageRgba8(image); // Push a new image to the thread // TODO: return this result - self.image_sender - .send(image) - .expect("Failed to send image update to painter"); + let _ = self.image_sender.send(image); } }