[PATCH] MdeModulePkg/Core/Dxe: Add lock protection in CoreLocateHandleBuffer()

["Ma, Hua" <hua.ma@intel.com>]

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

Currently, CoreLocateHandleBuffer() follows three steps:
1) get the size of protocol database firstly
2) allocate the buffer based on the size
3) get the protocol database into the buffer
There is no lock protection for the whole three steps. If a new protocol
added in step 2) by other task, e.g. (event timer handle USB device
hotplug). The size of protocol database may be increased and cannot fit
into the previous buffer in step 3). The protocol database cannot be
returned successfully, EFI_BUFFER_TOO_SMALL error will be returned.

This patch adds the lock to protect the whole three steps.
It can make sure the correct protocol database be returned.

Cc: Jian J Wang <jian.j.wang@intel.com>
Cc: Liming Gao <gaoliming@byosoft.com.cn>
Cc: Dandan Bi <dandan.bi@intel.com>
Signed-off-by: Hua Ma <hua.ma@intel.com>
---
 MdeModulePkg/Core/Dxe/Hand/Locate.c | 64 +++++++++++++++++++++++------
 1 file changed, 51 insertions(+), 13 deletions(-)

diff --git a/MdeModulePkg/Core/Dxe/Hand/Locate.c b/MdeModulePkg/Core/Dxe/Hand/Locate.c
index be17f4cbc3..4987c046c6 100644
--- a/MdeModulePkg/Core/Dxe/Hand/Locate.c
+++ b/MdeModulePkg/Core/Dxe/Hand/Locate.c
@@ -86,7 +86,8 @@ CoreGetNextLocateByProtocol (
 
 
 /**
-  Locates the requested handle(s) and returns them in Buffer.
+  Internal function for locating the requested handle(s) and returns them in Buffer.
+  The caller should already have acquired the ProtocolLock.
 
   @param  SearchType             The type of search to perform to locate the
                                  handles
@@ -104,8 +105,7 @@ CoreGetNextLocateByProtocol (
 
 **/
 EFI_STATUS
-EFIAPI
-CoreLocateHandle (
+InternalCoreLocateHandle (
   IN EFI_LOCATE_SEARCH_TYPE   SearchType,
   IN EFI_GUID                 *Protocol   OPTIONAL,
   IN VOID                     *SearchKey  OPTIONAL,
@@ -143,11 +143,6 @@ CoreLocateHandle (
   ResultBuffer = (IHANDLE **) Buffer;
   Status = EFI_SUCCESS;
 
-  //
-  // Lock the protocol database
-  //
-  CoreAcquireProtocolLock ();
-
   //
   // Get the search function based on type
   //
@@ -190,7 +185,6 @@ CoreLocateHandle (
   }
 
   if (EFI_ERROR(Status)) {
-    CoreReleaseProtocolLock ();
     return Status;
   }
 
@@ -247,10 +241,47 @@ CoreLocateHandle (
     }
   }
 
-  CoreReleaseProtocolLock ();
   return Status;
 }
 
+/**
+  Locates the requested handle(s) and returns them in Buffer.
+
+  @param  SearchType             The type of search to perform to locate the
+                                 handles
+  @param  Protocol               The protocol to search for
+  @param  SearchKey              Dependant on SearchType
+  @param  BufferSize             On input the size of Buffer.  On output the
+                                 size of data returned.
+  @param  Buffer                 The buffer to return the results in
+
+  @retval EFI_BUFFER_TOO_SMALL   Buffer too small, required buffer size is
+                                 returned in BufferSize.
+  @retval EFI_INVALID_PARAMETER  Invalid parameter
+  @retval EFI_SUCCESS            Successfully found the requested handle(s) and
+                                 returns them in Buffer.
+
+**/
+EFI_STATUS
+EFIAPI
+CoreLocateHandle (
+  IN EFI_LOCATE_SEARCH_TYPE   SearchType,
+  IN EFI_GUID                 *Protocol   OPTIONAL,
+  IN VOID                     *SearchKey  OPTIONAL,
+  IN OUT UINTN                *BufferSize,
+  OUT EFI_HANDLE              *Buffer
+  )
+{
+  EFI_STATUS Status;
+
+  //
+  // Lock the protocol database
+  //
+  CoreAcquireProtocolLock ();
+  Status = InternalCoreLocateHandle(SearchType, Protocol, SearchKey, BufferSize, Buffer);
+  CoreReleaseProtocolLock ();
+  return Status;
+}
 
 
 /**
@@ -610,7 +641,6 @@ Done:
   return Status;
 }
 
-
 /**
   Function returns an array of handles that support the requested protocol
   in a buffer allocated from pool. This is a version of CoreLocateHandle()
@@ -657,7 +687,12 @@ CoreLocateHandleBuffer (
   BufferSize = 0;
   *NumberHandles = 0;
   *Buffer = NULL;
-  Status = CoreLocateHandle (
+
+  //
+  // Lock the protocol database
+  //
+  CoreAcquireProtocolLock();
+  Status = InternalCoreLocateHandle (
              SearchType,
              Protocol,
              SearchKey,
@@ -674,15 +709,17 @@ CoreLocateHandleBuffer (
     if (Status != EFI_INVALID_PARAMETER) {
       Status = EFI_NOT_FOUND;
     }
+    CoreReleaseProtocolLock ();
     return Status;
   }
 
   *Buffer = AllocatePool (BufferSize);
   if (*Buffer == NULL) {
+    CoreReleaseProtocolLock ();
     return EFI_OUT_OF_RESOURCES;
   }
 
-  Status = CoreLocateHandle (
+  Status = InternalCoreLocateHandle (
              SearchType,
              Protocol,
              SearchKey,
@@ -695,6 +732,7 @@ CoreLocateHandleBuffer (
     *NumberHandles = 0;
   }
 
+  CoreReleaseProtocolLock ();
   return Status;
 }
 
-- 
2.32.0.windows.2