diff --git a/Cargo.lock b/Cargo.lock index d6bf5c0bb..dbe0da002 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2579,6 +2579,7 @@ dependencies = [ "bitflags 2.9.0", "bitvec", "bmp", + "bytemuck", "byteorder", "expect-test", "ironrdp-core", diff --git a/crates/ironrdp-graphics/Cargo.toml b/crates/ironrdp-graphics/Cargo.toml index 809d7eb61..c5200c013 100644 --- a/crates/ironrdp-graphics/Cargo.toml +++ b/crates/ironrdp-graphics/Cargo.toml @@ -30,6 +30,7 @@ yuvutils-rs = { version = "0.8", features = ["rdp"] } [dev-dependencies] bmp = "0.5" +bytemuck = "1.21" expect-test.workspace = true [lints] diff --git a/crates/ironrdp-graphics/src/diff.rs b/crates/ironrdp-graphics/src/diff.rs new file mode 100644 index 000000000..d1e523753 --- /dev/null +++ b/crates/ironrdp-graphics/src/diff.rs @@ -0,0 +1,333 @@ +#[derive(Debug, PartialEq, Eq, Clone, Hash)] +pub struct Rect { + pub x: usize, + pub y: usize, + pub width: usize, + pub height: usize, +} + +impl Rect { + pub fn new(x: usize, y: usize, width: usize, height: usize) -> Self { + Self { x, y, width, height } + } + + #[must_use] + pub fn add_xy(mut self, x: usize, y: usize) -> Self { + self.x += x; + self.y += y; + self + } + + fn intersect(&self, other: &Rect) -> Option { + let x = self.x.max(other.x); + let y = self.y.max(other.y); + let width = (self.x + self.width).min(other.x + other.width); + if width <= x { + return None; + } + let width = width - x; + let height = (self.y + self.height).min(other.y + other.height); + if height <= y { + return None; + } + let height = height - y; + + Some(Rect::new(x, y, width, height)) + } +} + +const TILE_SIZE: usize = 64; + +fn find_different_tiles( + image1: &[u8], + stride1: usize, + image2: &[u8], + stride2: usize, + width: usize, + height: usize, +) -> Vec { + assert!(stride1 >= width * BPP); + assert!(stride2 >= width * BPP); + assert!(image1.len() >= (height - 1) * stride1 + width * BPP); + assert!(image2.len() >= (height - 1) * stride2 + width * BPP); + + let tiles_x = width.div_ceil(TILE_SIZE); + let tiles_y = height.div_ceil(TILE_SIZE); + let mut tile_differences = vec![false; tiles_x * tiles_y]; + + tile_differences.iter_mut().enumerate().for_each(|(idx, diff)| { + let tile_start_x = (idx % tiles_x) * TILE_SIZE; + let tile_end_x = (tile_start_x + TILE_SIZE).min(width); + let tile_start_y = (idx / tiles_x) * TILE_SIZE; + let tile_end_y = (tile_start_y + TILE_SIZE).min(height); + + // Check for any difference in tile using slice comparisons + let has_diff = (tile_start_y..tile_end_y).any(|y| { + let row_start1 = y * stride1; + let row_start2 = y * stride2; + let tile_row_start1 = row_start1 + tile_start_x * BPP; + let tile_row_end1 = row_start1 + tile_end_x * BPP; + let tile_row_start2 = row_start2 + tile_start_x * BPP; + let tile_row_end2 = row_start2 + tile_end_x * BPP; + + image1[tile_row_start1..tile_row_end1] != image2[tile_row_start2..tile_row_end2] + }); + + *diff = has_diff; + }); + + tile_differences +} + +fn find_different_rects( + image1: &[u8], + stride1: usize, + image2: &[u8], + stride2: usize, + width: usize, + height: usize, +) -> Vec { + let mut tile_differences = find_different_tiles::(image1, stride1, image2, stride2, width, height); + + let mod_width = width % TILE_SIZE; + let mod_height = height % TILE_SIZE; + let tiles_x = width.div_ceil(TILE_SIZE); + let tiles_y = height.div_ceil(TILE_SIZE); + + let mut rectangles = Vec::new(); + let mut current_idx = 0; + let total_tiles = tiles_x * tiles_y; + + // Process tiles in linear fashion to find rectangular regions + while current_idx < total_tiles { + if !tile_differences[current_idx] { + current_idx += 1; + continue; + } + + let start_y = current_idx / tiles_x; + let start_x = current_idx % tiles_x; + + // Expand horizontally as much as possible + let mut max_width = 1; + while start_x + max_width < tiles_x && tile_differences[current_idx + max_width] { + max_width += 1; + } + + // Expand vertically as much as possible + let mut max_height = 1; + 'vertical: while start_y + max_height < tiles_y { + for x in 0..max_width { + let check_idx = (start_y + max_height) * tiles_x + start_x + x; + if !tile_differences[check_idx] { + break 'vertical; + } + } + max_height += 1; + } + + // Calculate pixel coordinates + let pixel_x = start_x * TILE_SIZE; + let pixel_y = start_y * TILE_SIZE; + + let pixel_width = if start_x + max_width == tiles_x && mod_width > 0 { + (max_width - 1) * TILE_SIZE + mod_width + } else { + max_width * TILE_SIZE + }; + + let pixel_height = if start_y + max_height == tiles_y && mod_height > 0 { + (max_height - 1) * TILE_SIZE + mod_height + } else { + max_height * TILE_SIZE + }; + + rectangles.push(Rect { + x: pixel_x, + y: pixel_y, + width: pixel_width, + height: pixel_height, + }); + + // Mark tiles as processed + for y in 0..max_height { + for x in 0..max_width { + let idx = (start_y + y) * tiles_x + start_x + x; + tile_differences[idx] = false; + } + } + + current_idx += max_width; + } + + rectangles +} + +/// Helper function to find different regions in two images. +/// +/// This function takes two images as input and returns a list of rectangles +/// representing the different regions between the two images, in image2 coordinates. +/// +/// ```text +/// ┌───────────────────────────────────────────┐ +/// │ image1 │ +/// │ │ +/// │ (x,y) │ +/// │ ┌───────────────┐ │ +/// │ │ image2 │ │ +/// │ │ │ │ +/// │ │ │ │ +/// │ │ │ │ +/// │ │ │ │ +/// │ │ │ │ +/// │ └───────────────┘ │ +/// │ │ +/// └───────────────────────────────────────────┘ +/// ``` +#[allow(clippy::too_many_arguments)] +pub fn find_different_rects_sub( + image1: &[u8], + stride1: usize, + width1: usize, + height1: usize, + image2: &[u8], + stride2: usize, + width2: usize, + height2: usize, + x: usize, + y: usize, +) -> Vec { + let rect1 = Rect::new(0, 0, width1, height1); + let rect2 = Rect::new(x, y, width2, height2); + let Some(inter) = rect1.intersect(&rect2) else { + return vec![]; + }; + + let image1 = &image1[y * stride1 + x * BPP..]; + find_different_rects::(image1, stride1, image2, stride2, inter.width, inter.height) +} + +#[cfg(test)] +mod tests { + use super::*; + use bytemuck::cast_slice; + + #[test] + fn test_intersect() { + let r1 = Rect::new(0, 0, 640, 480); + let r2 = Rect::new(10, 10, 10, 10); + let r3 = Rect::new(630, 470, 20, 20); + + assert_eq!(r1.intersect(&r1).as_ref(), Some(&r1)); + assert_eq!(r1.intersect(&r2).as_ref(), Some(&r2)); + assert_eq!(r1.intersect(&r3), Some(Rect::new(630, 470, 10, 10))); + assert_eq!(r2.intersect(&r3), None); + } + + #[test] + fn test_single_tile() { + const SIZE: usize = 128; + let image1 = vec![0u32; SIZE * SIZE]; + let mut image2 = vec![0u32; SIZE * SIZE]; + image2[65 * 128 + 65] = 1; + let result = + find_different_rects::<4>(cast_slice(&image1), SIZE * 4, cast_slice(&image2), SIZE * 4, SIZE, SIZE); + assert_eq!( + result, + vec![Rect { + x: 64, + y: 64, + width: 64, + height: 64 + }] + ); + } + + #[test] + fn test_adjacent_tiles() { + const SIZE: usize = 256; + let image1 = vec![0u32; SIZE * SIZE]; + let mut image2 = vec![0u32; SIZE * SIZE]; + // Modify two adjacent tiles + image2[65 * SIZE + 65] = 1; + image2[65 * SIZE + 129] = 1; + let result = + find_different_rects::<4>(cast_slice(&image1), SIZE * 4, cast_slice(&image2), SIZE * 4, SIZE, SIZE); + assert_eq!( + result, + vec![Rect { + x: 64, + y: 64, + width: 128, + height: 64 + }] + ); + } + + #[test] + fn test_edge_tiles() { + const SIZE: usize = 100; + let image1 = vec![0u32; SIZE * SIZE]; + let mut image2 = vec![0u32; SIZE * SIZE]; + image2[65 * SIZE + 65] = 1; + let result = + find_different_rects::<4>(cast_slice(&image1), SIZE * 4, cast_slice(&image2), SIZE * 4, SIZE, SIZE); + assert_eq!( + result, + vec![Rect { + x: 64, + y: 64, + width: 36, + height: 36 + }] + ); + } + + #[test] + fn test_large() { + const SIZE: usize = 4096; + let image1 = vec![0u32; SIZE * SIZE]; + let mut image2 = vec![0u32; SIZE * SIZE]; + image2[95 * 100 + 95] = 1; + let _result = + find_different_rects::<4>(cast_slice(&image1), SIZE * 4, cast_slice(&image2), SIZE * 4, SIZE, SIZE); + } + + #[test] + fn test_sub_diff() { + let image1 = vec![0u32; 2048 * 2048]; + let mut image2 = vec![0u32; 1024 * 1024]; + image2[0] = 1; + image2[1024 * 65 + 512 - 1] = 1; + + let res = find_different_rects_sub::<4>( + cast_slice(&image1), + 2048 * 4, + 2048, + 2048, + cast_slice(&image2), + 1024 * 4, + 512, + 512, + 1024, + 1024, + ); + assert_eq!( + res, + vec![ + Rect { + x: 0, + y: 0, + width: 64, + height: 64 + }, + Rect { + x: 448, + y: 64, + width: 64, + height: 64 + } + ] + ) + } +} diff --git a/crates/ironrdp-graphics/src/lib.rs b/crates/ironrdp-graphics/src/lib.rs index d5d411076..dcbed913d 100644 --- a/crates/ironrdp-graphics/src/lib.rs +++ b/crates/ironrdp-graphics/src/lib.rs @@ -7,6 +7,7 @@ #![allow(clippy::cast_sign_loss)] // FIXME: remove pub mod color_conversion; +pub mod diff; pub mod dwt; pub mod image_processing; pub mod pointer; diff --git a/crates/ironrdp-server/src/display.rs b/crates/ironrdp-server/src/display.rs index 4101aedc0..3f97fd451 100644 --- a/crates/ironrdp-server/src/display.rs +++ b/crates/ironrdp-server/src/display.rs @@ -3,6 +3,7 @@ use core::num::NonZeroU16; use anyhow::Result; use bytes::{Bytes, BytesMut}; use ironrdp_displaycontrol::pdu::DisplayControlMonitorLayout; +use ironrdp_graphics::diff; use ironrdp_pdu::pointer::PointerPositionAttribute; #[rustfmt::skip] @@ -74,12 +75,73 @@ impl TryInto for BitmapUpdate { width: self.width, height: self.height, format: self.format, - data: self.data.try_into_mut().map_err(|_| "BitmapUpdate is shared")?, + data: self.data.into(), stride: self.stride, }) } } +impl Framebuffer { + pub fn new(width: NonZeroU16, height: NonZeroU16, format: PixelFormat) -> Self { + let mut data = BytesMut::new(); + let w = usize::from(width.get()); + let h = usize::from(height.get()); + let bpp = usize::from(format.bytes_per_pixel()); + data.resize(bpp * w * h, 0); + + Self { + width, + height, + format, + data, + stride: bpp * w, + } + } + + pub fn update(&mut self, bitmap: &BitmapUpdate) { + if self.format != bitmap.format { + warn!("Bitmap format mismatch, unsupported"); + return; + } + let bpp = usize::from(self.format.bytes_per_pixel()); + let x = usize::from(bitmap.x); + let y = usize::from(bitmap.y); + let width = usize::from(bitmap.width.get()); + let height = usize::from(bitmap.height.get()); + + let data = &mut self.data; + let start = y * self.stride + x * bpp; + let end = start + (height - 1) * self.stride + width * bpp; + let dst = &mut data[start..end]; + + for y in 0..height { + let start = y * bitmap.stride; + let end = start + width * bpp; + let src = bitmap.data.slice(start..end); + + let start = y * self.stride; + let end = start + width * bpp; + let dst = &mut dst[start..end]; + + dst.copy_from_slice(&src); + } + } + + pub(crate) fn update_diffs(&mut self, bitmap: &BitmapUpdate, diffs: &[diff::Rect]) { + diffs + .iter() + .filter_map(|diff| { + let x = u16::try_from(diff.x).ok()?; + let y = u16::try_from(diff.y).ok()?; + let width = u16::try_from(diff.width).ok().and_then(NonZeroU16::new)?; + let height = u16::try_from(diff.height).ok().and_then(NonZeroU16::new)?; + + bitmap.sub(x, y, width, height) + }) + .for_each(|sub| self.update(&sub)); + } +} + /// Bitmap Display Update /// /// Bitmap updates are encoded using RDP 6.0 compression, fragmented and sent using @@ -231,3 +293,43 @@ pub trait RdpServerDisplay: Send { debug!(?layout, "Requesting layout") } } + +#[cfg(test)] +mod tests { + use super::{BitmapUpdate, Framebuffer}; + use core::num::NonZeroU16; + use ironrdp_graphics::{diff::Rect, image_processing::PixelFormat}; + + #[test] + fn framebuffer_update() { + let width = NonZeroU16::new(800).unwrap(); + let height = NonZeroU16::new(600).unwrap(); + let fmt = PixelFormat::ABgr32; + let bpp = usize::from(fmt.bytes_per_pixel()); + let mut fb = Framebuffer::new(width, height, fmt); + + let width = 15; + let stride = width * bpp; + let height = 20; + let data = vec![1u8; height * stride]; + let update = BitmapUpdate { + x: 1, + y: 2, + width: NonZeroU16::new(15).unwrap(), + height: NonZeroU16::new(20).unwrap(), + format: fmt, + data: data.into(), + stride, + }; + let diffs = vec![Rect::new(2, 3, 4, 5)]; + fb.update_diffs(&update, &diffs); + let data = fb.data; + for y in 5..10 { + for x in 3..7 { + for b in 0..bpp { + assert_eq!(data[y * fb.stride + x * bpp + b], 1); + } + } + } + } +} diff --git a/crates/ironrdp-server/src/encoder/mod.rs b/crates/ironrdp-server/src/encoder/mod.rs index 80f4cad85..288928c4d 100644 --- a/crates/ironrdp-server/src/encoder/mod.rs +++ b/crates/ironrdp-server/src/encoder/mod.rs @@ -1,7 +1,9 @@ use core::fmt; +use core::num::NonZeroU16; use anyhow::{Context, Result}; use ironrdp_acceptor::DesktopSize; +use ironrdp_graphics::diff::{find_different_rects_sub, Rect}; use ironrdp_pdu::encode_vec; use ironrdp_pdu::fast_path::UpdateCode; use ironrdp_pdu::geometry::ExclusiveRectangle; @@ -28,7 +30,6 @@ enum CodecId { pub(crate) struct UpdateEncoder { desktop_size: DesktopSize, - // FIXME: draw updates on the framebuffer framebuffer: Option, bitmap_updater: BitmapUpdater, } @@ -62,7 +63,7 @@ impl UpdateEncoder { pub(crate) fn update(&mut self, update: DisplayUpdate) -> EncoderIter<'_> { EncoderIter { encoder: self, - update: Some(update), + state: State::Start(update), } } @@ -121,14 +122,41 @@ impl UpdateEncoder { Ok(UpdateFragmenter::new(UpdateCode::PositionPointer, encode_vec(&pos)?)) } - async fn bitmap(&mut self, bitmap: BitmapUpdate) -> Result { - // Clone to satisfy spawn_blocking 'static requirement - // this should be cheap, even if using bitmap, since vec![] will be empty - let mut updater = self.bitmap_updater.clone(); - let (res, bitmap) = - tokio::task::spawn_blocking(move || time_warn!("Encoding bitmap", 10, (updater.handle(&bitmap), bitmap))) - .await - .unwrap(); + fn bitmap_diffs(&mut self, bitmap: &BitmapUpdate) -> Vec { + // TODO: we may want to make it optional for servers that already provide damaged regions + const USE_DIFFS: bool = true; + + if let Some(Framebuffer { + data, + stride, + width, + height, + .. + }) = USE_DIFFS.then_some(self.framebuffer.as_ref()).flatten() + { + find_different_rects_sub::<4>( + data, + *stride, + width.get().into(), + height.get().into(), + &bitmap.data, + bitmap.stride, + bitmap.width.get().into(), + bitmap.height.get().into(), + bitmap.x.into(), + bitmap.y.into(), + ) + } else { + vec![Rect { + x: 0, + y: 0, + width: bitmap.width.get().into(), + height: bitmap.height.get().into(), + }] + } + } + + fn bitmap_update_framebuffer(&mut self, bitmap: BitmapUpdate, diffs: &[Rect]) { if bitmap.x == 0 && bitmap.y == 0 && bitmap.width.get() == self.desktop_size.width @@ -138,32 +166,86 @@ impl UpdateEncoder { Ok(framebuffer) => self.framebuffer = Some(framebuffer), Err(err) => warn!("Failed to convert bitmap to framebuffer: {}", err), } + } else if let Some(fb) = self.framebuffer.as_mut() { + fb.update_diffs(&bitmap, diffs); } - res } + + async fn bitmap(&mut self, bitmap: BitmapUpdate) -> Result { + // Clone to satisfy spawn_blocking 'static requirement + // this should be cheap, even if using bitmap, since vec![] will be empty + let mut updater = self.bitmap_updater.clone(); + tokio::task::spawn_blocking(move || time_warn!("Encoding bitmap", 10, updater.handle(&bitmap))) + .await + .unwrap() + } +} + +#[derive(Debug, Default)] +enum State { + Start(DisplayUpdate), + BitmapDiffs { + diffs: Vec, + bitmap: BitmapUpdate, + pos: usize, + }, + #[default] + Ended, } pub(crate) struct EncoderIter<'a> { encoder: &'a mut UpdateEncoder, - update: Option, + state: State, } impl EncoderIter<'_> { pub(crate) async fn next(&mut self) -> Option> { - let update = self.update.take()?; - let encoder = &mut self.encoder; - - let res = match update { - DisplayUpdate::Bitmap(bitmap) => encoder.bitmap(bitmap).await, - DisplayUpdate::PointerPosition(pos) => UpdateEncoder::pointer_position(pos), - DisplayUpdate::RGBAPointer(ptr) => UpdateEncoder::rgba_pointer(ptr), - DisplayUpdate::ColorPointer(ptr) => UpdateEncoder::color_pointer(ptr), - DisplayUpdate::HidePointer => UpdateEncoder::hide_pointer(), - DisplayUpdate::DefaultPointer => UpdateEncoder::default_pointer(), - DisplayUpdate::Resize(_) => return None, - }; - - Some(res) + loop { + let state = core::mem::take(&mut self.state); + let encoder = &mut self.encoder; + + let res = match state { + State::Start(update) => match update { + DisplayUpdate::Bitmap(bitmap) => { + let diffs = encoder.bitmap_diffs(&bitmap); + self.state = State::BitmapDiffs { diffs, bitmap, pos: 0 }; + continue; + } + DisplayUpdate::PointerPosition(pos) => UpdateEncoder::pointer_position(pos), + DisplayUpdate::RGBAPointer(ptr) => UpdateEncoder::rgba_pointer(ptr), + DisplayUpdate::ColorPointer(ptr) => UpdateEncoder::color_pointer(ptr), + DisplayUpdate::HidePointer => UpdateEncoder::hide_pointer(), + DisplayUpdate::DefaultPointer => UpdateEncoder::default_pointer(), + DisplayUpdate::Resize(_) => return None, + }, + State::BitmapDiffs { diffs, bitmap, pos } => { + let Some(rect) = diffs.get(pos) else { + encoder.bitmap_update_framebuffer(bitmap, &diffs); + self.state = State::Ended; + return None; + }; + let Rect { x, y, width, height } = *rect; + let Some(sub) = bitmap.sub( + u16::try_from(x).unwrap(), + u16::try_from(y).unwrap(), + NonZeroU16::new(u16::try_from(width).unwrap()).unwrap(), + NonZeroU16::new(u16::try_from(height).unwrap()).unwrap(), + ) else { + warn!("Failed to extract bitmap subregion"); + return None; + }; + self.state = State::BitmapDiffs { + diffs, + bitmap, + pos: pos + 1, + }; + encoder.bitmap(sub).await + } + State::Ended => return None, + }; + + return Some(res); + } } }