diff --git a/Cargo.toml b/Cargo.toml index f9b639f..2f748c6 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" @@ -23,10 +23,12 @@ 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" tokio = { version = "1", features = ["full"] } +cancel-safe-futures = "0.1" [dev-dependencies] anyhow = "1.0" diff --git a/examples/server.rs b/examples/server.rs index affdf67..96feeed 100644 --- a/examples/server.rs +++ b/examples/server.rs @@ -4,21 +4,21 @@ // // Copyright 2022 Oxide Computer Company -use anyhow::{bail, Result}; +use anyhow::Result; use async_trait::async_trait; use clap::{Parser, ValueEnum}; use env_logger; use image::io::Reader as ImageReader; use image::GenericImageView; use log::info; -use rfb::encodings::RawEncoding; +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, -}; -use rfb::{ - pixel_formats::rgb_888, - server::{Server, VncServer, VncServerConfig, VncServerData}, + ConnectionContext, FramebufferUpdate, KeyEvent, PixelFormat, ProtoVersion, Rectangle, + SecurityType, SecurityTypes, }; +use rfb::server::{Server, VncServer, VncServerConfig, VncServerData}; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; const WIDTH: usize = 1024; @@ -43,21 +43,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 +62,7 @@ enum Image { #[derive(Clone)] struct ExampleServer { display: Image, - rgb_order: (u8, u8, u8), - big_endian: bool, + pixfmt: PixelFormat, } #[tokio::main] @@ -83,22 +70,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 +86,12 @@ 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(), + connection_context: ConnectionContext::default().into(), }; 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 +99,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,13 +158,18 @@ 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, pixels_width, pixels_height, - Box::new(RawEncoding::new(pixels)), + Box::new(ZRLEncoding::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 94dbf0d..7a25def 100644 --- a/src/encodings.rs +++ b/src/encodings.rs @@ -4,13 +4,30 @@ // // Copyright 2022 Oxide Computer Company -use crate::{ - pixel_formats::rgb_888, - rfb::{PixelFormat, Position, Resolution}, -}; +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; +mod raw; +// mod rre; +mod trle; +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; + #[derive(Debug)] #[allow(unused)] pub enum EncodingType { @@ -30,17 +47,25 @@ pub enum EncodingType { Other(i32), } -pub trait Encoding -where - Self: Send, -{ +struct Pixel { + bytes: Vec, +} + +#[async_trait] +pub trait Encoding: Send + Sync { 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); + + /// Return the pixel format of this encoding's data. + fn pixel_format(&self) -> &PixelFormat; - /// Translates this encoding type from an input pixel format to an output format. - fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box; + /// Transform this encoding from its representation into a byte sequence that can be passed to the client. + 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; } impl From for i32 { @@ -80,75 +105,7 @@ impl From for EncodingType { 21 => JPEG, 6 => Zlib, -314 => CursorWithAlpha, - v => EncodingType::Other(v), + v => Other(v), } } } - -/// 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 - } - - fn transform(&self, input: &PixelFormat, output: &PixelFormat) -> Box { - // XXX: This assumes the pixel formats are both rgb888. The server code verifies this - // before calling. - assert!(input.is_rgb_888()); - assert!(output.is_rgb_888()); - - Box::new(Self { - pixels: rgb_888::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>, -} - -#[allow(dead_code)] -enum HextileTile { - Raw(Vec), - Encoded(HextileTileEncoded), -} - -#[allow(dead_code)] -struct HextileTileEncoded { - background: Option, - foreground: Option, - // TODO: finish this -} diff --git a/src/encodings/hextile.rs b/src/encodings/hextile.rs new file mode 100644 index 0000000..4529317 --- /dev/null +++ b/src/encodings/hextile.rs @@ -0,0 +1,88 @@ +use crate::{ + encodings::{Encoding, EncodingType, Pixel}, + rfb::{ConnectionContext, PixelFormat}, +}; + +use super::RawEncodingRef; + +#[allow(dead_code)] +struct HextileEncoding { + tiles: Vec>, + width: u16, + height: u16, + pixfmt: PixelFormat, +} + +impl From<&RawEncodingRef<'_>> for HextileEncoding { + fn from(raw: &RawEncodingRef) -> Self { + let pixfmt = raw.pixel_format().to_owned(); + let (width, height) = raw.dimensions(); + Self { + tiles: todo!("create subrects. need dimensions"), + width, + height, + pixfmt, + } + } +} + +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, _ctx: &mut ConnectionContext) -> Box + '_> { + Box::new(self.tiles.iter().flat_map(|tile| { + let subencoding_mask = todo!(); + [todo!()].into_iter() + })) + } + + 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)] +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..f2e8505 --- /dev/null +++ b/src/encodings/raw.rs @@ -0,0 +1,137 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use futures::{ + stream::{self, BoxStream}, + StreamExt, +}; + +use crate::{ + encodings::{Encoding, EncodingType}, + pixel_formats::transform, + rfb::{ConnectionContext, PixelFormat}, +}; + +/// Section 7.7.1 +pub struct RawEncoding { + pub(crate) pixels: Vec, + pub(crate) width: u16, + height: u16, + pixfmt: PixelFormat, +} + +impl RawEncoding { + pub fn new(pixels: Vec, width: u16, height: u16, pixfmt: &PixelFormat) -> Self { + Self { + pixels, + width, + height, + pixfmt: pixfmt.clone(), + } + } + + pub(crate) fn raw_buffer(&self) -> &[u8] { + &self.pixels + } +} + +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(), + } + } +} + +#[async_trait] +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 + } + + async fn encode(&self, _ctx: Arc) -> BoxStream { + stream::iter(self.pixels.iter().copied()).boxed() + } + + 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] { + &self.pixels + } +} + +#[async_trait] +impl<'a> Encoding for RawEncodingRef<'a> { + fn get_type(&self) -> EncodingType { + EncodingType::Raw + } + + fn dimensions(&self) -> (u16, u16) { + (self.width, self.height) + } + + fn pixel_format(&self) -> &PixelFormat { + &self.pixfmt + } + + async fn encode(&self, _ctx: Arc) -> BoxStream { + stream::iter(self.pixels.iter().copied()).boxed() + } + + 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(), + }) + } +} diff --git a/src/encodings/rre.rs b/src/encodings/rre.rs new file mode 100644 index 0000000..b601b9c --- /dev/null +++ b/src/encodings/rre.rs @@ -0,0 +1,87 @@ +use crate::{ + encodings::{Encoding, EncodingType, Pixel}, + pixel_formats::transform, + rfb::{PixelFormat, Position, Resolution}, +}; + +use super::RawEncodingRef; + +struct RREncoding { + background_pixel: Pixel, + sub_rectangles: Vec, + width: u16, + height: u16, + pixfmt: PixelFormat, +} + +struct RRESubrectangle { + pixel: Pixel, + position: Position, + dimensions: Resolution, +} + +impl Encoding for RREncoding { + fn get_type(&self) -> EncodingType { + EncodingType::RRE + } + + fn encode(&self) -> Box + '_> { + Box::new(stream::iter( + (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, output: &PixelFormat) -> Box { + let input = &self.pixfmt; + 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, + 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 new file mode 100644 index 0000000..c40392f --- /dev/null +++ b/src/encodings/trle.rs @@ -0,0 +1,421 @@ +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; +use crate::rfb::{ConnectionContext, PixelFormat}; + +use super::RawEncodingRef; + +pub struct RLEncoding { + tiles: Vec>, + width: u16, + height: u16, + pixfmt: PixelFormat, +} + +const TRLE_TILE_PX_SIZE: usize = 16; +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 + } + + fn dimensions(&self) -> (u16, u16) { + self.0.dimensions() + } + + fn pixel_format(&self) -> &PixelFormat { + self.0.pixel_format() + } + + async fn encode(&self, ctx: Arc) -> BoxStream { + 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 { + Box::new(Self(self.0.transform_inner(output))) + } +} + +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; + + 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 (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, + 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); + +// 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) +#[derive(Clone)] +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 { runs: Vec<(CPixel, usize)> }, + /// 129 + PaletteRLEReused { runs: Vec<(u8, usize)> }, + /// 130-255 + PaletteRLE { + palette: Vec, + runs: Vec<(u8, usize)>, + }, +} + +fn rle(mut length: usize) -> impl Iterator + Send { + 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 + Send> { + 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. + /// 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 + Send + '_> { + 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)), + ), + } + } +} + +// TODO: [u8; 4] so we can derive Copy and go fast +#[derive(Clone)] +struct CPixel { + bytes: Vec, +} + +enum CPixelTransformType { + AsIs, + AppendZero, + PrependZero, +} + +impl CPixel { + fn which_padding(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::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::which_padding(output) { + CPixelTransformType::AsIs => (), + CPixelTransformType::AppendZero => out_bytes.push(0u8), + CPixelTransformType::PrependZero => out_bytes.insert(0, 0u8), + } + 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 } + } +} + +#[async_trait] +impl Encoding for RLEncoding { + fn get_type(&self) -> EncodingType { + EncodingType::TRLE + } + + 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 { + Box::new(self.transform_inner(output)) + } + + fn dimensions(&self) -> (u16, u16) { + (self.width, self.height) + } + + fn pixel_format(&self) -> &PixelFormat { + &self.pixfmt + } +} diff --git a/src/encodings/zlib.rs b/src/encodings/zlib.rs new file mode 100644 index 0000000..3ffb573 --- /dev/null +++ b/src/encodings/zlib.rs @@ -0,0 +1,117 @@ +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}; + +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 { + EncodingType::Zlib + } + + fn dimensions(&self) -> (u16, u16) { + self.raw.dimensions() + } + + fn pixel_format(&self) -> &crate::rfb::PixelFormat { + self.raw.pixel_format() + } + + 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.lock().await.unwrap().perform(|zlib| { + zlib.compress_vec(in_buf, &mut out_buf, flate2::FlushCompress::Sync) + .expect("zlib error") + }); + 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 { + 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 From for ZlibEncoding { + fn from(raw: RawEncoding) -> Self { + Self { raw } + } +} + +#[async_trait] +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() + } + + 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.lock().await.unwrap().perform(|zlib| { + zlib.compress_vec(in_buf, &mut out_buf, flate2::FlushCompress::Sync) + .expect("zlib error") + }); + 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 { + 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/pixel_formats.rs b/src/pixel_formats.rs index faa39eb..3312a46 100644 --- a/src/pixel_formats.rs +++ b/src/pixel_formats.rs @@ -53,10 +53,14 @@ //! - 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}")] UnsupportedFourCc(u32), + #[error("invalid fourcc name: {0}")] + InvalidFourCcString(String), } /// Utility functions and constants related to fourcc codes. @@ -68,102 +72,219 @@ 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; + use std::str::FromStr; + + #[derive(Copy, Clone, Debug)] + #[repr(u32)] + pub enum FourCC { + /// little-endian xRGB, 8:8:8:8 + XR24 = u32::from_le_bytes(*b"XR24"), + /// little-endian RGBx, 8:8:8:8 + RX24 = u32::from_le_bytes(*b"RX24"), + /// little-endian xBGR, 8:8:8:8 + XB24 = u32::from_le_bytes(*b"XB24"), + /// little-endian BGRx, 8:8:8:8 + BX24 = u32::from_le_bytes(*b"BX24"), + /// little-endian RGB, 5:6:5 + RG16 = u32::from_le_bytes(*b"RG16"), + /// little-endian BGR, 5:6:5 + BG16 = u32::from_le_bytes(*b"BG16"), + /// RGB, 3:3:2 + RGB8 = u32::from_le_bytes(*b"RGB8"), + /// BGR, 2:3:3 + 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_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; + 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 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) + } + } - // 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 + 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 +302,114 @@ 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 (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_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 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; - } + // 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); - buf - } -} + let mut buf = Vec::with_capacity(pixels.len() * in_bytes_pp / out_bytes_pp); -#[cfg(test)] -mod tests { - use crate::pixel_formats::rgb_888::{color_shift_to_index, rgbx_index}; + let ColorSpecification::ColorFormat(in_cf) = &input.color_spec else { + unimplemented!("converting from indexed color mode"); + }; - use super::{fourcc, rgb_888::transform}; + let mut i = 0; + while i < pixels.len() { + 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(in_bytes) >> in_be_shift + } else { + u32::from_le_bytes(in_bytes) + }; - #[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); - } + // 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; - #[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)); + // 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 + } - assert_eq!(rgbx_index(8, 16, 24, false), (1, 2, 3, 0)); - assert_eq!(rgbx_index(8, 16, 24, true), (2, 1, 0, 3)); + 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); - assert_eq!(rgbx_index(0, 16, 24, false), (0, 2, 3, 1)); - assert_eq!(rgbx_index(0, 16, 24, true), (3, 1, 0, 2)); + let or = or_raw << out_cf.red_shift; + let og = og_raw << out_cf.green_shift; + let ob = ob_raw << out_cf.blue_shift; - assert_eq!(rgbx_index(8, 16, 24, false), (1, 2, 3, 0)); - assert_eq!(rgbx_index(8, 16, 24, true), (2, 1, 0, 3)); + let word = or | og | ob; - assert_eq!(rgbx_index(0, 24, 8, false), (0, 3, 1, 2)); - assert_eq!(rgbx_index(0, 24, 8, true), (3, 0, 2, 1)); + 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; } + buf +} + +#[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 +429,128 @@ 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, 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); - - // 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); + #[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); } } + +// 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 870a55c..a0371a5 100644 --- a/src/rfb.rs +++ b/src/rfb.rs @@ -4,16 +4,33 @@ // // Copyright 2022 Oxide Computer Company +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::io::{AsyncReadExt, AsyncWrite, AsyncWriteExt}; use tokio::net::TcpStream; use crate::encodings::{Encoding, EncodingType}; use crate::keysym::KeySym; -use crate::pixel_formats::rgb_888; + +pub struct ConnectionContext { + pub zlib: RobustMutex, +} +impl Default for ConnectionContext { + fn default() -> Self { + Self { + zlib: RobustMutex::new(flate2::Compress::new(flate2::Compression::fast(), false)), + } + } +} #[derive(Debug, Error)] pub enum ProtocolError { @@ -42,8 +59,19 @@ pub trait ReadMessage { Self: Sized; } -pub trait WriteMessage { - fn write_to<'a>(self, stream: &'a mut TcpStream) -> BoxFuture<'a, Result<(), ProtocolError>>; +#[async_trait] +pub trait WriteMessage: Sized { + async fn encode(&self, ctx: Arc) -> BoxStream; + 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)] @@ -58,30 +86,40 @@ 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 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", - }; +impl TryFrom<[u8; 12]> for ProtoVersion { + type Error = ProtocolError; - Ok(stream.write_all(s).await?) + 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), } - .boxed() + } +} + +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", + } + } +} + +#[async_trait] +impl WriteMessage for ProtoVersion { + async fn encode(&self, _ctx: Arc) -> BoxStream { + let s: &[u8; 12] = self.into(); + stream::iter(s.iter().copied()).boxed() } } @@ -89,24 +127,23 @@ 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, - VncAuthentication, + 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 { - // 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( + stream::iter(self.0.iter()) + .then(move |t| t.encode(ctx.to_owned())) + .flatten(), + ) + .boxed() } } @@ -124,18 +161,10 @@ 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 { - let val = match self { - SecurityType::None => 0, - SecurityType::VncAuthentication => 1, - }; - stream.write_u8(val).await?; - - Ok(()) - } - .boxed() + async fn encode(&self, _ctx: Arc) -> BoxStream { + stream::iter([*self as u8].into_iter()).boxed() } } @@ -145,22 +174,15 @@ 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 { - 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(()) + async fn encode(&self, _ctx: Arc) -> BoxStream { + match self { + SecurityResult::Success => stream::iter(0u32.to_be_bytes().into_iter()).boxed(), + SecurityResult::Failure(s) => { + stream::iter(1u32.to_be_bytes().into_iter().chain(s.bytes())).boxed() + } } - .boxed() } } @@ -201,19 +223,20 @@ 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()) + .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() } } @@ -234,20 +257,19 @@ 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 } } } #[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 { @@ -264,8 +286,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 { @@ -280,13 +302,15 @@ 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(()) - } + async fn encode(&self, _ctx: Arc) -> BoxStream { + stream::iter( + self.width + .to_be_bytes() + .into_iter() + .chain(self.height.to_be_bytes().into_iter()), + ) .boxed() } } @@ -307,54 +331,49 @@ 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), } } } +#[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?; - - let data = self.data.encode(); - stream.write_all(data).await?; - - Ok(()) - } + async fn encode(&self, ctx: Arc) -> BoxStream { + stream::iter( + 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.clone()).await) .boxed() } } +#[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(()) - } + async fn encode(&self, ctx: Arc) -> BoxStream { + // number of rectangles + let n_rect = self.rectangles.len() as u16; + 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() } } @@ -391,7 +410,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 +421,7 @@ impl PixelFormat { blue_max: u16, ) -> Self { PixelFormat { - bits_per_pixel: bbp, + bits_per_pixel: bpp, depth, big_endian, color_spec: ColorSpecification::ColorFormat(ColorFormat { @@ -416,23 +435,59 @@ 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 { - 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; } } @@ -452,6 +507,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, @@ -463,26 +522,17 @@ 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) + .chain(stream::iter([0u8; 3].into_iter())) // 3 bytes of padding + .boxed() } } #[derive(Debug, Clone, PartialEq)] -#[allow(dead_code)] pub enum ColorSpecification { ColorFormat(ColorFormat), ColorMap(ColorMap), // TODO: implement @@ -500,17 +550,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?; @@ -536,28 +585,27 @@ 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) => { - unimplemented!() - } - }; - - Ok(()) + async fn encode(&self, _ctx: Arc) -> BoxStream { + match self { + ColorSpecification::ColorFormat(cf) => { + 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 + stream::iter(std::iter::repeat(0).take(10)).boxed() + } } - .boxed() } } @@ -701,7 +749,7 @@ impl KeyEvent { } bitflags! { - struct MouseButtons: u8 { + pub struct MouseButtons: u8 { const LEFT = 1 << 0; const MIDDLE = 1 << 1; const RIGHT = 1 << 2; @@ -713,10 +761,9 @@ bitflags! { } #[derive(Debug)] -#[allow(dead_code)] pub struct PointerEvent { - position: Position, - pressed: MouseButtons, + pub position: Position, + pub pressed: MouseButtons, } impl ReadMessage for PointerEvent { diff --git a/src/server.rs b/src/server.rs index 0ce98c1..722f709 100644 --- a/src/server.rs +++ b/src/server.rs @@ -10,19 +10,20 @@ use std::net::SocketAddr; use std::sync::Arc; use async_trait::async_trait; -use futures::FutureExt; +use cancel_safe_futures::sync::RobustMutex; use futures::future::Shared; +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::{Mutex, oneshot}; +use tokio::sync::oneshot; use crate::rfb::{ - ClientInit, ClientMessage, FramebufferUpdate, KeyEvent, PixelFormat, 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)] @@ -43,6 +44,9 @@ pub enum HandshakeError { #[error(transparent)] Protocol(#[from] ProtocolError), + + #[error(transparent)] + Io(#[from] std::io::Error), } /// Immutable state @@ -61,6 +65,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: Arc, } pub struct VncServer { @@ -68,19 +74,20 @@ 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] 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) {} } @@ -91,22 +98,34 @@ impl VncServer { "at least one security type must be defined" ); Arc::new(Self { - config: config, - data: Mutex::new(data), - server: server, - stop_ch: Mutex::new(None), + config, + data: Arc::new(RobustMutex::new(data)), + server, + 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( @@ -116,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); @@ -135,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 { @@ -152,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(()) } @@ -170,20 +193,27 @@ 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(()) } - 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 { @@ -196,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! { @@ -230,32 +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 - && data.input_pixel_format.is_rgb_888() - && output_pixel_format.is_rgb_888() - { - 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"); - } - - 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 @@ -270,6 +300,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); @@ -289,7 +320,15 @@ 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 + .unwrap() + .perform(|lock| lock.replace(close_tx)) + .is_none(), + "server already started" + ); let mut close_rx = close_rx.shared(); loop { @@ -309,14 +348,22 @@ 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; }); } } /// 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(()); } }