SecurityPkg/Library/DxeTpmMeasureBootLib/InternalUnitTest/DxeTpmMeasureBootLibSanitizationTest.c - edk2 - Git at Google

 /** @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 = TpmSanitizeEfiPartitionTableHeader (&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                                 = TpmSanitizeEfiPartitionTableHeader (&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                          = TpmSanitizeEfiPartitionTableHeader (&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                             = TpmSanitizeEfiPartitionTableHeader (&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 = TpmSanitizePrimaryHeaderAllocationSize (&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                                 = TpmSanitizePrimaryHeaderAllocationSize (&PrimaryHeader, &AllocationSize);
   UT_ASSERT_EQUAL (Status, EFI_BAD_BUFFER_SIZE);

   // Test the inverse
   PrimaryHeader.NumberOfPartitionEntries = 5;
   PrimaryHeader.SizeOfPartitionEntry     = MAX_UINT32;
   Status                                 = TpmSanitizePrimaryHeaderAllocationSize (&PrimaryHeader, &AllocationSize);
   UT_ASSERT_EQUAL (Status, EFI_BAD_BUFFER_SIZE);

   // Test the worst case scenario
   PrimaryHeader.NumberOfPartitionEntries = MAX_UINT32;
   PrimaryHeader.SizeOfPartitionEntry     = MAX_UINT32;
   Status                                 = TpmSanitizePrimaryHeaderAllocationSize (&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;

   // 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 = TpmSanitizePrimaryHeaderGptEventSize (&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 = TpmSanitizePrimaryHeaderGptEventSize (&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                             = TpmSanitizePrimaryHeaderGptEventSize (&PrimaryHeader, NumberOfPartition, &EventSize);
   UT_ASSERT_EQUAL (Status, EFI_BAD_BUFFER_SIZE);

   DEBUG ((DEBUG_INFO, "%a: Test passed\n", __func__));

   return UNIT_TEST_PASSED;
 }

 /**
   This function tests the SanitizePeImageEventSize function.
   It's intent is to test that the untrusted input from a file path for an
   EFI_IMAGE_LOAD_EVENT structure will not cause an overflow when calculating
   the event size when allocating space.

   @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
 TestSanitizePeImageEventSize (
   IN UNIT_TEST_CONTEXT  Context
   )
 {
   UINT32                    EventSize;
   UINTN                     ExistingLogicEventSize;
   UINT32                    FilePathSize;
   EFI_STATUS                Status;
   EFI_DEVICE_PATH_PROTOCOL  DevicePath;
   EFI_IMAGE_LOAD_EVENT      *ImageLoadEvent;
   UNIT_TEST_STATUS          TestStatus;

   TestStatus = UNIT_TEST_ERROR_TEST_FAILED;

   // Generate EFI_DEVICE_PATH_PROTOCOL test data
   DevicePath.Type      = 0;
   DevicePath.SubType   = 0;
   DevicePath.Length[0] = 0;
   DevicePath.Length[1] = 0;

   // Generate EFI_IMAGE_LOAD_EVENT test data
   ImageLoadEvent = AllocateZeroPool (sizeof (EFI_IMAGE_LOAD_EVENT) + sizeof (EFI_DEVICE_PATH_PROTOCOL));
   if (ImageLoadEvent == NULL) {
     DEBUG ((DEBUG_ERROR, "%a: AllocateZeroPool failed\n", __func__));
     goto Exit;
   }

   // Populate EFI_IMAGE_LOAD_EVENT54 test data
   ImageLoadEvent->ImageLocationInMemory = (EFI_PHYSICAL_ADDRESS)0x12345678;
   ImageLoadEvent->ImageLengthInMemory   = 0x1000;
   ImageLoadEvent->ImageLinkTimeAddress  = (UINTN)ImageLoadEvent;
   ImageLoadEvent->LengthOfDevicePath    = sizeof (EFI_DEVICE_PATH_PROTOCOL);
   CopyMem (ImageLoadEvent->DevicePath, &DevicePath, sizeof (EFI_DEVICE_PATH_PROTOCOL));

   FilePathSize = 255;

   // Test that a normal PE image passes validation
   Status = TpmSanitizePeImageEventSize (FilePathSize, &EventSize);
   if (EFI_ERROR (Status)) {
     UT_LOG_ERROR ("SanitizePeImageEventSize failed with %r\n", Status);
     goto Exit;
   }

   // Test that the event size is correct compared to the existing logic
   ExistingLogicEventSize  = OFFSET_OF (EFI_IMAGE_LOAD_EVENT, DevicePath) + FilePathSize;
   ExistingLogicEventSize += sizeof (TCG_PCR_EVENT_HDR);

   if (EventSize != ExistingLogicEventSize) {
     UT_LOG_ERROR ("SanitizePeImageEventSize returned an incorrect event size. Expected %u, got %u\n", ExistingLogicEventSize, EventSize);
     goto Exit;
   }

   // Test that the event size may not overflow
   Status = TpmSanitizePeImageEventSize (MAX_UINT32, &EventSize);
   if (Status != EFI_BAD_BUFFER_SIZE) {
     UT_LOG_ERROR ("SanitizePeImageEventSize succeded when it was supposed to fail with %r\n", Status);
     goto Exit;
   }

   TestStatus = UNIT_TEST_PASSED;
 Exit:

   if (ImageLoadEvent != NULL) {
     FreePool (ImageLoadEvent);
   }

   if (TestStatus == UNIT_TEST_ERROR_TEST_FAILED) {
     DEBUG ((DEBUG_ERROR, "%a: Test failed\n", __func__));
   } else {
     DEBUG ((DEBUG_INFO, "%a: Test passed\n", __func__));
   }

   return TestStatus;
 }

 // *--------------------------------------------------------------------*
 // *  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);
   AddTestCase (TcgMeasureBootLibValidationTestSuite, "Tests PE Image and FileSize checks for overflow", "Common.TcgMeasureBootLibValidation", TestSanitizePeImageEventSize, 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 ();
 }
	/** @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 = TpmSanitizeEfiPartitionTableHeader (&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 = TpmSanitizeEfiPartitionTableHeader (&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 = TpmSanitizeEfiPartitionTableHeader (&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 = TpmSanitizeEfiPartitionTableHeader (&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 = TpmSanitizePrimaryHeaderAllocationSize (&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 = TpmSanitizePrimaryHeaderAllocationSize (&PrimaryHeader, &AllocationSize);
	UT_ASSERT_EQUAL (Status, EFI_BAD_BUFFER_SIZE);

	// Test the inverse
	PrimaryHeader.NumberOfPartitionEntries = 5;
	PrimaryHeader.SizeOfPartitionEntry = MAX_UINT32;
	Status = TpmSanitizePrimaryHeaderAllocationSize (&PrimaryHeader, &AllocationSize);
	UT_ASSERT_EQUAL (Status, EFI_BAD_BUFFER_SIZE);

	// Test the worst case scenario
	PrimaryHeader.NumberOfPartitionEntries = MAX_UINT32;
	PrimaryHeader.SizeOfPartitionEntry = MAX_UINT32;
	Status = TpmSanitizePrimaryHeaderAllocationSize (&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;

	// 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 = TpmSanitizePrimaryHeaderGptEventSize (&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 = TpmSanitizePrimaryHeaderGptEventSize (&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 = TpmSanitizePrimaryHeaderGptEventSize (&PrimaryHeader, NumberOfPartition, &EventSize);
	UT_ASSERT_EQUAL (Status, EFI_BAD_BUFFER_SIZE);

	DEBUG ((DEBUG_INFO, "%a: Test passed\n", __func__));

	return UNIT_TEST_PASSED;
	}

	/**
	This function tests the SanitizePeImageEventSize function.
	It's intent is to test that the untrusted input from a file path for an
	EFI_IMAGE_LOAD_EVENT structure will not cause an overflow when calculating
	the event size when allocating space.

	@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
	TestSanitizePeImageEventSize (
	IN UNIT_TEST_CONTEXT Context
	)
	{
	UINT32 EventSize;
	UINTN ExistingLogicEventSize;
	UINT32 FilePathSize;
	EFI_STATUS Status;
	EFI_DEVICE_PATH_PROTOCOL DevicePath;
	EFI_IMAGE_LOAD_EVENT *ImageLoadEvent;
	UNIT_TEST_STATUS TestStatus;

	TestStatus = UNIT_TEST_ERROR_TEST_FAILED;

	// Generate EFI_DEVICE_PATH_PROTOCOL test data
	DevicePath.Type = 0;
	DevicePath.SubType = 0;
	DevicePath.Length[0] = 0;
	DevicePath.Length[1] = 0;

	// Generate EFI_IMAGE_LOAD_EVENT test data
	ImageLoadEvent = AllocateZeroPool (sizeof (EFI_IMAGE_LOAD_EVENT) + sizeof (EFI_DEVICE_PATH_PROTOCOL));
	if (ImageLoadEvent == NULL) {
	DEBUG ((DEBUG_ERROR, "%a: AllocateZeroPool failed\n", __func__));
	goto Exit;
	}

	// Populate EFI_IMAGE_LOAD_EVENT54 test data
	ImageLoadEvent->ImageLocationInMemory = (EFI_PHYSICAL_ADDRESS)0x12345678;
	ImageLoadEvent->ImageLengthInMemory = 0x1000;
	ImageLoadEvent->ImageLinkTimeAddress = (UINTN)ImageLoadEvent;
	ImageLoadEvent->LengthOfDevicePath = sizeof (EFI_DEVICE_PATH_PROTOCOL);
	CopyMem (ImageLoadEvent->DevicePath, &DevicePath, sizeof (EFI_DEVICE_PATH_PROTOCOL));

	FilePathSize = 255;

	// Test that a normal PE image passes validation
	Status = TpmSanitizePeImageEventSize (FilePathSize, &EventSize);
	if (EFI_ERROR (Status)) {
	UT_LOG_ERROR ("SanitizePeImageEventSize failed with %r\n", Status);
	goto Exit;
	}

	// Test that the event size is correct compared to the existing logic
	ExistingLogicEventSize = OFFSET_OF (EFI_IMAGE_LOAD_EVENT, DevicePath) + FilePathSize;
	ExistingLogicEventSize += sizeof (TCG_PCR_EVENT_HDR);

	if (EventSize != ExistingLogicEventSize) {
	UT_LOG_ERROR ("SanitizePeImageEventSize returned an incorrect event size. Expected %u, got %u\n", ExistingLogicEventSize, EventSize);
	goto Exit;
	}

	// Test that the event size may not overflow
	Status = TpmSanitizePeImageEventSize (MAX_UINT32, &EventSize);
	if (Status != EFI_BAD_BUFFER_SIZE) {
	UT_LOG_ERROR ("SanitizePeImageEventSize succeded when it was supposed to fail with %r\n", Status);
	goto Exit;
	}

	TestStatus = UNIT_TEST_PASSED;
	Exit:

	if (ImageLoadEvent != NULL) {
	FreePool (ImageLoadEvent);
	}

	if (TestStatus == UNIT_TEST_ERROR_TEST_FAILED) {
	DEBUG ((DEBUG_ERROR, "%a: Test failed\n", __func__));
	} else {
	DEBUG ((DEBUG_INFO, "%a: Test passed\n", __func__));
	}

	return TestStatus;
	}

	// --------------------------------------------------------------------
	// * 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);
	AddTestCase (TcgMeasureBootLibValidationTestSuite, "Tests PE Image and FileSize checks for overflow", "Common.TcgMeasureBootLibValidation", TestSanitizePeImageEventSize, 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 ();
	}