From 1c5a20fcd50d67efc49594f96dda206c88b29b21 Mon Sep 17 00:00:00 2001
From: OuyangHang33 <hank.ouyang@intel.com>
Date: Wed, 29 May 2024 14:10:58 +0800
Subject: [PATCH] Add length check when take(rsp_length)

Signed-off-by: OuyangHang33 <hank.ouyang@intel.com>
---
 spdmlib/src/message/vendor.rs | 25 +++++++++++++++----------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/spdmlib/src/message/vendor.rs b/spdmlib/src/message/vendor.rs
index adef223..ad1ddd5 100644
--- a/spdmlib/src/message/vendor.rs
+++ b/spdmlib/src/message/vendor.rs
@@ -160,17 +160,22 @@ impl Codec for VendorDefinedRspPayloadStruct {
 
     fn read(r: &mut Reader) -> Option<VendorDefinedRspPayloadStruct> {
         let rsp_length = u16::read(r)?;
-        let mut vendor_defined_rsp_payload = [0u8; MAX_SPDM_VENDOR_DEFINED_PAYLOAD_SIZE];
-        for d in vendor_defined_rsp_payload
-            .iter_mut()
-            .take(rsp_length as usize)
-        {
-            *d = u8::read(r)?;
+        if rsp_length as usize > MAX_SPDM_VENDOR_DEFINED_PAYLOAD_SIZE {
+            log::error!("invalid rsp length!!!\n");
+            None
+        } else {
+            let mut vendor_defined_rsp_payload = [0u8; MAX_SPDM_VENDOR_DEFINED_PAYLOAD_SIZE];
+            for d in vendor_defined_rsp_payload
+                .iter_mut()
+                .take(rsp_length as usize)
+            {
+                *d = u8::read(r)?;
+            }
+            Some(VendorDefinedRspPayloadStruct {
+                rsp_length,
+                vendor_defined_rsp_payload,
+            })
         }
-        Some(VendorDefinedRspPayloadStruct {
-            rsp_length,
-            vendor_defined_rsp_payload,
-        })
     }
 }