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

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

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

Signed-off-by: Doug Flick [MSFT] <doug.edk2@gmail.com>

reference: https://github.com/tianocore/edk2/pull/5264
Signed-off-by: yexiao <yexiao7@huawei.com>
---
 .../DxeTpm2MeasureBootLib.c                   |  69 ++--
 .../DxeTpm2MeasureBootLib.inf                 |   4 +-
 .../DxeTpm2MeasureBootLibSanitization.c       | 275 ++++++++++++++++
 .../DxeTpm2MeasureBootLibSanitization.h       | 113 +++++++
 .../DxeTpm2MeasureBootLibSanitizationTest.c   | 303 ++++++++++++++++++
 ...Tpm2MeasureBootLibSanitizationTestHost.inf |  28 ++
 SecurityPkg/SecurityPkg.ci.yaml               |   1 +
 SecurityPkg/Test/SecurityPkgHostTest.dsc      |   1 +
 8 files changed, 764 insertions(+), 30 deletions(-)
 create mode 100644 SecurityPkg/Library/DxeTpm2MeasureBootLib/DxeTpm2MeasureBootLibSanitization.c
 create mode 100644 SecurityPkg/Library/DxeTpm2MeasureBootLib/DxeTpm2MeasureBootLibSanitization.h
 create mode 100644 SecurityPkg/Library/DxeTpm2MeasureBootLib/InternalUnitTest/DxeTpm2MeasureBootLibSanitizationTest.c
 create mode 100644 SecurityPkg/Library/DxeTpm2MeasureBootLib/InternalUnitTest/DxeTpm2MeasureBootLibSanitizationTestHost.inf

diff --git a/SecurityPkg/Library/DxeTpm2MeasureBootLib/DxeTpm2MeasureBootLib.c b/SecurityPkg/Library/DxeTpm2MeasureBootLib/DxeTpm2MeasureBootLib.c
index 36a256a7..0a37a4b6 100644
--- a/SecurityPkg/Library/DxeTpm2MeasureBootLib/DxeTpm2MeasureBootLib.c
+++ b/SecurityPkg/Library/DxeTpm2MeasureBootLib/DxeTpm2MeasureBootLib.c
@@ -20,6 +20,8 @@ Copyright (c) 2013 - 2018, Intel Corporation. All rights reserved.<BR>
 (C) Copyright 2015 Hewlett Packard Enterprise Development LP<BR>
 SPDX-License-Identifier: BSD-2-Clause-Patent
 
+Copyright (c) Microsoft Corporation.<BR>
+SPDX-License-Identifier: BSD-2-Clause-Patent
 **/
 
 #include <PiDxe.h>
@@ -44,6 +46,8 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Library/HobLib.h>
 #include <Protocol/CcMeasurement.h>
 
