Skip to content

Commit

Permalink
feat: support for Server Deactivate All PDU (#418)
Browse files Browse the repository at this point in the history
  • Loading branch information
Isaiah Becker-Mayer authored Mar 27, 2024
1 parent 31ed09b commit e3aa8bc
Show file tree
Hide file tree
Showing 22 changed files with 725 additions and 423 deletions.
1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ tracing = { version = "0.1", features = ["log"] }
thiserror = "1.0"
png = "0.17"
bitflags = "2.4"
byteorder = "1.5"

[profile.dev]
opt-level = 1
Expand Down
50 changes: 4 additions & 46 deletions crates/ironrdp-acceptor/src/lib.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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?;
}
}

Expand All @@ -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)
}
52 changes: 5 additions & 47 deletions crates/ironrdp-async/src/connector.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 as _, ServerName, State as _,
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;
Expand All @@ -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)
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -177,45 +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 = 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)?
};

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(())
}
67 changes: 66 additions & 1 deletion crates/ironrdp-async/src/framed.rs
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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(())
}
1 change: 1 addition & 0 deletions crates/ironrdp-client/src/rdp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -280,6 +280,7 @@ async fn active_session(
ActiveStageOutput::PointerBitmap(_) => {
// Not applicable, because we use the software cursor rendering.
}
ActiveStageOutput::DeactivateAll(_) => todo!("DeactivateAll"),
ActiveStageOutput::Terminate(reason) => break 'outer reason,
}
}
Expand Down
Loading

0 comments on commit e3aa8bc

Please sign in to comment.