feat: support for Server Deactivate All PDU (#418)

This commit is contained in:
Isaiah Becker-Mayer 2024-03-27 01:39:35 -05:00 committed by GitHub
parent 31ed09b909
commit e3aa8bcf0d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
22 changed files with 722 additions and 420 deletions

View file

@ -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;
@ -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,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(())
}

View file

@ -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(())
}