From 210d155761f4e83df3d30bc3ddb5b9998ea71e7c Mon Sep 17 00:00:00 2001
From: Qiu Shumin <shumin.qiu@intel.com>
Date: Fri, 25 Dec 2015 02:12:21 +0000
Subject: ShellPkg: Fix memory leak when running Shell script.

When we run following script in Shell:
"
for %a run (1 200)
  echo %a
  memmap
endfor
"
We may find memory leak in system. This patch free buffer in 'BufferToFreeList' to avoid this issue.

(Sync patch r19521 from main trunk.)

Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Qiu Shumin <shumin.qiu@intel.com>
Reviewed-by: Ruiyu Ni <ruiyu.ni@intel.com>

git-svn-id: https://svn.code.sf.net/p/edk2/code/branches/UDK2015@19533 6f19259b-4bc3-4df7-8a09-765794883524
---
 ShellPkg/Application/Shell/Shell.c | 38 ++++++++++++++++++++++++++++++++++++++
 ShellPkg/Application/Shell/Shell.h | 23 +++++++++++++++++++++++
 2 files changed, 61 insertions(+)

(limited to 'ShellPkg/Application')

diff --git a/ShellPkg/Application/Shell/Shell.c b/ShellPkg/Application/Shell/Shell.c
index 97b8c8d6fb..465f4f3974 100644
--- a/ShellPkg/Application/Shell/Shell.c
+++ b/ShellPkg/Application/Shell/Shell.c
@@ -1276,6 +1276,36 @@ AddBufferToFreeList(
   return (Buffer);
 }
 
+
+/**
+  Create a new buffer list and stores the old one to OldBufferList 
+
+  @param OldBufferList   The temporary list head used to store the nodes in BufferToFreeList.
+**/
+VOID
+SaveBufferList (
+  OUT LIST_ENTRY     *OldBufferList
+  )
+{
+  CopyMem (OldBufferList, &ShellInfoObject.BufferToFreeList.Link, sizeof (LIST_ENTRY));
+  InitializeListHead (&ShellInfoObject.BufferToFreeList.Link);
+}
+
+/**
+  Restore previous nodes into BufferToFreeList .
+
+  @param OldBufferList   The temporary list head used to store the nodes in BufferToFreeList.
+**/
+VOID
+RestoreBufferList (
+  IN OUT LIST_ENTRY     *OldBufferList
+  )
+{
+  FreeBufferList (&ShellInfoObject.BufferToFreeList);
+  CopyMem (&ShellInfoObject.BufferToFreeList.Link, OldBufferList, sizeof (LIST_ENTRY));
+}
+
+
 /**
   Add a buffer to the Line History List
 
@@ -2661,6 +2691,7 @@ RunScriptFileHandle (
   CONST CHAR16        *CurDir;
   UINTN               LineCount;
   CHAR16              LeString[50];
+  LIST_ENTRY          OldBufferList;
 
   ASSERT(!ShellCommandGetScriptExit());
 
@@ -2763,6 +2794,8 @@ RunScriptFileHandle (
               PrintBuffSize/sizeof(CHAR16) - 1
               );
 
+    SaveBufferList(&OldBufferList);
+
     //
     // NULL out comments
     //
@@ -2897,15 +2930,19 @@ RunScriptFileHandle (
 
           ShellCommandRegisterExit(FALSE, 0);
           Status = EFI_SUCCESS;
+          RestoreBufferList(&OldBufferList);
           break;
         }
         if (ShellGetExecutionBreakFlag()) {
+          RestoreBufferList(&OldBufferList);
           break;
         }
         if (EFI_ERROR(Status)) {
+          RestoreBufferList(&OldBufferList);
           break;
         }
         if (ShellCommandGetExit()) {
+          RestoreBufferList(&OldBufferList);
           break;
         }
       }
@@ -2924,6 +2961,7 @@ RunScriptFileHandle (
         NewScriptFile->CurrentCommand->Reset = TRUE;
       }
     }
+    RestoreBufferList(&OldBufferList);
   }
 
 
diff --git a/ShellPkg/Application/Shell/Shell.h b/ShellPkg/Application/Shell/Shell.h
index 57263204e3..351b94188e 100644
--- a/ShellPkg/Application/Shell/Shell.h
+++ b/ShellPkg/Application/Shell/Shell.h
@@ -382,5 +382,28 @@ TrimSpaces(
   IN CHAR16 **String
   );
 
+/**
+  
+  Create a new buffer list and stores the old one to OldBufferList  
+
+  @param OldBufferList   The temporary list head used to store the nodes in BufferToFreeList.
+**/
+VOID
+SaveBufferList (
+  OUT LIST_ENTRY     *OldBufferList
+  );
+
+/**
+  Restore previous nodes into BufferToFreeList .
+
+  @param OldBufferList   The temporary list head used to store the nodes in BufferToFreeList.
+**/
+VOID
+RestoreBufferList (
+  IN OUT LIST_ENTRY     *OldBufferList
+  );
+
+
+
 #endif //_SHELL_INTERNAL_HEADER_
 
-- 
cgit v1.2.3