Prevent VariablePei Memory Bucket Fragmentation

MdeModulePkg/Universal/Variable/Pei/Variable.c

@@ -1433,7 +1433,6 @@ BuildVariableRuntimeCacheInfoHob (
 {
   VARIABLE_RUNTIME_CACHE_INFO  TempHobBuffer;
   VARIABLE_RUNTIME_CACHE_INFO  *VariableRuntimeCacheInfo;
-  EFI_STATUS                   Status;
   VOID                         *Buffer;
   UINTN                        BufferSize;
   BOOLEAN                      NvAuthFlag;
@@ -1442,18 +1441,11 @@ BuildVariableRuntimeCacheInfoHob (
   ZeroMem (&TempHobBuffer, sizeof (VARIABLE_RUNTIME_CACHE_INFO));
 
   //
-  // AllocateRuntimePages for CACHE_INFO_FLAG and unblock it.
+  // AllocatePages for CACHE_INFO_FLAG buffer (will be relocated to runtime by DXE).
   //
   Pages  = EFI_SIZE_TO_PAGES (sizeof (CACHE_INFO_FLAG));
-  Buffer = AllocateRuntimePages (Pages);
+  Buffer = AllocatePages (Pages);
   ASSERT (Buffer != NULL);
-  Status = MmUnblockMemoryRequest (
-             (EFI_PHYSICAL_ADDRESS)(UINTN)Buffer,
-             Pages
-             );
-  if ((Status != EFI_UNSUPPORTED) && EFI_ERROR (Status)) {
-    return Status;
-  }
 
   TempHobBuffer.CacheInfoFlagBuffer = (UINTN)Buffer;
   DEBUG ((
@@ -1464,20 +1456,13 @@ BuildVariableRuntimeCacheInfoHob (
     ));
 
   //
-  // AllocateRuntimePages for VolatileCache and unblock it.
+  // AllocatePages for VolatileCache buffer (will be relocated to runtime by DXE).
   //
   BufferSize = PcdGet32 (PcdVariableStoreSize);
   if (BufferSize > 0) {
     Pages  = EFI_SIZE_TO_PAGES (BufferSize);
-    Buffer = AllocateRuntimePages (Pages);
+    Buffer = AllocatePages (Pages);
     ASSERT (Buffer != NULL);
-    Status = MmUnblockMemoryRequest (
-               (EFI_PHYSICAL_ADDRESS)(UINTN)Buffer,
-               Pages
-               );
-    if ((Status != EFI_UNSUPPORTED) && EFI_ERROR (Status)) {
-      return Status;
-    }
 
     TempHobBuffer.RuntimeVolatileCacheBuffer = (UINTN)Buffer;
     TempHobBuffer.RuntimeVolatileCachePages  = Pages;
@@ -1491,20 +1476,13 @@ BuildVariableRuntimeCacheInfoHob (
     ));
 
   //
-  // AllocateRuntimePages for NVCache and unblock it.
+  // AllocatePages for NVCache buffer (will be relocated to runtime by DXE).
   //
   BufferSize = CalculateNvVariableCacheSize (&NvAuthFlag);
   if (BufferSize > 0) {
     Pages  = EFI_SIZE_TO_PAGES (BufferSize);
-    Buffer = AllocateRuntimePages (Pages);
+    Buffer = AllocatePages (Pages);
     ASSERT (Buffer != NULL);
-    Status = MmUnblockMemoryRequest (
-               (EFI_PHYSICAL_ADDRESS)(UINTN)Buffer,
-               Pages
-               );
-    if ((Status != EFI_UNSUPPORTED) && EFI_ERROR (Status)) {
-      return Status;
-    }
 
     TempHobBuffer.RuntimeNvCacheBuffer = (UINTN)Buffer;
     TempHobBuffer.RuntimeNvCachePages  = Pages;
@@ -1518,20 +1496,13 @@ BuildVariableRuntimeCacheInfoHob (
     ));
 
   //
-  // AllocateRuntimePages for HobCache and unblock it.
+  // AllocatePages for HobCache buffer (will be relocated to runtime by DXE).
   //
   BufferSize = CalculateHobVariableCacheSize (NvAuthFlag);
   if (BufferSize > 0) {
     Pages  = EFI_SIZE_TO_PAGES (BufferSize);
-    Buffer = AllocateRuntimePages (Pages);
+    Buffer = AllocatePages (Pages);
     ASSERT (Buffer != NULL);
-    Status = MmUnblockMemoryRequest (
-               (EFI_PHYSICAL_ADDRESS)(UINTN)Buffer,
-               Pages
-               );
-    if ((Status != EFI_UNSUPPORTED) && EFI_ERROR (Status)) {
-      return Status;
-    }
 
     TempHobBuffer.RuntimeHobCacheBuffer = (UINTN)Buffer;
     TempHobBuffer.RuntimeHobCachePages  = Pages;

MdeModulePkg/Universal/Variable/Pei/Variable.h

@@ -22,7 +22,6 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Library/PeiServicesLib.h>
 #include <Library/SafeIntLib.h>
 #include <Library/VariableFlashInfoLib.h>
-#include <Library/MmUnblockMemoryLib.h>
 #include <Library/MemoryAllocationLib.h>
 
 #include <Guid/VariableFormat.h>

MdeModulePkg/Universal/Variable/Pei/VariablePei.inf

@@ -41,7 +41,6 @@
   PeiServicesLib
   SafeIntLib
   VariableFlashInfoLib
-  MmUnblockMemoryLib
   MemoryAllocationLib
 
 [Guids]

MdeModulePkg/Universal/Variable/RuntimeDxe/VariableSmmRuntimeDxe.c

@@ -37,6 +37,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Library/UefiLib.h>
 #include <Library/BaseLib.h>
 #include <Library/HobLib.h>
+#include <Library/MmUnblockMemoryLib.h>
 
 #include <Guid/EventGroup.h>
 #include <Guid/SmmVariableCommon.h>
@@ -1620,6 +1621,21 @@ InitVariableStoreHeader (
   }
 }
 
+/**
+  Invalidates all runtime cache pointers to prevent access to boot services memory.
+
+**/
+VOID
+InvalidateAllRuntimeCaches (
+  VOID
+  )
+{
+  mVariableRtCacheInfo.CacheInfoFlagBuffer        = 0;
+  mVariableRtCacheInfo.RuntimeHobCacheBuffer      = 0;
+  mVariableRtCacheInfo.RuntimeNvCacheBuffer       = 0;
+  mVariableRtCacheInfo.RuntimeVolatileCacheBuffer = 0;
+}
+
 /**
   Initialize the runtime variable cache related content.
 
@@ -1642,6 +1658,15 @@ InitVariableCache (
   UINTN                        AllocatedNvCacheSize;
   UINTN                        AllocatedVolatileCacheSize;
   VARIABLE_RUNTIME_CACHE_INFO  *VariableRuntimeCacheInfo;
+  VOID                         *RuntimeBuffer;
+  UINTN                        Pages;
+
+  // Storage for temporary allocations - only committed to mVariableRtCacheInfo if all succeed
+  EFI_PHYSICAL_ADDRESS  TempCacheInfoBuffer             = 0;
+  EFI_PHYSICAL_ADDRESS  TempHobCacheBuffer              = 0;
+  EFI_PHYSICAL_ADDRESS  TempNvCacheBuffer               = 0;
+  EFI_PHYSICAL_ADDRESS  TempVolatileCacheBuffer         = 0;
+  BOOLEAN               AllRtBufferAllocationsSucceeded = TRUE;
 
   //
   // Get needed runtime cache buffer size and check if auth variables are to be used from SMM
@@ -1665,9 +1690,157 @@ InitVariableCache (
       );
 
     CopyMem (&mVariableRtCacheInfo, VariableRuntimeCacheInfo, sizeof (VARIABLE_RUNTIME_CACHE_INFO));
-    InitVariableStoreHeader ((VOID *)(UINTN)mVariableRtCacheInfo.RuntimeHobCacheBuffer, AllocatedHobCacheSize);
-    InitVariableStoreHeader ((VOID *)(UINTN)mVariableRtCacheInfo.RuntimeNvCacheBuffer, AllocatedNvCacheSize);
-    InitVariableStoreHeader ((VOID *)(UINTN)mVariableRtCacheInfo.RuntimeVolatileCacheBuffer, AllocatedVolatileCacheSize);
+
+    //
+    // Relocate runtime cache buffers from boot services to runtime services
+    // This prevents runtime memory fragmentation by consolidating allocations in DXE
+    //
+    // Relocate CacheInfoFlag buffer from boot services to runtime services
+    //
+    if (mVariableRtCacheInfo.CacheInfoFlagBuffer != 0) {
+      Pages  = EFI_SIZE_TO_PAGES (sizeof (CACHE_INFO_FLAG));
+      Status = gBS->AllocatePages (
+                      AllocateAnyPages,
+                      EfiRuntimeServicesData,
+                      Pages,
+                      &TempCacheInfoBuffer
+                      );
+      if (EFI_ERROR (Status)) {
+        DEBUG ((DEBUG_WARN, "Failed to allocate runtime CacheInfoFlag buffer: %r\n", Status));
+        AllRtBufferAllocationsSucceeded = FALSE;
+      } else {
+        Status = MmUnblockMemoryRequest (TempCacheInfoBuffer, Pages);
+        if ((Status != EFI_UNSUPPORTED) && EFI_ERROR (Status)) {
+          DEBUG ((DEBUG_WARN, "Failed to unblock CacheInfoFlag buffer: %r\n", Status));
+          gBS->FreePages (TempCacheInfoBuffer, Pages);
+          TempCacheInfoBuffer             = 0;
+          AllRtBufferAllocationsSucceeded = FALSE;
+        }
+      }
+    }
+
+    if (AllRtBufferAllocationsSucceeded && (mVariableRtCacheInfo.RuntimeHobCacheBuffer != 0) && (AllocatedHobCacheSize > 0)) {
+      Pages  = EFI_SIZE_TO_PAGES (AllocatedHobCacheSize);
+      Status = gBS->AllocatePages (
+                      AllocateAnyPages,
+                      EfiRuntimeServicesData,
+                      Pages,
+                      &TempHobCacheBuffer
+                      );
+      if (EFI_ERROR (Status)) {
+        DEBUG ((DEBUG_WARN, "Failed to allocate runtime HOB cache buffer: %r\n", Status));
+        AllRtBufferAllocationsSucceeded = FALSE;
+      } else {
+        Status = MmUnblockMemoryRequest (TempHobCacheBuffer, Pages);
+        if ((Status != EFI_UNSUPPORTED) && EFI_ERROR (Status)) {
+          DEBUG ((DEBUG_WARN, "Failed to unblock HOB cache buffer: %r\n", Status));
+          gBS->FreePages (TempHobCacheBuffer, Pages);
+          TempHobCacheBuffer              = 0;
+          AllRtBufferAllocationsSucceeded = FALSE;
+        }
+      }
+    }
+
+    if (AllRtBufferAllocationsSucceeded && (mVariableRtCacheInfo.RuntimeNvCacheBuffer != 0) && (AllocatedNvCacheSize > 0)) {
+      Pages  = EFI_SIZE_TO_PAGES (AllocatedNvCacheSize);
+      Status = gBS->AllocatePages (
+                      AllocateAnyPages,
+                      EfiRuntimeServicesData,
+                      Pages,
+                      &TempNvCacheBuffer
+                      );
+      if (EFI_ERROR (Status)) {
+        DEBUG ((DEBUG_WARN, "Failed to allocate runtime NV cache buffer: %r\n", Status));
+        AllRtBufferAllocationsSucceeded = FALSE;
+      } else {
+        Status = MmUnblockMemoryRequest (TempNvCacheBuffer, Pages);
+        if ((Status != EFI_UNSUPPORTED) && EFI_ERROR (Status)) {
+          DEBUG ((DEBUG_WARN, "Failed to unblock NV cache buffer: %r\n", Status));
+          gBS->FreePages (TempNvCacheBuffer, Pages);
+          TempNvCacheBuffer               = 0;
+          AllRtBufferAllocationsSucceeded = FALSE;
+        }
+      }
+    }
+
+    if (AllRtBufferAllocationsSucceeded && (mVariableRtCacheInfo.RuntimeVolatileCacheBuffer != 0) && (AllocatedVolatileCacheSize > 0)) {
+      Pages  = EFI_SIZE_TO_PAGES (AllocatedVolatileCacheSize);
+      Status = gBS->AllocatePages (
+                      AllocateAnyPages,
+                      EfiRuntimeServicesData,
+                      Pages,
+                      &TempVolatileCacheBuffer
+                      );
+      if (EFI_ERROR (Status)) {
+        DEBUG ((DEBUG_WARN, "Failed to allocate runtime volatile cache buffer: %r\n", Status));
+        AllRtBufferAllocationsSucceeded = FALSE;
+      } else {
+        Status = MmUnblockMemoryRequest (TempVolatileCacheBuffer, Pages);
+        if ((Status != EFI_UNSUPPORTED) && EFI_ERROR (Status)) {
+          DEBUG ((DEBUG_WARN, "Failed to unblock volatile cache buffer: %r\n", Status));
+          gBS->FreePages (TempVolatileCacheBuffer, Pages);
+          TempVolatileCacheBuffer         = 0;
+          AllRtBufferAllocationsSucceeded = FALSE;
+        }
+      }
+    }
+
+    if (AllRtBufferAllocationsSucceeded) {
+      if (TempCacheInfoBuffer != 0) {
+        RuntimeBuffer = (VOID *)(UINTN)TempCacheInfoBuffer;
+        CopyMem (RuntimeBuffer, (VOID *)(UINTN)mVariableRtCacheInfo.CacheInfoFlagBuffer, sizeof (CACHE_INFO_FLAG));
+        mVariableRtCacheInfo.CacheInfoFlagBuffer = TempCacheInfoBuffer;
+      }
+
+      if (TempHobCacheBuffer != 0) {
+        RuntimeBuffer = (VOID *)(UINTN)TempHobCacheBuffer;
+        CopyMem (RuntimeBuffer, (VOID *)(UINTN)mVariableRtCacheInfo.RuntimeHobCacheBuffer, AllocatedHobCacheSize);
+        mVariableRtCacheInfo.RuntimeHobCacheBuffer = TempHobCacheBuffer;
+        InitVariableStoreHeader (RuntimeBuffer, AllocatedHobCacheSize);
+      }
+
+      if (TempNvCacheBuffer != 0) {
+        RuntimeBuffer = (VOID *)(UINTN)TempNvCacheBuffer;
+        CopyMem (RuntimeBuffer, (VOID *)(UINTN)mVariableRtCacheInfo.RuntimeNvCacheBuffer, AllocatedNvCacheSize);
+        mVariableRtCacheInfo.RuntimeNvCacheBuffer = TempNvCacheBuffer;
+        InitVariableStoreHeader (RuntimeBuffer, AllocatedNvCacheSize);
+      }
+
+      if (TempVolatileCacheBuffer != 0) {
+        RuntimeBuffer = (VOID *)(UINTN)TempVolatileCacheBuffer;
+        CopyMem (RuntimeBuffer, (VOID *)(UINTN)mVariableRtCacheInfo.RuntimeVolatileCacheBuffer, AllocatedVolatileCacheSize);
+        mVariableRtCacheInfo.RuntimeVolatileCacheBuffer = TempVolatileCacheBuffer;
+        InitVariableStoreHeader (RuntimeBuffer, AllocatedVolatileCacheSize);
+      }
+
+      DEBUG ((DEBUG_INFO, "Runtime variable cache buffers successfully relocated\n"));
+    } else {
+      //
+      // One or more allocations failed - clean up any successful allocations
+      // and disable runtime caching entirely
+      //
+      if (TempCacheInfoBuffer != 0) {
+        gBS->FreePages (TempCacheInfoBuffer, EFI_SIZE_TO_PAGES (sizeof (CACHE_INFO_FLAG)));
+      }
+
+      if (TempHobCacheBuffer != 0) {
+        gBS->FreePages (TempHobCacheBuffer, EFI_SIZE_TO_PAGES (AllocatedHobCacheSize));
+      }
+
+      if (TempNvCacheBuffer != 0) {
+        gBS->FreePages (TempNvCacheBuffer, EFI_SIZE_TO_PAGES (AllocatedNvCacheSize));
+      }
+
+      if (TempVolatileCacheBuffer != 0) {
+        gBS->FreePages (TempVolatileCacheBuffer, EFI_SIZE_TO_PAGES (AllocatedVolatileCacheSize));
+      }
+
+      InvalidateAllRuntimeCaches ();
+
+      DEBUG ((DEBUG_WARN, "Runtime variable cache disabled due to allocation failure.\n"));
+
+      // Don't return failure - the variable service can function without the runtime cache
+    }
   }
 
   return Status;

MdeModulePkg/Universal/Variable/RuntimeDxe/VariableSmmRuntimeDxe.inf

@@ -61,6 +61,7 @@
   SafeIntLib
   PcdLib
   HobLib
+  MmUnblockMemoryLib
 
 [Protocols]
   gEfiVariableWriteArchProtocolGuid             ## PRODUCES