Merge branch 'feat/deactivate-all' into feat/dynamic-resize

Author: Isaiah Becker-Mayer

crates/ironrdp-acceptor/src/lib.rs

@@ -1,8 +1,8 @@
 #[macro_use]
 extern crate tracing;
 
-use ironrdp_async::{Framed, FramedRead, FramedWrite, StreamWrapper};
-use ironrdp_connector::{custom_err, ConnectorResult, Sequence, Written};
+use ironrdp_async::{single_sequence_step, Framed, FramedRead, FramedWrite, StreamWrapper};
+use ironrdp_connector::ConnectorResult;
 use ironrdp_pdu::write_buf::WriteBuf;
 
 mod channel_connection;
@@ -41,7 +41,7 @@ where
             return Ok(result);
         }
 
-        single_accept_state(&mut framed, acceptor, &mut buf).await?;
+        single_sequence_step(&mut framed, acceptor, &mut buf).await?;
     }
 }
 
@@ -59,48 +59,6 @@ where
             return Ok((framed, result));
         }
 
-        single_accept_state(&mut framed, acceptor, &mut buf).await?;
+        single_sequence_step(&mut framed, acceptor, &mut buf).await?;
     }
 }
-
-async fn single_accept_state<S>(
-    framed: &mut Framed<S>,
-    acceptor: &mut Acceptor,
-    buf: &mut WriteBuf,
-) -> ConnectorResult<Written>
-where
-    S: FramedRead + FramedWrite,
-{
-    buf.clear();
-
-    let written = if let Some(next_pdu_hint) = acceptor.next_pdu_hint() {
-        debug!(
-            acceptor.state = acceptor.state().name(),
-            hint = ?next_pdu_hint,
-            "Wait for PDU"
-        );
-
-        let pdu = framed
-            .read_by_hint(next_pdu_hint)
-            .await
-            .map_err(|e| custom_err!("read frame by hint", e))?;
-
-        trace!(length = pdu.len(), "PDU received");
-
-        acceptor.step(&pdu, buf)?
-    } else {
-        acceptor.step_no_input(buf)?
-    };
-
-    if let Some(response_len) = written.size() {
-        debug_assert_eq!(buf.filled_len(), response_len);
-        let response = buf.filled();
-        trace!(response_len, "Send response");
-        framed
-            .write_all(response)
-            .await
-            .map_err(|e| custom_err!("write all", e))?;
-    }
-
-    Ok(written)
-}

crates/ironrdp-async/src/connector.rs

@@ -2,13 +2,13 @@ use ironrdp_connector::credssp::{CredsspProcessGenerator, CredsspSequence, Kerbe
 use ironrdp_connector::sspi::credssp::ClientState;
 use ironrdp_connector::sspi::generator::GeneratorState;
 use ironrdp_connector::{
-    custom_err, ClientConnector, ClientConnectorState, ConnectionResult, ConnectorError, ConnectorResult, Sequence,
-    ServerName, State as _, Written,
+    custom_err, ClientConnector, ClientConnectorState, ConnectionResult, ConnectorError, ConnectorResult, ServerName,
+    State as _,
 };
 use ironrdp_pdu::write_buf::WriteBuf;
 
 use crate::framed::{Framed, FramedRead, FramedWrite};
-use crate::AsyncNetworkClient;
+use crate::{single_sequence_step, AsyncNetworkClient};
 
 #[non_exhaustive]
 pub struct ShouldUpgrade;
@@ -23,7 +23,7 @@ where
     info!("Begin connection procedure");
 
     while !connector.should_perform_security_upgrade() {
-        single_connect_step(framed, connector, &mut buf).await?;
+        single_sequence_step(framed, connector, &mut buf).await?;
     }
 
     Ok(ShouldUpgrade)
@@ -73,7 +73,7 @@ where
     }
 
     let result = loop {
-        single_connect_step(framed, &mut connector, &mut buf).await?;
+        single_sequence_step(framed, &mut connector, &mut buf).await?;
 
         if let ClientConnectorState::Connected { result } = connector.state {
             break result;
@@ -177,67 +177,3 @@ where
 
     Ok(())
 }
-
-pub async fn single_connect_step<S>(
-    framed: &mut Framed<S>,
-    connector: &mut ClientConnector,
-    buf: &mut WriteBuf,
-) -> ConnectorResult<()>
-where
-    S: FramedWrite + FramedRead,
-{
-    buf.clear();
-    let written = single_connect_step_read(framed, connector, buf).await?;
-    single_connect_step_write(framed, buf, written).await
-}
-
-pub async fn single_connect_step_read<S>(
-    framed: &mut Framed<S>,
-    connector: &mut dyn Sequence,
-    buf: &mut WriteBuf,
-) -> ConnectorResult<Written>
-where
-    S: FramedRead,
-{
-    buf.clear();
-
-    if let Some(next_pdu_hint) = connector.next_pdu_hint() {
-        debug!(
-            connector.state = connector.state().name(),
-            hint = ?next_pdu_hint,
-            "Wait for PDU"
-        );
-
-        let pdu = framed
-            .read_by_hint(next_pdu_hint)
-            .await
-            .map_err(|e| ironrdp_connector::custom_err!("read frame by hint", e))?;
-
-        trace!(length = pdu.len(), "PDU received");
-
-        connector.step(&pdu, buf)
-    } else {
-        connector.step_no_input(buf)
-    }
-}
-
-async fn single_connect_step_write<S>(
-    framed: &mut Framed<S>,
-    buf: &mut WriteBuf,
-    written: Written,
-) -> ConnectorResult<()>
-where
-    S: FramedWrite,
-{
-    if let Some(response_len) = written.size() {
-        debug_assert_eq!(buf.filled_len(), response_len);
-        let response = buf.filled();
-        trace!(response_len, "Send response");
-        framed
-            .write_all(response)
-            .await
-            .map_err(|e| ironrdp_connector::custom_err!("write all", e))?;
-    }
-
-    Ok(())
-}

crates/ironrdp-async/src/framed.rs

@@ -1,7 +1,8 @@
 use std::io;
 
 use bytes::{Bytes, BytesMut};
-use ironrdp_pdu::PduHint;
+use ironrdp_connector::{ConnectorResult, Sequence, Written};
+use ironrdp_pdu::{write_buf::WriteBuf, PduHint};
 
 // TODO: investigate if we could use static async fn / return position impl trait in traits when stabilized:
 // https://github.com/rust-lang/rust/issues/91611
@@ -213,3 +214,67 @@ where
         self.stream.write_all(buf).await
     }
 }
