From f500a8742c19be604d33907b56ab9597fe448b65 Mon Sep 17 00:00:00 2001
From: Gary Ching-Pang Lin <glin@suse.com>
Date: Tue, 27 May 2014 14:12:32 +0800
Subject: [PATCH 29/74] Fetch the netboot image from the same device

The previous strategy is to locate the first available PXE_BASE_CODE
protocol and to fetch the second stage image from it, and this may
cause shim to fetch the wrong second stage image, i.e. grub.efi.

Consider the machine with the following boot order:
1. PXE Boot
2. Hard Drive

Assume that the EFI image, e.g. bootx64.efi, in the PXE server is
broken, then "PXE Boot" will fail and fallback to "Hard Drive". While
shim.efi in "Hard Drive" is loaded, it will find the PXE protocol is
available and fetch grub.efi from the PXE server, not grub.efi in the
disk.

This commit checks the DeviceHandle from Loaded Image. If the device
supports PXE, then shim fetches grub.efi with the PXE protocol. Otherwise,
shim loads grub.efi from the disk.

Signed-off-by: Gary Ching-Pang Lin <glin@suse.com>
---
 netboot.c | 77 +++++++++++++--------------------------------------------------
 shim.c    |  2 +-
 2 files changed, 17 insertions(+), 62 deletions(-)

diff --git a/netboot.c b/netboot.c
index 07e2773..5ef53f7 100644
--- a/netboot.c
+++ b/netboot.c
@@ -85,78 +85,33 @@ translate_slashes(char *str)
  * Returns TRUE if we identify a protocol that is enabled and Providing us with
  * the needed information to fetch a grubx64.efi image
  */
-BOOLEAN findNetboot(EFI_HANDLE image_handle)
+BOOLEAN findNetboot(EFI_HANDLE device)
 {
-	UINTN bs = sizeof(EFI_HANDLE);
-	EFI_GUID pxe_base_code_protocol = EFI_PXE_BASE_CODE_PROTOCOL;
-	EFI_HANDLE *hbuf;
-	BOOLEAN rc = FALSE;
-	void *buffer = AllocatePool(bs);
-	UINTN errcnt = 0;
-	UINTN i;
 	EFI_STATUS status;
 
-	if (!buffer)
+	status = uefi_call_wrapper(BS->HandleProtocol, 3, device,
+				   &PxeBaseCodeProtocol, (VOID **)&pxe);
+	if (status != EFI_SUCCESS) {
+		pxe = NULL;
 		return FALSE;
-
-try_again:
-	status = uefi_call_wrapper(BS->LocateHandle,5, ByProtocol, 
-				   &pxe_base_code_protocol, NULL, &bs,
-				   buffer);
-
-	if (status == EFI_BUFFER_TOO_SMALL) {
-		errcnt++;
-		FreePool(buffer);
-		if (errcnt > 1)
-			return FALSE;
-		buffer = AllocatePool(bs);
-		if (!buffer)
-			return FALSE;
-		goto try_again;
 	}
 
-	if (status == EFI_NOT_FOUND) {
-		FreePool(buffer);
+	if (!pxe || !pxe->Mode) {
+		pxe = NULL;
 		return FALSE;
 	}
 
-	/*
- 	 * We have a list of pxe supporting protocols, lets see if any are
- 	 * active
- 	 */
-	hbuf = buffer;
-	pxe = NULL;
-	for (i=0; i < (bs / sizeof(EFI_HANDLE)); i++) {
-		status = uefi_call_wrapper(BS->OpenProtocol, 6, hbuf[i],
-					   &pxe_base_code_protocol,
-					   (void **)&pxe, image_handle, NULL,
-					   EFI_OPEN_PROTOCOL_GET_PROTOCOL);
-
-		if (status != EFI_SUCCESS) {
-			pxe = NULL;
-			continue;
-		}
-
-		if (!pxe || !pxe->Mode) {
-			pxe = NULL;
-			continue;
-		}
-
-		if (pxe->Mode->Started && pxe->Mode->DhcpAckReceived) {
-			/*
- 			 * We've located a pxe protocol handle thats been 
- 			 * started and has received an ACK, meaning its
- 			 * something we'll be able to get tftp server info
- 			 * out of
- 			 */
-			rc = TRUE;
-			break;
-		}
-			
+	if (!pxe->Mode->Started || !pxe->Mode->DhcpAckReceived) {
+		pxe = NULL;
+		return FALSE;
 	}
 
-	FreePool(buffer);
-	return rc;
+	/*
+	 * We've located a pxe protocol handle thats been started and has
+	 * received an ACK, meaning its something we'll be able to get
+	 * tftp server info out of
+	 */
+	return TRUE;
 }
 
 static CHAR8 *get_v6_bootfile_url(EFI_PXE_BASE_CODE_DHCPV6_PACKET *pkt)
diff --git a/shim.c b/shim.c
index 48a6f2f..d8699f9 100644
--- a/shim.c
+++ b/shim.c
@@ -1373,7 +1373,7 @@ EFI_STATUS start_image(EFI_HANDLE image_handle, CHAR16 *ImagePath)
 		goto done;
 	}
 
-	if (findNetboot(image_handle)) {
+	if (findNetboot(li->DeviceHandle)) {
 		efi_status = parseNetbootinfo(image_handle);
 		if (efi_status != EFI_SUCCESS) {
 			Print(L"Netboot parsing failed: %r\n", efi_status);
-- 
1.9.3