feat(egfx): skip ZGFX compression for H.264 surface commands

crates/ironrdp-egfx/src/server.rs

@@ -60,7 +60,9 @@ use std::time::Instant;
 
 use ironrdp_core::{Encode, EncodeResult, WriteCursor, decode, impl_as_any};
 use ironrdp_dvc::{DvcEncode, DvcMessage, DvcProcessor, DvcServerProcessor};
-use ironrdp_graphics::zgfx::{CompressionMode, Compressor, compress_and_wrap_egfx, wrap_uncompressed};
+use ironrdp_graphics::zgfx::{
+    CompressionMode, Compressor, compress_and_wrap_egfx, wrap_uncompressed, wrap_uncompressed_recorded,
+};
 use ironrdp_pdu::gcc::Monitor;
 use ironrdp_pdu::geometry::ExclusiveRectangle;
 use ironrdp_pdu::{PduResult, decode_err};
@@ -1164,7 +1166,12 @@ impl GraphicsPipelineServer {
     /// When `compression_mode` is `Auto` or `Always`, `drain_output()` will
     /// compress PDUs before ZGFX wrapping, reducing bandwidth at the cost
     /// of CPU. The compressor maintains a sliding history window across
-    /// frames for back-reference efficiency.
+    /// frames for back-reference efficiency. `WireToSurface1` PDUs carrying
+    /// AVC420, AVC444 or AVC444v2 are sent uncompressed in both modes, since
+    /// H.264 is already entropy coded; their bytes still enter the history.
+    /// The trade-off is that a byte-identical repeat of an earlier H.264 PDU,
+    /// such as a periodic keyframe of an unchanged screen, is sent in full where
+    /// ZGFX would have matched it against the history.
     pub fn with_compression(handler: Box<dyn GraphicsPipelineHandler>, compression_mode: CompressionMode) -> Self {
         let mut server = Self::new(handler);
         server.compression_mode = compression_mode;
@@ -2219,10 +2226,28 @@ impl GraphicsPipelineServer {
                 pdu.encode(&mut cursor).expect("GfxPdu encoding should not fail");
                 total_uncompressed += pdu_size;
 
+                // H.264 bitstreams are already entropy coded, so searching them
+                // for matches costs CPU on the caller's send path and almost never
+                // shrinks them. Send them uncompressed, but still record their
+                // bytes in the history (MS-RDPEGFX 3.1.9.1.2).
+                let is_h264 = matches!(
+                    &pdu,
+                    GfxPdu::WireToSurface1(wire)
+                        if matches!(wire.codec_id, Codec1Type::Avc420 | Codec1Type::Avc444 | Codec1Type::Avc444v2)
+                );
+
                 let wrapped = match mode {
                     CompressionMode::Never => wrap_uncompressed(&pdu_bytes),
-                    _ => compress_and_wrap_egfx(&pdu_bytes, &mut self.zgfx_compressor, mode)
-                        .unwrap_or_else(|_| wrap_uncompressed(&pdu_bytes)),
+                    CompressionMode::Auto | CompressionMode::Always if is_h264 => {
+                        wrap_uncompressed_recorded(&pdu_bytes, &mut self.zgfx_compressor)
+                    }
+                    CompressionMode::Auto | CompressionMode::Always => {
+                        compress_and_wrap_egfx(&pdu_bytes, &mut self.zgfx_compressor, mode)
+                            // None of the compressor's encoders returns an error today. If one ever did, the
+                            // tokens before it would already be in the history, so recording this PDU again
+                            // here would not bring the peer back in step. Both histories would need a reset.
+                            .unwrap_or_else(|_| wrap_uncompressed(&pdu_bytes))
+                    }
                 };
                 total_compressed += wrapped.len();
                 trace!(pdu_name, pdu_size, wrapped = wrapped.len(), mode = ?mode, "ZGFX output");

crates/ironrdp-graphics/src/zgfx/api.rs

@@ -5,9 +5,19 @@ use super::compressor::Compressor;
 use super::wrapper::{ZGFX_SEGMENTED_MAXSIZE, wrap_compressed, wrap_uncompressed};
 
 /// Controls whether ZGFX compression is applied.
+///
+/// The modes apply to PDUs passed to [`compress_and_wrap_egfx()`]. A caller that
+/// knows a PDU is already entropy coded, as the EGFX server does for H.264
+/// surface commands, sends it with [`wrap_uncompressed_recorded()`] instead,
+/// whatever the mode, so `Always` means every PDU the caller leaves to this
+/// module.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum CompressionMode {
     /// Send uncompressed (no CPU overhead).
+    ///
+    /// Nothing is recorded in the compressor history in this mode, so a [`Compressor`] used with
+    /// `Never` must not later be used with `Auto`, `Always` or [`wrap_uncompressed_recorded()`]:
+    /// the receiver records every segment, and the histories would no longer agree.
     Never,
     /// Compress and use the smaller result (bandwidth vs CPU trade-off).
     Auto,
@@ -60,6 +70,25 @@ pub fn compress_and_wrap_egfx(
     }
 }
 
+/// Wrap `data` in uncompressed ZGFX segments and record its bytes in the
+/// `compressor` history, for a PDU the caller sends uncompressed on purpose.
+///
+/// MS-RDPEGFX 3.1.9.1.2 requires every output byte, including the bytes of
+/// segments sent uncompressed, to be recorded in the history, because the
+/// receiver does so and later back-references point at those bytes. Recording
+/// and wrapping happen together here so that neither can be done without the
+/// other. Use it, in place of [`compress_and_wrap_egfx()`], for data that is
+/// already entropy coded such as H.264, where searching for matches costs CPU
+/// and almost never shrinks the PDU.
+///
+/// The `compressor` history must already hold every byte the receiver has seen
+/// on this stream, so don't use this after PDUs sent with
+/// [`CompressionMode::Never`] or [`wrap_uncompressed()`], which record nothing.
+pub fn wrap_uncompressed_recorded(data: &[u8], compressor: &mut Compressor) -> Vec<u8> {
+    compressor.record_uncompressed(data);
+    wrap_uncompressed(data)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -114,4 +143,43 @@ mod tests {
             assert_eq!(&output, data, "Round-trip failed for mode {mode:?}");
         }
     }
+
+    /// The receiver records an uncompressed segment in its history, so a PDU sent
+    /// through `wrap_uncompressed_recorded()` must be recorded here too, or the
+    /// back-reference in the PDU after it points at the wrong bytes.
+    #[test]
+    fn wrap_uncompressed_recorded_keeps_the_receivers_history_in_step() {
+        use super::super::Decompressor;
+
+        let text = b"a tile of text that repeats later in the session, ".repeat(40);
+        let mut state: u32 = 0x1234_5678;
+        let entropy_coded: Vec<u8> = core::iter::repeat_with(|| {
+            state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
+            u8::try_from(state >> 24).unwrap()
+        })
+        .take(8_000)
+        .collect();
+
+        let mut compressor = Compressor::new();
+        let mut decompressor = Decompressor::new();
+        let mut decoded = Vec::new();
+
+        let first = compress_and_wrap_egfx(&text, &mut compressor, CompressionMode::Always).unwrap();
+        decompressor.decompress(&first, &mut decoded).unwrap();
+        assert_eq!(decoded, text);
+
+        let second = wrap_uncompressed_recorded(&entropy_coded, &mut compressor);
+        assert_eq!(second[1], 0x04, "sent uncompressed");
+        decoded.clear();
+        decompressor.decompress(&second, &mut decoded).unwrap();
+        assert_eq!(decoded, entropy_coded);
+
+        // Reaches back over the recorded bytes to the first PDU, which only decodes
+        // correctly when both sides counted the uncompressed segment.
+        let third = compress_and_wrap_egfx(&text, &mut compressor, CompressionMode::Always).unwrap();
+        assert_eq!(third[1], 0x24, "sent compressed");
+        decoded.clear();
+        decompressor.decompress(&third, &mut decoded).unwrap();
+        assert_eq!(decoded, text);
+    }
 }

crates/ironrdp-graphics/src/zgfx/compressor.rs

@@ -73,6 +73,21 @@ impl Compressor {
         Ok(bit_writer.finish())
     }
 
+    /// Record `bytes` in the history without searching them for matches, for
+    /// data the caller sends in an uncompressed segment instead of passing it
+    /// to [`Compressor::compress`].
+    ///
+    /// MS-RDPEGFX 3.1.9.1.2 requires every output byte, including the bytes of
+    /// segments sent uncompressed, to be recorded in the history; the receiver
+    /// does so, and skipping them here would make later back-references point
+    /// at the wrong bytes. The bytes are not bulk indexed, so the search that
+    /// compression would spend on already entropy coded data is skipped. The
+    /// last two recorded bytes can still start a match, because the next
+    /// append indexes the positions that straddle the boundary.
+    pub(crate) fn record_uncompressed(&mut self, bytes: &[u8]) {
+        self.history.push(bytes);
+    }
+
     /// Extend the sliding window; the ring overwrites the oldest bytes once full.
     fn add_to_history(&mut self, bytes: &[u8]) {
         let base_pos = self.history.total();
@@ -106,10 +121,11 @@ impl Compressor {
     /// Remember that `prefix` occurs at the absolute stream position `pos`, dropping the
     /// prefix's oldest position when it already has `MAX_POSITIONS_PER_PREFIX`.
     ///
-    /// INVARIANT: each position is recorded once, when its third byte arrives. The main
-    /// loop covers the positions whose three bytes are all in the new chunk, and the
-    /// boundary loop the two positions that straddle the append point, so no position
-    /// is in both.
+    /// INVARIANT: no position is recorded twice. A position is recorded when its third byte
+    /// arrives: the main loop covers the positions whose three bytes are all in the new chunk,
+    /// and the boundary loop the two positions that straddle the append point, so no position
+    /// is in both. Bytes added with `record_uncompressed` are not indexed, apart from the
+    /// positions that straddle their end, which the next append indexes.
     fn record_position(&mut self, prefix: [u8; 3], pos: u64) {
         let positions = self.match_table.entry(prefix).or_default();
         debug_assert!(positions.last() != Some(&pos));

crates/ironrdp-graphics/src/zgfx/mod.rs

@@ -9,7 +9,7 @@ mod wrapper;
 use std::io::{self, Write as _};
 use std::sync::LazyLock;
 
-pub use api::{CompressionMode, compress_and_wrap_egfx};
+pub use api::{CompressionMode, compress_and_wrap_egfx, wrap_uncompressed_recorded};
 use bitvec::bits;
 use bitvec::field::BitField as _;
 use bitvec::order::Msb0;

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

@@ -8,7 +8,7 @@ use ironrdp_egfx::pdu::{
 use ironrdp_egfx::server::{
     Avc444SubStream, Avc444Tile, GraphicsPipelineHandler, GraphicsPipelineServer, MixedTilePayload, QoeMetrics, Surface,
 };
-use ironrdp_graphics::zgfx::Decompressor;
+use ironrdp_graphics::zgfx::{CompressionMode, Decompressor};
 use ironrdp_pdu::geometry::ExclusiveRectangle;
 use rstest::rstest;
 
@@ -354,6 +354,101 @@ fn avc444v2_sender_rejects_invalid_stream_shapes_without_queueing() {
     }
 }
 
+// ============================================================================
+// ZGFX Output Tests
+// ============================================================================
+
+/// Whether the first ZGFX segment of a wrapped DVC message is compressed
+/// (MS-RDPEGFX 2.2.5.1 to 2.2.5.3; PACKET_COMPRESSED is 0x20).
+fn first_segment_is_compressed(wrapped: &[u8]) -> bool {
+    match wrapped[0] {
+        0xE0 => wrapped[1] & 0x20 != 0,
+        0xE1 => wrapped[1 + 2 + 4 + 4] & 0x20 != 0,
+        descriptor => panic!("unexpected ZGFX descriptor {descriptor:#x}"),
+    }
+}
+
+/// H.264 surface commands go out uncompressed in both compressing modes, for each of the three
+/// H.264 codecs, and the receiver's history stays in step with the sender's. 70,000 bytes is
+/// above the 65,535 of a single ZGFX segment, so that case is multipart and records one large push.
+#[rstest]
+fn zgfx_sends_h264_uncompressed_and_keeps_history_in_step(
+    #[values(CompressionMode::Auto, CompressionMode::Always)] mode: CompressionMode,
+    #[values(Codec1Type::Avc420, Codec1Type::Avc444, Codec1Type::Avc444v2)] h264_codec: Codec1Type,
+    #[values(8_000, 70_000)] h264_len: usize,
+) {
+    let handler = Box::new(TestHandler::new());
+    let mut server = GraphicsPipelineServer::with_compression(handler, mode);
+    let client_caps_pdu = GfxPdu::CapabilitiesAdvertise(CapabilitiesAdvertisePdu::from_typed(&[CapabilitySet::V10 {
+        flags: CapabilitiesV10Flags::SMALL_CACHE,
+    }]));
+    server
+        .process(0, &encode_pdu(&client_caps_pdu))
+        .expect("process capabilities");
+    let surface_id = server.create_surface(64, 64).expect("create surface");
+
+    let mut decompressor = Decompressor::new();
+    let mut decode_all = |server: &mut GraphicsPipelineServer| -> Vec<(bool, GfxPdu)> {
+        server
+            .drain_output()
+            .iter()
+            .map(|message| {
+                let wrapped = encode_vec(message.as_ref()).expect("encode DVC message");
+                let mut decoded = Vec::new();
+                decompressor.decompress(&wrapped, &mut decoded).expect("decompress PDU");
+                let pdu = GfxPdu::decode(&mut ReadCursor::new(&decoded)).expect("decode PDU");
+                (first_segment_is_compressed(&wrapped), pdu)
+            })
+            .collect()
+    };
+    decode_all(&mut server);
+
+    let clearcodec_data = b"lossless tile bytes that repeat ".repeat(64);
+    let h264_data: Vec<u8> = (0..=u8::MAX).cycle().skip(3).step_by(7).take(h264_len).collect();
+    let destination = ExclusiveRectangle {
+        left: 0,
+        top: 0,
+        right: 16,
+        bottom: 16,
+    };
+
+    // ClearCodec, then H.264, then ClearCodec again: the last one only
+    // decodes if both sides recorded the H.264 bytes in their histories.
+    server
+        .send_clearcodec_frame(surface_id, destination.clone(), clearcodec_data.clone(), 42)
+        .expect("queue ClearCodec frame");
+    let regions = [Avc420Region::full_frame(64, 64, 22)];
+    match h264_codec {
+        Codec1Type::Avc420 => server.send_avc420_frame(surface_id, &h264_data, &regions, 42),
+        Codec1Type::Avc444 => server.send_avc444_frame(surface_id, &h264_data, &regions, None, None, 42),
+        Codec1Type::Avc444v2 => {
+            server.send_avc444v2_frame(surface_id, Encoding::LUMA, &h264_data, &regions, None, None, 42)
+        }
+        other => panic!("not an H.264 codec: {other:?}"),
+    }
+    .expect("queue H.264 frame");
+    server
+        .send_clearcodec_frame(surface_id, destination, clearcodec_data.clone(), 42)
+        .expect("queue ClearCodec frame");
+
+    let mut codecs = Vec::new();
+    for (compressed, pdu) in decode_all(&mut server) {
+        let GfxPdu::WireToSurface1(wire) = pdu else {
+            continue;
+        };
+        match wire.codec_id {
+            Codec1Type::ClearCodec => {
+                assert!(compressed, "ClearCodec PDU should be compressed");
+                assert_eq!(wire.bitmap_data, clearcodec_data);
+            }
+            codec if codec == h264_codec => assert!(!compressed, "H.264 PDU should be sent uncompressed"),
+            other => panic!("unexpected codec {other:?}"),
+        }
+        codecs.push(wire.codec_id);
+    }
+    assert_eq!(codecs, [Codec1Type::ClearCodec, h264_codec, Codec1Type::ClearCodec]);
+}
+
 // ============================================================================
 // Mixed-Codec Frame Tests
 // ============================================================================