[Rebase & FF] Adding MM communication buffer validation

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> // MU_CHANGE: Add SafeIntLib for safe integer operations
 
 #include <Protocol/MmCommunication2.h>
 #include <Protocol/MmCommunication3.h>
@@ -213,6 +214,7 @@ MmCommunicationCommon (
   EFI_MM_COMMUNICATE_HEADER     *CommunicateHeader;
   EFI_MM_COMMUNICATE_HEADER_V3  *CommunicateHeaderV3;
   UINTN                         BufferSize;
+  UINTN                         InputBufferSize;  // MU_CHANGE
   UINTN                         *MessageSize;
   UINTN                         HeaderSize;
   EFI_STATUS                    Status;
@@ -243,15 +245,41 @@ MmCommunicationCommon (
     BufferSize          = CommunicateHeaderV3->BufferSize;
     MessageSize         = &CommunicateHeaderV3->MessageSize;
     HeaderSize          = sizeof (EFI_MM_COMMUNICATE_HEADER_V3);
+
+    // MU_CHANGE Starts: Add a check to ensure that the buffer size in the header is sane
+    if (BufferSize < HeaderSize) {
+      return EFI_INVALID_PARAMETER;
+    }
+
+    if (BufferSize - HeaderSize < *MessageSize) {
+      return EFI_INVALID_PARAMETER;
+    }
+
+    // MU_CHANGE Ends
   } else {
-    BufferSize = CommunicateHeader->MessageLength +
-                 sizeof (CommunicateHeader->HeaderGuid) +
-                 sizeof (CommunicateHeader->MessageLength);
+    // MU_CHANGE: Use SafeIntLib for safe arithmetic operations
+    // BufferSize = CommunicateHeader->MessageLength +
+    //              sizeof (CommunicateHeader->HeaderGuid) +
+    //              sizeof (CommunicateHeader->MessageLength);
+    Status = SafeUintnAdd (
+               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) +
-                  sizeof (CommunicateHeader->MessageLength);
+    // MU_CHANGE: Avoid direct calculation of HeaderSize
+    // HeaderSize  = sizeof (CommunicateHeader->HeaderGuid) +
+    //               sizeof (CommunicateHeader->MessageLength);
+    HeaderSize = OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data);
   }
 
+  // MU_CHANGE: Record the input buffer size for later comparison with the returned buffer size
+  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
@@ -269,6 +297,9 @@ MmCommunicationCommon (
     if (*CommSize < BufferSize) {
       Status = EFI_INVALID_PARAMETER;
     }
+
+    // MU_CHANGE: Record the updated input buffer size for later comparison with the returned buffer size
+    InputBufferSize = *CommSize;
   }
 
   //
@@ -297,7 +328,8 @@ MmCommunicationCommon (
   }
 
   if (!EFI_ERROR (Status)) {
-    ZeroMem (CommBufferVirtual, BufferSize);
+    // MU_CHANGE: Do not clear the input buffer as we will copy the returned data to the caller's buffer below
+    // ZeroMem (CommBufferVirtual, BufferSize);
     // On successful return, the size of data being returned is inferred from
     // MessageLength + Header.
     CommunicateHeader = (EFI_MM_COMMUNICATE_HEADER *)mNsCommBuffMemRegion.VirtualBase;
@@ -320,20 +352,30 @@ MmCommunicationCommon (
       CommunicateHeaderV3 = (EFI_MM_COMMUNICATE_HEADER_V3 *)CommunicateHeader;
       BufferSize          = CommunicateHeaderV3->BufferSize;
     } else {
-      BufferSize = CommunicateHeader->MessageLength +
-                   sizeof (CommunicateHeader->HeaderGuid) +
-                   sizeof (CommunicateHeader->MessageLength);
+      // MU_CHANGE: Use SafeIntLib for safe arithmetic operations
+      // BufferSize = CommunicateHeader->MessageLength +
+      //              sizeof (CommunicateHeader->HeaderGuid) +
+      //              sizeof (CommunicateHeader->MessageLength);
+      Status = SafeUintnAdd (
+                 CommunicateHeader->MessageLength,
+                 OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data),
+                 &BufferSize
+                 );
+      if (EFI_ERROR (Status)) {
+        return EFI_INVALID_PARAMETER;
+      }
     }
 
