EmulatorPkg/Win/Host: Fix image unload regression

https://bugzilla.tianocore.org/show_bug.cgi?id=2104

When UEFI Applications or UEFI Drivers are unloaded,
the PeCoffLoaderUnloadImageExtraAction() needs to unload
the image using FreeLibrary() if the image was successfully
loaded using LoadLibrrayEx().

This is a regression from the Nt32Pkg that supported
unloading applications and drivers as well as loading
the same application or driver multiple times.

Cc: Jordan Justen <[email protected]>
Reviewed-by: Ray Ni <[email protected]>
Cc: Andrew Fish <[email protected]>
Tested-by: Tim Lewis <[email protected]>
Signed-off-by: Michael D Kinney <[email protected]>

EmulatorPkg/Win/Host/WinHost.c

@@ -19,6 +19,25 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #define SE_TIME_ZONE_NAME                 TEXT("SeTimeZonePrivilege")
 #endif
 
+//
+// The growth size for array of module handle entries
+//
+#define MAX_PDB_NAME_TO_MOD_HANDLE_ARRAY_SIZE 0x100
+
+//
+// Module handle entry structure
+//
+typedef struct {
+  CHAR8   *PdbPointer;
+  VOID    *ModHandle;
+} PDB_NAME_TO_MOD_HANDLE;
+
+//
+// An Array to hold the module handles
+//
+PDB_NAME_TO_MOD_HANDLE  *mPdbNameModHandleArray = NULL;
+UINTN                   mPdbNameModHandleArraySize = 0;
+
 //
 // Default information about where the FD is located.
 //  This array gets filled in with information from PcdWinNtFirmwareVolume
@@ -840,13 +859,128 @@ Routine Description:
   return Count;
 }
 
