1 files changed, 55 insertions, 1 deletions
diff --git a/MdePkg/Library/StandaloneMmDriverEntryPoint/StandaloneMmDriverEntryPoint.c b/MdePkg/Library/StandaloneMmDriverEntryPoint/StandaloneMmDriverEntryPoint.c
index 2c41e23a03..d74c9bdfed 100644
--- a/MdePkg/Library/StandaloneMmDriverEntryPoint/StandaloneMmDriverEntryPoint.c
+++ b/MdePkg/Library/StandaloneMmDriverEntryPoint/StandaloneMmDriverEntryPoint.c
@@ -1,7 +1,7 @@
 /** @file
   Entry point to a Standalone MM driver.
 
-Copyright (c) 2015, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2015 - 2021, Intel Corporation. All rights reserved.<BR>
 Copyright (c) 2016 - 2018, ARM Ltd. All rights reserved.<BR>
 Copyright (c) 2018, Linaro, Limited. All rights reserved.<BR>
 
@@ -11,12 +11,52 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 
 #include <PiMm.h>
 
+#include <Protocol/LoadedImage.h>
 #include <Library/BaseLib.h>
 #include <Library/DebugLib.h>
 #include <Library/MmServicesTableLib.h>
 #include <Library/StandaloneMmDriverEntryPoint.h>
 
 /**
+  Unloads an image from memory.
+
+  This function is a callback that a driver registers to do cleanup
+  when the UnloadImage boot service function is called.
+
+  @param  ImageHandle The handle to the image to unload.
+
+  @return Status returned by all unload().
+
+**/
+EFI_STATUS
+EFIAPI
+_DriverUnloadHandler (
+  EFI_HANDLE ImageHandle
+  )
+{
+  EFI_STATUS  Status;
+
+  //
+  // If an UnloadImage() handler is specified, then call it
+  //
+  Status = ProcessModuleUnloadList (ImageHandle);
+
+  //
+  // If the driver specific unload handler does not return an error, then call all of the
+  // library destructors.  If the unload handler returned an error, then the driver can not be
+  // unloaded, and the library destructors should not be called
+  //
+  if (!EFI_ERROR (Status)) {
+    ProcessLibraryDestructorList (ImageHandle, gMmst);
+  }
+
+  //
+  // Return the status from the driver specific unload handler
+  //
+  return Status;
+}
+
+/**
   The entry point of PE/COFF Image for a Standalone MM Driver.
 
   This function is the entry point for a Standalone MM Driver.
@@ -46,6 +86,7 @@ _ModuleEntryPoint (
   )
 {
   EFI_STATUS                 Status;
+  EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage;
 
   if (_gMmRevision != 0) {
     //
@@ -63,6 +104,19 @@ _ModuleEntryPoint (
   ProcessLibraryConstructorList (ImageHandle, MmSystemTable);
 
   //
+  //  Install unload handler...
+  //
+  if (_gDriverUnloadImageCount != 0) {
+    Status = gMmst->MmHandleProtocol (
+                      ImageHandle,
+                      &gEfiLoadedImageProtocolGuid,
+                      (VOID **)&LoadedImage
+                      );
+    ASSERT_EFI_ERROR (Status);
+    LoadedImage->Unload = _DriverUnloadHandler;
+  }
+
+  //
   // Call the driver entry point
   //
   Status = ProcessModuleEntryPointList (ImageHandle, MmSystemTable);