diff --git a/MdeModulePkg/Library/ArmFfaLib/ArmFfaCommon.c b/MdeModulePkg/Library/ArmFfaLib/ArmFfaCommon.c
index bf37f2fa0f..d8b04454b1 100644
--- a/MdeModulePkg/Library/ArmFfaLib/ArmFfaCommon.c
+++ b/MdeModulePkg/Library/ArmFfaLib/ArmFfaCommon.c
@@ -34,9 +34,6 @@
 
 #include "ArmFfaCommon.h"
 
-BOOLEAN  gFfaSupported;
-UINT16   gPartId;
-
 /**
   Convert EFI_STATUS to FFA return code.
 
@@ -190,22 +187,6 @@ ArmCallFfa (
   }
 }
 
-/**
-  Check FF-A support or not.
-
-  @retval TRUE                   Supported
-  @retval FALSE                  Not supported
-
-**/
-BOOLEAN
-EFIAPI
-IsFfaSupported (
-  IN VOID
-  )
-{
-  return gFfaSupported;
-}
-
 /**
   Get FF-A version.
 
@@ -837,15 +818,21 @@ ArmFfaLibMsgSendDirectReq (
 {
   EFI_STATUS    Status;
   ARM_FFA_ARGS  FfaArgs;
+  UINT16        PartId;
 
-  if ((DestPartId == gPartId) || (ImpDefArgs == NULL)) {
+  Status = ArmFfaLibGetPartId (&PartId);
+  if (EFI_ERROR (Status)) {
+    return Status;
+  }
+
+  if ((DestPartId == PartId) || (ImpDefArgs == NULL)) {
     return EFI_INVALID_PARAMETER;
   }
 
   ZeroMem (&FfaArgs, sizeof (ARM_FFA_ARGS));
 
   FfaArgs.Arg0 = ARM_FID_FFA_MSG_SEND_DIRECT_REQ;
-  FfaArgs.Arg1 = PACK_PARTITION_ID_INFO (gPartId, DestPartId);
+  FfaArgs.Arg1 = PACK_PARTITION_ID_INFO (PartId, DestPartId);
   FfaArgs.Arg2 = Flags;
   FfaArgs.Arg3 = ImpDefArgs->Arg0;
   FfaArgs.Arg4 = ImpDefArgs->Arg1;
@@ -892,6 +879,7 @@ ArmFfaLibMsgSendDirectReq2 (
   EFI_STATUS    Status;
   UINT64        Uuid[2];
   ARM_FFA_ARGS  FfaArgs;
+  UINT16        PartId;
 
   /*
    * Direct message request 2 is only supported on AArch64.
@@ -900,7 +888,12 @@ ArmFfaLibMsgSendDirectReq2 (
     return EFI_UNSUPPORTED;
   }
 
-  if ((DestPartId == gPartId) || (ImpDefArgs == NULL)) {
+  Status = ArmFfaLibGetPartId (&PartId);
+  if (EFI_ERROR (Status)) {
+    return Status;
+  }
+
+  if ((DestPartId == PartId) || (ImpDefArgs == NULL)) {
     return EFI_INVALID_PARAMETER;
   }
 
@@ -913,7 +906,7 @@ ArmFfaLibMsgSendDirectReq2 (
   ZeroMem (&FfaArgs, sizeof (ARM_FFA_ARGS));
 
   FfaArgs.Arg0  = ARM_FID_FFA_MSG_SEND_DIRECT_REQ2;
-  FfaArgs.Arg1  = PACK_PARTITION_ID_INFO (gPartId, DestPartId);
+  FfaArgs.Arg1  = PACK_PARTITION_ID_INFO (PartId, DestPartId);
   FfaArgs.Arg2  = Uuid[0];
   FfaArgs.Arg3  = Uuid[1];
   FfaArgs.Arg4  = ImpDefArgs->Arg0;
@@ -959,56 +952,39 @@ ArmFfaLibMsgSendDirectReq2 (
 /**
   Common ArmFfaLib init.
 
+  @param [out] PartId            PartitionId
+  @param [out] IsFfaSupported    FF-A supported flag
+
   @retval EFI_SUCCESS            Success
-  @retval EFI_UNSUPPORTED        FF-A isn't supported
+  @retval EFI_INVALID_PARAMETER  Invalid parameter
   @retval Others                 Error
 
 **/
 EFI_STATUS
 EFIAPI
 ArmFfaLibCommonInit (
-  IN VOID
+  OUT UINT16   *PartId,
+  OUT BOOLEAN  *IsFfaSupported
   )
 {
   EFI_STATUS  Status;
-  UINT16      CurrentMajorVersion;
-  UINT16      CurrentMinorVersion;
 
-  gFfaSupported = FALSE;
+  if ((PartId == NULL) || (IsFfaSupported == NULL)) {
+    return EFI_INVALID_PARAMETER;
+  }
 
-  Status = ArmFfaLibGetVersion (
-             ARM_FFA_MAJOR_VERSION,
-             ARM_FFA_MINOR_VERSION,
-             &CurrentMajorVersion,
-             &CurrentMinorVersion
-             );
+  Status = ArmFfaLibIsFfaSupported ();
   if (EFI_ERROR (Status)) {
-    return EFI_UNSUPPORTED;
+    return Status;
   }
 
-  if ((ARM_FFA_MAJOR_VERSION != CurrentMajorVersion) ||
-      (ARM_FFA_MINOR_VERSION > CurrentMinorVersion))
-  {
-    DEBUG ((
-      DEBUG_INFO,
-      "Incompatible FF-A Versions.\n" \
-      "Request Version: Major=0x%x, Minor=0x%x.\n" \
-      "Current Version: Major=0x%x, Minor>=0x%x.\n",
-      ARM_FFA_MAJOR_VERSION,
-      ARM_FFA_MINOR_VERSION,
-      CurrentMajorVersion,
-      CurrentMinorVersion
-      ));
-    return EFI_UNSUPPORTED;
-  }
+  *IsFfaSupported = TRUE;
 
-  Status = ArmFfaLibPartitionIdGet (&gPartId);
+  Status = ArmFfaLibPartitionIdGet (PartId);
   if (EFI_ERROR (Status)) {
     return Status;
   }
 
-  gFfaSupported = TRUE;
-
   return EFI_SUCCESS;
 }
 
