From a95b6045c32b8ea2c564c6848f00d785056dd56f Mon Sep 17 00:00:00 2001
From: ydong10 <ydong10@6f19259b-4bc3-4df7-8a09-765794883524>
Date: Fri, 18 May 2012 02:24:25 +0000
Subject: [PATCH] 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
---
 MdeModulePkg/Core/DxeIplPeim/DxeLoad.c   | 49 +++++++++++++++++++++++-
 Nt32Pkg/WinNtAutoScanPei/WinNtAutoScan.c | 49 +++++++++++++++++++++++-
 2 files changed, 96 insertions(+), 2 deletions(-)

diff --git a/MdeModulePkg/Core/DxeIplPeim/DxeLoad.c b/MdeModulePkg/Core/DxeIplPeim/DxeLoad.c
index d0baf2eb91..dea627e162 100644
--- a/MdeModulePkg/Core/DxeIplPeim/DxeLoad.c
+++ b/MdeModulePkg/Core/DxeIplPeim/DxeLoad.c
@@ -124,6 +124,53 @@ PeimInitializeDxeIpl (
   return Status;
 }
 
+/**
+   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. 
 
@@ -214,7 +261,7 @@ DxeLoadCore (
                          &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 90a97b9767..ab5fd0f3f1 100644
--- a/Nt32Pkg/WinNtAutoScanPei/WinNtAutoScan.c
+++ b/Nt32Pkg/WinNtAutoScanPei/WinNtAutoScan.c
@@ -49,6 +49,53 @@ EFI_MEMORY_TYPE_INFORMATION mDefaultMemoryTypeInformation[] = {
   { 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 @@ Returns:
                        &DataSize,
                        &MemoryData
                        );
-  if (EFI_ERROR (Status)) {
+  if (EFI_ERROR (Status) || !ValidateMemoryTypeInfoVariable(MemoryData, DataSize)) {
     //
     // Create Memory Type Information HOB
     //