8646 loader: replace EFI part devices.
Reviewed by: Robert Mustacchi <rm@joyent.com>
Reviewed by: Igor Kozhukhov <igor@dilos.org>
Approved by: Dan McDonald <danmcd@joyent.com>
diff --git a/usr/src/boot/sys/boot/efi/loader/main.c b/usr/src/boot/sys/boot/efi/loader/main.c
index 3ef635f..7db2db1 100644
--- a/usr/src/boot/sys/boot/efi/loader/main.c
+++ b/usr/src/boot/sys/boot/efi/loader/main.c
@@ -35,6 +35,7 @@
 #include <inttypes.h>
 #include <string.h>
 #include <setjmp.h>
+#include <disk.h>
 
 #include <efi.h>
 #include <efilib.h>
@@ -65,6 +66,7 @@
 extern void efi_getsmap(void);
 #ifdef EFI_ZFS_BOOT
 static void efi_zfs_probe(void);
+static uint64_t pool_guid;
 #endif
 
 static int
@@ -149,12 +151,107 @@
 	return retval;
 }
 
-static int
-find_currdev(EFI_LOADED_IMAGE *img, struct devsw **dev, int *unit,
-    uint64_t *extra)
+static void
+set_devdesc_currdev(struct devsw *dev, int unit)
 {
+	struct devdesc currdev;
+	char *devname;
+
+	currdev.d_dev = dev;
+	currdev.d_type = currdev.d_dev->dv_type;
+	currdev.d_unit = unit;
+	currdev.d_opendata = NULL;
+	devname = efi_fmtdev(&currdev);
+
+	env_setenv("currdev", EV_VOLATILE, devname, efi_setcurrdev,
+	    env_nounset);
+	env_setenv("loaddev", EV_VOLATILE, devname, env_noset, env_nounset);
+}
+
+static int
+find_currdev(EFI_LOADED_IMAGE *img)
+{
+	pdinfo_list_t *pdi_list;
+	pdinfo_t *dp, *pp;
 	EFI_DEVICE_PATH *devpath, *copy;
 	EFI_HANDLE h;
+	char *devname;
+	struct devsw *dev;
+	int unit;
+	uint64_t extra;
+
+	/* Did efi_zfs_probe() detect the boot pool? */
+	if (pool_guid != 0) {
+		struct zfs_devdesc currdev;
+
+		currdev.d_dev = &zfs_dev;
+		currdev.d_unit = 0;
+		currdev.d_type = currdev.d_dev->dv_type;
+		currdev.d_opendata = NULL;
+		currdev.pool_guid = pool_guid;
+		currdev.root_guid = 0;
+		devname = efi_fmtdev(&currdev);
+
+		env_setenv("currdev", EV_VOLATILE, devname, efi_setcurrdev,
+		    env_nounset);
+		env_setenv("loaddev", EV_VOLATILE, devname, env_noset,
+		    env_nounset);
+		return (0);
+	}
+
+	/* We have device lists for hd, cd, fd, walk them all. */
+	pdi_list = efiblk_get_pdinfo_list(&efipart_hddev);
+	STAILQ_FOREACH(dp, pdi_list, pd_link) {
+		struct disk_devdesc currdev;
+
+		currdev.d_dev = &efipart_hddev;
+		currdev.d_type = currdev.d_dev->dv_type;
+		currdev.d_unit = dp->pd_unit;
+		currdev.d_opendata = NULL;
+		currdev.d_slice = -1;
+		currdev.d_partition = -1;
+
+		if (dp->pd_handle == img->DeviceHandle) {
+			devname = efi_fmtdev(&currdev);
+
+			env_setenv("currdev", EV_VOLATILE, devname,
+			    efi_setcurrdev, env_nounset);
+			env_setenv("loaddev", EV_VOLATILE, devname,
+			    env_noset, env_nounset);
+			return (0);
+		}
+		/* Assuming GPT partitioning. */
+		STAILQ_FOREACH(pp, &dp->pd_part, pd_link) {
+			if (pp->pd_handle == img->DeviceHandle) {
+				currdev.d_slice = pp->pd_unit;
+				currdev.d_partition = 255;
+				devname = efi_fmtdev(&currdev);
+
+				env_setenv("currdev", EV_VOLATILE, devname,
+				    efi_setcurrdev, env_nounset);
+				env_setenv("loaddev", EV_VOLATILE, devname,
+				    env_noset, env_nounset);
+				return (0);
+			}
+		}
+	}
+
+	pdi_list = efiblk_get_pdinfo_list(&efipart_cddev);
+	STAILQ_FOREACH(dp, pdi_list, pd_link) {
+		if (dp->pd_handle == img->DeviceHandle ||
+		    dp->pd_alias == img->DeviceHandle) {
+			set_devdesc_currdev(&efipart_cddev, dp->pd_unit);
+			return (0);
+		}
+	}
+
+	pdi_list = efiblk_get_pdinfo_list(&efipart_fddev);
+	STAILQ_FOREACH(dp, pdi_list, pd_link) {
+		if (dp->pd_handle == img->DeviceHandle) {
+			set_devdesc_currdev(&efipart_fddev, dp->pd_unit);
+			return (0);
+		}
+	}
 
 	/*
 	 * Try the device handle from our loaded image first.  If that
@@ -162,8 +259,10 @@
 	 * any of the nodes in that path match one of the enumerated
 	 * handles.
 	 */
-	if (efi_handle_lookup(img->DeviceHandle, dev, unit, extra) == 0)
+	if (efi_handle_lookup(img->DeviceHandle, &dev, &unit, &extra) == 0) {
+		set_devdesc_currdev(dev, unit);
 		return (0);
+	}
 
 	copy = NULL;
 	devpath = efi_lookup_image_devpath(IH);
@@ -175,8 +274,10 @@
 		free(copy);
 		copy = NULL;
 
-		if (efi_handle_lookup(h, dev, unit, extra) == 0)
+		if (efi_handle_lookup(h, &dev, &unit, &extra) == 0) {
+			set_devdesc_currdev(dev, unit);
 			return (0);
+		}
 
 		devpath = efi_lookup_devpath(h);
 		if (devpath != NULL) {
@@ -195,9 +296,7 @@
 	char var[128];
 	EFI_LOADED_IMAGE *img;
 	EFI_GUID *guid;
-	int i, j, vargood, unit, howto;
-	struct devsw *dev;
-	uint64_t pool_guid;
+	int i, j, vargood, howto;
 	void *ptr;
 	UINTN k;
 	int has_kbd;
@@ -365,45 +464,9 @@
 	 */
 	BS->SetWatchdogTimer(0, 0, 0, NULL);
 
-	if (find_currdev(img, &dev, &unit, &pool_guid) != 0)
+	if (find_currdev(img) != 0)
 		return (EFI_NOT_FOUND);
 
-	switch (dev->dv_type) {
-#ifdef EFI_ZFS_BOOT
-	case DEVT_ZFS: {
-		struct zfs_devdesc currdev;
-
-		currdev.d_dev = dev;
-		currdev.d_unit = unit;
-		currdev.d_type = currdev.d_dev->dv_type;
-		currdev.d_opendata = NULL;
-		currdev.pool_guid = pool_guid;
-		currdev.root_guid = 0;
-		env_setenv("currdev", EV_VOLATILE, efi_fmtdev(&currdev),
-			   efi_setcurrdev, env_nounset);
-		env_setenv("loaddev", EV_VOLATILE, efi_fmtdev(&currdev), env_noset,
-			   env_nounset);
-#ifdef __FreeBSD__
-		init_zfs_bootenv(zfs_fmtdev(&currdev));
-#endif
-		break;
-	}
-#endif
-	default: {
-		struct devdesc currdev;
-
-		currdev.d_dev = dev;
-		currdev.d_unit = unit;
-		currdev.d_opendata = NULL;
-		currdev.d_type = currdev.d_dev->dv_type;
-		env_setenv("currdev", EV_VOLATILE, efi_fmtdev(&currdev),
-			   efi_setcurrdev, env_nounset);
-		env_setenv("loaddev", EV_VOLATILE, efi_fmtdev(&currdev), env_noset,
-			   env_nounset);
-		break;
-	}
-	}
-
 	efi_init_environment();
 	setenv("ISADIR", "amd64", 1);	/* we only build 64bit */
 	acpi_detect();
@@ -842,18 +905,31 @@
 static void
 efi_zfs_probe(void)
 {
-	EFI_HANDLE h;
-	u_int unit;
-	int i;
-	char dname[SPECNAMELEN + 1];
-	uint64_t guid;
+	pdinfo_list_t *hdi;
+	pdinfo_t *hd, *pd = NULL;
+	EFI_GUID imgid = LOADED_IMAGE_PROTOCOL;
+	EFI_LOADED_IMAGE *img;
+	char devname[SPECNAMELEN + 1];
 
-	unit = 0;
-	h = efi_find_handle(&efipart_dev, 0);
-	for (i = 0; h != NULL; h = efi_find_handle(&efipart_dev, ++i)) {
-		snprintf(dname, sizeof(dname), "%s%d:", efipart_dev.dv_name, i);
-		if (zfs_probe_dev(dname, &guid) == 0)
-			(void)efi_handle_update_dev(h, &zfs_dev, unit++, guid);
+	BS->HandleProtocol(IH, &imgid, (VOID**)&img);
+	hdi = efiblk_get_pdinfo_list(&efipart_hddev);
+
+	/*
+	 * Find the handle for the boot device. The boot1 did find the
+	 * device with loader binary, now we need to search for the
+	 * same device and if it is part of the zfs pool, we record the
+	 * pool GUID for currdev setup.
+	 */
+	STAILQ_FOREACH(hd, hdi, pd_link) {
+		STAILQ_FOREACH(pd, &hd->pd_part, pd_link) {
+
+			snprintf(devname, sizeof(devname), "%s%dp%d:",
+			    efipart_hddev.dv_name, hd->pd_unit, pd->pd_unit);
+			if (pd->pd_handle == img->DeviceHandle)
+				(void) zfs_probe_dev(devname, &pool_guid);
+			else
+				(void) zfs_probe_dev(devname, NULL);
+		}
 	}
 }