From e54fa5f4c872e09ab2f62a4bd40e6a8c285b2274 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:29:07 +0400
Subject: [PATCH] feat(async): teach single_sequence_step() to keep unmatched
 PDUs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The caller can gather the unmatching/unexpected PDUs as necessary.

Signed-off-by: Marc-André Lureau <marcandre.lureau@redhat.com>
---
 crates/ironrdp-acceptor/src/lib.rs    | 4 ++--
 crates/ironrdp-async/src/connector.rs | 4 ++--
 crates/ironrdp-async/src/framed.rs    | 6 ++++--
 crates/ironrdp-client/src/rdp.rs      | 7 ++++---
 crates/ironrdp-web/src/session.rs     | 2 +-
 5 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/crates/ironrdp-acceptor/src/lib.rs b/crates/ironrdp-acceptor/src/lib.rs
index 6faba115..169c239a 100644
--- a/crates/ironrdp-acceptor/src/lib.rs
+++ b/crates/ironrdp-acceptor/src/lib.rs
@@ -41,7 +41,7 @@ where
             return Ok(result);
         }
 
-        single_sequence_step(&mut framed, acceptor, &mut buf).await?;
+        single_sequence_step(&mut framed, acceptor, &mut buf, None).await?;
     }
 }
 
@@ -59,6 +59,6 @@ where
             return Ok((framed, result));
         }
 
-        single_sequence_step(&mut framed, acceptor, &mut buf).await?;
+        single_sequence_step(&mut framed, acceptor, &mut buf, None).await?;
     }
 }
diff --git a/crates/ironrdp-async/src/connector.rs b/crates/ironrdp-async/src/connector.rs
index 53285980..d18de69a 100644
--- a/crates/ironrdp-async/src/connector.rs
+++ b/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;
diff --git a/crates/ironrdp-async/src/framed.rs b/crates/ironrdp-async/src/framed.rs
index aa084901..2fcf0652 100644
--- a/crates/ironrdp-async/src/framed.rs
+++ b/crates/ironrdp-async/src/framed.rs
@@ -230,12 +230,13 @@ 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
 }
 
@@ -243,6 +244,7 @@ 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,
@@ -257,7 +259,7 @@ where
         );
 
         let pdu = framed
-            .read_by_hint(next_pdu_hint, None)
+            .read_by_hint(next_pdu_hint, unmatched)
             .await
             .map_err(|e| ironrdp_connector::custom_err!("read frame by hint", e))?;
 
diff --git a/crates/ironrdp-client/src/rdp.rs b/crates/ironrdp-client/src/rdp.rs
index c2c45fea..04f38eca 100644
--- a/crates/ironrdp-client/src/rdp.rs
+++ b/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| {
diff --git a/crates/ironrdp-web/src/session.rs b/crates/ironrdp-web/src/session.rs
index 717f624a..af1e8bd0 100644
--- a/crates/ironrdp-web/src/session.rs
+++ b/crates/ironrdp-web/src/session.rs
@@ -616,7 +616,7 @@ impl Session {
                         let mut buf = WriteBuf::new();
                         'activation_seq: loop {
                             let written =
-                                single_sequence_step_read(&mut framed, &mut *box_connection_activation, &mut buf)
+                                single_sequence_step_read(&mut framed, &mut *box_connection_activation, &mut buf, None)
                                     .await?;
 
                             if written.size().is_some() {