Server resize follow-up

crates/ironrdp-acceptor/src/lib.rs

@@ -1,6 +1,7 @@
 #[macro_use]
 extern crate tracing;
 
+use ironrdp_async::bytes::Bytes;
 use ironrdp_async::{single_sequence_step, Framed, FramedRead, FramedWrite, StreamWrapper};
 use ironrdp_connector::ConnectorResult;
 use ironrdp_pdu::write_buf::WriteBuf;
@@ -41,13 +42,14 @@ where
             return Ok(result);
         }
 
-        single_sequence_step(&mut framed, acceptor, &mut buf).await?;
+        single_sequence_step(&mut framed, acceptor, &mut buf, None).await?;
     }
 }
 
 pub async fn accept_finalize<S>(
     mut framed: Framed<S>,
     acceptor: &mut Acceptor,
+    mut unmatched: Option<&mut Vec<Bytes>>,
 ) -> ConnectorResult<(Framed<S>, AcceptorResult)>
 where
     S: FramedRead + FramedWrite,
@@ -58,7 +60,6 @@ where
         if let Some(result) = acceptor.get_result() {
             return Ok((framed, result));
         }
-
-        single_sequence_step(&mut framed, acceptor, &mut buf).await?;
+        single_sequence_step(&mut framed, acceptor, &mut buf, unmatched.as_deref_mut()).await?;
     }
 }

crates/ironrdp-async/src/connector.rs

@@ -23,7 +23,7 @@ where
     info!("Begin connection procedure");
 
     while !connector.should_perform_security_upgrade() {
-        single_sequence_step(framed, connector, &mut buf).await?;
+        single_sequence_step(framed, connector, &mut buf, None).await?;
     }
 
     Ok(ShouldUpgrade)
