Add logic to validate variable before use it.

Signed-off-by: Eric Dong <eric.dong@intel.com>
Reviewed-by: Jiewen Yao <jiewen,yao@intel.com>
Reviewed-by: Michael D. Kinney <michael.d.kinney@intel.com>
 

git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@13323 6f19259b-4bc3-4df7-8a09-765794883524
diff --git a/MdeModulePkg/Core/DxeIplPeim/DxeLoad.c b/MdeModulePkg/Core/DxeIplPeim/DxeLoad.c
index d0baf2e..dea627e 100644
--- a/MdeModulePkg/Core/DxeIplPeim/DxeLoad.c
+++ b/MdeModulePkg/Core/DxeIplPeim/DxeLoad.c
@@ -125,6 +125,53 @@
 }

 

 /**

+   Validate variable data for the MemoryTypeInformation. 

+

+   @param MemoryData       Variable data.

+   @param MemoryDataSize   Variable data length.

+

+   @return TRUE            The variable data is valid.

+   @return FALSE           The variable data is invalid.

+

+**/

+BOOLEAN

+ValidateMemoryTypeInfoVariable (

+  IN EFI_MEMORY_TYPE_INFORMATION      *MemoryData,

+  IN UINTN                            MemoryDataSize

+  )

+{

+  UINTN                       Count;

+  UINTN                       Index;

+

+  // Check the input parameter.

+  if (MemoryData == NULL) {

+    return FALSE;

+  }

+

+  // Get Count

+  Count = MemoryDataSize / sizeof (*MemoryData);

+

+  // Check Size

+  if (Count * sizeof(*MemoryData) != MemoryDataSize) {

+    return FALSE;

+  }

+

+  // Check last entry type filed.

+  if (MemoryData[Count - 1].Type != EfiMaxMemoryType) {

+    return FALSE;

+  }

+

+  // Check the type filed.

+  for (Index = 0; Index < Count - 1; Index++) {

+    if (MemoryData[Index].Type >= EfiMaxMemoryType) {

+      return FALSE;

+    }

+  }

+

+  return TRUE;

+}

+

+/**

    Main entry point to last PEIM. 

 

    This function finds DXE Core in the firmware volume and transfer the control to

@@ -214,7 +261,7 @@
                          &DataSize,

                          &MemoryData

                          );

-    if (!EFI_ERROR (Status)) {

+    if (!EFI_ERROR (Status) && ValidateMemoryTypeInfoVariable(MemoryData, DataSize)) {

       //

       // Build the GUID'd HOB for DXE

       //

diff --git a/Nt32Pkg/WinNtAutoScanPei/WinNtAutoScan.c b/Nt32Pkg/WinNtAutoScanPei/WinNtAutoScan.c
index 90a97b9..ab5fd0f 100644
--- a/Nt32Pkg/WinNtAutoScanPei/WinNtAutoScan.c
+++ b/Nt32Pkg/WinNtAutoScanPei/WinNtAutoScan.c
@@ -49,6 +49,53 @@
   { EfiMaxMemoryType,       0      }

 };

 

+/**

+   Validate variable data for the MemoryTypeInformation. 

+

+   @param MemoryData       Variable data.

+   @param MemoryDataSize   Variable data length.

+   

+   @return TRUE            The variable data is valid.

+   @return FALSE           The variable data is invalid.

+

+**/

+BOOLEAN

+ValidateMemoryTypeInfoVariable (

+  IN EFI_MEMORY_TYPE_INFORMATION      *MemoryData,

+  IN UINTN                            MemoryDataSize

+  )

+{

+  UINTN                       Count;

+  UINTN                       Index;

+

+  // Check the input parameter.

+  if (MemoryData == NULL) {

+    return FALSE;

+  }

+

+  // Get Count

+  Count = MemoryDataSize / sizeof (*MemoryData);

+

+  // Check Size

+  if (Count * sizeof(*MemoryData) != MemoryDataSize) {

+    return FALSE;

+  }

+

+  // Check last entry type filed.

+  if (MemoryData[Count - 1].Type != EfiMaxMemoryType) {

+    return FALSE;

+  }

+

+  // Check the type filed.

+  for (Index = 0; Index < Count - 1; Index++) {

+    if (MemoryData[Index].Type >= EfiMaxMemoryType) {

+      return FALSE;

+    }

+  }

+

+  return TRUE;

+}

+

 EFI_STATUS

 EFIAPI

 PeimInitializeWinNtAutoScan (

@@ -153,7 +200,7 @@
                        &DataSize,

                        &MemoryData

                        );

-  if (EFI_ERROR (Status)) {

+  if (EFI_ERROR (Status) || !ValidateMemoryTypeInfoVariable(MemoryData, DataSize)) {

     //

     // Create Memory Type Information HOB

     //