MdeModulePkg/Library: Fix memory leak for ArmFfaLib's Rx/Tx buffer and add ArmFfaLib for PeilessSec

MdeModulePkg/Include/Guid/ArmFfaRxTxBufferInfo.h

@@ -9,21 +9,29 @@
 #ifndef ARM_FFA_RX_TX_BUFFER_INFO_H_
 #define ARM_FFA_RX_TX_BUFFER_INFO_H_
 
+#include <Uefi/UefiBaseType.h>
+
 /**
  * Guid Hob Data for gArmFfaRxTxBufferInfoGuid Guid Hob.
  */
 typedef struct ArmFfaRxTxBuffersInfo {
   /// Tx Buffer Address.
-  VOID      *TxBufferAddr;
+  EFI_PHYSICAL_ADDRESS    TxBufferAddr;
 
   /// Tx Buffer Size.
-  UINT64    TxBufferSize;
+  UINT64                  TxBufferSize;
 
   /// Rx Buffer Address.
-  VOID      *RxBufferAddr;
+  EFI_PHYSICAL_ADDRESS    RxBufferAddr;
 
   /// Rx Buffer Size.
-  UINT64    RxBufferSize;
+  UINT64                  RxBufferSize;
+
+  /// Rx/Tx buffer should be remapped to permanent memory.
+  BOOLEAN                 RemapRequired;
+
+  /// Rx/Tx buffer offset from its allocation base.
+  UINT64                  RemapOffset;
 } ARM_FFA_RX_TX_BUFFER_INFO;
 
 extern EFI_GUID  gArmFfaRxTxBufferInfoGuid;

MdeModulePkg/Library/ArmFfaLib/ArmFfaCommon.c

@@ -21,12 +21,15 @@
 #include <Library/BaseLib.h>
 #include <Library/BaseMemoryLib.h>
 #include <Library/DebugLib.h>
+#include <Library/HobLib.h>
 #include <Library/PcdLib.h>
 
 #include <IndustryStandard/ArmFfaSvc.h>
 #include <IndustryStandard/ArmFfaPartInfo.h>
 #include <IndustryStandard/ArmStdSmc.h>
 
+#include <Guid/ArmFfaRxTxBufferInfo.h>
+
 #include "ArmFfaCommon.h"
 
 BOOLEAN  gFfaSupported;
@@ -752,3 +755,146 @@ ArmFfaLibCommonInit (
 
   return EFI_SUCCESS;
 }
