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 <[email protected]>

Author: bobm21

ShellPkg/Library/UefiHandleParsingLib/UefiHandleParsingLib.c

@@ -3761,7 +3761,7 @@ ParseHandleDatabaseForChildDevices (
   @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 @@ GetHandleListByProtocol (
   )
 {
   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);