From 98e7dbab998ef3e4d250fbf784bc217fd713902d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marc-Andr=C3=A9=20Lureau?= <marcandre.lureau@redhat.com>
Date: Wed, 14 Aug 2024 19:17:02 +0400
Subject: [PATCH] refactor(async): let read_by_hint() optionally accumulate
 unmatched bytes
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The caller can then decide what to do.

Signed-off-by: Marc-André Lureau <marcandre.lureau@redhat.com>
---
 crates/ironrdp-async/src/connector.rs |  2 +-
 crates/ironrdp-async/src/framed.rs    | 19 +++++++++++++++----
 crates/ironrdp-web/src/session.rs     |  2 +-
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/crates/ironrdp-async/src/connector.rs b/crates/ironrdp-async/src/connector.rs
index 1a2d9550..53285980 100644
--- a/crates/ironrdp-async/src/connector.rs
+++ b/crates/ironrdp-async/src/connector.rs
@@ -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))?;
 
diff --git a/crates/ironrdp-async/src/framed.rs b/crates/ironrdp-async/src/framed.rs
index 5efcb3e9..aa084901 100644
--- a/crates/ironrdp-async/src/framed.rs
+++ b/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((_matched, 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?;
@@ -246,7 +257,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))?;
 
diff --git a/crates/ironrdp-web/src/session.rs b/crates/ironrdp-web/src/session.rs
index 628b7941..717f624a 100644
--- a/crates/ironrdp-web/src/session.rs
+++ b/crates/ironrdp-web/src/session.rs
@@ -940,7 +940,7 @@ where
         // RDCleanPath response
 
         let rdcleanpath_res = framed
-            .read_by_hint(&RDCLEANPATH_HINT)
+            .read_by_hint(&RDCLEANPATH_HINT, None)
             .await
             .context("read RDCleanPath request")?;