+
+/**
+  Get first Rx/Tx Buffer allocation hob.
+  If UseGuid is TRUE, BufferAddr and BufferSize parameters are ignored.
+
+  @param[in]  BufferAddr       Buffer address
+  @param[in]  BufferSize       Buffer Size
+  @param[in]  UseGuid          Find MemoryAllocationHob using gArmFfaRxTxBufferInfoGuid.
+
+  @retval     NULL             Not found
+  @retval     Other            MemoryAllocationHob related to Rx/Tx buffer
+
+**/
+EFI_HOB_MEMORY_ALLOCATION *
+EFIAPI
+GetRxTxBufferAllocationHob (
+  IN EFI_PHYSICAL_ADDRESS  BufferAddr,
+  IN UINT64                BufferSize,
+  IN BOOLEAN               UseGuid
+  )
+{
+  EFI_PEI_HOB_POINTERS       Hob;
+  EFI_HOB_MEMORY_ALLOCATION  *MemoryAllocationHob;
+  EFI_PHYSICAL_ADDRESS       MemoryBase;
+  UINT64                     MemorySize;
+
+  if (!UseGuid && (BufferAddr == 0x00)) {
+    return NULL;
+  }
+
+  MemoryAllocationHob = NULL;
+  Hob.Raw             = GetFirstHob (EFI_HOB_TYPE_MEMORY_ALLOCATION);
+
+  while (Hob.Raw != NULL) {
+    if (Hob.MemoryAllocation->AllocDescriptor.MemoryType == EfiConventionalMemory) {
+      continue;
+    }
+
+    MemoryBase = Hob.MemoryAllocation->AllocDescriptor.MemoryBaseAddress;
+    MemorySize = Hob.MemoryAllocation->AllocDescriptor.MemoryLength;
+
+    if ((!UseGuid && (BufferAddr >= MemoryBase) &&
+         ((BufferAddr + BufferSize) <= (MemoryBase + MemorySize))) ||
+        (UseGuid && CompareGuid (
+                      &gArmFfaRxTxBufferInfoGuid,
+                      &Hob.MemoryAllocation->AllocDescriptor.Name
+                      )))
+    {
+      MemoryAllocationHob = (EFI_HOB_MEMORY_ALLOCATION *)Hob.Raw;
+      break;
+    }
+
+    Hob.Raw = GET_NEXT_HOB (Hob);
+    Hob.Raw = GetNextHob (EFI_HOB_TYPE_MEMORY_ALLOCATION, Hob.Raw);
+  }
+
+  return MemoryAllocationHob;
+}
+
+/**
+  Get Rx/Tx buffer MinSizeAndAign and MaxSize
+
+  @param[out] MinSizeAndAlign  Minimum size of Buffer.
+
+  @retval EFI_SUCCESS
+  @retval EFI_UNSUPPORTED          Wrong min size received from SPMC
+  @retval EFI_INVALID_PARAMETER    Wrong buffer size
+  @retval Others                   Failure of ArmFfaLibGetFeatures()
+
+**/
+EFI_STATUS
+EFIAPI
+GetRxTxBufferMinSizeAndAlign (
+  OUT UINTN  *MinSizeAndAlign
+  )
+{
+  EFI_STATUS  Status;
+  UINTN       MinAndAlign;
+  UINTN       MaxSize;
+  UINTN       Property1;
+  UINTN       Property2;
+
+  Status = ArmFfaLibGetFeatures (
+             ARM_FID_FFA_RXTX_MAP,
+             FFA_RXTX_MAP_INPUT_PROPERTY_DEFAULT,
+             &Property1,
+             &Property2
+             );
+  if (EFI_ERROR (Status)) {
+    DEBUG ((
+      DEBUG_ERROR,
+      "%a: Failed to get RX/TX buffer property... Status: %r\n",
+      __func__,
+      Status
+      ));
+    return Status;
+  }
+
+  MinAndAlign =
+    ((Property1 >>
+      ARM_FFA_BUFFER_MINSIZE_AND_ALIGN_SHIFT) &
+     ARM_FFA_BUFFER_MINSIZE_AND_ALIGN_MASK);
+
+  switch (MinAndAlign) {
+    case ARM_FFA_BUFFER_MINSIZE_AND_ALIGN_4K:
+      MinAndAlign = SIZE_4KB;
+      break;
+    case ARM_FFA_BUFFER_MINSIZE_AND_ALIGN_16K:
+      MinAndAlign = SIZE_16KB;
+      break;
+    case ARM_FFA_BUFFER_MINSIZE_AND_ALIGN_64K:
+      MinAndAlign = SIZE_64KB;
+      break;
+    default:
+      DEBUG ((DEBUG_ERROR, "%a: Invalid MinSizeAndAlign: 0x%x\n", __func__, MinAndAlign));
+      return EFI_UNSUPPORTED;
+  }
+
+  MaxSize =
+    (((Property1 >>
+       ARM_FFA_BUFFER_MAXSIZE_PAGE_COUNT_SHIFT) &
+      ARM_FFA_BUFFER_MAXSIZE_PAGE_COUNT_MASK));
+
+  MaxSize = ((MaxSize == 0) ? MAX_UINTN : (MaxSize * MinAndAlign));
+
+  if ((MinAndAlign > (PcdGet64 (PcdFfaTxRxPageCount) * EFI_PAGE_SIZE)) ||
+      (MaxSize < (PcdGet64 (PcdFfaTxRxPageCount) * EFI_PAGE_SIZE)))
+  {
+    DEBUG ((
+      DEBUG_ERROR,
+      "%a: Buffer is too small! MinSize: 0x%x, MaxSize: 0x%x, PageCount: %d\n",
+      __func__,
+      MinAndAlign,
+      MaxSize,
+      PcdGet64 (PcdFfaTxRxPageCount)
+      ));
+    return EFI_INVALID_PARAMETER;
+  }
+
+  *MinSizeAndAlign = MinAndAlign;
+
+  return EFI_SUCCESS;
+}

