efi/libstub: Simplify PCI I/O handle buffer traversal

Use LocateHandleBuffer() and a __free() cleanup helper to simplify the
PCI I/O handle buffer traversal code.

Signed-off-by: Ard Biesheuvel <ardb@kernel.org>

2 files changed, 13 insertions, 49 deletions

diff --git a/drivers/firmware/efi/libstub/pci.c b/drivers/firmware/efi/libstub/pci.c
index b0ba372c26c5..1dccf77958d3 100644
--- a/drivers/firmware/efi/libstub/pci.c
+++ b/drivers/firmware/efi/libstub/pci.c
@@ -16,36 +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;
 
-	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, efi_get_handle_num(pci_handle_size)) {
+	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;
 
@@ -81,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, efi_get_handle_num(pci_handle_size)) {
+	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,
@@ -107,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);
 }
diff --git a/drivers/firmware/efi/libstub/x86-stub.c b/drivers/firmware/efi/libstub/x86-stub.c
index 53da6b5be739..014ccf36a065 100644
--- a/drivers/firmware/efi/libstub/x86-stub.c
+++ b/drivers/firmware/efi/libstub/x86-stub.c
@@ -119,37 +119,23 @@ free_struct:
 static void setup_efi_pci(struct boot_params *params)
 {
 	efi_status_t status;
-	void **pci_handle = NULL;
+	efi_handle_t *pci_handle __free(efi_pool) = NULL;
 	efi_guid_t pci_proto = EFI_PCI_IO_PROTOCOL_GUID;
-	unsigned long size = 0;
 	struct setup_data *data;
+	unsigned long num;
 	efi_handle_t h;
 
-	status = efi_bs_call(locate_handle, EFI_LOCATE_BY_PROTOCOL,
-			     &pci_proto, NULL, &size, pci_handle);
-
-	if (status == EFI_BUFFER_TOO_SMALL) {
-		status = efi_bs_call(allocate_pool, EFI_LOADER_DATA, size,
-				     (void **)&pci_handle);
-
-		if (status != EFI_SUCCESS) {
-			efi_err("Failed to allocate memory for 'pci_handle'\n");
-			return;
-		}
-
-		status = efi_bs_call(locate_handle, EFI_LOCATE_BY_PROTOCOL,
-				     &pci_proto, NULL, &size, pci_handle);
-	}
-
+	status = efi_bs_call(locate_handle_buffer, EFI_LOCATE_BY_PROTOCOL,
+			     &pci_proto, NULL, &num, &pci_handle);
 	if (status != EFI_SUCCESS)
-		goto free_handle;
+		return;
 
 	data = (struct setup_data *)(unsigned long)params->hdr.setup_data;
 
 	while (data && data->next)
 		data = (struct setup_data *)(unsigned long)data->next;
 
-	for_each_efi_handle(h, pci_handle, efi_get_handle_num(size)) {
+	for_each_efi_handle(h, pci_handle, num) {
 		efi_pci_io_protocol_t *pci = NULL;
 		struct pci_setup_rom *rom;
 
@@ -169,9 +155,6 @@ static void setup_efi_pci(struct boot_params *params)
 
 		data = (struct setup_data *)rom;
 	}
-
-free_handle:
-	efi_bs_call(free_pool, pci_handle);
 }
 
 static void retrieve_apple_device_properties(struct boot_params *boot_params)