Fix NVMe library memory de-allocation issue (#667)

This patch added NVMe de-initialization function to stop the controller
and de-allocate all memory allocated.

Signed-off-by: Maurice Ma <maurice.ma@intel.com>

BootloaderCommonPkg/Library/NvmExpressLib/NvmExpress.c

@@ -228,6 +228,43 @@ DiscoverAllNamespaces (
   return EFI_SUCCESS;
 }
 
+/**
+  De-initialize the NVMe controller.
+
+  @param[in]  Private        NVMe controller private data pointer.
+
+  @retval EFI_SUCCESS        Always return EFI_SUCCESS.
+
+**/
+EFI_STATUS
+NvmeDeInitialize (
+  IN  NVME_CONTROLLER_PRIVATE_DATA        *Private
+  )
+{
+  UINTN   Index;
+
+  NvmeDisableController (Private);
+
+  for (Index = 0; Index < ARRAY_SIZE (mMultiNvmeDrive); Index++) {
+    if (mMultiNvmeDrive[Index] != NULL) {
+      FreePool (mMultiNvmeDrive[Index]);
+      mMultiNvmeDrive[Index] = NULL;
+    }
+  }
+
+  if (Private->Buffer != NULL) {
+    IoMmuFreeBuffer (6, Private->Buffer, Private->Mapping);
+  }
+
+  if (Private->ControllerData != NULL) {
+    FreePool (Private->ControllerData);
+  }
+
+  FreePool (Private);
+
+  return EFI_SUCCESS;
+}
+
 /**
   Starts a device controller or a bus controller.
 
@@ -263,14 +300,10 @@ NvmeInitialize (
   EFI_PHYSICAL_ADDRESS                MappedAddr;
 
   if (NvmeInitMode == DevDeinit) {
-    Private = mNvmeCtrlPrivate;
-    if ((Private != NULL) && (Private->ControllerData != NULL)) {
-      FreePool (Private->ControllerData);
+    if (mNvmeCtrlPrivate != NULL) {
+      NvmeDeInitialize (mNvmeCtrlPrivate);
+      mNvmeCtrlPrivate = NULL;
     }
-    if (Private != NULL) {
-      FreePool (Private);
-    }
-    mNvmeCtrlPrivate = NULL;
     return EFI_SUCCESS;
   }
 
@@ -333,7 +366,7 @@ NvmeInitialize (
   return EFI_SUCCESS;
 
 Exit:
-  if ((Private != NULL) && (Private->Mapping != NULL)) {
+  if ((Private != NULL) && (Private->Buffer != NULL)) {
     IoMmuFreeBuffer (6, Private->Buffer, Private->Mapping);
   }
 

BootloaderCommonPkg/Library/NvmExpressLib/NvmExpressHci.h

@@ -97,5 +97,19 @@ NvmeIdentifyNamespace (
   IN VOID                              *Buffer
   );
 
+/**
+  Disable the Nvm Express controller.
+
+  @param  Private          The pointer to the NVME_CONTROLLER_PRIVATE_DATA data structure.
+
+  @return EFI_SUCCESS      Successfully disable the controller.
+  @return EFI_DEVICE_ERROR Fail to disable the controller.
+
+**/
+EFI_STATUS
+NvmeDisableController (
+  IN NVME_CONTROLLER_PRIVATE_DATA     *Private
+  );
+
 #endif
 

BootloaderCommonPkg/Library/NvmExpressLib/NvmExpressPassthru.c

@@ -679,7 +679,7 @@ EXIT:
     IoMmuUnmap (MapMeta);
   }
 
-  if (MapPrpList != NULL) {
+  if (PrpListHost != NULL) {
     IoMmuFreeBuffer (PrpListNo, PrpListHost, MapPrpList);
   }