+#include "DxeTpm2MeasureBootLibSanitization.h"
+
 typedef struct {
   EFI_TCG2_PROTOCOL              *Tcg2Protocol;
   EFI_CC_MEASUREMENT_PROTOCOL    *CcProtocol;
@@ -144,10 +148,11 @@ Tcg2MeasureGptTable (
   EFI_TCG2_EVENT               *Tcg2Event;
   EFI_CC_EVENT                 *CcEvent;
   EFI_GPT_DATA                 *GptData;
-  UINT32                       EventSize;
+  UINT32                       TcgEventSize;
   EFI_TCG2_PROTOCOL            *Tcg2Protocol;
   EFI_CC_MEASUREMENT_PROTOCOL  *CcProtocol;
   EFI_CC_MR_INDEX              MrIndex;
+  UINT32                       AllocSize;
 
   if (mTcg2MeasureGptCount > 0) {
     return EFI_SUCCESS;
@@ -195,25 +200,22 @@ Tcg2MeasureGptTable (
                      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;
   }
 
   //
-  // PrimaryHeader->SizeOfPartitionEntry should not be zero
+  // Read the partition entry.
   //
-  if (PrimaryHeader->SizeOfPartitionEntry == 0) {
-    DEBUG ((DEBUG_ERROR, "SizeOfPartitionEntry should not be zero!\n"));
+  Status = SanitizePrimaryHeaderAllocationSize (PrimaryHeader, &AllocSize);
+  if (EFI_ERROR (Status)) {
     FreePool (PrimaryHeader);
     return EFI_BAD_BUFFER_SIZE;
   }
 
-  //
-  // Read the partition entry.
-  //
-  EntryPtr = (UINT8 *)AllocatePool (PrimaryHeader->NumberOfPartitionEntries * PrimaryHeader->SizeOfPartitionEntry);
+  EntryPtr = (UINT8 *)AllocatePool (AllocSize);
   if (EntryPtr == NULL) {
     FreePool (PrimaryHeader);
     return EFI_OUT_OF_RESOURCES;
@@ -223,7 +225,7 @@ Tcg2MeasureGptTable (
                      DiskIo,
                      BlockIo->Media->MediaId,
                      MultU64x32 (PrimaryHeader->PartitionEntryLBA, BlockIo->Media->BlockSize),
-                     PrimaryHeader->NumberOfPartitionEntries * PrimaryHeader->SizeOfPartitionEntry,
+                     AllocSize,
                      EntryPtr
                      );
   if (EFI_ERROR (Status)) {
@@ -248,16 +250,21 @@ Tcg2MeasureGptTable (
   //
   // Prepare Data for Measurement (CcProtocol and Tcg2Protocol)
   //
-  EventSize = (UINT32)(sizeof (EFI_GPT_DATA) - sizeof (GptData->Partitions)
-                       + NumberOfPartition * PrimaryHeader->SizeOfPartitionEntry);
-  EventPtr = (UINT8 *)AllocateZeroPool (EventSize + sizeof (EFI_TCG2_EVENT) - sizeof (Tcg2Event->Event));
+  Status = SanitizePrimaryHeaderGptEventSize (PrimaryHeader, NumberOfPartition, &TcgEventSize);
+  if (EFI_ERROR (Status)) {
+    FreePool (PrimaryHeader);
+    FreePool (EntryPtr);
+    return EFI_DEVICE_ERROR;
+  }
+
+  EventPtr = (UINT8 *)AllocateZeroPool (TcgEventSize);
   if (EventPtr == NULL) {
     Status = EFI_OUT_OF_RESOURCES;
     goto Exit;
   }
 
   Tcg2Event                       = (EFI_TCG2_EVENT *)EventPtr;
-  Tcg2Event->Size                 = EventSize + sizeof (EFI_TCG2_EVENT) - sizeof (Tcg2Event->Event);
+  Tcg2Event->Size                 = TcgEventSize;
   Tcg2Event->Header.HeaderSize    = sizeof (EFI_TCG2_EVENT_HEADER);
   Tcg2Event->Header.HeaderVersion = EFI_TCG2_EVENT_HEADER_VERSION;
   Tcg2Event->Header.PCRIndex      = 5;
@@ -310,7 +317,7 @@ Tcg2MeasureGptTable (
                                             CcProtocol,
                                             0,
                                             (EFI_PHYSICAL_ADDRESS)(UINTN)(VOID *)GptData,
-                                            (UINT64)EventSize,
+                                            (UINT64)TcgEventSize - OFFSET_OF (EFI_TCG2_EVENT, Event),
                                             CcEvent
                                             );
     if (!EFI_ERROR (Status)) {
@@ -326,7 +333,7 @@ Tcg2MeasureGptTable (
                              Tcg2Protocol,
                              0,
                              (EFI_PHYSICAL_ADDRESS)(UINTN)(VOID *)GptData,
-                             (UINT64)EventSize,
+                             (UINT64)TcgEventSize -  OFFSET_OF (EFI_TCG2_EVENT, Event),
                              Tcg2Event
                              );
     if (!EFI_ERROR (Status)) {
@@ -443,11 +450,13 @@ Tcg2MeasurePeImage (
       Tcg2Event->Header.PCRIndex  = 2;
       break;
     default:
-      DEBUG ((
-        DEBUG_ERROR,
-        "Tcg2MeasurePeImage: Unknown subsystem type %d",
-        ImageType
-        ));
+      DEBUG (
+        (
+         DEBUG_ERROR,
+         "Tcg2MeasurePeImage: Unknown subsystem type %d",
+         ImageType
+        )
+        );
       goto Finish;
   }
 
@@ -515,7 +524,7 @@ Finish:
 
   @param  MeasureBootProtocols  Pointer to the located measure boot protocol instances.
 
-  @retval EFI_SUCCESS           Sucessfully locate the measure boot protocol instances (at least one instance).
+  @retval EFI_SUCCESS           Successfully locate the measure boot protocol instances (at least one instance).
   @retval EFI_UNSUPPORTED       Measure boot is not supported.
 **/
 EFI_STATUS
@@ -646,12 +655,14 @@ DxeTpm2MeasureBootHandler (
     return EFI_SUCCESS;
   }
 
-  DEBUG ((
-    DEBUG_INFO,
-    "Tcg2Protocol = %p, CcMeasurementProtocol = %p\n",
-    MeasureBootProtocols.Tcg2Protocol,
-    MeasureBootProtocols.CcProtocol
-    ));
+  DEBUG (
+    (
+     DEBUG_INFO,
+     "Tcg2Protocol = %p, CcMeasurementProtocol = %p\n",
+     MeasureBootProtocols.Tcg2Protocol,
+     MeasureBootProtocols.CcProtocol
+    )
+    );
 
   //
   // Copy File Device Path
diff --git a/SecurityPkg/Library/DxeTpm2MeasureBootLib/DxeTpm2MeasureBootLib.inf b/SecurityPkg/Library/DxeTpm2MeasureBootLib/DxeTpm2MeasureBootLib.inf
index 6dca79a2..78395573 100644
--- a/SecurityPkg/Library/DxeTpm2MeasureBootLib/DxeTpm2MeasureBootLib.inf
+++ b/SecurityPkg/Library/DxeTpm2MeasureBootLib/DxeTpm2MeasureBootLib.inf
@@ -37,6 +37,8 @@
 
 [Sources]
   DxeTpm2MeasureBootLib.c
+  DxeTpm2MeasureBootLibSanitization.c
+  DxeTpm2MeasureBootLibSanitization.h
 
 [Packages]
   MdePkg/MdePkg.dec
@@ -46,6 +48,7 @@
 
 [LibraryClasses]
   BaseMemoryLib
+  SafeIntLib
   DebugLib
   MemoryAllocationLib
   DevicePathLib
@@ -65,4 +68,3 @@
   gEfiFirmwareVolumeBlockProtocolGuid   ## SOMETIMES_CONSUMES
   gEfiBlockIoProtocolGuid               ## SOMETIMES_CONSUMES
   gEfiDiskIoProtocolGuid                ## SOMETIMES_CONSUMES
-
diff --git a/SecurityPkg/Library/DxeTpm2MeasureBootLib/DxeTpm2MeasureBootLibSanitization.c b/SecurityPkg/Library/DxeTpm2MeasureBootLib/DxeTpm2MeasureBootLibSanitization.c
new file mode 100644
index 00000000..ef3b77b4
--- /dev/null
+++ b/SecurityPkg/Library/DxeTpm2MeasureBootLib/DxeTpm2MeasureBootLibSanitization.c
@@ -0,0 +1,275 @@
+/** @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:
+
+  DxeTpm2MeasureBootLibImageRead() 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 "DxeTpm2MeasureBootLibSanitization.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;
+  }
+
+  //
+  // The partition entries should all be before the first usable block
+  //
+  if (PrimaryHeader->FirstUsableLBA <= PrimaryHeader->PartitionEntryLBA) {
+    DEBUG ((DEBUG_ERROR, "GPT PartitionEntryLBA is not less than FirstUsableLBA!\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
+  (sizeof (EFI_TCG2_EVENT) - sizeof (Tcg2Event->Event)) . When hashing the buffer allocated with this
+  size, the caller must subtract the size of the (sizeof (EFI_TCG2_EVENT) - sizeof (Tcg2Event->Event))
+  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);
+  //
+  Status = SafeUint32Mult (SafeNumberOfPartitions, PrimaryHeader->SizeOfPartitionEntry, EventSize);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "Event Size would have overflowed!\n"));
+    return EFI_BAD_BUFFER_SIZE;
+  }
+
+  //
+  // Replacing logic:
+  // *EventSize + sizeof (EFI_TCG2_EVENT) - sizeof (Tcg2Event->Event);
+  //
+  Status = SafeUint32Add (
+             OFFSET_OF (EFI_TCG2_EVENT, Event) + 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/DxeTpm2MeasureBootLib/DxeTpm2MeasureBootLibSanitization.h b/SecurityPkg/Library/DxeTpm2MeasureBootLib/DxeTpm2MeasureBootLibSanitization.h
new file mode 100644
index 00000000..74af8c8b
--- /dev/null
+++ b/SecurityPkg/Library/DxeTpm2MeasureBootLib/DxeTpm2MeasureBootLibSanitization.h
@@ -0,0 +1,113 @@
+/** @file
+  This file includes the function prototypes for the sanitization functions.
+
+  These are those functions:
+
+  DxeTpm2MeasureBootLibImageRead() 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
+
+**/
+
+#ifndef DXE_TPM2_MEASURE_BOOT_LIB_SANITATION_
+#define DXE_TPM2_MEASURE_BOOT_LIB_SANITATION_
+
+#include <Uefi.h>
+#include <Uefi/UefiSpec.h>
+#include <Protocol/BlockIo.h>
+#include <IndustryStandard/UefiTcgPlatform.h>
+#include <Protocol/Tcg2Protocol.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
+  (sizeof (EFI_TCG2_EVENT) - sizeof (Tcg2Event->Event)) . When hashing the buffer allocated with this
+  size, the caller must subtract the size of the (sizeof (EFI_TCG2_EVENT) - sizeof (Tcg2Event->Event))
+  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_TPM2_MEASURE_BOOT_LIB_SANITATION_
diff --git a/SecurityPkg/Library/DxeTpm2MeasureBootLib/InternalUnitTest/DxeTpm2MeasureBootLibSanitizationTest.c b/SecurityPkg/Library/DxeTpm2MeasureBootLib/InternalUnitTest/DxeTpm2MeasureBootLibSanitizationTest.c
new file mode 100644
index 00000000..701d1b46
--- /dev/null
+++ b/SecurityPkg/Library/DxeTpm2MeasureBootLib/InternalUnitTest/DxeTpm2MeasureBootLibSanitizationTest.c
@@ -0,0 +1,303 @@
+/** @file
+  This file includes the unit test cases for the DxeTpm2MeasureBootLibSanitizationTest.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 <Protocol/Tcg2Protocol.h>
+
+#include "../DxeTpm2MeasureBootLibSanitization.h"
+
+#define UNIT_TEST_NAME     "DxeTpm2MeasureBootLibSanitizationTest"
+#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;
+  EFI_TCG2_EVENT              *Tcg2Event;
+
+  Tcg2Event = NULL;
+  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)(OFFSET_OF (EFI_TCG2_EVENT, Event) + 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.
+
+  @retval UNIT_TEST_PASSED  The test passed.
+  @retval UNIT_TEST_ERROR_TEST_FAILED  The test failed.
+  @retval others The test failed.
+**/
+EFI_STATUS
+EFIAPI
+UefiTestMain (
+  VOID
+  )
+{
+  EFI_STATUS                  Status;
+  UNIT_TEST_FRAMEWORK_HANDLE  Framework;
+  UNIT_TEST_SUITE_HANDLE      Tcg2MeasureBootLibValidationTestSuite;
+
+  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 (&Tcg2MeasureBootLibValidationTestSuite, Framework, "Tcg2MeasureBootLibValidationTestSuite", "Common.Tcg2MeasureBootLibValidation", NULL, NULL);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "%s: Failed in CreateUnitTestSuite for Tcg2MeasureBootLibValidationTestSuite\n", UNIT_TEST_NAME));
+    Status = EFI_OUT_OF_RESOURCES;
+    goto EXIT;
+  }
+
+  // -----------Suite---------------------------------Description----------------------------Class----------------------------------Test Function------------------------Pre---Clean-Context
+  AddTestCase (Tcg2MeasureBootLibValidationTestSuite, "Tests Validating EFI Partition Table", "Common.Tcg2MeasureBootLibValidation", TestSanitizeEfiPartitionTableHeader, NULL, NULL, NULL);
+  AddTestCase (Tcg2MeasureBootLibValidationTestSuite, "Tests Primary header gpt event checks for overflow", "Common.Tcg2MeasureBootLibValidation", TestSanitizePrimaryHeaderAllocationSize, NULL, NULL, NULL);
+  AddTestCase (Tcg2MeasureBootLibValidationTestSuite, "Tests Primary header allocation size checks for overflow", "Common.Tcg2MeasureBootLibValidation", 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 DxeTpm2MeasureBootLibUnitTestMain  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
+DxeTpm2MeasureBootLibUnitTestMain (
+  IN INT32  Argc,
+  IN CHAR8  *Argv[]
+  )
+{
+  return (INT32)UefiTestMain ();
+}
diff --git a/SecurityPkg/Library/DxeTpm2MeasureBootLib/InternalUnitTest/DxeTpm2MeasureBootLibSanitizationTestHost.inf b/SecurityPkg/Library/DxeTpm2MeasureBootLib/InternalUnitTest/DxeTpm2MeasureBootLibSanitizationTestHost.inf
new file mode 100644
index 00000000..4ea55787
--- /dev/null
+++ b/SecurityPkg/Library/DxeTpm2MeasureBootLib/InternalUnitTest/DxeTpm2MeasureBootLibSanitizationTestHost.inf
@@ -0,0 +1,28 @@
+## @file
+# This file builds the unit tests for DxeTpm2MeasureBootLib
+#
+# Copyright (C) Microsoft Corporation.<BR>
+# SPDX-License-Identifier: BSD-2-Clause-Patent
+##
+
+[Defines]
+  INF_VERSION                    = 0x00010006
+  BASE_NAME                      = DxeTpm2MeasuredBootLibTest
+  FILE_GUID                      = 144d757f-d423-484e-9309-a23695fad5bd
+  MODULE_TYPE                    = HOST_APPLICATION
+  VERSION_STRING                 = 1.0
+  ENTRY_POINT                    = main
+
+[Sources]
+  DxeTpm2MeasureBootLibSanitizationTest.c
+  ../DxeTpm2MeasureBootLibSanitization.c
+
+[Packages]
+  MdePkg/MdePkg.dec
+
+[LibraryClasses]
+  BaseLib
+  DebugLib
+  UnitTestLib
+  PrintLib
+  SafeIntLib
diff --git a/SecurityPkg/SecurityPkg.ci.yaml b/SecurityPkg/SecurityPkg.ci.yaml
index 2138b0a5..97137406 100644
--- a/SecurityPkg/SecurityPkg.ci.yaml
+++ b/SecurityPkg/SecurityPkg.ci.yaml
@@ -16,6 +16,7 @@
         ## ]
         "ExceptionList": [
             "8005", "gRT",
+            "8001", "DxeTpm2MeasureBootLibUnitTestMain",
         ],
         ## Both file path and directory path are accepted.
         "IgnoreFiles": [
diff --git a/SecurityPkg/Test/SecurityPkgHostTest.dsc b/SecurityPkg/Test/SecurityPkgHostTest.dsc
index ad5b4fc3..57ab3af0 100644
--- a/SecurityPkg/Test/SecurityPkgHostTest.dsc
+++ b/SecurityPkg/Test/SecurityPkgHostTest.dsc
@@ -26,6 +26,7 @@
   SecurityPkg/Library/SecureBootVariableLib/UnitTest/MockPlatformPKProtectionLib.inf
   SecurityPkg/Library/SecureBootVariableLib/UnitTest/MockUefiLib.inf
   SecurityPkg/Test/Mock/Library/GoogleTest/MockPlatformPKProtectionLib/MockPlatformPKProtectionLib.inf
+  SecurityPkg/Library/DxeTpm2MeasureBootLib/InternalUnitTest/DxeTpm2MeasureBootLibSanitizationTestHost.inf
 
   #
   # Build SecurityPkg HOST_APPLICATION Tests
-- 
2.33.0