-    if (BufferSize > mNsCommBuffMemRegion.Length) {
+    // MU_CHANGE: Add a check to ensure that the returned buffer size does not exceed the caller supplied buffer size
+    if (BufferSize > InputBufferSize) {
       // 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",
         __func__,
         BufferSize,
-        (UINTN)mNsCommBuffMemRegion.Length
+        InputBufferSize
         ));
     } else {
       CopyMem (

ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.inf

@@ -39,6 +39,7 @@
   DxeServicesTableLib
   HobLib
   UefiDriverEntryPoint
+  SafeIntLib  # MU_CHANGE
 
 [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> // MU_CHANGE: Use SafeIntLib for safe arithmetic operations
 
 //
 // Partition ID if FF-A support is enabled
@@ -342,6 +343,7 @@ MmCommunicationPeimCommon (
   EFI_MM_COMMUNICATE_HEADER_V3  *CommunicateHeaderV3;
   EFI_STATUS                    Status;
   UINTN                         BufferSize;
+  UINTN                         InputBufferSize; // MU_CHANGE: Use SafeIntLib for safe arithmetic operations
   UINTN                         HeaderSize;
 
   //
@@ -362,7 +364,31 @@ MmCommunicationPeimCommon (
     // This is a v3 header
     CommunicateHeaderV3 = (EFI_MM_COMMUNICATE_HEADER_V3 *)(UINTN)CommBuffer;
     HeaderSize          = sizeof (EFI_MM_COMMUNICATE_HEADER_V3);
-    BufferSize          = CommunicateHeaderV3->BufferSize;
+
+    // MU_CHANGE Starts: Add a check to ensure that the buffer size in the header is sane
+    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;
+    }
+
+    // MU_CHANGE Ends
   } else {
     // This is a v1 header, do some checks
     if (CommSize == NULL) {
@@ -390,35 +416,67 @@ MmCommunicationPeimCommon (
       return EFI_BAD_BUFFER_SIZE;
     }
 
-    HeaderSize =  sizeof (CommunicateHeader->HeaderGuid) +
-                 sizeof (CommunicateHeader->MessageLength);
+    // MU_CHANGE Starts: Avoid unsafe arithmetic operations
+    // HeaderSize =  sizeof (CommunicateHeader->HeaderGuid) +
+    //              sizeof (CommunicateHeader->MessageLength);
+    HeaderSize =  OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data);
+    // MU_CHANGE Ends
 
     // CommBuffer is a mandatory parameter. Hence, Rely on
     // MessageLength + Header to ascertain the
     // total size of the communication payload rather than
     // rely on optional CommSize parameter
-    BufferSize = CommunicateHeader->MessageLength +
-                 sizeof (CommunicateHeader->HeaderGuid) +
-                 sizeof (CommunicateHeader->MessageLength);
+    // MU_CHANGE: Use SafeIntLib for safe arithmetic operations
+    // BufferSize = CommunicateHeader->MessageLength +
+    //              sizeof (CommunicateHeader->HeaderGuid) +
+    //              sizeof (CommunicateHeader->MessageLength);
+    Status = SafeUintnAdd (CommunicateHeader->MessageLength, OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data), &InputBufferSize);
+    if (EFI_ERROR (Status)) {
+      DEBUG ((
+        DEBUG_ERROR,
+        "%a Overflow occurred while calculating input BufferSize!\n",
+        __func__
+        ));
+      return Status;
+    }
+
+    // MU_CHANGE Ends
+
     //
     // If CommSize is supplied it must match MessageLength + sizeof (EFI_MM_COMMUNICATE_HEADER);
     //
-    if (*CommSize != BufferSize) {
+    // MU_CHANGE: Record the input buffer size for later comparison with the returned buffer size
+    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;
     }
   }
 
+  // MU_CHANGE Starts: Add a check to ensure that the buffer size does not exceed the allocated MM Communication Buffer size
+  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;
+  }
+
+  // MU_CHANGE Ends
+
   // 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);
+  // MU_CHANGE: Record the input buffer size for later comparison with the returned buffer size
+  CopyMem (CommunicateHeader, CommBuffer, InputBufferSize);
   if (IsFfaSupported ()) {
     Status = SendFfaMmCommunicate ();
   } else {
@@ -447,20 +505,33 @@ MmCommunicationPeimCommon (
       CommunicateHeaderV3 = (EFI_MM_COMMUNICATE_HEADER_V3 *)CommunicateHeader;
       BufferSize          = CommunicateHeaderV3->BufferSize;
     } else {
-      BufferSize = CommunicateHeader->MessageLength +
-                   sizeof (CommunicateHeader->HeaderGuid) +
-                   sizeof (CommunicateHeader->MessageLength);
+      // MU_CHANGE: Use SafeIntLib for safe arithmetic operations
+      // BufferSize = CommunicateHeader->MessageLength +
+      //              sizeof (CommunicateHeader->HeaderGuid) +
+      //              sizeof (CommunicateHeader->MessageLength);
+      Status = SafeUintnAdd (CommunicateHeader->MessageLength, OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data), &BufferSize);
+      if (EFI_ERROR (Status)) {
+        DEBUG ((
+          DEBUG_ERROR,
+          "%a Overflow occurred while calculating returned BufferSize!\n",
+          __func__
+          ));
+        return Status;
+      }
+
+      // MU_CHANGE Ends
     }
 
-    if (BufferSize > (UINTN)PcdGet64 (PcdMmBufferSize)) {
+    // MU_CHANGE: Add a check to ensure that the returned buffer size does not exceed the caller supplied buffer size
+    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,
-        (UINTN)PcdGet64 (PcdMmBufferSize)
+        InputBufferSize,
+        BufferSize
         ));
     } else {
       CopyMem (CommBuffer, CommunicateHeader, BufferSize);

ArmPkg/Drivers/MmCommunicationPei/MmCommunicationPei.inf

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

ArmPkg/Library/ArmStandaloneMmCoreEntryPoint/ArmStandaloneMmCoreEntryPoint.c

@@ -504,6 +504,11 @@ ValidateMmCommBufferAddr (
         &gEfiMmCommunicateHeaderV3Guid
         ))
   {
+    // MU_CHANGE: Now this is supposed to be a V3 header, so we need to check the size of the buffer in the header as well
+    if (CommBufferRange < sizeof (EFI_MM_COMMUNICATE_HEADER_V3)) {
+      return EFI_ACCESS_DENIED;
+    }
+
     CommBufferHeaderV3 = (EFI_MM_COMMUNICATE_HEADER_V3 *)CommBufferAddr;
     Status             = SafeUint64Add (
                            CommBufferHeaderV3->MessageSize,
@@ -513,9 +518,24 @@ ValidateMmCommBufferAddr (
     if (EFI_ERROR (Status)) {
       return EFI_ACCESS_DENIED;
     }
+
+    // MU_CHANGE Starts: Make sure the buffer size in the header is also sane
+    if (BufferSize > CommBufferHeaderV3->BufferSize) {
+      return EFI_ACCESS_DENIED;
+    }
+
+    BufferSize = CommBufferHeaderV3->BufferSize;
+    // MU_CHANGE Ends
   } else {
-    BufferSize = ((EFI_MM_COMMUNICATE_HEADER *)CommBufferAddr)->MessageLength +
-                 OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data);
+    // MU_CHANGE: Use SafeIntLib for safe arithmetic operations
+    Status = SafeUint64Add (
+               ((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

@@ -50,6 +50,7 @@
   ArmFfaLib
   StackCheckLib
   HobLib
+  SafeIntLib  # MU_CHANGE
 
 [Guids]
   gMpInformationHobGuid