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 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(¤t_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}