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>

8 files changed, 716 insertions, 14 deletions

diff --git a/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLib.c b/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLib.c
index 220393d..669ab19 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 ebab6f7..414c654 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 0000000..a3fa46f
--- /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 0000000..0d9d00c
--- /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 0000000..eeb928c
--- /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 0000000..47b0811
--- /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 2438953..53e5b1f 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 788c1ab..1655e57 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