From 9dec81092bf99bbd28abd95c92023561fa04526f Mon Sep 17 00:00:00 2001
From: Kun Qin <kuqin@microsoft.com>
Date: Wed, 7 Aug 2024 09:39:07 -0700
Subject: [PATCH] MdeModulePkg: PiSmmCore: Added parser of 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.

Cc: Jian J Wang <jian.j.wang@intel.com>
Cc: Hao A Wu <hao.a.wu@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>

Signed-off-by: Kun Qin <kuqin12@gmail.com>
---
 MdeModulePkg/Core/PiSmmCore/PiSmmCore.c   | 64 ++++++++++++++++-------
 MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf |  1 +
 2 files changed, 46 insertions(+), 19 deletions(-)

diff --git a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
index e957cb67f8..1b2b8857c5 100644
--- a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
+++ b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
@@ -670,13 +670,16 @@ SmmEntryPoint (
   IN CONST EFI_SMM_ENTRY_CONTEXT  *SmmEntryContext
   )
 {
-  EFI_STATUS                  Status;
-  EFI_SMM_COMMUNICATE_HEADER  *CommunicateHeader;
-  BOOLEAN                     InLegacyBoot;
-  BOOLEAN                     IsOverlapped;
-  BOOLEAN                     IsOverUnderflow;
-  VOID                        *CommunicationBuffer;
-  UINTN                       BufferSize;
+  EFI_STATUS                    Status;
+  EFI_MM_COMMUNICATE_HEADER_V3  *CommunicateHeader;
+  EFI_SMM_COMMUNICATE_HEADER    *LegacyCommunicateHeader;
+  BOOLEAN                       InLegacyBoot;
+  BOOLEAN                       IsOverlapped;
+  VOID                          *CommunicationBuffer;
+  UINTN                         BufferSize;
+  EFI_GUID                      *CommGuid;
+  VOID                          *CommData;
+  UINTN                         CommHeaderSize;
 
   PERF_FUNCTION_BEGIN ();
 
@@ -730,10 +733,8 @@ SmmEntryPoint (
       //
       // Check for over or underflows
       //
-      IsOverUnderflow = EFI_ERROR (SafeUintnSub (BufferSize, OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data), &BufferSize));
-
       if (!SmmIsBufferOutsideSmmValid ((UINTN)CommunicationBuffer, BufferSize) ||
-          IsOverlapped || IsOverUnderflow)
+          IsOverlapped || (BufferSize < sizeof (EFI_SMM_COMMUNICATE_HEADER)))
       {
         //
         // If CommunicationBuffer is not in valid address scope,
@@ -744,25 +745,50 @@ SmmEntryPoint (
         gSmmCorePrivate->CommunicationBuffer = NULL;
         gSmmCorePrivate->ReturnStatus        = EFI_ACCESS_DENIED;
       } else {
-        CommunicateHeader = (EFI_SMM_COMMUNICATE_HEADER *)CommunicationBuffer;
-        // BufferSize was updated by the SafeUintnSub() call above.
-        Status = SmiManage (
-                   &CommunicateHeader->HeaderGuid,
-                   NULL,
-                   CommunicateHeader->Data,
-                   &BufferSize
-                   );
+        CommGuid = &((EFI_MM_COMMUNICATE_HEADER_V3 *)CommunicationBuffer)->HeaderGuid;
+        //
+        // Check if the signature matches EFI_MM_COMMUNICATE_HEADER_V3 definition
+        //
+        if (CompareGuid (CommGuid, &gEfiMmCommunicateHeaderV3Guid)) {
+          //
+          // If so, need to make sure the size is at least the size of the header
+          //
+          if (BufferSize < sizeof (EFI_MM_COMMUNICATE_HEADER_V3)) {
+            gSmmCorePrivate->CommunicationBuffer = NULL;
+            gSmmCorePrivate->ReturnStatus        = EFI_ACCESS_DENIED;
+            goto AsyncSmi;
+          }
+
+          CommunicateHeader = (EFI_MM_COMMUNICATE_HEADER_V3 *)CommunicationBuffer;
+          CommGuid          = &CommunicateHeader->MessageGuid;
+          CommData          = CommunicateHeader->MessageData;
+          CommHeaderSize    = sizeof (EFI_MM_COMMUNICATE_HEADER_V3);
+        } else {
+          LegacyCommunicateHeader = (EFI_SMM_COMMUNICATE_HEADER *)CommunicationBuffer;
+          CommGuid                = &LegacyCommunicateHeader->HeaderGuid;
+          CommData                = LegacyCommunicateHeader->Data;
+          CommHeaderSize          = OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data);
+        }
+
+        BufferSize -= CommHeaderSize;
+        Status      = SmiManage (
+                        CommGuid,
+                        NULL,
+                        CommData,
+                        &BufferSize
+                        );
         //
         // Update CommunicationBuffer, BufferSize and ReturnStatus
         // Communicate service finished, reset the pointer to CommBuffer to NULL
         //
-        gSmmCorePrivate->BufferSize          = BufferSize + OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data);
+        gSmmCorePrivate->BufferSize          = BufferSize + CommHeaderSize;
         gSmmCorePrivate->CommunicationBuffer = NULL;
         gSmmCorePrivate->ReturnStatus        = (Status == EFI_SUCCESS) ? EFI_SUCCESS : EFI_NOT_FOUND;
       }
     }
   }
 
+AsyncSmi:
   //
   // Process Asynchronous SMI sources
   //
diff --git a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf
index 75a5934f0c..22a1048d5b 100644
--- a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf
+++ b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf
@@ -120,6 +120,7 @@
   gSmiHandlerProfileGuid
   gEdkiiEndOfS3ResumeGuid ## SOMETIMES_PRODUCES ## GUID # Install protocol
   gEdkiiS3SmmInitDoneGuid ## SOMETIMES_PRODUCES ## GUID # Install protocol
+  gEfiMmCommunicateHeaderV3Guid    ## CONSUMES   ## GUID # Communicate header
 
 [UserExtensions.TianoCore."ExtraFiles"]
   PiSmmCoreExtra.uni