+/**
+  Store the ModHandle in an array indexed by the Pdb File name.
+  The ModHandle is needed to unload the image.
+  @param ImageContext - Input data returned from PE Laoder Library. Used to find the
+                 .PDB file name of the PE Image.
+  @param ModHandle    - Returned from LoadLibraryEx() and stored for call to
+                 FreeLibrary().
+  @return   return EFI_SUCCESS when ModHandle was stored.
+--*/
+EFI_STATUS
+AddModHandle (
+  IN  PE_COFF_LOADER_IMAGE_CONTEXT         *ImageContext,
+  IN  VOID                                 *ModHandle
+  )
+
+{
+  UINTN                   Index;
+  PDB_NAME_TO_MOD_HANDLE  *Array;
+  UINTN                   PreviousSize;
+  PDB_NAME_TO_MOD_HANDLE  *TempArray;
+  HANDLE                  Handle;
+  UINTN                   Size;
+
+  //
+  // Return EFI_ALREADY_STARTED if this DLL has already been loaded
+  //
+  Array = mPdbNameModHandleArray;
+  for (Index = 0; Index < mPdbNameModHandleArraySize; Index++, Array++) {
+    if (Array->PdbPointer != NULL && Array->ModHandle == ModHandle) {
+      return EFI_ALREADY_STARTED;
+    }
+  }
+
+  Array = mPdbNameModHandleArray;
+  for (Index = 0; Index < mPdbNameModHandleArraySize; Index++, Array++) {
+    if (Array->PdbPointer == NULL) {
+      //
+      // Make a copy of the stirng and store the ModHandle
+      //
+      Handle = GetProcessHeap ();
+      Size = AsciiStrLen (ImageContext->PdbPointer) + 1;
+      Array->PdbPointer = HeapAlloc ( Handle, HEAP_ZERO_MEMORY, Size);
+      ASSERT (Array->PdbPointer != NULL);
+
+      AsciiStrCpyS (Array->PdbPointer, Size, ImageContext->PdbPointer);
+      Array->ModHandle = ModHandle;
+      return EFI_SUCCESS;
+    }
+  }
+
+  //
+  // No free space in mPdbNameModHandleArray so grow it by
+  // MAX_PDB_NAME_TO_MOD_HANDLE_ARRAY_SIZE entires.
+  //
+  PreviousSize = mPdbNameModHandleArraySize * sizeof (PDB_NAME_TO_MOD_HANDLE);
+  mPdbNameModHandleArraySize += MAX_PDB_NAME_TO_MOD_HANDLE_ARRAY_SIZE;
+  //
+  // re-allocate a new buffer and copy the old values to the new locaiton.
+  //
+  TempArray = HeapAlloc (GetProcessHeap (),
+                                HEAP_ZERO_MEMORY,
+                                mPdbNameModHandleArraySize * sizeof (PDB_NAME_TO_MOD_HANDLE)
+                               );
+
+  CopyMem ((VOID *) (UINTN) TempArray, (VOID *) (UINTN)mPdbNameModHandleArray, PreviousSize);
+
+  HeapFree (GetProcessHeap (), 0, mPdbNameModHandleArray);
+
+  mPdbNameModHandleArray = TempArray;
+  if (mPdbNameModHandleArray == NULL) {
+    ASSERT (FALSE);
+    return EFI_OUT_OF_RESOURCES;
+  }
+
+  return AddModHandle (ImageContext, ModHandle);
+}
+
+/**
+  Return the ModHandle and delete the entry in the array.
+   @param  ImageContext - Input data returned from PE Laoder Library. Used to find the
+                 .PDB file name of the PE Image.
+  @return
+    ModHandle - ModHandle assoicated with ImageContext is returned
+    NULL      - No ModHandle associated with ImageContext
+**/
+VOID *
+RemoveModHandle (
+  IN  PE_COFF_LOADER_IMAGE_CONTEXT         *ImageContext
+  )
+{
+  UINTN                   Index;
+  PDB_NAME_TO_MOD_HANDLE  *Array;
+
+  if (ImageContext->PdbPointer == NULL) {
+    //
+    // If no PDB pointer there is no ModHandle so return NULL
+    //
+    return NULL;
+  }
+
+  Array = mPdbNameModHandleArray;
+  for (Index = 0; Index < mPdbNameModHandleArraySize; Index++, Array++) {
+    if ((Array->PdbPointer != NULL) && (AsciiStrCmp(Array->PdbPointer, ImageContext->PdbPointer) == 0)) {
+      //
+      // If you find a match return it and delete the entry
+      //
+      HeapFree (GetProcessHeap (), 0, Array->PdbPointer);
+      Array->PdbPointer = NULL;
+      return Array->ModHandle;
+    }
+  }
+
+  return NULL;
+}
 
 VOID
 EFIAPI
 PeCoffLoaderRelocateImageExtraAction (
   IN OUT PE_COFF_LOADER_IMAGE_CONTEXT         *ImageContext
   )
 {
+  EFI_STATUS        Status;
   VOID              *DllEntryPoint;
   CHAR16            *DllFileName;
   HMODULE           Library;
@@ -855,7 +989,7 @@ PeCoffLoaderRelocateImageExtraAction (
   ASSERT (ImageContext != NULL);
   //
   // If we load our own PE COFF images the Windows debugger can not source
-  //  level debug our code. If a valid PDB pointer exists usw it to load
+  //  level debug our code. If a valid PDB pointer exists use it to load
   //  the *.dll file as a library using Windows* APIs. This allows
   //  source level debug. The image is still loaded and relocated
   //  in the Framework memory space like on a real system (by the code above),
@@ -913,10 +1047,22 @@ PeCoffLoaderRelocateImageExtraAction (
     }
 
     if ((Library != NULL) && (DllEntryPoint != NULL)) {
-      ImageContext->EntryPoint  = (EFI_PHYSICAL_ADDRESS) (UINTN) DllEntryPoint;
-      SecPrint ("LoadLibraryEx (%S,\n               NULL, DONT_RESOLVE_DLL_REFERENCES)\n", DllFileName);
+      Status = AddModHandle (ImageContext, Library);
+      if (Status == EFI_ALREADY_STARTED) {
+        //
+        // If the DLL has already been loaded before, then this instance of the DLL can not be debugged.
+        //
+        ImageContext->PdbPointer = NULL;
+        SecPrint ("WARNING: DLL already loaded.  No source level debug %S.\n\r", DllFileName);
+      } else {
+        //
+        // This DLL is not already loaded, so source level debugging is supported.
+        //
+        ImageContext->EntryPoint  = (EFI_PHYSICAL_ADDRESS) (UINTN) DllEntryPoint;
+        SecPrint ("LoadLibraryEx (\n\r  %S,\n\r  NULL, DONT_RESOLVE_DLL_REFERENCES)\n\r", DllFileName);
+      }
     } else {
-      SecPrint ("WARNING: No source level debug %S. \n", DllFileName);
+      SecPrint ("WARNING: No source level debug %S. \n\r", DllFileName);
     }
 
     free (DllFileName);
@@ -926,12 +1072,21 @@ PeCoffLoaderRelocateImageExtraAction (
 VOID
 EFIAPI
 PeCoffLoaderUnloadImageExtraAction (
-  IN PE_COFF_LOADER_IMAGE_CONTEXT         *ImageContext
+  IN PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
 )
 {
+  VOID  *ModHandle;
+
   ASSERT (ImageContext != NULL);
-}
 
+  ModHandle = RemoveModHandle (ImageContext);
+  if (ModHandle != NULL) {
+    FreeLibrary (ModHandle);
+    SecPrint ("FreeLibrary (\n\r  %s)\n\r", ImageContext->PdbPointer);
+  } else {
+    SecPrint ("WARNING: Unload image without source level debug\n\r");
+  }
+}
 
 VOID
 _ModuleEntryPoint (