efi/libstub: Simplify GOP handling code

Use the LocateHandleBuffer() API and a __free() function to simplify the
logic that allocates a handle buffer to iterate over all GOP protocols
in the EFI database.

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

4 files changed, 37 insertions, 74 deletions

diff --git a/drivers/firmware/efi/libstub/efi-stub.c b/drivers/firmware/efi/libstub/efi-stub.c
index 382b54f40603..90e06a6b1a45 100644
--- a/drivers/firmware/efi/libstub/efi-stub.c
+++ b/drivers/firmware/efi/libstub/efi-stub.c
@@ -10,6 +10,7 @@
  */
 
 #include <linux/efi.h>
+#include <linux/screen_info.h>
 #include <asm/efi.h>
 
 #include "efistub.h"
@@ -53,25 +54,16 @@ void __weak free_screen_info(struct screen_info *si)
 
 static struct screen_info *setup_graphics(void)
 {
-	efi_guid_t gop_proto = EFI_GRAPHICS_OUTPUT_PROTOCOL_GUID;
-	efi_status_t status;
-	unsigned long size;
-	void **gop_handle = NULL;
-	struct screen_info *si = NULL;
+	struct screen_info *si, tmp = {};
 
-	size = 0;
-	status = efi_bs_call(locate_handle, EFI_LOCATE_BY_PROTOCOL,
-			     &gop_proto, NULL, &size, gop_handle);
-	if (status == EFI_BUFFER_TOO_SMALL) {
-		si = alloc_screen_info();
-		if (!si)
-			return NULL;
-		status = efi_setup_gop(si, &gop_proto, size);
-		if (status != EFI_SUCCESS) {
-			free_screen_info(si);
-			return NULL;
-		}
-	}
+	if (efi_setup_gop(&tmp) != EFI_SUCCESS)
+		return NULL;
+
+	si = alloc_screen_info();
+	if (!si)
+		return NULL;
+
+	*si = tmp;
 	return si;
 }
 
diff --git a/drivers/firmware/efi/libstub/efistub.h b/drivers/firmware/efi/libstub/efistub.h
index 55553038b850..d96d4494070d 100644
--- a/drivers/firmware/efi/libstub/efistub.h
+++ b/drivers/firmware/efi/libstub/efistub.h
@@ -4,6 +4,7 @@
 #define _DRIVERS_FIRMWARE_EFI_EFISTUB_H
 
 #include <linux/compiler.h>
+#include <linux/cleanup.h>
 #include <linux/efi.h>
 #include <linux/kernel.h>
 #include <linux/kern_levels.h>
