NetworkPkg: Update the PxeBcDhcp6GoogleTest due to underlying changes

This patch updates the PxeBcDhcp6GoogleTest due to the changes in the underlying code. The changes are as follows: - Random now comes from the RngLib Protocol - The TCP ISN is now generated by the hash function Cc: Saloni Kasbekar <[email protected]> Cc: Zachary Clark-williams <[email protected]> Signed-off-by: Doug Flick [MSFT] <[email protected]>
tianocore · May 6, 2024 · 0b67c6b · 0b67c6b
1 parent e69f2a8
commit 0b67c6b
Show file tree

Hide file tree

Showing 3 changed files with 100 additions and 6 deletions.
diff --git a/NetworkPkg/Test/NetworkPkgHostTest.dsc b/NetworkPkg/Test/NetworkPkgHostTest.dsc
@@ -30,6 +30,7 @@
   NetworkPkg/UefiPxeBcDxe/GoogleTest/UefiPxeBcDxeGoogleTest.inf {
     <LibraryClasses>
       UefiRuntimeServicesTableLib|MdePkg/Test/Mock/Library/GoogleTest/MockUefiRuntimeServicesTableLib/MockUefiRuntimeServicesTableLib.inf
+      UefiBootServicesTableLib|MdePkg/Test/Mock/Library/GoogleTest/MockUefiBootServicesTableLib/MockUefiBootServicesTableLib.inf
   }
 
 # Despite these library classes being listed in [LibraryClasses] below, they are not needed for the host-based unit tests.

diff --git a/NetworkPkg/UefiPxeBcDxe/GoogleTest/PxeBcDhcp6GoogleTest.cpp b/NetworkPkg/UefiPxeBcDxe/GoogleTest/PxeBcDhcp6GoogleTest.cpp
@@ -7,6 +7,8 @@
 #include <Library/GoogleTestLib.h>
 #include <GoogleTest/Library/MockUefiLib.h>
 #include <GoogleTest/Library/MockUefiRuntimeServicesTableLib.h>
