[PATCH 2/4] ArmVirtPkg/ArmVirtMemoryInitPeiLib: check for capsules before memory init

[Ard Biesheuvel <ard.biesheuvel@linaro.org>]

Look for any capsules left in memory by the OS across reset before
releasing the memory for normal use, so that they can be preserved and
processed later.

Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Ard Biesheuvel <ard.biesheuvel@linaro.org>
---
 ArmVirtPkg/Library/ArmVirtMemoryInitPeiLib/ArmVirtMemoryInitPeiLib.c   | 60 +++++++++++++++++++-
 ArmVirtPkg/Library/ArmVirtMemoryInitPeiLib/ArmVirtMemoryInitPeiLib.inf |  9 ++-
 2 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/ArmVirtPkg/Library/ArmVirtMemoryInitPeiLib/ArmVirtMemoryInitPeiLib.c b/ArmVirtPkg/Library/ArmVirtMemoryInitPeiLib/ArmVirtMemoryInitPeiLib.c
index 6f3e54b7afcb..7f55f634e4e5 100644
--- a/ArmVirtPkg/Library/ArmVirtMemoryInitPeiLib/ArmVirtMemoryInitPeiLib.c
+++ b/ArmVirtPkg/Library/ArmVirtMemoryInitPeiLib/ArmVirtMemoryInitPeiLib.c
@@ -17,11 +17,15 @@
 
 #include <Library/ArmMmuLib.h>
 #include <Library/ArmPlatformLib.h>
+#include <Library/CacheMaintenanceLib.h>
 #include <Library/DebugLib.h>
 #include <Library/HobLib.h>
 #include <Library/MemoryAllocationLib.h>
 #include <Library/PcdLib.h>
-#include <Library/CacheMaintenanceLib.h>
+#include <Library/PeiServicesLib.h>
+#include <Library/PeiServicesTablePointerLib.h>
+
+#include <Ppi/Capsule.h>
 
 VOID
 BuildMemoryTypeInformationHob (
@@ -49,6 +53,58 @@ InitMmu (
   }
 }
 
+STATIC
+VOID
+CheckCapsule (
+  IN EFI_PHYSICAL_ADDRESS               UefiMemoryBase,
+  IN UINT64                             UefiMemorySize
+  )
+{
+  EFI_STATUS         Status;
+  EFI_PEI_SERVICES   **PeiServices;
+  PEI_CAPSULE_PPI    *Capsule;
+  VOID               *CapsuleBuffer;
+  UINTN              CapsuleBufferLength;
+
+  PeiServices = (EFI_PEI_SERVICES **) GetPeiServicesTablePointer ();
+  ASSERT (PeiServices != NULL);
+
+  //
+  // Check for persistent capsules
+  //
+  Status = PeiServicesLocatePpi (&gPeiCapsulePpiGuid, 0, NULL,
+             (VOID **)&Capsule);
+  if (Status == EFI_SUCCESS) {
+    Status = Capsule->CheckCapsuleUpdate (PeiServices);
+    if (Status == EFI_SUCCESS) {
+
+      CapsuleBuffer = (VOID *)((UINTN)FixedPcdGet32 (PcdCPUCoresStackBase) +
+                               FixedPcdGet32 (PcdCPUCorePrimaryStackSize));
+      CapsuleBufferLength = (UINTN)UefiMemoryBase - (UINTN)CapsuleBuffer;
+
+      PeiServicesSetBootMode (BOOT_ON_FLASH_UPDATE);
+      Status = Capsule->Coalesce (PeiServices, &CapsuleBuffer,
+                          &CapsuleBufferLength);
+      if (!EFI_ERROR (Status)) {
+        DEBUG ((DEBUG_INFO, "%a: Coalesced capsule @ %p (0x%lx) capsule\n",
+                __FUNCTION__, CapsuleBuffer, CapsuleBufferLength));
+      } else {
+        DEBUG ((DEBUG_WARN, "%a: failed to coalesce() capsule (Status == %r)\n",
+                __FUNCTION__, Status));
+        return;
+      }
+
+      Status = Capsule->CreateState (PeiServices, CapsuleBuffer,
+                          CapsuleBufferLength);
+
+      if (EFI_ERROR (Status)) {
+        DEBUG ((DEBUG_WARN, "%a: Capsule->CreateState failed (Status == %r)\n",
+                __FUNCTION__, Status));
+      }
+    }
+  }
+}
+
 EFI_STATUS
 EFIAPI
 MemoryPeim (
@@ -109,6 +165,8 @@ MemoryPeim (
   // Build Memory Allocation Hob
   InitMmu ();
 
+  CheckCapsule (UefiMemoryBase, UefiMemorySize);
+
   if (FeaturePcdGet (PcdPrePiProduceMemoryTypeInformationHob)) {
     // Optional feature that helps prevent EFI memory map fragmentation.
     BuildMemoryTypeInformationHob ();
diff --git a/ArmVirtPkg/Library/ArmVirtMemoryInitPeiLib/ArmVirtMemoryInitPeiLib.inf b/ArmVirtPkg/Library/ArmVirtMemoryInitPeiLib/ArmVirtMemoryInitPeiLib.inf
index 028d6fb5ac28..4524afd2c7ed 100644
--- a/ArmVirtPkg/Library/ArmVirtMemoryInitPeiLib/ArmVirtMemoryInitPeiLib.inf
+++ b/ArmVirtPkg/Library/ArmVirtMemoryInitPeiLib/ArmVirtMemoryInitPeiLib.inf
@@ -37,6 +37,8 @@ [LibraryClasses]
   ArmMmuLib
   ArmPlatformLib
   CacheMaintenanceLib
+  PeiServicesLib
+  PeiServicesTablePointerLib
 
 [Guids]
   gEfiMemoryTypeInformationGuid
@@ -48,6 +50,8 @@ [FixedPcd]
   gArmTokenSpaceGuid.PcdFdSize
 
   gArmPlatformTokenSpaceGuid.PcdSystemMemoryUefiRegionSize
+  gArmPlatformTokenSpaceGuid.PcdCPUCoresStackBase
+  gArmPlatformTokenSpaceGuid.PcdCPUCorePrimaryStackSize
 
   gEmbeddedTokenSpaceGuid.PcdMemoryTypeEfiACPIReclaimMemory
   gEmbeddedTokenSpaceGuid.PcdMemoryTypeEfiACPIMemoryNVS
@@ -64,5 +68,8 @@ [Pcd]
   gArmTokenSpaceGuid.PcdSystemMemorySize
   gArmTokenSpaceGuid.PcdFdBaseAddress
 
+[Ppis]
+  gPeiCapsulePpiGuid
+
 [Depex]
-  TRUE
+   gPeiCapsulePpiGuid
-- 
2.7.4