[edk2-platforms][PATCH V1 17/20] ArmPkg/MmCommunicationDxe: Unmap FF-A RX/TX buffers during ExitBootServices

["Nishant Sharma" <nishant.sharma@arm.com>]

From: Achin Gupta <achin.gupta@arm.com>

An FF-A partition can map only a single RX/TX buffer pair with the
framework. The DXE MM communication driver maps its pair before
ExitBootServices is called. The OS cannot re-use this pair once it boots
subsequently and loads its own FF-A driver. This patch ensures that the
DXE MM communication driver unmaps its buffer pair when ExitBootServices
is called so that the OS can register its own pair if required.

Signed-off-by: Achin Gupta <achin.gupta@arm.com>
Signed-off-by: Nishant Sharma <nishant.sharma@arm.com>
---
 ArmPkg/Include/IndustryStandard/ArmFfaSvc.h         | 10 ++++
 ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c | 49 ++++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/ArmPkg/Include/IndustryStandard/ArmFfaSvc.h b/ArmPkg/Include/IndustryStandard/ArmFfaSvc.h
index ebdf29e8d69a..f78442a465e1 100644
--- a/ArmPkg/Include/IndustryStandard/ArmFfaSvc.h
+++ b/ArmPkg/Include/IndustryStandard/ArmFfaSvc.h
@@ -122,6 +122,16 @@
   (((content) & FFA_BOOT_INFO_FLAG_CONTENT_MASK)  \
    << FFA_BOOT_INFO_FLAG_CONTENT_SHIFT)
 
+/* Fromat SP ID info. */
+#define FFA_PARTITION_ID_SHIFT          16
+#define FFA_PARTITION_ID_WIDTH          16
+#define FFA_PARTITION_ID_MASK           \
+  (((1U << FFA_PARTITION_ID_WIDTH) - 1) \
+   << FFA_PARTITION_ID_SHIFT)
+#define FFA_PARTITION_ID(partid)        \
+  ((partid << FFA_PARTITION_ID_SHIFT) & \
+   FFA_PARTITION_ID_MASK)
+
 // Descriptor to pass boot information as per the FF-A v1.1 spec.
 typedef struct {
   UINT32 Name[4];
diff --git a/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c b/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c
index 8a4d46e4f80a..39a1b329b9ea 100644
--- a/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c
+++ b/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c
@@ -43,6 +43,9 @@ STATIC ARM_MEMORY_REGION_DESCRIPTOR  mNsCommBuffMemRegion;
 // Notification event when virtual address map is set.
 STATIC EFI_EVENT  mSetVirtualAddressMapEvent;
 
+// Notification event when exit boot services is called.
+STATIC EFI_EVENT  mExitBootServicesEvent;
+
 //
 // Handle to install the MM Communication Protocol
 //
@@ -255,6 +258,39 @@ NotifySetVirtualAddressMap (
   }
 }
 
+/**
+  Notification callback on ExitBootServices event.
+
+  This function notifies the MM communication protocol interface on
+  ExitBootServices event and releases the FF-A RX/TX buffer.
+
+  @param  Event          ExitBootServices event.
+  @param  Context        A context when the ExitBootServices triggered.
+
+  @retval EFI_SUCCESS    The function executed successfully.
+  @retval Other          Some error occurred when executing this function.
+
+**/
+STATIC
+VOID
+EFIAPI
+NotifyExitBootServices (
+  IN EFI_EVENT  Event,
+  IN VOID      *Context
+  )
+{
+  ARM_SMC_ARGS SmcArgs = {0};
+
+  SmcArgs.Arg0 = ARM_SVC_ID_FFA_RXTX_UNMAP_AARCH32;
+  SmcArgs.Arg1 = FFA_PARTITION_ID(mFfaPartId);
+  ArmCallSmc (&SmcArgs);
+
+  // We do not bother checking the error code of the RXTX_UNMAP invocation
+  // since we did map the buffers and this call must succeed.
+  return;
+
+}
+
 STATIC
 EFI_STATUS
 GetMmCompatibility (
@@ -452,6 +488,19 @@ MmCommunication2Initialize (
     goto CleanAddedMemorySpace;
   }
 
+  // Register notification callback when ExitBootservices is called to
+  // unregister the FF-A RX/TX buffer pair. This allows the OS to register its
+  // own buffer pair.
+  if (FixedPcdGet32 (PcdFfaEnable) != 0) {
+    Status = gBS->CreateEvent (
+                    EVT_SIGNAL_EXIT_BOOT_SERVICES,
+                    TPL_NOTIFY,
+                    NotifyExitBootServices,
+                    NULL,
+                    &mExitBootServicesEvent
+                    );
+    ASSERT_EFI_ERROR (Status);
+  }
   // Register notification callback when virtual address is associated
   // with the physical address.
   // Create a Set Virtual Address Map event.
-- 
2.34.1