@@ -73,7 +73,7 @@ where
     }
 
     let result = loop {
-        single_sequence_step(framed, &mut connector, &mut buf).await?;
+        single_sequence_step(framed, &mut connector, &mut buf, None).await?;
 
         if let ClientConnectorState::Connected { result } = connector.state {
             break result;
@@ -171,7 +171,7 @@ where
         );
 
         let pdu = framed
-            .read_by_hint(next_pdu_hint)
+            .read_by_hint(next_pdu_hint, None)
             .await
             .map_err(|e| ironrdp_connector::custom_err!("read frame by hint", e))?;
 

crates/ironrdp-async/src/framed.rs

@@ -164,14 +164,25 @@ where
     /// `tokio::select!` statement and some other branch
     /// completes first, then it is safe to drop the future and re-create it later.
     /// Data may have been read, but it will be stored in the internal buffer.
-    pub async fn read_by_hint(&mut self, hint: &dyn PduHint) -> io::Result<Bytes> {
+    pub async fn read_by_hint(
+        &mut self,
+        hint: &dyn PduHint,
+        mut unmatched: Option<&mut Vec<Bytes>>,
+    ) -> io::Result<Bytes> {
         loop {
             match hint
                 .find_size(self.peek())
                 .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?
             {
-                Some(length) => {
-                    return Ok(self.read_exact(length).await?.freeze());
+                Some((matched, length)) => {
+                    let bytes = self.read_exact(length).await?.freeze();
+                    if matched {
+                        return Ok(bytes);
+                    } else if let Some(ref mut unmatched) = unmatched {
+                        unmatched.push(bytes);
+                    } else {
+                        warn!("Received and lost an unexpected PDU");
+                    }
                 }
                 None => {
                     let len = self.read().await?;
@@ -219,19 +230,21 @@ pub async fn single_sequence_step<S>(
     framed: &mut Framed<S>,
     sequence: &mut dyn Sequence,
     buf: &mut WriteBuf,
+    unmatched: Option<&mut Vec<Bytes>>,
 ) -> ConnectorResult<()>
 where
     S: FramedWrite + FramedRead,
 {
     buf.clear();
-    let written = single_sequence_step_read(framed, sequence, buf).await?;
+    let written = single_sequence_step_read(framed, sequence, buf, unmatched).await?;
     single_sequence_step_write(framed, buf, written).await
 }
 
 pub async fn single_sequence_step_read<S>(
     framed: &mut Framed<S>,
     sequence: &mut dyn Sequence,
     buf: &mut WriteBuf,
+    unmatched: Option<&mut Vec<Bytes>>,
 ) -> ConnectorResult<Written>
 where
     S: FramedRead,
@@ -246,7 +259,7 @@ where
         );
 
         let pdu = framed
-            .read_by_hint(next_pdu_hint)
+            .read_by_hint(next_pdu_hint, unmatched)
             .await
             .map_err(|e| ironrdp_connector::custom_err!("read frame by hint", e))?;
 

crates/ironrdp-blocking/src/connector.rs

@@ -1,5 +1,6 @@
 use std::io::{Read, Write};
 
+use bytes::Bytes;
 use ironrdp_connector::credssp::{CredsspProcessGenerator, CredsspSequence, KerberosConfig};
 use ironrdp_connector::sspi::credssp::ClientState;
 use ironrdp_connector::sspi::generator::GeneratorState;
@@ -25,7 +26,7 @@ where
     info!("Begin connection procedure");
 
     while !connector.should_perform_security_upgrade() {
-        single_sequence_step(framed, connector, &mut buf)?;
+        single_sequence_step(framed, connector, &mut buf, None)?;
     }
 
     Ok(ShouldUpgrade)
@@ -78,7 +79,7 @@ where
     debug!("Remaining of connection sequence");
 
     let result = loop {
-        single_sequence_step(framed, &mut connector, &mut buf)?;
+        single_sequence_step(framed, &mut connector, &mut buf, None)?;
 
         if let ClientConnectorState::Connected { result } = connector.state {
             break result;
@@ -167,7 +168,7 @@ where
         );
 
         let pdu = framed
-            .read_by_hint(next_pdu_hint)
+            .read_by_hint(next_pdu_hint, None)
             .map_err(|e| ironrdp_connector::custom_err!("read frame by hint", e))?;
 
         trace!(length = pdu.len(), "PDU received");
@@ -188,6 +189,7 @@ pub fn single_sequence_step<S>(
     framed: &mut Framed<S>,
     connector: &mut ClientConnector,
     buf: &mut WriteBuf,
+    unmatched: Option<&mut Vec<Bytes>>,
 ) -> ConnectorResult<()>
 where
     S: Read + Write,
@@ -202,7 +204,7 @@ where
         );
 
         let pdu = framed
-            .read_by_hint(next_pdu_hint)
+            .read_by_hint(next_pdu_hint, unmatched)
             .map_err(|e| ironrdp_connector::custom_err!("read frame by hint", e))?;
 
         trace!(length = pdu.len(), "PDU received");

crates/ironrdp-blocking/src/framed.rs

@@ -87,14 +87,21 @@ where
     }
 
     /// Reads a frame using the provided PduHint.
-    pub fn read_by_hint(&mut self, hint: &dyn PduHint) -> io::Result<Bytes> {
+    pub fn read_by_hint(&mut self, hint: &dyn PduHint, mut unmatched: Option<&mut Vec<Bytes>>) -> io::Result<Bytes> {
         loop {
             match hint
                 .find_size(self.peek())
                 .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?
             {
-                Some(length) => {
-                    return Ok(self.read_exact(length)?.freeze());
+                Some((matched, length)) => {
+                    let bytes = self.read_exact(length)?.freeze();
+                    if matched {
+                        return Ok(bytes);
+                    } else if let Some(ref mut unmatched) = unmatched {
+                        unmatched.push(bytes);
+                    } else {
+                        warn!("Received and lost an unexpected PDU");
+                    }
                 }
                 None => {
                     let len = self.read()?;

crates/ironrdp-client/src/rdp.rs

@@ -294,9 +294,10 @@ async fn active_session(
                     debug!("Received Server Deactivate All PDU, executing Deactivation-Reactivation Sequence");
                     let mut buf = WriteBuf::new();
                     'activation_seq: loop {
-                        let written = single_sequence_step_read(&mut framed, &mut *connection_activation, &mut buf)
-                            .await
-                            .map_err(|e| session::custom_err!("read deactivation-reactivation sequence step", e))?;
+                        let written =
+                            single_sequence_step_read(&mut framed, &mut *connection_activation, &mut buf, None)
+                                .await
+                                .map_err(|e| session::custom_err!("read deactivation-reactivation sequence step", e))?;
 
                         if written.size().is_some() {
                             framed.write_all(buf.filled()).await.map_err(|e| {

crates/ironrdp-connector/src/credssp.rs

@@ -43,9 +43,9 @@ struct CredsspTsRequestHint;
 const CREDSSP_TS_REQUEST_HINT: CredsspTsRequestHint = CredsspTsRequestHint;
 
 impl PduHint for CredsspTsRequestHint {
-    fn find_size(&self, bytes: &[u8]) -> ironrdp_pdu::PduResult<Option<usize>> {
+    fn find_size(&self, bytes: &[u8]) -> ironrdp_pdu::PduResult<Option<(bool, usize)>> {
         match credssp::TsRequest::read_length(bytes) {
-            Ok(length) => Ok(Some(length)),
+            Ok(length) => Ok(Some((true, length))),
             Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => Ok(None),
             Err(e) => Err(ironrdp_pdu::custom_err!("CredsspTsRequestHint", e)),
         }
@@ -58,8 +58,8 @@ struct CredsspEarlyUserAuthResultHint;
 const CREDSSP_EARLY_USER_AUTH_RESULT_HINT: CredsspEarlyUserAuthResultHint = CredsspEarlyUserAuthResultHint;
 
 impl PduHint for CredsspEarlyUserAuthResultHint {
-    fn find_size(&self, _: &[u8]) -> ironrdp_pdu::PduResult<Option<usize>> {
-        Ok(Some(credssp::EARLY_USER_AUTH_RESULT_PDU_SIZE))
+    fn find_size(&self, _: &[u8]) -> ironrdp_pdu::PduResult<Option<(bool, usize)>> {
+        Ok(Some((true, credssp::EARLY_USER_AUTH_RESULT_PDU_SIZE)))
     }
 }
 

crates/ironrdp-glutin-renderer/src/draw.rs

@@ -4,10 +4,7 @@ use std::slice::from_raw_parts;
 use std::sync::Arc;
 
 use glow::*;
-use ironrdp::pdu::geometry::{
-    Rectangle as _,
-    InclusiveRectangle,
-}
+use ironrdp::pdu::geometry::{InclusiveRectangle, Rectangle as _};
 
 fn cast_as_bytes<T>(input: &[T]) -> &[u8] {
     unsafe { from_raw_parts(input.as_ptr() as *const u8, input.len() * size_of::<T>()) }

crates/ironrdp-pdu/src/basic_output/bitmap.rs

@@ -244,6 +244,12 @@ impl PduEncode for CompressedDataHeader {
     fn encode(&self, dst: &mut crate::cursor::WriteCursor<'_>) -> PduResult<()> {
         ensure_fixed_part_size!(in: dst);
 
+        if self.scan_width % 4 != 0 {
+            return Err(invalid_message_err!(
+                "cbScanWidth",
+                "The width of the bitmap must be divisible by 4"
+            ));
+        }
         dst.write_u16(FIRST_ROW_SIZE_VALUE);
         dst.write_u16(self.main_body_size);
         dst.write_u16(self.scan_width);

crates/ironrdp-pdu/src/lib.rs

@@ -333,7 +333,10 @@ pub fn find_size(bytes: &[u8]) -> PduResult<Option<PduInfo>> {
 
 pub trait PduHint: Send + Sync + fmt::Debug + 'static {
     /// Finds next PDU size by reading the next few bytes.
-    fn find_size(&self, bytes: &[u8]) -> PduResult<Option<usize>>;
+    ///
+    /// Returns `Some((hint_matching, size))` if the size is known.
+    /// Returns `None` if the size cannot be determined yet.
+    fn find_size(&self, bytes: &[u8]) -> PduResult<Option<(bool, usize)>>;
 }
 
 // Matches both X224 and FastPath pdus
@@ -343,8 +346,8 @@ pub struct RdpHint;
 pub const RDP_HINT: RdpHint = RdpHint;
 
 impl PduHint for RdpHint {
-    fn find_size(&self, bytes: &[u8]) -> PduResult<Option<usize>> {
-        find_size(bytes).map(|opt| opt.map(|info| info.length))
+    fn find_size(&self, bytes: &[u8]) -> PduResult<Option<(bool, usize)>> {
+        find_size(bytes).map(|opt| opt.map(|info| (true, info.length)))
     }
 }
 
@@ -354,11 +357,11 @@ pub struct X224Hint;
 pub const X224_HINT: X224Hint = X224Hint;
 
 impl PduHint for X224Hint {
-    fn find_size(&self, bytes: &[u8]) -> PduResult<Option<usize>> {
+    fn find_size(&self, bytes: &[u8]) -> PduResult<Option<(bool, usize)>> {
         match find_size(bytes)? {
             Some(pdu_info) => {
-                debug_assert_eq!(pdu_info.action, Action::X224);
-                Ok(Some(pdu_info.length))
+                let res = (pdu_info.action == Action::X224, pdu_info.length);
+                Ok(Some(res))
             }
             None => Ok(None),
         }
@@ -371,11 +374,11 @@ pub struct FastPathHint;
 pub const FAST_PATH_HINT: FastPathHint = FastPathHint;
 
 impl PduHint for FastPathHint {
-    fn find_size(&self, bytes: &[u8]) -> PduResult<Option<usize>> {
+    fn find_size(&self, bytes: &[u8]) -> PduResult<Option<(bool, usize)>> {
         match find_size(bytes)? {
             Some(pdu_info) => {
-                debug_assert_eq!(pdu_info.action, Action::FastPath);
-                Ok(Some(pdu_info.length))
+                let res = (pdu_info.action == Action::FastPath, pdu_info.length);
+                Ok(Some(res))
             }
             None => Ok(None),
         }

crates/ironrdp-server/src/encoder/bitmap.rs

@@ -3,7 +3,7 @@ use ironrdp_graphics::rdp6::{ABgrChannels, ARgbChannels, BgrAChannels, BitmapStr
 use ironrdp_pdu::bitmap::{self, BitmapData, BitmapUpdateData, Compression};
 use ironrdp_pdu::cursor::WriteCursor;
 use ironrdp_pdu::geometry::InclusiveRectangle;
-use ironrdp_pdu::{PduEncode, PduError};
+use ironrdp_pdu::{invalid_message_err, PduEncode, PduError};
 
 use crate::{BitmapUpdate, PixelOrder};
 
@@ -20,29 +20,41 @@ impl BitmapEncoder {
     }
 
     pub(crate) fn encode(&mut self, bitmap: &BitmapUpdate, output: &mut [u8]) -> Result<usize, PduError> {
-        let row_len = usize::from(bitmap.width.get()) * usize::from(bitmap.format.bytes_per_pixel());
+        // FIXME: support non-multiple of 4 widths.
+        //
+        // It’s not clear how to achieve that yet, but generally, server uses multiple of 4-widths,
+        // and client has surface capabilities, so this path is unlikely.
+        if bitmap.width.get() % 4 != 0 {
+            return Err(invalid_message_err!("bitmap", "Width must be a multiple of 4"));
+        }
+
+        let bytes_per_pixel = usize::from(bitmap.format.bytes_per_pixel());
+        let row_len = usize::from(bitmap.width.get()) * bytes_per_pixel;
         let chunk_height = usize::from(u16::MAX) / row_len;
 
         let mut cursor = WriteCursor::new(output);
-        let chunks = bitmap.data.chunks(row_len * chunk_height);
+        let chunks = bitmap.data.chunks(bitmap.stride * chunk_height);
 
-        let total = u16::try_from(chunks.clone().count()).unwrap();
+        let total = u16::try_from(chunks.size_hint().0).unwrap();
         BitmapUpdateData::encode_header(total, &mut cursor)?;
 
         for (i, chunk) in chunks.enumerate() {
-            let height = chunk.len() / row_len;
+            let height = chunk.len() / bitmap.stride;
             let top = usize::from(bitmap.top) + i * chunk_height;
 
             let encoder = BitmapStreamEncoder::new(usize::from(bitmap.width.get()), height);
 
             let len = match bitmap.order {
                 PixelOrder::BottomToTop => {
-                    Self::encode_slice(encoder, bitmap.format, chunk, self.buffer.as_mut_slice())
+                    Self::encode_slice(encoder, bitmap.format, &chunk[..row_len], self.buffer.as_mut_slice())
                 }
 
                 PixelOrder::TopToBottom => {
-                    let bytes_per_pixel = usize::from(bitmap.format.bytes_per_pixel());
-                    let pixels = chunk.chunks(row_len).rev().flat_map(|row| row.chunks(bytes_per_pixel));
+                    let pixels = chunk
+                        .chunks(bitmap.stride)
+                        .map(|row| &row[..row_len])
+                        .rev()
+                        .flat_map(|row| row.chunks(bytes_per_pixel));
 
                     Self::encode_iter(encoder, bitmap.format, pixels, self.buffer.as_mut_slice())
                 }
@@ -62,7 +74,7 @@ impl BitmapEncoder {
                 compressed_data_header: Some(bitmap::CompressedDataHeader {
                     main_body_size: u16::try_from(len).unwrap(),
                     scan_width: u16::from(bitmap.width),
-                    uncompressed_size: u16::try_from(chunk.len()).unwrap(),
+                    uncompressed_size: u16::try_from(height * row_len).unwrap(),
                 }),
                 bitmap_data: &self.buffer[..len],
             };