Add logic to validate variable before use it.

[mirror_edk2.git] / Nt32Pkg / WinNtAutoScanPei / WinNtAutoScan.c

diff --git a/Nt32Pkg/WinNtAutoScanPei/WinNtAutoScan.c b/Nt32Pkg/WinNtAutoScanPei/WinNtAutoScan.c
index 90a97b97676393035c12fcc3151e93d869b3d682..ab5fd0f3f1ac279b55358feb4dcc0f44efbdebcb 100644 (file)
--- a/Nt32Pkg/WinNtAutoScanPei/WinNtAutoScan.c
+++ b/Nt32Pkg/WinNtAutoScanPei/WinNtAutoScan.c
@@ -49,6 +49,53 @@ EFI_MEMORY_TYPE_INFORMATION mDefaultMemoryTypeInformation[] = {
   { EfiMaxMemoryType,       0      }\r
 };\r
 \r
+/**\r
+   Validate variable data for the MemoryTypeInformation. \r
+\r
+   @param MemoryData       Variable data.\r
+   @param MemoryDataSize   Variable data length.\r
+   \r
+   @return TRUE            The variable data is valid.\r
+   @return FALSE           The variable data is invalid.\r
+\r
+**/\r
+BOOLEAN\r
+ValidateMemoryTypeInfoVariable (\r
+  IN EFI_MEMORY_TYPE_INFORMATION      *MemoryData,\r
+  IN UINTN                            MemoryDataSize\r
+  )\r
+{\r
+  UINTN                       Count;\r
+  UINTN                       Index;\r
+\r
+  // Check the input parameter.\r
+  if (MemoryData == NULL) {\r
+    return FALSE;\r
+  }\r
+\r
+  // Get Count\r
+  Count = MemoryDataSize / sizeof (*MemoryData);\r
+\r
+  // Check Size\r
+  if (Count * sizeof(*MemoryData) != MemoryDataSize) {\r
+    return FALSE;\r
+  }\r
+\r
+  // Check last entry type filed.\r
+  if (MemoryData[Count - 1].Type != EfiMaxMemoryType) {\r
+    return FALSE;\r
+  }\r
+\r
+  // Check the type filed.\r
+  for (Index = 0; Index < Count - 1; Index++) {\r
+    if (MemoryData[Index].Type >= EfiMaxMemoryType) {\r
+      return FALSE;\r
+    }\r
+  }\r
+\r
+  return TRUE;\r
+}\r
+\r
 EFI_STATUS\r
 EFIAPI\r
 PeimInitializeWinNtAutoScan (\r
@@ -153,7 +200,7 @@ Returns:
                        &DataSize,\r
                        &MemoryData\r
                        );\r
-  if (EFI_ERROR (Status)) {\r
+  if (EFI_ERROR (Status) || !ValidateMemoryTypeInfoVariable(MemoryData, DataSize)) {\r
     //\r
     // Create Memory Type Information HOB\r
     //\r