+
+pub async fn single_sequence_step<S>(
+    framed: &mut Framed<S>,
+    sequence: &mut dyn Sequence,
+    buf: &mut WriteBuf,
+) -> ConnectorResult<()>
+where
+    S: FramedWrite + FramedRead,
+{
+    buf.clear();
+    let written = single_sequence_step_read(framed, sequence, buf).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,
+) -> ConnectorResult<Written>
+where
+    S: FramedRead,
+{
+    buf.clear();
+
+    if let Some(next_pdu_hint) = sequence.next_pdu_hint() {
+        debug!(
+            connector.state = sequence.state().name(),
+            hint = ?next_pdu_hint,
+            "Wait for PDU"
+        );
+
+        let pdu = framed
+            .read_by_hint(next_pdu_hint)
+            .await
+            .map_err(|e| ironrdp_connector::custom_err!("read frame by hint", e))?;
+
+        trace!(length = pdu.len(), "PDU received");
+
+        sequence.step(&pdu, buf)
+    } else {
+        sequence.step_no_input(buf)
+    }
+}
+
+async fn single_sequence_step_write<S>(
+    framed: &mut Framed<S>,
+    buf: &mut WriteBuf,
+    written: Written,
+) -> ConnectorResult<()>
+where
+    S: FramedWrite,
+{
+    if let Some(response_len) = written.size() {
+        debug_assert_eq!(buf.filled_len(), response_len);
+        let response = buf.filled();
+        trace!(response_len, "Send response");
+        framed
+            .write_all(response)
+            .await
+            .map_err(|e| ironrdp_connector::custom_err!("write all", e))?;
+    }
+
+    Ok(())
+}

crates/ironrdp-pdu/src/rdp/headers.rs

@@ -520,16 +520,23 @@ bitflags! {
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct ServerDeactivateAll;
 
+impl ServerDeactivateAll {
+    const FIXED_PART_SIZE: usize = 2 /* length_source_descriptor */ + 1 /* source_descriptor */;
+}
+
 impl PduDecode<'_> for ServerDeactivateAll {
     fn decode(src: &mut ReadCursor<'_>) -> PduResult<Self> {
+        ensure_fixed_part_size!(in: src);
         let length_source_descriptor = src.read_u16();
+        ensure_size!(in: src, size: length_source_descriptor.into());
         let _ = src.read_slice(length_source_descriptor.into());
-        Ok(Self)
+        Ok(Self {})
     }
 }
 
 impl PduEncode for ServerDeactivateAll {
     fn encode(&self, dst: &mut WriteCursor<'_>) -> PduResult<()> {
+        ensure_fixed_part_size!(in: dst);
         // A 16-bit, unsigned integer. The size in bytes of the sourceDescriptor field.
         dst.write_u16(1);
         // Variable number of bytes. The source descriptor. This field SHOULD be set to 0x00.
@@ -542,6 +549,6 @@ impl PduEncode for ServerDeactivateAll {
     }
 
     fn size(&self) -> usize {
-        2 /* length_source_descriptor */ + 1 /* source_descriptor */
+        Self::FIXED_PART_SIZE
     }
 }