Merge e333559c03 into cc9a0425a2

ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c

@@ -19,6 +19,7 @@
 #include <Library/UefiBootServicesTableLib.h>
 #include <Library/UefiLib.h>
 #include <Library/UefiRuntimeServicesTableLib.h>
+#include <Library/SafeIntLib.h>
 
 #include <Protocol/MmCommunication2.h>
 #include <Protocol/MmCommunication3.h>
@@ -209,6 +210,7 @@ MmCommunicationCommon (
   EFI_MM_COMMUNICATE_HEADER     *CommunicateHeader;
   EFI_MM_COMMUNICATE_HEADER_V3  *CommunicateHeaderV3;
   UINTN                         BufferSize;
+  UINTN                         InputBufferSize;
   UINTN                         *MessageSize;
   UINTN                         HeaderSize;
   EFI_STATUS                    Status;
@@ -239,15 +241,30 @@ MmCommunicationCommon (
     BufferSize          = CommunicateHeaderV3->BufferSize;
     MessageSize         = &CommunicateHeaderV3->MessageSize;
     HeaderSize          = sizeof (EFI_MM_COMMUNICATE_HEADER_V3);
+
+    if (BufferSize < HeaderSize) {
+      return EFI_INVALID_PARAMETER;
+    }
+
+    if (BufferSize - HeaderSize < *MessageSize) {
+      return EFI_INVALID_PARAMETER;
+    }
   } else {
-    BufferSize = CommunicateHeader->MessageLength +
+    Status = SafeUintnAdd (
-                 sizeof (CommunicateHeader->HeaderGuid) +
+               CommunicateHeader->MessageLength,
-                 sizeof (CommunicateHeader->MessageLength);
+               OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data),
+               &BufferSize
+               );
+    if (EFI_ERROR (Status)) {
+      return EFI_INVALID_PARAMETER;
+    }
+
     MessageSize = &CommunicateHeader->MessageLength;
-    HeaderSize  = sizeof (CommunicateHeader->HeaderGuid) +
+    HeaderSize  = OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data);
-                  sizeof (CommunicateHeader->MessageLength);
   }
 
