From a9a2895b62a6a28178f2222f39d42df1b7a8e07c Mon Sep 17 00:00:00 2001 From: timvisee Date: Thu, 28 Dec 2023 18:57:45 +0100 Subject: [PATCH] Add support for binary pixels PB --- src/arg_handler.rs | 37 +++++++++++++++++++++++++------- src/color.rs | 10 ++++----- src/main.rs | 7 +++--- src/pix/canvas.rs | 22 ++++++++++++------- src/pix/client.rs | 53 ++++++++++++++++++++++++++++++++++------------ 5 files changed, 91 insertions(+), 38 deletions(-) diff --git a/src/arg_handler.rs b/src/arg_handler.rs index 2ee3ddc..1ea0f72 100644 --- a/src/arg_handler.rs +++ b/src/arg_handler.rs @@ -23,7 +23,8 @@ impl<'a: 'b, 'b> ArgHandler<'a> { .help("The host to pwn \"host:port\"") .required(true) .index(1), - ).arg( + ) + .arg( Arg::with_name("image") .short("i") .long("image") @@ -34,7 +35,8 @@ impl<'a: 'b, 'b> ArgHandler<'a> { .multiple(true) .display_order(1) .takes_value(true), - ).arg( + ) + .arg( Arg::with_name("width") .short("w") .long("width") @@ -42,7 +44,8 @@ impl<'a: 'b, 'b> ArgHandler<'a> { .help("Draw width (def: screen width)") .display_order(2) .takes_value(true), - ).arg( + ) + .arg( Arg::with_name("height") .short("h") .long("height") @@ -50,21 +53,24 @@ impl<'a: 'b, 'b> ArgHandler<'a> { .help("Draw height (def: screen height)") .display_order(3) .takes_value(true), - ).arg( + ) + .arg( Arg::with_name("x") .short("x") .value_name("PIXELS") .help("Draw X offset (def: 0)") .display_order(4) .takes_value(true), - ).arg( + ) + .arg( Arg::with_name("y") .short("y") .value_name("PIXELS") .help("Draw Y offset (def: 0)") .display_order(5) .takes_value(true), - ).arg( + ) + .arg( Arg::with_name("count") .short("c") .long("count") @@ -74,7 +80,8 @@ impl<'a: 'b, 'b> ArgHandler<'a> { .help("Number of concurrent threads (def: CPUs)") .display_order(6) .takes_value(true), - ).arg( + ) + .arg( Arg::with_name("fps") .short("r") .long("fps") @@ -82,7 +89,16 @@ impl<'a: 'b, 'b> ArgHandler<'a> { .help("Frames per second with multiple images (def: 1)") .display_order(7) .takes_value(true), - ).get_matches(); + ) + .arg( + Arg::with_name("binary") + .short("b") + .long("binary") + .help("Use binary mode to set pixels (PB)") + .display_order(7) + .takes_value(false), + ) + .get_matches(); // Instantiate ArgHandler { matches } @@ -147,4 +163,9 @@ impl<'a: 'b, 'b> ArgHandler<'a> { .map(|fps| fps.parse::().expect("Invalid frames per second rate")) .unwrap_or(DEFAULT_IMAGE_FPS) } + + /// Whether to use binary mode. + pub fn binary(&self) -> bool { + self.matches.is_present("binary") + } } diff --git a/src/color.rs b/src/color.rs index 41d5f6e..1a5fdf7 100644 --- a/src/color.rs +++ b/src/color.rs @@ -3,10 +3,10 @@ /// Represents a color with RGB values from 0 to 255. #[derive(Copy, Clone)] pub struct Color { - r: u8, - g: u8, - b: u8, - a: u8, + pub(crate) r: u8, + pub(crate) g: u8, + pub(crate) b: u8, + pub(crate) a: u8, } impl Color { @@ -14,7 +14,7 @@ impl Color { /// /// The color channels must be between 0 and 255. pub fn from(r: u8, g: u8, b: u8, a: u8) -> Color { - Color { r, g, b, a} + Color { r, g, b, a } } /// Get a hexadecimal representation of the color, diff --git a/src/main.rs b/src/main.rs index dad0b7e..5a2d972 100644 --- a/src/main.rs +++ b/src/main.rs @@ -32,8 +32,8 @@ fn start<'a>(arg_handler: &ArgHandler<'a>) { println!("Starting... (use CTRL+C to stop)"); // Gather facts about the host - let screen_size = gather_host_facts(&arg_handler) - .expect("Failed to gather facts about pixelflut server"); + let screen_size = + gather_host_facts(&arg_handler).expect("Failed to gather facts about pixelflut server"); // Determine the size to use let size = arg_handler.size(Some(screen_size)); @@ -44,6 +44,7 @@ fn start<'a>(arg_handler: &ArgHandler<'a>) { arg_handler.count(), size, arg_handler.offset(), + arg_handler.binary(), ); // Load the image manager @@ -56,7 +57,7 @@ fn start<'a>(arg_handler: &ArgHandler<'a>) { /// Gather important facts about the host. fn gather_host_facts(arg_handler: &ArgHandler) -> Result<(u32, u32), Error> { // Set up a client, and get the screen size - let size = Client::connect(arg_handler.host().to_string())?.read_screen_size()?; + let size = Client::connect(arg_handler.host().to_string(), false)?.read_screen_size()?; // Print status println!("Gathered screen size: {}x{}", size.0, size.1); diff --git a/src/pix/canvas.rs b/src/pix/canvas.rs index 14bfe9e..984e8e0 100644 --- a/src/pix/canvas.rs +++ b/src/pix/canvas.rs @@ -1,6 +1,6 @@ use std::sync::mpsc::{self, Receiver, Sender}; -use std::thread::sleep; use std::thread; +use std::thread::sleep; use std::time::Duration; use image::DynamicImage; @@ -21,7 +21,13 @@ pub struct Canvas { impl Canvas { /// Create a new pixelflut canvas. - pub fn new(host: &str, painter_count: usize, size: (u32, u32), offset: (u32, u32)) -> Canvas { + pub fn new( + host: &str, + painter_count: usize, + size: (u32, u32), + offset: (u32, u32), + binary: bool, + ) -> Canvas { // Initialize the object let mut canvas = Canvas { host: host.to_string(), @@ -35,14 +41,14 @@ impl Canvas { println!("Starting painter threads..."); // Spawn some painters - canvas.spawn_painters(); + canvas.spawn_painters(binary); // Return the canvas canvas } /// Spawn the painters for this canvas - fn spawn_painters(&mut self) { + fn spawn_painters(&mut self, binary: bool) { // Spawn some painters for i in 0..self.painter_count { // Determine the slice width @@ -52,12 +58,12 @@ impl Canvas { let painter_area = Rect::from((i as u32) * width, 0, width, self.size.1); // Spawn the painter - self.spawn_painter(painter_area); + self.spawn_painter(painter_area, binary); } } /// Spawn a single painter in a thread. - fn spawn_painter(&mut self, area: Rect) { + fn spawn_painter(&mut self, area: Rect, binary: bool) { // Get the host that will be used let host = self.host.to_string(); @@ -76,12 +82,12 @@ impl Canvas { // The painting loop 'paint: loop { // Connect - let client = match Client::connect(host.clone()) { + let client = match Client::connect(host.clone(), binary) { Ok(client) => client, Err(e) => { eprintln!("Painter failed to connect: {}", e); break 'paint; - }, + } }; painter.set_client(Some(client)); diff --git a/src/pix/client.rs b/src/pix/client.rs index aeb34ac..80ad752 100644 --- a/src/pix/client.rs +++ b/src/pix/client.rs @@ -27,33 +27,57 @@ const PIX_SERVER_SIZE_REGEX: &str = r"^(?i)\s*SIZE\s+([[:digit:]]+)\s+([[:digit: /// to the pixelflut panel. pub struct Client { stream: BufStream, + + /// Whether to use binary mode (PB) instead of (PX). + binary: bool, } impl Client { /// Create a new client instance. - pub fn new(stream: TcpStream) -> Client { + pub fn new(stream: TcpStream, binary: bool) -> Client { Client { stream: BufStream::new(stream), + binary, } } /// Create a new client instane from the given host, and connect to it. - pub fn connect(host: String) -> Result { + pub fn connect(host: String, binary: bool) -> Result { // Create a new stream, and instantiate the client - Ok(Client::new(create_stream(host)?)) + Ok(Client::new(create_stream(host)?, binary)) } /// Write a pixel to the given stream. pub fn write_pixel(&mut self, x: u32, y: u32, color: Color) -> Result<(), Error> { - // Write the command to set a pixel - self.write_command(&format!("PX {} {} {}", x, y, color.as_hex())) + if self.binary { + self.write_command( + &[ + b'P', + b'B', + x as u8, + (x >> 8) as u8, + y as u8, + (y >> 8) as u8, + color.r, + color.g, + color.b, + color.a, + ], + false, + ) + } else { + self.write_command( + format!("PX {} {} {}", x, y, color.as_hex()).as_bytes(), + true, + ) + } } /// Read the size of the screen. pub fn read_screen_size(&mut self) -> Result<(u32, u32), Error> { // Read the screen size let data = self - .write_read_command("SIZE".into()) + .write_read_command(b"SIZE") .expect("Failed to read screen size"); // Build a regex to parse the screen size @@ -77,26 +101,27 @@ impl Client { } /// Write the given command to the given stream. - fn write_command(&mut self, cmd: &str) -> Result<(), Error> { + fn write_command(&mut self, cmd: &[u8], newline: bool) -> Result<(), Error> { // Write the pixels and a new line - self.stream.write_all(cmd.as_bytes())?; - self.stream.write_all(b"\n")?; + self.stream.write_all(cmd)?; + if newline { + self.stream.write_all(b"\n")?; + } // Flush, make sure to clear the send buffer // TODO: only flush each 100 pixels? // TODO: make flushing configurable? // TODO: make buffer size configurable? - self.stream - .flush()?; + self.stream.flush()?; // Everything seems to be ok Ok(()) } /// Write the given command to the given stream, and read the output. - fn write_read_command(&mut self, cmd: &str) -> Result { + fn write_read_command(&mut self, cmd: &[u8]) -> Result { // Write the command - self.write_command(cmd)?; + self.write_command(cmd, true)?; // Flush the pipe, ensure the command is actually sent self.stream.flush()?; @@ -114,7 +139,7 @@ impl Client { impl Drop for Client { /// Nicely drop the connection when the client is disconnected. fn drop(&mut self) { - let _ = self.write_command("\nQUIT".into()); + let _ = self.write_command(b"\nQUIT", true); } }