From d757066d88fc557f0d72cdc718d0ab3657bb9f72 Mon Sep 17 00:00:00 2001 From: lif <> Date: Wed, 27 Mar 2024 15:20:36 -0700 Subject: [PATCH 01/20] wip - pixfmts - ground work for supporting non-32-bit color depths --- src/encodings.rs | 21 ++- src/pixel_formats.rs | 429 +++++++++++++++++++++++++------------------ src/server.rs | 20 +- 3 files changed, 278 insertions(+), 192 deletions(-) diff --git a/src/encodings.rs b/src/encodings.rs index 94dbf0d..9d6539a 100644 --- a/src/encodings.rs +++ b/src/encodings.rs @@ -5,10 +5,11 @@ // Copyright 2022 Oxide Computer Company use crate::{ - pixel_formats::rgb_888, + pixel_formats::transform, rfb::{PixelFormat, Position, Resolution}, }; +use crate::rfb::ColorSpecification; use EncodingType::*; #[derive(Debug)] @@ -80,7 +81,7 @@ impl From for EncodingType { 21 => JPEG, 6 => Zlib, -314 => CursorWithAlpha, - v => EncodingType::Other(v), + v => Other(v), } } } @@ -98,7 +99,7 @@ impl RawEncoding { impl Encoding for RawEncoding { fn get_type(&self) -> EncodingType { - EncodingType::Raw + Raw } fn encode(&self) -> &Vec { @@ -106,13 +107,19 @@ impl Encoding for RawEncoding { } fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box { - // XXX: This assumes the pixel formats are both rgb888. The server code verifies this + // XXX: This assumes the pixel formats are both rgb. The server code verifies this // before calling. - assert!(input.is_rgb_888()); - assert!(output.is_rgb_888()); + assert!(matches!( + &input.color_spec, + ColorSpecification::ColorFormat(_) + )); + assert!(matches!( + &output.color_spec, + ColorSpecification::ColorFormat(_) + )); Box::new(Self { - pixels: rgb_888::transform(&self.pixels, &input, &output), + pixels: transform(&self.pixels, &input, &output), }) } } diff --git a/src/pixel_formats.rs b/src/pixel_formats.rs index faa39eb..b1e5baa 100644 --- a/src/pixel_formats.rs +++ b/src/pixel_formats.rs @@ -53,6 +53,8 @@ //! - blue = pixel\[1\] & 255 = 0x03 //! +use crate::rfb::{ColorFormat, ColorSpecification, PixelFormat}; + #[derive(Debug, thiserror::Error)] pub enum PixelFormatError { #[error("unsupported or unknown fourcc: 0x{0:x}")] @@ -68,102 +70,192 @@ pub enum PixelFormatError { /// A good reference for mapping common fourcc codes to their corresponding pixel formats is the /// drm_fourcc.h header file in the linux source code. pub mod fourcc { - use super::{rgb_888, PixelFormatError}; - use crate::rfb::{ColorFormat, ColorSpecification, PixelFormat}; + use super::{ColorConstants, PixelFormatError}; + use crate::pixel_formats::{Rgb332Formats, Rgb565Formats, Rgb888Formats}; + use crate::rfb::PixelFormat; + + #[repr(u32)] + pub enum FourCC { + /// little-endian xRGB, 8:8:8:8 + XR24 = u32::from_ne_bytes(*b"XR24"), + /// little-endian RGBx, 8:8:8:8 + RX24 = u32::from_ne_bytes(*b"RX24"), + /// little-endian xBGR, 8:8:8:8 + XB24 = u32::from_ne_bytes(*b"XB24"), + /// little-endian BGRx, 8:8:8:8 + BX24 = u32::from_ne_bytes(*b"BX24"), + /// little-endian RGB, 5:6:5 + RG16 = u32::from_ne_bytes(*b"RG16"), + /// little-endian BGR, 5:6:5 + BG16 = u32::from_ne_bytes(*b"BG16"), + /// RGB, 3:3:2 + RGB8 = u32::from_ne_bytes(*b"RGB8"), + /// BGR, 2:3:3 + BGR8 = u32::from_ne_bytes(*b"BGR8"), + } - // XXX: it might make sense to turn fourcc values into a type (such as an enum or collection of - // enums) - pub const FOURCC_XR24: u32 = 0x34325258; // little-endian xRGB, 8:8:8:8 - pub const FOURCC_RX24: u32 = 0x34325852; // little-endian RGBx, 8:8:8:8 - pub const FOURCC_BX24: u32 = 0x34325842; // little-endian BGRx, 8:8:8:8 - pub const FOURCC_XB24: u32 = 0x34324258; // little-endian xBGR, 8:8:8:8 + pub const FOURCC_XR24: u32 = FourCC::XR24 as u32; + pub const FOURCC_RX24: u32 = FourCC::RX24 as u32; + pub const FOURCC_BX24: u32 = FourCC::BX24 as u32; + pub const FOURCC_XB24: u32 = FourCC::XB24 as u32; + pub const FOURCC_RG16: u32 = FourCC::RG16 as u32; + pub const FOURCC_BG16: u32 = FourCC::BG16 as u32; + pub const FOURCC_RGB8: u32 = FourCC::RGB8 as u32; + pub const FOURCC_BGR8: u32 = FourCC::BGR8 as u32; + + impl TryFrom for FourCC { + type Error = PixelFormatError; + + fn try_from(value: u32) -> Result { + match value { + FOURCC_XR24 => Ok(FourCC::XR24), + FOURCC_RX24 => Ok(FourCC::RX24), + FOURCC_XB24 => Ok(FourCC::XB24), + FOURCC_BX24 => Ok(FourCC::BX24), + FOURCC_RG16 => Ok(FourCC::RG16), + FOURCC_BG16 => Ok(FourCC::BG16), + FOURCC_RGB8 => Ok(FourCC::RGB8), + FOURCC_BGR8 => Ok(FourCC::BGR8), + v => Err(PixelFormatError::UnsupportedFourCc(v)), + } + } + } + + impl From<&FourCC> for PixelFormat { + fn from(value: &FourCC) -> Self { + match value { + FourCC::XR24 => Rgb888Formats::to_pix_fmt(false, 0), + FourCC::RX24 => Rgb888Formats::to_pix_fmt(false, 8), + FourCC::XB24 => Rgb888Formats::to_pix_fmt(true, 0), + FourCC::BX24 => Rgb888Formats::to_pix_fmt(true, 8), + FourCC::RG16 => Rgb565Formats::to_pix_fmt(false, 0), + FourCC::BG16 => Rgb565Formats::to_pix_fmt(true, 0), + FourCC::RGB8 => Rgb332Formats::to_pix_fmt(false, 0), + FourCC::BGR8 => Rgb332Formats::to_pix_fmt(true, 0), + } + } + } pub fn fourcc_to_pixel_format(fourcc: u32) -> Result { - match fourcc { - // little-endian xRGB - FOURCC_XR24 => Ok(PixelFormat { - bits_per_pixel: rgb_888::BITS_PER_PIXEL, - depth: rgb_888::DEPTH, - big_endian: false, - color_spec: ColorSpecification::ColorFormat(ColorFormat { - red_max: rgb_888::MAX_VALUE, - green_max: rgb_888::MAX_VALUE, - blue_max: rgb_888::MAX_VALUE, - red_shift: rgb_888::BITS_PER_COLOR * 2, - green_shift: rgb_888::BITS_PER_COLOR * 1, - blue_shift: rgb_888::BITS_PER_COLOR * 0, - }), - }), + FourCC::try_from(fourcc).map(|fmt| PixelFormat::from(&fmt)) + } +} - // little-endian RGBx - FOURCC_RX24 => Ok(PixelFormat { - bits_per_pixel: rgb_888::BITS_PER_PIXEL, - depth: rgb_888::DEPTH, - big_endian: false, - color_spec: ColorSpecification::ColorFormat(ColorFormat { - red_max: rgb_888::MAX_VALUE, - green_max: rgb_888::MAX_VALUE, - blue_max: rgb_888::MAX_VALUE, - red_shift: rgb_888::BITS_PER_COLOR * 3, - green_shift: rgb_888::BITS_PER_COLOR * 2, - blue_shift: rgb_888::BITS_PER_COLOR * 1, - }), - }), +trait ColorConstants { + const BYTES_PER_PIXEL: usize = (Self::DEPTH as usize).next_power_of_two() / 8; + const BITS_PER_PIXEL: u8 = (Self::BYTES_PER_PIXEL * 8) as u8; - // little-endian BGRx - FOURCC_BX24 => Ok(PixelFormat { - bits_per_pixel: rgb_888::BITS_PER_PIXEL, - depth: rgb_888::DEPTH, + /// Number of bits used for color in a pixel + const DEPTH: u8 = Self::RED_BITS + Self::GREEN_BITS + Self::BLUE_BITS; + + /// Number of bits used for red channel value + const RED_BITS: u8; + /// Number of bits used for green channel value + const GREEN_BITS: u8; + /// Number of bits used for blue channel value + const BLUE_BITS: u8; + + /// Max value for red channel + const RED_MAX: u16 = (1u16 << Self::RED_BITS) - 1; + /// Max value for green channel + const GREEN_MAX: u16 = (1u16 << Self::GREEN_BITS) - 1; + /// Max value for blue channel + const BLUE_MAX: u16 = (1u16 << Self::BLUE_BITS) - 1; + + /// Returns true if a shift as specified in a pixel format is valid for described formats. + fn valid_shift(shift: u8) -> bool; + + /// Construct an appropriate PixelFormat definition for the given channel + /// ordering and base shift (e.g. BGRx 8:8:8:8 would be (true, 8)) + fn to_pix_fmt(bgr_order: bool, base_shift: u8) -> PixelFormat { + if bgr_order { + PixelFormat { + bits_per_pixel: Self::BITS_PER_PIXEL, + depth: Self::DEPTH, big_endian: false, color_spec: ColorSpecification::ColorFormat(ColorFormat { - red_max: rgb_888::MAX_VALUE, - green_max: rgb_888::MAX_VALUE, - blue_max: rgb_888::MAX_VALUE, - red_shift: rgb_888::BITS_PER_COLOR * 1, - green_shift: rgb_888::BITS_PER_COLOR * 2, - blue_shift: rgb_888::BITS_PER_COLOR * 3, + red_max: Self::RED_MAX, + green_max: Self::GREEN_MAX, + blue_max: Self::BLUE_MAX, + red_shift: base_shift, + green_shift: base_shift + Self::RED_BITS, + blue_shift: base_shift + Self::RED_BITS + Self::GREEN_BITS, }), - }), - - // little-endian xBGR - FOURCC_XB24 => Ok(PixelFormat { - bits_per_pixel: rgb_888::BITS_PER_PIXEL, - depth: rgb_888::DEPTH, + } + } else { + PixelFormat { + bits_per_pixel: Self::BITS_PER_PIXEL, + depth: Self::DEPTH, big_endian: false, color_spec: ColorSpecification::ColorFormat(ColorFormat { - red_max: rgb_888::MAX_VALUE, - green_max: rgb_888::MAX_VALUE, - blue_max: rgb_888::MAX_VALUE, - red_shift: rgb_888::BITS_PER_COLOR * 0, - green_shift: rgb_888::BITS_PER_COLOR * 1, - blue_shift: rgb_888::BITS_PER_COLOR * 2, + red_max: Self::RED_MAX, + green_max: Self::GREEN_MAX, + blue_max: Self::BLUE_MAX, + red_shift: base_shift + Self::GREEN_BITS + Self::BLUE_BITS, + green_shift: base_shift + Self::BLUE_BITS, + blue_shift: base_shift, }), - }), - - v => Err(PixelFormatError::UnsupportedFourCc(v)), + } } } } +struct Rgb888Formats; +struct Rgb565Formats; +struct Rgb332Formats; + +impl ColorConstants for Rgb888Formats { + const RED_BITS: u8 = 8; + const GREEN_BITS: u8 = 8; + const BLUE_BITS: u8 = 8; + + fn valid_shift(shift: u8) -> bool { + shift == 0 || shift == 8 || shift == 16 || shift == 24 + } +} + +impl ColorConstants for Rgb565Formats { + const RED_BITS: u8 = 5; + const GREEN_BITS: u8 = 6; + const BLUE_BITS: u8 = 5; + + fn valid_shift(shift: u8) -> bool { + shift == 0 || shift == 5 || shift == 11 + } +} + +impl ColorConstants for Rgb332Formats { + const RED_BITS: u8 = 3; + const GREEN_BITS: u8 = 3; + const BLUE_BITS: u8 = 2; + + // not the most thorough + fn valid_shift(shift: u8) -> bool { + shift == 0 || shift == 2 || shift == 3 || shift == 5 || shift == 6 + } +} + /// Utility functions for 32-bit RGB pixel formats, with 8-bits used per color. +#[deprecated] pub mod rgb_888 { - use crate::rfb::{ColorSpecification, PixelFormat}; + pub use super::transform; + use crate::pixel_formats::{ColorConstants, Rgb888Formats}; - pub const BYTES_PER_PIXEL: usize = 4; - pub const BITS_PER_PIXEL: u8 = 32; + pub const BYTES_PER_PIXEL: usize = Rgb888Formats::BYTES_PER_PIXEL; + pub const BITS_PER_PIXEL: u8 = Rgb888Formats::BITS_PER_PIXEL; /// Number of bits used for color in a pixel - pub const DEPTH: u8 = 24; + pub const DEPTH: u8 = Rgb888Formats::DEPTH; /// Number of bits used for a single color value - pub const BITS_PER_COLOR: u8 = 8; + pub const BITS_PER_COLOR: u8 = Rgb888Formats::RED_BITS; /// Max value for a given color - pub const MAX_VALUE: u16 = 255; + pub const MAX_VALUE: u16 = Rgb888Formats::RED_MAX; /// Returns true if a shift as specified in a pixel format is valid for rgb888 formats. pub fn valid_shift(shift: u8) -> bool { - shift == 0 || shift == 8 || shift == 16 || shift == 24 + Rgb888Formats::valid_shift(shift) } /// Returns the byte index into a 4-byte pixel vector for a given color shift, accounting for endianness. @@ -181,120 +273,79 @@ pub mod rgb_888 { pub fn unused_index(r: usize, g: usize, b: usize) -> usize { (3 + 2 + 1) - r - g - b } +} - /// Given a set of red/green/blue shifts from a pixel format and its endianness, determine - /// which byte index in a 4-byte vector representing a pixel maps to which color. - /// - /// For example, for the shifts red=0, green=8, blue=16, and a little-endian format, the - /// indices would be red=0, green=1, blue=2, and x=3. - pub fn rgbx_index( - red_shift: u8, - green_shift: u8, - blue_shift: u8, - big_endian: bool, - ) -> (usize, usize, usize, usize) { - let r = color_shift_to_index(red_shift, big_endian); - let g = color_shift_to_index(green_shift, big_endian); - let b = color_shift_to_index(blue_shift, big_endian); - let x = unused_index(r, g, b); - - (r, g, b, x) +/// Translate between RGB formats. +pub fn transform(pixels: &[u8], input: &PixelFormat, output: &PixelFormat) -> Vec { + if input == output { + return pixels.to_vec(); } - /// Translate between RGB888 formats. The input and output format must both be RGB888. - pub fn transform(pixels: &Vec, input: &PixelFormat, output: &PixelFormat) -> Vec { - assert!(input.is_rgb_888()); - assert!(output.is_rgb_888()); - - //let mut buf = Vec::with_capacity(pixels.len()); - //buf.resize(pixels.len(), 0x0u8); - let mut buf = vec![0; pixels.len()]; - - let (ir, ig, ib, ix) = match &input.color_spec { - ColorSpecification::ColorFormat(cf) => rgbx_index( - cf.red_shift, - cf.green_shift, - cf.blue_shift, - input.big_endian, - ), - ColorSpecification::ColorMap(_) => { - unreachable!(); - } - }; + let in_bytes_pp = input.bits_per_pixel.next_power_of_two() as usize / 8; + let out_bytes_pp = output.bits_per_pixel.next_power_of_two() as usize / 8; - let (or, og, ob, ox) = match &output.color_spec { - ColorSpecification::ColorFormat(cf) => rgbx_index( - cf.red_shift, - cf.green_shift, - cf.blue_shift, - output.big_endian, - ), - ColorSpecification::ColorMap(_) => { - unreachable!(); - } - }; + let in_be_shift = 8 * (4 - in_bytes_pp); + let out_be_shift = 8 * (4 - out_bytes_pp); - let mut i = 0; - while i < pixels.len() { - // Get the value for each color from the input... - let r = pixels[i + ir]; - let g = pixels[i + ig]; - let b = pixels[i + ib]; - let x = pixels[i + ix]; - - // and assign it to the right spot in the output pixel - buf[i + or] = r; - buf[i + og] = g; - buf[i + ob] = b; - buf[i + ox] = x; - - i += 4; - } + let mut buf = Vec::with_capacity(pixels.len() * in_bytes_pp / out_bytes_pp); - buf - } -} + let ColorSpecification::ColorFormat(in_cf) = &input.color_spec else { + unimplemented!("converting from indexed color mode"); + }; + let ColorSpecification::ColorFormat(out_cf) = &input.color_spec else { + unimplemented!("converting to indexed color mode"); + }; -#[cfg(test)] -mod tests { - use crate::pixel_formats::rgb_888::{color_shift_to_index, rgbx_index}; + let mut i = 0; + while i < pixels.len() { + let mut bytes = [0u8; 4]; + bytes.copy_from_slice(&pixels[i..i + 4]); + let word = if input.big_endian { + u32::from_be_bytes(bytes) >> in_be_shift + } else { + u32::from_le_bytes(bytes) + }; - use super::{fourcc, rgb_888::transform}; + // shift and mask + let ir_raw = (word >> in_cf.red_shift) & in_cf.red_max as u32; + let ig_raw = (word >> in_cf.green_shift) & in_cf.green_max as u32; + let ib_raw = (word >> in_cf.blue_shift) & in_cf.blue_max as u32; + + // convert to new range + let ir = ir_raw * out_cf.red_max as u32 / in_cf.red_max as u32; + let ig = ig_raw * out_cf.green_max as u32 / in_cf.green_max as u32; + let ib = ib_raw * out_cf.blue_max as u32 / in_cf.blue_max as u32; + + let or = ir << out_cf.red_shift; + let og = ig << out_cf.green_shift; + let ob = ib << out_cf.blue_shift; + let word = or | og | ob; + let bytes = if output.big_endian { + (word << out_be_shift).to_be_bytes() + } else { + word.to_le_bytes() + }; + buf.extend(&bytes[..out_bytes_pp]); - #[test] - fn test_color_shift_to_index() { - assert_eq!(color_shift_to_index(0, false), 0); - assert_eq!(color_shift_to_index(8, false), 1); - assert_eq!(color_shift_to_index(16, false), 2); - assert_eq!(color_shift_to_index(24, false), 3); - - assert_eq!(color_shift_to_index(0, true), 3); - assert_eq!(color_shift_to_index(8, true), 2); - assert_eq!(color_shift_to_index(16, true), 1); - assert_eq!(color_shift_to_index(24, true), 0); + i += in_bytes_pp; } - #[test] - fn test_rgbx_index() { - assert_eq!(rgbx_index(0, 8, 16, false), (0, 1, 2, 3)); - assert_eq!(rgbx_index(0, 8, 16, true), (3, 2, 1, 0)); - - assert_eq!(rgbx_index(8, 16, 24, false), (1, 2, 3, 0)); - assert_eq!(rgbx_index(8, 16, 24, true), (2, 1, 0, 3)); - - assert_eq!(rgbx_index(0, 16, 24, false), (0, 2, 3, 1)); - assert_eq!(rgbx_index(0, 16, 24, true), (3, 1, 0, 2)); - - assert_eq!(rgbx_index(8, 16, 24, false), (1, 2, 3, 0)); - assert_eq!(rgbx_index(8, 16, 24, true), (2, 1, 0, 3)); + buf +} - assert_eq!(rgbx_index(0, 24, 8, false), (0, 3, 1, 2)); - assert_eq!(rgbx_index(0, 24, 8, true), (3, 0, 2, 1)); - } +#[cfg(test)] +mod tests { + use super::{fourcc, transform}; #[test] - fn test_rgb888_transform() { - let pixels = vec![0u8, 1u8, 2u8, 3u8]; + fn test_rgb_transform() { + #[rustfmt::skip] + let pixels = vec![ + 0x00, 0x00, 0x00, 0x00, + 0x12, 0x34, 0x56, 0x78, + 0x9A, 0xBC, 0xDE, 0xF0, + 0xFF, 0xFF, 0xFF, 0xFF, + ]; // little-endian xRGB let xrgb_le = fourcc::fourcc_to_pixel_format(fourcc::FOURCC_XR24).unwrap(); @@ -314,32 +365,50 @@ mod tests { assert_eq!(transform(&pixels, &bgrx_le, &bgrx_le), pixels); assert_eq!(transform(&pixels, &xbgr_le, &xbgr_le), pixels); + // in all examples below, the 'x' non-channel value is dropped (zeroed) + // little-endian xRGB -> little-endian RGBx // B G R x x B G R - // [0, 1, 2, 3] -> [3, 0, 1, 2] - let p2 = vec![3u8, 0u8, 1u8, 2u8]; + // [0, 1, 2, 3] -> [0, 0, 1, 2] + #[rustfmt::skip] + let p2 = vec![ + 0x00, 0x00, 0x00, 0x00, + 0x00, 0x12, 0x34, 0x56, + 0x00, 0x9A, 0xBC, 0xDE, + 0x00, 0xFF, 0xFF, 0xFF, + ]; assert_eq!(transform(&pixels, &xrgb_le, &rgbx_le), p2); // little-endian RGBx -> little-endian xRGB // x B G R B G R x // [0, 1, 2, 3] -> [1, 2, 3, 0] - let p3 = vec![1u8, 2u8, 3u8, 0u8]; + #[rustfmt::skip] + let p3 = vec![ + 0x00, 0x00, 0x00, 0x00, + 0x34, 0x56, 0x78, 0x00, + 0xBC, 0xDE, 0xF0, 0x00, + 0xFF, 0xFF, 0xFF, 0x00, + ]; assert_eq!(transform(&pixels, &rgbx_le, &xrgb_le), p3); + // little-endian BGRx -> little-endian xBGR + // x R G B R G B x + // [0, 1, 2, 3] -> [1, 2, 3, 0] + assert_eq!(transform(&pixels, &bgrx_le, &xbgr_le), p3); // little-endian xRGB -> little-endian BGRx // B G R x x R G B - // [0, 1, 2, 3] -> [3, 2, 1, 0] - let p4 = vec![3u8, 2u8, 1u8, 0u8]; + // [0, 1, 2, 3] -> [0, 2, 1, 0] + #[rustfmt::skip] + let p4 = vec![ + 0x00, 0x00, 0x00, 0x00, + 0x00, 0x56, 0x34, 0x12, + 0x00, 0xF0, 0xDE, 0xBC, + 0x00, 0xFF, 0xFF, 0xFF, + ]; assert_eq!(transform(&pixels, &xrgb_le, &bgrx_le), p4); // little-endian BGRx -> little-endian xRGB // x R G B B G R x // [0, 1, 2, 3] -> [3, 2, 1, 0] assert_eq!(transform(&pixels, &bgrx_le, &xrgb_le), p4); - - // little-endian BGRx -> little-endian xBGR - // x R G B R G B x - // [0, 1, 2, 3] -> [1, 2, 3, 0] - let p5 = vec![1u8, 2u8, 3u8, 0u8]; - assert_eq!(transform(&pixels, &bgrx_le, &xbgr_le), p5); } } diff --git a/src/server.rs b/src/server.rs index 0ce98c1..cdbaf0a 100644 --- a/src/server.rs +++ b/src/server.rs @@ -10,14 +10,14 @@ use std::net::SocketAddr; use std::sync::Arc; use async_trait::async_trait; -use futures::FutureExt; use futures::future::Shared; +use futures::FutureExt; use log::{debug, error, info, trace}; use thiserror::Error; use tokio::io::AsyncWriteExt; use tokio::net::{TcpListener, TcpStream}; use tokio::select; -use tokio::sync::{Mutex, oneshot}; +use tokio::sync::{oneshot, Mutex}; use crate::rfb::{ ClientInit, ClientMessage, FramebufferUpdate, KeyEvent, PixelFormat, ProtoVersion, @@ -183,7 +183,12 @@ impl VncServer { Ok(()) } - async fn handle_conn(&self, s: &mut TcpStream, addr: SocketAddr, mut close_ch: Shared>) { + async fn handle_conn( + &self, + s: &mut TcpStream, + addr: SocketAddr, + mut close_ch: Shared>, + ) { info!("[{:?}] new connection", addr); if let Err(e) = self.rfb_handshake(s, addr).await { @@ -289,7 +294,10 @@ impl VncServer { // Create a channel to signal the server to stop. let (close_tx, close_rx) = oneshot::channel(); - assert!(self.stop_ch.lock().await.replace(close_tx).is_none(), "server already started"); + assert!( + self.stop_ch.lock().await.replace(close_tx).is_none(), + "server already started" + ); let mut close_rx = close_rx.shared(); loop { @@ -309,7 +317,9 @@ impl VncServer { let close_rx = close_rx.clone(); let server = self.clone(); tokio::spawn(async move { - server.handle_conn(&mut client_sock, client_addr, close_rx).await; + server + .handle_conn(&mut client_sock, client_addr, close_rx) + .await; }); } } From 7ebf78a59aefd0dce390d1757022032e32f06d93 Mon Sep 17 00:00:00 2001 From: lif <> Date: Wed, 27 Mar 2024 17:15:35 -0700 Subject: [PATCH 02/20] wip - pixfmts - fix bugs --- examples/server.rs | 156 +++++++++++++------------------------------ src/pixel_formats.rs | 79 ++++++++++++++++------ src/rfb.rs | 46 ++++++++----- src/server.rs | 13 ++-- 4 files changed, 138 insertions(+), 156 deletions(-) diff --git a/examples/server.rs b/examples/server.rs index affdf67..7a8ab1c 100644 --- a/examples/server.rs +++ b/examples/server.rs @@ -4,7 +4,7 @@ // // Copyright 2022 Oxide Computer Company -use anyhow::{bail, Result}; +use anyhow::Result; use async_trait::async_trait; use clap::{Parser, ValueEnum}; use env_logger; @@ -12,13 +12,12 @@ use image::io::Reader as ImageReader; use image::GenericImageView; use log::info; use rfb::encodings::RawEncoding; +use rfb::pixel_formats::fourcc::FourCC; +use rfb::pixel_formats::transform; use rfb::rfb::{ FramebufferUpdate, KeyEvent, PixelFormat, ProtoVersion, Rectangle, SecurityType, SecurityTypes, }; -use rfb::{ - pixel_formats::rgb_888, - server::{Server, VncServer, VncServerConfig, VncServerData}, -}; +use rfb::server::{Server, VncServer, VncServerConfig, VncServerData}; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; const WIDTH: usize = 1024; @@ -43,21 +42,9 @@ struct Args { #[clap(value_enum, short, long, default_value_t = Image::Oxide)] image: Image, - /// Pixel endianness - #[clap(long, default_value_t = false, action = clap::ArgAction::Set)] - big_endian: bool, - - /// Byte mapping to red (4-byte RGB pixel, endian-agnostic) - #[clap(short, long, default_value_t = 0)] - red_order: u8, - - /// Byte mapping to green (4-byte RGB pixel, endian-agnostic) - #[clap(short, long, default_value_t = 1)] - green_order: u8, - - /// Byte mapping to blue (4-byte RGB pixel, endian-agnostic) - #[clap(short, long, default_value_t = 2)] - blue_order: u8, + /// Pixel format + #[clap(short, long, default_value = "XR24", action = clap::ArgAction::Set)] + fourcc: FourCC, } #[derive(ValueEnum, Debug, Copy, Clone)] @@ -74,8 +61,7 @@ enum Image { #[derive(Clone)] struct ExampleServer { display: Image, - rgb_order: (u8, u8, u8), - big_endian: bool, + pixfmt: PixelFormat, } #[tokio::main] @@ -83,22 +69,11 @@ async fn main() -> Result<()> { env_logger::init(); let args = Args::parse(); - validate_order(args.red_order, args.green_order, args.blue_order)?; - - let pf = PixelFormat::new_colorformat( - rgb_888::BITS_PER_PIXEL, - rgb_888::DEPTH, - args.big_endian, - order_to_shift(args.red_order), - rgb_888::MAX_VALUE, - order_to_shift(args.green_order), - rgb_888::MAX_VALUE, - order_to_shift(args.blue_order), - rgb_888::MAX_VALUE, - ); + + let pixfmt = PixelFormat::from(&args.fourcc); info!( "Starting server: image: {:?}, pixel format; {:#?}", - args.image, pf + args.image, pixfmt ); let config = VncServerConfig { @@ -110,12 +85,11 @@ async fn main() -> Result<()> { let data = VncServerData { width: WIDTH as u16, height: HEIGHT as u16, - input_pixel_format: pf.clone(), + input_pixel_format: pixfmt.clone(), }; let server = ExampleServer { display: args.image, - rgb_order: (args.red_order, args.green_order, args.blue_order), - big_endian: args.big_endian, + pixfmt, }; let s = VncServer::new(server, config, data); s.start().await?; @@ -123,97 +97,57 @@ async fn main() -> Result<()> { Ok(()) } -fn validate_order(r: u8, g: u8, b: u8) -> Result<()> { - if r > 3 || g > 3 || b > 3 { - bail!("r/g/b must have ordering of 0, 1, 2, or 3"); - } - - if r == g || r == b || g == b { - bail!("r/g/b must have unique orderings"); - } - - Ok(()) -} - -fn order_to_shift(order: u8) -> u8 { - assert!(order <= 3); - (3 - order) * rgb_888::BITS_PER_COLOR -} - -fn order_to_index(order: u8, big_endian: bool) -> u8 { - assert!(order <= 3); - - if big_endian { - order - } else { - 4 - order - 1 - } -} - -fn generate_color(index: u8, big_endian: bool) -> Vec { - const LEN: usize = WIDTH * HEIGHT * rgb_888::BYTES_PER_PIXEL; - let mut pixels = vec![0x0u8; LEN]; - - let idx = order_to_index(index, big_endian); - - let mut x = 0; - for i in 0..pixels.len() { - if x == idx { - pixels[i] = 0xff; - } +fn generate_color(img: Image, pixfmt: &PixelFormat) -> Vec { + let bytes_pp = pixfmt.bits_per_pixel as usize / 8; + let len = WIDTH * HEIGHT * bytes_pp; + let mut pixels = Vec::with_capacity(len); + + let color = match img { + Image::Red => 0xFF000000u32.to_le_bytes(), + Image::Green => 0x00FF0000u32.to_le_bytes(), + Image::Blue => 0x0000FF00u32.to_le_bytes(), + Image::White => 0xFFFFFF00u32.to_le_bytes(), + Image::Black => 0u32.to_le_bytes(), + _ => unreachable!(), + }; + let bytes = transform(&color, &PixelFormat::from(&FourCC::RX24), pixfmt); - if x == 3 { - x = 0; - } else { - x += 1; - } + while pixels.len() < len { + pixels.extend(&bytes); } pixels } -fn generate_image(name: &str, big_endian: bool, rgb_order: (u8, u8, u8)) -> Vec { - const LEN: usize = WIDTH * HEIGHT * rgb_888::BYTES_PER_PIXEL; +fn generate_image(name: &str, pixfmt: &PixelFormat) -> Vec { + const RGBX24_BYTES_PP: usize = 4; + const LEN: usize = WIDTH * HEIGHT * RGBX24_BYTES_PP; + let mut pixels = vec![0xffu8; LEN]; let img = ImageReader::open(name).unwrap().decode().unwrap(); - let (r, g, b) = rgb_order; - let r_idx = order_to_index(r, big_endian) as usize; - let g_idx = order_to_index(g, big_endian) as usize; - let b_idx = order_to_index(b, big_endian) as usize; - let x_idx = rgb_888::unused_index(r_idx, g_idx, b_idx); - // Convert the input image pixels to the requested pixel format. for (x, y, pixel) in img.pixels() { let ux = x as usize; let uy = y as usize; - let y_offset = WIDTH * rgb_888::BYTES_PER_PIXEL; - let x_offset = ux * rgb_888::BYTES_PER_PIXEL; + let y_offset = WIDTH * RGBX24_BYTES_PP; + let x_offset = ux * RGBX24_BYTES_PP; + let offset = uy * y_offset + x_offset; - pixels[uy * y_offset + x_offset + r_idx] = pixel[0]; - pixels[uy * y_offset + x_offset + g_idx] = pixel[1]; - pixels[uy * y_offset + x_offset + b_idx] = pixel[2]; - pixels[uy * y_offset + x_offset + x_idx] = pixel[3]; + pixels[offset..offset + 4].copy_from_slice(&pixel.0); } - - pixels + transform(&pixels, &PixelFormat::from(&FourCC::XB24), pixfmt) } -fn generate_pixels(img: Image, big_endian: bool, rgb_order: (u8, u8, u8)) -> Vec { - const LEN: usize = WIDTH * HEIGHT * rgb_888::BYTES_PER_PIXEL; - - let (r, g, b) = rgb_order; - +fn generate_pixels(img: Image, pixfmt: &PixelFormat) -> Vec { match img { - Image::Oxide => generate_image("example-images/oxide.jpg", big_endian, rgb_order), - Image::TestTubes => generate_image("example-images/test-tubes.jpg", big_endian, rgb_order), - Image::Red => generate_color(r, big_endian), - Image::Green => generate_color(g, big_endian), - Image::Blue => generate_color(b, big_endian), - Image::White => vec![0xffu8; LEN], - Image::Black => vec![0x0u8; LEN], + Image::Oxide => generate_image("example-images/oxide.jpg", pixfmt), + Image::TestTubes => generate_image("example-images/test-tubes.jpg", pixfmt), + Image::Red | Image::Green | Image::Blue | Image::White | Image::Black => { + generate_color(img, pixfmt) + } } } @@ -222,7 +156,7 @@ impl Server for ExampleServer { async fn get_framebuffer_update(&self) -> FramebufferUpdate { let pixels_width = 1024; let pixels_height = 768; - let pixels = generate_pixels(self.display, self.big_endian, self.rgb_order); + let pixels = generate_pixels(self.display, &self.pixfmt); let r = Rectangle::new( 0, 0, diff --git a/src/pixel_formats.rs b/src/pixel_formats.rs index b1e5baa..9839fcb 100644 --- a/src/pixel_formats.rs +++ b/src/pixel_formats.rs @@ -59,6 +59,8 @@ use crate::rfb::{ColorFormat, ColorSpecification, PixelFormat}; pub enum PixelFormatError { #[error("unsupported or unknown fourcc: 0x{0:x}")] UnsupportedFourCc(u32), + #[error("invalid fourcc name: {0}")] + InvalidFourCcString(String), } /// Utility functions and constants related to fourcc codes. @@ -73,31 +75,44 @@ pub mod fourcc { use super::{ColorConstants, PixelFormatError}; use crate::pixel_formats::{Rgb332Formats, Rgb565Formats, Rgb888Formats}; use crate::rfb::PixelFormat; + use std::str::FromStr; + #[derive(Copy, Clone, Debug)] #[repr(u32)] pub enum FourCC { /// little-endian xRGB, 8:8:8:8 - XR24 = u32::from_ne_bytes(*b"XR24"), + XR24 = u32::from_le_bytes(*b"XR24"), /// little-endian RGBx, 8:8:8:8 - RX24 = u32::from_ne_bytes(*b"RX24"), + RX24 = u32::from_le_bytes(*b"RX24"), /// little-endian xBGR, 8:8:8:8 - XB24 = u32::from_ne_bytes(*b"XB24"), + XB24 = u32::from_le_bytes(*b"XB24"), /// little-endian BGRx, 8:8:8:8 - BX24 = u32::from_ne_bytes(*b"BX24"), + BX24 = u32::from_le_bytes(*b"BX24"), /// little-endian RGB, 5:6:5 - RG16 = u32::from_ne_bytes(*b"RG16"), + RG16 = u32::from_le_bytes(*b"RG16"), /// little-endian BGR, 5:6:5 - BG16 = u32::from_ne_bytes(*b"BG16"), + BG16 = u32::from_le_bytes(*b"BG16"), /// RGB, 3:3:2 - RGB8 = u32::from_ne_bytes(*b"RGB8"), + RGB8 = u32::from_le_bytes(*b"RGB8"), /// BGR, 2:3:3 - BGR8 = u32::from_ne_bytes(*b"BGR8"), + BGR8 = u32::from_le_bytes(*b"BGR8"), } + pub const SUPPORTED: &[FourCC] = &[ + FourCC::XR24, + FourCC::RX24, + FourCC::XB24, + FourCC::BX24, + FourCC::RG16, + FourCC::BG16, + FourCC::RGB8, + FourCC::BGR8, + ]; + pub const FOURCC_XR24: u32 = FourCC::XR24 as u32; pub const FOURCC_RX24: u32 = FourCC::RX24 as u32; - pub const FOURCC_BX24: u32 = FourCC::BX24 as u32; pub const FOURCC_XB24: u32 = FourCC::XB24 as u32; + pub const FOURCC_BX24: u32 = FourCC::BX24 as u32; pub const FOURCC_RG16: u32 = FourCC::RG16 as u32; pub const FOURCC_BG16: u32 = FourCC::BG16 as u32; pub const FOURCC_RGB8: u32 = FourCC::RGB8 as u32; @@ -121,6 +136,20 @@ pub mod fourcc { } } + impl FromStr for FourCC { + type Err = PixelFormatError; + + fn from_str(s: &str) -> Result { + let mut bytes = [0u8; 4]; + if s.as_bytes().len() != 4 { + return Err(PixelFormatError::InvalidFourCcString(s.to_string())); + } + bytes.copy_from_slice(s.as_bytes()); + let word = u32::from_le_bytes(bytes); + Self::try_from(word) + } + } + impl From<&FourCC> for PixelFormat { fn from(value: &FourCC) -> Self { match value { @@ -292,18 +321,22 @@ pub fn transform(pixels: &[u8], input: &PixelFormat, output: &PixelFormat) -> Ve let ColorSpecification::ColorFormat(in_cf) = &input.color_spec else { unimplemented!("converting from indexed color mode"); }; - let ColorSpecification::ColorFormat(out_cf) = &input.color_spec else { + let ColorSpecification::ColorFormat(out_cf) = &output.color_spec else { unimplemented!("converting to indexed color mode"); }; let mut i = 0; while i < pixels.len() { - let mut bytes = [0u8; 4]; - bytes.copy_from_slice(&pixels[i..i + 4]); + let mut in_bytes = [0u8; 4]; + let mut k = 0; + for j in i..(i + 4).min(pixels.len()) { + in_bytes[k] = pixels[j]; + k += 1; + } let word = if input.big_endian { - u32::from_be_bytes(bytes) >> in_be_shift + u32::from_be_bytes(in_bytes) >> in_be_shift } else { - u32::from_le_bytes(bytes) + u32::from_le_bytes(in_bytes) }; // shift and mask @@ -312,20 +345,22 @@ pub fn transform(pixels: &[u8], input: &PixelFormat, output: &PixelFormat) -> Ve let ib_raw = (word >> in_cf.blue_shift) & in_cf.blue_max as u32; // convert to new range - let ir = ir_raw * out_cf.red_max as u32 / in_cf.red_max as u32; - let ig = ig_raw * out_cf.green_max as u32 / in_cf.green_max as u32; - let ib = ib_raw * out_cf.blue_max as u32 / in_cf.blue_max as u32; + let or_raw = ir_raw * out_cf.red_max as u32 / in_cf.red_max as u32; + let og_raw = ig_raw * out_cf.green_max as u32 / in_cf.green_max as u32; + let ob_raw = ib_raw * out_cf.blue_max as u32 / in_cf.blue_max as u32; + + let or = or_raw << out_cf.red_shift; + let og = og_raw << out_cf.green_shift; + let ob = ob_raw << out_cf.blue_shift; - let or = ir << out_cf.red_shift; - let og = ig << out_cf.green_shift; - let ob = ib << out_cf.blue_shift; let word = or | og | ob; - let bytes = if output.big_endian { + + let out_bytes = if output.big_endian { (word << out_be_shift).to_be_bytes() } else { word.to_le_bytes() }; - buf.extend(&bytes[..out_bytes_pp]); + buf.extend(&out_bytes[..out_bytes_pp]); i += in_bytes_pp; } diff --git a/src/rfb.rs b/src/rfb.rs index 870a55c..9945e5d 100644 --- a/src/rfb.rs +++ b/src/rfb.rs @@ -13,7 +13,6 @@ use tokio::net::TcpStream; use crate::encodings::{Encoding, EncodingType}; use crate::keysym::KeySym; -use crate::pixel_formats::rgb_888; #[derive(Debug, Error)] pub enum ProtocolError { @@ -391,7 +390,7 @@ pub struct PixelFormat { impl PixelFormat { /// Constructor for a PixelFormat that uses a color format to specify colors. pub fn new_colorformat( - bbp: u8, + bpp: u8, depth: u8, big_endian: bool, red_shift: u8, @@ -402,7 +401,7 @@ impl PixelFormat { blue_max: u16, ) -> Self { PixelFormat { - bits_per_pixel: bbp, + bits_per_pixel: bpp, depth, big_endian, color_spec: ColorSpecification::ColorFormat(ColorFormat { @@ -417,22 +416,40 @@ impl PixelFormat { } /// Returns true if the pixel format is RGB888 (8-bits per color and 32 bits per pixel). + #[deprecated] pub fn is_rgb_888(&self) -> bool { - if self.bits_per_pixel != rgb_888::BITS_PER_PIXEL || self.depth != rgb_888::DEPTH { - return false; + #[allow(deprecated)] + { + use crate::pixel_formats::rgb_888; + + if self.bits_per_pixel != rgb_888::BITS_PER_PIXEL || self.depth != rgb_888::DEPTH { + return false; + } + + match &self.color_spec { + ColorSpecification::ColorFormat(cf) => { + (cf.red_max == rgb_888::MAX_VALUE) + && (cf.green_max == rgb_888::MAX_VALUE) + && (cf.blue_max == rgb_888::MAX_VALUE) + && (rgb_888::valid_shift(cf.red_shift)) + && (rgb_888::valid_shift(cf.green_shift)) + && (rgb_888::valid_shift(cf.blue_shift)) + } + ColorSpecification::ColorMap(_) => false, + } } + } - match &self.color_spec { - ColorSpecification::ColorFormat(cf) => { - (cf.red_max == rgb_888::MAX_VALUE) - && (cf.green_max == rgb_888::MAX_VALUE) - && (cf.blue_max == rgb_888::MAX_VALUE) - && (rgb_888::valid_shift(cf.red_shift)) - && (rgb_888::valid_shift(cf.green_shift)) - && (rgb_888::valid_shift(cf.blue_shift)) + /// Returns true if the pixel format is supported (currently only certain + /// variants of RGB888, RGB565, and RGB332). + pub fn is_supported(&self) -> bool { + for fcc in crate::pixel_formats::fourcc::SUPPORTED { + let fmt: PixelFormat = fcc.into(); + if *self == fmt { + return true; } - ColorSpecification::ColorMap(_) => false, } + return false; } } @@ -482,7 +499,6 @@ impl WriteMessage for PixelFormat { } #[derive(Debug, Clone, PartialEq)] -#[allow(dead_code)] pub enum ColorSpecification { ColorFormat(ColorFormat), ColorMap(ColorMap), // TODO: implement diff --git a/src/server.rs b/src/server.rs index cdbaf0a..2a2c4a1 100644 --- a/src/server.rs +++ b/src/server.rs @@ -243,21 +243,18 @@ impl VncServer { // For now, we only support transformations between 4-byte RGB formats, so // if the requested format isn't one of those, we'll just leave the pixels // as is. - if data.input_pixel_format != output_pixel_format - && data.input_pixel_format.is_rgb_888() - && output_pixel_format.is_rgb_888() + if data.input_pixel_format == output_pixel_format { + debug!("no input transformation needed"); + } else if data.input_pixel_format.is_supported() + && output_pixel_format.is_supported() { debug!( "transforming: input={:#?}, output={:#?}", data.input_pixel_format, output_pixel_format ); fbu = fbu.transform(&data.input_pixel_format, &output_pixel_format); - } else if !(data.input_pixel_format.is_rgb_888() - && output_pixel_format.is_rgb_888()) - { - debug!("cannot transform between pixel formats (not rgb888): input.is_rgb_888()={}, output.is_rgb_888()={}", data.input_pixel_format.is_rgb_888(), output_pixel_format.is_rgb_888()); } else { - debug!("no input transformation needed"); + debug!("cannot transform between pixel formats: input.is_supported()={}, output.is_supported()={}", data.input_pixel_format.is_supported(), output_pixel_format.is_supported()); } if let Err(e) = fbu.write_to(s).await { From 2b46f1880250faf6d31ab892b1fb2f09748fa8bf Mon Sep 17 00:00:00 2001 From: lif <> Date: Thu, 28 Mar 2024 03:59:59 +0000 Subject: [PATCH 03/20] wip - pixfmts - fix test --- src/pixel_formats.rs | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/pixel_formats.rs b/src/pixel_formats.rs index 9839fcb..90ace90 100644 --- a/src/pixel_formats.rs +++ b/src/pixel_formats.rs @@ -437,13 +437,21 @@ mod tests { let p4 = vec![ 0x00, 0x00, 0x00, 0x00, 0x00, 0x56, 0x34, 0x12, - 0x00, 0xF0, 0xDE, 0xBC, + 0x00, 0xDE, 0xBC, 0x9A, 0x00, 0xFF, 0xFF, 0xFF, ]; assert_eq!(transform(&pixels, &xrgb_le, &bgrx_le), p4); + // little-endian BGRx -> little-endian xRGB // x R G B B G R x // [0, 1, 2, 3] -> [3, 2, 1, 0] - assert_eq!(transform(&pixels, &bgrx_le, &xrgb_le), p4); + #[rustfmt::skip] + let p5 = vec![ + 0x00, 0x00, 0x00, 0x00, + 0x78, 0x56, 0x34, 0x00, + 0xF0, 0xDE, 0xBC, 0x00, + 0xFF, 0xFF, 0xFF, 0x00, + ]; + assert_eq!(transform(&pixels, &bgrx_le, &xrgb_le), p5); } } From 465a6d86f6d76ffad75a6012ebfed0ee8b90b3fd Mon Sep 17 00:00:00 2001 From: lif <> Date: Thu, 28 Mar 2024 04:00:10 +0000 Subject: [PATCH 04/20] wip - mouse - add pointer_event to interface --- src/server.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/server.rs b/src/server.rs index 2a2c4a1..f77e234 100644 --- a/src/server.rs +++ b/src/server.rs @@ -20,9 +20,9 @@ use tokio::select; use tokio::sync::{oneshot, Mutex}; use crate::rfb::{ - ClientInit, ClientMessage, FramebufferUpdate, KeyEvent, PixelFormat, ProtoVersion, - ProtocolError, ReadMessage, SecurityResult, SecurityType, SecurityTypes, ServerInit, - WriteMessage, + ClientInit, ClientMessage, FramebufferUpdate, KeyEvent, PixelFormat, PointerEvent, + ProtoVersion, ProtocolError, ReadMessage, SecurityResult, SecurityType, SecurityTypes, + ServerInit, WriteMessage, }; #[derive(Debug, Error)] @@ -81,6 +81,7 @@ pub struct VncServer { pub trait Server: Sync + Send + 'static { async fn get_framebuffer_update(&self) -> FramebufferUpdate; async fn key_event(&self, _ke: KeyEvent) {} + async fn pointer_event(&self, _pe: PointerEvent) {} async fn stop(&self) {} } @@ -272,6 +273,7 @@ impl VncServer { } ClientMessage::PointerEvent(pe) => { trace!("Rx [{:?}: PointerEvent={:?}", addr, pe); + self.server.pointer_event(pe).await; } ClientMessage::ClientCutText(t) => { trace!("Rx [{:?}: ClientCutText={:?}", addr, t); From ffbb311956879e60ce8b55ab4166abfe4f4a1680 Mon Sep 17 00:00:00 2001 From: lif <> Date: Thu, 28 Mar 2024 04:45:29 +0000 Subject: [PATCH 05/20] wip - mouse - expose pointer event data --- src/rfb.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/rfb.rs b/src/rfb.rs index 9945e5d..123ac99 100644 --- a/src/rfb.rs +++ b/src/rfb.rs @@ -244,9 +244,9 @@ impl FramebufferUpdate { } #[derive(Debug, Copy, Clone)] -pub(crate) struct Position { - x: u16, - y: u16, +pub struct Position { + pub x: u16, + pub y: u16, } impl ReadMessage for Position { @@ -717,7 +717,7 @@ impl KeyEvent { } bitflags! { - struct MouseButtons: u8 { + pub struct MouseButtons: u8 { const LEFT = 1 << 0; const MIDDLE = 1 << 1; const RIGHT = 1 << 2; @@ -729,10 +729,9 @@ bitflags! { } #[derive(Debug)] -#[allow(dead_code)] pub struct PointerEvent { - position: Position, - pressed: MouseButtons, + pub position: Position, + pub pressed: MouseButtons, } impl ReadMessage for PointerEvent { From f6a5fc8a1652b4982ba6019cac5f7d8ad3349d4a Mon Sep 17 00:00:00 2001 From: lif <> Date: Wed, 3 Apr 2024 22:04:38 +0000 Subject: [PATCH 06/20] wip - pixfmts - round when changing range --- src/pixel_formats.rs | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/src/pixel_formats.rs b/src/pixel_formats.rs index 90ace90..aa96f5d 100644 --- a/src/pixel_formats.rs +++ b/src/pixel_formats.rs @@ -313,6 +313,7 @@ pub fn transform(pixels: &[u8], input: &PixelFormat, output: &PixelFormat) -> Ve let in_bytes_pp = input.bits_per_pixel.next_power_of_two() as usize / 8; let out_bytes_pp = output.bits_per_pixel.next_power_of_two() as usize / 8; + // regardless of byteorder, we still want to let in_be_shift = 8 * (4 - in_bytes_pp); let out_be_shift = 8 * (4 - out_bytes_pp); @@ -344,10 +345,16 @@ pub fn transform(pixels: &[u8], input: &PixelFormat, output: &PixelFormat) -> Ve let ig_raw = (word >> in_cf.green_shift) & in_cf.green_max as u32; let ib_raw = (word >> in_cf.blue_shift) & in_cf.blue_max as u32; - // convert to new range - let or_raw = ir_raw * out_cf.red_max as u32 / in_cf.red_max as u32; - let og_raw = ig_raw * out_cf.green_max as u32 / in_cf.green_max as u32; - let ob_raw = ib_raw * out_cf.blue_max as u32 / in_cf.blue_max as u32; + // convert to new range (with rounding) + fn convert(c_in: u32, in_max: u16, out_max: u16) -> u32 { + let in_max = in_max as u32; + let out_max = out_max as u32; + ((c_in * out_max) + (in_max / 2)) / in_max + } + + let or_raw = convert(ir_raw, in_cf.red_max, out_cf.red_max); + let og_raw = convert(ig_raw, in_cf.green_max, out_cf.green_max); + let ob_raw = convert(ib_raw, in_cf.blue_max, out_cf.blue_max); let or = or_raw << out_cf.red_shift; let og = og_raw << out_cf.green_shift; From 03607a2b5ce208e22950354b7303e8477753a039 Mon Sep 17 00:00:00 2001 From: lif <> Date: Sat, 20 Jul 2024 01:32:53 -0700 Subject: [PATCH 07/20] wip - pixfmts - work on not crashing when a client requests a palette mode --- src/pixel_formats.rs | 126 +++++++++++++++++++++++++++++++++++++------ src/rfb.rs | 18 ++++--- 2 files changed, 121 insertions(+), 23 deletions(-) diff --git a/src/pixel_formats.rs b/src/pixel_formats.rs index aa96f5d..3312a46 100644 --- a/src/pixel_formats.rs +++ b/src/pixel_formats.rs @@ -322,9 +322,6 @@ pub fn transform(pixels: &[u8], input: &PixelFormat, output: &PixelFormat) -> Ve let ColorSpecification::ColorFormat(in_cf) = &input.color_spec else { unimplemented!("converting from indexed color mode"); }; - let ColorSpecification::ColorFormat(out_cf) = &output.color_spec else { - unimplemented!("converting to indexed color mode"); - }; let mut i = 0; while i < pixels.len() { @@ -352,23 +349,48 @@ pub fn transform(pixels: &[u8], input: &PixelFormat, output: &PixelFormat) -> Ve ((c_in * out_max) + (in_max / 2)) / in_max } - let or_raw = convert(ir_raw, in_cf.red_max, out_cf.red_max); - let og_raw = convert(ig_raw, in_cf.green_max, out_cf.green_max); - let ob_raw = convert(ib_raw, in_cf.blue_max, out_cf.blue_max); - - let or = or_raw << out_cf.red_shift; - let og = og_raw << out_cf.green_shift; - let ob = ob_raw << out_cf.blue_shift; + match &output.color_spec { + ColorSpecification::ColorFormat(out_cf) => { + let or_raw = convert(ir_raw, in_cf.red_max, out_cf.red_max); + let og_raw = convert(ig_raw, in_cf.green_max, out_cf.green_max); + let ob_raw = convert(ib_raw, in_cf.blue_max, out_cf.blue_max); - let word = or | og | ob; + let or = or_raw << out_cf.red_shift; + let og = og_raw << out_cf.green_shift; + let ob = ob_raw << out_cf.blue_shift; - let out_bytes = if output.big_endian { - (word << out_be_shift).to_be_bytes() - } else { - word.to_le_bytes() - }; - buf.extend(&out_bytes[..out_bytes_pp]); + let word = or | og | ob; + let out_bytes = if output.big_endian { + (word << out_be_shift).to_be_bytes() + } else { + word.to_le_bytes() + }; + buf.extend(&out_bytes[..out_bytes_pp]); + } + ColorSpecification::ColorMap(cmap) => { + let ir = convert(ir_raw, in_cf.red_max, 255) as i16; + let ig = convert(ig_raw, in_cf.green_max, 255) as i16; + let ib = convert(ib_raw, in_cf.blue_max, 255) as i16; + + let color_index = VESA_VGA_256_COLOR_PALETTE + .iter() + .enumerate() + .min_by_key(|(_, (pr, pg, pb))| { + // min by RGB distance formula. not the best by today's + // understanding of color perception, but neither is + // the VGA palette -- for now we're just supporting + // the protocol without crashing or misbehaving here. + (*pr as i16 - ir).pow(2) + + (*pg as i16 - ig).pow(2) + + (*pb as i16 - ib).pow(2) + }) + .unwrap() // constant array of nonzero size always has a min + .0; + + buf.push(color_index as u8); + } + } i += in_bytes_pp; } @@ -462,3 +484,73 @@ mod tests { assert_eq!(transform(&pixels, &bgrx_le, &xrgb_le), p5); } } + +// for compatibility with VNC clients configured to requesting ColorMap modes, +// we'll punt for now (rather than doing anything adaptive) and use this look-up table +// https://upload.wikimedia.org/wikipedia/commons/6/66/VGA_palette_with_black_borders.svg +#[rustfmt::skip] +const VESA_VGA_256_COLOR_PALETTE: [(u8, u8, u8); 256] = [ + (0x00, 0x00, 0x00), (0x00, 0x00, 0xaa), (0x00, 0xaa, 0x00), (0x00, 0xaa, 0xaa), + (0xaa, 0x00, 0x00), (0xaa, 0x00, 0xaa), (0xaa, 0x55, 0x00), (0xaa, 0xaa, 0xaa), + (0x55, 0x55, 0x55), (0x55, 0x55, 0xff), (0x55, 0xff, 0x55), (0x55, 0xff, 0xff), + (0xff, 0x55, 0x55), (0xff, 0x55, 0xff), (0xff, 0xff, 0x55), (0xff, 0xff, 0xff), + (0x00, 0x00, 0x00), (0x10, 0x10, 0x10), (0x20, 0x20, 0x20), (0x35, 0x35, 0x35), + (0x45, 0x45, 0x45), (0x55, 0x55, 0x55), (0x65, 0x65, 0x65), (0x75, 0x75, 0x75), + (0x8a, 0x8a, 0x8a), (0x9a, 0x9a, 0x9a), (0xaa, 0xaa, 0xaa), (0xba, 0xba, 0xba), + (0xca, 0xca, 0xca), (0xdf, 0xdf, 0xdf), (0xef, 0xef, 0xef), (0xff, 0xff, 0xff), + (0x00, 0x00, 0xff), (0x41, 0x00, 0xff), (0x82, 0x00, 0xff), (0xbe, 0x00, 0xff), + (0xff, 0x00, 0xff), (0xff, 0x00, 0xbe), (0xff, 0x00, 0x82), (0xff, 0x00, 0x41), + (0xff, 0x00, 0x00), (0xff, 0x41, 0x00), (0xff, 0x82, 0x00), (0xff, 0xbe, 0x00), + (0xff, 0xff, 0x00), (0xbe, 0xff, 0x00), (0x82, 0xff, 0x00), (0x41, 0xff, 0x00), + (0x00, 0xff, 0x00), (0x00, 0xff, 0x41), (0x00, 0xff, 0x82), (0x00, 0xff, 0xbe), + (0x00, 0xff, 0xff), (0x00, 0xbe, 0xff), (0x00, 0x82, 0xff), (0x00, 0x41, 0xff), + (0x82, 0x82, 0xff), (0x9e, 0x82, 0xff), (0xbe, 0x82, 0xff), (0xdf, 0x82, 0xff), + (0xff, 0x82, 0xff), (0xff, 0x82, 0xdf), (0xff, 0x82, 0xbe), (0xff, 0x82, 0x9e), + (0xff, 0x82, 0x82), (0xff, 0x9e, 0x82), (0xff, 0xbe, 0x82), (0xff, 0xdf, 0x82), + (0xff, 0xff, 0x82), (0xdf, 0xff, 0x82), (0xbe, 0xff, 0x82), (0x9e, 0xff, 0x82), + (0x82, 0xff, 0x82), (0x82, 0xff, 0x9e), (0x82, 0xff, 0xbe), (0x82, 0xff, 0xdf), + (0x82, 0xff, 0xff), (0x82, 0xdf, 0xff), (0x82, 0xbe, 0xff), (0x82, 0x9e, 0xff), + (0xba, 0xba, 0xff), (0xca, 0xba, 0xff), (0xdf, 0xba, 0xff), (0xef, 0xba, 0xff), + (0xff, 0xba, 0xff), (0xff, 0xba, 0xef), (0xff, 0xba, 0xdf), (0xff, 0xba, 0xca), + (0xff, 0xba, 0xba), (0xff, 0xca, 0xba), (0xff, 0xdf, 0xba), (0xff, 0xef, 0xba), + (0xff, 0xff, 0xba), (0xef, 0xff, 0xba), (0xdf, 0xff, 0xba), (0xca, 0xff, 0xba), + (0xba, 0xff, 0xba), (0xba, 0xff, 0xca), (0xba, 0xff, 0xdf), (0xba, 0xff, 0xef), + (0xba, 0xff, 0xff), (0xba, 0xef, 0xff), (0xba, 0xdf, 0xff), (0xba, 0xca, 0xff), + (0x00, 0x00, 0x71), (0x1c, 0x00, 0x71), (0x39, 0x00, 0x71), (0x55, 0x00, 0x71), + (0x71, 0x00, 0x71), (0x71, 0x00, 0x55), (0x71, 0x00, 0x39), (0x71, 0x00, 0x1c), + (0x71, 0x00, 0x00), (0x71, 0x1c, 0x00), (0x71, 0x39, 0x00), (0x71, 0x55, 0x00), + (0x71, 0x71, 0x00), (0x55, 0x71, 0x00), (0x39, 0x71, 0x00), (0x1c, 0x71, 0x00), + (0x00, 0x71, 0x00), (0x00, 0x71, 0x1c), (0x00, 0x71, 0x39), (0x00, 0x71, 0x55), + (0x00, 0x71, 0x71), (0x00, 0x55, 0x71), (0x00, 0x39, 0x71), (0x00, 0x1c, 0x71), + (0x39, 0x39, 0x71), (0x45, 0x39, 0x71), (0x55, 0x39, 0x71), (0x61, 0x39, 0x71), + (0x71, 0x39, 0x71), (0x71, 0x39, 0x61), (0x71, 0x39, 0x55), (0x71, 0x39, 0x45), + (0x71, 0x39, 0x39), (0x71, 0x45, 0x39), (0x71, 0x55, 0x39), (0x71, 0x61, 0x39), + (0x71, 0x71, 0x39), (0x61, 0x71, 0x39), (0x55, 0x71, 0x39), (0x45, 0x71, 0x39), + (0x39, 0x71, 0x39), (0x39, 0x71, 0x45), (0x39, 0x71, 0x55), (0x39, 0x71, 0x61), + (0x39, 0x71, 0x71), (0x39, 0x61, 0x71), (0x39, 0x55, 0x71), (0x39, 0x45, 0x71), + (0x51, 0x51, 0x71), (0x59, 0x51, 0x71), (0x61, 0x51, 0x71), (0x69, 0x51, 0x71), + (0x71, 0x51, 0x71), (0x71, 0x51, 0x69), (0x71, 0x51, 0x61), (0x71, 0x51, 0x59), + (0x71, 0x51, 0x51), (0x71, 0x59, 0x51), (0x71, 0x61, 0x51), (0x71, 0x69, 0x51), + (0x71, 0x71, 0x51), (0x69, 0x71, 0x51), (0x61, 0x71, 0x51), (0x59, 0x71, 0x51), + (0x51, 0x71, 0x51), (0x51, 0x71, 0x59), (0x51, 0x71, 0x61), (0x51, 0x71, 0x69), + (0x51, 0x71, 0x71), (0x51, 0x69, 0x71), (0x51, 0x61, 0x71), (0x51, 0x59, 0x71), + (0x00, 0x00, 0x41), (0x10, 0x00, 0x41), (0x20, 0x00, 0x41), (0x31, 0x00, 0x41), + (0x41, 0x00, 0x41), (0x41, 0x00, 0x31), (0x41, 0x00, 0x20), (0x41, 0x00, 0x10), + (0x41, 0x00, 0x00), (0x41, 0x10, 0x00), (0x41, 0x20, 0x00), (0x41, 0x31, 0x00), + (0x41, 0x41, 0x00), (0x31, 0x41, 0x00), (0x20, 0x41, 0x00), (0x10, 0x41, 0x00), + (0x00, 0x41, 0x00), (0x00, 0x41, 0x10), (0x00, 0x41, 0x20), (0x00, 0x41, 0x31), + (0x00, 0x41, 0x41), (0x00, 0x31, 0x41), (0x00, 0x20, 0x41), (0x00, 0x10, 0x41), + (0x20, 0x20, 0x41), (0x28, 0x20, 0x41), (0x31, 0x20, 0x41), (0x39, 0x20, 0x41), + (0x41, 0x20, 0x41), (0x41, 0x20, 0x39), (0x41, 0x20, 0x31), (0x41, 0x20, 0x28), + (0x41, 0x20, 0x20), (0x41, 0x28, 0x20), (0x41, 0x31, 0x20), (0x41, 0x39, 0x20), + (0x41, 0x41, 0x20), (0x39, 0x41, 0x20), (0x31, 0x41, 0x20), (0x28, 0x41, 0x20), + (0x20, 0x41, 0x20), (0x20, 0x41, 0x28), (0x20, 0x41, 0x31), (0x20, 0x41, 0x39), + (0x20, 0x41, 0x41), (0x20, 0x39, 0x41), (0x20, 0x31, 0x41), (0x20, 0x28, 0x41), + (0x2d, 0x2d, 0x41), (0x31, 0x2d, 0x41), (0x35, 0x2d, 0x41), (0x3d, 0x2d, 0x41), + (0x41, 0x2d, 0x41), (0x41, 0x2d, 0x3d), (0x41, 0x2d, 0x35), (0x41, 0x2d, 0x31), + (0x41, 0x2d, 0x2d), (0x41, 0x31, 0x2d), (0x41, 0x35, 0x2d), (0x41, 0x3d, 0x2d), + (0x41, 0x41, 0x2d), (0x3d, 0x41, 0x2d), (0x35, 0x41, 0x2d), (0x31, 0x41, 0x2d), + (0x2d, 0x41, 0x2d), (0x2d, 0x41, 0x31), (0x2d, 0x41, 0x35), (0x2d, 0x41, 0x3d), + (0x2d, 0x41, 0x41), (0x2d, 0x3d, 0x41), (0x2d, 0x35, 0x41), (0x2d, 0x31, 0x41), + (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), +]; diff --git a/src/rfb.rs b/src/rfb.rs index 123ac99..2659654 100644 --- a/src/rfb.rs +++ b/src/rfb.rs @@ -469,6 +469,10 @@ impl ReadMessage for PixelFormat { let mut buf = [0u8; 3]; stream.read_exact(&mut buf).await?; + if let ColorSpecification::ColorMap(..) = &color_spec { + todo!("SetColorMapEntries"); //.write_to(stream).await?; + } + Ok(Self { bits_per_pixel, depth, @@ -516,17 +520,16 @@ pub struct ColorFormat { } #[derive(Debug, Clone, PartialEq)] -pub struct ColorMap {} +pub struct ColorMap { + // we currently just use VESA_VGA_256_COLOR_PALETTE +} impl ReadMessage for ColorSpecification { fn read_from<'a>(stream: &'a mut TcpStream) -> BoxFuture<'a, Result> { async { let tc_flag = stream.read_u8().await?; match tc_flag { - 0 => { - // ColorMap - unimplemented!() - } + 0 => Ok(ColorSpecification::ColorMap(ColorMap {})), _ => { // ColorFormat let red_max = stream.read_u16().await?; @@ -567,7 +570,10 @@ impl WriteMessage for ColorSpecification { stream.write_u8(cf.blue_shift).await?; } ColorSpecification::ColorMap(_cm) => { - unimplemented!() + // first 0 byte is true-color-flag = false; + // the remaining 9 are the above max/shift fields, + // which aren't relevant in ColorMap mode + stream.write_all(&[0u8; 10]).await?; } }; From 024905e23993ab5b0d9352cb1bd83c67b4c653c4 Mon Sep 17 00:00:00 2001 From: lif <> Date: Fri, 9 Aug 2024 05:31:51 +0000 Subject: [PATCH 08/20] wip - starting on hextile --- src/encodings.rs | 55 +++++++++++++++++++++++++++++++++++++++++++++--- src/rfb.rs | 2 +- 2 files changed, 53 insertions(+), 4 deletions(-) diff --git a/src/encodings.rs b/src/encodings.rs index 9d6539a..933e22e 100644 --- a/src/encodings.rs +++ b/src/encodings.rs @@ -38,7 +38,7 @@ where fn get_type(&self) -> EncodingType; /// Transform this encoding from its representation into a byte vector that can be passed to the client. - fn encode(&self) -> &Vec; + fn encode(&self) -> Vec; /// Translates this encoding type from an input pixel format to an output format. fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box; @@ -102,8 +102,8 @@ impl Encoding for RawEncoding { Raw } - fn encode(&self) -> &Vec { - &self.pixels + fn encode(&self) -> Vec { + self.pixels.clone() } fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box { @@ -147,6 +147,42 @@ struct HextileEncoding { tiles: Vec>, } +impl HextileEncoding { + pub fn from_raw(raw: Vec) -> Self { + Self { + tiles: todo!("create subrects"), + } + } +} + +bitflags::bitflags! { + pub struct HextileSubencMask: u8 { + const RAW = 1 << 0; + const BACKGROUND_SPECIFIED = 1 << 1; + const FOREGROUND_SPECIFIED = 1 << 2; + const ANY_SUBRECTS = 1 << 3; + const SUBRECTS_COLORED = 1 << 4; + } +} + +impl Encoding for HextileEncoding { + fn get_type(&self) -> EncodingType { + EncodingType::Hextile + } + + fn encode(&self) -> Vec { + let mut v = Vec::new(); + for tile in self.tiles { + v.push(subencoding_mask) + } + v + } + + fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box { + todo!() + } +} + #[allow(dead_code)] enum HextileTile { Raw(Vec), @@ -159,3 +195,16 @@ struct HextileTileEncoded { foreground: Option, // TODO: finish this } + +impl HextileTileEncoded { + fn subenc_mask(&self) -> HextileSubencMask { + let mut x = HextileSubencMask::empty(); + if self.background.is_some() { + x |= HextileSubencMask::BACKGROUND_SPECIFIED; + } + if self.foreground.is_some() { + x |= HextileSubencMask::FOREGROUND_SPECIFIED; + } + x + } +} diff --git a/src/rfb.rs b/src/rfb.rs index 2659654..7a9d2a1 100644 --- a/src/rfb.rs +++ b/src/rfb.rs @@ -326,7 +326,7 @@ impl WriteMessage for Rectangle { stream.write_i32(encoding_type).await?; let data = self.data.encode(); - stream.write_all(data).await?; + stream.write_all(&data).await?; Ok(()) } From 50ce328fe18284f44e773b1d08f435aff46052b4 Mon Sep 17 00:00:00 2001 From: lif <> Date: Tue, 13 Aug 2024 02:48:25 +0000 Subject: [PATCH 09/20] start on TRLE, give each encoding its own module --- src/encodings.rs | 142 ++++----------------------------------- src/encodings/hextile.rs | 72 ++++++++++++++++++++ src/encodings/raw.rs | 43 ++++++++++++ src/encodings/rre.rs | 67 ++++++++++++++++++ src/encodings/trle.rs | 82 ++++++++++++++++++++++ src/rfb.rs | 4 +- 6 files changed, 280 insertions(+), 130 deletions(-) create mode 100644 src/encodings/hextile.rs create mode 100644 src/encodings/raw.rs create mode 100644 src/encodings/rre.rs create mode 100644 src/encodings/trle.rs diff --git a/src/encodings.rs b/src/encodings.rs index 933e22e..2f9c988 100644 --- a/src/encodings.rs +++ b/src/encodings.rs @@ -4,14 +4,19 @@ // // Copyright 2022 Oxide Computer Company -use crate::{ - pixel_formats::transform, - rfb::{PixelFormat, Position, Resolution}, -}; +use crate::rfb::PixelFormat; -use crate::rfb::ColorSpecification; use EncodingType::*; +#[allow(unused)] +mod hextile; +mod raw; +#[allow(unused)] +mod rre; +mod trle; + +pub use raw::RawEncoding; + #[derive(Debug)] #[allow(unused)] pub enum EncodingType { @@ -31,6 +36,10 @@ pub enum EncodingType { Other(i32), } +struct Pixel { + bytes: Vec, +} + pub trait Encoding where Self: Send, @@ -85,126 +94,3 @@ impl From for EncodingType { } } } - -/// Section 7.7.1 -pub struct RawEncoding { - pixels: Vec, -} - -impl RawEncoding { - pub fn new(pixels: Vec) -> Self { - Self { pixels } - } -} - -impl Encoding for RawEncoding { - fn get_type(&self) -> EncodingType { - Raw - } - - fn encode(&self) -> Vec { - self.pixels.clone() - } - - fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box { - // XXX: This assumes the pixel formats are both rgb. The server code verifies this - // before calling. - assert!(matches!( - &input.color_spec, - ColorSpecification::ColorFormat(_) - )); - assert!(matches!( - &output.color_spec, - ColorSpecification::ColorFormat(_) - )); - - Box::new(Self { - pixels: transform(&self.pixels, &input, &output), - }) - } -} - -#[allow(dead_code)] -struct RREncoding { - background_pixel: Pixel, - sub_rectangles: Vec, -} - -#[allow(dead_code)] -struct Pixel { - bytes: Vec, -} - -#[allow(dead_code)] -struct RRESubrectangle { - pixel: Pixel, - position: Position, - dimensions: Resolution, -} - -#[allow(dead_code)] -struct HextileEncoding { - tiles: Vec>, -} - -impl HextileEncoding { - pub fn from_raw(raw: Vec) -> Self { - Self { - tiles: todo!("create subrects"), - } - } -} - -bitflags::bitflags! { - pub struct HextileSubencMask: u8 { - const RAW = 1 << 0; - const BACKGROUND_SPECIFIED = 1 << 1; - const FOREGROUND_SPECIFIED = 1 << 2; - const ANY_SUBRECTS = 1 << 3; - const SUBRECTS_COLORED = 1 << 4; - } -} - -impl Encoding for HextileEncoding { - fn get_type(&self) -> EncodingType { - EncodingType::Hextile - } - - fn encode(&self) -> Vec { - let mut v = Vec::new(); - for tile in self.tiles { - v.push(subencoding_mask) - } - v - } - - fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box { - todo!() - } -} - -#[allow(dead_code)] -enum HextileTile { - Raw(Vec), - Encoded(HextileTileEncoded), -} - -#[allow(dead_code)] -struct HextileTileEncoded { - background: Option, - foreground: Option, - // TODO: finish this -} - -impl HextileTileEncoded { - fn subenc_mask(&self) -> HextileSubencMask { - let mut x = HextileSubencMask::empty(); - if self.background.is_some() { - x |= HextileSubencMask::BACKGROUND_SPECIFIED; - } - if self.foreground.is_some() { - x |= HextileSubencMask::FOREGROUND_SPECIFIED; - } - x - } -} diff --git a/src/encodings/hextile.rs b/src/encodings/hextile.rs new file mode 100644 index 0000000..11e3634 --- /dev/null +++ b/src/encodings/hextile.rs @@ -0,0 +1,72 @@ +use crate::{ + encodings::{Encoding, EncodingType, Pixel}, + rfb::PixelFormat, +}; + +#[allow(dead_code)] +struct HextileEncoding { + tiles: Vec>, +} + +impl HextileEncoding { + pub fn from_raw(raw: Vec) -> Self { + Self { + tiles: todo!("create subrects. need dimensions"), + } + } +} + +bitflags::bitflags! { + pub struct HextileSubencMask: u8 { + const RAW = 1 << 0; + const BACKGROUND_SPECIFIED = 1 << 1; + const FOREGROUND_SPECIFIED = 1 << 2; + const ANY_SUBRECTS = 1 << 3; + const SUBRECTS_COLORED = 1 << 4; + } +} + +impl Encoding for HextileEncoding { + fn get_type(&self) -> EncodingType { + EncodingType::Hextile + } + + fn encode(&self) -> Vec { + let mut v = Vec::new(); + for tile in &self.tiles { + let subencoding_mask = todo!(); + v.push(subencoding_mask) + } + v + } + + fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box { + todo!() + } +} + +#[allow(dead_code)] +enum HextileTile { + Raw(Vec), + Encoded(HextileTileEncoded), +} + +#[allow(dead_code)] +struct HextileTileEncoded { + background: Option, + foreground: Option, + // TODO: finish this +} + +impl HextileTileEncoded { + fn subenc_mask(&self) -> HextileSubencMask { + let mut x = HextileSubencMask::empty(); + if self.background.is_some() { + x |= HextileSubencMask::BACKGROUND_SPECIFIED; + } + if self.foreground.is_some() { + x |= HextileSubencMask::FOREGROUND_SPECIFIED; + } + x + } +} diff --git a/src/encodings/raw.rs b/src/encodings/raw.rs new file mode 100644 index 0000000..1631940 --- /dev/null +++ b/src/encodings/raw.rs @@ -0,0 +1,43 @@ +use crate::{ + encodings::{Encoding, EncodingType}, + pixel_formats::transform, + rfb::{ColorSpecification, PixelFormat}, +}; + +/// Section 7.7.1 +pub struct RawEncoding { + pixels: Vec, +} + +impl RawEncoding { + pub fn new(pixels: Vec) -> Self { + Self { pixels } + } +} + +impl Encoding for RawEncoding { + fn get_type(&self) -> EncodingType { + EncodingType::Raw + } + + fn encode(&self) -> Vec { + self.pixels.clone() + } + + fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box { + // XXX: This assumes the pixel formats are both rgb. The server code verifies this + // before calling. + assert!(matches!( + &input.color_spec, + ColorSpecification::ColorFormat(_) + )); + assert!(matches!( + &output.color_spec, + ColorSpecification::ColorFormat(_) + )); + + Box::new(Self { + pixels: transform(&self.pixels, &input, &output), + }) + } +} diff --git a/src/encodings/rre.rs b/src/encodings/rre.rs new file mode 100644 index 0000000..fa84de2 --- /dev/null +++ b/src/encodings/rre.rs @@ -0,0 +1,67 @@ +use crate::{ + encodings::{Encoding, EncodingType, Pixel}, + pixel_formats::transform, + rfb::{PixelFormat, Position, Resolution}, +}; + +struct RREncoding { + background_pixel: Pixel, + sub_rectangles: Vec, +} + +struct RRESubrectangle { + pixel: Pixel, + position: Position, + dimensions: Resolution, +} + +impl Encoding for RREncoding { + fn get_type(&self) -> EncodingType { + EncodingType::RRE + } + + fn encode(&self) -> Vec { + let mut buf = Vec::with_capacity( + (self.sub_rectangles.len() + 1) * (self.background_pixel.bytes.len() + 8) - 4, + ); + buf.extend_from_slice(&(self.sub_rectangles.len() as u32).to_be_bytes()); + buf.extend_from_slice(&self.background_pixel.bytes); + for sr in &self.sub_rectangles { + buf.extend_from_slice(&sr.pixel.bytes); + buf.extend_from_slice(&sr.position.x.to_be_bytes()); + buf.extend_from_slice(&sr.position.y.to_be_bytes()); + buf.extend_from_slice(&sr.dimensions.width.to_be_bytes()); + buf.extend_from_slice(&sr.dimensions.height.to_be_bytes()); + } + buf + } + + fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box { + let background_pixel = Pixel { + bytes: transform(&self.background_pixel.bytes, input, output), + }; + let sub_rectangles = self + .sub_rectangles + .iter() + .map(|sr| { + let RRESubrectangle { + pixel, + position, + dimensions, + } = sr; + let pixel = Pixel { + bytes: transform(&pixel.bytes, input, output), + }; + RRESubrectangle { + pixel, + position: *position, + dimensions: *dimensions, + } + }) + .collect(); + Box::new(Self { + background_pixel, + sub_rectangles, + }) + } +} diff --git a/src/encodings/trle.rs b/src/encodings/trle.rs new file mode 100644 index 0000000..5f2761b --- /dev/null +++ b/src/encodings/trle.rs @@ -0,0 +1,82 @@ +use crate::encodings::{Encoding, EncodingType}; +use crate::rfb::PixelFormat; + +pub struct TRLEncoding { + tiles: Vec>, +} + +#[repr(transparent)] +struct PackedIndeces(u8); + +// impl From<&[u8; 2]> for PackedIndeces { +// fn from(&[left, right]: &[u8; 2]) -> Self { +// Self((left << 4) | (right & 0xF)) +// } +// } +impl PackedIndeces { + // fn new(indeces: &[u8; N]) -> Self { + // const BPP: usize = 16 / size_of; + // } + fn new_4bpp(&[left, right]: &[u8; 2]) -> Self { + Self((left << 4) | (right & 0xF)) + } + fn new_2bpp(indeces: &[u8]) -> Self { + let mut x = 0; + assert!(indeces.len() <= 4); + for (pos, ci) in indeces.iter().copied().enumerate() { + x |= ((ci << 6) & 0xC0) >> (pos * 2); + } + Self(x) + } + fn new_1bpp(indeces: &[u8]) -> Self { + let mut x = 0; + assert!(indeces.len() <= 16); + for (pos, ci) in indeces.iter().copied().enumerate() { + x |= ((ci << 7) & 0x80) >> pos; + } + Self(x) + } +} + +// may be able to reuse this for ZRLE? (64px instead of 16px) +enum TRLETile { + /// 0 + Raw { pixels: Vec }, + /// 1 + SolidColor { color: CPixel }, + /// 2-16 + PackedPalette { + palette: Vec, + packed_pixels: Vec, + }, + /// 127 + PackedPaletteReused { packed_pixels: Vec }, + /// 128 + PlainRLE { color: CPixel, length: usize }, + /// 129 + PaletteRLEReused { pixels: Vec }, + /// 130-255 + PaletteRLE { + palette: Vec, + pixels: Vec, + }, +} + +struct CPixel { + format: PixelFormat, + bytes: Vec, +} + +impl Encoding for TRLEncoding { + fn get_type(&self) -> EncodingType { + EncodingType::TRLE + } + + fn encode(&self) -> Vec { + todo!() + } + + fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box { + todo!() + } +} diff --git a/src/rfb.rs b/src/rfb.rs index 7a9d2a1..228e71d 100644 --- a/src/rfb.rs +++ b/src/rfb.rs @@ -263,8 +263,8 @@ impl ReadMessage for Position { #[derive(Debug, Copy, Clone)] pub(crate) struct Resolution { - width: u16, - height: u16, + pub width: u16, + pub height: u16, } impl ReadMessage for Resolution { From 9af7eb1d147f32d54fa673de9dc993c90fe6124e Mon Sep 17 00:00:00 2001 From: lif <> Date: Thu, 15 Aug 2024 04:44:23 +0000 Subject: [PATCH 10/20] implement subencodings of TRLE --- src/encodings/trle.rs | 80 ++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 76 insertions(+), 4 deletions(-) diff --git a/src/encodings/trle.rs b/src/encodings/trle.rs index 5f2761b..cc865f6 100644 --- a/src/encodings/trle.rs +++ b/src/encodings/trle.rs @@ -1,3 +1,5 @@ +use std::iter::{from_fn, once}; + use crate::encodings::{Encoding, EncodingType}; use crate::rfb::PixelFormat; @@ -52,16 +54,83 @@ enum TRLETile { /// 127 PackedPaletteReused { packed_pixels: Vec }, /// 128 - PlainRLE { color: CPixel, length: usize }, + PlainRLE { runs: Vec<(CPixel, usize)> }, /// 129 - PaletteRLEReused { pixels: Vec }, + PaletteRLEReused { runs: Vec<(u8, usize)> }, /// 130-255 PaletteRLE { palette: Vec, - pixels: Vec, + runs: Vec<(u8, usize)>, }, } +fn rle(mut length: usize) -> impl Iterator { + from_fn(move || { + if length == 0 { + None + } else if length > 0xFF { + length -= 0xFF; + Some(0xFF) + } else { + let byte = (length - 1) as u8; + length = 0; + Some(byte) + } + }) +} + +fn pal_rle((index, length): &(u8, usize)) -> Box> { + if *length == 1 { + Box::new(once(*index)) + } else { + Box::new(once(*index | 0x80).chain(rle(*length))) + } +} + +impl TRLETile { + /// Subencoding of the tile according to RFB 6143 7.7.5 + fn encode(&self) -> Box + '_> { + match self { + TRLETile::Raw { pixels } => { + Box::new(once(0u8).chain(pixels.iter().flat_map(|c| c.bytes.iter().copied()))) + } + TRLETile::SolidColor { color } => { + Box::new(once(1u8).chain(color.bytes.iter().copied())) + } + TRLETile::PackedPalette { + palette, + packed_pixels, + } => Box::new( + once(palette.len() as u8) + .chain(palette.iter().flat_map(|c| c.bytes.iter().copied())) + .chain(packed_pixels.iter().map(|p| p.0)), + ), + TRLETile::PackedPaletteReused { packed_pixels } => { + Box::new(once(127u8).chain(packed_pixels.iter().map(|p| p.0))) + } + TRLETile::PlainRLE { runs } => { + Box::new(once(128).chain( + runs.iter().flat_map(|(color, length)| { + color.bytes.iter().copied().chain(rle(*length)) + }), + )) + } + TRLETile::PaletteRLEReused { runs } => { + Box::new(once(129).chain(runs.iter().flat_map(pal_rle))) + } + TRLETile::PaletteRLE { palette, runs } => Box::new( + once( + (palette.len() + 128) + .try_into() + .expect("TRLE tile palette too large!"), + ) + .chain(palette.iter().flat_map(|c| c.bytes.iter().copied())) + .chain(runs.iter().flat_map(pal_rle)), + ), + } + } +} + struct CPixel { format: PixelFormat, bytes: Vec, @@ -73,7 +142,10 @@ impl Encoding for TRLEncoding { } fn encode(&self) -> Vec { - todo!() + self.tiles + .iter() + .flat_map(|row| row.iter().flat_map(|tile| tile.encode())) + .collect() } fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box { From b0e05794092e17749abb7dc33628884837c3213b Mon Sep 17 00:00:00 2001 From: lif <> Date: Thu, 22 Aug 2024 20:41:07 +0000 Subject: [PATCH 11/20] trle transform --- src/encodings/trle.rs | 103 +++++++++++++++++++++++++++++++++++++++++- src/rfb.rs | 20 ++++++++ 2 files changed, 121 insertions(+), 2 deletions(-) diff --git a/src/encodings/trle.rs b/src/encodings/trle.rs index cc865f6..07dee46 100644 --- a/src/encodings/trle.rs +++ b/src/encodings/trle.rs @@ -1,13 +1,21 @@ use std::iter::{from_fn, once}; use crate::encodings::{Encoding, EncodingType}; +use crate::pixel_formats; use crate::rfb::PixelFormat; pub struct TRLEncoding { tiles: Vec>, } +impl TRLEncoding { + pub fn new(pixels: Vec, width: usize, height: usize) { + todo!() + } +} + #[repr(transparent)] +#[derive(Copy, Clone)] struct PackedIndeces(u8); // impl From<&[u8; 2]> for PackedIndeces { @@ -41,6 +49,7 @@ impl PackedIndeces { } // may be able to reuse this for ZRLE? (64px instead of 16px) +#[derive(Clone)] enum TRLETile { /// 0 Raw { pixels: Vec }, @@ -131,11 +140,57 @@ impl TRLETile { } } +// TODO: [u8; 4] so we can derive Copy and go fast +#[derive(Clone)] struct CPixel { - format: PixelFormat, bytes: Vec, } +enum CPixelTransformType { + AsIs, + AppendZero, + PrependZero, +} + +impl CPixel { + fn asdf(pixfmt: &PixelFormat) -> CPixelTransformType { + if pixfmt.depth <= 24 && pixfmt.bits_per_pixel == 32 { + let mask = pixfmt + .value_mask() + .expect("colormap not supported in cpixel"); + let should_append = if mask.trailing_zeros() >= 8 { + false + } else if mask.leading_zeros() >= 8 { + true + } else { + return CPixelTransformType::AsIs; + } ^ pixfmt.big_endian; + if should_append { + CPixelTransformType::AppendZero + } else { + CPixelTransformType::PrependZero + } + } else { + CPixelTransformType::AsIs + } + } + + fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Self { + let in_bytes = match Self::asdf(input) { + CPixelTransformType::AsIs => &self.bytes, + CPixelTransformType::AppendZero => &self.bytes[0..=2], + CPixelTransformType::PrependZero => &self.bytes[1..=3], + }; + let mut out_bytes = pixel_formats::transform(&in_bytes, input, output); + match Self::asdf(output) { + CPixelTransformType::AsIs => (), + CPixelTransformType::AppendZero => out_bytes.push(0u8), + CPixelTransformType::PrependZero => out_bytes.insert(0, 0u8), + } + Self { bytes: out_bytes } + } +} + impl Encoding for TRLEncoding { fn get_type(&self) -> EncodingType { EncodingType::TRLE @@ -149,6 +204,50 @@ impl Encoding for TRLEncoding { } fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box { - todo!() + let tiles = self + .tiles + .iter() + .map(|row| { + row.iter() + .map(|tile| match tile { + TRLETile::Raw { pixels } => TRLETile::Raw { + pixels: pixels + .iter() + .map(|cp| cp.transform(input, output)) + .collect(), + }, + TRLETile::SolidColor { color } => TRLETile::SolidColor { + color: color.transform(input, output), + }, + TRLETile::PackedPalette { + palette, + packed_pixels, + } => TRLETile::PackedPalette { + palette: palette + .iter() + .map(|cp| cp.transform(input, output)) + .collect(), + packed_pixels: packed_pixels.clone(), + }, + TRLETile::PlainRLE { runs } => TRLETile::PlainRLE { + runs: runs + .iter() + .map(|(cp, len)| (cp.transform(input, output), *len)) + .collect(), + }, + TRLETile::PaletteRLE { palette, runs } => TRLETile::PaletteRLE { + palette: palette + .iter() + .map(|cp| cp.transform(input, output)) + .collect(), + runs: runs.clone(), + }, + TRLETile::PackedPaletteReused { .. } + | TRLETile::PaletteRLEReused { .. } => tile.clone(), + }) + .collect() + }) + .collect(); + Box::new(Self { tiles }) } } diff --git a/src/rfb.rs b/src/rfb.rs index 228e71d..0f9ce42 100644 --- a/src/rfb.rs +++ b/src/rfb.rs @@ -4,6 +4,8 @@ // // Copyright 2022 Oxide Computer Company +use std::ops::BitOr; + use bitflags::bitflags; use futures::future::BoxFuture; use futures::FutureExt; @@ -415,6 +417,24 @@ impl PixelFormat { } } + pub fn value_mask(&self) -> Option { + match self.color_spec { + ColorSpecification::ColorFormat(ColorFormat { + red_max, + green_max, + blue_max, + red_shift, + green_shift, + blue_shift, + }) => Some( + ((red_max as u64) << red_shift) + | ((green_max as u64) << green_shift) + | ((blue_max as u64) << blue_shift), + ), + ColorSpecification::ColorMap(_) => None, + } + } + /// Returns true if the pixel format is RGB888 (8-bits per color and 32 bits per pixel). #[deprecated] pub fn is_rgb_888(&self) -> bool { From ed7d574de556623d4741ee7b8e9a2556c7a216bf Mon Sep 17 00:00:00 2001 From: lif <> Date: Fri, 23 Aug 2024 02:55:46 +0000 Subject: [PATCH 12/20] changing the trait and bumping to 0.2.0 accordingly because if we keep copying huge Vecs everywhere we're gonna have a *lot* of resource churn --- Cargo.toml | 2 +- examples/server.rs | 7 +++++- src/encodings.rs | 14 ++++++++---- src/encodings/hextile.rs | 34 ++++++++++++++++++++-------- src/encodings/raw.rs | 38 ++++++++++++++++++++++++++----- src/encodings/rre.rs | 48 +++++++++++++++++++++++++++------------ src/encodings/trle.rs | 49 +++++++++++++++++++++++++++++----------- src/rfb.rs | 9 ++++---- 8 files changed, 148 insertions(+), 53 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index f9b639f..299f10d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "rfb" -version = "0.1.0" +version = "0.2.0" description = "Implementation of the RFB protocol (RFC 6143), including a server implementation." repository = "https://github.com/oxidecomputer/rfb" readme = "README.md" diff --git a/examples/server.rs b/examples/server.rs index 7a8ab1c..9695a54 100644 --- a/examples/server.rs +++ b/examples/server.rs @@ -162,7 +162,12 @@ impl Server for ExampleServer { 0, pixels_width, pixels_height, - Box::new(RawEncoding::new(pixels)), + Box::new(RawEncoding::new( + pixels, + pixels_width, + pixels_height, + &self.pixfmt, + )), ); FramebufferUpdate::new(vec![r]) } diff --git a/src/encodings.rs b/src/encodings.rs index 2f9c988..8dcc86a 100644 --- a/src/encodings.rs +++ b/src/encodings.rs @@ -46,11 +46,17 @@ where { fn get_type(&self) -> EncodingType; - /// Transform this encoding from its representation into a byte vector that can be passed to the client. - fn encode(&self) -> Vec; + /// Return the width and height in pixels of the encoded screen region. + fn dimensions(&self) -> (u16, u16); - /// Translates this encoding type from an input pixel format to an output format. - fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box; + /// Return the pixel format of this encoding's data. + fn pixel_format(&self) -> &PixelFormat; + + /// Transform this encoding from its representation into a byte sequence that can be passed to the client. + fn encode(&self) -> Box + '_>; + + /// Translates this encoding type from its current pixel format to the given format. + fn transform(&self, output: &PixelFormat) -> Box; } impl From for i32 { diff --git a/src/encodings/hextile.rs b/src/encodings/hextile.rs index 11e3634..9012d19 100644 --- a/src/encodings/hextile.rs +++ b/src/encodings/hextile.rs @@ -3,15 +3,25 @@ use crate::{ rfb::PixelFormat, }; +use super::RawEncoding; + #[allow(dead_code)] struct HextileEncoding { tiles: Vec>, + width: u16, + height: u16, + pixfmt: PixelFormat, } -impl HextileEncoding { - pub fn from_raw(raw: Vec) -> Self { +impl From<&RawEncoding> for HextileEncoding { + fn from(raw: &RawEncoding) -> Self { + let pixfmt = raw.pixel_format().to_owned(); + let (width, height) = raw.dimensions(); Self { tiles: todo!("create subrects. need dimensions"), + width, + height, + pixfmt, } } } @@ -31,18 +41,24 @@ impl Encoding for HextileEncoding { EncodingType::Hextile } - fn encode(&self) -> Vec { - let mut v = Vec::new(); - for tile in &self.tiles { + fn encode(&self) -> Box + '_> { + Box::new(self.tiles.iter().flat_map(|tile| { let subencoding_mask = todo!(); - v.push(subencoding_mask) - } - v + [todo!()].into_iter() + })) } - fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box { + fn transform(&self, output: &PixelFormat) -> Box { todo!() } + + fn dimensions(&self) -> (u16, u16) { + (self.width, self.height) + } + + fn pixel_format(&self) -> &PixelFormat { + &self.pixfmt + } } #[allow(dead_code)] diff --git a/src/encodings/raw.rs b/src/encodings/raw.rs index 1631940..2c0e30d 100644 --- a/src/encodings/raw.rs +++ b/src/encodings/raw.rs @@ -6,12 +6,25 @@ use crate::{ /// Section 7.7.1 pub struct RawEncoding { - pixels: Vec, + pixels: Vec, // TODO: self-lifetime-bound slice to avoid copy + width: u16, + height: u16, + pixfmt: PixelFormat, } impl RawEncoding { - pub fn new(pixels: Vec) -> Self { - Self { pixels } + pub fn new(pixels: Vec, width: u16, height: u16, pixfmt: &PixelFormat) -> Self { + Self { + pixels, + width, + height, + pixfmt: pixfmt.clone(), + } + } + + // useful for transforming into other encodings + pub(crate) fn raw_buffer(&self) -> &[u8] { + &self.pixels } } @@ -20,11 +33,21 @@ impl Encoding for RawEncoding { EncodingType::Raw } - fn encode(&self) -> Vec { - self.pixels.clone() + fn dimensions(&self) -> (u16, u16) { + (self.width, self.height) + } + + fn pixel_format(&self) -> &PixelFormat { + &self.pixfmt } - fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box { + fn encode(&self) -> Box + '_> { + Box::new(self.pixels.iter().copied()) + } + + fn transform(&self, output: &PixelFormat) -> Box { + let input = &self.pixfmt; + // XXX: This assumes the pixel formats are both rgb. The server code verifies this // before calling. assert!(matches!( @@ -38,6 +61,9 @@ impl Encoding for RawEncoding { Box::new(Self { pixels: transform(&self.pixels, &input, &output), + width: self.width, + height: self.height, + pixfmt: output.clone(), }) } } diff --git a/src/encodings/rre.rs b/src/encodings/rre.rs index fa84de2..ea497be 100644 --- a/src/encodings/rre.rs +++ b/src/encodings/rre.rs @@ -7,6 +7,9 @@ use crate::{ struct RREncoding { background_pixel: Pixel, sub_rectangles: Vec, + width: u16, + height: u16, + pixfmt: PixelFormat, } struct RRESubrectangle { @@ -20,23 +23,27 @@ impl Encoding for RREncoding { EncodingType::RRE } - fn encode(&self) -> Vec { - let mut buf = Vec::with_capacity( - (self.sub_rectangles.len() + 1) * (self.background_pixel.bytes.len() + 8) - 4, - ); - buf.extend_from_slice(&(self.sub_rectangles.len() as u32).to_be_bytes()); - buf.extend_from_slice(&self.background_pixel.bytes); - for sr in &self.sub_rectangles { - buf.extend_from_slice(&sr.pixel.bytes); - buf.extend_from_slice(&sr.position.x.to_be_bytes()); - buf.extend_from_slice(&sr.position.y.to_be_bytes()); - buf.extend_from_slice(&sr.dimensions.width.to_be_bytes()); - buf.extend_from_slice(&sr.dimensions.height.to_be_bytes()); - } - buf + fn encode(&self) -> Box + '_> { + Box::new( + (self.sub_rectangles.len() as u32) + .to_be_bytes() + .into_iter() + .chain(self.background_pixel.bytes.iter().copied()) + .chain(self.sub_rectangles.iter().flat_map(|sr| { + sr.pixel + .bytes + .iter() + .copied() + .chain(sr.position.x.to_be_bytes().into_iter()) + .chain(sr.position.y.to_be_bytes().into_iter()) + .chain(sr.dimensions.width.to_be_bytes().into_iter()) + .chain(sr.dimensions.height.to_be_bytes().into_iter()) + })), + ) } - fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box { + fn transform(&self, output: &PixelFormat) -> Box { + let input = &self.pixfmt; let background_pixel = Pixel { bytes: transform(&self.background_pixel.bytes, input, output), }; @@ -62,6 +69,17 @@ impl Encoding for RREncoding { Box::new(Self { background_pixel, sub_rectangles, + width: self.width, + height: self.height, + pixfmt: output.to_owned(), }) } + + fn dimensions(&self) -> (u16, u16) { + (self.width, self.height) + } + + fn pixel_format(&self) -> &PixelFormat { + &self.pixfmt + } } diff --git a/src/encodings/trle.rs b/src/encodings/trle.rs index 07dee46..d79484d 100644 --- a/src/encodings/trle.rs +++ b/src/encodings/trle.rs @@ -4,12 +4,18 @@ use crate::encodings::{Encoding, EncodingType}; use crate::pixel_formats; use crate::rfb::PixelFormat; +use super::RawEncoding; + pub struct TRLEncoding { tiles: Vec>, + width: u16, + height: u16, + pixfmt: PixelFormat, } -impl TRLEncoding { - pub fn new(pixels: Vec, width: usize, height: usize) { +impl From<&RawEncoding> for TRLEncoding { + fn from(raw: &RawEncoding) -> Self { + raw.raw_buffer(); todo!() } } @@ -97,7 +103,9 @@ fn pal_rle((index, length): &(u8, usize)) -> Box> { } impl TRLETile { - /// Subencoding of the tile according to RFB 6143 7.7.5 + /// Subencoding of the tile according to RFB 6143 7.7.5. + /// To the extent possible, this function is a translation of that + /// section of the RFB RFC from English into chained iterators. fn encode(&self) -> Box + '_> { match self { TRLETile::Raw { pixels } => { @@ -153,7 +161,7 @@ enum CPixelTransformType { } impl CPixel { - fn asdf(pixfmt: &PixelFormat) -> CPixelTransformType { + fn which_padding(pixfmt: &PixelFormat) -> CPixelTransformType { if pixfmt.depth <= 24 && pixfmt.bits_per_pixel == 32 { let mask = pixfmt .value_mask() @@ -176,13 +184,13 @@ impl CPixel { } fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Self { - let in_bytes = match Self::asdf(input) { + let in_bytes = match Self::which_padding(input) { CPixelTransformType::AsIs => &self.bytes, CPixelTransformType::AppendZero => &self.bytes[0..=2], CPixelTransformType::PrependZero => &self.bytes[1..=3], }; let mut out_bytes = pixel_formats::transform(&in_bytes, input, output); - match Self::asdf(output) { + match Self::which_padding(output) { CPixelTransformType::AsIs => (), CPixelTransformType::AppendZero => out_bytes.push(0u8), CPixelTransformType::PrependZero => out_bytes.insert(0, 0u8), @@ -196,14 +204,16 @@ impl Encoding for TRLEncoding { EncodingType::TRLE } - fn encode(&self) -> Vec { - self.tiles - .iter() - .flat_map(|row| row.iter().flat_map(|tile| tile.encode())) - .collect() + fn encode(&self) -> Box + '_> { + Box::new( + self.tiles + .iter() + .flat_map(|row| row.iter().flat_map(|tile| tile.encode())), + ) } - fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box { + fn transform(&self, output: &PixelFormat) -> Box { + let input = &self.pixfmt; let tiles = self .tiles .iter() @@ -248,6 +258,19 @@ impl Encoding for TRLEncoding { .collect() }) .collect(); - Box::new(Self { tiles }) + Box::new(Self { + tiles, + width: self.width, + height: self.height, + pixfmt: output.to_owned(), + }) + } + + fn dimensions(&self) -> (u16, u16) { + (self.width, self.height) + } + + fn pixel_format(&self) -> &PixelFormat { + &self.pixfmt } } diff --git a/src/rfb.rs b/src/rfb.rs index 0f9ce42..7adf524 100644 --- a/src/rfb.rs +++ b/src/rfb.rs @@ -4,8 +4,6 @@ // // Copyright 2022 Oxide Computer Company -use std::ops::BitOr; - use bitflags::bitflags; use futures::future::BoxFuture; use futures::FutureExt; @@ -308,10 +306,12 @@ impl Rectangle { } pub fn transform(&self, input_pf: &PixelFormat, output_pf: &PixelFormat) -> Self { + // TODO: refactor out of method args here? + assert_eq!(input_pf, self.data.pixel_format()); Rectangle { position: self.position, dimensions: self.dimensions, - data: self.data.transform(input_pf, output_pf), + data: self.data.transform(output_pf), } } } @@ -327,7 +327,8 @@ impl WriteMessage for Rectangle { stream.write_u16(self.dimensions.height).await?; stream.write_i32(encoding_type).await?; - let data = self.data.encode(); + // TODO: avoid collect alloc? + let data: Vec<_> = self.data.encode().collect(); stream.write_all(&data).await?; Ok(()) From fd2d8e68de61280ec2a91719c9af4dc08a1a9eef Mon Sep 17 00:00:00 2001 From: lif <> Date: Fri, 23 Aug 2024 03:46:38 +0000 Subject: [PATCH 13/20] borrowed-data type for representing raw framebuffer data with a known size and format internally --- src/encodings.rs | 1 + src/encodings/hextile.rs | 6 +-- src/encodings/raw.rs | 92 +++++++++++++++++++++++++++++++--------- src/encodings/trle.rs | 6 +-- 4 files changed, 80 insertions(+), 25 deletions(-) diff --git a/src/encodings.rs b/src/encodings.rs index 8dcc86a..67d8a15 100644 --- a/src/encodings.rs +++ b/src/encodings.rs @@ -16,6 +16,7 @@ mod rre; mod trle; pub use raw::RawEncoding; +pub use raw::RawEncodingRef; #[derive(Debug)] #[allow(unused)] diff --git a/src/encodings/hextile.rs b/src/encodings/hextile.rs index 9012d19..3633be0 100644 --- a/src/encodings/hextile.rs +++ b/src/encodings/hextile.rs @@ -3,7 +3,7 @@ use crate::{ rfb::PixelFormat, }; -use super::RawEncoding; +use super::RawEncodingRef; #[allow(dead_code)] struct HextileEncoding { @@ -13,8 +13,8 @@ struct HextileEncoding { pixfmt: PixelFormat, } -impl From<&RawEncoding> for HextileEncoding { - fn from(raw: &RawEncoding) -> Self { +impl From<&RawEncodingRef<'_>> for HextileEncoding { + fn from(raw: &RawEncodingRef) -> Self { let pixfmt = raw.pixel_format().to_owned(); let (width, height) = raw.dimensions(); Self { diff --git a/src/encodings/raw.rs b/src/encodings/raw.rs index 2c0e30d..92ac16e 100644 --- a/src/encodings/raw.rs +++ b/src/encodings/raw.rs @@ -1,12 +1,12 @@ use crate::{ encodings::{Encoding, EncodingType}, pixel_formats::transform, - rfb::{ColorSpecification, PixelFormat}, + rfb::PixelFormat, }; /// Section 7.7.1 pub struct RawEncoding { - pixels: Vec, // TODO: self-lifetime-bound slice to avoid copy + pixels: Vec, width: u16, height: u16, pixfmt: PixelFormat, @@ -21,6 +21,73 @@ impl RawEncoding { pixfmt: pixfmt.clone(), } } +} + +impl<'a> From<&RawEncodingRef<'a>> for RawEncoding { + fn from(raw_ref: &RawEncodingRef<'a>) -> Self { + RawEncoding { + pixels: raw_ref.pixels.to_vec(), + width: raw_ref.width, + height: raw_ref.height, + pixfmt: raw_ref.pixfmt.to_owned(), + } + } +} + +impl<'a> From<&'a RawEncoding> for RawEncodingRef<'a> { + fn from(raw_owned: &'a RawEncoding) -> Self { + Self { + pixels: &raw_owned.pixels, + width: raw_owned.width, + height: raw_owned.height, + pixfmt: raw_owned.pixfmt.to_owned(), + } + } +} + +impl Encoding for RawEncoding { + fn get_type(&self) -> EncodingType { + EncodingType::Raw + } + + fn dimensions(&self) -> (u16, u16) { + (self.width, self.height) + } + + fn pixel_format(&self) -> &PixelFormat { + &self.pixfmt + } + + fn encode(&self) -> Box + '_> { + Box::new(self.pixels.iter().copied()) + } + + fn transform(&self, output: &PixelFormat) -> Box { + Box::new(RawEncoding { + pixels: transform(&self.pixels, &self.pixfmt, &output), + width: self.width, + height: self.height, + pixfmt: output.to_owned(), + }) + } +} + +pub struct RawEncodingRef<'a> { + pixels: &'a [u8], + width: u16, + height: u16, + pixfmt: PixelFormat, +} + +impl<'a> RawEncodingRef<'a> { + pub fn new(pixels: &'a [u8], width: u16, height: u16, pixfmt: &PixelFormat) -> Self { + Self { + pixels, + width, + height, + pixfmt: pixfmt.clone(), + } + } // useful for transforming into other encodings pub(crate) fn raw_buffer(&self) -> &[u8] { @@ -28,7 +95,7 @@ impl RawEncoding { } } -impl Encoding for RawEncoding { +impl<'a> Encoding for RawEncodingRef<'a> { fn get_type(&self) -> EncodingType { EncodingType::Raw } @@ -46,24 +113,11 @@ impl Encoding for RawEncoding { } fn transform(&self, output: &PixelFormat) -> Box { - let input = &self.pixfmt; - - // XXX: This assumes the pixel formats are both rgb. The server code verifies this - // before calling. - assert!(matches!( - &input.color_spec, - ColorSpecification::ColorFormat(_) - )); - assert!(matches!( - &output.color_spec, - ColorSpecification::ColorFormat(_) - )); - - Box::new(Self { - pixels: transform(&self.pixels, &input, &output), + Box::new(RawEncoding { + pixels: transform(&self.pixels, &self.pixfmt, &output), width: self.width, height: self.height, - pixfmt: output.clone(), + pixfmt: output.to_owned(), }) } } diff --git a/src/encodings/trle.rs b/src/encodings/trle.rs index d79484d..4781db6 100644 --- a/src/encodings/trle.rs +++ b/src/encodings/trle.rs @@ -4,7 +4,7 @@ use crate::encodings::{Encoding, EncodingType}; use crate::pixel_formats; use crate::rfb::PixelFormat; -use super::RawEncoding; +use super::RawEncodingRef; pub struct TRLEncoding { tiles: Vec>, @@ -13,8 +13,8 @@ pub struct TRLEncoding { pixfmt: PixelFormat, } -impl From<&RawEncoding> for TRLEncoding { - fn from(raw: &RawEncoding) -> Self { +impl<'a> From<&RawEncodingRef<'a>> for TRLEncoding { + fn from(raw: &RawEncodingRef) -> Self { raw.raw_buffer(); todo!() } From 591acc1ec3378c09c4ff28db165d2b0125e10249 Mon Sep 17 00:00:00 2001 From: lif <> Date: Fri, 23 Aug 2024 05:28:55 +0000 Subject: [PATCH 14/20] start encoding trle tiles (note: trle is not the end goal here, *zrle* is, which is based on trle) --- examples/server.rs | 8 +++--- src/encodings.rs | 1 + src/encodings/rre.rs | 2 ++ src/encodings/trle.rs | 63 +++++++++++++++++++++++++++++++++++++++++-- 4 files changed, 68 insertions(+), 6 deletions(-) diff --git a/examples/server.rs b/examples/server.rs index 9695a54..5423c6b 100644 --- a/examples/server.rs +++ b/examples/server.rs @@ -11,7 +11,7 @@ use env_logger; use image::io::Reader as ImageReader; use image::GenericImageView; use log::info; -use rfb::encodings::RawEncoding; +use rfb::encodings::{RawEncoding, RawEncodingRef, TRLEncoding}; use rfb::pixel_formats::fourcc::FourCC; use rfb::pixel_formats::transform; use rfb::rfb::{ @@ -162,12 +162,12 @@ impl Server for ExampleServer { 0, pixels_width, pixels_height, - Box::new(RawEncoding::new( - pixels, + Box::new(TRLEncoding::from(&RawEncodingRef::new( + &pixels, pixels_width, pixels_height, &self.pixfmt, - )), + ))), ); FramebufferUpdate::new(vec![r]) } diff --git a/src/encodings.rs b/src/encodings.rs index 67d8a15..bea51c6 100644 --- a/src/encodings.rs +++ b/src/encodings.rs @@ -17,6 +17,7 @@ mod trle; pub use raw::RawEncoding; pub use raw::RawEncodingRef; +pub use trle::TRLEncoding; #[derive(Debug)] #[allow(unused)] diff --git a/src/encodings/rre.rs b/src/encodings/rre.rs index ea497be..1bdf8fd 100644 --- a/src/encodings/rre.rs +++ b/src/encodings/rre.rs @@ -4,6 +4,8 @@ use crate::{ rfb::{PixelFormat, Position, Resolution}, }; +use super::RawEncodingRef; + struct RREncoding { background_pixel: Pixel, sub_rectangles: Vec, diff --git a/src/encodings/trle.rs b/src/encodings/trle.rs index 4781db6..8e163e5 100644 --- a/src/encodings/trle.rs +++ b/src/encodings/trle.rs @@ -13,10 +13,55 @@ pub struct TRLEncoding { pixfmt: PixelFormat, } +const TILE_PIXEL_SIZE: usize = 16; impl<'a> From<&RawEncodingRef<'a>> for TRLEncoding { fn from(raw: &RawEncodingRef) -> Self { - raw.raw_buffer(); - todo!() + let (w16, h16) = raw.dimensions(); + let (width, height) = (w16 as usize, h16 as usize); + + let pixfmt = raw.pixel_format(); + let bytes_per_px = (pixfmt.bits_per_pixel as usize + 7) / 8; + + let buf = raw.raw_buffer(); + + // if rect isn't a multiple of TILE_SIZE, we still encode the + // last partial tile. but if it *is* a multiple of TILE_SIZE, + // we don't -- hence inclusive range, but minus one before divide + let last_tile_row = (height - 1) / TILE_PIXEL_SIZE; + let last_tile_col = (width - 1) / TILE_PIXEL_SIZE; + let tiles = (0..=last_tile_row) + .into_iter() + .map(|tile_row_idx| { + let y_start = tile_row_idx * TILE_PIXEL_SIZE; + let y_end = height.min((tile_row_idx + 1) * TILE_PIXEL_SIZE); + (0..=last_tile_col) + .into_iter() + .map(move |tile_col_idx| { + let x_start = tile_col_idx * TILE_PIXEL_SIZE; + let x_end = width.min((tile_col_idx + 1) * TILE_PIXEL_SIZE); + let tile_pixels = (y_start..y_end).into_iter().flat_map(move |y| { + (x_start..x_end).into_iter().map(move |x| { + let px_start = (y * width + x) * bytes_per_px; + let px_end = (y * width + x + 1) * bytes_per_px; + &buf[px_start..px_end] + }) + }); + // TODO: other encodings + TRLETile::Raw { + pixels: tile_pixels + .map(|px_bytes| CPixel::from_raw(px_bytes, pixfmt)) + .collect(), + } + }) + .collect() + }) + .collect(); + Self { + tiles, + width: w16, + height: h16, + pixfmt: pixfmt.clone(), + } } } @@ -197,6 +242,20 @@ impl CPixel { } Self { bytes: out_bytes } } + + fn from_raw<'a>(raw_bytes: &[u8], pixfmt: &PixelFormat) -> Self { + let mut bytes = raw_bytes.to_vec(); + match Self::which_padding(pixfmt) { + CPixelTransformType::AsIs => (), + CPixelTransformType::AppendZero => { + bytes.pop(); + } + CPixelTransformType::PrependZero => { + bytes.remove(0); + } + } + Self { bytes } + } } impl Encoding for TRLEncoding { From f8e7d5d82a587eb5c186251fa29caaadb3d69364 Mon Sep 17 00:00:00 2001 From: lif <> Date: Fri, 23 Aug 2024 06:51:36 +0000 Subject: [PATCH 15/20] you know, nothing even supports TRLE any more to test against so we may as well just go straight to it --- Cargo.toml | 1 + src/encodings/trle.rs | 155 +++++++++++++++++++++++++----------------- 2 files changed, 95 insertions(+), 61 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 299f10d..7b2e128 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,7 @@ ascii = { version = "1.0", default-features = false } async-trait = "0.1.53" bitflags = "1.3.2" env_logger = "0.9.0" +flate2 = "1.0.32" futures = "0.3.21" log = "0.4.17" thiserror = "1.0" diff --git a/src/encodings/trle.rs b/src/encodings/trle.rs index 8e163e5..a8454fc 100644 --- a/src/encodings/trle.rs +++ b/src/encodings/trle.rs @@ -6,15 +6,98 @@ use crate::rfb::PixelFormat; use super::RawEncodingRef; -pub struct TRLEncoding { +pub struct RLEncoding { tiles: Vec>, width: u16, height: u16, pixfmt: PixelFormat, } -const TILE_PIXEL_SIZE: usize = 16; -impl<'a> From<&RawEncodingRef<'a>> for TRLEncoding { +pub type TRLEncoding = RLEncoding<16>; +pub struct ZRLEncoding(RLEncoding<64>); + +impl Encoding for ZRLEncoding { + fn get_type(&self) -> EncodingType { + EncodingType::ZRLE + } + + fn dimensions(&self) -> (u16, u16) { + self.0.dimensions() + } + + fn pixel_format(&self) -> &PixelFormat { + self.0.pixel_format() + } + + fn encode(&self /*, ctx: &mut ConnectionContext*/) -> Box + '_> { + todo!("flate2 with zlib stream shared with stream (but flushed to byte boundary at end of this fn)"); + todo!("also disable re-use of palettes in zrle mode") + } + + fn transform(&self, output: &PixelFormat) -> Box { + Box::new(Self(self.0.transform_inner(output))) + } +} + +impl RLEncoding { + const TILE_PIXEL_SIZE: usize = PX; + + fn transform_inner(&self, output: &PixelFormat) -> RLEncoding { + let input = &self.pixfmt; + let tiles = self + .tiles + .iter() + .map(|row| { + row.iter() + .map(|tile| match tile { + TRLETile::Raw { pixels } => TRLETile::Raw { + pixels: pixels + .iter() + .map(|cp| cp.transform(input, output)) + .collect(), + }, + TRLETile::SolidColor { color } => TRLETile::SolidColor { + color: color.transform(input, output), + }, + TRLETile::PackedPalette { + palette, + packed_pixels, + } => TRLETile::PackedPalette { + palette: palette + .iter() + .map(|cp| cp.transform(input, output)) + .collect(), + packed_pixels: packed_pixels.clone(), + }, + TRLETile::PlainRLE { runs } => TRLETile::PlainRLE { + runs: runs + .iter() + .map(|(cp, len)| (cp.transform(input, output), *len)) + .collect(), + }, + TRLETile::PaletteRLE { palette, runs } => TRLETile::PaletteRLE { + palette: palette + .iter() + .map(|cp| cp.transform(input, output)) + .collect(), + runs: runs.clone(), + }, + TRLETile::PackedPaletteReused { .. } + | TRLETile::PaletteRLEReused { .. } => tile.clone(), + }) + .collect() + }) + .collect(); + Self { + tiles, + width: self.width, + height: self.height, + pixfmt: output.to_owned(), + } + } +} + +impl<'a, const PX: usize> From<&RawEncodingRef<'a>> for RLEncoding { fn from(raw: &RawEncodingRef) -> Self { let (w16, h16) = raw.dimensions(); let (width, height) = (w16 as usize, h16 as usize); @@ -27,18 +110,18 @@ impl<'a> From<&RawEncodingRef<'a>> for TRLEncoding { // if rect isn't a multiple of TILE_SIZE, we still encode the // last partial tile. but if it *is* a multiple of TILE_SIZE, // we don't -- hence inclusive range, but minus one before divide - let last_tile_row = (height - 1) / TILE_PIXEL_SIZE; - let last_tile_col = (width - 1) / TILE_PIXEL_SIZE; + let last_tile_row = (height - 1) / Self::TILE_PIXEL_SIZE; + let last_tile_col = (width - 1) / Self::TILE_PIXEL_SIZE; let tiles = (0..=last_tile_row) .into_iter() .map(|tile_row_idx| { - let y_start = tile_row_idx * TILE_PIXEL_SIZE; - let y_end = height.min((tile_row_idx + 1) * TILE_PIXEL_SIZE); + let y_start = tile_row_idx * Self::TILE_PIXEL_SIZE; + let y_end = height.min((tile_row_idx + 1) * Self::TILE_PIXEL_SIZE); (0..=last_tile_col) .into_iter() .map(move |tile_col_idx| { - let x_start = tile_col_idx * TILE_PIXEL_SIZE; - let x_end = width.min((tile_col_idx + 1) * TILE_PIXEL_SIZE); + let x_start = tile_col_idx * Self::TILE_PIXEL_SIZE; + let x_end = width.min((tile_col_idx + 1) * Self::TILE_PIXEL_SIZE); let tile_pixels = (y_start..y_end).into_iter().flat_map(move |y| { (x_start..x_end).into_iter().map(move |x| { let px_start = (y * width + x) * bytes_per_px; @@ -258,7 +341,7 @@ impl CPixel { } } -impl Encoding for TRLEncoding { +impl Encoding for RLEncoding { fn get_type(&self) -> EncodingType { EncodingType::TRLE } @@ -272,57 +355,7 @@ impl Encoding for TRLEncoding { } fn transform(&self, output: &PixelFormat) -> Box { - let input = &self.pixfmt; - let tiles = self - .tiles - .iter() - .map(|row| { - row.iter() - .map(|tile| match tile { - TRLETile::Raw { pixels } => TRLETile::Raw { - pixels: pixels - .iter() - .map(|cp| cp.transform(input, output)) - .collect(), - }, - TRLETile::SolidColor { color } => TRLETile::SolidColor { - color: color.transform(input, output), - }, - TRLETile::PackedPalette { - palette, - packed_pixels, - } => TRLETile::PackedPalette { - palette: palette - .iter() - .map(|cp| cp.transform(input, output)) - .collect(), - packed_pixels: packed_pixels.clone(), - }, - TRLETile::PlainRLE { runs } => TRLETile::PlainRLE { - runs: runs - .iter() - .map(|(cp, len)| (cp.transform(input, output), *len)) - .collect(), - }, - TRLETile::PaletteRLE { palette, runs } => TRLETile::PaletteRLE { - palette: palette - .iter() - .map(|cp| cp.transform(input, output)) - .collect(), - runs: runs.clone(), - }, - TRLETile::PackedPaletteReused { .. } - | TRLETile::PaletteRLEReused { .. } => tile.clone(), - }) - .collect() - }) - .collect(); - Box::new(Self { - tiles, - width: self.width, - height: self.height, - pixfmt: output.to_owned(), - }) + Box::new(self.transform_inner(output)) } fn dimensions(&self) -> (u16, u16) { From 427504bae1829f46ce42f48378ae4c6d7191852d Mon Sep 17 00:00:00 2001 From: lif <> Date: Fri, 23 Aug 2024 08:19:35 +0000 Subject: [PATCH 16/20] start adding whole-screen zlib encoding too --- examples/server.rs | 4 +- src/encodings.rs | 3 ++ src/encodings/raw.rs | 8 ++- src/encodings/trle.rs | 117 ++++++++++++++++++++++++++---------------- src/encodings/zlib.rs | 0 5 files changed, 83 insertions(+), 49 deletions(-) create mode 100644 src/encodings/zlib.rs diff --git a/examples/server.rs b/examples/server.rs index 5423c6b..f41df00 100644 --- a/examples/server.rs +++ b/examples/server.rs @@ -11,7 +11,7 @@ use env_logger; use image::io::Reader as ImageReader; use image::GenericImageView; use log::info; -use rfb::encodings::{RawEncoding, RawEncodingRef, TRLEncoding}; +use rfb::encodings::{RawEncodingRef, ZRLEncoding}; use rfb::pixel_formats::fourcc::FourCC; use rfb::pixel_formats::transform; use rfb::rfb::{ @@ -162,7 +162,7 @@ impl Server for ExampleServer { 0, pixels_width, pixels_height, - Box::new(TRLEncoding::from(&RawEncodingRef::new( + Box::new(ZRLEncoding::from(&RawEncodingRef::new( &pixels, pixels_width, pixels_height, diff --git a/src/encodings.rs b/src/encodings.rs index bea51c6..c69dca9 100644 --- a/src/encodings.rs +++ b/src/encodings.rs @@ -14,10 +14,13 @@ mod raw; #[allow(unused)] mod rre; mod trle; +mod zlib; pub use raw::RawEncoding; pub use raw::RawEncodingRef; pub use trle::TRLEncoding; +pub use trle::ZRLEncoding; +pub use zlib::ZlibEncoding; #[derive(Debug)] #[allow(unused)] diff --git a/src/encodings/raw.rs b/src/encodings/raw.rs index 92ac16e..453c5aa 100644 --- a/src/encodings/raw.rs +++ b/src/encodings/raw.rs @@ -6,8 +6,8 @@ use crate::{ /// Section 7.7.1 pub struct RawEncoding { - pixels: Vec, - width: u16, + pub(crate) pixels: Vec, + pub(crate) width: u16, height: u16, pixfmt: PixelFormat, } @@ -21,6 +21,10 @@ impl RawEncoding { pixfmt: pixfmt.clone(), } } + + pub(crate) fn raw_buffer(&self) -> &[u8] { + &self.pixels + } } impl<'a> From<&RawEncodingRef<'a>> for RawEncoding { diff --git a/src/encodings/trle.rs b/src/encodings/trle.rs index a8454fc..d87d06d 100644 --- a/src/encodings/trle.rs +++ b/src/encodings/trle.rs @@ -13,8 +13,11 @@ pub struct RLEncoding { pixfmt: PixelFormat, } -pub type TRLEncoding = RLEncoding<16>; -pub struct ZRLEncoding(RLEncoding<64>); +const TRLE_TILE_PX_SIZE: usize = 16; +const ZRLE_TILE_PX_SIZE: usize = 64; + +pub type TRLEncoding = RLEncoding; +pub struct ZRLEncoding(RLEncoding); impl Encoding for ZRLEncoding { fn get_type(&self) -> EncodingType { @@ -39,6 +42,20 @@ impl Encoding for ZRLEncoding { } } +impl<'a> From<&RawEncodingRef<'a>> for ZRLEncoding { + fn from(raw: &RawEncodingRef) -> Self { + let (width, height) = raw.dimensions(); + let tiles = from_rawenc_inner(raw, ZRLE_TILE_PX_SIZE, true); + let pixfmt = raw.pixel_format().to_owned(); + Self(RLEncoding { + tiles, + width, + height, + pixfmt, + }) + } +} + impl RLEncoding { const TILE_PIXEL_SIZE: usize = PX; @@ -99,55 +116,65 @@ impl RLEncoding { impl<'a, const PX: usize> From<&RawEncodingRef<'a>> for RLEncoding { fn from(raw: &RawEncodingRef) -> Self { - let (w16, h16) = raw.dimensions(); - let (width, height) = (w16 as usize, h16 as usize); - - let pixfmt = raw.pixel_format(); - let bytes_per_px = (pixfmt.bits_per_pixel as usize + 7) / 8; - - let buf = raw.raw_buffer(); - - // if rect isn't a multiple of TILE_SIZE, we still encode the - // last partial tile. but if it *is* a multiple of TILE_SIZE, - // we don't -- hence inclusive range, but minus one before divide - let last_tile_row = (height - 1) / Self::TILE_PIXEL_SIZE; - let last_tile_col = (width - 1) / Self::TILE_PIXEL_SIZE; - let tiles = (0..=last_tile_row) - .into_iter() - .map(|tile_row_idx| { - let y_start = tile_row_idx * Self::TILE_PIXEL_SIZE; - let y_end = height.min((tile_row_idx + 1) * Self::TILE_PIXEL_SIZE); - (0..=last_tile_col) - .into_iter() - .map(move |tile_col_idx| { - let x_start = tile_col_idx * Self::TILE_PIXEL_SIZE; - let x_end = width.min((tile_col_idx + 1) * Self::TILE_PIXEL_SIZE); - let tile_pixels = (y_start..y_end).into_iter().flat_map(move |y| { - (x_start..x_end).into_iter().map(move |x| { - let px_start = (y * width + x) * bytes_per_px; - let px_end = (y * width + x + 1) * bytes_per_px; - &buf[px_start..px_end] - }) - }); - // TODO: other encodings - TRLETile::Raw { - pixels: tile_pixels - .map(|px_bytes| CPixel::from_raw(px_bytes, pixfmt)) - .collect(), - } - }) - .collect() - }) - .collect(); + let (width, height) = raw.dimensions(); + let tiles = from_rawenc_inner(raw, Self::TILE_PIXEL_SIZE, true); + let pixfmt = raw.pixel_format().to_owned(); Self { tiles, - width: w16, - height: h16, - pixfmt: pixfmt.clone(), + width, + height, + pixfmt, } } } +fn from_rawenc_inner( + raw: &RawEncodingRef<'_>, + tile_px_size: usize, + allow_pal_reuse: bool, +) -> Vec> { + let (w16, h16) = raw.dimensions(); + let (width, height) = (w16 as usize, h16 as usize); + + let pixfmt = raw.pixel_format(); + let bytes_per_px = (pixfmt.bits_per_pixel as usize + 7) / 8; + + let buf = raw.raw_buffer(); + + // if rect isn't a multiple of TILE_SIZE, we still encode the + // last partial tile. but if it *is* a multiple of TILE_SIZE, + // we don't -- hence inclusive range, but minus one before divide + let last_tile_row = (height - 1) / tile_px_size; + let last_tile_col = (width - 1) / tile_px_size; + (0..=last_tile_row) + .into_iter() + .map(|tile_row_idx| { + let y_start = tile_row_idx * tile_px_size; + let y_end = height.min((tile_row_idx + 1) * tile_px_size); + (0..=last_tile_col) + .into_iter() + .map(move |tile_col_idx| { + let x_start = tile_col_idx * tile_px_size; + let x_end = width.min((tile_col_idx + 1) * tile_px_size); + let tile_pixels = (y_start..y_end).into_iter().flat_map(move |y| { + (x_start..x_end).into_iter().map(move |x| { + let px_start = (y * width + x) * bytes_per_px; + let px_end = (y * width + x + 1) * bytes_per_px; + &buf[px_start..px_end] + }) + }); + // TODO: other encodings + TRLETile::Raw { + pixels: tile_pixels + .map(|px_bytes| CPixel::from_raw(px_bytes, pixfmt)) + .collect(), + } + }) + .collect() + }) + .collect() +} + #[repr(transparent)] #[derive(Copy, Clone)] struct PackedIndeces(u8); diff --git a/src/encodings/zlib.rs b/src/encodings/zlib.rs new file mode 100644 index 0000000..e69de29 From 6bba99c0a4f98994331340f15565d19a4a8a4d6e Mon Sep 17 00:00:00 2001 From: lif <> Date: Wed, 28 Aug 2024 06:23:35 +0000 Subject: [PATCH 17/20] wip: start plumbing zlib stream context & start decoupling from TcpStream so we can more easily websockify --- examples/server.rs | 4 +- src/encodings.rs | 9 +- src/encodings/hextile.rs | 4 +- src/encodings/raw.rs | 6 +- src/encodings/trle.rs | 6 +- src/encodings/zlib.rs | 81 ++++++++++++++++++ src/rfb.rs | 176 ++++++++++++++++++++++++++++++++------- src/server.rs | 12 +-- 8 files changed, 251 insertions(+), 47 deletions(-) diff --git a/examples/server.rs b/examples/server.rs index f41df00..623e2a9 100644 --- a/examples/server.rs +++ b/examples/server.rs @@ -15,7 +15,8 @@ use rfb::encodings::{RawEncodingRef, ZRLEncoding}; use rfb::pixel_formats::fourcc::FourCC; use rfb::pixel_formats::transform; use rfb::rfb::{ - FramebufferUpdate, KeyEvent, PixelFormat, ProtoVersion, Rectangle, SecurityType, SecurityTypes, + ConnectionContext, FramebufferUpdate, KeyEvent, PixelFormat, ProtoVersion, Rectangle, + SecurityType, SecurityTypes, }; use rfb::server::{Server, VncServer, VncServerConfig, VncServerData}; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; @@ -86,6 +87,7 @@ async fn main() -> Result<()> { width: WIDTH as u16, height: HEIGHT as u16, input_pixel_format: pixfmt.clone(), + connection_context: ConnectionContext::default(), }; let server = ExampleServer { display: args.image, diff --git a/src/encodings.rs b/src/encodings.rs index c69dca9..bdaa0b0 100644 --- a/src/encodings.rs +++ b/src/encodings.rs @@ -4,15 +4,14 @@ // // Copyright 2022 Oxide Computer Company +use crate::rfb::ConnectionContext; use crate::rfb::PixelFormat; use EncodingType::*; -#[allow(unused)] -mod hextile; +// mod hextile; mod raw; -#[allow(unused)] -mod rre; +// mod rre; mod trle; mod zlib; @@ -58,7 +57,7 @@ where fn pixel_format(&self) -> &PixelFormat; /// Transform this encoding from its representation into a byte sequence that can be passed to the client. - fn encode(&self) -> Box + '_>; + fn encode(&self, ctx: &mut ConnectionContext) -> Box + '_>; /// Translates this encoding type from its current pixel format to the given format. fn transform(&self, output: &PixelFormat) -> Box; diff --git a/src/encodings/hextile.rs b/src/encodings/hextile.rs index 3633be0..65821aa 100644 --- a/src/encodings/hextile.rs +++ b/src/encodings/hextile.rs @@ -1,6 +1,6 @@ use crate::{ encodings::{Encoding, EncodingType, Pixel}, - rfb::PixelFormat, + rfb::{ConnectionContext, PixelFormat}, }; use super::RawEncodingRef; @@ -41,7 +41,7 @@ impl Encoding for HextileEncoding { EncodingType::Hextile } - fn encode(&self) -> Box + '_> { + fn encode(&self, _ctx: &mut ConnectionContext) -> Box + '_> { Box::new(self.tiles.iter().flat_map(|tile| { let subencoding_mask = todo!(); [todo!()].into_iter() diff --git a/src/encodings/raw.rs b/src/encodings/raw.rs index 453c5aa..8e5067a 100644 --- a/src/encodings/raw.rs +++ b/src/encodings/raw.rs @@ -1,7 +1,7 @@ use crate::{ encodings::{Encoding, EncodingType}, pixel_formats::transform, - rfb::PixelFormat, + rfb::{ConnectionContext, PixelFormat}, }; /// Section 7.7.1 @@ -62,7 +62,7 @@ impl Encoding for RawEncoding { &self.pixfmt } - fn encode(&self) -> Box + '_> { + fn encode(&self, _ctx: &mut ConnectionContext) -> Box + '_> { Box::new(self.pixels.iter().copied()) } @@ -112,7 +112,7 @@ impl<'a> Encoding for RawEncodingRef<'a> { &self.pixfmt } - fn encode(&self) -> Box + '_> { + fn encode(&self, _ctx: &mut ConnectionContext) -> Box + '_> { Box::new(self.pixels.iter().copied()) } diff --git a/src/encodings/trle.rs b/src/encodings/trle.rs index d87d06d..3e6ac70 100644 --- a/src/encodings/trle.rs +++ b/src/encodings/trle.rs @@ -2,7 +2,7 @@ use std::iter::{from_fn, once}; use crate::encodings::{Encoding, EncodingType}; use crate::pixel_formats; -use crate::rfb::PixelFormat; +use crate::rfb::{ConnectionContext, PixelFormat}; use super::RawEncodingRef; @@ -32,7 +32,7 @@ impl Encoding for ZRLEncoding { self.0.pixel_format() } - fn encode(&self /*, ctx: &mut ConnectionContext*/) -> Box + '_> { + fn encode(&self, ctx: &mut ConnectionContext) -> Box + '_> { todo!("flate2 with zlib stream shared with stream (but flushed to byte boundary at end of this fn)"); todo!("also disable re-use of palettes in zrle mode") } @@ -373,7 +373,7 @@ impl Encoding for RLEncoding { EncodingType::TRLE } - fn encode(&self) -> Box + '_> { + fn encode(&self, _ctx: &mut ConnectionContext) -> Box + '_> { Box::new( self.tiles .iter() diff --git a/src/encodings/zlib.rs b/src/encodings/zlib.rs index e69de29..97857f1 100644 --- a/src/encodings/zlib.rs +++ b/src/encodings/zlib.rs @@ -0,0 +1,81 @@ +use crate::{pixel_formats::transform, rfb::ConnectionContext}; + +use super::{Encoding, EncodingType, RawEncoding, RawEncodingRef}; + +pub struct ZlibEncodingRef<'a> { + raw: RawEncodingRef<'a>, +} + +impl<'a> Encoding for ZlibEncodingRef<'a> { + fn get_type(&self) -> EncodingType { + EncodingType::Zlib + } + + fn dimensions(&self) -> (u16, u16) { + self.raw.dimensions() + } + + fn pixel_format(&self) -> &crate::rfb::PixelFormat { + self.raw.pixel_format() + } + + fn encode(&self, ctx: &mut ConnectionContext) -> Box + '_> { + let in_buf = self.raw.raw_buffer(); + let mut out_buf = Vec::with_capacity(in_buf.len()); + ctx.zlib + .compress_vec(in_buf, &mut out_buf, flate2::FlushCompress::Sync) + .expect("zlib error"); + Box::new(out_buf.into_iter()) + } + + fn transform(&self, output: &crate::rfb::PixelFormat) -> Box { + let (width, height) = self.raw.dimensions(); + Box::new(ZlibEncoding { + raw: RawEncoding::new( + transform(self.raw.raw_buffer(), self.pixel_format(), output), + width, + height, + output, + ), + }) + } +} + +pub struct ZlibEncoding { + raw: RawEncoding, +} + +impl Encoding for ZlibEncoding { + fn get_type(&self) -> EncodingType { + EncodingType::Zlib + } + + fn dimensions(&self) -> (u16, u16) { + self.raw.dimensions() + } + + fn pixel_format(&self) -> &crate::rfb::PixelFormat { + self.raw.pixel_format() + } + + fn encode(&self, ctx: &mut ConnectionContext) -> Box + '_> { + let in_buf = self.raw.raw_buffer(); + let mut out_buf = Vec::with_capacity(in_buf.len()); + ctx.zlib + .compress_vec(in_buf, &mut out_buf, flate2::FlushCompress::Sync) + .expect("zlib error"); + Box::new(out_buf.into_iter()) + } + + fn transform(&self, output: &crate::rfb::PixelFormat) -> Box { + let (width, height) = self.raw.dimensions(); + Box::new(Self { + raw: RawEncoding::new( + transform(self.raw.raw_buffer(), self.pixel_format(), output), + width, + height, + output, + ), + }) + } +} diff --git a/src/rfb.rs b/src/rfb.rs index 7adf524..0534036 100644 --- a/src/rfb.rs +++ b/src/rfb.rs @@ -4,6 +4,8 @@ // // Copyright 2022 Oxide Computer Company +use std::iter::once; + use bitflags::bitflags; use futures::future::BoxFuture; use futures::FutureExt; @@ -14,6 +16,17 @@ use tokio::net::TcpStream; use crate::encodings::{Encoding, EncodingType}; use crate::keysym::KeySym; +pub struct ConnectionContext { + pub zlib: flate2::Compress, +} +impl Default for ConnectionContext { + fn default() -> Self { + Self { + zlib: flate2::Compress::new(flate2::Compression::fast(), false), + } + } +} + #[derive(Debug, Error)] pub enum ProtocolError { #[error("invalid protocol version message")] @@ -42,6 +55,7 @@ pub trait ReadMessage { } pub trait WriteMessage { + fn encode(self, ctx: &mut ConnectionContext) -> Box + '_>; fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>>; } @@ -57,31 +71,48 @@ impl ReadMessage for ProtoVersion { async move { let mut buf = [0u8; 12]; stream.read_exact(&mut buf).await?; - - match &buf { - b"RFB 003.003\n" => Ok(ProtoVersion::Rfb33), - b"RFB 003.007\n" => Ok(ProtoVersion::Rfb37), - b"RFB 003.008\n" => Ok(ProtoVersion::Rfb38), - _ => Err(ProtocolError::InvalidProtocolVersion), - } + Self::try_from(buf) } .boxed() } } +impl TryFrom<[u8; 12]> for ProtoVersion { + type Error = ProtocolError; + + fn try_from(buf: [u8; 12]) -> Result { + match &buf { + b"RFB 003.003\n" => Ok(ProtoVersion::Rfb33), + b"RFB 003.007\n" => Ok(ProtoVersion::Rfb37), + b"RFB 003.008\n" => Ok(ProtoVersion::Rfb38), + _ => Err(ProtocolError::InvalidProtocolVersion), + } + } +} + +impl Into<&'static [u8; 12]> for ProtoVersion { + fn into(self) -> &'static [u8; 12] { + match self { + ProtoVersion::Rfb33 => b"RFB 003.003\n", + ProtoVersion::Rfb37 => b"RFB 003.007\n", + ProtoVersion::Rfb38 => b"RFB 003.008\n", + } + } +} + impl WriteMessage for ProtoVersion { fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { async move { - let s = match self { - ProtoVersion::Rfb33 => b"RFB 003.003\n", - ProtoVersion::Rfb37 => b"RFB 003.007\n", - ProtoVersion::Rfb38 => b"RFB 003.008\n", - }; - + let s: &[u8; 12] = self.into(); Ok(stream.write_all(s).await?) } .boxed() } + + fn encode(self, _ctx: &mut ConnectionContext) -> Box + '_> { + let s: &[u8; 12] = self.into(); + Box::new(s.iter().copied()) + } } // Section 7.1.2 @@ -89,9 +120,10 @@ impl WriteMessage for ProtoVersion { pub struct SecurityTypes(pub Vec); #[derive(Clone, PartialEq, Debug)] +#[repr(u8)] pub enum SecurityType { - None, - VncAuthentication, + None = 0, + VncAuthentication = 1, } impl WriteMessage for SecurityTypes { @@ -107,6 +139,13 @@ impl WriteMessage for SecurityTypes { } .boxed() } + + fn encode(self, ctx: &mut ConnectionContext) -> Box + '_> { + Box::new( + once(self.0.len() as u8) // TODO: fix cast + .chain(self.0.into_iter().flat_map(|t| t.encode(ctx))), + ) + } } impl ReadMessage for SecurityType { @@ -126,16 +165,15 @@ impl ReadMessage for SecurityType { impl WriteMessage for SecurityType { fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { async move { - let val = match self { - SecurityType::None => 0, - SecurityType::VncAuthentication => 1, - }; - stream.write_u8(val).await?; - + stream.write_u8(self as u8).await?; Ok(()) } .boxed() } + + fn encode(self, _ctx: &mut ConnectionContext) -> Box + '_> { + Box::new(once(self as u8)) + } } // Section 7.1.3 @@ -161,6 +199,15 @@ impl WriteMessage for SecurityResult { } .boxed() } + + fn encode(self, _ctx: &mut ConnectionContext) -> Box + '_> { + match self { + SecurityResult::Success => Box::new(0u32.to_be_bytes().iter().copied()), + SecurityResult::Failure(s) => { + Box::new(1u32.to_be_bytes().iter().copied().chain(s.bytes())) + } + } + } } // Section 7.3.1 @@ -214,6 +261,16 @@ impl WriteMessage for ServerInit { } .boxed() } + + fn encode(self, ctx: &mut ConnectionContext) -> Box + '_> { + Box::new( + self.initial_res + .encode(ctx) + .chain(self.pixel_format.encode(ctx)) + .chain((self.name.len() as u32).to_be_bytes().iter().copied()) + .chain(self.name.bytes()), + ) + } } pub enum _ServerMessage { @@ -233,12 +290,11 @@ impl FramebufferUpdate { } pub fn transform(&self, input_pf: &PixelFormat, output_pf: &PixelFormat) -> Self { - let mut rectangles = Vec::new(); - - for r in self.rectangles.iter() { - rectangles.push(r.transform(input_pf, output_pf)); - } - + let rectangles = self + .rectangles + .iter() + .map(|r| r.transform(input_pf, output_pf)) + .collect(); FramebufferUpdate { rectangles } } } @@ -288,6 +344,16 @@ impl WriteMessage for Resolution { } .boxed() } + + fn encode(self, ctx: &mut ConnectionContext) -> Box + '_> { + Box::new( + self.width + .to_be_bytes() + .iter() + .copied() + .chain(self.height.to_be_bytes().iter().copied()), + ) + } } pub struct Rectangle { @@ -328,13 +394,27 @@ impl WriteMessage for Rectangle { stream.write_i32(encoding_type).await?; // TODO: avoid collect alloc? - let data: Vec<_> = self.data.encode().collect(); + let data: Vec<_> = self.data.encode(todo!()).collect(); stream.write_all(&data).await?; Ok(()) } .boxed() } + + fn encode(self, ctx: &mut ConnectionContext) -> Box + '_> { + Box::new( + self.position + .x + .to_be_bytes() + .into_iter() + .chain(self.position.y.to_be_bytes().into_iter()) + .chain(self.dimensions.width.to_be_bytes().into_iter()) + .chain(self.dimensions.height.to_be_bytes().into_iter()) + .chain(i32::from(self.data.get_type()).to_be_bytes().into_iter()) + .chain(self.data.encode(ctx)), + ) + } } impl WriteMessage for FramebufferUpdate { @@ -359,6 +439,17 @@ impl WriteMessage for FramebufferUpdate { } .boxed() } + + fn encode(self, ctx: &mut ConnectionContext) -> Box + '_> { + // number of rectangles + let n_rect = self.rectangles.len() as u16; + Box::new( + once(0u8) // TODO: type function? + .chain(once(0u8)) // 1 byte of padding + .chain(n_rect.to_be_bytes().into_iter()) + .chain(self.rectangles.into_iter().flat_map(|r| r.encode(ctx))), + ) + } } #[derive(Debug)] @@ -521,6 +612,15 @@ impl WriteMessage for PixelFormat { } .boxed() } + + fn encode(self, ctx: &mut ConnectionContext) -> Box + '_> { + Box::new( + [self.bits_per_pixel, self.depth, self.big_endian as u8] + .into_iter() + .chain(self.color_spec.encode(ctx)) + .chain([0u8; 3].into_iter()), // 3 bytes of padding + ) + } } #[derive(Debug, Clone, PartialEq)] @@ -602,6 +702,26 @@ impl WriteMessage for ColorSpecification { } .boxed() } + + fn encode(self, ctx: &mut ConnectionContext) -> Box + '_> { + match self { + ColorSpecification::ColorFormat(cf) => { + Box::new( + once(1u8) // true color + .chain(cf.red_max.to_be_bytes().into_iter()) + .chain(cf.green_max.to_be_bytes().into_iter()) + .chain(cf.blue_max.to_be_bytes().into_iter()) + .chain([cf.red_shift, cf.green_shift, cf.blue_shift].into_iter()), + ) + } + ColorSpecification::ColorMap(_cm) => { + // first 0 byte is true-color-flag = false; + // the remaining 9 are the above max/shift fields, + // which aren't relevant in ColorMap mode + Box::new(std::iter::repeat(0).take(10)) + } + } + } } // Section 7.5 diff --git a/src/server.rs b/src/server.rs index f77e234..d0c4599 100644 --- a/src/server.rs +++ b/src/server.rs @@ -20,9 +20,9 @@ use tokio::select; use tokio::sync::{oneshot, Mutex}; use crate::rfb::{ - ClientInit, ClientMessage, FramebufferUpdate, KeyEvent, PixelFormat, PointerEvent, - ProtoVersion, ProtocolError, ReadMessage, SecurityResult, SecurityType, SecurityTypes, - ServerInit, WriteMessage, + ClientInit, ClientMessage, ConnectionContext, FramebufferUpdate, KeyEvent, PixelFormat, + PointerEvent, ProtoVersion, ProtocolError, ReadMessage, SecurityResult, SecurityType, + SecurityTypes, ServerInit, WriteMessage, }; #[derive(Debug, Error)] @@ -61,6 +61,8 @@ pub struct VncServerData { /// The pixel format of the framebuffer data passed in to the server via /// get_framebuffer_update. pub input_pixel_format: PixelFormat, + /// State used during encoding, such as the Zlib stream (which is shared between rectangles). + pub connection_context: ConnectionContext, } pub struct VncServer { @@ -92,9 +94,9 @@ impl VncServer { "at least one security type must be defined" ); Arc::new(Self { - config: config, + config, data: Mutex::new(data), - server: server, + server, stop_ch: Mutex::new(None), }) } From 5f49c745a4541c5455871f57673101e23f3aa9d7 Mon Sep 17 00:00:00 2001 From: lif <> Date: Thu, 29 Aug 2024 04:18:52 +0000 Subject: [PATCH 18/20] chat, what do we think about using streams here --- Cargo.toml | 1 + src/encodings.rs | 12 ++-- src/encodings/hextile.rs | 2 +- src/encodings/raw.rs | 18 +++-- src/encodings/rre.rs | 6 +- src/encodings/trle.rs | 20 ++++-- src/encodings/zlib.rs | 32 ++++++--- src/rfb.rs | 144 +++++++++++++++++++++++---------------- 8 files changed, 146 insertions(+), 89 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 7b2e128..2f748c6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,6 +28,7 @@ futures = "0.3.21" log = "0.4.17" thiserror = "1.0" tokio = { version = "1", features = ["full"] } +cancel-safe-futures = "0.1" [dev-dependencies] anyhow = "1.0" diff --git a/src/encodings.rs b/src/encodings.rs index bdaa0b0..3fc0e28 100644 --- a/src/encodings.rs +++ b/src/encodings.rs @@ -4,9 +4,13 @@ // // Copyright 2022 Oxide Computer Company +use std::sync::Arc; + use crate::rfb::ConnectionContext; use crate::rfb::PixelFormat; +use async_trait::async_trait; +use futures::stream::BoxStream; use EncodingType::*; // mod hextile; @@ -44,10 +48,8 @@ struct Pixel { bytes: Vec, } -pub trait Encoding -where - Self: Send, -{ +#[async_trait] +pub trait Encoding: Send + Sync { fn get_type(&self) -> EncodingType; /// Return the width and height in pixels of the encoded screen region. @@ -57,7 +59,7 @@ where fn pixel_format(&self) -> &PixelFormat; /// Transform this encoding from its representation into a byte sequence that can be passed to the client. - fn encode(&self, ctx: &mut ConnectionContext) -> Box + '_>; + async fn encode(&self, ctx: Arc) -> BoxStream; /// Translates this encoding type from its current pixel format to the given format. fn transform(&self, output: &PixelFormat) -> Box; diff --git a/src/encodings/hextile.rs b/src/encodings/hextile.rs index 65821aa..4529317 100644 --- a/src/encodings/hextile.rs +++ b/src/encodings/hextile.rs @@ -41,7 +41,7 @@ impl Encoding for HextileEncoding { EncodingType::Hextile } - fn encode(&self, _ctx: &mut ConnectionContext) -> Box + '_> { + fn encode(&self, _ctx: &mut ConnectionContext) -> Box + '_> { Box::new(self.tiles.iter().flat_map(|tile| { let subencoding_mask = todo!(); [todo!()].into_iter() diff --git a/src/encodings/raw.rs b/src/encodings/raw.rs index 8e5067a..f2e8505 100644 --- a/src/encodings/raw.rs +++ b/src/encodings/raw.rs @@ -1,3 +1,11 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use futures::{ + stream::{self, BoxStream}, + StreamExt, +}; + use crate::{ encodings::{Encoding, EncodingType}, pixel_formats::transform, @@ -49,6 +57,7 @@ impl<'a> From<&'a RawEncoding> for RawEncodingRef<'a> { } } +#[async_trait] impl Encoding for RawEncoding { fn get_type(&self) -> EncodingType { EncodingType::Raw @@ -62,8 +71,8 @@ impl Encoding for RawEncoding { &self.pixfmt } - fn encode(&self, _ctx: &mut ConnectionContext) -> Box + '_> { - Box::new(self.pixels.iter().copied()) + async fn encode(&self, _ctx: Arc) -> BoxStream { + stream::iter(self.pixels.iter().copied()).boxed() } fn transform(&self, output: &PixelFormat) -> Box { @@ -99,6 +108,7 @@ impl<'a> RawEncodingRef<'a> { } } +#[async_trait] impl<'a> Encoding for RawEncodingRef<'a> { fn get_type(&self) -> EncodingType { EncodingType::Raw @@ -112,8 +122,8 @@ impl<'a> Encoding for RawEncodingRef<'a> { &self.pixfmt } - fn encode(&self, _ctx: &mut ConnectionContext) -> Box + '_> { - Box::new(self.pixels.iter().copied()) + async fn encode(&self, _ctx: Arc) -> BoxStream { + stream::iter(self.pixels.iter().copied()).boxed() } fn transform(&self, output: &PixelFormat) -> Box { diff --git a/src/encodings/rre.rs b/src/encodings/rre.rs index 1bdf8fd..b601b9c 100644 --- a/src/encodings/rre.rs +++ b/src/encodings/rre.rs @@ -25,8 +25,8 @@ impl Encoding for RREncoding { EncodingType::RRE } - fn encode(&self) -> Box + '_> { - Box::new( + fn encode(&self) -> Box + '_> { + Box::new(stream::iter( (self.sub_rectangles.len() as u32) .to_be_bytes() .into_iter() @@ -41,7 +41,7 @@ impl Encoding for RREncoding { .chain(sr.dimensions.width.to_be_bytes().into_iter()) .chain(sr.dimensions.height.to_be_bytes().into_iter()) })), - ) + )) } fn transform(&self, output: &PixelFormat) -> Box { diff --git a/src/encodings/trle.rs b/src/encodings/trle.rs index 3e6ac70..b438ab5 100644 --- a/src/encodings/trle.rs +++ b/src/encodings/trle.rs @@ -1,4 +1,9 @@ use std::iter::{from_fn, once}; +use std::sync::Arc; + +use async_trait::async_trait; +use futures::stream::{self, BoxStream}; +use futures::StreamExt; use crate::encodings::{Encoding, EncodingType}; use crate::pixel_formats; @@ -19,6 +24,7 @@ const ZRLE_TILE_PX_SIZE: usize = 64; pub type TRLEncoding = RLEncoding; pub struct ZRLEncoding(RLEncoding); +#[async_trait] impl Encoding for ZRLEncoding { fn get_type(&self) -> EncodingType { EncodingType::ZRLE @@ -32,7 +38,7 @@ impl Encoding for ZRLEncoding { self.0.pixel_format() } - fn encode(&self, ctx: &mut ConnectionContext) -> Box + '_> { + async fn encode(&self, ctx: Arc) -> BoxStream { todo!("flate2 with zlib stream shared with stream (but flushed to byte boundary at end of this fn)"); todo!("also disable re-use of palettes in zrle mode") } @@ -234,7 +240,7 @@ enum TRLETile { }, } -fn rle(mut length: usize) -> impl Iterator { +fn rle(mut length: usize) -> impl Iterator + Send { from_fn(move || { if length == 0 { None @@ -249,7 +255,7 @@ fn rle(mut length: usize) -> impl Iterator { }) } -fn pal_rle((index, length): &(u8, usize)) -> Box> { +fn pal_rle((index, length): &(u8, usize)) -> Box + Send> { if *length == 1 { Box::new(once(*index)) } else { @@ -261,7 +267,7 @@ impl TRLETile { /// Subencoding of the tile according to RFB 6143 7.7.5. /// To the extent possible, this function is a translation of that /// section of the RFB RFC from English into chained iterators. - fn encode(&self) -> Box + '_> { + fn encode(&self) -> Box + Send + '_> { match self { TRLETile::Raw { pixels } => { Box::new(once(0u8).chain(pixels.iter().flat_map(|c| c.bytes.iter().copied()))) @@ -368,17 +374,19 @@ impl CPixel { } } +#[async_trait] impl Encoding for RLEncoding { fn get_type(&self) -> EncodingType { EncodingType::TRLE } - fn encode(&self, _ctx: &mut ConnectionContext) -> Box + '_> { - Box::new( + async fn encode(&self, _ctx: Arc) -> BoxStream { + stream::iter( self.tiles .iter() .flat_map(|row| row.iter().flat_map(|tile| tile.encode())), ) + .boxed() } fn transform(&self, output: &PixelFormat) -> Box { diff --git a/src/encodings/zlib.rs b/src/encodings/zlib.rs index 97857f1..e1c1df5 100644 --- a/src/encodings/zlib.rs +++ b/src/encodings/zlib.rs @@ -1,3 +1,11 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use futures::{ + stream::{self, BoxStream}, + StreamExt, +}; + use crate::{pixel_formats::transform, rfb::ConnectionContext}; use super::{Encoding, EncodingType, RawEncoding, RawEncodingRef}; @@ -6,6 +14,7 @@ pub struct ZlibEncodingRef<'a> { raw: RawEncodingRef<'a>, } +#[async_trait] impl<'a> Encoding for ZlibEncodingRef<'a> { fn get_type(&self) -> EncodingType { EncodingType::Zlib @@ -19,13 +28,14 @@ impl<'a> Encoding for ZlibEncodingRef<'a> { self.raw.pixel_format() } - fn encode(&self, ctx: &mut ConnectionContext) -> Box + '_> { + async fn encode(&self, ctx: Arc) -> BoxStream { let in_buf = self.raw.raw_buffer(); let mut out_buf = Vec::with_capacity(in_buf.len()); - ctx.zlib - .compress_vec(in_buf, &mut out_buf, flate2::FlushCompress::Sync) - .expect("zlib error"); - Box::new(out_buf.into_iter()) + ctx.zlib.lock().await.unwrap().perform(|zlib| { + zlib.compress_vec(in_buf, &mut out_buf, flate2::FlushCompress::Sync) + .expect("zlib error") + }); + stream::iter(out_buf.into_iter()).boxed() } fn transform(&self, output: &crate::rfb::PixelFormat) -> Box { @@ -45,6 +55,7 @@ pub struct ZlibEncoding { raw: RawEncoding, } +#[async_trait] impl Encoding for ZlibEncoding { fn get_type(&self) -> EncodingType { EncodingType::Zlib @@ -58,13 +69,14 @@ impl Encoding for ZlibEncoding { self.raw.pixel_format() } - fn encode(&self, ctx: &mut ConnectionContext) -> Box + '_> { + async fn encode(&self, ctx: Arc) -> BoxStream { let in_buf = self.raw.raw_buffer(); let mut out_buf = Vec::with_capacity(in_buf.len()); - ctx.zlib - .compress_vec(in_buf, &mut out_buf, flate2::FlushCompress::Sync) - .expect("zlib error"); - Box::new(out_buf.into_iter()) + ctx.zlib.lock().await.unwrap().perform(|zlib| { + zlib.compress_vec(in_buf, &mut out_buf, flate2::FlushCompress::Sync) + .expect("zlib error") + }); + stream::iter(out_buf.into_iter()).boxed() } fn transform(&self, output: &crate::rfb::PixelFormat) -> Box { diff --git a/src/rfb.rs b/src/rfb.rs index 0534036..38d8b29 100644 --- a/src/rfb.rs +++ b/src/rfb.rs @@ -4,11 +4,16 @@ // // Copyright 2022 Oxide Computer Company -use std::iter::once; +use std::sync::Arc; +use async_trait::async_trait; use bitflags::bitflags; -use futures::future::BoxFuture; -use futures::FutureExt; +use cancel_safe_futures::sync::RobustMutex; +use futures::{ + future::BoxFuture, + stream::{self, BoxStream}, + FutureExt, StreamExt, +}; use thiserror::Error; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpStream; @@ -17,12 +22,12 @@ use crate::encodings::{Encoding, EncodingType}; use crate::keysym::KeySym; pub struct ConnectionContext { - pub zlib: flate2::Compress, + pub zlib: RobustMutex, } impl Default for ConnectionContext { fn default() -> Self { Self { - zlib: flate2::Compress::new(flate2::Compression::fast(), false), + zlib: RobustMutex::new(flate2::Compress::new(flate2::Compression::fast(), false)), } } } @@ -54,8 +59,9 @@ pub trait ReadMessage { Self: Sized; } +#[async_trait] pub trait WriteMessage { - fn encode(self, ctx: &mut ConnectionContext) -> Box + '_>; + async fn encode(&self, ctx: Arc) -> BoxStream; fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>>; } @@ -90,7 +96,7 @@ impl TryFrom<[u8; 12]> for ProtoVersion { } } -impl Into<&'static [u8; 12]> for ProtoVersion { +impl Into<&'static [u8; 12]> for &ProtoVersion { fn into(self) -> &'static [u8; 12] { match self { ProtoVersion::Rfb33 => b"RFB 003.003\n", @@ -100,18 +106,18 @@ impl Into<&'static [u8; 12]> for ProtoVersion { } } +#[async_trait] impl WriteMessage for ProtoVersion { fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { async move { - let s: &[u8; 12] = self.into(); + let s: &[u8; 12] = (&self).into(); Ok(stream.write_all(s).await?) } .boxed() } - - fn encode(self, _ctx: &mut ConnectionContext) -> Box + '_> { + async fn encode(&self, _ctx: Arc) -> BoxStream { let s: &[u8; 12] = self.into(); - Box::new(s.iter().copied()) + stream::iter(s.iter().copied()).boxed() } } @@ -119,13 +125,14 @@ impl WriteMessage for ProtoVersion { #[derive(Debug, Clone)] pub struct SecurityTypes(pub Vec); -#[derive(Clone, PartialEq, Debug)] +#[derive(Copy, Clone, PartialEq, Debug)] #[repr(u8)] pub enum SecurityType { None = 0, VncAuthentication = 1, } +#[async_trait] impl WriteMessage for SecurityTypes { fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { async move { @@ -139,12 +146,14 @@ impl WriteMessage for SecurityTypes { } .boxed() } - - fn encode(self, ctx: &mut ConnectionContext) -> Box + '_> { - Box::new( - once(self.0.len() as u8) // TODO: fix cast - .chain(self.0.into_iter().flat_map(|t| t.encode(ctx))), - ) + async fn encode(&self, ctx: Arc) -> BoxStream { + stream::iter([self.0.len() as u8].into_iter()) // TODO: fix cast + .chain( + stream::iter(self.0.iter()) + .then(move |t| t.encode(ctx.to_owned())) + .flatten(), + ) + .boxed() } } @@ -162,6 +171,7 @@ impl ReadMessage for SecurityType { } } +#[async_trait] impl WriteMessage for SecurityType { fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { async move { @@ -171,8 +181,8 @@ impl WriteMessage for SecurityType { .boxed() } - fn encode(self, _ctx: &mut ConnectionContext) -> Box + '_> { - Box::new(once(self as u8)) + async fn encode(&self, _ctx: Arc) -> BoxStream { + stream::iter([*self as u8].into_iter()).boxed() } } @@ -182,6 +192,7 @@ pub enum SecurityResult { Failure(String), } +#[async_trait] impl WriteMessage for SecurityResult { fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { async move { @@ -199,12 +210,11 @@ impl WriteMessage for SecurityResult { } .boxed() } - - fn encode(self, _ctx: &mut ConnectionContext) -> Box + '_> { + async fn encode(&self, _ctx: Arc) -> BoxStream { match self { - SecurityResult::Success => Box::new(0u32.to_be_bytes().iter().copied()), + SecurityResult::Success => stream::iter(0u32.to_be_bytes().into_iter()).boxed(), SecurityResult::Failure(s) => { - Box::new(1u32.to_be_bytes().iter().copied().chain(s.bytes())) + stream::iter(1u32.to_be_bytes().into_iter().chain(s.bytes())).boxed() } } } @@ -247,6 +257,7 @@ impl ServerInit { } } +#[async_trait] impl WriteMessage for ServerInit { fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { async move { @@ -262,14 +273,18 @@ impl WriteMessage for ServerInit { .boxed() } - fn encode(self, ctx: &mut ConnectionContext) -> Box + '_> { - Box::new( - self.initial_res - .encode(ctx) - .chain(self.pixel_format.encode(ctx)) - .chain((self.name.len() as u32).to_be_bytes().iter().copied()) - .chain(self.name.bytes()), - ) + async fn encode(&self, ctx: Arc) -> BoxStream { + self.initial_res + .encode(ctx.clone()) + .await + .chain(self.pixel_format.encode(ctx).await) + .chain(stream::iter( + (self.name.len() as u32) + .to_be_bytes() + .into_iter() + .chain(self.name.bytes()), + )) + .boxed() } } @@ -335,6 +350,7 @@ impl ReadMessage for Resolution { } } +#[async_trait] impl WriteMessage for Resolution { fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { async move { @@ -345,14 +361,14 @@ impl WriteMessage for Resolution { .boxed() } - fn encode(self, ctx: &mut ConnectionContext) -> Box + '_> { - Box::new( + async fn encode(&self, _ctx: Arc) -> BoxStream { + stream::iter( self.width .to_be_bytes() - .iter() - .copied() - .chain(self.height.to_be_bytes().iter().copied()), + .into_iter() + .chain(self.height.to_be_bytes().into_iter()), ) + .boxed() } } @@ -382,6 +398,7 @@ impl Rectangle { } } +#[async_trait] impl WriteMessage for Rectangle { fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { async move { @@ -394,7 +411,7 @@ impl WriteMessage for Rectangle { stream.write_i32(encoding_type).await?; // TODO: avoid collect alloc? - let data: Vec<_> = self.data.encode(todo!()).collect(); + let data: Vec<_> = self.data.encode(todo!()).await.collect().await; stream.write_all(&data).await?; Ok(()) @@ -402,8 +419,8 @@ impl WriteMessage for Rectangle { .boxed() } - fn encode(self, ctx: &mut ConnectionContext) -> Box + '_> { - Box::new( + async fn encode(&self, ctx: Arc) -> BoxStream { + stream::iter( self.position .x .to_be_bytes() @@ -411,12 +428,14 @@ impl WriteMessage for Rectangle { .chain(self.position.y.to_be_bytes().into_iter()) .chain(self.dimensions.width.to_be_bytes().into_iter()) .chain(self.dimensions.height.to_be_bytes().into_iter()) - .chain(i32::from(self.data.get_type()).to_be_bytes().into_iter()) - .chain(self.data.encode(ctx)), + .chain(i32::from(self.data.get_type()).to_be_bytes().into_iter()), ) + .chain(self.data.encode(ctx.clone()).await) + .boxed() } } +#[async_trait] impl WriteMessage for FramebufferUpdate { fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { async move { @@ -440,15 +459,20 @@ impl WriteMessage for FramebufferUpdate { .boxed() } - fn encode(self, ctx: &mut ConnectionContext) -> Box + '_> { + async fn encode(&self, ctx: Arc) -> BoxStream { // number of rectangles let n_rect = self.rectangles.len() as u16; - Box::new( - once(0u8) // TODO: type function? - .chain(once(0u8)) // 1 byte of padding - .chain(n_rect.to_be_bytes().into_iter()) - .chain(self.rectangles.into_iter().flat_map(|r| r.encode(ctx))), + stream::iter( + std::iter::once(0u8) // TODO: type function? + .chain(std::iter::once(0u8)) // 1 byte of padding + .chain(n_rect.to_be_bytes().into_iter()), + ) + .chain( + stream::iter(self.rectangles.iter()) + .then(move |r| r.encode(ctx.clone())) + .flatten(), ) + .boxed() } } @@ -596,6 +620,7 @@ impl ReadMessage for PixelFormat { } } +#[async_trait] impl WriteMessage for PixelFormat { fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { async move { @@ -612,14 +637,11 @@ impl WriteMessage for PixelFormat { } .boxed() } - - fn encode(self, ctx: &mut ConnectionContext) -> Box + '_> { - Box::new( - [self.bits_per_pixel, self.depth, self.big_endian as u8] - .into_iter() - .chain(self.color_spec.encode(ctx)) - .chain([0u8; 3].into_iter()), // 3 bytes of padding - ) + async fn encode(&self, ctx: Arc) -> BoxStream { + stream::iter([self.bits_per_pixel, self.depth, self.big_endian as u8].into_iter()) + .chain(self.color_spec.encode(ctx).await) + .chain(stream::iter([0u8; 3].into_iter())) // 3 bytes of padding + .boxed() } } @@ -676,6 +698,7 @@ impl ReadMessage for ColorSpecification { } } +#[async_trait] impl WriteMessage for ColorSpecification { fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { async move { @@ -703,22 +726,23 @@ impl WriteMessage for ColorSpecification { .boxed() } - fn encode(self, ctx: &mut ConnectionContext) -> Box + '_> { + async fn encode(&self, _ctx: Arc) -> BoxStream { match self { ColorSpecification::ColorFormat(cf) => { - Box::new( - once(1u8) // true color + stream::iter( + std::iter::once(1u8) // true color .chain(cf.red_max.to_be_bytes().into_iter()) .chain(cf.green_max.to_be_bytes().into_iter()) .chain(cf.blue_max.to_be_bytes().into_iter()) .chain([cf.red_shift, cf.green_shift, cf.blue_shift].into_iter()), ) + .boxed() } ColorSpecification::ColorMap(_cm) => { // first 0 byte is true-color-flag = false; // the remaining 9 are the above max/shift fields, // which aren't relevant in ColorMap mode - Box::new(std::iter::repeat(0).take(10)) + stream::iter(std::iter::repeat(0).take(10)).boxed() } } } From 15c1ce96968d6848bcf7d69fca99822e32284554 Mon Sep 17 00:00:00 2001 From: lif <> Date: Thu, 29 Aug 2024 05:18:52 +0000 Subject: [PATCH 19/20] continue weaving changes into VncServer impl --- examples/server.rs | 9 +-- src/encodings.rs | 5 +- src/encodings/zlib.rs | 12 ++++ src/rfb.rs | 163 ++++-------------------------------------- src/server.rs | 140 ++++++++++++++++++++++-------------- 5 files changed, 121 insertions(+), 208 deletions(-) diff --git a/examples/server.rs b/examples/server.rs index 623e2a9..394f271 100644 --- a/examples/server.rs +++ b/examples/server.rs @@ -11,7 +11,7 @@ use env_logger; use image::io::Reader as ImageReader; use image::GenericImageView; use log::info; -use rfb::encodings::{RawEncodingRef, ZRLEncoding}; +use rfb::encodings::{RawEncoding, RawEncodingRef, ZRLEncoding, ZlibEncoding, ZlibEncodingRef}; use rfb::pixel_formats::fourcc::FourCC; use rfb::pixel_formats::transform; use rfb::rfb::{ @@ -87,7 +87,7 @@ async fn main() -> Result<()> { width: WIDTH as u16, height: HEIGHT as u16, input_pixel_format: pixfmt.clone(), - connection_context: ConnectionContext::default(), + connection_context: ConnectionContext::default().into(), }; let server = ExampleServer { display: args.image, @@ -164,12 +164,13 @@ impl Server for ExampleServer { 0, pixels_width, pixels_height, - Box::new(ZRLEncoding::from(&RawEncodingRef::new( + // TODO: untangle so we can ZlibEncodingRef... + Box::new(ZlibEncoding::from(RawEncoding::from(&RawEncodingRef::new( &pixels, pixels_width, pixels_height, &self.pixfmt, - ))), + )))), ); FramebufferUpdate::new(vec![r]) } diff --git a/src/encodings.rs b/src/encodings.rs index 3fc0e28..7a25def 100644 --- a/src/encodings.rs +++ b/src/encodings.rs @@ -21,9 +21,12 @@ mod zlib; pub use raw::RawEncoding; pub use raw::RawEncodingRef; + +pub use zlib::ZlibEncoding; +pub use zlib::ZlibEncodingRef; + pub use trle::TRLEncoding; pub use trle::ZRLEncoding; -pub use zlib::ZlibEncoding; #[derive(Debug)] #[allow(unused)] diff --git a/src/encodings/zlib.rs b/src/encodings/zlib.rs index e1c1df5..5e13bea 100644 --- a/src/encodings/zlib.rs +++ b/src/encodings/zlib.rs @@ -14,6 +14,12 @@ pub struct ZlibEncodingRef<'a> { raw: RawEncodingRef<'a>, } +impl<'a> From> for ZlibEncodingRef<'a> { + fn from(raw: RawEncodingRef<'a>) -> Self { + Self { raw } + } +} + #[async_trait] impl<'a> Encoding for ZlibEncodingRef<'a> { fn get_type(&self) -> EncodingType { @@ -55,6 +61,12 @@ pub struct ZlibEncoding { raw: RawEncoding, } +impl From for ZlibEncoding { + fn from(raw: RawEncoding) -> Self { + Self { raw } + } +} + #[async_trait] impl Encoding for ZlibEncoding { fn get_type(&self) -> EncodingType { diff --git a/src/rfb.rs b/src/rfb.rs index 38d8b29..a0371a5 100644 --- a/src/rfb.rs +++ b/src/rfb.rs @@ -15,7 +15,7 @@ use futures::{ FutureExt, StreamExt, }; use thiserror::Error; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::io::{AsyncReadExt, AsyncWrite, AsyncWriteExt}; use tokio::net::TcpStream; use crate::encodings::{Encoding, EncodingType}; @@ -60,9 +60,18 @@ pub trait ReadMessage { } #[async_trait] -pub trait WriteMessage { +pub trait WriteMessage: Sized { async fn encode(&self, ctx: Arc) -> BoxStream; - fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>>; + async fn write_to<'a, W: AsyncWrite + Unpin + Send>( + self, + writer: &'a mut W, + ctx: Arc, + ) -> Result<(), ProtocolError> { + writer + .write_all(&self.encode(ctx).await.collect::>().await) + .await?; + Ok(()) + } } #[derive(Debug, Copy, Clone, PartialEq, PartialOrd)] @@ -108,13 +117,6 @@ impl Into<&'static [u8; 12]> for &ProtoVersion { #[async_trait] impl WriteMessage for ProtoVersion { - fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { - async move { - let s: &[u8; 12] = (&self).into(); - Ok(stream.write_all(s).await?) - } - .boxed() - } async fn encode(&self, _ctx: Arc) -> BoxStream { let s: &[u8; 12] = self.into(); stream::iter(s.iter().copied()).boxed() @@ -134,18 +136,6 @@ pub enum SecurityType { #[async_trait] impl WriteMessage for SecurityTypes { - fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { - async move { - // TODO: fix cast - stream.write_u8(self.0.len() as u8).await?; - for t in self.0.into_iter() { - t.write_to(stream).await?; - } - - Ok(()) - } - .boxed() - } async fn encode(&self, ctx: Arc) -> BoxStream { stream::iter([self.0.len() as u8].into_iter()) // TODO: fix cast .chain( @@ -173,14 +163,6 @@ impl ReadMessage for SecurityType { #[async_trait] impl WriteMessage for SecurityType { - fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { - async move { - stream.write_u8(self as u8).await?; - Ok(()) - } - .boxed() - } - async fn encode(&self, _ctx: Arc) -> BoxStream { stream::iter([*self as u8].into_iter()).boxed() } @@ -194,22 +176,6 @@ pub enum SecurityResult { #[async_trait] impl WriteMessage for SecurityResult { - fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { - async move { - match self { - SecurityResult::Success => { - stream.write_u32(0).await?; - } - SecurityResult::Failure(s) => { - stream.write_u32(1).await?; - stream.write_all(s.as_bytes()).await?; - } - }; - - Ok(()) - } - .boxed() - } async fn encode(&self, _ctx: Arc) -> BoxStream { match self { SecurityResult::Success => stream::iter(0u32.to_be_bytes().into_iter()).boxed(), @@ -259,20 +225,6 @@ impl ServerInit { #[async_trait] impl WriteMessage for ServerInit { - fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { - async move { - self.initial_res.write_to(stream).await?; - self.pixel_format.write_to(stream).await?; - - // TODO: cast properly - stream.write_u32(self.name.len() as u32).await?; - stream.write_all(self.name.as_bytes()).await?; - - Ok(()) - } - .boxed() - } - async fn encode(&self, ctx: Arc) -> BoxStream { self.initial_res .encode(ctx.clone()) @@ -352,15 +304,6 @@ impl ReadMessage for Resolution { #[async_trait] impl WriteMessage for Resolution { - fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { - async move { - stream.write_u16(self.width).await?; - stream.write_u16(self.height).await?; - Ok(()) - } - .boxed() - } - async fn encode(&self, _ctx: Arc) -> BoxStream { stream::iter( self.width @@ -400,25 +343,6 @@ impl Rectangle { #[async_trait] impl WriteMessage for Rectangle { - fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { - async move { - let encoding_type: i32 = self.data.get_type().into(); - - stream.write_u16(self.position.x).await?; - stream.write_u16(self.position.y).await?; - stream.write_u16(self.dimensions.width).await?; - stream.write_u16(self.dimensions.height).await?; - stream.write_i32(encoding_type).await?; - - // TODO: avoid collect alloc? - let data: Vec<_> = self.data.encode(todo!()).await.collect().await; - stream.write_all(&data).await?; - - Ok(()) - } - .boxed() - } - async fn encode(&self, ctx: Arc) -> BoxStream { stream::iter( self.position @@ -437,28 +361,6 @@ impl WriteMessage for Rectangle { #[async_trait] impl WriteMessage for FramebufferUpdate { - fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { - async move { - // TODO: type function? - stream.write_u8(0).await?; - - // 1 byte of padding - stream.write_u8(0).await?; - - // number of rectangles - let n_rect = self.rectangles.len() as u16; - stream.write_u16(n_rect).await?; - - // rectangles - for r in self.rectangles.into_iter() { - r.write_to(stream).await?; - } - - Ok(()) - } - .boxed() - } - async fn encode(&self, ctx: Arc) -> BoxStream { // number of rectangles let n_rect = self.rectangles.len() as u16; @@ -622,21 +524,6 @@ impl ReadMessage for PixelFormat { #[async_trait] impl WriteMessage for PixelFormat { - fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { - async move { - stream.write_u8(self.bits_per_pixel).await?; - stream.write_u8(self.depth).await?; - stream.write_u8(if self.big_endian { 1 } else { 0 }).await?; - self.color_spec.write_to(stream).await?; - - // 3 bytes of padding - let buf = [0u8; 3]; - stream.write_all(&buf).await?; - - Ok(()) - } - .boxed() - } async fn encode(&self, ctx: Arc) -> BoxStream { stream::iter([self.bits_per_pixel, self.depth, self.big_endian as u8].into_iter()) .chain(self.color_spec.encode(ctx).await) @@ -700,32 +587,6 @@ impl ReadMessage for ColorSpecification { #[async_trait] impl WriteMessage for ColorSpecification { - fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>> { - async move { - match self { - ColorSpecification::ColorFormat(cf) => { - stream.write_u8(1).await?; // true color - stream.write_u16(cf.red_max).await?; - stream.write_u16(cf.green_max).await?; - stream.write_u16(cf.blue_max).await?; - - stream.write_u8(cf.red_shift).await?; - stream.write_u8(cf.green_shift).await?; - stream.write_u8(cf.blue_shift).await?; - } - ColorSpecification::ColorMap(_cm) => { - // first 0 byte is true-color-flag = false; - // the remaining 9 are the above max/shift fields, - // which aren't relevant in ColorMap mode - stream.write_all(&[0u8; 10]).await?; - } - }; - - Ok(()) - } - .boxed() - } - async fn encode(&self, _ctx: Arc) -> BoxStream { match self { ColorSpecification::ColorFormat(cf) => { diff --git a/src/server.rs b/src/server.rs index d0c4599..722f709 100644 --- a/src/server.rs +++ b/src/server.rs @@ -10,14 +10,15 @@ use std::net::SocketAddr; use std::sync::Arc; use async_trait::async_trait; +use cancel_safe_futures::sync::RobustMutex; use futures::future::Shared; -use futures::FutureExt; +use futures::{FutureExt, StreamExt}; use log::{debug, error, info, trace}; use thiserror::Error; use tokio::io::AsyncWriteExt; use tokio::net::{TcpListener, TcpStream}; use tokio::select; -use tokio::sync::{oneshot, Mutex}; +use tokio::sync::oneshot; use crate::rfb::{ ClientInit, ClientMessage, ConnectionContext, FramebufferUpdate, KeyEvent, PixelFormat, @@ -43,6 +44,9 @@ pub enum HandshakeError { #[error(transparent)] Protocol(#[from] ProtocolError), + + #[error(transparent)] + Io(#[from] std::io::Error), } /// Immutable state @@ -62,7 +66,7 @@ pub struct VncServerData { /// get_framebuffer_update. pub input_pixel_format: PixelFormat, /// State used during encoding, such as the Zlib stream (which is shared between rectangles). - pub connection_context: ConnectionContext, + pub connection_context: Arc, } pub struct VncServer { @@ -70,13 +74,13 @@ pub struct VncServer { config: VncServerConfig, /// VNC runtime mutable state - data: Mutex, + data: Arc>, /// The underlying [`Server`] implementation pub server: S, /// One-shot channel used to signal that the server should shut down. - stop_ch: Mutex>>, + stop_ch: RobustMutex>>, } #[async_trait] @@ -95,21 +99,33 @@ impl VncServer { ); Arc::new(Self { config, - data: Mutex::new(data), + data: Arc::new(RobustMutex::new(data)), server, - stop_ch: Mutex::new(None), + stop_ch: RobustMutex::new(None), }) } pub async fn set_pixel_format(&self, pixel_format: PixelFormat) { - let mut locked = self.data.lock().await; - locked.input_pixel_format = pixel_format; + self.data + .lock() + .await + .unwrap() + .perform(move |locked| locked.input_pixel_format = pixel_format); } pub async fn set_resolution(&self, width: u16, height: u16) { - let mut locked = self.data.lock().await; - locked.width = width; - locked.height = height; + self.data.lock().await.unwrap().perform(move |locked| { + locked.width = width; + locked.height = height; + }); + } + + async fn get_ctx(&self) -> Arc { + self.data + .lock() + .await + .unwrap() + .perform(|data| data.connection_context.to_owned()) } async fn rfb_handshake( @@ -119,7 +135,10 @@ impl VncServer { ) -> Result<(), HandshakeError> { // ProtocolVersion handshake info!("Tx [{:?}]: ProtoVersion={:?}", addr, self.config.version); - self.config.version.write_to(s).await?; + let ctx = self.get_ctx().await; + + self.config.version.write_to(s, ctx.to_owned()).await?; + let client_version = ProtoVersion::read_from(s).await?; info!("Rx [{:?}]: ClientVersion={:?}", addr, client_version); @@ -138,13 +157,14 @@ impl VncServer { // Security Handshake let supported_types = self.config.sec_types.clone(); info!("Tx [{:?}]: SecurityTypes={:?}", addr, supported_types); - supported_types.write_to(s).await?; + supported_types.write_to(s, ctx.to_owned()).await?; + let client_choice = SecurityType::read_from(s).await?; info!("Rx [{:?}]: SecurityType Choice={:?}", addr, client_choice); if !self.config.sec_types.0.contains(&client_choice) { info!("Tx [{:?}]: SecurityResult=Failure", addr); let failure = SecurityResult::Failure("unsupported security type".to_string()); - failure.write_to(s).await?; + failure.write_to(s, ctx).await?; let err_str = format!("invalid security choice={:?}", client_choice); error!("{}", err_str); return Err(HandshakeError::IncompatibleSecurityTypes { @@ -155,7 +175,7 @@ impl VncServer { let res = SecurityResult::Success; info!("Tx: SecurityResult=Success"); - res.write_to(s).await?; + res.write_to(s, ctx).await?; Ok(()) } @@ -173,15 +193,17 @@ impl VncServer { false => {} } - let data = self.data.lock().await; - let server_init = ServerInit::new( - data.width, - data.height, - self.config.name.clone(), - data.input_pixel_format.clone(), - ); + let server_init = self.data.lock().await.unwrap().perform(|data| { + ServerInit::new( + data.width, + data.height, + self.config.name.clone(), + data.input_pixel_format.clone(), + ) + }); + info!("Tx [{:?}]: ServerInit={:#?}", addr, server_init); - server_init.write_to(s).await?; + server_init.write_to(s, self.get_ctx().await).await?; Ok(()) } @@ -204,9 +226,12 @@ impl VncServer { return; } - let data = self.data.lock().await; - let mut output_pixel_format = data.input_pixel_format.clone(); - drop(data); + let mut output_pixel_format = self + .data + .lock() + .await + .unwrap() + .perform(|data| data.input_pixel_format.clone()); loop { let req = select! { @@ -238,29 +263,29 @@ impl VncServer { let mut fbu = self.server.get_framebuffer_update().await; - let data = self.data.lock().await; - - // We only need to change pixel formats if the client requested a different - // one than what's specified in the input. - // - // For now, we only support transformations between 4-byte RGB formats, so - // if the requested format isn't one of those, we'll just leave the pixels - // as is. - if data.input_pixel_format == output_pixel_format { - debug!("no input transformation needed"); - } else if data.input_pixel_format.is_supported() - && output_pixel_format.is_supported() - { - debug!( - "transforming: input={:#?}, output={:#?}", - data.input_pixel_format, output_pixel_format - ); - fbu = fbu.transform(&data.input_pixel_format, &output_pixel_format); - } else { - debug!("cannot transform between pixel formats: input.is_supported()={}, output.is_supported()={}", data.input_pixel_format.is_supported(), output_pixel_format.is_supported()); - } - - if let Err(e) = fbu.write_to(s).await { + self.data.lock().await.unwrap().perform(|data| { + // We only need to change pixel formats if the client requested a different + // one than what's specified in the input. + // + // For now, we only support transformations between 4-byte RGB formats, so + // if the requested format isn't one of those, we'll just leave the pixels + // as is. + if data.input_pixel_format == output_pixel_format { + debug!("no input transformation needed"); + } else if data.input_pixel_format.is_supported() + && output_pixel_format.is_supported() + { + debug!( + "transforming: input={:#?}, output={:#?}", + data.input_pixel_format, output_pixel_format + ); + fbu = fbu.transform(&data.input_pixel_format, &output_pixel_format); + } else { + debug!("cannot transform between pixel formats: input.is_supported()={}, output.is_supported()={}", data.input_pixel_format.is_supported(), output_pixel_format.is_supported()); + } + }); + + if let Err(e) = fbu.write_to(s, self.get_ctx().await).await { error!( "[{:?}] could not write FramebufferUpdateRequest: {:?}", addr, e @@ -296,7 +321,12 @@ impl VncServer { // Create a channel to signal the server to stop. let (close_tx, close_rx) = oneshot::channel(); assert!( - self.stop_ch.lock().await.replace(close_tx).is_none(), + self.stop_ch + .lock() + .await + .unwrap() + .perform(|lock| lock.replace(close_tx)) + .is_none(), "server already started" ); let mut close_rx = close_rx.shared(); @@ -327,7 +357,13 @@ impl VncServer { /// Stop the server (and disconnect any client) if it's running. pub async fn stop(self: &Arc) { - if let Some(close_tx) = self.stop_ch.lock().await.take() { + if let Some(close_tx) = self + .stop_ch + .lock() + .await + .unwrap() + .perform(|lock| lock.take()) + { let _ = close_tx.send(()); } } From a279967533d84484c4d18097676403c3b8103d1c Mon Sep 17 00:00:00 2001 From: lif <> Date: Thu, 29 Aug 2024 05:45:19 +0000 Subject: [PATCH 20/20] include size ahead of zlib encodings. getting inflate errors in clients though.. --- examples/server.rs | 7 +++---- src/encodings/trle.rs | 22 ++++++++++++++++++++-- src/encodings/zlib.rs | 16 ++++++++++++++-- 3 files changed, 37 insertions(+), 8 deletions(-) diff --git a/examples/server.rs b/examples/server.rs index 394f271..96feeed 100644 --- a/examples/server.rs +++ b/examples/server.rs @@ -11,7 +11,7 @@ use env_logger; use image::io::Reader as ImageReader; use image::GenericImageView; use log::info; -use rfb::encodings::{RawEncoding, RawEncodingRef, ZRLEncoding, ZlibEncoding, ZlibEncodingRef}; +use rfb::encodings::{RawEncodingRef, ZRLEncoding}; use rfb::pixel_formats::fourcc::FourCC; use rfb::pixel_formats::transform; use rfb::rfb::{ @@ -164,13 +164,12 @@ impl Server for ExampleServer { 0, pixels_width, pixels_height, - // TODO: untangle so we can ZlibEncodingRef... - Box::new(ZlibEncoding::from(RawEncoding::from(&RawEncodingRef::new( + Box::new(ZRLEncoding::from(&RawEncodingRef::new( &pixels, pixels_width, pixels_height, &self.pixfmt, - )))), + ))), ); FramebufferUpdate::new(vec![r]) } diff --git a/src/encodings/trle.rs b/src/encodings/trle.rs index b438ab5..c40392f 100644 --- a/src/encodings/trle.rs +++ b/src/encodings/trle.rs @@ -39,8 +39,26 @@ impl Encoding for ZRLEncoding { } async fn encode(&self, ctx: Arc) -> BoxStream { - todo!("flate2 with zlib stream shared with stream (but flushed to byte boundary at end of this fn)"); - todo!("also disable re-use of palettes in zrle mode") + let in_buf = self + .0 + .encode(ctx.to_owned()) + .await + .collect::>() + .await; + let out_buf = ctx.zlib.lock().await.unwrap().perform(|zlib| { + let mut out_buf = Vec::with_capacity(in_buf.len()); + zlib.compress_vec(&in_buf, &mut out_buf, flate2::FlushCompress::Sync) + .expect("zlib error"); + out_buf + }); + stream::iter( + (out_buf.len() as u32) + .to_be_bytes() + .into_iter() + .chain(out_buf.into_iter()), + ) + .boxed() + // todo!("also disable re-use of palettes in zrle mode") } fn transform(&self, output: &PixelFormat) -> Box { diff --git a/src/encodings/zlib.rs b/src/encodings/zlib.rs index 5e13bea..3ffb573 100644 --- a/src/encodings/zlib.rs +++ b/src/encodings/zlib.rs @@ -41,7 +41,13 @@ impl<'a> Encoding for ZlibEncodingRef<'a> { zlib.compress_vec(in_buf, &mut out_buf, flate2::FlushCompress::Sync) .expect("zlib error") }); - stream::iter(out_buf.into_iter()).boxed() + stream::iter( + (out_buf.len() as u32) + .to_be_bytes() + .into_iter() + .chain(out_buf.into_iter()), + ) + .boxed() } fn transform(&self, output: &crate::rfb::PixelFormat) -> Box { @@ -88,7 +94,13 @@ impl Encoding for ZlibEncoding { zlib.compress_vec(in_buf, &mut out_buf, flate2::FlushCompress::Sync) .expect("zlib error") }); - stream::iter(out_buf.into_iter()).boxed() + stream::iter( + (out_buf.len() as u32) + .to_be_bytes() + .into_iter() + .chain(out_buf.into_iter()), + ) + .boxed() } fn transform(&self, output: &crate::rfb::PixelFormat) -> Box {