@@ -1043,6 +1019,7 @@ ArmFfaLibGetPartitionInfo (
   UINT64      RxBufferSize;
   UINT32      Count;
   UINT32      Size;
+  UINT16      PartId;
 
   if (PartDesc == NULL) {
     return EFI_INVALID_PARAMETER;
@@ -1081,13 +1058,18 @@ ArmFfaLibGetPartitionInfo (
     return Status;
   }
 
+  Status = ArmFfaLibGetPartId (&PartId);
+  if (EFI_ERROR (Status)) {
+    return Status;
+  }
+
   if ((Size < sizeof (EFI_FFA_PART_INFO_DESC))) {
-    ArmFfaLibRxRelease (gPartId);
+    ArmFfaLibRxRelease (PartId);
     return EFI_INVALID_PARAMETER;
   }
 
   CopyMem (PartDesc, RxBuffer, sizeof (EFI_FFA_PART_INFO_DESC));
-  ArmFfaLibRxRelease (gPartId);
+  ArmFfaLibRxRelease (PartId);
 
   return EFI_SUCCESS;
 }
@@ -1234,3 +1216,49 @@ GetRxTxBufferMinSizeAndAlign (
 
   return EFI_SUCCESS;
 }
+
+/**
+  Determine if FF-A is supported
+
+  @retval EFI_SUCCESS            FF-A is supported
+  @retval EFI_UNSUPPORTED        FF-A isn't supported
+
+**/
+EFI_STATUS
+EFIAPI
+ArmFfaLibIsFfaSupported (
+  IN VOID
+  )
+{
+  EFI_STATUS  Status;
+  UINT16      CurrentMajorVersion;
+  UINT16      CurrentMinorVersion;
+
+  Status = ArmFfaLibGetVersion (
+             ARM_FFA_MAJOR_VERSION,
+             ARM_FFA_MINOR_VERSION,
+             &CurrentMajorVersion,
+             &CurrentMinorVersion
+             );
+  if (EFI_ERROR (Status)) {
+    return EFI_UNSUPPORTED;
+  }
+
+  if ((ARM_FFA_MAJOR_VERSION != CurrentMajorVersion) ||
+      (ARM_FFA_MINOR_VERSION > CurrentMinorVersion))
+  {
+    DEBUG ((
+      DEBUG_INFO,
+      "Incompatible FF-A Versions.\n" \
+      "Request Version: Major=0x%x, Minor=0x%x.\n" \
+      "Current Version: Major=0x%x, Minor>=0x%x.\n",
+      ARM_FFA_MAJOR_VERSION,
+      ARM_FFA_MINOR_VERSION,
+      CurrentMajorVersion,
+      CurrentMinorVersion
+      ));
+    return EFI_UNSUPPORTED;
+  }
+
+  return EFI_SUCCESS;
+}
diff --git a/MdeModulePkg/Library/ArmFfaLib/ArmFfaCommon.h b/MdeModulePkg/Library/ArmFfaLib/ArmFfaCommon.h
index 1994dd8eef..b1f7af88f7 100644
--- a/MdeModulePkg/Library/ArmFfaLib/ArmFfaCommon.h
+++ b/MdeModulePkg/Library/ArmFfaLib/ArmFfaCommon.h
@@ -19,9 +19,6 @@
 
 #include <Library/ArmFfaLib.h>
 
-extern BOOLEAN  gFfaSupported;
-extern UINT16   gPartId;
-
 /**
   Convert FfArgs to EFI_STATUS.
 
@@ -37,16 +34,21 @@ FfaArgsToEfiStatus (
   );
 
 /**
-  Common ArmFfaLib Constructor.
+  Common ArmFfaLib init.
 
-  @retval EFI_SUCCESS
-  @retval Others                  Error
+  @param [out] PartId            PartitionId
+  @param [out] IsFfaSupported    FF-A supported flag
+
+  @retval EFI_SUCCESS            Success
+  @retval EFI_INVALID_PARAMETER  Invalid parameter
+  @retval Others                 Error
 
 **/
 EFI_STATUS
 EFIAPI
 ArmFfaLibCommonInit (
-  IN VOID
+  OUT UINT16   *PartId,
+  OUT BOOLEAN  *IsFfaSupported
   );
 
 /**
@@ -86,4 +88,34 @@ GetRxTxBufferMinSizeAndAlign (
   OUT UINTN  *MinSizeAndAlign
   );
 
+/**
+  Determine if FF-A is supported
+
+  @param[out] IsSupported  If FF-A is supported
+
+  @retval EFI_SUCCESS            FF-A is supported
+  @retval EFI_UNSUPPORTED        FF-A isn't supported
+
+**/
+EFI_STATUS
+EFIAPI
+ArmFfaLibIsFfaSupported (
+  IN VOID
+  );
+
+/**
+  Return partition or VM ID
+
+  @param[out] PartId  The partition or VM ID
+
+  @retval EFI_SUCCESS
+  @retval Others       Errors
+
+**/
+EFI_STATUS
+EFIAPI
+ArmFfaLibGetPartId (
+  OUT UINT16  *PartId
+  );
+
 #endif
diff --git a/MdeModulePkg/Library/ArmFfaLib/ArmFfaDxeLib.c b/MdeModulePkg/Library/ArmFfaLib/ArmFfaDxeLib.c
index 94a4a21560..f4a4774fdf 100644
--- a/MdeModulePkg/Library/ArmFfaLib/ArmFfaDxeLib.c
+++ b/MdeModulePkg/Library/ArmFfaLib/ArmFfaDxeLib.c
@@ -34,6 +34,8 @@
 #include "ArmFfaRxTxMap.h"
 
 STATIC EFI_EVENT  mFfaExitBootServiceEvent;
+STATIC UINT16     mPartId;
+STATIC BOOLEAN    mIsFfaSupported;
 
 /**
   Unmap RX/TX buffer on Exit Boot Service.
@@ -78,7 +80,7 @@ ArmFfaDxeLibConstructor (
   UINTN                      Property1;
   UINTN                      Property2;
 
-  Status = ArmFfaLibCommonInit ();
+  Status = ArmFfaLibCommonInit (&mPartId, &mIsFfaSupported);
   if (EFI_ERROR (Status)) {
     if (Status == EFI_UNSUPPORTED) {
       /*
@@ -194,3 +196,53 @@ ArmFfaDxeLibConstructor (
 
   return Status;
 }
+
+/**
+  Return partition or VM ID
+
+  @retval EFI_SUCCESS
+  @retval Others       Errors
+
+**/
+EFI_STATUS
+EFIAPI
+ArmFfaLibGetPartId (
+  IN UINT16  *PartId
+  )
+{
+  if (PartId != NULL) {
+    *PartId = mPartId;
+  }
+
+  return EFI_SUCCESS;
+}
+
+/**
+  Check FF-A support or not.
+
+  @retval TRUE                   Supported
+  @retval FALSE                  Not supported
+
+**/
+BOOLEAN
+EFIAPI
+IsFfaSupported (
+  IN VOID
+  )
+{
+  return mIsFfaSupported;
+}
+
+/**
+  Callback for when Unmap is called to handle any post unmap
+  functionality.
+
+**/
+VOID
+EFIAPI
+UnmapCallback (
+  IN VOID
+  )
+{
+  // Do nothing
+}
diff --git a/MdeModulePkg/Library/ArmFfaLib/ArmFfaLibBase.c b/MdeModulePkg/Library/ArmFfaLib/ArmFfaLibBase.c
index aadc07158b..3875b70a18 100644
--- a/MdeModulePkg/Library/ArmFfaLib/ArmFfaLibBase.c
+++ b/MdeModulePkg/Library/ArmFfaLib/ArmFfaLibBase.c
@@ -23,6 +23,9 @@
 
 #include "ArmFfaCommon.h"
 
+STATIC UINT16   mPartId;
+STATIC BOOLEAN  mIsFfaSupported;
+
 /**
   ArmFfaLib Constructor.
 
@@ -38,7 +41,10 @@ ArmFfaLibBaseConstructor (
 {
   EFI_STATUS  Status;
 
-  Status = ArmFfaLibCommonInit ();
+  mPartId         = 0xFFFF;
+  mIsFfaSupported = FALSE;
+
+  Status = ArmFfaLibCommonInit (&mPartId, &mIsFfaSupported);
   if (Status == EFI_UNSUPPORTED) {
     /*
      * EFI_UNSUPPORTED means FF-A interface isn't available.
@@ -55,3 +61,53 @@ ArmFfaLibBaseConstructor (
 
   return Status;
 }
+
+/**
+  Return partition or VM ID
+
+  @retval EFI_SUCCESS
+  @retval Others       Errors
+
+**/
+EFI_STATUS
+EFIAPI
+ArmFfaLibGetPartId (
+  IN UINT16  *PartId
+  )
+{
+  if (PartId != NULL) {
+    *PartId = mPartId;
+  }
+
+  return EFI_SUCCESS;
+}
+
+/**
+  Check FF-A support or not.
+
+  @retval TRUE                   Supported
+  @retval FALSE                  Not supported
+
+**/
+BOOLEAN
+EFIAPI
+IsFfaSupported (
+  IN VOID
+  )
+{
+  return mIsFfaSupported;
+}
+
+/**
+  Callback for when Unmap is called to handle any post unmap
+  functionality.
+
+**/
+VOID
+EFIAPI
+UnmapCallback (
+  IN VOID
+  )
+{
+  // Do nothing
+}
diff --git a/MdeModulePkg/Library/ArmFfaLib/ArmFfaPeiLib.c b/MdeModulePkg/Library/ArmFfaLib/ArmFfaPeiLib.c
index c21d5f2ba0..6faec34d57 100644
--- a/MdeModulePkg/Library/ArmFfaLib/ArmFfaPeiLib.c
+++ b/MdeModulePkg/Library/ArmFfaLib/ArmFfaPeiLib.c
@@ -102,8 +102,10 @@ ArmFfaPeiLibConstructor (
   EFI_HOB_MEMORY_ALLOCATION  *RxTxBufferAllocationHob;
   UINTN                      Property1;
   UINTN                      Property2;
+  UINT16                     PartId;
+  BOOLEAN                    IsFfaSupported;
 
-  Status = ArmFfaLibCommonInit ();
+  Status = ArmFfaLibCommonInit (&PartId, &IsFfaSupported);
   if (EFI_ERROR (Status)) {
     if (Status == EFI_UNSUPPORTED) {
       /*
@@ -220,3 +222,72 @@ ArmFfaPeiLibConstructor (
 
   return EFI_SUCCESS;
 }
+
+/**
+  Return partition or VM ID
+
+  @retval EFI_SUCCESS
+  @retval Others       Errors
+
+**/
+EFI_STATUS
+EFIAPI
+ArmFfaLibGetPartId (
+  IN UINT16  *PartId
+  )
+{
+  return ArmFfaLibPartitionIdGet (PartId);
+}
+
+/**
+  Check FF-A support or not.
+
+  @retval TRUE                   Supported
+  @retval FALSE                  Not supported
+
+**/
+BOOLEAN
+EFIAPI
+IsFfaSupported (
+  IN VOID
+  )
+{
+  EFI_STATUS  Status;
+
+  Status = ArmFfaLibIsFfaSupported ();
+  if (EFI_ERROR (Status)) {
+    return FALSE;
+  }
+
+  return TRUE;
+}
+
+/**
+  Callback for when Unmap is called to handle any post unmap
+  functionality. In PEI the Rx/Tx buffer HOB needs to be
+  invalidated.
+
+**/
+VOID
+EFIAPI
+UnmapCallback (
+  IN VOID
+  )
+{
+  EFI_HOB_MEMORY_ALLOCATION  *RxTxBufferAllocationHob;
+
+  /*
+   * Rx/Tx buffers are allocated with continuous pages.
+   * See ArmFfaLibRxTxMap(). If HOB is not found, the Rx/Tx
+   * buffers were not successfully mapped.
+   */
+  RxTxBufferAllocationHob = FindRxTxBufferAllocationHob (TRUE);
+  if (RxTxBufferAllocationHob == NULL) {
+    return;
+  }
+
+  /*
+   * Invalidate the HOB.
+   */
+  ZeroMem (RxTxBufferAllocationHob, sizeof (ARM_FFA_RX_TX_BUFFER_INFO));
+}
diff --git a/MdeModulePkg/Library/ArmFfaLib/ArmFfaRxTxMap.c b/MdeModulePkg/Library/ArmFfaLib/ArmFfaRxTxMap.c
index 4f0e1a4e8f..411d880e25 100644
--- a/MdeModulePkg/Library/ArmFfaLib/ArmFfaRxTxMap.c
+++ b/MdeModulePkg/Library/ArmFfaLib/ArmFfaRxTxMap.c
@@ -189,11 +189,14 @@ ArmFfaLibRxTxUnmap (
   EFI_STATUS    Status;
   ARM_FFA_ARGS  FfaArgs;
   VOID          *Buffers;
+  UINT16        PartId;
+
+  ArmFfaLibGetPartId (&PartId);
 
   ZeroMem (&FfaArgs, sizeof (ARM_FFA_ARGS));
 
   FfaArgs.Arg0 = ARM_FID_FFA_RXTX_UNMAP;
-  FfaArgs.Arg1 = (gPartId << ARM_FFA_SOURCE_EP_SHIFT);
+  FfaArgs.Arg1 = (PartId << ARM_FFA_SOURCE_EP_SHIFT);
 
   ArmCallFfa (&FfaArgs);
 
@@ -215,6 +218,8 @@ ArmFfaLibRxTxUnmap (
   PcdSet64S (PcdFfaTxBuffer, 0x00);
   PcdSet64S (PcdFfaRxBuffer, 0x00);
 
+  UnmapCallback ();
+
   return EFI_SUCCESS;
 }
 
@@ -283,6 +288,9 @@ RemapFfaRxTxBuffer (
   UINTN                      NewRxBuffer;
   UINTN                      MinSizeAndAlign;
   EFI_HOB_MEMORY_ALLOCATION  *RxTxBufferAllocationHob;
+  UINT16                     PartId;
+
+  ArmFfaLibGetPartId (&PartId);
 
   RxTxBufferAllocationHob = FindRxTxBufferAllocationHob (TRUE);
   if (RxTxBufferAllocationHob == NULL) {
@@ -293,7 +301,7 @@ RemapFfaRxTxBuffer (
 
   ZeroMem (&FfaArgs, sizeof (ARM_FFA_ARGS));
   FfaArgs.Arg0 = ARM_FID_FFA_RXTX_UNMAP;
-  FfaArgs.Arg1 = (gPartId << ARM_FFA_SOURCE_EP_SHIFT);
+  FfaArgs.Arg1 = (PartId << ARM_FFA_SOURCE_EP_SHIFT);
 
   ArmCallFfa (&FfaArgs);
   Status = FfaArgsToEfiStatus (&FfaArgs);
diff --git a/MdeModulePkg/Library/ArmFfaLib/ArmFfaRxTxMap.h b/MdeModulePkg/Library/ArmFfaLib/ArmFfaRxTxMap.h
index 5140865b5f..87e2545b0a 100644
--- a/MdeModulePkg/Library/ArmFfaLib/ArmFfaRxTxMap.h
+++ b/MdeModulePkg/Library/ArmFfaLib/ArmFfaRxTxMap.h
@@ -96,4 +96,15 @@ RemapFfaRxTxBuffer (
   IN OUT ARM_FFA_RX_TX_BUFFER_INFO  *BufferInfo
   );
 
+/**
+  Callback for when Unmap is called to handle any post unmap
+  functionality.
+
+**/
+VOID
+EFIAPI
+UnmapCallback (
+  IN VOID
+  );
+
 #endif
diff --git a/MdeModulePkg/Library/ArmFfaLib/ArmFfaSecLib.c b/MdeModulePkg/Library/ArmFfaLib/ArmFfaSecLib.c
index 4ee471d795..30017aaf90 100644
--- a/MdeModulePkg/Library/ArmFfaLib/ArmFfaSecLib.c
+++ b/MdeModulePkg/Library/ArmFfaLib/ArmFfaSecLib.c
@@ -48,13 +48,13 @@ ArmFfaSecLibConstructor (
   IN VOID
   )
 {
-  EFI_STATUS                 Status;
-  ARM_FFA_RX_TX_BUFFER_INFO  *BufferInfo;
-  EFI_HOB_MEMORY_ALLOCATION  *RxTxBufferAllocationHob;
-  UINTN                      Property1;
-  UINTN                      Property2;
+  EFI_STATUS  Status;
+  UINTN       Property1;
+  UINTN       Property2;
+  UINT16      PartId;
+  BOOLEAN     IsFfaSupported;
 
-  Status = ArmFfaLibCommonInit ();
+  Status = ArmFfaLibCommonInit (&PartId, &IsFfaSupported);
   if (EFI_ERROR (Status)) {
     if (Status == EFI_UNSUPPORTED) {
       /*
@@ -91,39 +91,80 @@ ArmFfaSecLibConstructor (
                &Property2
                );
     if (!EFI_ERROR (Status)) {
-      DEBUG ((DEBUG_INFO, "%a Rx/Tx buffer doesn't support.\n", __func__));
+      DEBUG ((DEBUG_INFO, "%a Rx/Tx buffer isn't supported.\n", __func__));
     }
 
     return Status;
   }
 
-  BufferInfo = BuildGuidHob (
-                 &gArmFfaRxTxBufferInfoGuid,
-                 sizeof (ARM_FFA_RX_TX_BUFFER_INFO)
-                 );
-  if (BufferInfo == NULL) {
-    DEBUG ((DEBUG_ERROR, "%a: Failed to create Rx/Tx Buffer Info Hob\n", __func__));
-    ArmFfaLibRxTxUnmap ();
-    return EFI_OUT_OF_RESOURCES;
+  return EFI_SUCCESS;
+}
+
+/**
+  Return partition or VM ID
+
+  @retval EFI_SUCCESS
+  @retval Others       Errors
+
+**/
+EFI_STATUS
+EFIAPI
+ArmFfaLibGetPartId (
+  IN UINT16  *PartId
+  )
+{
+  return ArmFfaLibPartitionIdGet (PartId);
+}
+
+/**
+  Check FF-A support or not.
+
+  @retval TRUE                   Supported
+  @retval FALSE                  Not supported
+
+**/
+BOOLEAN
+EFIAPI
+IsFfaSupported (
+  IN VOID
+  )
+{
+  EFI_STATUS  Status;
+
+  Status = ArmFfaLibIsFfaSupported ();
+  if (EFI_ERROR (Status)) {
+    return FALSE;
   }
 
-  RxTxBufferAllocationHob = FindRxTxBufferAllocationHob (FALSE);
-  ASSERT (RxTxBufferAllocationHob != NULL);
+  return TRUE;
+}
+
+/**
+  Callback for when Unmap is called to handle any post unmap
+  functionality. In SEC, the Rx/Tx buffer HOB needs to be
+  invalidated.
+
+**/
+VOID
+EFIAPI
+UnmapCallback (
+  IN VOID
+  )
+{
+  EFI_HOB_MEMORY_ALLOCATION  *RxTxBufferAllocationHob;
 
   /*
-   * Set then Name with gArmFfaRxTxBufferInfoGuid, so that ArmFfaPeiLib or
-   * ArmFfaDxeLib can find the Rx/Tx buffer allocation area.
+   * Rx/Tx buffers are allocated with continuous pages.
+   * See ArmFfaLibRxTxMap(). If HOB is not found, the Rx/Tx
+   * buffers were not successfully mapped.
    */
-  CopyGuid (
-    &RxTxBufferAllocationHob->AllocDescriptor.Name,
-    &gArmFfaRxTxBufferInfoGuid
-    );
-
-  UpdateRxTxBufferInfo (BufferInfo);
-  BufferInfo->RemapOffset =
-    (UINTN)(BufferInfo->TxBufferAddr -
-            RxTxBufferAllocationHob->AllocDescriptor.MemoryBaseAddress);
-  BufferInfo->RemapRequired = TRUE;
+  RxTxBufferAllocationHob = FindRxTxBufferAllocationHob (TRUE);
+  if (RxTxBufferAllocationHob == NULL) {
+    return;
+  }
 
-  return EFI_SUCCESS;
+  /*
+   * Invalidate the HOB.
+   */
+  ZeroMem (RxTxBufferAllocationHob, sizeof (ARM_FFA_RX_TX_BUFFER_INFO));
 }
diff --git a/MdeModulePkg/Library/ArmFfaLib/ArmFfaSecRxTxMap.c b/MdeModulePkg/Library/ArmFfaLib/ArmFfaSecRxTxMap.c
index 7c44e80822..8b88621516 100644
--- a/MdeModulePkg/Library/ArmFfaLib/ArmFfaSecRxTxMap.c
+++ b/MdeModulePkg/Library/ArmFfaLib/ArmFfaSecRxTxMap.c
@@ -31,9 +31,6 @@
 #include "ArmFfaCommon.h"
 #include "ArmFfaRxTxMap.h"
 
-STATIC VOID  *mTxBuffer;
-STATIC VOID  *mRxBuffer;
-
 /**
   Get mapped Rx/Tx buffers.
 
@@ -55,24 +52,34 @@ ArmFfaLibGetRxTxBuffers (
   OUT UINT64  *RxBufferSize OPTIONAL
   )
 {
-  if ((mTxBuffer == NULL) || (mRxBuffer == NULL)) {
+  EFI_HOB_MEMORY_ALLOCATION  *RxTxBufferAllocationHob;
+  UINT64                     BufferSize;
+
+  /*
+   * Attempt to find the allocation HOB, if it exists then the
+   * Rx/Tx buffers have already been mapped.
+   */
+  RxTxBufferAllocationHob = FindRxTxBufferAllocationHob (TRUE);
+  if (RxTxBufferAllocationHob == NULL) {
     return EFI_NOT_READY;
   }
 
+  BufferSize = PcdGet64 (PcdFfaTxRxPageCount) * EFI_PAGE_SIZE;
+
   if (TxBuffer != NULL) {
-    *TxBuffer = mTxBuffer;
+    *TxBuffer = (VOID *)RxTxBufferAllocationHob->AllocDescriptor.MemoryBaseAddress;
   }
 
   if (TxBufferSize != NULL) {
-    *TxBufferSize = PcdGet64 (PcdFfaTxRxPageCount) * EFI_PAGE_SIZE;
+    *TxBufferSize = BufferSize;
   }
 
   if (RxBuffer != NULL) {
-    *RxBuffer = mRxBuffer;
+    *RxBuffer = (VOID *)RxTxBufferAllocationHob->AllocDescriptor.MemoryBaseAddress + BufferSize;
   }
 
   if (RxBufferSize != NULL) {
-    *RxBufferSize = PcdGet64 (PcdFfaTxRxPageCount) * EFI_PAGE_SIZE;
+    *RxBufferSize = BufferSize;
   }
 
   return EFI_SUCCESS;
@@ -96,21 +103,25 @@ ArmFfaLibRxTxMap (
   IN VOID
   )
 {
-  EFI_STATUS    Status;
-  ARM_FFA_ARGS  FfaArgs;
-  UINTN         Property1;
-  UINTN         Property2;
-  UINTN         MinSizeAndAlign;
-  UINTN         MaxSize;
-  VOID          *Buffers;
-  VOID          *TxBuffer;
-  VOID          *RxBuffer;
+  EFI_STATUS                 Status;
+  ARM_FFA_ARGS               FfaArgs;
+  UINTN                      Property1;
+  UINTN                      Property2;
+  UINTN                      MinSizeAndAlign;
+  UINTN                      MaxSize;
+  VOID                       *Buffers;
+  VOID                       *TxBuffer;
+  VOID                       *RxBuffer;
+  ARM_FFA_RX_TX_BUFFER_INFO  *BufferInfo;
+  EFI_HOB_MEMORY_ALLOCATION  *RxTxBufferAllocationHob;
+  UINT16                     PartId;
 
   /*
-   * If someone already mapped Rx/Tx Buffers, return EFI_ALREADY_STARTED.
-   * return EFI_ALREADY_STARTED.
+   * Attempt to find the allocation HOB, if it exists then the
+   * Rx/Tx buffers have already been mapped.
    */
-  if ((mTxBuffer != NULL) && (mRxBuffer != NULL)) {
+  RxTxBufferAllocationHob = FindRxTxBufferAllocationHob (TRUE);
+  if (RxTxBufferAllocationHob != NULL) {
     return EFI_ALREADY_STARTED;
   }
 
@@ -123,7 +134,7 @@ ArmFfaLibRxTxMap (
   if (EFI_ERROR (Status)) {
     DEBUG ((
       DEBUG_ERROR,
-      "%a: Failed to get RX/TX buffer property... Status: %r\n",
+      "%a: Failed to get Rx/Tx buffer property... Status: %r\n",
       __func__,
       Status
       ));
@@ -205,11 +216,50 @@ ArmFfaLibRxTxMap (
     goto ErrorHandler;
   }
 
-  mTxBuffer = TxBuffer;
-  mRxBuffer = RxBuffer;
+  BufferInfo = BuildGuidHob (&gArmFfaRxTxBufferInfoGuid, sizeof (ARM_FFA_RX_TX_BUFFER_INFO));
+  if (BufferInfo == NULL) {
+    DEBUG ((DEBUG_ERROR, "%a: Failed to create Rx/Tx Buffer Info Hob\n", __func__));
+    Status = EFI_OUT_OF_RESOURCES;
+    goto Unmap;
+  }
+
+  RxTxBufferAllocationHob = GetRxTxBufferAllocationHob ((EFI_PHYSICAL_ADDRESS)((UINTN)TxBuffer), PcdGet64 (PcdFfaTxRxPageCount) * EFI_PAGE_SIZE * 2, FALSE);
+  ASSERT (RxTxBufferAllocationHob != NULL);
+
+  /*
+   * Set then Name with gArmFfaRxTxBufferInfoGuid, so that ArmFfaPeiLib or
+   * ArmFfaDxeLib can find the Rx/Tx buffer allocation area.
+   */
+  CopyGuid (&RxTxBufferAllocationHob->AllocDescriptor.Name, &gArmFfaRxTxBufferInfoGuid);
+
+  BufferInfo->TxBufferAddr  = (UINTN)RxTxBufferAllocationHob->AllocDescriptor.MemoryBaseAddress;
+  BufferInfo->TxBufferSize  = PcdGet64 (PcdFfaTxRxPageCount) * EFI_PAGE_SIZE;
+  BufferInfo->RxBufferAddr  = (UINTN)RxTxBufferAllocationHob->AllocDescriptor.MemoryBaseAddress + BufferInfo->TxBufferSize;
+  BufferInfo->RxBufferSize  = PcdGet64 (PcdFfaTxRxPageCount) * EFI_PAGE_SIZE;
+  BufferInfo->RemapOffset   = (UINTN)(BufferInfo->TxBufferAddr - RxTxBufferAllocationHob->AllocDescriptor.MemoryBaseAddress);
+  BufferInfo->RemapRequired = TRUE;
 
   return EFI_SUCCESS;
 
+Unmap:
+  Status = ArmFfaLibPartitionIdGet (&PartId);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "%a, Failed to acquire partition ID. Status: %r\n", __func__, Status));
+    goto ErrorHandler;
+  }
+
+  ZeroMem (&FfaArgs, sizeof (ARM_FFA_ARGS));
+
+  FfaArgs.Arg0 = ARM_FID_FFA_RXTX_UNMAP;
+  FfaArgs.Arg1 = (PartId << ARM_FFA_SOURCE_EP_SHIFT);
+
+  ArmCallFfa (&FfaArgs);
+
+  Status = FfaArgsToEfiStatus (&FfaArgs);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "%a, Failed to unmap Rx/Tx buffer. Status: %r\n", __func__, Status));
+  }
+
 ErrorHandler:
   FreePages (Buffers, (PcdGet64 (PcdFfaTxRxPageCount) * 2));
 
@@ -222,8 +272,8 @@ ArmFfaLibRxTxMap (
   Rx/Tx buffer is registered only once per partition.
 
   @retval EFI_SUCCESS
-  @retval EFI_INVALID_PARAMETERS               Already unregistered
-  @retval EFI_UNSUPPORTED                      Not supported
+  @retval EFI_INVALID_PARAMETER  Already unregistered
+  @retval EFI_UNSUPPORTED        Not supported
 
 **/
 EFI_STATUS
@@ -232,14 +282,32 @@ ArmFfaLibRxTxUnmap (
   IN VOID
   )
 {
-  EFI_STATUS    Status;
-  ARM_FFA_ARGS  FfaArgs;
-  VOID          *Buffers;
+  EFI_STATUS                 Status;
+  ARM_FFA_ARGS               FfaArgs;
+  VOID                       *Buffers;
+  EFI_HOB_MEMORY_ALLOCATION  *RxTxBufferAllocationHob;
+  UINT16                     PartId;
+
+  /*
+   * Rx/Tx buffers are allocated with continuous pages.
+   * See ArmFfaLibRxTxMap(). If HOB is not found, the Rx/Tx
+   * buffers were not successfully mapped.
+   */
+  RxTxBufferAllocationHob = FindRxTxBufferAllocationHob (TRUE);
+  if (RxTxBufferAllocationHob == NULL) {
+    return EFI_INVALID_PARAMETER;
+  }
+
+  Status = ArmFfaLibPartitionIdGet (&PartId);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "%a, Failed to acquire partition ID. Status: %r\n", __func__, Status));
+    return EFI_INVALID_PARAMETER;
+  }
 
   ZeroMem (&FfaArgs, sizeof (ARM_FFA_ARGS));
 
   FfaArgs.Arg0 = ARM_FID_FFA_RXTX_UNMAP;
-  FfaArgs.Arg1 = (gPartId << ARM_FFA_SOURCE_EP_SHIFT);
+  FfaArgs.Arg1 = (PartId << ARM_FFA_SOURCE_EP_SHIFT);
 
   ArmCallFfa (&FfaArgs);
 
@@ -248,26 +316,20 @@ ArmFfaLibRxTxUnmap (
     return Status;
   }
 
-  /*
-   * Rx/Tx Buffer are allocated with continuous pages.
-   * and start address of these pages is set on PcdFfaTxBuffer.
-   * See ArmFfaLibRxTxMap().
-   */
-  Buffers = mTxBuffer;
+  Buffers = (VOID *)RxTxBufferAllocationHob->AllocDescriptor.MemoryBaseAddress;
   if (Buffers != NULL) {
     FreePages (Buffers, (PcdGet64 (PcdFfaTxRxPageCount) * 2));
   }
 
-  mTxBuffer = NULL;
-  mRxBuffer = NULL;
+  UnmapCallback ();
 
   return EFI_SUCCESS;
 }
 
 /**
-  Update Rx/TX buffer information.
+  Update Rx/Tx buffer information.
 
-  @param  BufferInfo            Rx/Tx buffer information.
+  @param  BufferInfo  Rx/Tx buffer information.
 
 **/
 VOID
@@ -276,14 +338,24 @@ UpdateRxTxBufferInfo (
   OUT ARM_FFA_RX_TX_BUFFER_INFO  *BufferInfo
   )
 {
-  BufferInfo->TxBufferAddr = (UINTN)mTxBuffer;
-  BufferInfo->TxBufferSize = PcdGet64 (PcdFfaTxRxPageCount) * EFI_PAGE_SIZE;
-  BufferInfo->RxBufferAddr = (UINTN)mRxBuffer;
-  BufferInfo->RxBufferSize = PcdGet64 (PcdFfaTxRxPageCount) * EFI_PAGE_SIZE;
+  EFI_HOB_MEMORY_ALLOCATION  *RxTxBufferAllocationHob;
+
+  /*
+   * Attempt to find the HOB containing the Rx/Tx buffer information.
+   * If it is not found, the Rx/Tx buffers were not successfully mapped.
+   */
+  RxTxBufferAllocationHob = FindRxTxBufferAllocationHob (TRUE);
+
+  if ((RxTxBufferAllocationHob != NULL) && (BufferInfo != NULL)) {
+    BufferInfo->TxBufferAddr = (UINTN)RxTxBufferAllocationHob->AllocDescriptor.MemoryBaseAddress;
+    BufferInfo->TxBufferSize = PcdGet64 (PcdFfaTxRxPageCount) * EFI_PAGE_SIZE;
+    BufferInfo->RxBufferAddr = (UINTN)RxTxBufferAllocationHob->AllocDescriptor.MemoryBaseAddress + BufferInfo->TxBufferSize;
+    BufferInfo->RxBufferSize = PcdGet64 (PcdFfaTxRxPageCount) * EFI_PAGE_SIZE;
+  }
 }
 
 /**
-  Find Rx/TX buffer memory allocation hob.
+  Find Rx/Tx buffer memory allocation hob.
 
   @param  UseGuid             Find MemoryAllocationHob using Guid.
 
@@ -297,17 +369,18 @@ FindRxTxBufferAllocationHob (
   IN BOOLEAN  UseGuid
   )
 {
-  EFI_PHYSICAL_ADDRESS  BufferBase;
-  UINT64                BufferSize;
-
-  BufferBase = (EFI_PHYSICAL_ADDRESS)((UINTN)mTxBuffer);
-  BufferSize = PcdGet64 (PcdFfaTxRxPageCount) * EFI_PAGE_SIZE * 2;
+  (void)UseGuid;
 
-  return GetRxTxBufferAllocationHob (BufferBase, BufferSize, UseGuid);
+  /*
+   * Mapping the Rx/Tx buffer should have generated the HOB, if not,
+   * then the allocation HOB will not be found and NULL will be returned.
+   * Note that BufferAddr and BufferSize are ignored when UseGuid is TRUE.
+   */
+  return GetRxTxBufferAllocationHob (0, 0, TRUE);
 }
 
 /**
-  Remap Rx/TX buffer with converted Rx/Tx Buffer address after
+  Remap Rx/Tx buffer with converted Rx/Tx Buffer address after
   using permanent memory.
 
   @param[out] BufferInfo    BufferInfo
diff --git a/MdeModulePkg/Library/ArmFfaLib/ArmFfaStandaloneMmLib.c b/MdeModulePkg/Library/ArmFfaLib/ArmFfaStandaloneMmLib.c
index 72f9a80b27..dc8cc86df6 100644
--- a/MdeModulePkg/Library/ArmFfaLib/ArmFfaStandaloneMmLib.c
+++ b/MdeModulePkg/Library/ArmFfaLib/ArmFfaStandaloneMmLib.c
@@ -24,6 +24,9 @@
 #include "ArmFfaCommon.h"
 #include "ArmFfaRxTxMap.h"
 
+STATIC UINT16   mPartId;
+STATIC BOOLEAN  mIsFfaSupported;
+
 /**
   ArmFfaLib Constructor.
 
@@ -43,7 +46,10 @@ ArmFfaStandaloneMmLibConstructor (
 {
   EFI_STATUS  Status;
 
-  Status = ArmFfaLibCommonInit ();
+  mPartId         = 0xFFFF;
+  mIsFfaSupported = FALSE;
+
+  Status = ArmFfaLibCommonInit (&mPartId, &mIsFfaSupported);
   if (Status == EFI_UNSUPPORTED) {
     /*
      * EFI_UNSUPPORTED means FF-A interface isn't available.
@@ -73,3 +79,53 @@ ArmFfaStandaloneMmLibConstructor (
 
   return Status;
 }
+
+/**
+  Return partition or VM ID
+
+  @retval EFI_SUCCESS
+  @retval Others       Errors
+
+**/
+EFI_STATUS
+EFIAPI
+ArmFfaLibGetPartId (
+  IN UINT16  *PartId
+  )
+{
+  if (PartId != NULL) {
+    *PartId = mPartId;
+  }
+
+  return EFI_SUCCESS;
+}
+
+/**
+  Check FF-A support or not.
+
+  @retval TRUE                   Supported
+  @retval FALSE                  Not supported
+
+**/
+BOOLEAN
+EFIAPI
+IsFfaSupported (
+  IN VOID
+  )
+{
+  return mIsFfaSupported;
+}
+
+/**
+  Callback for when Unmap is called to handle any post unmap
+  functionality.
+
+**/
+VOID
+EFIAPI
+UnmapCallback (
+  IN VOID
+  )
+{
+  // Do nothing
+}
diff --git a/MdeModulePkg/Library/ArmFfaLib/ArmFfaStandaloneMmRxTxMap.c b/MdeModulePkg/Library/ArmFfaLib/ArmFfaStandaloneMmRxTxMap.c
index d368c329ee..93a0a7ee2b 100644
--- a/MdeModulePkg/Library/ArmFfaLib/ArmFfaStandaloneMmRxTxMap.c
+++ b/MdeModulePkg/Library/ArmFfaLib/ArmFfaStandaloneMmRxTxMap.c
@@ -214,6 +214,9 @@ ArmFfaLibRxTxUnmap (
   EFI_STATUS    Status;
   ARM_FFA_ARGS  FfaArgs;
   VOID          *Buffers;
+  UINT16        SourceId;
+
+  ArmFfaLibGetPartId (&SourceId);
 
   if (mArmFfaRxTxBufferStmmInfoHandle == NULL) {
     // This means that the agent tried to unmap the buffers before even know them.
@@ -224,7 +227,7 @@ ArmFfaLibRxTxUnmap (
   ZeroMem (&FfaArgs, sizeof (ARM_FFA_ARGS));
 
   FfaArgs.Arg0 = ARM_FID_FFA_RXTX_UNMAP;
-  FfaArgs.Arg1 = (gPartId << ARM_FFA_SOURCE_EP_SHIFT);
+  FfaArgs.Arg1 = (SourceId << ARM_FFA_SOURCE_EP_SHIFT);
 
   ArmCallFfa (&FfaArgs);