From 2918cc8f6e594b7afc5c0304f7da057c69b08528 Mon Sep 17 00:00:00 2001
From: "Douglas Flick [MSFT]" <doug.edk2@gmail.com>
Date: Fri, 12 Jan 2024 02:16:02 +0800
Subject: [PATCH 02/10] SecurityPkg: DxeTpmMeasureBootLib: SECURITY PATCH 4117
 - CVE 2022-36763

This commit contains the patch files and tests for DxeTpmMeasureBootLib
CVE 2022-36763.

Cc: Jiewen Yao <jiewen.yao@intel.com>

Signed-off-by: Doug Flick [MSFT] <doug.edk2@gmail.com>
Reviewed-by: Jiewen Yao <jiewen.yao@intel.com>

reference: https://github.com/tianocore/edk2/pull/5264
Signed-off-by: yexiao <yexiao7@huawei.com>
---
 .../DxeTpmMeasureBootLib.c                    |  40 ++-
 .../DxeTpmMeasureBootLib.inf                  |   4 +-
 .../DxeTpmMeasureBootLibSanitization.c        | 241 ++++++++++++++
 .../DxeTpmMeasureBootLibSanitization.h        | 114 +++++++
 .../DxeTpmMeasureBootLibSanitizationTest.c    | 301 ++++++++++++++++++
 ...eTpmMeasureBootLibSanitizationTestHost.inf |  28 ++
 SecurityPkg/SecurityPkg.ci.yaml               |   1 +
 SecurityPkg/Test/SecurityPkgHostTest.dsc      |   1 +
 8 files changed, 716 insertions(+), 14 deletions(-)
 create mode 100644 SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLibSanitization.c
 create mode 100644 SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLibSanitization.h
 create mode 100644 SecurityPkg/Library/DxeTpmMeasureBootLib/InternalUnitTest/DxeTpmMeasureBootLibSanitizationTest.c
 create mode 100644 SecurityPkg/Library/DxeTpmMeasureBootLib/InternalUnitTest/DxeTpmMeasureBootLibSanitizationTestHost.inf

diff --git a/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLib.c b/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLib.c
index 220393dd..f31b49f8 100644
--- a/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLib.c
+++ b/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLib.c
@@ -18,6 +18,8 @@
 Copyright (c) 2009 - 2018, Intel Corporation. All rights reserved.<BR>
 SPDX-License-Identifier: BSD-2-Clause-Patent
 
+Copyright (c) Microsoft Corporation.<BR>
+SPDX-License-Identifier: BSD-2-Clause-Patent
 **/
 
 #include <PiDxe.h>
@@ -40,6 +42,8 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Library/SecurityManagementLib.h>
 #include <Library/HobLib.h>
 
+#include "DxeTpmMeasureBootLibSanitization.h"
+
 //
 // Flag to check GPT partition. It only need be measured once.
 //
