fix(server): leave room for the surface-bits framing when tiling updates

crates/ironrdp-server/src/encoder/mod.rs

@@ -4,7 +4,6 @@ use core::num::NonZeroU16;
 use ironrdp_acceptor::DesktopSize;
 use ironrdp_graphics::diff::{Rect, find_different_rects_sub};
 use ironrdp_pdu::codecs::rfx::Quant;
-use ironrdp_pdu::encode_vec;
 use ironrdp_pdu::fast_path::UpdateCode;
 use ironrdp_pdu::geometry::ExclusiveRectangle;
 use ironrdp_pdu::pointer::{
@@ -13,6 +12,7 @@ use ironrdp_pdu::pointer::{
 };
 use ironrdp_pdu::rdp::capability_sets::{CmdFlags, EntropyBits, LargePointerSupportFlags};
 use ironrdp_pdu::surface_commands::{ExtendedBitmapDataPdu, SurfaceBitsPdu, SurfaceCommand};
+use ironrdp_pdu::{Encode as _, encode_vec};
 use tracing::{debug, warn};
 
 use self::bitmap::BitmapEncoder;
@@ -124,10 +124,14 @@ pub(crate) struct UpdateEncoder {
     desktop_size: DesktopSize,
     framebuffer: Option<Framebuffer>,
     bitmap_updater: Option<BitmapUpdater>,
-    /// Negotiated MultifragmentUpdate reassembly buffer size. Used to split
-    /// oversized bitmaps into strips that fit within the limit when sent as
-    /// uncompressed surface commands.
-    max_request_size: usize,
+    /// Largest uncompressed tile, in bytes of pixels, whose surface-bits update
+    /// still fits the client's MultifragmentUpdate reassembly buffer: the
+    /// advertised `max_request_size` minus the framing `set_surface` puts
+    /// around the pixels. The client counts the whole reassembled update,
+    /// framing included, against that buffer and drops the connection when it
+    /// is exceeded, so a tile that fills the buffer with pixels alone is one
+    /// header too large.
+    tile_budget: usize,
     /// Client's advertised New Pointer Update cache size (MS-RDPBCGR 2.2.7.1.5
     /// `pointerCacheSize`). Zero means the client did not advertise support for the
     /// New Pointer Update; `RGBAPointer`/`CachedPointer` emission is skipped in that
@@ -191,8 +195,9 @@ impl UpdateEncoder {
             desktop_size,
             framebuffer: None,
             bitmap_updater: Some(bitmap_updater),
-            max_request_size: usize::try_from(max_request_size)
-                .map_err(|e| ServerError::custom("max_request_size", e))?,
+            tile_budget: usize::try_from(max_request_size)
+                .map_err(|e| ServerError::custom("max_request_size", e))?
+                .saturating_sub(surface_bits_framing_size()),
             pointer_cache_size,
             large_pointer_flags,
         })
@@ -363,11 +368,11 @@ impl UpdateEncoder {
             }]
         };
 
