From 3236b614ee778d5fec0b58098bb5b7b5e88a9bca Mon Sep 17 00:00:00 2001
From: Doug Flick <dougflick@microsoft.com>
Date: Fri, 26 Jan 2024 05:54:56 +0800
Subject: [PATCH 14/19] NetworkPkg: UefiPxeBcDxe: SECURITY PATCH CVE-2023-45235
 Unit Tests

REF:https://bugzilla.tianocore.org/show_bug.cgi?id=4540

Unit tests to confirm that the bug..

Buffer overflow when handling Server ID option from a DHCPv6 proxy
Advertise message

..has been patched.

This patch contains unit tests for the following functions:
PxeBcRequestBootService
PxeBcDhcp6Discover

Cc: Saloni Kasbekar <saloni.kasbekar@intel.com>
Cc: Zachary Clark-williams <zachary.clark-williams@intel.com>

Signed-off-by: Doug Flick [MSFT] <doug.edk2@gmail.com>
Reviewed-by: Saloni Kasbekar <saloni.kasbekar@intel.com>

reference: https://github.com/tianocore/edk2/pull/5352
Signed-off-by: yexiao <yexiao7@huawei.com>
---
 NetworkPkg/Test/NetworkPkgHostTest.dsc        |   5 +-
 .../GoogleTest/PxeBcDhcp6GoogleTest.cpp       | 278 +++++++++++++++++-
 .../GoogleTest/PxeBcDhcp6GoogleTest.h         |  18 ++
 3 files changed, 298 insertions(+), 3 deletions(-)

diff --git a/NetworkPkg/Test/NetworkPkgHostTest.dsc b/NetworkPkg/Test/NetworkPkgHostTest.dsc
index e26a2598..f313a605 100644
--- a/NetworkPkg/Test/NetworkPkgHostTest.dsc
+++ b/NetworkPkg/Test/NetworkPkgHostTest.dsc
@@ -27,7 +27,10 @@
   #
   NetworkPkg/Dhcp6Dxe/GoogleTest/Dhcp6DxeGoogleTest.inf
   NetworkPkg/Ip6Dxe/GoogleTest/Ip6DxeGoogleTest.inf
-  NetworkPkg/UefiPxeBcDxe/GoogleTest/UefiPxeBcDxeGoogleTest.inf
+  NetworkPkg/UefiPxeBcDxe/GoogleTest/UefiPxeBcDxeGoogleTest.inf {
+    <LibraryClasses>
+      UefiRuntimeServicesTableLib|MdePkg/Test/Mock/Library/GoogleTest/MockUefiRuntimeServicesTableLib/MockUefiRuntimeServicesTableLib.inf
+  }
 
 # Despite these library classes being listed in [LibraryClasses] below, they are not needed for the host-based unit tests.
 [LibraryClasses]
diff --git a/NetworkPkg/UefiPxeBcDxe/GoogleTest/PxeBcDhcp6GoogleTest.cpp b/NetworkPkg/UefiPxeBcDxe/GoogleTest/PxeBcDhcp6GoogleTest.cpp
index c66a60bc..d447ac02 100644
--- a/NetworkPkg/UefiPxeBcDxe/GoogleTest/PxeBcDhcp6GoogleTest.cpp
+++ b/NetworkPkg/UefiPxeBcDxe/GoogleTest/PxeBcDhcp6GoogleTest.cpp
@@ -4,7 +4,9 @@
   Copyright (c) Microsoft Corporation
   SPDX-License-Identifier: BSD-2-Clause-Patent
 **/