MdeModulePkg/Library/ArmFfaLib/ArmFfaCommon.h

@@ -49,4 +49,41 @@ ArmFfaLibCommonInit (
   IN VOID
   );
 
+/**
+  Get first Rx/Tx Buffer allocation hob.
+  If UseGuid is TRUE, BufferAddr and BufferSize parameters are ignored.
+
+  @param[in]  BufferAddr       Buffer address
+  @param[in]  BufferSize       Buffer Size
+  @param[in]  UseGuid          Find MemoryAllocationHob using gArmFfaRxTxBufferInfoGuid.
+
+  @retval     NULL             Not found
+  @retval     Other            MemoryAllocationHob related to Rx/Tx buffer
+
+**/
+EFI_HOB_MEMORY_ALLOCATION *
+EFIAPI
+GetRxTxBufferAllocationHob (
+  IN EFI_PHYSICAL_ADDRESS  BufferAddr,
+  IN UINT64                BufferSize,
+  IN BOOLEAN               UseGuid
+  );
+
+/**
+  Get Rx/Tx buffer MinSizeAndAign and MaxSize
+
+  @param[out] MinSizeAndAlign  Minimum size of Buffer.
+
+  @retval EFI_SUCCESS
+  @retval EFI_UNSUPPORTED          Wrong min size received from SPMC
+  @retval EFI_INVALID_PARAMETER    Wrong buffer size
+  @retval Others                   Failure of ArmFfaLibGetFeatures()
+
+**/
+EFI_STATUS
+EFIAPI
+GetRxTxBufferMinSizeAndAlign (
+  OUT UINTN  *MinSizeAndAlign
+  );
+
 #endif

MdeModulePkg/Library/ArmFfaLib/ArmFfaDxeLib.c

@@ -98,15 +98,32 @@ ArmFfaDxeLibConstructor (
     return EFI_SUCCESS;
   }
 
-  /*
-   * If PEIM uses ArmFfaPeiLib, the Rx/Tx buffers is already mapped in PEI phase.
-   * In this case, get Rx/Tx buffer info from Hob.
-   */
   RxTxBufferHob = GetFirstGuidHob (&gArmFfaRxTxBufferInfoGuid);
   if (RxTxBufferHob != NULL) {
     BufferInfo = GET_GUID_HOB_DATA (RxTxBufferHob);
-    PcdSet64S (PcdFfaTxBuffer, (UINTN)BufferInfo->TxBufferAddr);
-    PcdSet64S (PcdFfaRxBuffer, (UINTN)BufferInfo->RxBufferAddr);
+    if (!BufferInfo->RemapRequired) {
+      /*
+       * ArmFfaPeiLib handles the Rx/Tx buffer Remap and update the
+       * BufferInfo with permanant memory. So use it as it is.
+       */
+      PcdSet64S (PcdFfaTxBuffer, (UINTN)BufferInfo->TxBufferAddr);
+      PcdSet64S (PcdFfaRxBuffer, (UINTN)BufferInfo->RxBufferAddr);
+    } else {
+      /*
+       * SEC maps Rx/Tx buffer, But no PEIM module doesn't use
+       * ArmFfaPeiLib. In this case, the BufferInfo includes
+       * temporary Rx/Tx buffer address.
+       *
+       * Therefore, remap Rx/Tx buffer with migrated address again.
+       */
+      Status = RemapFfaRxTxBuffer (BufferInfo);
+      if (EFI_ERROR (Status)) {
+        DEBUG ((DEBUG_ERROR, "%a: Failed to remap Rx/Tx buffer... Status: %r\n", __func__, Status));
+        return Status;
+      }
+
+      BufferInfo->RemapRequired = FALSE;
+    }
   } else {
     Status = ArmFfaLibRxTxMap ();