From 46b703e813a06e3272541f69bf2613138d27a11f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marc-Andr=C3=A9=20Lureau?= Date: Wed, 14 Aug 2024 18:49:48 +0400 Subject: [PATCH] refactor(pdu): return whether hint::find_size() matches the expected hint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When compiled in debug mode, the code checks the expected Action hint. But in release mode, no checks are done and the it will have to fail later. Instead, return whether the PDU is matching the hint, so the caller can decide what to do in this case. Signed-off-by: Marc-André Lureau --- crates/ironrdp-async/src/framed.rs | 2 +- crates/ironrdp-blocking/src/framed.rs | 2 +- crates/ironrdp-connector/src/credssp.rs | 8 +++---- crates/ironrdp-pdu/src/lib.rs | 21 +++++++++++-------- crates/ironrdp-web/src/session.rs | 4 ++-- .../Generated/RawPduHint.cs | 2 +- ffi/src/connector/mod.rs | 3 ++- 7 files changed, 23 insertions(+), 19 deletions(-) diff --git a/crates/ironrdp-async/src/framed.rs b/crates/ironrdp-async/src/framed.rs index e142a881..5efcb3e9 100644 --- a/crates/ironrdp-async/src/framed.rs +++ b/crates/ironrdp-async/src/framed.rs @@ -170,7 +170,7 @@ where .find_size(self.peek()) .map_err(|e| io::Error::new(io::ErrorKind::Other, e))? { - Some(length) => { + Some((_matched, length)) => { return Ok(self.read_exact(length).await?.freeze()); } None => { diff --git a/crates/ironrdp-blocking/src/framed.rs b/crates/ironrdp-blocking/src/framed.rs index 9c1bcb4f..c46109bb 100644 --- a/crates/ironrdp-blocking/src/framed.rs +++ b/crates/ironrdp-blocking/src/framed.rs @@ -93,7 +93,7 @@ where .find_size(self.peek()) .map_err(|e| io::Error::new(io::ErrorKind::Other, e))? { - Some(length) => { + Some((_matched, length)) => { return Ok(self.read_exact(length)?.freeze()); } None => { diff --git a/crates/ironrdp-connector/src/credssp.rs b/crates/ironrdp-connector/src/credssp.rs index cba1d554..1ae80cdd 100644 --- a/crates/ironrdp-connector/src/credssp.rs +++ b/crates/ironrdp-connector/src/credssp.rs @@ -43,9 +43,9 @@ struct CredsspTsRequestHint; const CREDSSP_TS_REQUEST_HINT: CredsspTsRequestHint = CredsspTsRequestHint; impl PduHint for CredsspTsRequestHint { - fn find_size(&self, bytes: &[u8]) -> ironrdp_pdu::PduResult> { + fn find_size(&self, bytes: &[u8]) -> ironrdp_pdu::PduResult> { match credssp::TsRequest::read_length(bytes) { - Ok(length) => Ok(Some(length)), + Ok(length) => Ok(Some((true, length))), Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => Ok(None), Err(e) => Err(ironrdp_pdu::custom_err!("CredsspTsRequestHint", e)), } @@ -58,8 +58,8 @@ struct CredsspEarlyUserAuthResultHint; const CREDSSP_EARLY_USER_AUTH_RESULT_HINT: CredsspEarlyUserAuthResultHint = CredsspEarlyUserAuthResultHint; impl PduHint for CredsspEarlyUserAuthResultHint { - fn find_size(&self, _: &[u8]) -> ironrdp_pdu::PduResult> { - Ok(Some(credssp::EARLY_USER_AUTH_RESULT_PDU_SIZE)) + fn find_size(&self, _: &[u8]) -> ironrdp_pdu::PduResult> { + Ok(Some((true, credssp::EARLY_USER_AUTH_RESULT_PDU_SIZE))) } } diff --git a/crates/ironrdp-pdu/src/lib.rs b/crates/ironrdp-pdu/src/lib.rs index d07bd63f..0ccafc6a 100644 --- a/crates/ironrdp-pdu/src/lib.rs +++ b/crates/ironrdp-pdu/src/lib.rs @@ -333,7 +333,10 @@ pub fn find_size(bytes: &[u8]) -> PduResult> { pub trait PduHint: Send + Sync + fmt::Debug + 'static { /// Finds next PDU size by reading the next few bytes. - fn find_size(&self, bytes: &[u8]) -> PduResult>; + /// + /// Returns `Some((hint_matching, size))` if the size is known. + /// Returns `None` if the size cannot be determined yet. + fn find_size(&self, bytes: &[u8]) -> PduResult>; } // Matches both X224 and FastPath pdus @@ -343,8 +346,8 @@ pub struct RdpHint; pub const RDP_HINT: RdpHint = RdpHint; impl PduHint for RdpHint { - fn find_size(&self, bytes: &[u8]) -> PduResult> { - find_size(bytes).map(|opt| opt.map(|info| info.length)) + fn find_size(&self, bytes: &[u8]) -> PduResult> { + find_size(bytes).map(|opt| opt.map(|info| (true, info.length))) } } @@ -354,11 +357,11 @@ pub struct X224Hint; pub const X224_HINT: X224Hint = X224Hint; impl PduHint for X224Hint { - fn find_size(&self, bytes: &[u8]) -> PduResult> { + fn find_size(&self, bytes: &[u8]) -> PduResult> { match find_size(bytes)? { Some(pdu_info) => { - debug_assert_eq!(pdu_info.action, Action::X224); - Ok(Some(pdu_info.length)) + let res = (pdu_info.action == Action::X224, pdu_info.length); + Ok(Some(res)) } None => Ok(None), } @@ -371,11 +374,11 @@ pub struct FastPathHint; pub const FAST_PATH_HINT: FastPathHint = FastPathHint; impl PduHint for FastPathHint { - fn find_size(&self, bytes: &[u8]) -> PduResult> { + fn find_size(&self, bytes: &[u8]) -> PduResult> { match find_size(bytes)? { Some(pdu_info) => { - debug_assert_eq!(pdu_info.action, Action::FastPath); - Ok(Some(pdu_info.length)) + let res = (pdu_info.action == Action::FastPath, pdu_info.length); + Ok(Some(res)) } None => Ok(None), } diff --git a/crates/ironrdp-web/src/session.rs b/crates/ironrdp-web/src/session.rs index 278080fa..628b7941 100644 --- a/crates/ironrdp-web/src/session.rs +++ b/crates/ironrdp-web/src/session.rs @@ -894,9 +894,9 @@ where const RDCLEANPATH_HINT: RDCleanPathHint = RDCleanPathHint; impl ironrdp::pdu::PduHint for RDCleanPathHint { - fn find_size(&self, bytes: &[u8]) -> ironrdp::pdu::PduResult> { + fn find_size(&self, bytes: &[u8]) -> ironrdp::pdu::PduResult> { match ironrdp_rdcleanpath::RDCleanPathPdu::detect(bytes) { - ironrdp_rdcleanpath::DetectionResult::Detected { total_length, .. } => Ok(Some(total_length)), + ironrdp_rdcleanpath::DetectionResult::Detected { total_length, .. } => Ok(Some((true, total_length))), ironrdp_rdcleanpath::DetectionResult::NotEnoughBytes => Ok(None), ironrdp_rdcleanpath::DetectionResult::Failed => Err(ironrdp::pdu::other_err!( "RDCleanPathHint", diff --git a/ffi/dotnet/Devolutions.IronRdp/Generated/RawPduHint.cs b/ffi/dotnet/Devolutions.IronRdp/Generated/RawPduHint.cs index 6252967e..35e5603d 100644 --- a/ffi/dotnet/Devolutions.IronRdp/Generated/RawPduHint.cs +++ b/ffi/dotnet/Devolutions.IronRdp/Generated/RawPduHint.cs @@ -17,7 +17,7 @@ public partial struct PduHint private const string NativeLib = "DevolutionsIronRdp"; [DllImport(NativeLib, CallingConvention = CallingConvention.Cdecl, EntryPoint = "PduHint_find_size", ExactSpelling = true)] - public static unsafe extern ConnectorFfiResultBoxOptionalUsizeBoxIronRdpError FindSize(PduHint* self, byte* bytes, nuint bytesSz); + public static unsafe extern ConnectorFfiResultBoxOptionalUsizeBoxIronRdpError FindSize(PduHint* self, byte* bytes, nuint bytesSz, bool* matched); [DllImport(NativeLib, CallingConvention = CallingConvention.Cdecl, EntryPoint = "PduHint_destroy", ExactSpelling = true)] public static unsafe extern void Destroy(PduHint* self); diff --git a/ffi/src/connector/mod.rs b/ffi/src/connector/mod.rs index ac6bc13e..b4ad9e50 100644 --- a/ffi/src/connector/mod.rs +++ b/ffi/src/connector/mod.rs @@ -160,7 +160,8 @@ pub mod ffi { impl<'a> PduHint<'a> { pub fn find_size(&'a self, bytes: &[u8]) -> Result, Box> { let pdu_hint = self.0; - let size = pdu_hint.find_size(bytes)?; + // TODO C# NuGet is only used on client-side so we probably don’t need to break the ABI for that just now. + let size = pdu_hint.find_size(bytes)?.map(|(_match, size)| size); Ok(Box::new(crate::utils::ffi::OptionalUsize(size))) } }