1 files changed, 7 insertions, 27 deletions
diff --git a/drivers/firmware/efi/libstub/pci.c b/drivers/firmware/efi/libstub/pci.c
index 99fb25d2bcf5..1dccf77958d3 100644
--- a/drivers/firmware/efi/libstub/pci.c
+++ b/drivers/firmware/efi/libstub/pci.c
@@ -16,37 +16,20 @@
 void efi_pci_disable_bridge_busmaster(void)
 {
 	efi_guid_t pci_proto = EFI_PCI_IO_PROTOCOL_GUID;
-	unsigned long pci_handle_size = 0;
-	efi_handle_t *pci_handle = NULL;
+	efi_handle_t *pci_handle __free(efi_pool) = NULL;
+	unsigned long pci_handle_num;
 	efi_handle_t handle;
 	efi_status_t status;
 	u16 class, command;
-	int i;
 
-	status = efi_bs_call(locate_handle, EFI_LOCATE_BY_PROTOCOL, &pci_proto,
-			     NULL, &pci_handle_size, NULL);
-
-	if (status != EFI_BUFFER_TOO_SMALL) {
-		if (status != EFI_SUCCESS && status != EFI_NOT_FOUND)
-			efi_err("Failed to locate PCI I/O handles'\n");
-		return;
-	}
-
-	status = efi_bs_call(allocate_pool, EFI_LOADER_DATA, pci_handle_size,
-			     (void **)&pci_handle);
+	status = efi_bs_call(locate_handle_buffer, EFI_LOCATE_BY_PROTOCOL,
+			     &pci_proto, NULL, &pci_handle_num, &pci_handle);
 	if (status != EFI_SUCCESS) {
-		efi_err("Failed to allocate memory for 'pci_handle'\n");
+		efi_err("Failed to locate PCI I/O handles\n");
 		return;
 	}
 
-	status = efi_bs_call(locate_handle, EFI_LOCATE_BY_PROTOCOL, &pci_proto,
-			     NULL, &pci_handle_size, pci_handle);
-	if (status != EFI_SUCCESS) {
-		efi_err("Failed to locate PCI I/O handles'\n");
-		goto free_handle;
-	}
-
-	for_each_efi_handle(handle, pci_handle, pci_handle_size, i) {
+	for_each_efi_handle(handle, pci_handle, pci_handle_num) {
 		efi_pci_io_protocol_t *pci;
 		unsigned long segment_nr, bus_nr, device_nr, func_nr;
 
@@ -82,7 +65,7 @@ void efi_pci_disable_bridge_busmaster(void)
 		efi_bs_call(disconnect_controller, handle, NULL, NULL);
 	}
 
-	for_each_efi_handle(handle, pci_handle, pci_handle_size, i) {
+	for_each_efi_handle(handle, pci_handle, pci_handle_num) {
 		efi_pci_io_protocol_t *pci;
 
 		status = efi_bs_call(handle_protocol, handle, &pci_proto,
@@ -108,7 +91,4 @@ void efi_pci_disable_bridge_busmaster(void)
 		if (status != EFI_SUCCESS)
 			efi_err("Failed to disable PCI busmastering\n");
 	}
-
-free_handle:
-	efi_bs_call(free_pool, pci_handle);
 }