+#include <GoogleTest/Library/MockUefiBootServicesTableLib.h>
+#include <GoogleTest/Protocol/MockRng.h>
 
 extern "C" {
   #include <Uefi.h>
@@ -165,7 +167,7 @@ class PxeBcHandleDhcp6OfferTest : public ::testing::Test {
 // Note:
 // Testing PxeBcHandleDhcp6Offer() is difficult because it depends on a
 // properly setup Private structure. Attempting to properly test this function
-// without a signficant refactor is a fools errand. Instead, we will test
+// without a significant refactor is a fools errand. Instead, we will test
 // that we can prevent an overflow in the function.
 TEST_F (PxeBcHandleDhcp6OfferTest, BasicUsageTest) {
   PXEBC_DHCP6_PACKET_CACHE  *Cache6 = NULL;
@@ -238,6 +240,7 @@ TEST_F (PxeBcCacheDnsServerAddressesTest, BasicUsageTest) {
     FreePool (Option);
   }
 }
+
 // Test Description
 // Test that we can prevent an overflow in the function
 TEST_F (PxeBcCacheDnsServerAddressesTest, AttemptOverflowTest) {
@@ -470,10 +473,15 @@ TEST_F (PxeBcRequestBootServiceTest, AttemptRequestOverFlowExpectFailure) {
 class PxeBcDhcp6DiscoverTest : public ::testing::Test {
 public:
   PXEBC_PRIVATE_DATA Private = { 0 };
+  // create a mock md5 hash
+  UINT8 Md5Hash[16] = { 0 };
+
   EFI_UDP6_PROTOCOL Udp6Read;
 
 protected:
   MockUefiRuntimeServicesTableLib RtServicesMock;
+  MockUefiBootServicesTableLib BsMock;
+  MockRng RngMock;
 
   // Add any setup code if needed
   virtual void
@@ -527,8 +535,21 @@ TEST_F (PxeBcDhcp6DiscoverTest, BasicOverflowTest) {
 
   Private.Dhcp6Request->Length = (UINT16)(Cursor - (UINT8 *)Private.Dhcp6Request);
 
-  EXPECT_CALL (RtServicesMock, gRT_GetTime)
-    .WillOnce (::testing::Return (0));
+  EXPECT_CALL (BsMock, gBS_LocateProtocol)
+    .WillOnce (
+       ::testing::DoAll (
+                    ::testing::SetArgPointee<2> (::testing::ByRef (gRngProtocol)),
+                    ::testing::Return (EFI_SUCCESS)
+                    )
+       );
+
+  EXPECT_CALL (RngMock, GetRng)
+    .WillOnce (
+       ::testing::DoAll (
+                    ::testing::SetArgPointee<3> (::testing::ByRef (Md5Hash[0])),
+                    ::testing::Return (EFI_SUCCESS)
+                    )
+       );
 
   ASSERT_EQ (
     PxeBcDhcp6Discover (
@@ -558,8 +579,21 @@ TEST_F (PxeBcDhcp6DiscoverTest, BasicUsageTest) {
 
   Private.Dhcp6Request->Length = (UINT16)(Cursor - (UINT8 *)Private.Dhcp6Request);
 
-  EXPECT_CALL (RtServicesMock, gRT_GetTime)
-    .WillOnce (::testing::Return (0));
+  EXPECT_CALL (BsMock, gBS_LocateProtocol)
+    .WillOnce (
+       ::testing::DoAll (
+                    ::testing::SetArgPointee<2> (::testing::ByRef (gRngProtocol)),
+                    ::testing::Return (EFI_SUCCESS)
+                    )
+       );
+
+  EXPECT_CALL (RngMock, GetRng)
+    .WillOnce (
+       ::testing::DoAll (
+                    ::testing::SetArgPointee<3> (::testing::ByRef (Md5Hash[0])),
+                    ::testing::Return (EFI_SUCCESS)
+                    )
+       );
 
   ASSERT_EQ (
     PxeBcDhcp6Discover (
@@ -572,3 +606,61 @@ TEST_F (PxeBcDhcp6DiscoverTest, BasicUsageTest) {
     EFI_SUCCESS
     );
 }
+
+TEST_F (PxeBcDhcp6DiscoverTest, MultipleRequestsAttemptOverflow) {
+  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 (REQUEST_OPTION_LENGTH); // this length would overflow without a check
+  UINT8  RequestOptBuffer[REQUEST_OPTION_LENGTH] = { 0 };
+
+  // make sure we have enough space for 10 of these options
+  ASSERT_TRUE (REQUEST_OPTION_LENGTH * 10 <= PACKET_SIZE);
+
+  UINT8             Index   = 0;
+  EFI_DHCP6_PACKET  *Packet = (EFI_DHCP6_PACKET *)&Private.Dhcp6Request[Index];
+  UINT8             *Cursor = (UINT8 *)(Packet->Dhcp6.Option);
+
+  // let's add 10 of these options - this should overflow
+  for (UINT8 i = 0; i < 10; i++) {
+    CopyMem (Cursor, &RequestOpt, sizeof (RequestOpt));
+    Cursor += sizeof (RequestOpt) - 1;
+    CopyMem (Cursor, RequestOptBuffer, REQUEST_OPTION_LENGTH);
+    Cursor += REQUEST_OPTION_LENGTH;
+  }
+
+  // Update the packet length
+  Packet->Length = (UINT16)(Cursor - (UINT8 *)Packet);
+  Packet->Size   = PACKET_SIZE;
+
+  // Make sure we're larger than the buffer we're trying to write into
+  ASSERT_TRUE (Packet->Length > sizeof (EFI_PXE_BASE_CODE_DHCPV6_PACKET));
+
+  EXPECT_CALL (BsMock, gBS_LocateProtocol)
+    .WillOnce (
+       ::testing::DoAll (
+                    ::testing::SetArgPointee<2> (::testing::ByRef (gRngProtocol)),
+                    ::testing::Return (EFI_SUCCESS)
+                    )
+       );
+
+  EXPECT_CALL (RngMock, GetRng)
+    .WillOnce (
+       ::testing::DoAll (
+                    ::testing::SetArgPointee<3> (::testing::ByRef (Md5Hash[0])),
+                    ::testing::Return (EFI_SUCCESS)
+                    )
+       );
+
+  ASSERT_EQ (
+    PxeBcDhcp6Discover (
+      &(PxeBcDhcp6DiscoverTest::Private),
+      0,
+      NULL,
+      FALSE,
+      (EFI_IP_ADDRESS *)&DestIp
+      ),
+    EFI_OUT_OF_RESOURCES
+    );
+}
diff --git a/NetworkPkg/UefiPxeBcDxe/GoogleTest/UefiPxeBcDxeGoogleTest.inf b/NetworkPkg/UefiPxeBcDxe/GoogleTest/UefiPxeBcDxeGoogleTest.inf
@@ -14,7 +14,7 @@ VERSION_STRING = 1.0
 #
 # The following information is for reference only and not required by the build tools.
 #
-#  VALID_ARCHITECTURES           = IA32 X64
+#  VALID_ARCHITECTURES           = IA32 X64 AARCH64
 #
 
 [Sources]
@@ -23,6 +23,7 @@ VERSION_STRING = 1.0
   PxeBcDhcp6GoogleTest.h
   ../PxeBcDhcp6.c
   ../PxeBcSupport.c
+  ../../../MdePkg/Test/Mock/Library/GoogleTest/Protocol/MockRng.cpp
 
 [Packages]
   MdePkg/MdePkg.dec