[CHERRY-PICK] MdeModulePkg: Improved error handling in RuntimeDxe Variable drivers [Rebase & FF]

MdeModulePkg/Universal/Variable/RuntimeDxe/PrivilegePolymorphic.h

@@ -100,6 +100,15 @@ VariableNotifySmmReady (
   VOID
   );
 
+/**
+  Revert the variable ready notification.
+  This function will be called when an error happens in variable initializing process.
+**/
+VOID
+VariableClearNotifySmmReady (
+  VOID
+  );
+
 /**
   Notify the system that the SMM variable write driver is ready.
 **/
@@ -114,7 +123,12 @@ VariableNotifySmmWriteReady (
   for variable read and write services being available. It also registers
   a notification function for an EVT_SIGNAL_VIRTUAL_ADDRESS_CHANGE event.
 
-  @retval EFI_SUCCESS       Variable service successfully initialized.
+  @retval EFI_SUCCESS           Variable service successfully initialized.
+  @retval EFI_OUT_OF_RESOURCES  Insufficient memory to allocate variable
+                                storage or communication buffers.
+  @retval EFI_VOLUME_CORRUPTED  The non-volatile variable store is corrupted.
+  @retval Others                An error from protocol installation, SMI handler
+                                registration, or protocol notification registration.
 **/
 EFI_STATUS
 EFIAPI

MdeModulePkg/Universal/Variable/RuntimeDxe/Variable.c

@@ -3843,6 +3843,9 @@ GetHobVariableStore (
 /**
   Initializes variable store area for non-volatile and volatile variable.
 
+  On any failure, the caller must invoke VariableCommonUninitialize() to
+  release resources that may have been partially allocated.
+
   @retval EFI_SUCCESS           Function successfully executed.
   @retval EFI_OUT_OF_RESOURCES  Fail to allocate enough memory resource.
 
@@ -3872,7 +3875,6 @@ VariableCommonInitialize (
   //
   Status = InitNonVolatileVariableStore ();
   if (EFI_ERROR (Status)) {
-    FreePool (mVariableModuleGlobal);
     return Status;
   }
 
@@ -3898,11 +3900,6 @@ VariableCommonInitialize (
   //
   Status = GetHobVariableStore (VariableGuid);
   if (EFI_ERROR (Status)) {
-    if (mNvFvHeaderCache != NULL) {
-      FreePool (mNvFvHeaderCache);
-    }
-
-    FreePool (mVariableModuleGlobal);
     return Status;
   }
 
@@ -3917,15 +3914,6 @@ VariableCommonInitialize (
   mVariableModuleGlobal->ScratchBufferSize = ScratchSize;
   VolatileVariableStore                    = AllocateRuntimePool (PcdGet32 (PcdVariableStoreSize) + ScratchSize);
   if (VolatileVariableStore == NULL) {
-    if (mVariableModuleGlobal->VariableGlobal.HobVariableBase != 0) {
-      FreePool ((VOID *)(UINTN)mVariableModuleGlobal->VariableGlobal.HobVariableBase);
-    }
-
-    if (mNvFvHeaderCache != NULL) {
-      FreePool (mNvFvHeaderCache);
-    }
-
-    FreePool (mVariableModuleGlobal);
     return EFI_OUT_OF_RESOURCES;
   }
 
@@ -3947,6 +3935,47 @@ VariableCommonInitialize (
   return EFI_SUCCESS;
 }
 
+/**
+  Uninitialize variable store area, freeing all resources allocated by
+  VariableCommonInitialize().
+
+**/
+VOID
+VariableCommonUninitialize (
+  VOID
+  )
+{
+  if (mVariableModuleGlobal != NULL) {
+    if (mVariableModuleGlobal->VariableGlobal.HobVariableBase != 0) {
+      FreePool ((VOID *)(UINTN)mVariableModuleGlobal->VariableGlobal.HobVariableBase);
+      mVariableModuleGlobal->VariableGlobal.HobVariableBase = 0;
+    }
+
+    if (mVariableModuleGlobal->VariableGlobal.VolatileVariableBase != 0) {
+      FreePool ((VOID *)(UINTN)mVariableModuleGlobal->VariableGlobal.VolatileVariableBase);
+    }
+
+    FreePool (mVariableModuleGlobal);
+    mVariableModuleGlobal = NULL;
+  }
+
+  if (mNvFvHeaderCache != NULL) {
+    //
+    // In real NV mode, mNvVariableCache points into the same allocation as mNvFvHeaderCache.
+    //
+    ASSERT ((UINTN)mNvVariableCache == (UINTN)mNvFvHeaderCache + mNvFvHeaderCache->HeaderLength);
+    FreePool (mNvFvHeaderCache);
+    mNvFvHeaderCache = NULL;
+    mNvVariableCache = NULL;
+  } else if ((mNvVariableCache != NULL) && (PcdGet64 (PcdEmuVariableNvStoreReserved) == 0)) {
+    //
+    // In emulated NV mode without a pre-reserved store, mNvVariableCache is a dynamic allocation.
+    //
+    FreePool (mNvVariableCache);
+    mNvVariableCache = NULL;
+  }
+}
+
 /**
   Get the proper fvb handle and/or fvb protocol by the given Flash address.
 

MdeModulePkg/Universal/Variable/RuntimeDxe/Variable.h

@@ -368,6 +368,16 @@ VariableCommonInitialize (
   VOID
   );
 
+/**
+  Uninitialize variable store area, freeing all resources allocated by
+  VariableCommonInitialize().
+
+**/
+VOID
+VariableCommonUninitialize (
+  VOID
+  );
+
 /**
   This function reclaims variable storage if free size is below the threshold.