2 files changed, 98 insertions, 23 deletions

diff --git a/sys/boot/i386/loader/main.c b/sys/boot/i386/loader/main.c
index 08cf16b09fe7..143241b766b7 100644
--- a/sys/boot/i386/loader/main.c
+++ b/sys/boot/i386/loader/main.c
@@ -356,25 +356,17 @@ static void
 i386_zfs_probe(void)
 {
     char devname[32];
-    int unit, slice;
+    int unit;
 
     /*
      * Open all the disks we can find and see if we can reconstruct
-     * ZFS pools from them. Bogusly assumes that the disks are named
-     * diskN, diskNpM or diskNsM.
+     * ZFS pools from them.
      */
     for (unit = 0; unit < MAXBDDEV; unit++) {
+	if (bd_unit2bios(unit) == -1)
+	    break;
 	sprintf(devname, "disk%d:", unit);
-	if (zfs_probe_dev(devname, NULL) == ENXIO)
-	    continue;
-	for (slice = 1; slice <= 128; slice++) {
-	    sprintf(devname, "disk%dp%d:", unit, slice);
-	    zfs_probe_dev(devname, NULL);
-	}
-	for (slice = 1; slice <= 4; slice++) {
-	    sprintf(devname, "disk%ds%d:", unit, slice);
-	    zfs_probe_dev(devname, NULL);
-	}
+	zfs_probe_dev(devname, NULL);
     }
 }
 #endif
diff --git a/sys/boot/zfs/zfs.c b/sys/boot/zfs/zfs.c
index cffb19ec8b50..d7ea878ca520 100644
--- a/sys/boot/zfs/zfs.c
+++ b/sys/boot/zfs/zfs.c
@@ -33,10 +33,11 @@ __FBSDID("$FreeBSD$");
  *	Stand-alone file reading package.
  */
 
+#include <sys/disk.h>
 #include <sys/param.h>
-#include <sys/disklabel.h>
 #include <sys/time.h>
 #include <sys/queue.h>
+#include <part.h>
 #include <stddef.h>
 #include <stdarg.h>
 #include <string.h>
@@ -376,21 +377,103 @@ zfs_dev_init(void)
 	return (0);
 }
 
-int
-zfs_probe_dev(const char *devname, uint64_t *pool_guid)
+struct zfs_probe_args {
+	int		fd;
+	const char	*devname;
+	uint64_t	*pool_guid;
+	uint16_t	secsz;
+};
+
+static int
+zfs_diskread(void *arg, void *buf, size_t blocks, off_t offset)
+{
+	struct zfs_probe_args *ppa;
+
+	ppa = (struct zfs_probe_args *)arg;
+	return (vdev_read(NULL, (void *)(uintptr_t)ppa->fd,
+	    offset * ppa->secsz, buf, blocks * ppa->secsz));
+}
+
+static int
+zfs_probe(int fd, uint64_t *pool_guid)
 {
 	spa_t *spa;
-	int fd;
 	int ret;
 
-	fd = open(devname, O_RDONLY);
-	if (fd == -1)
-		return (ENXIO);
 	ret = vdev_probe(vdev_read, (void *)(uintptr_t)fd, &spa);
-	if (ret != 0)
-		close(fd);
-	else if (pool_guid != NULL)
+	if (ret == 0 && pool_guid != NULL)
 		*pool_guid = spa->spa_guid;
+	return (ret);
+}
+
+static void
+zfs_probe_partition(void *arg, const char *partname,
+    const struct ptable_entry *part)
+{
+	struct zfs_probe_args *ppa, pa;
+	struct ptable *table;
+	char devname[32];
+	int ret;
+
+	/* Probe only freebsd-zfs and freebsd partitions */
+	if (part->type != PART_FREEBSD &&
+	    part->type != PART_FREEBSD_ZFS)
+		return;
+
+	ppa = (struct zfs_probe_args *)arg;
+	strncpy(devname, ppa->devname, strlen(ppa->devname) - 1);
+	sprintf(devname, "%s%s:", devname, partname);
+	pa.fd = open(devname, O_RDONLY);
+	if (pa.fd == -1)
+		return;
+	ret = zfs_probe(pa.fd, ppa->pool_guid);
+	if (ret == 0)
+		return;
+	/* Do we have BSD label here? */
+	if (part->type == PART_FREEBSD) {
+		pa.devname = devname;
+		pa.pool_guid = ppa->pool_guid;
+		pa.secsz = ppa->secsz;
+		table = ptable_open(&pa, part->end - part->start + 1,
+		    ppa->secsz, zfs_diskread);
+		if (table != NULL) {
+			ptable_iterate(table, &pa, zfs_probe_partition);
+			ptable_close(table);
+		}
+	}
+	close(pa.fd);
+}
+
+int
+zfs_probe_dev(const char *devname, uint64_t *pool_guid)
+{
+	struct ptable *table;
+	struct zfs_probe_args pa;
+	off_t mediasz;
+	int ret;
+
+	pa.fd = open(devname, O_RDONLY);
+	if (pa.fd == -1)
+		return (ENXIO);
+	/* Probe the whole disk */
+	ret = zfs_probe(pa.fd, pool_guid);
+	if (ret == 0)
+		return (0);
+	/* Probe each partition */
+	ret = ioctl(pa.fd, DIOCGMEDIASIZE, &mediasz);
+	if (ret == 0)
+		ret = ioctl(pa.fd, DIOCGSECTORSIZE, &pa.secsz);
+	if (ret == 0) {
+		pa.devname = devname;
+		pa.pool_guid = pool_guid;
+		table = ptable_open(&pa, mediasz / pa.secsz, pa.secsz,
+		    zfs_diskread);
+		if (table != NULL) {
+			ptable_iterate(table, &pa, zfs_probe_partition);
+			ptable_close(table);
+		}
+	}
+	close(pa.fd);
 	return (0);
 }