From 2c7556ba1e419fc82e06ea4095f7711eb95d4a6c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marc-Andr=C3=A9=20Lureau?= Date: Tue, 25 Feb 2025 15:37:43 +0400 Subject: [PATCH] refactor(server): make UpdateEncoder::update() an iterator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A single display update can now result in multiple commands / update code (FastPathUpdate). The update dispatching and bitmap encoding is now done by the UpdateEncoder itself. Signed-off-by: Marc-André Lureau --- crates/ironrdp-server/src/encoder/bitmap.rs | 1 + crates/ironrdp-server/src/encoder/mod.rs | 65 ++++++++++++++++----- crates/ironrdp-server/src/encoder/rfx.rs | 2 +- crates/ironrdp-server/src/server.rs | 55 ++++++++--------- 4 files changed, 74 insertions(+), 49 deletions(-) diff --git a/crates/ironrdp-server/src/encoder/bitmap.rs b/crates/ironrdp-server/src/encoder/bitmap.rs index 71a1ec17..725232a0 100644 --- a/crates/ironrdp-server/src/encoder/bitmap.rs +++ b/crates/ironrdp-server/src/encoder/bitmap.rs @@ -7,6 +7,7 @@ use ironrdp_pdu::geometry::InclusiveRectangle; use crate::BitmapUpdate; // PERF: we could also remove the need for this buffer +#[derive(Clone)] pub(crate) struct BitmapEncoder { buffer: Vec, } diff --git a/crates/ironrdp-server/src/encoder/mod.rs b/crates/ironrdp-server/src/encoder/mod.rs index 89be5e3b..012e013c 100644 --- a/crates/ironrdp-server/src/encoder/mod.rs +++ b/crates/ironrdp-server/src/encoder/mod.rs @@ -12,7 +12,7 @@ use ironrdp_pdu::surface_commands::{ExtendedBitmapDataPdu, SurfaceBitsPdu, Surfa use self::bitmap::BitmapEncoder; use self::rfx::RfxEncoder; use super::BitmapUpdate; -use crate::{ColorPointer, Framebuffer, RGBAPointer}; +use crate::{time_warn, ColorPointer, DisplayUpdate, Framebuffer, RGBAPointer}; mod bitmap; mod fast_path; @@ -59,12 +59,18 @@ impl UpdateEncoder { } } + pub(crate) fn update(&mut self, update: DisplayUpdate) -> EncoderIter<'_> { + EncoderIter { + encoder: self, + update: Some(update), + } + } + pub(crate) fn set_desktop_size(&mut self, size: DesktopSize) { self.desktop_size = size; } - #[allow(clippy::unused_self)] - pub(crate) fn rgba_pointer(&mut self, ptr: RGBAPointer) -> Result { + fn rgba_pointer(ptr: RGBAPointer) -> Result { let xor_mask = ptr.data; let hot_spot = Point16 { @@ -86,8 +92,7 @@ impl UpdateEncoder { Ok(UpdateFragmenter::new(UpdateCode::NewPointer, encode_vec(&ptr)?)) } - #[allow(clippy::unused_self)] - pub(crate) fn color_pointer(&mut self, ptr: ColorPointer) -> Result { + fn color_pointer(ptr: ColorPointer) -> Result { let hot_spot = Point16 { x: ptr.hot_x, y: ptr.hot_y, @@ -103,23 +108,26 @@ impl UpdateEncoder { Ok(UpdateFragmenter::new(UpdateCode::ColorPointer, encode_vec(&ptr)?)) } - #[allow(clippy::unused_self)] - pub(crate) fn default_pointer(&mut self) -> Result { + fn default_pointer() -> Result { Ok(UpdateFragmenter::new(UpdateCode::DefaultPointer, vec![])) } - #[allow(clippy::unused_self)] - pub(crate) fn hide_pointer(&mut self) -> Result { + fn hide_pointer() -> Result { Ok(UpdateFragmenter::new(UpdateCode::HiddenPointer, vec![])) } - #[allow(clippy::unused_self)] - pub(crate) fn pointer_position(&mut self, pos: PointerPositionAttribute) -> Result { + fn pointer_position(pos: PointerPositionAttribute) -> Result { Ok(UpdateFragmenter::new(UpdateCode::PositionPointer, encode_vec(&pos)?)) } - pub(crate) fn bitmap(&mut self, bitmap: BitmapUpdate) -> Result { - let res = self.bitmap_updater.handle(&bitmap); + 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(); if bitmap.x == 0 && bitmap.y == 0 && bitmap.width.get() == self.desktop_size.width @@ -134,7 +142,31 @@ impl UpdateEncoder { } } -#[derive(Debug)] +pub(crate) struct EncoderIter<'a> { + encoder: &'a mut UpdateEncoder, + update: Option, +} + +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) + } +} + +#[derive(Debug, Clone)] enum BitmapUpdater { None(NoneHandler), Bitmap(BitmapHandler), @@ -155,7 +187,7 @@ trait BitmapUpdateHandler { fn handle(&mut self, bitmap: &BitmapUpdate) -> Result; } -#[derive(Debug)] +#[derive(Clone, Debug)] struct NoneHandler; impl BitmapUpdateHandler for NoneHandler { @@ -169,6 +201,7 @@ impl BitmapUpdateHandler for NoneHandler { } } +#[derive(Clone)] struct BitmapHandler { bitmap: BitmapEncoder, } @@ -209,7 +242,7 @@ impl BitmapUpdateHandler for BitmapHandler { } } -#[derive(Debug)] +#[derive(Debug, Clone)] struct RemoteFxHandler { remotefx: RfxEncoder, codec_id: u8, diff --git a/crates/ironrdp-server/src/encoder/rfx.rs b/crates/ironrdp-server/src/encoder/rfx.rs index 68566c0e..784775c0 100644 --- a/crates/ironrdp-server/src/encoder/rfx.rs +++ b/crates/ironrdp-server/src/encoder/rfx.rs @@ -11,7 +11,7 @@ use ironrdp_pdu::WriteCursor; use crate::BitmapUpdate; -#[derive(Debug)] +#[derive(Debug, Clone)] pub(crate) struct RfxEncoder { entropy_algorithm: rfx::EntropyAlgorithm, } diff --git a/crates/ironrdp-server/src/server.rs b/crates/ironrdp-server/src/server.rs index 2a95272b..fabf9d6e 100644 --- a/crates/ironrdp-server/src/server.rs +++ b/crates/ironrdp-server/src/server.rs @@ -32,7 +32,7 @@ use crate::clipboard::CliprdrServerFactory; use crate::display::{DisplayUpdate, RdpServerDisplay}; use crate::encoder::UpdateEncoder; use crate::handler::RdpServerInputHandler; -use crate::{builder, capabilities, time_warn, SoundServerFactory}; +use crate::{builder, capabilities, SoundServerFactory}; #[derive(Clone)] pub struct RdpServerOptions { @@ -423,39 +423,30 @@ impl RdpServer { buffer: &mut Vec, mut encoder: UpdateEncoder, ) -> Result<(RunState, UpdateEncoder)> { - let mut fragmenter = match update { - DisplayUpdate::Bitmap(bitmap) => { - let (enc, res) = task::spawn_blocking(move || { - let res = time_warn!("Encoding bitmap", 10, encoder.bitmap(bitmap)); - (encoder, res) - }) - .await?; - encoder = enc; - res - } - DisplayUpdate::PointerPosition(pos) => encoder.pointer_position(pos), - DisplayUpdate::Resize(desktop_size) => { - debug!(?desktop_size, "Display resize"); - encoder.set_desktop_size(desktop_size); - deactivate_all(io_channel_id, user_channel_id, writer).await?; - return Ok((RunState::DeactivationReactivation { desktop_size }, encoder)); - } - DisplayUpdate::RGBAPointer(ptr) => encoder.rgba_pointer(ptr), - DisplayUpdate::ColorPointer(ptr) => encoder.color_pointer(ptr), - DisplayUpdate::HidePointer => encoder.hide_pointer(), - DisplayUpdate::DefaultPointer => encoder.default_pointer(), - } - .context("error during update encoding")?; - - if fragmenter.size_hint() > buffer.len() { - buffer.resize(fragmenter.size_hint(), 0); + if let DisplayUpdate::Resize(desktop_size) = update { + debug!(?desktop_size, "Display resize"); + encoder.set_desktop_size(desktop_size); + deactivate_all(io_channel_id, user_channel_id, writer).await?; + return Ok((RunState::DeactivationReactivation { desktop_size }, encoder)); } - while let Some(len) = fragmenter.next(buffer) { - writer - .write_all(&buffer[..len]) - .await - .context("failed to write display update")?; + let mut encoder_iter = encoder.update(update); + loop { + let Some(fragmenter) = encoder_iter.next().await else { + break; + }; + + let mut fragmenter = fragmenter.context("error while encoding")?; + if fragmenter.size_hint() > buffer.len() { + buffer.resize(fragmenter.size_hint(), 0); + } + + while let Some(len) = fragmenter.next(buffer) { + writer + .write_all(&buffer[..len]) + .await + .context("failed to write display update")?; + } } Ok((RunState::Continue, encoder))