diff --git a/crates/ironrdp-graphics/src/zgfx/circular_buffer.rs b/crates/ironrdp-graphics/src/zgfx/circular_buffer.rs index 56b503a756..68ca8ce0ce 100644 --- a/crates/ironrdp-graphics/src/zgfx/circular_buffer.rs +++ b/crates/ironrdp-graphics/src/zgfx/circular_buffer.rs @@ -59,6 +59,15 @@ impl FixedCircularBuffer { Ok(()) } + + /// The byte `distance` positions before the write position. + /// + /// INVARIANT: `0 < distance <= self.buffer.len()`. The callers are the compressor's history, + /// which only asks for positions it computed itself, so this is checked in debug builds only. + pub(crate) fn byte_back(&self, distance: usize) -> u8 { + debug_assert!(0 < distance && distance <= self.buffer.len()); + self.buffer[(self.position + self.buffer.len() - distance) % self.buffer.len()] + } } impl io::Write for FixedCircularBuffer { diff --git a/crates/ironrdp-graphics/src/zgfx/compressor.rs b/crates/ironrdp-graphics/src/zgfx/compressor.rs index f9c11585d0..309d5eb930 100644 --- a/crates/ironrdp-graphics/src/zgfx/compressor.rs +++ b/crates/ironrdp-graphics/src/zgfx/compressor.rs @@ -7,11 +7,11 @@ //! //! [\[MS-RDPEGFX\] 2.2.1.1.1]: https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-rdpegfx/ -use std::collections::HashMap; +use std::{collections::HashMap, io::Write as _}; use bitvec::prelude::*; -use super::{HISTORY_SIZE, TOKEN_TABLE, ZgfxError}; +use super::{FixedCircularBuffer, HISTORY_SIZE, TOKEN_TABLE, ZgfxError}; const MIN_MATCH_LENGTH: usize = 3; const MAX_MATCH_LENGTH: usize = 65535; /// Maximum back-reference distance (last token in MS-RDPEGFX table) @@ -32,15 +32,17 @@ const COMPACT_TARGET_ENTRIES: usize = MAX_HASH_TABLE_ENTRIES / 2; /// ZGFX compressor maintaining a 2.5 MB history buffer and prefix hash table. pub struct Compressor { - history: Vec, - /// 3-byte prefix → history positions, for O(1) match candidate lookup - match_table: HashMap<[u8; 3], Vec>, + history: History, + /// 3-byte prefix → absolute stream positions, for O(1) match candidate lookup. + /// Positions are never rebased; candidates that have fallen out of reach are + /// skipped at lookup time. + match_table: HashMap<[u8; 3], Vec>, } impl Compressor { pub fn new() -> Self { Self { - history: Vec::with_capacity(HISTORY_SIZE), + history: History::new(), match_table: HashMap::new(), } } @@ -71,47 +73,17 @@ impl Compressor { Ok(bit_writer.finish()) } - /// Extend the sliding window, evicting oldest bytes when full. + /// Extend the sliding window; the ring overwrites the oldest bytes once full. fn add_to_history(&mut self, bytes: &[u8]) { - if self.history.len() + bytes.len() > HISTORY_SIZE { - let overflow = (self.history.len() + bytes.len()) - HISTORY_SIZE; - - self.history.drain(..overflow); - - // Shift all stored positions to account for evicted bytes - for positions in self.match_table.values_mut() { - positions.retain_mut(|pos| { - if *pos >= overflow { - *pos -= overflow; - true - } else { - false - } - }); - } - self.match_table.retain(|_, positions| !positions.is_empty()); - } - - let base_pos = self.history.len(); - self.history.extend_from_slice(bytes); + let base_pos = self.history.total(); + self.history.push(bytes); // For large chunks (match replays), sample every 4th position to keep // the hash table manageable without sacrificing much compression ratio let step_size = if bytes.len() > 256 { 4 } else { 1 }; for i in (0..bytes.len().saturating_sub(MIN_MATCH_LENGTH - 1)).step_by(step_size) { - let pos = base_pos + i; - let prefix = [self.history[pos], self.history[pos + 1], self.history[pos + 2]]; - - let entry = self.match_table.entry(prefix).or_default(); - - if entry.len() < MAX_POSITIONS_PER_PREFIX { - entry.push(pos); - } else { - // Evict oldest position to keep recent data preferred - entry.remove(0); - entry.push(pos); - } + self.record_position([bytes[i], bytes[i + 1], bytes[i + 2]], base_pos + as_u64(i)); } if self.match_table.len() > MAX_HASH_TABLE_ENTRIES { @@ -120,23 +92,34 @@ impl Compressor { // Index positions spanning the old/new data boundary so matches // that straddle the append point can be found - for offset in [2, 1] { - if base_pos >= offset && bytes.len() + offset > 2 { - let pos = base_pos - offset; - if pos + MIN_MATCH_LENGTH <= self.history.len() { - let prefix = [self.history[pos], self.history[pos + 1], self.history[pos + 2]]; - let entry = self.match_table.entry(prefix).or_default(); - if entry.last() != Some(&pos) { - if entry.len() >= MAX_POSITIONS_PER_PREFIX { - entry.remove(0); - } - entry.push(pos); - } - } - } + for offset in [2u64, 1] { + let Some(pos) = base_pos.checked_sub(offset) else { + continue; + }; + let Some(prefix) = self.history.prefix_at(pos) else { + continue; + }; + self.record_position(prefix, pos); } } + /// 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. + 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)); + if MAX_POSITIONS_PER_PREFIX <= positions.len() { + // Evict oldest position to keep recent data preferred + positions.remove(0); + } + positions.push(pos); + } + /// Halve stored positions per prefix, then evict whole prefixes down to /// `COMPACT_TARGET_ENTRIES` when the table is over the cap. /// @@ -164,7 +147,7 @@ impl Compressor { // MAX_MATCH_DISTANCE regardless. Each history position belongs to a // single prefix, so these newest positions are distinct and the cutoff // retains exactly COMPACT_TARGET_ENTRIES entries. - let mut newest: Vec = self + let mut newest: Vec = self .match_table .values() .map(|positions| positions.last().copied().unwrap_or(0)) @@ -178,7 +161,7 @@ impl Compressor { /// Search hash table for the longest match at `input[pos..]`. fn find_best_match(&self, input: &[u8], pos: usize) -> Option { let remaining = input.len() - pos; - if remaining < MIN_MATCH_LENGTH || self.history.is_empty() { + if remaining < MIN_MATCH_LENGTH { return None; } @@ -187,22 +170,26 @@ impl Compressor { let max_match_len = remaining.min(MAX_MATCH_LENGTH); let mut best_match: Option = None; - let search_limit = self.history.len().min(MAX_MATCH_DISTANCE); + let search_limit = as_u64(self.history.len().min(MAX_MATCH_DISTANCE)); + let total = self.history.total(); // Most recent candidates first — better locality, often longer matches for &hist_pos in candidates.iter().rev().take(MAX_CANDIDATES) { - let distance = self.history.len() - hist_pos; - - if distance > search_limit { + // Positions are absolute, so anything older than the reachable + // window is simply skipped rather than rebased on every append. + let Some(distance) = total.checked_sub(hist_pos).filter(|d| *d <= search_limit) else { continue; - } + }; + let distance = usize::try_from(distance).expect("distance is at most MAX_MATCH_DISTANCE"); - // Prefix already matched via hash table; extend from byte 3 onward + // Prefix already matched via hash table; extend from byte 3 onward. + // INVARIANT: match_len < distance, so the match reads only bytes + // already in the history. let mut match_len = MIN_MATCH_LENGTH; while match_len < max_match_len - && hist_pos + match_len < self.history.len() - && self.history[hist_pos + match_len] == input[pos + match_len] + && match_len < distance + && self.history.byte_back(distance - match_len) == input[pos + match_len] { match_len += 1; } @@ -339,6 +326,71 @@ impl Default for Compressor { } } +fn as_u64(value: usize) -> u64 { + u64::try_from(value).expect("usize fits in u64") +} + +/// The last `HISTORY_SIZE` bytes of the stream, kept in a fixed ring so that +/// appending never moves existing bytes. +struct History { + ring: FixedCircularBuffer, + /// Bytes appended since the compressor was created; the absolute stream + /// position of the next byte. + total: u64, +} + +impl History { + fn new() -> Self { + Self { + ring: FixedCircularBuffer::new(HISTORY_SIZE), + total: 0, + } + } + + /// Number of valid bytes in the ring, at most `HISTORY_SIZE`. + fn len(&self) -> usize { + usize::try_from(self.total.min(as_u64(HISTORY_SIZE))).expect("at most HISTORY_SIZE") + } + + fn total(&self) -> u64 { + self.total + } + + fn push(&mut self, bytes: &[u8]) { + // The ring keeps only the last HISTORY_SIZE bytes of an input larger than itself. + self.ring.write_all(bytes).expect("writing to the ring cannot fail"); + self.total += as_u64(bytes.len()); + } + + /// The byte `distance` positions before the end of the stream. + /// + /// INVARIANT: `0 < distance <= self.len()`. Both callers establish it + /// themselves, which is why the check is a `debug_assert!` and not a branch + /// on every extended byte: `find_best_match` only keeps candidates within + /// the window and extends while `match_len < distance`, and `prefix_at` + /// checks the window before it reads. Tests run in debug builds, where a + /// violation panics, and the round-trip tests would catch a wrong byte. + fn byte_back(&self, distance: usize) -> u8 { + debug_assert!(0 < distance && distance <= self.len()); + self.ring.byte_back(distance) + } + + /// The three bytes starting at absolute position `pos`, or `None` when they + /// are not all in the ring: fewer than three bytes follow `pos`, or `pos` + /// has fallen out of the window. + fn prefix_at(&self, pos: u64) -> Option<[u8; 3]> { + let distance = usize::try_from(self.total.checked_sub(pos)?).ok()?; + if !(MIN_MATCH_LENGTH..=self.len()).contains(&distance) { + return None; + } + Some([ + self.byte_back(distance), + self.byte_back(distance - 1), + self.byte_back(distance - 2), + ]) + } +} + #[derive(Debug, Clone, Copy)] struct Match { distance: usize, @@ -523,6 +575,114 @@ mod tests { ); } + #[test] + fn compress_round_trips_across_history_wrap() { + use super::super::Decompressor; + + // More than twice HISTORY_SIZE through one compressor, in segments that + // repeat earlier content, so matches reach back across the ring wrap + // point and the decompressor's history has to agree byte for byte. + let mut state: u32 = 0x0BAD_F00D; + let mut next_byte = move || { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + u8::try_from(state >> 24).unwrap() + }; + let blocks: Vec> = + core::iter::repeat_with(|| core::iter::repeat_with(&mut next_byte).take(4096).collect()) + .take(64) + .collect(); + + let mut compressor = Compressor::new(); + let mut decompressor = Decompressor::new(); + let mut total = 0; + let mut segment_index = 0usize; + while total < 2 * HISTORY_SIZE + 100_000 { + let mut segment = Vec::with_capacity(60_000); + for k in 0..14 { + let block = &blocks[(segment_index * 7 + k * 5) % blocks.len()]; + segment.extend_from_slice(block); + segment.push(next_byte()); + } + segment_index += 1; + + let compressed = compressor.compress(&segment).unwrap(); + let mut output = Vec::new(); + decompressor.decompress_segment(&compressed, &mut output).unwrap(); + assert_eq!(output, segment, "segment {segment_index} after {total} bytes"); + total += segment.len(); + } + assert!(compressor.history.total() > u64::try_from(2 * HISTORY_SIZE).unwrap()); + } + + #[test] + fn a_prefix_keeps_only_its_newest_positions() { + let mut compressor = Compressor::new(); + let recorded = MAX_POSITIONS_PER_PREFIX + 8; + for pos in 0..recorded { + compressor.record_position(*b"abc", as_u64(pos)); + } + + let expected: Vec = (8..recorded).map(as_u64).collect(); + assert_eq!(compressor.match_table[b"abc"], expected); + } + + #[test] + fn positions_are_recorded_inside_a_chunk_and_across_the_append_point() { + let mut compressor = Compressor::new(); + compressor.add_to_history(b"abcd"); + compressor.add_to_history(b"ef"); + + // The first chunk holds "abc" at 0 and "bcd" at 1. The second chunk completes "cde" at 2 and + // "def" at 3, which start before the append point. + let table = &compressor.match_table; + assert_eq!(table[b"abc"], [0]); + assert_eq!(table[b"bcd"], [1]); + assert_eq!(table[b"cde"], [2]); + assert_eq!(table[b"def"], [3]); + assert_eq!(table.len(), 4); + } + + #[test] + fn history_ring_keeps_the_last_bytes_in_order() { + let mut history = History::new(); + let chunk: Vec = (0..=255u8).cycle().take(HISTORY_SIZE - 10).collect(); + history.push(&chunk); + history.push(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]); + + assert_eq!(history.len(), HISTORY_SIZE); + assert_eq!(history.byte_back(1), 15); + assert_eq!(history.byte_back(15), 1); + assert_eq!(history.byte_back(16), *chunk.last().unwrap()); + assert_eq!(history.prefix_at(history.total() - 3), Some([13, 14, 15])); + assert_eq!( + history.prefix_at(history.total() - 2), + None, + "fewer than three bytes follow the position" + ); + let oldest = history.total() - as_u64(HISTORY_SIZE); + assert_eq!( + history.prefix_at(oldest), + Some([ + history.byte_back(HISTORY_SIZE), + history.byte_back(HISTORY_SIZE - 1), + history.byte_back(HISTORY_SIZE - 2), + ]), + "the oldest position still in the ring is reachable" + ); + assert_eq!( + history.prefix_at(oldest - 1), + None, + "one position older has left the ring" + ); + + // A push larger than the ring keeps only its tail. + let big: Vec = (0..HISTORY_SIZE + 7).map(|i| u8::try_from(i % 251).unwrap()).collect(); + history.push(&big); + assert_eq!(history.len(), HISTORY_SIZE); + assert_eq!(history.byte_back(1), *big.last().unwrap()); + assert_eq!(history.byte_back(HISTORY_SIZE), big[7]); + } + #[test] fn bit_writer_basic() { let mut writer = BitWriter::new();