feat(server): find & send the damaged tiles

Keep a framebuffer and tile-diff against it, to save from
encoding/sending the same bitmap data regions.

Signed-off-by: Marc-André Lureau <[email protected]>

Author: elmarco

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<Framebuffer> 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);
+                }
+            }
+        }
+    }
+}

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;
@@ -29,7 +31,6 @@ enum CodecId {
 
 pub(crate) struct UpdateEncoder {
     desktop_size: DesktopSize,
-    // FIXME: draw updates on the framebuffer
     framebuffer: Option<Framebuffer>,
     bitmap_updater: BitmapUpdater,
 }
@@ -63,7 +64,7 @@ impl UpdateEncoder {
     pub(crate) fn update(&mut self, update: DisplayUpdate) -> EncoderIter<'_> {
         EncoderIter {
             encoder: self,
-            update: Some(update),
+            state: State::Start(update),
         }
     }
 
@@ -122,14 +123,41 @@ impl UpdateEncoder {
         Ok(UpdateFragmenter::new(UpdateCode::PositionPointer, encode_vec(&pos)?))
     }
 
-    async fn bitmap(&mut self, bitmap: BitmapUpdate) -> Result<UpdateFragmenter> {
-        // 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<Rect> {
+        // 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
@@ -139,32 +167,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<UpdateFragmenter> {
+        // 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<Rect>,
+        bitmap: BitmapUpdate,
+        pos: usize,
+    },
+    #[default]
+    Ended,
 }
 
 pub(crate) struct EncoderIter<'a> {
     encoder: &'a mut UpdateEncoder,
-    update: Option<DisplayUpdate>,
+    state: State,
 }
 
 impl EncoderIter<'_> {
     pub(crate) async fn next(&mut self) -> Option<Result<UpdateFragmenter>> {
-        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);
+        }
     }
 }