[PATCH edk2 v1 2/2] edk2/MM: Fix MM Communicate return wrong status issue

[Ming Huang <huangming@linux.alibaba.com>]

When edk2 driver call EFI_MM_COMMUNICATION_PROTOCOL.Communicate,
the status of mm event can't return exactly to edk2 driver.

Signed-off-by: Ming Huang <huangming@linux.alibaba.com>
---
 ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c           | 2 --
 MdePkg/Include/Pi/PiMmCis.h                                   | 2 +-
 StandaloneMmPkg/Core/StandaloneMmCore.c                       | 5 +++--
 StandaloneMmPkg/Drivers/StandaloneMmCpu/AArch64/EventHandle.c | 5 +++--
 4 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c b/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c
index 9457eaf1d8..b5b074b56e 100644
--- a/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c
+++ b/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c
@@ -172,12 +172,10 @@ MmCommunication2Communicate (
     // Unexpected error since the CommSize was checked for zero length
     // prior to issuing the SMC
     Status = EFI_OUT_OF_RESOURCES;
-    ASSERT (0);
     break;
 
   default:
     Status = EFI_ACCESS_DENIED;
-    ASSERT (0);
   }
 
   return Status;
diff --git a/MdePkg/Include/Pi/PiMmCis.h b/MdePkg/Include/Pi/PiMmCis.h
index fdf0591a03..762a1fbce5 100644
--- a/MdePkg/Include/Pi/PiMmCis.h
+++ b/MdePkg/Include/Pi/PiMmCis.h
@@ -237,7 +237,7 @@ typedef struct _EFI_MM_ENTRY_CONTEXT {
   @param[in] MmEntryContext  Processor information and functionality needed by MM Foundation.
 **/
 typedef
-VOID
+EFI_STATUS
 (EFIAPI *EFI_MM_ENTRY_POINT)(
   IN CONST EFI_MM_ENTRY_CONTEXT  *MmEntryContext
   );
diff --git a/StandaloneMmPkg/Core/StandaloneMmCore.c b/StandaloneMmPkg/Core/StandaloneMmCore.c
index ac3e2c0b1b..aa59ceb5d7 100644
--- a/StandaloneMmPkg/Core/StandaloneMmCore.c
+++ b/StandaloneMmPkg/Core/StandaloneMmCore.c
@@ -332,7 +332,7 @@ MmEndOfDxeHandler (
                                     needed by MM Foundation.
 
 **/
-VOID
+EFI_STATUS
 EFIAPI
 MmEntryPoint (
   IN CONST EFI_MM_ENTRY_CONTEXT  *MmEntryContext
@@ -398,7 +398,7 @@ MmEntryPoint (
   //
   // Process Asynchronous MMI sources
   //
-  MmiManage (NULL, NULL, NULL, NULL);
+  Status = MmiManage (NULL, NULL, NULL, NULL);
 
   //
   // TBD: Do not use private data structure ?
@@ -410,6 +410,7 @@ MmEntryPoint (
   gMmCorePrivate->InMm = FALSE;
 
   DEBUG ((DEBUG_INFO, "MmEntryPoint Done\n"));
+  return Status;
 }
 
 EFI_STATUS
diff --git a/StandaloneMmPkg/Drivers/StandaloneMmCpu/AArch64/EventHandle.c b/StandaloneMmPkg/Drivers/StandaloneMmCpu/AArch64/EventHandle.c
index 3730ee6379..6951d8ae50 100644
--- a/StandaloneMmPkg/Drivers/StandaloneMmCpu/AArch64/EventHandle.c
+++ b/StandaloneMmPkg/Drivers/StandaloneMmCpu/AArch64/EventHandle.c
@@ -68,6 +68,7 @@ PiMmStandaloneArmTfCpuDriverEntry (
   EFI_MM_COMMUNICATE_HEADER *GuidedEventContext = NULL;
   EFI_MM_ENTRY_CONTEXT        MmEntryPointContext = {0};
   EFI_STATUS                  Status;
+  EFI_STATUS                  MmRetStatus;
   UINTN                       NsCommBufferSize;
 
   DEBUG ((DEBUG_INFO, "Received event - 0x%x on cpu %d\n", EventId, CpuNumber));
@@ -135,7 +136,7 @@ PiMmStandaloneArmTfCpuDriverEntry (
     return EFI_UNSUPPORTED;
   }
 
-  mMmEntryPoint (&MmEntryPointContext);
+  MmRetStatus = mMmEntryPoint (&MmEntryPointContext);
 
   // Free the memory allocation done earlier and reset the per-cpu context
   ASSERT (GuidedEventContext);
@@ -147,7 +148,7 @@ PiMmStandaloneArmTfCpuDriverEntry (
   }
   PerCpuGuidedEventContext[CpuNumber] = NULL;
 
-  return Status;
+  return MmRetStatus;
 }
 
 EFI_STATUS
-- 
2.17.1