ironrdp_session/
rfx.rs

1use core::cmp::min;
2
3use ironrdp_graphics::color_conversion::{self, YCbCrBuffer};
4use ironrdp_graphics::image_processing::PixelFormat;
5use ironrdp_graphics::rectangle_processing::Region;
6use ironrdp_graphics::{dwt, quantization, rlgr, subband_reconstruction};
7use ironrdp_pdu::codecs::rfx::{self, EntropyAlgorithm, Quant, RfxRectangle, Tile};
8use ironrdp_pdu::geometry::{InclusiveRectangle, Rectangle as _};
9use ironrdp_pdu::{decode_cursor, Decode as _, ReadCursor};
10use tracing::{instrument, trace};
11
12use crate::image::DecodedImage;
13use crate::{custom_err, general_err, reason_err, SessionResult};
14
15const TILE_SIZE: u16 = 64;
16
17pub type FrameId = u32;
18
19pub struct DecodingContext {
20    context: rfx::ContextPdu,
21    channels: rfx::ChannelsPdu,
22    decoding_tiles: DecodingTileContext,
23}
24
25impl Default for DecodingContext {
26    fn default() -> Self {
27        Self {
28            context: rfx::ContextPdu {
29                flags: rfx::OperatingMode::empty(),
30                entropy_algorithm: EntropyAlgorithm::Rlgr1,
31            },
32            channels: rfx::ChannelsPdu(Vec::new()),
33            decoding_tiles: DecodingTileContext::new(),
34        }
35    }
36}
37
38impl DecodingContext {
39    pub fn new() -> Self {
40        Self::default()
41    }
42
43    pub fn decode(
44        &mut self,
45        image: &mut DecodedImage,
46        destination: &InclusiveRectangle,
47        input: &mut ReadCursor<'_>,
48    ) -> SessionResult<(FrameId, InclusiveRectangle)> {
49        loop {
50            let block = rfx::Block::decode(input).map_err(|e| custom_err!("decode block", e))?;
51
52            match block {
53                rfx::Block::Sync(_) => {
54                    self.process_sync(input)?;
55                }
56                rfx::Block::CodecChannel(rfx::CodecChannel::FrameBegin(f)) => {
57                    return self.process_frame(f, input, image, destination);
58                }
59                _ => {
60                    return Err(reason_err!(
61                        "rfx::DecodingContext",
62                        "unexpected RFX block type: {:?}",
63                        block.block_type()
64                    ));
65                }
66            }
67        }
68    }
69
70    fn process_sync(&mut self, input: &mut ReadCursor<'_>) -> SessionResult<()> {
71        self.process_headers(input)
72    }
73
74    fn process_headers(&mut self, input: &mut ReadCursor<'_>) -> SessionResult<()> {
75        let mut context = None;
76        let mut channels = None;
77
78        // headers can appear in any order: CodecVersions, Channels, Context
79        for _ in 0..3 {
80            match decode_cursor(input).map_err(|e| custom_err!("decode headers", e))? {
81                rfx::Block::CodecChannel(rfx::CodecChannel::Context(c)) => context = Some(c),
82                rfx::Block::Channels(c) => channels = Some(c),
83                rfx::Block::CodecVersions(_) => (),
84                _ => {
85                    return Err(general_err!("unexpected RFX block type"));
86                }
87            }
88        }
89
90        let context = context.ok_or_else(|| general_err!("context header is missing"))?;
91        let channels = channels.ok_or_else(|| general_err!("channels header is missing"))?;
92
93        if channels.0.is_empty() {
94            return Err(general_err!("no RFX channel announced"));
95        }
96
97        self.context = context;
98        self.channels = channels;
99
100        Ok(())
101    }
102
103    #[instrument(skip_all)]
104    fn process_frame(
105        &mut self,
106        frame_begin: rfx::FrameBeginPdu,
107        input: &mut ReadCursor<'_>,
108        image: &mut DecodedImage,
109        destination: &InclusiveRectangle,
110    ) -> SessionResult<(FrameId, InclusiveRectangle)> {
111        let channel = self
112            .channels
113            .0
114            .first()
115            .ok_or_else(|| general_err!("no RFX channel found"))?;
116        let width = channel.width.try_into().map_err(|_| general_err!("invalid width"))?;
117        let height = channel.height.try_into().map_err(|_| general_err!("invalid height"))?;
118        let entropy_algorithm = self.context.entropy_algorithm;
119
120        let region: rfx::Block<'_> = decode_cursor(input).map_err(|e| custom_err!("decode region", e))?;
121        let mut region = match region {
122            rfx::Block::CodecChannel(rfx::CodecChannel::Region(region)) => region,
123            _ => return Err(general_err!("unexpected block type")),
124        };
125        let tile_set: rfx::Block<'_> = decode_cursor(input).map_err(|e| custom_err!("decode tile_set", e))?;
126        let tile_set = match tile_set {
127            rfx::Block::CodecChannel(rfx::CodecChannel::TileSet(t)) => t,
128            _ => return Err(general_err!("unexpected block type")),
129        };
130        let frame_end: rfx::Block<'_> = decode_cursor(input).map_err(|e| custom_err!("decode frame_end", e))?;
131        if !matches!(frame_end, rfx::Block::CodecChannel(rfx::CodecChannel::FrameEnd(_))) {
132            return Err(general_err!("unexpected block type"));
133        }
134
135        if region.rectangles.is_empty() {
136            region.rectangles = vec![RfxRectangle {
137                x: 0,
138                y: 0,
139                width,
140                height,
141            }];
142        }
143        let region = region;
144
145        trace!(frame_index = frame_begin.index);
146        trace!(destination_rectangle = ?destination);
147        trace!(context = ?self.context);
148        trace!(channels = ?self.channels);
149        trace!(?region);
150
151        let clipping_rectangles = clipping_rectangles(region.rectangles.as_slice(), destination, width, height);
152        trace!("Clipping rectangles: {:?}", clipping_rectangles);
153
154        let mut final_update_rectangle = clipping_rectangles.extents.clone();
155
156        for (update_rectangle, tile_data) in tiles_to_rectangles(tile_set.tiles.as_slice(), destination)
157            .zip(map_tiles_data(tile_set.tiles.as_slice(), tile_set.quants.as_slice()))
158        {
159            decode_tile(
160                &tile_data,
161                entropy_algorithm,
162                self.decoding_tiles.tile_output.as_mut(),
163                self.decoding_tiles.ycbcr_buffer.as_mut(),
164                self.decoding_tiles.ycbcr_temp_buffer.as_mut(),
165            )?;
166
167            let current_update_rectangle = image.apply_tile(
168                &self.decoding_tiles.tile_output,
169                PixelFormat::RgbA32,
170                &clipping_rectangles,
171                &update_rectangle,
172            )?;
173
174            final_update_rectangle = final_update_rectangle.union(&current_update_rectangle);
175        }
176
177        Ok((frame_begin.index, final_update_rectangle))
178    }
179}
180
181#[derive(Debug, Clone)]
182struct DecodingTileContext {
183    tile_output: Vec<u8>,
184    ycbcr_buffer: Vec<Vec<i16>>,
185    ycbcr_temp_buffer: Vec<i16>,
186}
187
188impl DecodingTileContext {
189    fn new() -> Self {
190        let tile_size = usize::from(TILE_SIZE);
191        Self {
192            tile_output: vec![0; tile_size * tile_size * 4],
193            ycbcr_buffer: vec![vec![0; tile_size * tile_size]; 3],
194            ycbcr_temp_buffer: vec![0; tile_size * tile_size],
195        }
196    }
197}
198
199fn decode_tile(
200    tile: &TileData<'_>,
201    entropy_algorithm: EntropyAlgorithm,
202    output: &mut [u8],
203    ycbcr_temp: &mut [Vec<i16>],
204    temp: &mut [i16],
205) -> SessionResult<()> {
206    for ((quant, data), ycbcr_buffer) in tile.quants.iter().zip(tile.data.iter()).zip(ycbcr_temp.iter_mut()) {
207        decode_component(quant, entropy_algorithm, data, ycbcr_buffer.as_mut_slice(), temp)?;
208    }
209
210    let ycbcr_buffer = YCbCrBuffer {
211        y: ycbcr_temp[0].as_slice(),
212        cb: ycbcr_temp[1].as_slice(),
213        cr: ycbcr_temp[2].as_slice(),
214    };
215
216    color_conversion::ycbcr_to_rgba(ycbcr_buffer, output).map_err(|e| custom_err!("decode_tile", e))?;
217
218    Ok(())
219}
220
221fn decode_component(
222    quant: &Quant,
223    entropy_algorithm: EntropyAlgorithm,
224    data: &[u8],
225    output: &mut [i16],
226    temp: &mut [i16],
227) -> SessionResult<()> {
228    rlgr::decode(entropy_algorithm, data, output).map_err(|e| custom_err!("decode_component", e))?;
229    subband_reconstruction::decode(&mut output[4032..]);
230    quantization::decode(output, quant);
231    dwt::decode(output, temp);
232
233    Ok(())
234}
235
236fn clipping_rectangles(
237    rectangles: &[RfxRectangle],
238    destination: &InclusiveRectangle,
239    width: u16,
240    height: u16,
241) -> Region {
242    let mut clipping_rectangles = Region::new();
243
244    rectangles
245        .iter()
246        .map(|r| InclusiveRectangle {
247            left: min(destination.left + r.x, width - 1),
248            top: min(destination.top + r.y, height - 1),
249            right: min(destination.left + r.x + r.width - 1, width - 1),
250            bottom: min(destination.top + r.y + r.height - 1, height - 1),
251        })
252        .for_each(|r| clipping_rectangles.union_rectangle(r));
253
254    clipping_rectangles
255}
256
257fn tiles_to_rectangles<'a>(
258    tiles: &'a [Tile<'_>],
259    destination: &'a InclusiveRectangle,
260) -> impl Iterator<Item = InclusiveRectangle> + 'a {
261    tiles.iter().map(|t| InclusiveRectangle {
262        left: destination.left + t.x * TILE_SIZE,
263        top: destination.top + t.y * TILE_SIZE,
264        right: destination.left + t.x * TILE_SIZE + TILE_SIZE - 1,
265        bottom: destination.top + t.y * TILE_SIZE + TILE_SIZE - 1,
266    })
267}
268
269fn map_tiles_data<'a>(tiles: &[Tile<'a>], quants: &[Quant]) -> Vec<TileData<'a>> {
270    tiles
271        .iter()
272        .map(|t| TileData {
273            quants: [
274                quants[usize::from(t.y_quant_index)].clone(),
275                quants[usize::from(t.cb_quant_index)].clone(),
276                quants[usize::from(t.cr_quant_index)].clone(),
277            ],
278            data: [t.y_data, t.cb_data, t.cr_data],
279        })
280        .collect()
281}
282
283struct TileData<'a> {
284    quants: [Quant; 3],
285    data: [&'a [u8]; 3],
286}