ironrdp_graphics/zgfx/
control_messages.rs

1use bit_field::BitField as _;
2use bitflags::bitflags;
3use byteorder::{LittleEndian, ReadBytesExt as _};
4use num_derive::FromPrimitive;
5use num_traits::FromPrimitive as _;
6
7use super::ZgfxError;
8
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub(crate) enum SegmentedDataPdu<'a> {
11    Single(BulkEncodedData<'a>),
12    Multipart {
13        uncompressed_size: usize,
14        segments: Vec<BulkEncodedData<'a>>,
15    },
16}
17
18impl<'a> SegmentedDataPdu<'a> {
19    pub(crate) fn from_buffer(mut buffer: &'a [u8]) -> Result<Self, ZgfxError> {
20        let descriptor =
21            SegmentedDescriptor::from_u8(buffer.read_u8()?).ok_or(ZgfxError::InvalidSegmentedDescriptor)?;
22
23        match descriptor {
24            SegmentedDescriptor::Single => Ok(SegmentedDataPdu::Single(BulkEncodedData::from_buffer(buffer)?)),
25            SegmentedDescriptor::Multipart => {
26                let segment_count = buffer.read_u16::<LittleEndian>()? as usize;
27                let uncompressed_size = buffer.read_u32::<LittleEndian>()? as usize;
28
29                let mut segments = Vec::with_capacity(segment_count);
30                for _ in 0..segment_count {
31                    let size = buffer.read_u32::<LittleEndian>()? as usize;
32                    let (segment_data, new_buffer) = buffer.split_at(size);
33                    buffer = new_buffer;
34
35                    segments.push(BulkEncodedData::from_buffer(segment_data)?);
36                }
37
38                Ok(SegmentedDataPdu::Multipart {
39                    uncompressed_size,
40                    segments,
41                })
42            }
43        }
44    }
45}
46
47#[derive(Debug, Clone, PartialEq, Eq)]
48pub(crate) struct BulkEncodedData<'a> {
49    pub(crate) compression_flags: CompressionFlags,
50    pub(crate) data: &'a [u8],
51}
52
53impl<'a> BulkEncodedData<'a> {
54    pub(crate) fn from_buffer(mut buffer: &'a [u8]) -> Result<Self, ZgfxError> {
55        let compression_type_and_flags = buffer.read_u8()?;
56        let _compression_type = CompressionType::from_u8(compression_type_and_flags.get_bits(..4))
57            .ok_or(ZgfxError::InvalidCompressionType)?;
58        let compression_flags = CompressionFlags::from_bits_truncate(compression_type_and_flags.get_bits(4..));
59
60        Ok(Self {
61            compression_flags,
62            data: buffer,
63        })
64    }
65}
66
67#[derive(Debug, Copy, Clone, PartialEq, FromPrimitive)]
68enum SegmentedDescriptor {
69    Single = 0xe0,
70    Multipart = 0xe1,
71}
72
73#[derive(Debug, Copy, Clone, PartialEq, FromPrimitive)]
74enum CompressionType {
75    Rdp8 = 0x4,
76}
77
78bitflags! {
79    #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
80    pub struct CompressionFlags: u8 {
81        const COMPRESSED = 0x2;
82    }
83}
84
85#[cfg(test)]
86mod test {
87    use lazy_static::lazy_static;
88
89    use super::*;
90
91    const SINGLE_SEGMENTED_DATA_PDU_BUFFER: [u8; 17] = [
92        0xe0, // descriptor
93        0x24, // flags and compression type
94        0x09, 0xe3, 0x18, 0x0a, 0x44, 0x8d, 0x37, 0xf4, 0xc6, 0xe8, 0xa0, 0x20, 0xc6, 0x30, 0x01, // data
95    ];
96
97    const MULTIPART_SEGMENTED_DATA_PDU_BUFFER: [u8; 66] = [
98        0xE1, // descriptor
99        0x03, 0x00, // segment count
100        0x2B, 0x00, 0x00, 0x00, // uncompressed size
101        0x11, 0x00, 0x00, 0x00, // size of the first segment
102        0x04, // the first segment: flags and compression type
103        0x54, 0x68, 0x65, 0x20, 0x71, 0x75, 0x69, 0x63, 0x6B, 0x20, 0x62, 0x72, 0x6F, 0x77, 0x6E,
104        0x20, // the first segment: data
105        0x0E, 0x00, 0x00, 0x00, // size of the second segment
106        0x04, // the second segment: flags and compression type
107        0x66, 0x6F, 0x78, 0x20, 0x6A, 0x75, 0x6D, 0x70, 0x73, 0x20, 0x6F, 0x76, 0x65, // the second segment: data
108        0x10, 0x00, 0x00, 0x00, // size of the third segment
109        0x24, // the third segment: flags and compression type
110        0x39, 0x08, 0x0E, 0x91, 0xF8, 0xD8, 0x61, 0x3D, 0x1E, 0x44, 0x06, 0x43, 0x79, 0x9C,
111        0x02, // the third segment: data
112    ];
113
114    lazy_static! {
115        static ref SINGLE_SEGMENTED_DATA_PDU: SegmentedDataPdu<'static> = SegmentedDataPdu::Single(BulkEncodedData {
116            compression_flags: CompressionFlags::COMPRESSED,
117            data: &SINGLE_SEGMENTED_DATA_PDU_BUFFER[2..],
118        });
119        static ref MULTIPART_SEGMENTED_DATA_PDU: SegmentedDataPdu<'static> = SegmentedDataPdu::Multipart {
120            uncompressed_size: 0x2B,
121            segments: vec![
122                BulkEncodedData {
123                    compression_flags: CompressionFlags::empty(),
124                    data: &MULTIPART_SEGMENTED_DATA_PDU_BUFFER[12..12 + 16]
125                },
126                BulkEncodedData {
127                    compression_flags: CompressionFlags::empty(),
128                    data: &MULTIPART_SEGMENTED_DATA_PDU_BUFFER[33..33 + 13]
129                },
130                BulkEncodedData {
131                    compression_flags: CompressionFlags::COMPRESSED,
132                    data: &MULTIPART_SEGMENTED_DATA_PDU_BUFFER[51..]
133                },
134            ],
135        };
136    }
137
138    #[test]
139    fn from_buffer_correctly_parses_zgfx_single_segmented_data_pdu() {
140        let buffer = SINGLE_SEGMENTED_DATA_PDU_BUFFER.as_ref();
141
142        assert_eq!(
143            *SINGLE_SEGMENTED_DATA_PDU,
144            SegmentedDataPdu::from_buffer(buffer).unwrap()
145        );
146    }
147
148    #[test]
149    fn from_buffer_correctly_parses_zgfx_multipart_segmented_data_pdu() {
150        let buffer = MULTIPART_SEGMENTED_DATA_PDU_BUFFER.as_ref();
151
152        assert_eq!(
153            *MULTIPART_SEGMENTED_DATA_PDU,
154            SegmentedDataPdu::from_buffer(buffer).unwrap()
155        );
156    }
157}