-        // Subdivide diff rects whose uncompressed size would exceed the
-        // MultifragmentUpdate reassembly buffer.
+        // Subdivide diff rects whose uncompressed size would not leave room
+        // for the update framing in the MultifragmentUpdate reassembly buffer.
         let mut tiled = Vec::with_capacity(diffs.len());
         for rect in diffs {
-            if rect.width * rect.height * 4 <= self.max_request_size {
+            if rect.width * rect.height * 4 <= self.tile_budget {
                 tiled.push(rect);
             } else {
                 let rects = self.split_diff(rect);
@@ -377,19 +382,19 @@ impl UpdateEncoder {
         tiled
     }
 
-    /// Split a rect into tiles that fit within `max_request_size`.
+    /// Split a rect into tiles that fit within the tile budget.
     /// Splits by height first, then by width within each horizontal strip.
     fn split_diff(&self, rect: Rect) -> Vec<Rect> {
         let mut rects = Vec::new();
 
-        let max_height = (self.max_request_size / (rect.width * 4)).max(1);
+        let max_height = (self.tile_budget / (rect.width * 4)).max(1);
         let mut y = rect.y;
         let y_end = rect.y + rect.height;
         while y < y_end {
             let h = (y_end - y).min(max_height);
             // Width splitting is unlikely in practice (would require
             // max_request_size < ~256 KB), but ensures correctness.
-            let max_width = (self.max_request_size / (h * 4)).max(1);
+            let max_width = (self.tile_budget / (h * 4)).max(1);
             let mut x = rect.x;
             let x_end = rect.x + rect.width;
             while x < x_end {
@@ -943,3 +948,89 @@ fn set_surface(bitmap: &BitmapUpdate, codec_id: u8, data: &[u8]) -> ServerResult
         encode_vec(&cmd).map_err(ServerError::encode)?,
     ))
 }
+
+/// Size of everything `set_surface` puts on the wire besides the pixels: the
+/// `TS_SURFCMD_SET_SURF_BITS` header and its `TS_BITMAP_DATA_EX`. The client
+/// counts it against the MultifragmentUpdate reassembly buffer together with
+/// the pixel data.
+fn surface_bits_framing_size() -> usize {
+    let extended_bitmap_data = ExtendedBitmapDataPdu {
+        bpp: 32,
+        width: 0,
+        height: 0,
+        codec_id: CodecId::None.as_u8(),
+        header: None,
+        data: &[],
+    };
+    let pdu = SurfaceBitsPdu {
+        destination: ExclusiveRectangle {
+            left: 0,
+            top: 0,
+            right: 0,
+            bottom: 0,
+        },
+        extended_bitmap_data,
+    };
+    SurfaceCommand::SetSurfaceBits(pdu).size()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const EIGHT_MIB: usize = 8 * 1024 * 1024;
+
+    fn new_encoder(width: u16, height: u16) -> UpdateEncoder {
+        UpdateEncoder::new(
+            DesktopSize { width, height },
+            CmdFlags::SET_SURFACE_BITS,
+            UpdateEncoderCodecs::new(),
+            u32::try_from(EIGHT_MIB).unwrap(),
+            0,
+            LargePointerSupportFlags::empty(),
+        )
+        .unwrap()
+    }
+
+    fn full_rect(width: u16, height: u16) -> Rect {
+        Rect {
+            x: 0,
+            y: 0,
+            width: width.into(),
+            height: height.into(),
+        }
+    }
+
+    #[test]
+    fn surface_bits_framing_is_the_set_surface_header() {
+        // TS_SURFCMD_SET_SURF_BITS (cmdType 2 + destRect 8) + TS_BITMAP_DATA_EX (12).
+        assert_eq!(surface_bits_framing_size(), 22);
+    }
+
+    #[test]
+    fn strips_leave_room_for_the_framing_when_the_width_divides_the_buffer() {
+        let encoder = new_encoder(2048, 1080);
+        let tiles = encoder.split_diff(full_rect(2048, 1080));
+
+        // 1024 rows would fill the buffer exactly; one fewer leaves room for the header.
+        assert_eq!(tiles[0].height, 1023);
+        for tile in &tiles {
+            assert!(
+                tile.width * tile.height * 4 + surface_bits_framing_size() <= EIGHT_MIB,
+                "a {}x{} strip at y={} does not leave room for the framing",
+                tile.width,
+                tile.height,
+                tile.y
+            );
+        }
+        assert_eq!(tiles.iter().map(|tile| tile.height).sum::<usize>(), 1080);
+    }
+
+    #[test]
+    fn strips_with_slack_keep_their_height() {
+        // 3840 * 4 leaves 2 KiB unused in a 546-row strip, so nothing changes.
+        let encoder = new_encoder(3840, 2160);
+        let tiles = encoder.split_diff(full_rect(3840, 2160));
+        assert_eq!(tiles[0].height, 546);
+    }
+}

crates/ironrdp-testsuite-core/Cargo.toml

@@ -33,6 +33,7 @@ openh264 = { version = "0.9", optional = true, default-features = false, feature
 visibility = { version = "0.1", optional = true }
 
 [dev-dependencies]
+bytes = "1"
 anyhow = "1"
 async-trait = "0.1"
 expect-test.workspace = true
@@ -62,7 +63,7 @@ ironrdp-rdpeudp = { path = "../ironrdp-rdpeudp", features = ["std"] }
 ironrdp-rdpeusb.path = "../ironrdp-rdpeusb"
 ironrdp-usb.path = "../ironrdp-usb"
 ironrdp-rdpsnd = { path = "../ironrdp-rdpsnd", features = ["__test"] }
-ironrdp-server.path = "../ironrdp-server"
+ironrdp-server = { path = "../ironrdp-server", features = ["__bench"] }
 ironrdp-session = { path = "../ironrdp-session", features = ["qoi", "__test"] }
 ironrdp-cfg.path = "../ironrdp-cfg"
 ironrdp-propertyset.path = "../ironrdp-propertyset"

crates/ironrdp-testsuite-core/tests/server/encoder.rs

@@ -0,0 +1,103 @@
+use core::num::{NonZeroU16, NonZeroUsize};
+
+use bytes::Bytes;
+use ironrdp_pdu::rdp::capability_sets::{CmdFlags, LargePointerSupportFlags};
+use ironrdp_server::bench::encoder::{UpdateEncoder, UpdateEncoderCodecs};
+use ironrdp_server::{BitmapUpdate, DesktopSize, DisplayUpdate, PixelFormat};
+
+const EIGHT_MIB: u32 = 8 * 1024 * 1024;
+
+/// Framing `set_surface` adds around the pixels of one uncompressed surface-bits update:
+/// `TS_SURFCMD_SET_SURF_BITS` (cmdType 2 + destRect 8) + `TS_BITMAP_DATA_EX` (12).
+const SURFACE_BITS_FRAMING: usize = 22;
+
+fn uncompressed_encoder(width: u16, height: u16, max_request_size: u32) -> UpdateEncoder {
+    // Surface bits allowed and no codec: the `NoneHandler` path, one surface command per tile.
+    UpdateEncoder::new(
+        DesktopSize { width, height },
+        CmdFlags::SET_SURFACE_BITS,
+        UpdateEncoderCodecs::new(),
+        max_request_size,
+        0,
+        LargePointerSupportFlags::empty(),
+    )
+    .unwrap()
+}
+
+fn frame(width: u16, height: u16) -> DisplayUpdate {
+    let stride = usize::from(width) * 4;
+    DisplayUpdate::Bitmap(BitmapUpdate {
+        x: 0,
+        y: 0,
+        width: NonZeroU16::new(width).unwrap(),
+        height: NonZeroU16::new(height).unwrap(),
+        format: PixelFormat::XRgb32,
+        data: Bytes::from(vec![0; stride * usize::from(height)]),
+        stride: NonZeroUsize::new(stride).unwrap(),
+    })
+}
+
+/// Sizes of the updates the encoder emits for one frame: exactly what the client reassembles and
+/// checks against `MultifragMaxRequestSize`.
+async fn update_sizes(encoder: &mut UpdateEncoder, update: DisplayUpdate) -> Vec<usize> {
+    let mut updates = encoder.update(update);
+    let mut sizes = Vec::new();
+    while let Some(fragmenter) = updates.next().await {
+        sizes.push(fragmenter.unwrap().data.len());
+    }
+    sizes
+}
+
+fn strip_bytes(width: u16, rows: u16) -> usize {
+    usize::from(width) * usize::from(rows) * 4 + SURFACE_BITS_FRAMING
+}
+
+#[tokio::test]
+async fn a_width_that_divides_the_buffer_leaves_room_for_the_framing() {
+    // 2048 * 4 divides 8 MiB exactly: 1024 rows filled the buffer with pixels alone and the
+    // framing pushed the update past what the server advertised. One row fewer fits.
+    let mut encoder = uncompressed_encoder(2048, 1080, EIGHT_MIB);
+    let sizes = update_sizes(&mut encoder, frame(2048, 1080)).await;
+
+    assert_eq!(sizes, vec![strip_bytes(2048, 1023), strip_bytes(2048, 57)]);
+    assert!(sizes.iter().all(|&size| size <= usize::try_from(EIGHT_MIB).unwrap()));
+}
+
+#[tokio::test]
+async fn tall_frames_split_into_strips_that_fit() {
+    let mut encoder = uncompressed_encoder(4096, 2160, EIGHT_MIB);
+    let sizes = update_sizes(&mut encoder, frame(4096, 2160)).await;
+
+    let mut expected = vec![strip_bytes(4096, 511); 4];
+    expected.push(strip_bytes(4096, 2160 - 4 * 511));
+    assert_eq!(sizes, expected);
+}
+
+#[tokio::test]
+async fn a_frame_that_fits_with_its_framing_is_sent_whole() {
+    // The budget is inclusive: pixels plus framing equal to the buffer size is one update.
+    let (width, height) = (640u16, 480u16);
+    let cap = u32::try_from(strip_bytes(width, height)).unwrap();
+    let mut encoder = uncompressed_encoder(width, height, cap);
+    assert_eq!(
+        update_sizes(&mut encoder, frame(width, height)).await,
+        vec![strip_bytes(width, height)]
+    );
+
+    // One byte less and the frame has to be split.
+    let mut encoder = uncompressed_encoder(width, height, cap - 1);
+    let sizes = update_sizes(&mut encoder, frame(width, height)).await;
+    assert!(1 < sizes.len());
+    assert!(sizes.iter().all(|&size| size < usize::try_from(cap).unwrap()));
+}
+
+#[tokio::test]
+async fn a_budget_below_one_row_splits_by_width_too() {
+    // 2048 * 4 bytes per row do not fit a 1000-byte budget, so the strips are one row high and
+    // are split along the width as well; every piece still fits with its framing.
+    let mut encoder = uncompressed_encoder(2048, 8, 1000);
+    let sizes = update_sizes(&mut encoder, frame(2048, 8)).await;
+
+    assert!(1 < sizes.len());
+    assert!(sizes.iter().all(|&size| size <= 1000));
+}

crates/ironrdp-testsuite-core/tests/server/mod.rs

@@ -1,6 +1,7 @@
 mod acceptor;
 mod autodetect;
 mod credential_validator;
+mod encoder;
 mod fast_path;
 mod finalize_timeout;
 mod rdpdr;