+  InputBufferSize = BufferSize;
+
   // If CommSize is not omitted, perform size inspection before proceeding.
   if (CommSize != NULL) {
     // This case can be used by the consumer of this driver to find out the
@@ -265,6 +282,8 @@ MmCommunicationCommon (
     if (*CommSize < BufferSize) {
       Status = EFI_INVALID_PARAMETER;
     }
+
+    InputBufferSize = *CommSize;
   }
 
   //
@@ -293,7 +312,6 @@ MmCommunicationCommon (
   }
 
   if (!EFI_ERROR (Status)) {
-    ZeroMem (CommBufferVirtual, BufferSize);
     // On successful return, the size of data being returned is inferred from
     // MessageLength + Header.
     CommunicateHeader = (EFI_MM_COMMUNICATE_HEADER *)mNsCommBuffMemRegion.VirtualBase;
@@ -316,12 +334,17 @@ MmCommunicationCommon (
       CommunicateHeaderV3 = (EFI_MM_COMMUNICATE_HEADER_V3 *)CommunicateHeader;
       BufferSize          = CommunicateHeaderV3->BufferSize;
     } else {
-      BufferSize = CommunicateHeader->MessageLength +
+      Status = SafeUintnAdd (
-                   sizeof (CommunicateHeader->HeaderGuid) +
+                 CommunicateHeader->MessageLength,
-                   sizeof (CommunicateHeader->MessageLength);
+                 OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data),
+                 &BufferSize
+                 );
+      if (EFI_ERROR (Status)) {
+        return EFI_INVALID_PARAMETER;
+      }
     }
 
-    if (BufferSize > mNsCommBuffMemRegion.Length) {
+    if (BufferSize > InputBufferSize) {
       // Something bad has happened, we should have landed in ARM_SMC_MM_RET_NO_MEMORY
       Status = EFI_BAD_BUFFER_SIZE;
       DEBUG ((
@@ -329,7 +352,7 @@ MmCommunicationCommon (
         "%a Returned buffer exceeds communication buffer limit. Has: 0x%llx vs. max: 0x%llx!\n",
         __func__,
         BufferSize,
-        (UINTN)mNsCommBuffMemRegion.Length
+        InputBufferSize
         ));
     } else {
       CopyMem (

ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.inf

@@ -39,6 +39,7 @@
   DxeServicesTableLib
   HobLib
   UefiDriverEntryPoint
+  SafeIntLib
 
 [Protocols]
   gEfiDxeMmReadyToLockProtocolGuid              ## UNDEFINED # SmiHandlerRegister

ArmPkg/Drivers/MmCommunicationPei/MmCommunicationPei.c

@@ -22,6 +22,7 @@
 #include <Library/PcdLib.h>
 #include <Library/PeimEntryPoint.h>
 #include <Library/PeiServicesLib.h>
+#include <Library/SafeIntLib.h>
 
 //
 // Partition ID if FF-A support is enabled
@@ -323,6 +324,8 @@ MmCommunicationPeimCommon (
   EFI_MM_COMMUNICATE_HEADER_V3  *CommunicateHeaderV3;
   EFI_STATUS                    Status;
   UINTN                         BufferSize;
+  UINTN                         InputBufferSize;
+  UINTN                         HeaderSize;
 
   //
   // Check parameters
@@ -341,7 +344,28 @@ MmCommunicationPeimCommon (
   {
     // This is a v3 header
     CommunicateHeaderV3 = (EFI_MM_COMMUNICATE_HEADER_V3 *)(UINTN)CommBuffer;
-    BufferSize          = CommunicateHeaderV3->BufferSize;
+    HeaderSize          = sizeof (EFI_MM_COMMUNICATE_HEADER_V3);
+    InputBufferSize     = CommunicateHeaderV3->BufferSize;
+
+    if (InputBufferSize < HeaderSize) {
+      DEBUG ((
+        DEBUG_ERROR,
+        "%a Invalid BufferSize value 0x%llx!\n",
+        __func__,
+        InputBufferSize
+        ));
+      return EFI_INVALID_PARAMETER;
+    }
+
+    if (InputBufferSize - HeaderSize < CommunicateHeaderV3->MessageSize) {
+      DEBUG ((
+        DEBUG_ERROR,
+        "%a Invalid BufferSize value 0x%llx!\n",
+        __func__,
+        InputBufferSize
+        ));
+      return EFI_INVALID_PARAMETER;
+    }
   } else {
     // This is a v1 header, do some checks
     if (CommSize == NULL) {
@@ -373,28 +397,46 @@ MmCommunicationPeimCommon (
     // MessageLength + Header to ascertain the
     // total size of the communication payload rather than
     // rely on optional CommSize parameter
-    BufferSize = CommunicateHeader->MessageLength +
+    Status = SafeUintnAdd (CommunicateHeader->MessageLength, OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data), &InputBufferSize);
-                 sizeof (CommunicateHeader->HeaderGuid) +
+    if (EFI_ERROR (Status)) {
-                 sizeof (CommunicateHeader->MessageLength);
+      DEBUG ((
+        DEBUG_ERROR,
+        "%a Overflow occurred while calculating input BufferSize!\n",
+        __func__
+        ));
+      return Status;
+    }
+
     //
     // If CommSize is supplied it must match MessageLength + sizeof (EFI_MM_COMMUNICATE_HEADER);
     //
-    if (*CommSize != BufferSize) {
+    if (*CommSize != InputBufferSize) {
       DEBUG ((
         DEBUG_ERROR,
         "%a Unexpected CommSize value, has: 0x%llx vs. expected: 0x%llx!\n",
         __func__,
         *CommSize,
-        BufferSize
+        InputBufferSize
         ));
       return EFI_INVALID_PARAMETER;
     }
   }
 
+  if (InputBufferSize > (UINTN)PcdGet64 (PcdMmBufferSize)) {
+    DEBUG ((
+      DEBUG_ERROR,
+      "%a Input buffer exceeds communication buffer limit. Has: 0x%llx vs. max: 0x%llx!\n",
+      __func__,
+      InputBufferSize,
+      (UINTN)PcdGet64 (PcdMmBufferSize)
+      ));
+    return EFI_BAD_BUFFER_SIZE;
+  }
+
   // Now we know that the size is something we can handle, copy it over to the designated comm buffer.
   CommunicateHeader = (EFI_MM_COMMUNICATE_HEADER *)(UINTN)(PcdGet64 (PcdMmBufferBase));
 
-  CopyMem (CommunicateHeader, CommBuffer, BufferSize);
+  CopyMem (CommunicateHeader, CommBuffer, InputBufferSize);
   if (IsFfaSupported ()) {
     Status = SendFfaMmCommunicate ();
   } else {
@@ -423,20 +465,26 @@ MmCommunicationPeimCommon (
       CommunicateHeaderV3 = (EFI_MM_COMMUNICATE_HEADER_V3 *)CommunicateHeader;
       BufferSize          = CommunicateHeaderV3->BufferSize;
     } else {
-      BufferSize = CommunicateHeader->MessageLength +
+      Status = SafeUintnAdd (CommunicateHeader->MessageLength, OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data), &BufferSize);
-                   sizeof (CommunicateHeader->HeaderGuid) +
+      if (EFI_ERROR (Status)) {
-                   sizeof (CommunicateHeader->MessageLength);
+        DEBUG ((
+          DEBUG_ERROR,
+          "%a Overflow occurred while calculating returned BufferSize!\n",
+          __func__
+          ));
+        return Status;
+      }
     }
 
-    if (BufferSize > (UINTN)PcdGet64 (PcdMmBufferSize)) {
+    if (InputBufferSize < BufferSize) {
       // Something bad has happened, we should have landed in ARM_SMC_MM_RET_NO_MEMORY
       Status = EFI_BAD_BUFFER_SIZE;
       DEBUG ((
         DEBUG_ERROR,
-        "%a Returned buffer exceeds communication buffer limit. Has: 0x%llx vs. max: 0x%llx!\n",
+        "%a Returned buffer size is larger than input buffer size. Input: 0x%llx vs. returned: 0x%llx!\n",
         __func__,
-        BufferSize,
+        InputBufferSize,
-        (UINTN)PcdGet64 (PcdMmBufferSize)
+        BufferSize
         ));
     } else {
       CopyMem (CommBuffer, CommunicateHeader, BufferSize);

ArmPkg/Drivers/MmCommunicationPei/MmCommunicationPei.inf

@@ -29,6 +29,7 @@
   PeimEntryPoint
   PeiServicesLib
   HobLib
+  SafeIntLib
 
 [Pcd]
   gArmTokenSpaceGuid.PcdMmBufferBase

ArmPkg/Library/ArmStandaloneMmCoreEntryPoint/ArmStandaloneMmCoreEntryPoint.c

@@ -527,6 +527,10 @@ ValidateMmCommBufferAddr (
         &gEfiMmCommunicateHeaderV3Guid
         ))
   {
+    if (CommBufferRange < sizeof (EFI_MM_COMMUNICATE_HEADER_V3)) {
+      return EFI_ACCESS_DENIED;
+    }
+
     CommBufferHeaderV3 = (EFI_MM_COMMUNICATE_HEADER_V3 *)CommBufferAddr;
     Status             = SafeUint64Add (
                            CommBufferHeaderV3->MessageSize,
@@ -536,9 +540,21 @@ ValidateMmCommBufferAddr (
     if (EFI_ERROR (Status)) {
       return EFI_ACCESS_DENIED;
     }
+
+    if (BufferSize > CommBufferHeaderV3->BufferSize) {
+      return EFI_ACCESS_DENIED;
+    }
+
+    BufferSize = CommBufferHeaderV3->BufferSize;
   } else {
-    BufferSize = ((EFI_MM_COMMUNICATE_HEADER *)CommBufferAddr)->MessageLength +
+    Status = SafeUint64Add (
-                 OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data);
+               ((EFI_MM_COMMUNICATE_HEADER *)CommBufferAddr)->MessageLength,
+               OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data),
+               &BufferSize
+               );
+    if (EFI_ERROR (Status)) {
+      return EFI_ACCESS_DENIED;
+    }
   }
 
   Status = SafeUint64Add (

ArmPkg/Library/ArmStandaloneMmCoreEntryPoint/ArmStandaloneMmCoreEntryPoint.inf

@@ -44,6 +44,7 @@
   ArmFfaLib
   StackCheckLib
   HobLib
+  SafeIntLib
 
 [Guids]
   gMpInformationHobGuid

StandaloneMmPkg/Core/StandaloneMmCore.c

@@ -513,6 +513,8 @@ MmEntryPoint (
   EFI_GUID                      *CommGuid;
   UINTN                         CommGuidOffset;
   UINTN                         CommHeaderSize;
+  EFI_STATUS                    SafeIntStatus;
+  UINTN                         MaxBufferSize;
 
   //
   // Update MMST using the context
@@ -560,10 +562,31 @@ MmEntryPoint (
         LegacyCommunicateHeader = (EFI_MM_COMMUNICATE_HEADER *)(UINTN)mMmCommunicationBuffer->PhysicalStart;
         CommGuidOffset          = OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, HeaderGuid);
         CommHeaderSize          = OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data);
-        BufferSize              = OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data) + LegacyCommunicateHeader->MessageLength;
+
+        SafeIntStatus = SafeUintnAdd (
+                          OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data),
+                          LegacyCommunicateHeader->MessageLength,
+                          &BufferSize
+                          );
+        if (EFI_ERROR (SafeIntStatus)) {
+          DEBUG ((DEBUG_ERROR, "Failed to calculate buffer size: %r\n", SafeIntStatus));
+          ASSERT_EFI_ERROR (SafeIntStatus);
+          return;
+        }
       }
 
-      if (BufferSize <= EFI_PAGES_TO_SIZE (mMmCommunicationBuffer->NumberOfPages)) {
+      SafeIntStatus = SafeUintnMult (
+                        EFI_PAGE_SIZE,
+                        mMmCommunicationBuffer->NumberOfPages,
+                        &MaxBufferSize
+                        );
+      if (EFI_ERROR (SafeIntStatus)) {
+        DEBUG ((DEBUG_ERROR, "Failed to convert number of pages to bytes: %r\n", SafeIntStatus));
+        ASSERT_EFI_ERROR (SafeIntStatus);
+        return;
+      }
+
+      if (BufferSize <= MaxBufferSize) {
         //
         // Shadow the data from MM Communication Buffer to internal buffer
         //
@@ -574,19 +597,39 @@ MmEntryPoint (
           );
         ZeroMem (
           (UINT8 *)mInternalCommBufferCopy + BufferSize,
-          EFI_PAGES_TO_SIZE (mMmCommunicationBuffer->NumberOfPages) - BufferSize
+          MaxBufferSize - BufferSize
           );
 
-        BufferSize -= CommHeaderSize;
+        SafeIntStatus = SafeUintnSub (
-        Status      = MmiManage (
+                          BufferSize,
-                        (EFI_GUID *)((UINT8 *)mInternalCommBufferCopy + CommGuidOffset),
+                          CommHeaderSize,
-                        NULL,
+                          &BufferSize
-                        (UINT8 *)mInternalCommBufferCopy + CommHeaderSize,
+                          );
-                        &BufferSize
+        if (EFI_ERROR (SafeIntStatus)) {
-                        );
+          DEBUG ((DEBUG_ERROR, "Failed to subtract header from buffer size: %r\n", SafeIntStatus));
+          ASSERT_EFI_ERROR (SafeIntStatus);
+          return;
+        }
 
-        BufferSize = BufferSize + CommHeaderSize;
+        Status = MmiManage (
-        if (BufferSize <= EFI_PAGES_TO_SIZE (mMmCommunicationBuffer->NumberOfPages)) {
+                   (EFI_GUID *)((UINT8 *)mInternalCommBufferCopy + CommGuidOffset),
+                   NULL,
+                   (UINT8 *)mInternalCommBufferCopy + CommHeaderSize,
+                   &BufferSize
+                   );
+
+        SafeIntStatus = SafeUintnAdd (
+                          BufferSize,
+                          CommHeaderSize,
+                          &BufferSize
+                          );
+        if (EFI_ERROR (SafeIntStatus)) {
+          DEBUG ((DEBUG_ERROR, "Failed to calculate total buffer size: %r\n", SafeIntStatus));
+          ASSERT_EFI_ERROR (SafeIntStatus);
+          return;
+        }
+
+        if (BufferSize <= MaxBufferSize) {
           //
           // Copy the data back to MM Communication Buffer
           //

StandaloneMmPkg/Core/StandaloneMmCore.h

@@ -47,6 +47,7 @@
 #include <Library/StandaloneMmMemLib.h>
 #include <Library/HobLib.h>
 #include <Library/PerformanceLib.h>
+#include <Library/SafeIntLib.h>
 
 #include "StandaloneMmCorePrivateData.h"
 

StandaloneMmPkg/Core/StandaloneMmCore.inf

@@ -54,6 +54,7 @@
   HobPrintLib
   ImagePropertiesRecordLib
   PerformanceLib
+  SafeIntLib
 
 [Protocols]
   gEfiDxeMmReadyToLockProtocolGuid             ## UNDEFINED # SmiHandlerRegister

StandaloneMmPkg/Drivers/MmCommunicationDxe/GoogleTest/MmCommunicationDxeGoogleTest.cpp

@@ -0,0 +1,206 @@
+/** @file
+  GoogleTest for MmCommunicationDxe.
+
+  Copyright (c) 2024, Intel Corporation. All rights reserved.<BR>
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+**/
+
+#include <Library/GoogleTestLib.h>
+
+extern "C" {
+  #include <Uefi.h>
+  #include <Library/BaseLib.h>
+  #include <Library/BaseMemoryLib.h>
+  #include <Protocol/MmControl.h>
+  #include <Protocol/SmmControl2.h>
+  #include <Protocol/MmCommunication.h>
+  #include <Pi/PiMultiPhase.h>
+  #include <Guid/MmCommBuffer.h>
+
+  //
+  // Declare the function under test (defined in MmCommunicationDxe.c)
+  //
+  EFI_STATUS
+  EFIAPI
+  ProcessCommunicationBuffer (
+    IN OUT VOID   *CommBuffer,
+    IN OUT UINTN  *CommSize OPTIONAL
+    );
+
+  //
+  // Globals defined in MmCommunicationDxe.c that we need to set up
+  //
+  extern MM_COMM_BUFFER             mMmCommonBuffer;
+  extern EFI_SMM_CONTROL2_PROTOCOL  *mSmmControl2;
+}
+
+////////////////////////////////////////////////////////////////////////
+// Symbol Definitions
+// These functions are not directly under test - but required to compile
+////////////////////////////////////////////////////////////////////////
+
+//
+// Mock Trigger function for SmmControl2 protocol
+//
+static EFI_STATUS EFIAPI
+MockTrigger (
+  IN CONST EFI_MM_CONTROL_PROTOCOL  *This,
+  IN OUT UINT8                      *CommandPort       OPTIONAL,
+  IN OUT UINT8                      *DataPort          OPTIONAL,
+  IN BOOLEAN                        Periodic           OPTIONAL,
+  IN UINTN                          ActivationInterval OPTIONAL
+  )
+{
+  return EFI_SUCCESS;
+}
+
+//
+// Mock Deactivate (Clear) function
+//
+static EFI_STATUS EFIAPI
+MockClear (
+  IN CONST EFI_MM_CONTROL_PROTOCOL  *This,
+  IN BOOLEAN                        Periodic OPTIONAL
+  )
+{
+  return EFI_SUCCESS;
+}
+
+////////////////////////////////////////////////////////////////////////
+// Defines
+////////////////////////////////////////////////////////////////////////
+
+#define COMM_BUFFER_PAGES  4  // 16 KiB common buffer
+#define COMM_BUFFER_SIZE   EFI_PAGES_TO_SIZE (COMM_BUFFER_PAGES)
+
+////////////////////////////////////////////////////////////////////////
+// MmCommunicationOverflowTest Tests
+////////////////////////////////////////////////////////////////////////
+
+class MmCommunicationOverflowTest : public ::testing::Test {
+public:
+  UINT8 mCommonBuffer[COMM_BUFFER_SIZE];
+  MM_COMM_BUFFER_STATUS mCommonBufferStatus;
+  EFI_MM_CONTROL_PROTOCOL mMockSmmControl2;
+  UINT8 mCommBuffer[COMM_BUFFER_SIZE];
+
+protected:
+  virtual void
+  SetUp (
+    )
+  {
+    ZeroMem (mCommonBuffer, sizeof (mCommonBuffer));
+    ZeroMem (&mCommonBufferStatus, sizeof (mCommonBufferStatus));
+    ZeroMem (mCommBuffer, sizeof (mCommBuffer));
+
+    // Set up the global MM common buffer struct
+    mMmCommonBuffer.PhysicalStart = (EFI_PHYSICAL_ADDRESS)(UINTN)mCommonBuffer;
+    mMmCommonBuffer.NumberOfPages = COMM_BUFFER_PAGES;
+    mMmCommonBuffer.Status        = (EFI_PHYSICAL_ADDRESS)(UINTN)&mCommonBufferStatus;
+
+    // Set up the mock SmmControl2 protocol
+    mMockSmmControl2.Trigger              = MockTrigger;
+    mMockSmmControl2.Clear                = MockClear;
+    mMockSmmControl2.MinimumTriggerPeriod = 0;
+    mSmmControl2                          = &mMockSmmControl2;
+  }
+
+  virtual void
+  TearDown (
+    )
+  {
+    // Clean up any resources or variables
+  }
+};
+
+// Test Description:
+// Normal V1 path works with small MessageLength
+TEST_F (MmCommunicationOverflowTest, V1NormalMessageLengthSucceeds) {
+  EFI_MM_COMMUNICATE_HEADER  *Header = (EFI_MM_COMMUNICATE_HEADER *)mCommBuffer;
+
+  ZeroMem (&Header->HeaderGuid, sizeof (EFI_GUID));
+  Header->MessageLength = 64;
+
+  UINTN  CommSize = OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data) + 64;
+
+  EFI_STATUS  Status = ProcessCommunicationBuffer (mCommBuffer, &CommSize);
+
+  // BufferSize = 24 + 64 = 88, well within 16KiB
+  ASSERT_EQ (Status, EFI_SUCCESS);
+}
+
+// Test Description:
+// V1 path - CWE-190 integer overflow wraps BufferSize to 0.
+// MessageLength = MAX_UINT64 - OFFSET_OF(...) + 1 causes addition to wrap to 0.
+TEST_F (MmCommunicationOverflowTest, V1MessageLengthOverflowWrapsToZero) {
+  EFI_MM_COMMUNICATE_HEADER  *Header = (EFI_MM_COMMUNICATE_HEADER *)mCommBuffer;
+
+  ZeroMem (&Header->HeaderGuid, sizeof (EFI_GUID));
+
+  //
+  // Craft MessageLength so that OFFSET_OF(Data) + MessageLength = 0 (wraps)
+  // OFFSET_OF(EFI_MM_COMMUNICATE_HEADER, Data) = 0x18 (24)
+  // MAX_UINT64 - 0x18 + 1 = 0xFFFFFFFFFFFFFFE8
+  //
+  Header->MessageLength = MAX_UINT64 - OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data) + 1;
+
+  EFI_STATUS  Status = ProcessCommunicationBuffer (mCommBuffer, NULL);
+
+  // Without fix: BufferSize wraps to 0, passes bounds check, returns EFI_SUCCESS
+  ASSERT_EQ (Status, EFI_INVALID_PARAMETER);
+}
+
+// Test Description:
+// V1 path - Integer overflow wraps BufferSize to a small nonzero value (0x08)
+TEST_F (MmCommunicationOverflowTest, V1MessageLengthOverflowWrapsToSmallValue) {
+  EFI_MM_COMMUNICATE_HEADER  *Header = (EFI_MM_COMMUNICATE_HEADER *)mCommBuffer;
+
+  ZeroMem (&Header->HeaderGuid, sizeof (EFI_GUID));
+
+  //
+  // MessageLength = MAX_UINT64 - 0x18 + 1 + 8 = MAX_UINT64 - 0x0F
+  // Result: OFFSET_OF(Data) + MessageLength = 0x18 + (MAX_UINT64 - 0x0F) = 0x08 (wraps)
+  //
+  Header->MessageLength = MAX_UINT64 - OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data) + 1 + 8;
+
+  EFI_STATUS  Status = ProcessCommunicationBuffer (mCommBuffer, NULL);
+
+  // Without fix: BufferSize = 8, passes bounds check, returns EFI_SUCCESS
+  ASSERT_EQ (Status, EFI_INVALID_PARAMETER);
+}
+
+// Test Description:
+// V1 path - Large MessageLength without overflow is rejected by bounds check
+TEST_F (MmCommunicationOverflowTest, V1LargeMessageLengthWithoutOverflowIsRejected) {
+  EFI_MM_COMMUNICATE_HEADER  *Header = (EFI_MM_COMMUNICATE_HEADER *)mCommBuffer;
+
+  ZeroMem (&Header->HeaderGuid, sizeof (EFI_GUID));
+
+  //
+  // Set MessageLength larger than common buffer but not large enough to overflow.
+  // BufferSize = 0x18 + 0x10000 = 0x10018, which exceeds 16KiB (0x4000).
+  //
+  Header->MessageLength = 0x10000;
+
+  EFI_STATUS  Status = ProcessCommunicationBuffer (mCommBuffer, NULL);
+
+  ASSERT_EQ (Status, EFI_INVALID_PARAMETER);
+}
+
+// Test Description:
+// NULL CommBuffer is rejected
+TEST_F (MmCommunicationOverflowTest, NullCommBufferReturnsInvalidParameter) {
+  EFI_STATUS  Status = ProcessCommunicationBuffer (NULL, NULL);
+
+  ASSERT_EQ (Status, EFI_INVALID_PARAMETER);
+}
+
+int
+main (
+  int   argc,
+  char  *argv[]
+  )
+{
+  testing::InitGoogleTest (&argc, argv);
+  return RUN_ALL_TESTS ();
+}

StandaloneMmPkg/Drivers/MmCommunicationDxe/GoogleTest/MmCommunicationDxeGoogleTest.inf

@@ -0,0 +1,51 @@
+## @file
+# GoogleTest for MmCommunicationDxe.
+#
+# Copyright (c) 2024, Intel Corporation. All rights reserved.<BR>
+# SPDX-License-Identifier: BSD-2-Clause-Patent
+##
+
+[Defines]
+  INF_VERSION         = 0x0001001A
+  BASE_NAME           = MmCommunicationDxeGoogleTest
+  FILE_GUID           = 5bf153cc-013c-4c61-95b7-cd86b263bd9f
+  VERSION_STRING      = 1.0
+  MODULE_TYPE         = HOST_APPLICATION
+
+[Sources]
+  MmCommunicationDxeGoogleTest.cpp
+  ../MmCommunicationDxe.c
+
+[Packages]
+  MdePkg/MdePkg.dec
+  MdeModulePkg/MdeModulePkg.dec
+  UefiCpuPkg/UefiCpuPkg.dec
+  StandaloneMmPkg/StandaloneMmPkg.dec
+  UnitTestFrameworkPkg/UnitTestFrameworkPkg.dec
+
+[LibraryClasses]
+  GoogleTestLib
+  BaseLib
+  BaseMemoryLib
+  DebugLib
+  HobLib
+  MemoryAllocationLib
+  ReportStatusCodeLib
+  UefiBootServicesTableLib
+  UefiLib
+  UefiRuntimeLib
+  SafeIntLib
+[Protocols]
+  gEfiMmCommunication3ProtocolGuid
+  gEfiMmCommunication2ProtocolGuid
+  gEfiMmCommunicationProtocolGuid
+  gEfiSmmControl2ProtocolGuid
+  gEfiSmmAccess2ProtocolGuid
+
+[Guids]
+  gMmCommBufferHobGuid
+  gEfiEventVirtualAddressChangeGuid
+  gEfiMmCommunicateHeaderV3Guid
+
+[BuildOptions]
+  MSFT:*_*_*_CC_FLAGS = /EHsc

StandaloneMmPkg/Drivers/MmCommunicationDxe/MmCommunicationDxe.c

@@ -102,13 +102,18 @@ ProcessCommunicationBuffer (
   CommunicateHeader = (EFI_MM_COMMUNICATE_HEADER *)CommBuffer;
   if (CompareGuid (&CommunicateHeader->HeaderGuid, &gEfiMmCommunicateHeaderV3Guid)) {
     CommunicateHeaderV3 = (EFI_MM_COMMUNICATE_HEADER_V3 *)CommBuffer;
-    if (CommunicateHeaderV3->BufferSize < sizeof (EFI_MM_COMMUNICATE_HEADER_V3) + CommunicateHeaderV3->MessageSize) {
+
+    Status = SafeUintnAdd (sizeof (EFI_MM_COMMUNICATE_HEADER_V3), CommunicateHeaderV3->MessageSize, &BufferSize);
+    if (EFI_ERROR (Status) || (CommunicateHeaderV3->BufferSize < BufferSize)) {
       return EFI_INVALID_PARAMETER;
     }
 
     BufferSize = ((EFI_MM_COMMUNICATE_HEADER_V3 *)CommBuffer)->BufferSize;
   } else {
-    BufferSize = OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data) + CommunicateHeader->MessageLength;
+    Status = SafeUintnAdd (OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data), (UINTN)CommunicateHeader->MessageLength, &BufferSize);
+    if (EFI_ERROR (Status)) {
+      return EFI_INVALID_PARAMETER;
+    }
   }
 
   if (BufferSize > EFI_PAGES_TO_SIZE (mMmCommonBuffer.NumberOfPages)) {
@@ -141,6 +146,10 @@ ProcessCommunicationBuffer (
     return EFI_UNSUPPORTED;
   }
 
+  if (CommonBufferStatus->ReturnBufferSize > BufferSize) {
+    return EFI_BAD_BUFFER_SIZE;
+  }
+
   //
   // Copy the returned data to the non-mmram buffer (CommBuffer)
   //

StandaloneMmPkg/Drivers/MmCommunicationDxe/MmCommunicationDxe.h

@@ -19,6 +19,7 @@
 #include <Library/MemoryAllocationLib.h>
 #include <Library/UefiRuntimeLib.h>
 #include <Library/ReportStatusCodeLib.h>
+#include <Library/SafeIntLib.h>
 
 #include <Protocol/SmmControl2.h>
 #include <Protocol/MmCommunication3.h>

StandaloneMmPkg/Drivers/MmCommunicationDxe/MmCommunicationDxe.inf

@@ -38,7 +38,7 @@
   UefiLib
   UefiRuntimeLib
   ReportStatusCodeLib
-
+  SafeIntLib
 [Guids]
   gMmCommBufferHobGuid
   gEfiEventVirtualAddressChangeGuid

StandaloneMmPkg/Drivers/StandaloneMmIplPei/StandaloneMmIplPei.c

@@ -61,6 +61,7 @@ Communicate (
   EFI_HOB_GUID_TYPE       *GuidHob;
   MM_COMM_BUFFER          *MmCommBuffer;
   MM_COMM_BUFFER_STATUS   *MmCommBufferStatus;
+  UINTN                   MaxBufferSize;
 
   DEBUG ((DEBUG_INFO, "StandaloneMmIpl Communicate Enter\n"));
 
@@ -92,8 +93,14 @@ Communicate (
     }
   }
 
-  if (TempCommSize > EFI_PAGES_TO_SIZE (MmCommBuffer->NumberOfPages)) {
+  Status = SafeUintnMult (MmCommBuffer->NumberOfPages, EFI_PAGE_SIZE, &MaxBufferSize);
-    DEBUG ((DEBUG_ERROR, "Communicate buffer size (%d) is over MAX (%d) size!", TempCommSize, EFI_PAGES_TO_SIZE (MmCommBuffer->NumberOfPages)));
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "Overflow occurred while calculating MaxBufferSize!\n"));
+    return EFI_INVALID_PARAMETER;
+  }
+
+  if (TempCommSize > MaxBufferSize) {
+    DEBUG ((DEBUG_ERROR, "Communicate buffer size (%d) is over MAX (%d) size!", TempCommSize, MaxBufferSize));
     return EFI_INVALID_PARAMETER;
   }
 
@@ -119,7 +126,17 @@ Communicate (
   //
   // Return status from software SMI
   //
-  *CommSize = (UINTN)MmCommBufferStatus->ReturnBufferSize;
+  Status = SafeUint64ToUintn (MmCommBufferStatus->ReturnBufferSize, CommSize);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "Overflow occurred while converting ReturnBufferSize to CommSize!\n"));
+    return EFI_BAD_BUFFER_SIZE;
+  }
+
+  if (*CommSize > TempCommSize) {
+    DEBUG ((DEBUG_ERROR, "Returned buffer size is larger than the Communication Buffer, TempCommSize: 0x%llx, ReturnBufferSize: 0x%llx\n", TempCommSize, *CommSize));
+    ASSERT (*CommSize <= TempCommSize);
+    return EFI_BAD_BUFFER_SIZE;
+  }
 
   //
   // Copy the returned data to the non-mmram buffer (CommBuffer)
@@ -157,12 +174,13 @@ Communicate3 (
   EFI_STATUS                    Status;
   EFI_PEI_MM_CONTROL_PPI        *MmControl;
   UINT8                         SmiCommand;
-  UINTN                         Size;
+  UINT64                        Size;
-  UINTN                         TempCommSize;
+  UINT64                        TempCommSize;
   EFI_HOB_GUID_TYPE             *GuidHob;
   MM_COMM_BUFFER                *MmCommBuffer;
   MM_COMM_BUFFER_STATUS         *MmCommBufferStatus;
   EFI_MM_COMMUNICATE_HEADER_V3  *CommunicateHeader;
+  UINT64                        MaxBufferSize;
 
   DEBUG ((DEBUG_INFO, "StandaloneMmIpl Communicate Enter\n"));
 
@@ -204,8 +222,14 @@ Communicate3 (
     }
   }
 
-  if (TempCommSize > EFI_PAGES_TO_SIZE (MmCommBuffer->NumberOfPages)) {
+  Status = SafeUint64Mult (MmCommBuffer->NumberOfPages, EFI_PAGE_SIZE, &MaxBufferSize);
-    DEBUG ((DEBUG_ERROR, "Communicate buffer size (%d) is over MAX (%d) size!", TempCommSize, EFI_PAGES_TO_SIZE (MmCommBuffer->NumberOfPages)));
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "Overflow occurred while calculating MaxBufferSize!\n"));
+    return EFI_INVALID_PARAMETER;
+  }
+
+  if (TempCommSize > MaxBufferSize) {
+    DEBUG ((DEBUG_ERROR, "Communicate buffer size (%d) is over MAX (%d) size!", TempCommSize, MaxBufferSize));
     return EFI_INVALID_PARAMETER;
   }
 
@@ -231,7 +255,13 @@ Communicate3 (
   //
   // Return status from software SMI
   //
-  TempCommSize = (UINTN)MmCommBufferStatus->ReturnBufferSize;
+  if (MmCommBufferStatus->ReturnBufferSize > TempCommSize) {
+    DEBUG ((DEBUG_ERROR, "Returned buffer size is larger than the Communication Buffer, TempCommSize: 0x%llx, ReturnBufferSize: 0x%llx\n", TempCommSize, MmCommBufferStatus->ReturnBufferSize));
+    ASSERT (MmCommBufferStatus->ReturnBufferSize <= TempCommSize);
+    return EFI_BAD_BUFFER_SIZE;
+  }
+
+  TempCommSize = MmCommBufferStatus->ReturnBufferSize;
 
   //
   // Copy the returned data to the non-mmram buffer (CommBuffer)

StandaloneMmPkg/Drivers/StandaloneMmIplPei/StandaloneMmIplPei.h

@@ -27,6 +27,7 @@
 #include <Ppi/MmCoreFvLocationPpi.h>
 #include <Protocol/MmCommunication.h>
 #include <Library/MmPlatformHobProducerLib.h>
+#include <Library/SafeIntLib.h>
 
 /**
   Communicates with a registered handler.

StandaloneMmPkg/Drivers/StandaloneMmIplPei/StandaloneMmIplPei.inf

@@ -44,6 +44,7 @@
   PeCoffLib
   CacheMaintenanceLib
   MmPlatformHobProducerLib
+  SafeIntLib
 
 [Guids]
   gMmCommBufferHobGuid

StandaloneMmPkg/StandaloneMmPkg.ci.yaml

@@ -28,7 +28,7 @@
 
     ## options defined .pytool/Plugin/HostUnitTestCompilerPlugin
     "HostUnitTestCompilerPlugin": {
-        "DscPath": "" # Don't support this test
+        "DscPath": "Test/StandaloneMmPkgHostTest.dsc"
     },
 
     ## options defined .pytool/Plugin/CharEncodingCheck
@@ -64,7 +64,7 @@
     ## options defined .pytool/Plugin/HostUnitTestDscCompleteCheck
     "HostUnitTestDscCompleteCheck": {
         "IgnoreInf": [""],
-        "DscPath": "" # Don't support this test
+        "DscPath": "Test/StandaloneMmPkgHostTest.dsc"
     },
 
     ## options defined .pytool/Plugin/GuidCheck

StandaloneMmPkg/Test/StandaloneMmPkgHostTest.dsc

@@ -0,0 +1,39 @@
+## @file
+# StandaloneMmPkgHostTest DSC file used to build host-based unit tests.
+#
+# Copyright (c) 2024, Intel Corporation. All rights reserved.<BR>
+# SPDX-License-Identifier: BSD-2-Clause-Patent
+#
+##
+[Defines]
+  PLATFORM_NAME           = StandaloneMmPkgHostTest
+  PLATFORM_GUID           = A3D4B87C-5E2F-4A19-8C6D-1F7E09B23D56
+  PLATFORM_VERSION        = 0.1
+  DSC_SPECIFICATION       = 0x00010005
+  OUTPUT_DIRECTORY        = Build/StandaloneMmPkg/HostTest
+  SUPPORTED_ARCHITECTURES = IA32|X64|AARCH64
+  BUILD_TARGETS           = NOOPT
+  SKUID_IDENTIFIER        = DEFAULT
+
+!include UnitTestFrameworkPkg/UnitTestFrameworkPkgHost.dsc.inc
+
+[Packages]
+  MdePkg/MdePkg.dec
+  UnitTestFrameworkPkg/UnitTestFrameworkPkg.dec
+
+[Components]
+  #
+  # Build HOST_APPLICATION that tests StandaloneMmPkg
+  #
+  StandaloneMmPkg/Drivers/MmCommunicationDxe/GoogleTest/MmCommunicationDxeGoogleTest.inf
+
+[LibraryClasses]
+  DevicePathLib|MdePkg/Library/UefiDevicePathLib/UefiDevicePathLib.inf
+  DxeServicesTableLib|MdePkg/Library/DxeServicesTableLib/DxeServicesTableLib.inf
+  HobLib|MdePkg/Library/DxeHobLib/DxeHobLib.inf
+  ReportStatusCodeLib|MdePkg/Library/BaseReportStatusCodeLibNull/BaseReportStatusCodeLibNull.inf
+  UefiDriverEntryPoint|MdePkg/Library/UefiDriverEntryPoint/UefiDriverEntryPoint.inf
+  UefiLib|MdePkg/Library/UefiLib/UefiLib.inf
+  UefiRuntimeLib|MdePkg/Library/UefiRuntimeLib/UefiRuntimeLib.inf
+  UefiRuntimeServicesTableLib|MdePkg/Library/UefiRuntimeServicesTableLib/UefiRuntimeServicesTableLib.inf
+  DebugPrintErrorLevelLib|MdePkg/Library/BaseDebugPrintErrorLevelLib/BaseDebugPrintErrorLevelLib.inf