@@ -136,6 +140,9 @@ TcgMeasureGptTable (
   UINT32                      EventSize;
   UINT32                      EventNumber;
   EFI_PHYSICAL_ADDRESS        EventLogLastEntry;
+  UINT32                      AllocSize;
+
+  GptData = NULL;
 
   if (mMeasureGptCount > 0) {
     return EFI_SUCCESS;
@@ -166,8 +173,8 @@ TcgMeasureGptTable (
                      BlockIo->Media->BlockSize,
                      (UINT8 *)PrimaryHeader
                      );
-  if (EFI_ERROR (Status)) {
-    DEBUG ((DEBUG_ERROR, "Failed to Read Partition Table Header!\n"));
+  if (EFI_ERROR (Status) || EFI_ERROR (SanitizeEfiPartitionTableHeader (PrimaryHeader, BlockIo))) {
+    DEBUG ((DEBUG_ERROR, "Failed to read Partition Table Header or invalid Partition Table Header!\n"));
     FreePool (PrimaryHeader);
     return EFI_DEVICE_ERROR;
   }
@@ -175,7 +182,13 @@ TcgMeasureGptTable (
   //
   // Read the partition entry.
   //
-  EntryPtr = (UINT8 *)AllocatePool (PrimaryHeader->NumberOfPartitionEntries * PrimaryHeader->SizeOfPartitionEntry);
+  Status = SanitizePrimaryHeaderAllocationSize (PrimaryHeader, &AllocSize);
+  if (EFI_ERROR (Status)) {
+    FreePool (PrimaryHeader);
+    return EFI_DEVICE_ERROR;
+  }
+
+  EntryPtr = (UINT8 *)AllocatePool (AllocSize);
   if (EntryPtr == NULL) {
     FreePool (PrimaryHeader);
     return EFI_OUT_OF_RESOURCES;
@@ -185,7 +198,7 @@ TcgMeasureGptTable (
                      DiskIo,
                      BlockIo->Media->MediaId,
                      MultU64x32 (PrimaryHeader->PartitionEntryLBA, BlockIo->Media->BlockSize),
-                     PrimaryHeader->NumberOfPartitionEntries * PrimaryHeader->SizeOfPartitionEntry,
+                     AllocSize,
                      EntryPtr
                      );
   if (EFI_ERROR (Status)) {
@@ -210,9 +223,8 @@ TcgMeasureGptTable (
   //
   // Prepare Data for Measurement
   //
-  EventSize = (UINT32)(sizeof (EFI_GPT_DATA) - sizeof (GptData->Partitions)
-                       + NumberOfPartition * PrimaryHeader->SizeOfPartitionEntry);
-  TcgEvent = (TCG_PCR_EVENT *)AllocateZeroPool (EventSize + sizeof (TCG_PCR_EVENT_HDR));
+  Status   = SanitizePrimaryHeaderGptEventSize (PrimaryHeader, NumberOfPartition, &EventSize);
+  TcgEvent = (TCG_PCR_EVENT *)AllocateZeroPool (EventSize);
   if (TcgEvent == NULL) {
     FreePool (PrimaryHeader);
     FreePool (EntryPtr);
@@ -221,7 +233,7 @@ TcgMeasureGptTable (
 
   TcgEvent->PCRIndex  = 5;
   TcgEvent->EventType = EV_EFI_GPT_EVENT;
-  TcgEvent->EventSize = EventSize;
+  TcgEvent->EventSize = EventSize - sizeof (TCG_PCR_EVENT_HDR);
   GptData             = (EFI_GPT_DATA *)TcgEvent->Event;
 
   //
@@ -361,11 +373,13 @@ TcgMeasurePeImage (
       TcgEvent->PCRIndex  = 2;
       break;
     default:
-      DEBUG ((
-        DEBUG_ERROR,
-        "TcgMeasurePeImage: Unknown subsystem type %d",
-        ImageType
-        ));
+      DEBUG (
+        (
+         DEBUG_ERROR,
+         "TcgMeasurePeImage: Unknown subsystem type %d",
+         ImageType
+        )
+        );
       goto Finish;
   }
 
diff --git a/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLib.inf b/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLib.inf
index ebab6f7c..21bfd570 100644
--- a/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLib.inf
+++ b/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLib.inf
@@ -32,6 +32,8 @@
 
 [Sources]
   DxeTpmMeasureBootLib.c
+  DxeTpmMeasureBootLibSanitization.c
+  DxeTpmMeasureBootLibSanitization.h
 
 [Packages]
   MdePkg/MdePkg.dec
@@ -41,6 +43,7 @@
 
 [LibraryClasses]
   BaseMemoryLib
+  SafeIntLib
   DebugLib
   MemoryAllocationLib
   DevicePathLib
@@ -59,4 +62,3 @@
   gEfiFirmwareVolumeBlockProtocolGuid   ## SOMETIMES_CONSUMES
   gEfiBlockIoProtocolGuid               ## SOMETIMES_CONSUMES
   gEfiDiskIoProtocolGuid                ## SOMETIMES_CONSUMES
-
diff --git a/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLibSanitization.c b/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLibSanitization.c
new file mode 100644
index 00000000..e009ce46
--- /dev/null
+++ b/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLibSanitization.c
@@ -0,0 +1,241 @@
+/** @file
+  The library instance provides security service of TPM2 measure boot and
+  Confidential Computing (CC) measure boot.
+
+  Caution: This file requires additional review when modified.
+  This library will have external input - PE/COFF image and GPT partition.
+  This external input must be validated carefully to avoid security issue like
+  buffer overflow, integer overflow.
+
+  This file will pull out the validation logic from the following functions, in an
+  attempt to validate the untrusted input in the form of unit tests
+
+  These are those functions:
+
+  DxeTpmMeasureBootLibImageRead() function will make sure the PE/COFF image content
+  read is within the image buffer.
+
+  Tcg2MeasureGptTable() function will receive untrusted GPT partition table, and parse
+  partition data carefully.
+
+  Copyright (c) Microsoft Corporation.<BR>
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+**/
+#include <Uefi.h>
+#include <Uefi/UefiSpec.h>
+#include <Library/SafeIntLib.h>
+#include <Library/UefiLib.h>
+#include <Library/DebugLib.h>
+#include <Library/BaseLib.h>
+#include <IndustryStandard/UefiTcgPlatform.h>
+#include <Protocol/BlockIo.h>
+#include <Library/MemoryAllocationLib.h>
+
+#include "DxeTpmMeasureBootLibSanitization.h"
+
+#define GPT_HEADER_REVISION_V1  0x00010000
+
+/**
+  This function will validate the EFI_PARTITION_TABLE_HEADER structure is safe to parse
+  However this function will not attempt to verify the validity of the GPT partition
+  It will check the following:
+    - Signature
+    - Revision
+    - AlternateLBA
+    - FirstUsableLBA
+    - LastUsableLBA
+    - PartitionEntryLBA
+    - NumberOfPartitionEntries
+    - SizeOfPartitionEntry
+    - BlockIo
+
+  @param[in] PrimaryHeader
+    Pointer to the EFI_PARTITION_TABLE_HEADER structure.
+
+  @param[in] BlockIo
+    Pointer to the EFI_BLOCK_IO_PROTOCOL structure.
+
+  @retval EFI_SUCCESS
+    The EFI_PARTITION_TABLE_HEADER structure is valid.
+
+  @retval EFI_INVALID_PARAMETER
+    The EFI_PARTITION_TABLE_HEADER structure is invalid.
+**/
+EFI_STATUS
+EFIAPI
+SanitizeEfiPartitionTableHeader (
+  IN CONST EFI_PARTITION_TABLE_HEADER  *PrimaryHeader,
+  IN CONST EFI_BLOCK_IO_PROTOCOL       *BlockIo
+  )
+{
+  // Verify that the input parameters are safe to use
+  if (PrimaryHeader == NULL) {
+    DEBUG ((DEBUG_ERROR, "Invalid Partition Table Header!\n"));
+    return EFI_INVALID_PARAMETER;
+  }
+
+  if ((BlockIo == NULL) || (BlockIo->Media == NULL)) {
+    DEBUG ((DEBUG_ERROR, "Invalid BlockIo!\n"));
+    return EFI_INVALID_PARAMETER;
+  }
+
+  // The signature must be EFI_PTAB_HEADER_ID ("EFI PART" in ASCII)
+  if (PrimaryHeader->Header.Signature != EFI_PTAB_HEADER_ID) {
+    DEBUG ((DEBUG_ERROR, "Invalid Partition Table Header!\n"));
+    return EFI_DEVICE_ERROR;
+  }
+
+  // The version must be GPT_HEADER_REVISION_V1 (0x00010000)
+  if (PrimaryHeader->Header.Revision != GPT_HEADER_REVISION_V1) {
+    DEBUG ((DEBUG_ERROR, "Invalid Partition Table Header Revision!\n"));
+    return EFI_DEVICE_ERROR;
+  }
+
+  // The HeaderSize must be greater than or equal to 92 and must be less than or equal to the logical block size
+  if ((PrimaryHeader->Header.HeaderSize < sizeof (EFI_PARTITION_TABLE_HEADER)) || (PrimaryHeader->Header.HeaderSize > BlockIo->Media->BlockSize)) {
+    DEBUG ((DEBUG_ERROR, "Invalid Partition Table Header HeaderSize!\n"));
+    return EFI_DEVICE_ERROR;
+  }
+
+  // check that the PartitionEntryLBA greater than the Max LBA
+  // This will be used later for multiplication
+  if (PrimaryHeader->PartitionEntryLBA > DivU64x32 (MAX_UINT64, BlockIo->Media->BlockSize)) {
+    DEBUG ((DEBUG_ERROR, "Invalid Partition Table Header PartitionEntryLBA!\n"));
+    return EFI_DEVICE_ERROR;
+  }
+
+  // Check that the number of partition entries is greater than zero
+  if (PrimaryHeader->NumberOfPartitionEntries == 0) {
+    DEBUG ((DEBUG_ERROR, "Invalid Partition Table Header NumberOfPartitionEntries!\n"));
+    return EFI_DEVICE_ERROR;
+  }
+
+  // SizeOfPartitionEntry must be 128, 256, 512... improper size may lead to accessing uninitialized memory
+  if ((PrimaryHeader->SizeOfPartitionEntry < 128) || ((PrimaryHeader->SizeOfPartitionEntry & (PrimaryHeader->SizeOfPartitionEntry - 1)) != 0)) {
+    DEBUG ((DEBUG_ERROR, "SizeOfPartitionEntry shall be set to a value of 128 x 2^n where n is an integer greater than or equal to zero (e.g., 128, 256, 512, etc.)!\n"));
+    return EFI_DEVICE_ERROR;
+  }
+
+  // This check is to prevent overflow when calculating the allocation size for the partition entries
+  // This check will be used later for multiplication
+  if (PrimaryHeader->NumberOfPartitionEntries > DivU64x32 (MAX_UINT64, PrimaryHeader->SizeOfPartitionEntry)) {
+    DEBUG ((DEBUG_ERROR, "Invalid Partition Table Header NumberOfPartitionEntries!\n"));
+    return EFI_DEVICE_ERROR;
+  }
+
+  return EFI_SUCCESS;
+}
+
+/**
+  This function will validate that the allocation size from the primary header is sane
+  It will check the following:
+    - AllocationSize does not overflow
+
+  @param[in] PrimaryHeader
+    Pointer to the EFI_PARTITION_TABLE_HEADER structure.
+
+  @param[out] AllocationSize
+    Pointer to the allocation size.
+
+  @retval EFI_SUCCESS
+    The allocation size is valid.
+
+  @retval EFI_OUT_OF_RESOURCES
+    The allocation size is invalid.
+**/
+EFI_STATUS
+EFIAPI
+SanitizePrimaryHeaderAllocationSize (
+  IN CONST EFI_PARTITION_TABLE_HEADER  *PrimaryHeader,
+  OUT UINT32                           *AllocationSize
+  )
+{
+  EFI_STATUS  Status;
+
+  if (PrimaryHeader == NULL) {
+    return EFI_INVALID_PARAMETER;
+  }
+
+  if (AllocationSize == NULL) {
+    return EFI_INVALID_PARAMETER;
+  }
+
+  // Replacing logic:
+  // PrimaryHeader->NumberOfPartitionEntries * PrimaryHeader->SizeOfPartitionEntry;
+  Status = SafeUint32Mult (PrimaryHeader->NumberOfPartitionEntries, PrimaryHeader->SizeOfPartitionEntry, AllocationSize);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "Allocation Size would have overflowed!\n"));
+    return EFI_BAD_BUFFER_SIZE;
+  }
+
+  return EFI_SUCCESS;
+}
+
+/**
+  This function will validate that the Gpt Event Size calculated from the primary header is sane
+  It will check the following:
+    - EventSize does not overflow
+
+  Important: This function includes the entire length of the allocated space, including the
+  TCG_PCR_EVENT_HDR. When hashing the buffer allocated with this size, the caller must subtract
+  the size of the TCG_PCR_EVENT_HDR from the size of the buffer before hashing.
+
+  @param[in] PrimaryHeader - Pointer to the EFI_PARTITION_TABLE_HEADER structure.
+  @param[in] NumberOfPartition - Number of partitions.
+  @param[out] EventSize - Pointer to the event size.
+
+  @retval EFI_SUCCESS
+    The event size is valid.
+
+  @retval EFI_OUT_OF_RESOURCES
+    Overflow would have occurred.
+
+  @retval EFI_INVALID_PARAMETER
+    One of the passed parameters was invalid.
+**/
+EFI_STATUS
+SanitizePrimaryHeaderGptEventSize (
+  IN  CONST EFI_PARTITION_TABLE_HEADER  *PrimaryHeader,
+  IN  UINTN                             NumberOfPartition,
+  OUT UINT32                            *EventSize
+  )
+{
+  EFI_STATUS  Status;
+  UINT32      SafeNumberOfPartitions;
+
+  if (PrimaryHeader == NULL) {
+    return EFI_INVALID_PARAMETER;
+  }
+
+  if (EventSize == NULL) {
+    return EFI_INVALID_PARAMETER;
+  }
+
+  // We shouldn't even attempt to perform the multiplication if the number of partitions is greater than the maximum value of UINT32
+  Status = SafeUintnToUint32 (NumberOfPartition, &SafeNumberOfPartitions);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "NumberOfPartition would have overflowed!\n"));
+    return EFI_INVALID_PARAMETER;
+  }
+
+  // Replacing logic:
+  // (UINT32)(sizeof (EFI_GPT_DATA) - sizeof (GptData->Partitions) + NumberOfPartition * PrimaryHeader.SizeOfPartitionEntry + sizeof (TCG_PCR_EVENT_HDR));
+  Status = SafeUint32Mult (SafeNumberOfPartitions, PrimaryHeader->SizeOfPartitionEntry, EventSize);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "Event Size would have overflowed!\n"));
+    return EFI_BAD_BUFFER_SIZE;
+  }
+
+  Status = SafeUint32Add (
+             sizeof (TCG_PCR_EVENT_HDR) +
+             OFFSET_OF (EFI_GPT_DATA, Partitions),
+             *EventSize,
+             EventSize
+             );
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "Event Size would have overflowed because of GPTData!\n"));
+    return EFI_BAD_BUFFER_SIZE;
+  }
+
+  return EFI_SUCCESS;
+}
diff --git a/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLibSanitization.h b/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLibSanitization.h
new file mode 100644
index 00000000..57e05e48
--- /dev/null
+++ b/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLibSanitization.h
@@ -0,0 +1,114 @@
+/** @file
+  This file includes the function prototypes for the sanitization functions.
+
+  These are those functions:
+
+  DxeTpmMeasureBootLibImageRead() function will make sure the PE/COFF image content
+  read is within the image buffer.
+
+  TcgMeasurePeImage() function will accept untrusted PE/COFF image and validate its
+  data structure within this image buffer before use.
+
+  TcgMeasureGptTable() function will receive untrusted GPT partition table, and parse
+  partition data carefully.
+
+  Copyright (c) Microsoft Corporation.<BR>
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#ifndef DXE_TPM_MEASURE_BOOT_LIB_VALIDATION_
+#define DXE_TPM_MEASURE_BOOT_LIB_VALIDATION_
+
+#include <Uefi.h>
+#include <Uefi/UefiSpec.h>
+#include <Protocol/BlockIo.h>
+#include <IndustryStandard/UefiTcgPlatform.h>
+
+/**
+  This function will validate the EFI_PARTITION_TABLE_HEADER structure is safe to parse
+  However this function will not attempt to verify the validity of the GPT partition
+  It will check the following:
+    - Signature
+    - Revision
+    - AlternateLBA
+    - FirstUsableLBA
+    - LastUsableLBA
+    - PartitionEntryLBA
+    - NumberOfPartitionEntries
+    - SizeOfPartitionEntry
+    - BlockIo
+
+  @param[in] PrimaryHeader
+    Pointer to the EFI_PARTITION_TABLE_HEADER structure.
+
+  @param[in] BlockIo
+    Pointer to the EFI_BLOCK_IO_PROTOCOL structure.
+
+  @retval EFI_SUCCESS
+    The EFI_PARTITION_TABLE_HEADER structure is valid.
+
+  @retval EFI_INVALID_PARAMETER
+    The EFI_PARTITION_TABLE_HEADER structure is invalid.
+**/
+EFI_STATUS
+EFIAPI
+SanitizeEfiPartitionTableHeader (
+  IN CONST EFI_PARTITION_TABLE_HEADER  *PrimaryHeader,
+  IN CONST EFI_BLOCK_IO_PROTOCOL       *BlockIo
+  );
+
+/**
+  This function will validate that the allocation size from the primary header is sane
+  It will check the following:
+    - AllocationSize does not overflow
+
+  @param[in] PrimaryHeader
+    Pointer to the EFI_PARTITION_TABLE_HEADER structure.
+
+  @param[out] AllocationSize
+    Pointer to the allocation size.
+
+  @retval EFI_SUCCESS
+    The allocation size is valid.
+
+  @retval EFI_OUT_OF_RESOURCES
+    The allocation size is invalid.
+**/
+EFI_STATUS
+EFIAPI
+SanitizePrimaryHeaderAllocationSize (
+  IN CONST EFI_PARTITION_TABLE_HEADER  *PrimaryHeader,
+  OUT UINT32                           *AllocationSize
+  );
+
+/**
+  This function will validate that the Gpt Event Size calculated from the primary header is sane
+  It will check the following:
+    - EventSize does not overflow
+
+  Important: This function includes the entire length of the allocated space, including the
+  TCG_PCR_EVENT_HDR. When hashing the buffer allocated with this size, the caller must subtract
+  the size of the TCG_PCR_EVENT_HDR from the size of the buffer before hashing.
+
+  @param[in] PrimaryHeader - Pointer to the EFI_PARTITION_TABLE_HEADER structure.
+  @param[in] NumberOfPartition - Number of partitions.
+  @param[out] EventSize - Pointer to the event size.
+
+  @retval EFI_SUCCESS
+    The event size is valid.
+
+  @retval EFI_OUT_OF_RESOURCES
+    Overflow would have occurred.
+
+  @retval EFI_INVALID_PARAMETER
+    One of the passed parameters was invalid.
+**/
+EFI_STATUS
+SanitizePrimaryHeaderGptEventSize (
+  IN  CONST EFI_PARTITION_TABLE_HEADER  *PrimaryHeader,
+  IN  UINTN                             NumberOfPartition,
+  OUT UINT32                            *EventSize
+  );
+
+#endif // DXE_TPM_MEASURE_BOOT_LIB_VALIDATION_
diff --git a/SecurityPkg/Library/DxeTpmMeasureBootLib/InternalUnitTest/DxeTpmMeasureBootLibSanitizationTest.c b/SecurityPkg/Library/DxeTpmMeasureBootLib/InternalUnitTest/DxeTpmMeasureBootLibSanitizationTest.c
new file mode 100644
index 00000000..8cd97480
--- /dev/null
+++ b/SecurityPkg/Library/DxeTpmMeasureBootLib/InternalUnitTest/DxeTpmMeasureBootLibSanitizationTest.c
@@ -0,0 +1,301 @@
+/** @file
+This file includes the unit test cases for the DxeTpmMeasureBootLibSanitizationTest.c.
+
+Copyright (c) Microsoft Corporation.<BR>
+SPDX-License-Identifier: BSD-2-Clause-Patent
+**/
+
+#include <Uefi.h>
+#include <Library/UefiLib.h>
+#include <Library/DebugLib.h>
+#include <Library/UnitTestLib.h>
+#include <Protocol/BlockIo.h>
+#include <Library/MemoryAllocationLib.h>
+#include <Library/BaseMemoryLib.h>
+#include <IndustryStandard/UefiTcgPlatform.h>
+
+#include "../DxeTpmMeasureBootLibSanitization.h"
+
+#define UNIT_TEST_NAME     "DxeTpmMeasureBootLibSanitizationTest"
+#define UNIT_TEST_VERSION  "1.0"
+
+#define DEFAULT_PRIMARY_TABLE_HEADER_REVISION                     0x00010000
+#define DEFAULT_PRIMARY_TABLE_HEADER_NUMBER_OF_PARTITION_ENTRIES  1
+#define DEFAULT_PRIMARY_TABLE_HEADER_SIZE_OF_PARTITION_ENTRY      128
+
+/**
+  This function tests the SanitizeEfiPartitionTableHeader function.
+  It's intent is to test that a malicious EFI_PARTITION_TABLE_HEADER
+  structure will not cause undefined or unexpected behavior.
+
+  In general the TPM should still be able to measure the data, but
+  be the header should be sanitized to prevent any unexpected behavior.
+
+  @param[in] Context  The unit test context.
+
+  @retval UNIT_TEST_PASSED  The test passed.
+  @retval UNIT_TEST_ERROR_TEST_FAILED  The test failed.
+**/
+UNIT_TEST_STATUS
+EFIAPI
+TestSanitizeEfiPartitionTableHeader (
+  IN UNIT_TEST_CONTEXT  Context
+  )
+{
+  EFI_STATUS                  Status;
+  EFI_PARTITION_TABLE_HEADER  PrimaryHeader;
+  EFI_BLOCK_IO_PROTOCOL       BlockIo;
+  EFI_BLOCK_IO_MEDIA          BlockMedia;
+
+  // Generate EFI_BLOCK_IO_MEDIA test data
+  BlockMedia.MediaId          = 1;
+  BlockMedia.RemovableMedia   = FALSE;
+  BlockMedia.MediaPresent     = TRUE;
+  BlockMedia.LogicalPartition = FALSE;
+  BlockMedia.ReadOnly         = FALSE;
+  BlockMedia.WriteCaching     = FALSE;
+  BlockMedia.BlockSize        = 512;
+  BlockMedia.IoAlign          = 1;
+  BlockMedia.LastBlock        = 0;
+
+  // Generate EFI_BLOCK_IO_PROTOCOL test data
+  BlockIo.Revision    = 1;
+  BlockIo.Media       = &BlockMedia;
+  BlockIo.Reset       = NULL;
+  BlockIo.ReadBlocks  = NULL;
+  BlockIo.WriteBlocks = NULL;
+  BlockIo.FlushBlocks = NULL;
+
+  // Geneate EFI_PARTITION_TABLE_HEADER test data
+  PrimaryHeader.Header.Signature         = EFI_PTAB_HEADER_ID;
+  PrimaryHeader.Header.Revision          = DEFAULT_PRIMARY_TABLE_HEADER_REVISION;
+  PrimaryHeader.Header.HeaderSize        = sizeof (EFI_PARTITION_TABLE_HEADER);
+  PrimaryHeader.MyLBA                    = 1;
+  PrimaryHeader.AlternateLBA             = 2;
+  PrimaryHeader.FirstUsableLBA           = 3;
+  PrimaryHeader.LastUsableLBA            = 4;
+  PrimaryHeader.PartitionEntryLBA        = 5;
+  PrimaryHeader.NumberOfPartitionEntries = DEFAULT_PRIMARY_TABLE_HEADER_NUMBER_OF_PARTITION_ENTRIES;
+  PrimaryHeader.SizeOfPartitionEntry     = DEFAULT_PRIMARY_TABLE_HEADER_SIZE_OF_PARTITION_ENTRY;
+  PrimaryHeader.PartitionEntryArrayCRC32 = 0; // Purposely invalid
+
+  // Calculate the CRC32 of the PrimaryHeader
+  PrimaryHeader.Header.CRC32 = CalculateCrc32 ((UINT8 *)&PrimaryHeader, PrimaryHeader.Header.HeaderSize);
+
+  // Test that a normal PrimaryHeader passes validation
+  Status = SanitizeEfiPartitionTableHeader (&PrimaryHeader, &BlockIo);
+  UT_ASSERT_NOT_EFI_ERROR (Status);
+
+  // Test that when number of partition entries is 0, the function returns EFI_DEVICE_ERROR
+  // Should print "Invalid Partition Table Header NumberOfPartitionEntries!""
+  PrimaryHeader.NumberOfPartitionEntries = 0;
+  Status                                 = SanitizeEfiPartitionTableHeader (&PrimaryHeader, &BlockIo);
+  UT_ASSERT_EQUAL (Status, EFI_DEVICE_ERROR);
+  PrimaryHeader.NumberOfPartitionEntries = DEFAULT_PRIMARY_TABLE_HEADER_SIZE_OF_PARTITION_ENTRY;
+
+  // Test that when the header size is too small, the function returns EFI_DEVICE_ERROR
+  // Should print "Invalid Partition Table Header Size!"
+  PrimaryHeader.Header.HeaderSize = 0;
+  Status                          = SanitizeEfiPartitionTableHeader (&PrimaryHeader, &BlockIo);
+  UT_ASSERT_EQUAL (Status, EFI_DEVICE_ERROR);
+  PrimaryHeader.Header.HeaderSize = sizeof (EFI_PARTITION_TABLE_HEADER);
+
+  // Test that when the SizeOfPartitionEntry is too small, the function returns EFI_DEVICE_ERROR
+  // should print: "SizeOfPartitionEntry shall be set to a value of 128 x 2^n where n is an integer greater than or equal to zero (e.g., 128, 256, 512, etc.)!"
+  PrimaryHeader.SizeOfPartitionEntry = 1;
+  Status                             = SanitizeEfiPartitionTableHeader (&PrimaryHeader, &BlockIo);
+  UT_ASSERT_EQUAL (Status, EFI_DEVICE_ERROR);
+
+  DEBUG ((DEBUG_INFO, "%a: Test passed\n", __func__));
+
+  return UNIT_TEST_PASSED;
+}
+
+/**
+  This function tests the SanitizePrimaryHeaderAllocationSize function.
+  It's intent is to test that the untrusted input from a EFI_PARTITION_TABLE_HEADER
+  structure will not cause an overflow when calculating the allocation size.
+
+  @param[in] Context  The unit test context.
+
+  @retval UNIT_TEST_PASSED  The test passed.
+  @retval UNIT_TEST_ERROR_TEST_FAILED  The test failed.
+**/
+UNIT_TEST_STATUS
+EFIAPI
+TestSanitizePrimaryHeaderAllocationSize (
+  IN UNIT_TEST_CONTEXT  Context
+  )
+{
+  UINT32  AllocationSize;
+
+  EFI_STATUS                  Status;
+  EFI_PARTITION_TABLE_HEADER  PrimaryHeader;
+
+  // Test that a normal PrimaryHeader passes validation
+  PrimaryHeader.NumberOfPartitionEntries = 5;
+  PrimaryHeader.SizeOfPartitionEntry     = DEFAULT_PRIMARY_TABLE_HEADER_SIZE_OF_PARTITION_ENTRY;
+
+  Status = SanitizePrimaryHeaderAllocationSize (&PrimaryHeader, &AllocationSize);
+  UT_ASSERT_NOT_EFI_ERROR (Status);
+
+  // Test that the allocation size is correct compared to the existing logic
+  UT_ASSERT_EQUAL (AllocationSize, PrimaryHeader.NumberOfPartitionEntries * PrimaryHeader.SizeOfPartitionEntry);
+
+  // Test that an overflow is detected
+  PrimaryHeader.NumberOfPartitionEntries = MAX_UINT32;
+  PrimaryHeader.SizeOfPartitionEntry     = 5;
+  Status                                 = SanitizePrimaryHeaderAllocationSize (&PrimaryHeader, &AllocationSize);
+  UT_ASSERT_EQUAL (Status, EFI_BAD_BUFFER_SIZE);
+
+  // Test the inverse
+  PrimaryHeader.NumberOfPartitionEntries = 5;
+  PrimaryHeader.SizeOfPartitionEntry     = MAX_UINT32;
+  Status                                 = SanitizePrimaryHeaderAllocationSize (&PrimaryHeader, &AllocationSize);
+  UT_ASSERT_EQUAL (Status, EFI_BAD_BUFFER_SIZE);
+
+  // Test the worst case scenario
+  PrimaryHeader.NumberOfPartitionEntries = MAX_UINT32;
+  PrimaryHeader.SizeOfPartitionEntry     = MAX_UINT32;
+  Status                                 = SanitizePrimaryHeaderAllocationSize (&PrimaryHeader, &AllocationSize);
+  UT_ASSERT_EQUAL (Status, EFI_BAD_BUFFER_SIZE);
+
+  DEBUG ((DEBUG_INFO, "%a: Test passed\n", __func__));
+
+  return UNIT_TEST_PASSED;
+}
+
+/**
+  This function tests the SanitizePrimaryHeaderGptEventSize function.
+  It's intent is to test that the untrusted input from a EFI_GPT_DATA structure
+  will not cause an overflow when calculating the event size.
+
+  @param[in] Context  The unit test context.
+
+  @retval UNIT_TEST_PASSED  The test passed.
+  @retval UNIT_TEST_ERROR_TEST_FAILED  The test failed.
+**/
+UNIT_TEST_STATUS
+EFIAPI
+TestSanitizePrimaryHeaderGptEventSize (
+  IN UNIT_TEST_CONTEXT  Context
+  )
+{
+  UINT32                      EventSize;
+  UINT32                      ExistingLogicEventSize;
+  EFI_STATUS                  Status;
+  EFI_PARTITION_TABLE_HEADER  PrimaryHeader;
+  UINTN                       NumberOfPartition;
+  EFI_GPT_DATA                *GptData;
+
+  GptData = NULL;
+
+  // Test that a normal PrimaryHeader passes validation
+  PrimaryHeader.NumberOfPartitionEntries = 5;
+  PrimaryHeader.SizeOfPartitionEntry     = DEFAULT_PRIMARY_TABLE_HEADER_SIZE_OF_PARTITION_ENTRY;
+
+  // set the number of partitions
+  NumberOfPartition = 13;
+
+  // that the primary event size is correct
+  Status = SanitizePrimaryHeaderGptEventSize (&PrimaryHeader, NumberOfPartition, &EventSize);
+  UT_ASSERT_NOT_EFI_ERROR (Status);
+
+  // Calculate the existing logic event size
+  ExistingLogicEventSize = (UINT32)(sizeof (TCG_PCR_EVENT_HDR) + OFFSET_OF (EFI_GPT_DATA, Partitions)
+                                    + NumberOfPartition * PrimaryHeader.SizeOfPartitionEntry);
+
+  // Check that the event size is correct
+  UT_ASSERT_EQUAL (EventSize, ExistingLogicEventSize);
+
+  // Tests that the primary event size may not overflow
+  Status = SanitizePrimaryHeaderGptEventSize (&PrimaryHeader, MAX_UINT32, &EventSize);
+  UT_ASSERT_EQUAL (Status, EFI_BAD_BUFFER_SIZE);
+
+  // Test that the size of partition entries may not overflow
+  PrimaryHeader.SizeOfPartitionEntry = MAX_UINT32;
+  Status                             = SanitizePrimaryHeaderGptEventSize (&PrimaryHeader, NumberOfPartition, &EventSize);
+  UT_ASSERT_EQUAL (Status, EFI_BAD_BUFFER_SIZE);
+
+  DEBUG ((DEBUG_INFO, "%a: Test passed\n", __func__));
+
+  return UNIT_TEST_PASSED;
+}
+
+// *--------------------------------------------------------------------*
+// *  Unit Test Code Main Function
+// *--------------------------------------------------------------------*
+
+/**
+  This function acts as the entry point for the unit tests.
+
+  @param argc - The number of command line arguments
+  @param argv - The command line arguments
+
+  @return int - The status of the test
+**/
+EFI_STATUS
+EFIAPI
+UefiTestMain (
+  VOID
+  )
+{
+  EFI_STATUS                  Status;
+  UNIT_TEST_FRAMEWORK_HANDLE  Framework;
+  UNIT_TEST_SUITE_HANDLE      TcgMeasureBootLibValidationTestSuite;
+
+  Framework = NULL;
+
+  DEBUG ((DEBUG_INFO, "%a: TestMain() - Start\n", UNIT_TEST_NAME));
+
+  Status = InitUnitTestFramework (&Framework, UNIT_TEST_NAME, gEfiCallerBaseName, UNIT_TEST_VERSION);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "%a: Failed in InitUnitTestFramework. Status = %r\n", UNIT_TEST_NAME, Status));
+    goto EXIT;
+  }
+
+  Status = CreateUnitTestSuite (&TcgMeasureBootLibValidationTestSuite, Framework, "TcgMeasureBootLibValidationTestSuite", "Common.TcgMeasureBootLibValidation", NULL, NULL);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "%s: Failed in CreateUnitTestSuite for TcgMeasureBootLibValidationTestSuite\n", UNIT_TEST_NAME));
+    Status = EFI_OUT_OF_RESOURCES;
+    goto EXIT;
+  }
+
+  // -----------Suite---------------------------------Description----------------------------Class----------------------------------Test Function------------------------Pre---Clean-Context
+  AddTestCase (TcgMeasureBootLibValidationTestSuite, "Tests Validating EFI Partition Table", "Common.TcgMeasureBootLibValidation", TestSanitizeEfiPartitionTableHeader, NULL, NULL, NULL);
+  AddTestCase (TcgMeasureBootLibValidationTestSuite, "Tests Primary header gpt event checks for overflow", "Common.TcgMeasureBootLibValidation", TestSanitizePrimaryHeaderAllocationSize, NULL, NULL, NULL);
+  AddTestCase (TcgMeasureBootLibValidationTestSuite, "Tests Primary header allocation size checks for overflow", "Common.TcgMeasureBootLibValidation", TestSanitizePrimaryHeaderGptEventSize, NULL, NULL, NULL);
+
+  Status = RunAllTestSuites (Framework);
+
+EXIT:
+  if (Framework != NULL) {
+    FreeUnitTestFramework (Framework);
+  }
+
+  DEBUG ((DEBUG_INFO, "%a: TestMain() - End\n", UNIT_TEST_NAME));
+  return Status;
+}
+
+///
+/// Avoid ECC error for function name that starts with lower case letter
+///
+#define DxeTpmMeasureBootLibUnitTestMain  main
+
+/**
+  Standard POSIX C entry point for host based unit test execution.
+
+  @param[in] Argc  Number of arguments
+  @param[in] Argv  Array of pointers to arguments
+
+  @retval 0      Success
+  @retval other  Error
+**/
+INT32
+DxeTpmMeasureBootLibUnitTestMain (
+  IN INT32  Argc,
+  IN CHAR8  *Argv[]
+  )
+{
+  return (INT32)UefiTestMain ();
+}
diff --git a/SecurityPkg/Library/DxeTpmMeasureBootLib/InternalUnitTest/DxeTpmMeasureBootLibSanitizationTestHost.inf b/SecurityPkg/Library/DxeTpmMeasureBootLib/InternalUnitTest/DxeTpmMeasureBootLibSanitizationTestHost.inf
new file mode 100644
index 00000000..ba4c58cb
--- /dev/null
+++ b/SecurityPkg/Library/DxeTpmMeasureBootLib/InternalUnitTest/DxeTpmMeasureBootLibSanitizationTestHost.inf
@@ -0,0 +1,28 @@
+## @file
+# This file builds the unit tests for DxeTpmMeasureBootLib
+#
+# Copyright (C) Microsoft Corporation.<BR>
+# SPDX-License-Identifier: BSD-2-Clause-Patent
+##
+
+[Defines]
+  INF_VERSION                    = 0x00010006
+  BASE_NAME                      = DxeTpmMeasuredBootLibTest
+  FILE_GUID                      = eb01bc38-309c-4d3e-967e-9f078c90772f
+  MODULE_TYPE                    = HOST_APPLICATION
+  VERSION_STRING                 = 1.0
+  ENTRY_POINT                    = main
+
+[Sources]
+  DxeTpmMeasureBootLibSanitizationTest.c
+  ../DxeTpmMeasureBootLibSanitization.c
+
+[Packages]
+  MdePkg/MdePkg.dec
+
+[LibraryClasses]
+  BaseLib
+  DebugLib
+  UnitTestLib
+  PrintLib
+  SafeIntLib
diff --git a/SecurityPkg/SecurityPkg.ci.yaml b/SecurityPkg/SecurityPkg.ci.yaml
index 97137406..c6716cb1 100644
--- a/SecurityPkg/SecurityPkg.ci.yaml
+++ b/SecurityPkg/SecurityPkg.ci.yaml
@@ -17,6 +17,7 @@
         "ExceptionList": [
             "8005", "gRT",
             "8001", "DxeTpm2MeasureBootLibUnitTestMain",
+            "8001", "DxeTpmMeasureBootLibUnitTestMain"
         ],
         ## Both file path and directory path are accepted.
         "IgnoreFiles": [
diff --git a/SecurityPkg/Test/SecurityPkgHostTest.dsc b/SecurityPkg/Test/SecurityPkgHostTest.dsc
index 57ab3af0..470bba07 100644
--- a/SecurityPkg/Test/SecurityPkgHostTest.dsc
+++ b/SecurityPkg/Test/SecurityPkgHostTest.dsc
@@ -27,6 +27,7 @@
   SecurityPkg/Library/SecureBootVariableLib/UnitTest/MockUefiLib.inf
   SecurityPkg/Test/Mock/Library/GoogleTest/MockPlatformPKProtectionLib/MockPlatformPKProtectionLib.inf
   SecurityPkg/Library/DxeTpm2MeasureBootLib/InternalUnitTest/DxeTpm2MeasureBootLibSanitizationTestHost.inf
+  SecurityPkg/Library/DxeTpmMeasureBootLib/InternalUnitTest/DxeTpmMeasureBootLibSanitizationTestHost.inf
 
   #
   # Build SecurityPkg HOST_APPLICATION Tests
-- 
2.33.0