From 99e8b6437609b2e1ab4b0caada9d09f365be947c Mon Sep 17 00:00:00 2001
From: Vladyslav Nikonov <mail@pacmancoder.xyz>
Date: Thu, 21 Mar 2024 18:33:09 +0200
Subject: [PATCH] fix(fuzz): fix failing fuzzing after pdu encode/decode
 refactoring

---
 crates/ironrdp-fuzzing/src/oracles/mod.rs                   | 6 ++++++
 crates/ironrdp-pdu/src/ber.rs                               | 4 ++++
 crates/ironrdp-pdu/src/gcc/security_data.rs                 | 1 +
 .../server_license/client_platform_challenge_response.rs    | 1 +
 .../src/rdp/server_license/server_license_request/cert.rs   | 1 +
 5 files changed, 13 insertions(+)

diff --git a/crates/ironrdp-fuzzing/src/oracles/mod.rs b/crates/ironrdp-fuzzing/src/oracles/mod.rs
index 179e8ef5f..0b4cc69ea 100644
--- a/crates/ironrdp-fuzzing/src/oracles/mod.rs
+++ b/crates/ironrdp-fuzzing/src/oracles/mod.rs
@@ -157,3 +157,9 @@ pub fn channel_process(input: &[u8]) {
 
     let _ = rdpdr.process(input);
 }
+
+#[test]
+fn test_pdu_decode() {
+    const DATA: &'static [u8] = &[130, 14, 239, 6, 21, 130, 0, 48, 9, 0, 1, 0, 77, 9, 0, 1, 0, 0, 0, 0, 42, 0, 0, 2, 0, 63, 0, 16, 241, 241, 241, 241, 0, 0, 50, 241, 4, 0, 0, 0, 47, 16, 0, 241, 4, 0, 0, 0, 47, 16, 0, 17, 13, 0, 0, 0, 1, 0, 0, 0, 3, 0, 0, 15, 2, 0, 0, 0];
+    pdu_decode(DATA);
+}
\ No newline at end of file
diff --git a/crates/ironrdp-pdu/src/ber.rs b/crates/ironrdp-pdu/src/ber.rs
index 838084ce9..40214e411 100644
--- a/crates/ironrdp-pdu/src/ber.rs
+++ b/crates/ironrdp-pdu/src/ber.rs
@@ -133,6 +133,7 @@ pub(crate) fn read_enumerated(stream: &mut ReadCursor<'_>, count: u8) -> PduResu
         return Err(invalid_message_err!("len", "invalid enumerated len"));
     }
 
+    ensure_size!(in: stream, size: 1);
     let enumerated = stream.read_u8();
     if enumerated == u8::MAX || enumerated + 1 > count {
         return Err(invalid_message_err!("enumerated", "invalid enumerated value"));
@@ -203,6 +204,8 @@ pub(crate) fn write_bool(stream: &mut WriteCursor<'_>, value: bool) -> PduResult
     let mut size = 0;
     size += write_universal_tag(stream, Tag::Boolean, Pc::Primitive)?;
     size += write_length(stream, 1)?;
+
+    ensure_size!(in: stream, size: 1);
     stream.write_u8(if value { 0xFF } else { 0x00 });
     size += 1;
 
@@ -217,6 +220,7 @@ pub(crate) fn read_bool(stream: &mut ReadCursor<'_>) -> PduResult<bool> {
         return Err(invalid_message_err!("len", "invalid integer len"));
     }
 
+    ensure_size!(in: stream, size: 1);
     Ok(stream.read_u8() != 0)
 }
 
diff --git a/crates/ironrdp-pdu/src/gcc/security_data.rs b/crates/ironrdp-pdu/src/gcc/security_data.rs
index 278dbfad1..2fd3eb350 100644
--- a/crates/ironrdp-pdu/src/gcc/security_data.rs
+++ b/crates/ironrdp-pdu/src/gcc/security_data.rs
@@ -171,6 +171,7 @@ impl<'de> PduDecode<'de> for ServerSecurityData {
                 ));
             }
 
+            ensure_size!(in: src, size: SERVER_RANDOM_LEN);
             let server_random = src.read_array();
 
             ensure_size!(in: src, size: server_cert_len);
diff --git a/crates/ironrdp-pdu/src/rdp/server_license/client_platform_challenge_response.rs b/crates/ironrdp-pdu/src/rdp/server_license/client_platform_challenge_response.rs
index 02dd9fb6a..4c22c14ef 100644
--- a/crates/ironrdp-pdu/src/rdp/server_license/client_platform_challenge_response.rs
+++ b/crates/ironrdp-pdu/src/rdp/server_license/client_platform_challenge_response.rs
@@ -152,6 +152,7 @@ impl<'de> PduDecode<'de> for ClientPlatformChallengeResponse {
         ensure_size!(in: src, size: encrypted_hwid_blob.length);
         let encrypted_hwid = src.read_slice(encrypted_hwid_blob.length).into();
 
+        ensure_size!(in: src, size: MAC_SIZE);
         let mac_data = src.read_slice(MAC_SIZE).into();
 
         Ok(Self {
diff --git a/crates/ironrdp-pdu/src/rdp/server_license/server_license_request/cert.rs b/crates/ironrdp-pdu/src/rdp/server_license/server_license_request/cert.rs
index 21ef332c6..82e4746cf 100644
--- a/crates/ironrdp-pdu/src/rdp/server_license/server_license_request/cert.rs
+++ b/crates/ironrdp-pdu/src/rdp/server_license/server_license_request/cert.rs
@@ -69,6 +69,7 @@ impl PduEncode for X509CertificateChain {
 
 impl<'de> PduDecode<'de> for X509CertificateChain {
     fn decode(src: &mut ReadCursor<'de>) -> PduResult<Self> {
+        ensure_size!(in: src, size: 4);
         let certificate_count = cast_length!("certArrayLen", src.read_u32())?;
         if !(MIN_CERTIFICATE_AMOUNT..MAX_CERTIFICATE_AMOUNT).contains(&certificate_count) {
             return Err(invalid_message_err!("certArrayLen", "invalid x509 certificate amount"));