ShellPkg/UefiHandleParsingLib: support handle list growth

Refactor GetHandleListByProtocol() to support
additions to the handle list during its execution.
Replace LocateHandle() with LocateHandleBuffer() to
avoid the possibility that the buffer allocated for
LocateHandle() is too small if additional handles
are added during GetHandleListByProtocol() execution.

Note that the previous implementation did not detect
the handle list growth and would cause memory
corruption when writing the terminating NULL handle
to the allocated buffer.

Signed-off-by: Bob Morgan <bobm@nvidia.com>

ShellPkg/Library/UefiHandleParsingLib/UefiHandleParsingLib.c

diff --git a/ShellPkg/Library/UefiHandleParsingLib/UefiHandleParsingLib.c b/ShellPkg/Library/UefiHandleParsingLib/UefiHandleParsingLib.c
index 023a4bc..0c65d36 100644
--- a/ShellPkg/Library/UefiHandleParsingLib/UefiHandleParsingLib.c
+++ b/ShellPkg/Library/UefiHandleParsingLib/UefiHandleParsingLib.c
@@ -3761,7 +3761,7 @@
   @param[in] ProtocolGuid The guid of the protocol to get handles for.  If NULL

                           then the function will return all handles.

 

-  @retval NULL            A memory allocation failed.

+  @retval NULL            Could not get handles or memory allocation failed.

   @return                 A NULL terminated list of handles.

 **/

 EFI_HANDLE *

@@ -3771,45 +3771,43 @@
   )

 {

   EFI_HANDLE  *HandleList;

-  UINTN       Size;

+  EFI_HANDLE  *OriginalHandleList;

+  UINTN       OriginalHandleCount;

   EFI_STATUS  Status;

 

-  Size       = 0;

-  HandleList = NULL;

-

-  //

-  // We cannot use LocateHandleBuffer since we need that NULL item on the ends of the list!

-  //

-  if (ProtocolGuid == NULL) {

-    Status = gBS->LocateHandle (AllHandles, NULL, NULL, &Size, HandleList);

-    if (Status == EFI_BUFFER_TOO_SMALL) {

-      HandleList = AllocateZeroPool (Size + sizeof (EFI_HANDLE));

-      if (HandleList == NULL) {

-        return (NULL);

-      }

-

-      Status                               = gBS->LocateHandle (AllHandles, NULL, NULL, &Size, HandleList);

-      HandleList[Size/sizeof (EFI_HANDLE)] = NULL;

+  OriginalHandleList = NULL;

+  Status             = gBS->LocateHandleBuffer (

+                              (ProtocolGuid == NULL) ? AllHandles : ByProtocol,

+                              (EFI_GUID *)ProtocolGuid,

+                              NULL,

+                              &OriginalHandleCount,

+                              &OriginalHandleList

+                              );

+  if (EFI_ERROR (Status)) {

+    if (Status != EFI_NOT_FOUND) {

+      DEBUG ((

+        DEBUG_ERROR,

+        "%a: LocateHandleBuffer %a failed: %r\n",

+        __func__,

+        (ProtocolGuid == NULL) ? "AllHandles" : "ByProtocol",

+        Status

+        ));

     }

-  } else {

-    Status = gBS->LocateHandle (ByProtocol, (EFI_GUID *)ProtocolGuid, NULL, &Size, HandleList);

-    if (Status == EFI_BUFFER_TOO_SMALL) {

-      HandleList = AllocateZeroPool (Size + sizeof (EFI_HANDLE));

-      if (HandleList == NULL) {

-        return (NULL);

-      }

 

-      Status                               = gBS->LocateHandle (ByProtocol, (EFI_GUID *)ProtocolGuid, NULL, &Size, HandleList);

-      HandleList[Size/sizeof (EFI_HANDLE)] = NULL;

-    }

+    return NULL;

   }

 

-  if (EFI_ERROR (Status)) {

-    if (HandleList != NULL) {

-      FreePool (HandleList);

+  // create new list with one more slot for the NULL terminator

+  HandleList = ReallocatePool (

+                 OriginalHandleCount * sizeof (EFI_HANDLE),

+                 (OriginalHandleCount + 1) * sizeof (EFI_HANDLE),

+                 OriginalHandleList

+                 );

+  if (HandleList == NULL) {

+    DEBUG ((DEBUG_ERROR, "%a: reallocate failed\n", __func__));

+    if (OriginalHandleList != NULL) {

+      FreePool (OriginalHandleList);

     }

-

-    return (NULL);

   }

 

   return (HandleList);