-#include <gtest/gtest.h>
+#include <Library/GoogleTestLib.h>
+#include <GoogleTest/Library/MockUefiLib.h>
+#include <GoogleTest/Library/MockUefiRuntimeServicesTableLib.h>
 
 extern "C" {
   #include <Uefi.h>
@@ -19,7 +21,8 @@ extern "C" {
 // Definitions
 ///////////////////////////////////////////////////////////////////////////////
 
-#define PACKET_SIZE  (1500)
+#define PACKET_SIZE            (1500)
+#define REQUEST_OPTION_LENGTH  (120)
 
 typedef struct {
   UINT16    OptionCode;   // The option code for DHCP6_OPT_SERVER_ID (e.g., 0x03)
@@ -76,6 +79,26 @@ MockConfigure (
 }
 
 // Needed by PxeBcSupport
+EFI_STATUS
+PxeBcDns6 (
+  IN PXEBC_PRIVATE_DATA  *Private,
+  IN     CHAR16          *HostName,
+  OUT EFI_IPv6_ADDRESS   *IpAddress
+  )
+{
+  return EFI_SUCCESS;
+}
+
+UINT32
+PxeBcBuildDhcp6Options (
+  IN  PXEBC_PRIVATE_DATA       *Private,
+  OUT EFI_DHCP6_PACKET_OPTION  **OptList,
+  IN  UINT8                    *Buffer
+  )
+{
+  return EFI_SUCCESS;
+}
+
 EFI_STATUS
 EFIAPI
 QueueDpc (
@@ -159,6 +182,10 @@ TEST_F (PxeBcHandleDhcp6OfferTest, BasicUsageTest) {
   ASSERT_EQ (PxeBcHandleDhcp6Offer (&(PxeBcHandleDhcp6OfferTest::Private)), EFI_DEVICE_ERROR);
 }
 
+///////////////////////////////////////////////////////////////////////////////
+// PxeBcCacheDnsServerAddresses Tests
+///////////////////////////////////////////////////////////////////////////////
+
 class PxeBcCacheDnsServerAddressesTest : public ::testing::Test {
 public:
   PXEBC_PRIVATE_DATA Private = { 0 };
@@ -298,3 +325,250 @@ TEST_F (PxeBcCacheDnsServerAddressesTest, MultipleDnsEntries) {
     FreePool (Private.DnsServer);
   }
 }
+
+///////////////////////////////////////////////////////////////////////////////
+// PxeBcRequestBootServiceTest Test Cases
+///////////////////////////////////////////////////////////////////////////////
+
+class PxeBcRequestBootServiceTest : public ::testing::Test {
+public:
+  PXEBC_PRIVATE_DATA Private = { 0 };
+  EFI_UDP6_PROTOCOL Udp6Read;
+
+protected:
+  // Add any setup code if needed
+  virtual void
+  SetUp (
+    )
+  {
+    Private.Dhcp6Request = (EFI_DHCP6_PACKET *)AllocateZeroPool (PACKET_SIZE);
+
+    // Need to setup the EFI_PXE_BASE_CODE_PROTOCOL
+    // The function under test really only needs the following:
+    //  UdpWrite
+    //  UdpRead
+
+    Private.PxeBc.UdpWrite = (EFI_PXE_BASE_CODE_UDP_WRITE)MockUdpWrite;
+    Private.PxeBc.UdpRead  = (EFI_PXE_BASE_CODE_UDP_READ)MockUdpRead;
+
+    // Need to setup EFI_UDP6_PROTOCOL
+    // The function under test really only needs the following:
+    //  Configure
+
+    Udp6Read.Configure = (EFI_UDP6_CONFIGURE)MockConfigure;
+    Private.Udp6Read   = &Udp6Read;
+  }
+
+  // Add any cleanup code if needed
+  virtual void
+  TearDown (
+    )
+  {
+    if (Private.Dhcp6Request != NULL) {
+      FreePool (Private.Dhcp6Request);
+    }
+
+    // Clean up any resources or variables
+  }
+};
+
+TEST_F (PxeBcRequestBootServiceTest, ServerDiscoverBasicUsageTest) {
+  PxeBcRequestBootServiceTest::Private.OfferBuffer[0].Dhcp6.OfferType = PxeOfferTypeProxyBinl;
+
+  DHCP6_OPTION_SERVER_ID  Server = { 0 };
+
+  Server.OptionCode =  HTONS (DHCP6_OPT_SERVER_ID);
+  Server.OptionLen  = HTONS (16); // valid length
+  UINT8  Index = 0;
+
+  EFI_DHCP6_PACKET  *Packet = (EFI_DHCP6_PACKET *)&Private.OfferBuffer[Index].Dhcp6.Packet.Offer;
+
+  UINT8  *Cursor = (UINT8 *)(Packet->Dhcp6.Option);
+
+  CopyMem (Cursor, &Server, sizeof (Server));
+  Cursor += sizeof (Server);
+
+  // Update the packet length
+  Packet->Length = (UINT16)(Cursor - (UINT8 *)Packet);
+  Packet->Size   = PACKET_SIZE;
+
+  ASSERT_EQ (PxeBcRequestBootService (&(PxeBcRequestBootServiceTest::Private), Index), EFI_SUCCESS);
+}
+
+TEST_F (PxeBcRequestBootServiceTest, AttemptDiscoverOverFlowExpectFailure) {
+  PxeBcRequestBootServiceTest::Private.OfferBuffer[0].Dhcp6.OfferType = PxeOfferTypeProxyBinl;
+
+  DHCP6_OPTION_SERVER_ID  Server = { 0 };
+
+  Server.OptionCode =  HTONS (DHCP6_OPT_SERVER_ID);
+  Server.OptionLen  = HTONS (1500); // This length would overflow without a check
+  UINT8  Index = 0;
+
+  EFI_DHCP6_PACKET  *Packet = (EFI_DHCP6_PACKET *)&Private.OfferBuffer[Index].Dhcp6.Packet.Offer;
+
+  UINT8  *Cursor = (UINT8 *)(Packet->Dhcp6.Option);
+
+  CopyMem (Cursor, &Server, sizeof (Server));
+  Cursor += sizeof (Server);
+
+  // Update the packet length
+  Packet->Length = (UINT16)(Cursor - (UINT8 *)Packet);
+  Packet->Size   = PACKET_SIZE;
+
+  // This is going to be stopped by the duid overflow check
+  ASSERT_EQ (PxeBcRequestBootService (&(PxeBcRequestBootServiceTest::Private), Index), EFI_INVALID_PARAMETER);
+}
+
+TEST_F (PxeBcRequestBootServiceTest, RequestBasicUsageTest) {
+  EFI_DHCP6_PACKET_OPTION  RequestOpt = { 0 }; // the data section doesn't really matter
+
+  RequestOpt.OpCode = HTONS (0x1337);
+  RequestOpt.OpLen  = 0; // valid length
+
+  UINT8  Index = 0;
+
+  EFI_DHCP6_PACKET  *Packet = (EFI_DHCP6_PACKET *)&Private.Dhcp6Request[Index];
+
+  UINT8  *Cursor = (UINT8 *)(Packet->Dhcp6.Option);
+
+  CopyMem (Cursor, &RequestOpt, sizeof (RequestOpt));
+  Cursor += sizeof (RequestOpt);
+
+  // Update the packet length
+  Packet->Length = (UINT16)(Cursor - (UINT8 *)Packet);
+  Packet->Size   = PACKET_SIZE;
+
+  ASSERT_EQ (PxeBcRequestBootService (&(PxeBcRequestBootServiceTest::Private), Index), EFI_SUCCESS);
+}
+
+TEST_F (PxeBcRequestBootServiceTest, AttemptRequestOverFlowExpectFailure) {
+  EFI_DHCP6_PACKET_OPTION  RequestOpt = { 0 }; // the data section doesn't really matter
+
+  RequestOpt.OpCode = HTONS (0x1337);
+  RequestOpt.OpLen  = 1500; // this length would overflow without a check
+
+  UINT8  Index = 0;
+
+  EFI_DHCP6_PACKET  *Packet = (EFI_DHCP6_PACKET *)&Private.Dhcp6Request[Index];
+
+  UINT8  *Cursor = (UINT8 *)(Packet->Dhcp6.Option);
+
+  CopyMem (Cursor, &RequestOpt, sizeof (RequestOpt));
+  Cursor += sizeof (RequestOpt);
+
+  // Update the packet length
+  Packet->Length = (UINT16)(Cursor - (UINT8 *)Packet);
+  Packet->Size   = PACKET_SIZE;
+
+  ASSERT_EQ (PxeBcRequestBootService (&(PxeBcRequestBootServiceTest::Private), Index), EFI_OUT_OF_RESOURCES);
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// PxeBcDhcp6Discover Test
+///////////////////////////////////////////////////////////////////////////////
+
+class PxeBcDhcp6DiscoverTest : public ::testing::Test {
+public:
+  PXEBC_PRIVATE_DATA Private = { 0 };
+  EFI_UDP6_PROTOCOL Udp6Read;
+
+protected:
+  MockUefiRuntimeServicesTableLib RtServicesMock;
+
+  // Add any setup code if needed
+  virtual void
+  SetUp (
+    )
+  {
+    Private.Dhcp6Request = (EFI_DHCP6_PACKET *)AllocateZeroPool (PACKET_SIZE);
+
+    // Need to setup the EFI_PXE_BASE_CODE_PROTOCOL
+    // The function under test really only needs the following:
+    //  UdpWrite
+    //  UdpRead
+
+    Private.PxeBc.UdpWrite = (EFI_PXE_BASE_CODE_UDP_WRITE)MockUdpWrite;
+    Private.PxeBc.UdpRead  = (EFI_PXE_BASE_CODE_UDP_READ)MockUdpRead;
+
+    // Need to setup EFI_UDP6_PROTOCOL
+    // The function under test really only needs the following:
+    //  Configure
+
+    Udp6Read.Configure = (EFI_UDP6_CONFIGURE)MockConfigure;
+    Private.Udp6Read   = &Udp6Read;
+  }
+
+  // Add any cleanup code if needed
+  virtual void
+  TearDown (
+    )
+  {
+    if (Private.Dhcp6Request != NULL) {
+      FreePool (Private.Dhcp6Request);
+    }
+
+    // Clean up any resources or variables
+  }
+};
+
+// Test Description
+// This will cause an overflow by an untrusted packet during the option parsing
+TEST_F (PxeBcDhcp6DiscoverTest, BasicOverflowTest) {
+  EFI_IPv6_ADDRESS         DestIp     = { 0 };
+  EFI_DHCP6_PACKET_OPTION  RequestOpt = { 0 }; // the data section doesn't really matter
+
+  RequestOpt.OpCode = HTONS (0x1337);
+  RequestOpt.OpLen  = HTONS (0xFFFF); // overflow
+
+  UINT8  *Cursor = (UINT8 *)(Private.Dhcp6Request->Dhcp6.Option);
+
+  CopyMem (Cursor, &RequestOpt, sizeof (RequestOpt));
+  Cursor += sizeof (RequestOpt);
+
+  Private.Dhcp6Request->Length = (UINT16)(Cursor - (UINT8 *)Private.Dhcp6Request);
+
+  EXPECT_CALL (RtServicesMock, gRT_GetTime)
+    .WillOnce (::testing::Return (0));
+
+  ASSERT_EQ (
+    PxeBcDhcp6Discover (
+      &(PxeBcDhcp6DiscoverTest::Private),
+      0,
+      NULL,
+      FALSE,
+      (EFI_IP_ADDRESS *)&DestIp
+      ),
+    EFI_OUT_OF_RESOURCES
+    );
+}
+
+// Test Description
+// This will test that we can handle a packet with a valid option length
+TEST_F (PxeBcDhcp6DiscoverTest, BasicUsageTest) {
+  EFI_IPv6_ADDRESS         DestIp     = { 0 };
+  EFI_DHCP6_PACKET_OPTION  RequestOpt = { 0 }; // the data section doesn't really matter
+
+  RequestOpt.OpCode = HTONS (0x1337);
+  RequestOpt.OpLen  = HTONS (0x30);
+
+  UINT8  *Cursor = (UINT8 *)(Private.Dhcp6Request->Dhcp6.Option);
+
+  CopyMem (Cursor, &RequestOpt, sizeof (RequestOpt));
+  Cursor += sizeof (RequestOpt);
+
+  Private.Dhcp6Request->Length = (UINT16)(Cursor - (UINT8 *)Private.Dhcp6Request);
+
+  EXPECT_CALL (RtServicesMock, gRT_GetTime)
+    .WillOnce (::testing::Return (0));
+
+  ASSERT_EQ (
+    PxeBcDhcp6Discover (
+      &(PxeBcDhcp6DiscoverTest::Private),
+      0,
+      NULL,
+      FALSE,
+      (EFI_IP_ADDRESS *)&DestIp
+      ),
+    EFI_SUCCESS
+    );
+}
diff --git a/NetworkPkg/UefiPxeBcDxe/GoogleTest/PxeBcDhcp6GoogleTest.h b/NetworkPkg/UefiPxeBcDxe/GoogleTest/PxeBcDhcp6GoogleTest.h
index 48ab238d..4c6173de 100644
--- a/NetworkPkg/UefiPxeBcDxe/GoogleTest/PxeBcDhcp6GoogleTest.h
+++ b/NetworkPkg/UefiPxeBcDxe/GoogleTest/PxeBcDhcp6GoogleTest.h
@@ -47,4 +47,22 @@ PxeBcCacheDnsServerAddresses (
   IN PXEBC_DHCP6_PACKET_CACHE  *Cache6
   );
 
+/**
+  Build and send out the request packet for the bootfile, and parse the reply.
+
+  @param[in]  Private               The pointer to PxeBc private data.
+  @param[in]  Index                 PxeBc option boot item type.
+
+  @retval     EFI_SUCCESS           Successfully discovered the boot file.
+  @retval     EFI_OUT_OF_RESOURCES  Failed to allocate resources.
+  @retval     EFI_NOT_FOUND         Can't get the PXE reply packet.
+  @retval     Others                Failed to discover the boot file.
+
+**/
+EFI_STATUS
+PxeBcRequestBootService (
+  IN  PXEBC_PRIVATE_DATA  *Private,
+  IN  UINT32              Index
+  );
+
 #endif // PXE_BC_DHCP6_GOOGLE_TEST_H_
-- 
2.33.0