diff --git a/crates/ironrdp-acceptor/src/connection.rs b/crates/ironrdp-acceptor/src/connection.rs index f4777e8e..d91f94c4 100644 --- a/crates/ironrdp-acceptor/src/connection.rs +++ b/crates/ironrdp-acceptor/src/connection.rs @@ -34,6 +34,7 @@ pub struct Acceptor { static_channels: StaticChannelSet, saved_for_reactivation: AcceptorState, pub(crate) creds: Option, + reactivation: bool, } #[derive(Debug)] @@ -62,6 +63,7 @@ impl Acceptor { static_channels: StaticChannelSet::new(), saved_for_reactivation: Default::default(), creds, + reactivation: false, } } @@ -98,6 +100,7 @@ impl Acceptor { static_channels: StaticChannelSet::new(), saved_for_reactivation, creds: consumed.creds, + reactivation: true, } } @@ -291,7 +294,9 @@ impl Sequence for Acceptor { } fn step(&mut self, input: &[u8], output: &mut WriteBuf) -> ConnectorResult { - let (written, next_state) = match mem::take(&mut self.state) { + let prev_state = mem::take(&mut self.state); + + let (written, next_state) = match prev_state { AcceptorState::InitiationWaitRequest => { let connection_request = decode::>(input) .map_err(ConnectorError::decode) @@ -639,15 +644,38 @@ impl Sequence for Acceptor { ) } - AcceptorState::CapabilitiesWaitConfirm { channels } => { + AcceptorState::CapabilitiesWaitConfirm { ref channels } => { let message = decode::>>(input) .map_err(ConnectorError::decode) - .map(|p| p.0)?; - + .map(|p| p.0); + let message = match message { + Ok(msg) => msg, + Err(e) => { + if self.reactivation { + debug!("Dropping unexpected PDU during reactivation"); + self.state = prev_state; + return Ok(Written::Nothing); + } else { + return Err(e); + } + } + }; match message { mcs::McsMessage::SendDataRequest(data) => { let capabilities_confirm = decode::(data.user_data.as_ref()) - .map_err(ConnectorError::decode)?; + .map_err(ConnectorError::decode); + let capabilities_confirm = match capabilities_confirm { + Ok(capabilities_confirm) => capabilities_confirm, + Err(e) => { + if self.reactivation { + debug!("Dropping unexpected PDU during reactivation"); + self.state = prev_state; + return Ok(Written::Nothing); + } else { + return Err(e); + } + } + }; debug!(message = ?capabilities_confirm, "Received"); @@ -659,7 +687,7 @@ impl Sequence for Acceptor { ( Written::Nothing, AcceptorState::ConnectionFinalization { - channels, + channels: channels.clone(), finalization: FinalizationSequence::new(self.user_channel_id, self.io_channel_id), client_capabilities: confirm.pdu.capability_sets, }, @@ -673,7 +701,7 @@ impl Sequence for Acceptor { _ => { warn!(?message, "Unexpected MCS message received"); - (Written::Nothing, AcceptorState::CapabilitiesWaitConfirm { channels }) + (Written::Nothing, prev_state) } } } @@ -684,6 +712,7 @@ impl Sequence for Acceptor { client_capabilities, } => { let written = finalization.step(input, output)?; + let state = if finalization.is_done() { AcceptorState::Accepted { channels, diff --git a/crates/ironrdp-acceptor/src/lib.rs b/crates/ironrdp-acceptor/src/lib.rs index a8091768..e42566f3 100644 --- a/crates/ironrdp-acceptor/src/lib.rs +++ b/crates/ironrdp-acceptor/src/lib.rs @@ -4,7 +4,6 @@ #[macro_use] extern crate tracing; -use ironrdp_async::bytes::Bytes; use ironrdp_async::{single_sequence_step, Framed, FramedRead, FramedWrite, StreamWrapper}; use ironrdp_connector::credssp::KerberosConfig; use ironrdp_connector::sspi::credssp::EarlyUserAuthResult; @@ -50,7 +49,7 @@ where return Ok(result); } - single_sequence_step(&mut framed, acceptor, &mut buf, None).await?; + single_sequence_step(&mut framed, acceptor, &mut buf).await?; } } @@ -84,7 +83,6 @@ where pub async fn accept_finalize( mut framed: Framed, acceptor: &mut Acceptor, - mut unmatched: Option<&mut Vec>, ) -> ConnectorResult<(Framed, AcceptorResult)> where S: FramedRead + FramedWrite, @@ -95,7 +93,7 @@ where if let Some(result) = acceptor.get_result() { return Ok((framed, result)); } - single_sequence_step(&mut framed, acceptor, &mut buf, unmatched.as_deref_mut()).await?; + single_sequence_step(&mut framed, acceptor, &mut buf).await?; } } @@ -152,7 +150,7 @@ where ); let pdu = framed - .read_by_hint(next_pdu_hint, None) + .read_by_hint(next_pdu_hint) .await .map_err(|e| ironrdp_connector::custom_err!("read frame by hint", e))?; diff --git a/crates/ironrdp-async/src/connector.rs b/crates/ironrdp-async/src/connector.rs index 23a300f4..4d1b085d 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, None).await?; + single_sequence_step(framed, connector, &mut buf).await?; } Ok(ShouldUpgrade) @@ -73,7 +73,7 @@ where } let result = loop { - single_sequence_step(framed, &mut connector, &mut buf, None).await?; + single_sequence_step(framed, &mut connector, &mut buf).await?; if let ClientConnectorState::Connected { result } = connector.state { break result; @@ -171,7 +171,7 @@ where ); let pdu = framed - .read_by_hint(next_pdu_hint, None) + .read_by_hint(next_pdu_hint) .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 370532c9..f628338b 100644 --- a/crates/ironrdp-async/src/framed.rs +++ b/crates/ironrdp-async/src/framed.rs @@ -165,11 +165,7 @@ 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, - mut unmatched: Option<&mut Vec>, - ) -> io::Result { + pub async fn read_by_hint(&mut self, hint: &dyn PduHint) -> io::Result { loop { match hint .find_size(self.peek()) @@ -179,10 +175,8 @@ where 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"); + debug!("Received and lost an unexpected PDU"); } } None => { @@ -236,13 +230,12 @@ pub async fn single_sequence_step( framed: &mut Framed, sequence: &mut dyn Sequence, buf: &mut WriteBuf, - unmatched: Option<&mut Vec>, ) -> ConnectorResult<()> where S: FramedWrite + FramedRead, { buf.clear(); - let written = single_sequence_step_read(framed, sequence, buf, unmatched).await?; + let written = single_sequence_step_read(framed, sequence, buf).await?; single_sequence_step_write(framed, buf, written).await } @@ -250,7 +243,6 @@ pub async fn single_sequence_step_read( framed: &mut Framed, sequence: &mut dyn Sequence, buf: &mut WriteBuf, - unmatched: Option<&mut Vec>, ) -> ConnectorResult where S: FramedRead, @@ -265,7 +257,7 @@ where ); let pdu = framed - .read_by_hint(next_pdu_hint, unmatched) + .read_by_hint(next_pdu_hint) .await .map_err(|e| ironrdp_connector::custom_err!("read frame by hint", e))?; diff --git a/crates/ironrdp-blocking/src/connector.rs b/crates/ironrdp-blocking/src/connector.rs index f98915c9..765560e6 100644 --- a/crates/ironrdp-blocking/src/connector.rs +++ b/crates/ironrdp-blocking/src/connector.rs @@ -1,6 +1,5 @@ use std::io::{Read, Write}; -use bytes::Bytes; use ironrdp_connector::credssp::{CredsspProcessGenerator, CredsspSequence, KerberosConfig}; use ironrdp_connector::sspi::credssp::ClientState; use ironrdp_connector::sspi::generator::GeneratorState; @@ -26,7 +25,7 @@ where info!("Begin connection procedure"); while !connector.should_perform_security_upgrade() { - single_sequence_step(framed, connector, &mut buf, None)?; + single_sequence_step(framed, connector, &mut buf)?; } Ok(ShouldUpgrade) @@ -79,7 +78,7 @@ where debug!("Remaining of connection sequence"); let result = loop { - single_sequence_step(framed, &mut connector, &mut buf, None)?; + single_sequence_step(framed, &mut connector, &mut buf)?; if let ClientConnectorState::Connected { result } = connector.state { break result; @@ -168,7 +167,7 @@ where ); let pdu = framed - .read_by_hint(next_pdu_hint, None) + .read_by_hint(next_pdu_hint) .map_err(|e| ironrdp_connector::custom_err!("read frame by hint", e))?; trace!(length = pdu.len(), "PDU received"); @@ -189,7 +188,6 @@ pub fn single_sequence_step( framed: &mut Framed, connector: &mut ClientConnector, buf: &mut WriteBuf, - unmatched: Option<&mut Vec>, ) -> ConnectorResult<()> where S: Read + Write, @@ -204,7 +202,7 @@ where ); let pdu = framed - .read_by_hint(next_pdu_hint, unmatched) + .read_by_hint(next_pdu_hint) .map_err(|e| ironrdp_connector::custom_err!("read frame by hint", e))?; trace!(length = pdu.len(), "PDU received"); diff --git a/crates/ironrdp-blocking/src/framed.rs b/crates/ironrdp-blocking/src/framed.rs index bcc2242d..e12483a9 100644 --- a/crates/ironrdp-blocking/src/framed.rs +++ b/crates/ironrdp-blocking/src/framed.rs @@ -87,7 +87,7 @@ where } /// Reads a frame using the provided PduHint. - pub fn read_by_hint(&mut self, hint: &dyn PduHint, mut unmatched: Option<&mut Vec>) -> io::Result { + pub fn read_by_hint(&mut self, hint: &dyn PduHint) -> io::Result { loop { match hint .find_size(self.peek()) @@ -97,10 +97,8 @@ where let bytes = self.read_exact(length)?.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"); + debug!("Received and lost an unexpected PDU"); } } None => { diff --git a/crates/ironrdp-client/src/rdp.rs b/crates/ironrdp-client/src/rdp.rs index 6c5cf28c..fedd012e 100644 --- a/crates/ironrdp-client/src/rdp.rs +++ b/crates/ironrdp-client/src/rdp.rs @@ -295,10 +295,9 @@ 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 reader, &mut *connection_activation, &mut buf, None) - .await - .map_err(|e| session::custom_err!("read deactivation-reactivation sequence step", e))?; + let written = single_sequence_step_read(&mut reader, &mut *connection_activation, &mut buf) + .await + .map_err(|e| session::custom_err!("read deactivation-reactivation sequence step", e))?; if written.size().is_some() { writer.write_all(buf.filled()).await.map_err(|e| { diff --git a/crates/ironrdp-server/src/server.rs b/crates/ironrdp-server/src/server.rs index c6c5ec64..483d355a 100644 --- a/crates/ironrdp-server/src/server.rs +++ b/crates/ironrdp-server/src/server.rs @@ -901,32 +901,22 @@ impl RdpServer { where S: AsyncRead + AsyncWrite + Sync + Send + Unpin, { - let mut other_pdus = None; - loop { - let (new_framed, result) = ironrdp_acceptor::accept_finalize(framed, &mut acceptor, other_pdus.as_mut()) + let (new_framed, result) = ironrdp_acceptor::accept_finalize(framed, &mut acceptor) .await .context("failed to accept client during finalize")?; - let (stream, mut leftover) = new_framed.into_inner(); - - if let Some(pdus) = other_pdus.take() { - let unmatched_frames = pdus.into_iter().flatten(); - let previous_leftover = leftover.split(); - leftover.extend(unmatched_frames); - leftover.extend_from_slice(&previous_leftover); - } - - let (mut reader, mut writer) = split_tokio_framed(TokioFramed::new_with_leftover(stream, leftover)); + let (mut reader, mut writer) = split_tokio_framed(new_framed); match self.client_accepted(&mut reader, &mut writer, result).await? { RunState::Continue => { unreachable!(); } RunState::DeactivationReactivation { desktop_size } => { - other_pdus = Some(Vec::new()); - acceptor = Acceptor::new_deactivation_reactivation(acceptor, desktop_size); - self.attach_channels(&mut acceptor); + acceptor = Acceptor::new_deactivation_reactivation( + acceptor, + desktop_size, + ); framed = unsplit_tokio_framed(reader, writer); continue; } diff --git a/crates/ironrdp-web/src/session.rs b/crates/ironrdp-web/src/session.rs index 2105ae67..65251817 100644 --- a/crates/ironrdp-web/src/session.rs +++ b/crates/ironrdp-web/src/session.rs @@ -653,7 +653,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, None) + single_sequence_step_read(&mut framed, &mut *box_connection_activation, &mut buf) .await?; if written.size().is_some() { @@ -1018,7 +1018,7 @@ where // RDCleanPath response let rdcleanpath_res = framed - .read_by_hint(&RDCLEANPATH_HINT, None) + .read_by_hint(&RDCLEANPATH_HINT) .await .context("read RDCleanPath request")?;