@@ -313,7 +314,9 @@ union efi_boot_services {
 		void *close_protocol;
 		void *open_protocol_information;
 		void *protocols_per_handle;
-		void *locate_handle_buffer;
+		efi_status_t (__efiapi *locate_handle_buffer)(int, efi_guid_t *,
+							      void *, unsigned long *,
+							      efi_handle_t **);
 		efi_status_t (__efiapi *locate_protocol)(efi_guid_t *, void *,
 							 void **);
 		efi_status_t (__efiapi *install_multiple_protocol_interfaces)(efi_handle_t *, ...);
@@ -1052,6 +1055,7 @@ void efi_puts(const char *str);
 __printf(1, 2) int efi_printk(char const *fmt, ...);
 
 void efi_free(unsigned long size, unsigned long addr);
+DEFINE_FREE(efi_pool, void *, if (_T) efi_bs_call(free_pool, _T));
 
 void efi_apply_loadoptions_quirk(const void **load_options, u32 *load_options_size);
 
@@ -1081,8 +1085,7 @@ efi_status_t efi_parse_options(char const *cmdline);
 
 void efi_parse_option_graphics(char *option);
 
-efi_status_t efi_setup_gop(struct screen_info *si, efi_guid_t *proto,
-			   unsigned long size);
+efi_status_t efi_setup_gop(struct screen_info *si);
 
 efi_status_t handle_cmdline_files(efi_loaded_image_t *image,
 				  const efi_char16_t *optstr,
diff --git a/drivers/firmware/efi/libstub/gop.c b/drivers/firmware/efi/libstub/gop.c
index 8eef63c48288..a4b3d18f91da 100644
--- a/drivers/firmware/efi/libstub/gop.c
+++ b/drivers/firmware/efi/libstub/gop.c
@@ -461,25 +461,25 @@ setup_pixel_info(struct screen_info *si, u32 pixels_per_scan_line,
 	}
 }
 
-static efi_graphics_output_protocol_t *
-find_gop(efi_guid_t *proto, unsigned long size, void **handles)
+static efi_graphics_output_protocol_t *find_gop(unsigned long num,
+						const efi_handle_t handles[])
 {
 	efi_graphics_output_protocol_t *first_gop;
 	efi_handle_t h;
 
 	first_gop = NULL;
 
-	for_each_efi_handle(h, handles, efi_get_handle_num(size)) {
+	for_each_efi_handle(h, handles, num) {
 		efi_status_t status;
 
 		efi_graphics_output_protocol_t *gop;
 		efi_graphics_output_protocol_mode_t *mode;
 		efi_graphics_output_mode_info_t *info;
-
-		efi_guid_t conout_proto = EFI_CONSOLE_OUT_DEVICE_GUID;
 		void *dummy = NULL;
 
-		status = efi_bs_call(handle_protocol, h, proto, (void **)&gop);
+		status = efi_bs_call(handle_protocol, h,
+				     &EFI_GRAPHICS_OUTPUT_PROTOCOL_GUID,
+				     (void **)&gop);
 		if (status != EFI_SUCCESS)
 			continue;
 
@@ -499,7 +499,8 @@ find_gop(efi_guid_t *proto, unsigned long size, void **handles)
 		 * Once we've found a GOP supporting ConOut,
 		 * don't bother looking any further.
 		 */
-		status = efi_bs_call(handle_protocol, h, &conout_proto, &dummy);
+		status = efi_bs_call(handle_protocol, h,
+				     &EFI_CONSOLE_OUT_DEVICE_GUID, &dummy);
 		if (status == EFI_SUCCESS)
 			return gop;
 
@@ -510,16 +511,22 @@ find_gop(efi_guid_t *proto, unsigned long size, void **handles)
 	return first_gop;
 }
 
-static efi_status_t setup_gop(struct screen_info *si, efi_guid_t *proto,
-			      unsigned long size, void **handles)
+efi_status_t efi_setup_gop(struct screen_info *si)
 {
-	efi_graphics_output_protocol_t *gop;
+	efi_handle_t *handles __free(efi_pool) = NULL;
 	efi_graphics_output_protocol_mode_t *mode;
 	efi_graphics_output_mode_info_t *info;
+	efi_graphics_output_protocol_t *gop;
+	efi_status_t status;
+	unsigned long num;
 
-	gop = find_gop(proto, size, handles);
+	status = efi_bs_call(locate_handle_buffer, EFI_LOCATE_BY_PROTOCOL,
+			      &EFI_GRAPHICS_OUTPUT_PROTOCOL_GUID, NULL, &num,
+			      &handles);
+	if (status != EFI_SUCCESS)
+		return status;
 
-	/* Did we find any GOPs? */
+	gop = find_gop(num, handles);
 	if (!gop)
 		return EFI_NOT_FOUND;
 
@@ -551,29 +558,3 @@ static efi_status_t setup_gop(struct screen_info *si, efi_guid_t *proto,
 
 	return EFI_SUCCESS;
 }
-
-/*
- * See if we have Graphics Output Protocol
- */
-efi_status_t efi_setup_gop(struct screen_info *si, efi_guid_t *proto,
-			   unsigned long size)
-{
-	efi_status_t status;
-	void **gop_handle = NULL;
-
-	status = efi_bs_call(allocate_pool, EFI_LOADER_DATA, size,
-			     (void **)&gop_handle);
-	if (status != EFI_SUCCESS)
-		return status;
-
-	status = efi_bs_call(locate_handle, EFI_LOCATE_BY_PROTOCOL, proto, NULL,
-			     &size, gop_handle);
-	if (status != EFI_SUCCESS)
-		goto free_handle;
-
-	status = setup_gop(si, proto, size, gop_handle);
-
-free_handle:
-	efi_bs_call(free_pool, gop_handle);
-	return status;
-}
diff --git a/drivers/firmware/efi/libstub/x86-stub.c b/drivers/firmware/efi/libstub/x86-stub.c
index 71173471faf6..53da6b5be739 100644
--- a/drivers/firmware/efi/libstub/x86-stub.c
+++ b/drivers/firmware/efi/libstub/x86-stub.c
@@ -406,24 +406,11 @@ static void setup_quirks(struct boot_params *boot_params)
 
 static void setup_graphics(struct boot_params *boot_params)
 {
-	efi_guid_t graphics_proto = EFI_GRAPHICS_OUTPUT_PROTOCOL_GUID;
-	struct screen_info *si;
-	efi_status_t status;
-	unsigned long size;
-	void **gop_handle = NULL;
-
-	si = &boot_params->screen_info;
-	memset(si, 0, sizeof(*si));
-
-	size = 0;
-	status = efi_bs_call(locate_handle, EFI_LOCATE_BY_PROTOCOL,
-			     &graphics_proto, NULL, &size, gop_handle);
-	if (status == EFI_BUFFER_TOO_SMALL)
-		status = efi_setup_gop(si, &graphics_proto, size);
+	struct screen_info *si = memset(&boot_params->screen_info, 0, sizeof(*si));
 
+	efi_setup_gop(si);
 }
 
-
 static void __noreturn efi_exit(efi_handle_t handle, efi_status_t status)
 {
 	efi_bs_call(exit, handle, status, 0, NULL);