Removes the explicit state machine from DecodingContext. State comes to us from the server via the BlockType in the BlockHeader, so it needn't be rigidly tracked explicitly in DecodingContext. This makes the RFX pipeline more flexible, which in turn allows us to seamlessly use it when we get another sync sequence mid-stream due to having fielded a Server Deactivate PDU.

Author: Isaiah Becker-Mayer

crates/ironrdp-pdu/src/codecs/rfx.rs

@@ -76,6 +76,11 @@ pub struct BlockHeader {
 }
 
 impl BlockHeader {
+    pub fn from_buffer_consume(buffer: &mut &[u8]) -> Result<Self, RfxError> {
+        let ty = BlockType::from_buffer(buffer)?;
+        Self::from_buffer_consume_with_type(buffer, ty)
+    }
+
     fn from_buffer_consume_with_type(buffer: &mut &[u8], ty: BlockType) -> Result<Self, RfxError> {
         let block_length = buffer.read_u32::<LittleEndian>()? as usize;
 
@@ -98,8 +103,7 @@ impl BlockHeader {
     }
 
     fn from_buffer_consume_with_expected_type(buffer: &mut &[u8], expected_type: BlockType) -> Result<Self, RfxError> {
-        let ty = buffer.read_u16::<LittleEndian>()?;
-        let ty = BlockType::from_u16(ty).ok_or(RfxError::InvalidBlockType(ty))?;
+        let ty = BlockType::from_buffer(buffer)?;
         if ty != expected_type {
             return Err(RfxError::UnexpectedBlockType {
                 expected: expected_type,
@@ -220,6 +224,14 @@ pub enum BlockType {
     Extension = 0xCCC7,
 }
 
+impl BlockType {
+    fn from_buffer(buffer: &mut &[u8]) -> Result<Self, RfxError> {
+        let ty = buffer.read_u16::<LittleEndian>()?;
+        let ty = BlockType::from_u16(ty).ok_or(RfxError::InvalidBlockType(ty))?;
+        Ok(ty)
+    }
+}
+
 #[derive(Debug, Error)]
 pub enum RfxError {
     #[error("IO error")]

crates/ironrdp-pdu/src/codecs/rfx/data_messages.rs

@@ -120,11 +120,8 @@ pub struct FrameBeginPdu {
     pub number_of_regions: i16,
 }
 
-impl PduBufferParsing<'_> for FrameBeginPdu {
-    type Error = RfxError;
-
-    fn from_buffer_consume(buffer: &mut &[u8]) -> Result<Self, Self::Error> {
-        let header = BlockHeader::from_buffer_consume_with_expected_type(buffer, BlockType::FrameBegin)?;
+impl FrameBeginPdu {
+    pub fn from_buffer_consume_with_header(buffer: &mut &[u8], header: BlockHeader) -> Result<Self, RfxError> {
         CodecChannelHeader::from_buffer_consume_with_type(buffer, BlockType::FrameBegin)?;
         let mut buffer = buffer.split_to(header.data_length);
 
@@ -136,6 +133,15 @@ impl PduBufferParsing<'_> for FrameBeginPdu {
             number_of_regions,
         })
     }
+}
+
+impl PduBufferParsing<'_> for FrameBeginPdu {
+    type Error = RfxError;
+
+    fn from_buffer_consume(buffer: &mut &[u8]) -> Result<Self, Self::Error> {
+        let header = BlockHeader::from_buffer_consume_with_expected_type(buffer, BlockType::FrameBegin)?;
+        Self::from_buffer_consume_with_header(buffer, header)
+    }
 
     fn to_buffer_consume(&self, buffer: &mut &mut [u8]) -> Result<(), Self::Error> {
         let header = BlockHeader {

crates/ironrdp-pdu/src/codecs/rfx/header_messages.rs

@@ -16,11 +16,8 @@ const CHANNEL_SIZE: usize = 5;
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct SyncPdu;
 
-impl PduBufferParsing<'_> for SyncPdu {
-    type Error = RfxError;
-
-    fn from_buffer_consume(buffer: &mut &[u8]) -> Result<Self, Self::Error> {
-        let header = BlockHeader::from_buffer_consume_with_expected_type(buffer, BlockType::Sync)?;
+impl SyncPdu {
+    pub fn from_buffer_consume_with_header(buffer: &mut &[u8], header: BlockHeader) -> Result<Self, RfxError> {
         let mut buffer = buffer.split_to(header.data_length);
 
         let magic = buffer.read_u32::<LittleEndian>()?;
@@ -34,6 +31,15 @@ impl PduBufferParsing<'_> for SyncPdu {
             Ok(Self)
         }
     }
+}
+
+impl PduBufferParsing<'_> for SyncPdu {
+    type Error = RfxError;
+
+    fn from_buffer_consume(buffer: &mut &[u8]) -> Result<Self, Self::Error> {
+        let header = BlockHeader::from_buffer_consume_with_expected_type(buffer, BlockType::Sync)?;
+        Self::from_buffer_consume_with_header(buffer, header)
+    }
 
     fn to_buffer_consume(&self, buffer: &mut &mut [u8]) -> Result<(), Self::Error> {
         let header = BlockHeader {

crates/ironrdp-session/src/rfx.rs

@@ -15,7 +15,6 @@ const TILE_SIZE: u16 = 64;
 pub type FrameId = u32;
 
 pub struct DecodingContext {
-    state: SequenceState,
     context: rfx::ContextPdu,
     channels: rfx::ChannelsPdu,
     decoding_tiles: DecodingTileContext,
@@ -24,7 +23,6 @@ pub struct DecodingContext {
 impl Default for DecodingContext {
     fn default() -> Self {
         Self {
-            state: SequenceState::HeaderMessages,
             context: rfx::ContextPdu {
                 flags: rfx::OperatingMode::empty(),
                 entropy_algorithm: rfx::EntropyAlgorithm::Rlgr1,
@@ -47,20 +45,31 @@ impl DecodingContext {
         input: &mut &[u8],
     ) -> SessionResult<(FrameId, InclusiveRectangle)> {
         loop {
-            match self.state {
-                SequenceState::HeaderMessages => {
-                    self.process_headers(input)?;
+            let block_header = rfx::BlockHeader::from_buffer_consume(input)?;
+            match block_header.ty {
+                rfx::BlockType::Sync => {
+                    self.process_sync(input, block_header)?;
                 }
-                SequenceState::DataMessages => {
-                    return self.process_data_messages(image, destination, input);
+                rfx::BlockType::FrameBegin => {
+                    return self.process_frame(input, block_header, image, destination);
+                }
+                _ => {
+                    return Err(reason_err!(
+                        "rfx::DecodingContext",
+                        "unexpected RFX block type: {:?}",
+                        block_header.ty
+                    ));
                 }
             }
         }
     }
 
-    fn process_headers(&mut self, input: &mut &[u8]) -> SessionResult<()> {
-        let _sync = rfx::SyncPdu::from_buffer_consume(input)?;
+    fn process_sync(&mut self, input: &mut &[u8], header: rfx::BlockHeader) -> SessionResult<()> {
+        let _sync = rfx::SyncPdu::from_buffer_consume_with_header(input, header)?;
+        self.process_headers(input)
+    }
 
+    fn process_headers(&mut self, input: &mut &[u8]) -> SessionResult<()> {
         let mut context = None;
         let mut channels = None;
 
@@ -81,24 +90,24 @@ impl DecodingContext {
 
         self.context = context;
         self.channels = channels;
-        self.state = SequenceState::DataMessages;
 
         Ok(())
     }
 
     #[instrument(skip_all)]
-    fn process_data_messages(
+    fn process_frame(
         &mut self,
+        input: &mut &[u8],
+        header: rfx::BlockHeader,
         image: &mut DecodedImage,
         destination: &InclusiveRectangle,
-        input: &mut &[u8],
     ) -> SessionResult<(FrameId, InclusiveRectangle)> {
         let channel = self.channels.0.first().unwrap();
         let width = channel.width.as_u16();
         let height = channel.height.as_u16();
         let entropy_algorithm = self.context.entropy_algorithm;
 
-        let frame_begin = rfx::FrameBeginPdu::from_buffer_consume(input)?;
+        let frame_begin = rfx::FrameBeginPdu::from_buffer_consume_with_header(input, header)?;
         let mut region = rfx::RegionPdu::from_buffer_consume(input)?;
         let tile_set = rfx::TileSetPdu::from_buffer_consume(input)?;
         let _frame_end = rfx::FrameEndPdu::from_buffer_consume(input)?;
@@ -145,10 +154,6 @@ impl DecodingContext {
             final_update_rectangle = final_update_rectangle.union(&current_update_rectangle);
         }
 
-        if self.context.flags.contains(rfx::OperatingMode::IMAGE_MODE) {
-            self.state = SequenceState::HeaderMessages;
-        }
-
         Ok((frame_begin.index, final_update_rectangle))
     }
 }
@@ -258,8 +263,3 @@ struct TileData<'a> {
     quants: [Quant; 3],
     data: [&'a [u8]; 3],
 }
-
-enum SequenceState {
-    HeaderMessages,
-    DataMessages,
-}

crates/ironrdp-session/src/x224/mod.rs

@@ -30,7 +30,10 @@ pub enum ProcessorOutput {
     ResponseFrame(Vec<u8>),
     /// A graceful disconnect notification. Client should close the connection upon receiving this.
     Disconnect(DisconnectReason),
-    /// Received a [`ironrdp_pdu::rdp::headers::ServerDeactivateAll`] PDU.
+    /// Received a [`ironrdp_pdu::rdp::headers::ServerDeactivateAll`] PDU. Client should execute the
+    /// [Deactivation-Reactivation Sequence].
+    ///
+    /// [Deactivation-Reactivation Sequence]: https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-rdpbcgr/dfc234ce-481a-4674-9a5d-2a7bafb14432
     DeactivateAll(ConnectionActivationSequence),
 }