From f2333c707dd04f069c102bcae004561ec590a3a0 Mon Sep 17 00:00:00 2001
From: Suman Prakash <suman.p@samsung.com>
Date: Mon, 20 Mar 2017 16:34:55 +0800
Subject: [PATCH] MdeModulePkg/NvmExpressDxe: Memory leak fix in async code
 flow

For async commands, the buffer allocated for Prp list is
not getting freed, which will cause memory leak for async
read write command. For example testing async command flow
with custom application to send multiple read write commands
were resulting in decrease of available memory page in memmap,
which eventually resulted in system hang. Hence freeing
AsyncRequest->MapData, AsyncRequest->MapMeta, AsyncRequest->MapPrpList and
AsyncRequest->PrpListHost when async command is completed.

Cc: Feng Tian <feng.tian@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Suman Prakash <suman.p@samsung.com.com>
Reviewed-by: Hao Wu <hao.a.wu@intel.com>
---
 .../Bus/Pci/NvmExpressDxe/NvmExpress.c        | 22 ++++++++++++++++++-
 .../Bus/Pci/NvmExpressDxe/NvmExpress.h        |  5 +++++
 .../Pci/NvmExpressDxe/NvmExpressPassthru.c    |  5 +++++
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/MdeModulePkg/Bus/Pci/NvmExpressDxe/NvmExpress.c b/MdeModulePkg/Bus/Pci/NvmExpressDxe/NvmExpress.c
index 39f49bd774..de5c2a05ea 100644
--- a/MdeModulePkg/Bus/Pci/NvmExpressDxe/NvmExpress.c
+++ b/MdeModulePkg/Bus/Pci/NvmExpressDxe/NvmExpress.c
@@ -548,6 +548,7 @@ ProcessAsyncTaskList (
   QueueId    = 2;
   Cq         = Private->CqBuffer[QueueId] + Private->CqHdbl[QueueId].Cqh;
   HasNewItem = FALSE;
+  PciIo      = Private->PciIo;
 
   //
   // Submit asynchronous subtasks to the NVMe Submission Queue
@@ -644,6 +645,26 @@ ProcessAsyncTaskList (
           sizeof(EFI_NVM_EXPRESS_COMPLETION)
           );
 
+        //
+        // Free the resources allocated before cmd submission
+        //
+        if (AsyncRequest->MapData != NULL) {
+          PciIo->Unmap (PciIo, AsyncRequest->MapData);
+        }
+        if (AsyncRequest->MapMeta != NULL) {
+          PciIo->Unmap (PciIo, AsyncRequest->MapMeta);
+        }
+        if (AsyncRequest->MapPrpList != NULL) {
+          PciIo->Unmap (PciIo, AsyncRequest->MapPrpList);
+        }
+        if (AsyncRequest->PrpListHost != NULL) {
+          PciIo->FreeBuffer (
+                   PciIo,
+                   AsyncRequest->PrpListNo,
+                   AsyncRequest->PrpListHost
+                   );
+        }
+
         RemoveEntryList (Link);
         gBS->SignalEvent (AsyncRequest->CallerEvent);
         FreePool (AsyncRequest);
@@ -666,7 +687,6 @@ ProcessAsyncTaskList (
   }
 
   if (HasNewItem) {
-    PciIo = Private->PciIo;
     Data  = ReadUnaligned32 ((UINT32*)&Private->CqHdbl[QueueId]);
     PciIo->Mem.Write (
                  PciIo,
diff --git a/MdeModulePkg/Bus/Pci/NvmExpressDxe/NvmExpress.h b/MdeModulePkg/Bus/Pci/NvmExpressDxe/NvmExpress.h
index 6a1c2575a3..fa4a34ab53 100644
--- a/MdeModulePkg/Bus/Pci/NvmExpressDxe/NvmExpress.h
+++ b/MdeModulePkg/Bus/Pci/NvmExpressDxe/NvmExpress.h
@@ -292,6 +292,11 @@ typedef struct {
 
   EFI_NVM_EXPRESS_PASS_THRU_COMMAND_PACKET *Packet;
   UINT16                                   CommandId;
+  VOID                                     *MapPrpList;
+  UINTN                                    PrpListNo;
+  VOID                                     *PrpListHost;
+  VOID                                     *MapData;
+  VOID                                     *MapMeta;
   EFI_EVENT                                CallerEvent;
 } NVME_PASS_THRU_ASYNC_REQ;
 
diff --git a/MdeModulePkg/Bus/Pci/NvmExpressDxe/NvmExpressPassthru.c b/MdeModulePkg/Bus/Pci/NvmExpressDxe/NvmExpressPassthru.c
index 2c30009945..ef3d772cc2 100644
--- a/MdeModulePkg/Bus/Pci/NvmExpressDxe/NvmExpressPassthru.c
+++ b/MdeModulePkg/Bus/Pci/NvmExpressDxe/NvmExpressPassthru.c
@@ -627,6 +627,11 @@ NvmExpressPassThru (
     AsyncRequest->Packet        = Packet;
     AsyncRequest->CommandId     = Sq->Cid;
     AsyncRequest->CallerEvent   = Event;
+    AsyncRequest->MapData       = MapData;
+    AsyncRequest->MapMeta       = MapMeta;
+    AsyncRequest->MapPrpList    = MapPrpList;
+    AsyncRequest->PrpListNo     = PrpListNo;
+    AsyncRequest->PrpListHost   = PrpListHost;
 
     OldTpl = gBS->RaiseTPL (TPL_NOTIFY);
     InsertTailList (&Private->AsyncPassThruQueue, &AsyncRequest->Link);