[PATCH 3/3] ShellPkg: Update CWD and current mapping when commands return

[Ruiyu Ni <ruiyu.ni@intel.com>]

From: Huajing Li <huajing.li@intel.com>

Some commands may cause current directory or even current mapping
become invalid. For example, "MAP" after unplugging the USB thumb
key may cause current mapping "FS0:" disappear.
This patch updates the CWD and current mapping when commands return.
It also causes the command prompt change to default "Shell>".

Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Huajing Li <huajing.li@intel.com>
Reviewed-by: Ruiyu Ni <ruiyu.ni@intel.com>
---
 ShellPkg/Application/Shell/Shell.c | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/ShellPkg/Application/Shell/Shell.c b/ShellPkg/Application/Shell/Shell.c
index fcf47b5c73..5471930ba1 100644
--- a/ShellPkg/Application/Shell/Shell.c
+++ b/ShellPkg/Application/Shell/Shell.c
@@ -2614,6 +2614,7 @@ RunShellCommand(
   CHAR16                    *FirstParameter;
   CHAR16                    *TempWalker;
   SHELL_OPERATION_TYPES     Type;
+  CONST CHAR16              *CurDir;
 
   ASSERT(CmdLine != NULL);
   if (StrLen(CmdLine) == 0) {
@@ -2705,7 +2706,22 @@ RunShellCommand(
     ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_SHELL_NOT_FOUND), ShellInfoObject.HiiHandle, FirstParameter);
     SetLastError(SHELL_NOT_FOUND);
   }
- 
+  //
+  // Check whether the current file system still exists. If not exist, we need update "cwd" and gShellCurMapping.
+  //
+  CurDir = EfiShellGetCurDir (NULL);
+  if (CurDir != NULL) {
+    if (EFI_ERROR(ShellFileExists (CurDir))) {
+      //
+      // EfiShellSetCurDir() cannot set current directory to NULL.
+      // EfiShellSetEnv() is not allowed to set the "cwd" variable.
+      // Only InternalEfiShellSetEnv () is allowed setting the "cwd" variable.
+      //
+      InternalEfiShellSetEnv (L"cwd", NULL, TRUE);
+      gShellCurMapping = NULL;
+    }
+  }
+
   SHELL_FREE_NON_NULL(CleanOriginal);
   SHELL_FREE_NON_NULL(FirstParameter);
 
-- 
2.12.2.windows.2