From 249ca10af0f4930f63601eb4e4a1f2106e2db2d9 Mon Sep 17 00:00:00 2001
From: Kun Qin <kuqin@microsoft.com>
Date: Wed, 12 Mar 2025 22:45:06 -0700
Subject: [PATCH] StandaloneMmPkg: StandaloneMmCore: Parsing new MM communicate
 header

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

MM communicate protocols are expanded with EFI_MM_COMMUNICATE_HEADER_V3
structure that cooperates with updated field types and flexible array.
The PiSmmCore implementation is updated to detect and process incoming
data accordingly.

Two checks are also performed to prevent legacy communicate data or
unsupported data is fed into MM core under agreed header guid.

Signed-off-by: Kun Qin <kuqin12@gmail.com>
---
 StandaloneMmPkg/Core/StandaloneMmCore.c   | 49 +++++++++++++++--------
 StandaloneMmPkg/Core/StandaloneMmCore.inf |  1 +
 2 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/StandaloneMmPkg/Core/StandaloneMmCore.c b/StandaloneMmPkg/Core/StandaloneMmCore.c
index ba36c5cdb6..154285fef6 100644
--- a/StandaloneMmPkg/Core/StandaloneMmCore.c
+++ b/StandaloneMmPkg/Core/StandaloneMmCore.c
@@ -502,11 +502,15 @@ MmEntryPoint (
   IN CONST EFI_MM_ENTRY_CONTEXT  *MmEntryContext
   )
 {
-  EFI_STATUS                 Status;
-  EFI_MM_COMMUNICATE_HEADER  *CommunicateHeader;
-  MM_COMM_BUFFER_STATUS      *CommunicationStatus;
-  UINTN                      BufferSize;
-  EFI_HANDLE                 MmHandle;
+  EFI_STATUS                    Status;
+  EFI_MM_COMMUNICATE_HEADER_V3  *CommunicateHeader;
+  EFI_MM_COMMUNICATE_HEADER     *LegacyCommunicateHeader;
+  MM_COMM_BUFFER_STATUS         *CommunicationStatus;
+  UINTN                         BufferSize;
+  EFI_HANDLE                    MmHandle;
+  EFI_GUID                      *CommGuid;
+  UINTN                         CommGuidOffset;
+  UINTN                         CommHeaderSize;
 
   DEBUG ((DEBUG_INFO, "MmEntryPoint ...\n"));
 
@@ -543,8 +547,22 @@ MmEntryPoint (
       //
       // Synchronous MMI for MM Core or request from Communicate protocol
       //
-      CommunicateHeader = (EFI_MM_COMMUNICATE_HEADER *)(UINTN)mMmCommunicationBuffer->PhysicalStart;
-      BufferSize        = OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data) + CommunicateHeader->MessageLength;
+      CommGuid = &((EFI_MM_COMMUNICATE_HEADER_V3 *)(UINTN)mMmCommunicationBuffer->PhysicalStart)->HeaderGuid;
+      //
+      // Check if the signature matches EFI_MM_COMMUNICATE_HEADER_V3 definition
+      //
+      if (CompareGuid (CommGuid, &gEfiMmCommunicateHeaderV3Guid)) {
+        CommunicateHeader = (EFI_MM_COMMUNICATE_HEADER_V3 *)(UINTN)mMmCommunicationBuffer->PhysicalStart;
+        CommGuidOffset    = OFFSET_OF (EFI_MM_COMMUNICATE_HEADER_V3, MessageGuid);
+        CommHeaderSize    = sizeof (EFI_MM_COMMUNICATE_HEADER_V3);
+        BufferSize        = CommunicateHeader->BufferSize;
+      } else {
+        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;
+      }
+
       if (BufferSize <= EFI_PAGES_TO_SIZE (mMmCommunicationBuffer->NumberOfPages)) {
         //
         // Shadow the data from MM Communication Buffer to internal buffer
@@ -559,16 +577,15 @@ MmEntryPoint (
           EFI_PAGES_TO_SIZE (mMmCommunicationBuffer->NumberOfPages) - BufferSize
           );
 
-        CommunicateHeader = (EFI_MM_COMMUNICATE_HEADER *)mInternalCommBufferCopy;
-        BufferSize        = CommunicateHeader->MessageLength;
-        Status            = MmiManage (
-                              &CommunicateHeader->HeaderGuid,
-                              NULL,
-                              CommunicateHeader->Data,
-                              &BufferSize
-                              );
+        BufferSize -= CommHeaderSize;
+        Status      = MmiManage (
+                        (EFI_GUID *)((UINT8 *)mInternalCommBufferCopy + CommGuidOffset),
+                        NULL,
+                        (UINT8 *)mInternalCommBufferCopy + CommHeaderSize,
+                        &BufferSize
+                        );
 
-        BufferSize = BufferSize + OFFSET_OF (EFI_MM_COMMUNICATE_HEADER, Data);
+        BufferSize = BufferSize + CommHeaderSize;
         if (BufferSize <= EFI_PAGES_TO_SIZE (mMmCommunicationBuffer->NumberOfPages)) {
           //
           // Copy the data back to MM Communication Buffer
diff --git a/StandaloneMmPkg/Core/StandaloneMmCore.inf b/StandaloneMmPkg/Core/StandaloneMmCore.inf
index f3340d26f9..e9ded96bf6 100644
--- a/StandaloneMmPkg/Core/StandaloneMmCore.inf
+++ b/StandaloneMmPkg/Core/StandaloneMmCore.inf
@@ -84,6 +84,7 @@
   gEfiSmmSmramMemoryGuid
   gEdkiiPiSmmMemoryAttributesTableGuid
   gEfiMmPeiMmramMemoryReserveGuid
+  gEfiMmCommunicateHeaderV3Guid    ## CONSUMES   ## GUID # Communicate header
 
 [Pcd]
   gStandaloneMmPkgTokenSpaceGuid.PcdFwVolMmMaxEncapsulationDepth    ##CONSUMES