Introduce the ZFS Boot Environments menu to the loader menu

If the system was booted with ZFS, a new menu item (#7) appears
It contains an autogenerated list of ZFS Boot Environments

This allows the user to switch to an alternate root file system
Use Cases:
 - Revert a failed upgrade
 - Concurrently run different versions of FreeBSD with common home directory
 - Easier integration with the sysadmin/beadm utility

Requested by:	many
Reviewed by:	dteske
MFC after:	10 days
Relnotes:	yes
Sponsored by:	ScaleEngine Inc.
Differential Revision:	https://reviews.freebsd.org/D3167

8 files changed, 423 insertions, 7 deletions

diff --git a/sys/boot/forth/loader.conf b/sys/boot/forth/loader.conf
index e2338ab..c13e90e 100644
--- a/sys/boot/forth/loader.conf
+++ b/sys/boot/forth/loader.conf
@@ -90,6 +90,7 @@ ram_blacklist_type="ram_blacklist"	# Required for the kernel to find
 #password=""			# Prevent changes to boot options
 #bootlock_password=""		# Prevent booting (see check-password.4th(8))
 #geom_eli_passphrase_prompt="NO" # Prompt for geli(8) passphrase to mount root
+bootenv_autolist="YES"		# Auto populate the list of ZFS Boot Environments
 #beastie_disable="NO"		# Turn the beastie boot menu on and off
 #kernels="kernel kernel.old"	# Kernels to display in the boot menu
 #loader_logo="orbbw"		# Desired logo: orbbw, orb, fbsdbw, beastiebw, beastie, none
diff --git a/sys/boot/forth/menu-commands.4th b/sys/boot/forth/menu-commands.4th
index 0784676..9adf30a 100644
--- a/sys/boot/forth/menu-commands.4th
+++ b/sys/boot/forth/menu-commands.4th
@@ -351,4 +351,68 @@ also menu-namespace also menu-command-helpers
 	2 goto_menu
 ;
 
+\ 
+\ Set boot environment defaults
+\ 
+
+: init_bootenv ( -- )
+	s" set menu_caption[1]=${bemenu_current}${vfs.root.mountfrom}" evaluate
+	s" set ansi_caption[1]=${beansi_current}${vfs.root.mountfrom}" evaluate
+	s" set menu_caption[2]=${bemenu_bootfs}${zfs_be_active}" evaluate
+	s" set ansi_caption[2]=${beansi_bootfs}${zfs_be_active}" evaluate
+	s" set menu_caption[3]=${bemenu_page}${zfs_be_currpage}${bemenu_pageof}${zfs_be_pages}" evaluate
+	s" set ansi_caption[3]=${beansi_page}${zfs_be_currpage}${bemenu_pageof}${zfs_be_pages}" evaluate
+;
+
+\
+\ Redraw the entire screen. A long BE name can corrupt the menu
+\ 
+
+: be_draw_screen
+	clear		\ Clear the screen (in screen.4th)
+	print_version	\ print version string (bottom-right; see version.4th)
+	draw-beastie	\ Draw FreeBSD logo at right (in beastie.4th)
+	draw-brand	\ Draw brand.4th logo at top (in brand.4th)
+	menu-init	\ Initialize menu and draw bounding box (in menu.4th)
+;
+
+\
+\ Select a boot environment
+\ 
+
+: set_bootenv ( N -- N TRUE )
+	dup s" set vfs.root.mountfrom=${bootenv_root[E]}" 38 +c! evaluate
+	s" set currdev=${vfs.root.mountfrom}:" evaluate
+	s" unload" evaluate
+	free-module-options
+	s" /boot/defaults/loader.conf" read-conf
+	s" /boot/loader.conf" read-conf
+	s" /boot/loader.conf.local" read-conf
+	init_bootenv
+	be_draw_screen
+	menu-redraw
+	TRUE
+;
+
+\
+\ Switch to the next page of boot environments
+\
+
+: set_be_page ( N -- N TRUE )
+	s" zfs_be_currpage" getenv dup -1 = if
+		drop s" 1"
+	else
+		0 s>d 2swap
+		>number ( ud caddr/u -- ud' caddr'/u' )	\ convert string to numbers
+		2drop					\ drop the string
+		1 um/mod ( ud u1 -- u2 u3 ) 		\ convert double ud' to single u3' and remainder u2
+		swap drop ( ud2 u3 -- u3 )		\ drop the remainder u2
+		1+					\ increment the page number
+		s>d <# #s #>				\ convert back to a string
+	then
+	s" zfs_be_currpage" setenv
+	s" reloadbe" evaluate
+	3 goto_menu
+;
+
 only forth definitions
diff --git a/sys/boot/forth/menu.rc b/sys/boot/forth/menu.rc
index 7ffeef4..3c7de71 100644
--- a/sys/boot/forth/menu.rc
+++ b/sys/boot/forth/menu.rc
@@ -68,6 +68,13 @@ set mainmenu_command[6]="2 goto_menu"
 set mainmenu_keycode[6]=111
 set mainansi_caption[6]="Configure Boot ^[1mO^[mptions..."
 
+s" currdev" getenv dup 0> [if] drop 4 s" zfs:" compare 0= [if]
+    set mainmenu_caption[7]="Select Boot [E]nvironment..."
+    set mainmenu_command[7]="3 goto_menu"
+    set mainmenu_keycode[7]=101
+    set mainansi_caption[7]="Select Boot ^[1mE^[37mnvironment..."
+[then] [else] drop [then]
+
 \ 
 \ BOOT OPTIONS MENU
 \ 
@@ -119,6 +126,37 @@ set optionsmenu_keycode[6]=118
 set optionsansi_caption[6]="^[1mV^[merbose..... ^[34;1mOff^[m"
 set optionstoggled_ansi[6]="^[1mV^[merbose..... ^[32;7mOn^[m"
 
+\ 
+\ BOOT ENVIRONMENT MENU
+\ 
+
+set menuset_name3="bootenv"
+
+set bemenu_current="Active: "
+set beansi_current="^[1m${bemenu_current}^[m"
+set bemenu_bootfs="bootfs: "
+set beansi_bootfs="^[1m${bemenu_bootfs}^[m"
+set bemenu_page="[P]age: "
+set beansi_page="^[1mP^[mage: "
+set bemenu_pageof=" of "
+set beansi_pageof="${bemenu_pageof}"
+set zfs_be_currpage=1
+
+set bootenvmenu_init="init_bootenv"
+
+set bootenvmenu_command[1]="be_draw_screen 1 goto_menu"
+set bootenvmenu_keycode[1]=8
+
+set bootenvmenu_command[2]="set_bootenv"
+set bootenvmenu_keycode[2]=97
+set bootenv_root[2]="${zfs_be_active}"
+
+set bootenvmenu_command[3]="set_be_page"
+set bootenvmenu_keycode[3]=112
+
+set bootenvmenu_options=4
+set bootenvmenu_optionstext="Boot Environments:"
+
 \ Enable automatic booting (add ``autoboot_delay=N'' to loader.conf(5) to
 \ customize the timeout; default is 10-seconds)
 \ 
@@ -128,6 +166,21 @@ set menu_timeout_command="boot"
 \ 
 try-include /boot/menu.rc.local
 
+\ Initialize boot environment variables
+\
+s" reloadbe" sfind ( xt|0 bool ) [if]
+    s" bootenv_autolist" getenv dup -1 = [if]
+	drop s" execute" evaluate		\ Use evaluate to avoid passing
+						\ reloadbe an optional parameter
+    [else]
+	s" YES" compare-insensitive 0= [if]
+	    s" execute" evaluate
+	[then]
+    [then]
+[else]
+    drop ( xt=0 )
+[then]
+
 \ Display the main menu (see `menu.4th')
 set menuset_initial=1
 menuset-loadinitial
diff --git a/sys/boot/forth/support.4th b/sys/boot/forth/support.4th
index 5a005b4..3a8e1f2 100644
--- a/sys/boot/forth/support.4th
+++ b/sys/boot/forth/support.4th
@@ -930,6 +930,30 @@ only forth definitions also support-functions
   repeat
 ;
 
+: free-one-module { addr -- addr }
+  addr module.name strfree
+  addr module.loadname strfree
+  addr module.type strfree
+  addr module.args strfree
+  addr module.beforeload strfree
+  addr module.afterload strfree
+  addr module.loaderror strfree
+  addr
+;
+
+: free-module-options
+  module_options @
+  begin
+    ?dup
+  while
+    free-one-module
+    dup module.next @
+    swap free-memory
+  repeat
+  0 module_options !
+  0 last_module_option !
+;
+
 only forth also support-functions definitions
 
 \ Variables used for processing multiple conf files
diff --git a/sys/boot/i386/loader/main.c b/sys/boot/i386/loader/main.c
index 0087338..d4073e7 100644
--- a/sys/boot/i386/loader/main.c
+++ b/sys/boot/i386/loader/main.c
@@ -69,6 +69,7 @@ static int		isa_inb(int port);
 static void		isa_outb(int port, int value);
 void			exit(int code);
 #ifdef LOADER_ZFS_SUPPORT
+static void		init_zfs_bootenv(char *currdev);
 static void		i386_zfs_probe(void);
 #endif
 
@@ -294,12 +295,40 @@ extract_currdev(void)
 	new_currdev.d_unit = 0;
     }
 
+#ifdef LOADER_ZFS_SUPPORT
+    init_zfs_bootenv(zfs_fmtdev(&new_currdev));
+#endif
+
     env_setenv("currdev", EV_VOLATILE, i386_fmtdev(&new_currdev),
 	       i386_setcurrdev, env_nounset);
     env_setenv("loaddev", EV_VOLATILE, i386_fmtdev(&new_currdev), env_noset,
 	       env_nounset);
 }
 
+#ifdef LOADER_ZFS_SUPPORT
+static void
+init_zfs_bootenv(char *currdev)
+{
+	char *beroot;
+
+	/* Remove the trailing : */
+	currdev[strlen(currdev) - 1] = '\0';
+	setenv("zfs_be_active", currdev, 1);
+	/* Do not overwrite if already set */
+	setenv("vfs.root.mountfrom", currdev, 0);
+	/* Forward past zfs: */
+	currdev = strchr(currdev, ':');
+	currdev++;
+	/* Remove the last element (current bootenv) */
+	beroot = strrchr(currdev, '/');
+	beroot[0] = '\0';
+
+	beroot = currdev;
+	
+	setenv("zfs_be_root", beroot, 1);
+}
+#endif
+
 COMMAND_SET(reboot, "reboot", "reboot the system", command_reboot);
 
 static int
@@ -353,6 +382,34 @@ command_lszfs(int argc, char *argv[])
 	command_errmsg = strerror(err);
 	return (CMD_ERROR);
     }
+
+    return (CMD_OK);
+}
+
+COMMAND_SET(reloadbe, "reloadbe", "refresh the list of ZFS Boot Environments",
+    command_reloadbe);
+
+static int
+command_reloadbe(int argc, char *argv[])
+{
+    int err;
+
+    if (argc > 2) {
+	command_errmsg = "wrong number of arguments";
+	return (CMD_ERROR);
+    }
+
+    if (argc == 2) {
+	err = zfs_bootenv(argv[1]);
+    } else {
+	err = zfs_bootenv(getenv("zfs_be_root"));
+    }
+
+    if (err != 0) {
+	command_errmsg = strerror(err);
+	return (CMD_ERROR);
+    }
+
     return (CMD_OK);
 }
 #endif
diff --git a/sys/boot/zfs/libzfs.h b/sys/boot/zfs/libzfs.h
index 6834f8b..b289849 100644
--- a/sys/boot/zfs/libzfs.h
+++ b/sys/boot/zfs/libzfs.h
@@ -62,6 +62,9 @@ int	zfs_parsedev(struct zfs_devdesc *dev, const char *devspec,
 char	*zfs_fmtdev(void *vdev);
 int	zfs_probe_dev(const char *devname, uint64_t *pool_guid);
 int	zfs_list(const char *name);
+int	zfs_bootenv(const char *name);
+int	zfs_belist_add(const char *name);
+int	zfs_set_env(void);
 
 extern struct devsw zfs_dev;
 extern struct fs_ops zfs_fsops;
diff --git a/sys/boot/zfs/zfs.c b/sys/boot/zfs/zfs.c
index 64c738d..0e15ac4 100644
--- a/sys/boot/zfs/zfs.c
+++ b/sys/boot/zfs/zfs.c
@@ -48,6 +48,10 @@ __FBSDID("$FreeBSD$");
 
 #include "zfsimpl.c"
 
+/* Define the range of indexes to be populated with ZFS Boot Environments */
+#define		ZFS_BE_FIRST	4
+#define		ZFS_BE_LAST	8
+
 static int	zfs_open(const char *path, struct open_file *f);
 static int	zfs_write(struct open_file *f, void *buf, size_t size, size_t *resid);
 static int	zfs_close(struct open_file *f);
@@ -80,6 +84,16 @@ struct file {
 	zap_leaf_phys_t	*f_zap_leaf;	/* zap leaf buffer */
 };
 
+static int	zfs_env_index;
+static int	zfs_env_count;
+
+SLIST_HEAD(zfs_be_list, zfs_be_entry) zfs_be_head = SLIST_HEAD_INITIALIZER(zfs_be_head);
+struct zfs_be_list *zfs_be_headp;
+struct zfs_be_entry {
+	const char *name;
+	SLIST_ENTRY(zfs_be_entry) entries;
+} *zfs_be, *zfs_be_tmp;
+
 /*
  * Open a file.
  */
@@ -691,6 +705,156 @@ zfs_list(const char *name)
 	rv = zfs_lookup_dataset(spa, dsname, &objid);
 	if (rv != 0)
 		return (rv);
-	rv = zfs_list_dataset(spa, objid);
+
+	return (zfs_list_dataset(spa, objid));
+}
+
+int
+zfs_bootenv(const char *name)
+{
+	static char	poolname[ZFS_MAXNAMELEN], *dsname;
+	char		becount[4];
+	uint64_t	objid;
+	spa_t		*spa;
+	int		len, rv, pages, perpage, currpage;
+
+	if (strcmp(name, getenv("zfs_be_root")) != 0) {
+		if (setenv("zfs_be_root", name, 1) != 0)
+			return (ENOMEM);
+	}
+
+	SLIST_INIT(&zfs_be_head);
+	zfs_env_count = 0;
+	len = strlen(name);
+	dsname = strchr(name, '/');
+	if (dsname != NULL) {
+		len = dsname - name;
+		dsname++;
+	} else
+		dsname = "";
+	memcpy(poolname, name, len);
+	poolname[len] = '\0';
+
+	spa = spa_find_by_name(poolname);
+	if (!spa)
+		return (ENXIO);
+	rv = zfs_lookup_dataset(spa, dsname, &objid);
+	if (rv != 0)
+		return (rv);
+	rv = zfs_callback_dataset(spa, objid, zfs_belist_add);
+
+	/* Calculate and store the number of pages of BEs */
+	perpage = (ZFS_BE_LAST - ZFS_BE_FIRST + 1);
+	pages = (zfs_env_count / perpage) + ((zfs_env_count % perpage) > 0 ? 1 : 0);
+	snprintf(becount, 4, "%d", pages);
+	if (setenv("zfs_be_pages", becount, 1) != 0)
+		return (ENOMEM);
+
+	/* Roll over the page counter if it has exceeded the maximum */
+	currpage = strtol(getenv("zfs_be_currpage"), NULL, 10);
+	if (currpage > pages) {
+		if (setenv("zfs_be_currpage", "1", 1) != 0)
+			return (ENOMEM);
+	}
+
+	/* Populate the menu environment variables */
+	zfs_set_env();
+
+	/* Clean up the SLIST of ZFS BEs */
+	while (!SLIST_EMPTY(&zfs_be_head)) {
+		zfs_be = SLIST_FIRST(&zfs_be_head);
+		SLIST_REMOVE_HEAD(&zfs_be_head, entries);
+		free(zfs_be);
+	}
+
 	return (rv);
 }
+
+int
+zfs_belist_add(const char *name)
+{
+
+	/* Add the boot environment to the head of the SLIST */
+	zfs_be = malloc(sizeof(struct zfs_be_entry));
+	zfs_be->name = name;
+	SLIST_INSERT_HEAD(&zfs_be_head, zfs_be, entries);
+	zfs_env_count++;
+
+	return (0);
+}
+
+int
+zfs_set_env(void)
+{
+	char envname[32], envval[256];
+	char *beroot, *pagenum;
+	int rv, page, ctr;
+
+	beroot = getenv("zfs_be_root");
+	if (beroot == NULL) {
+		return (1);
+	}
+
+	pagenum = getenv("zfs_be_currpage");
+	if (pagenum != NULL) {
+		page = strtol(pagenum, NULL, 10);
+	} else {
+		page = 1;
+	}
+
+	ctr = 1;
+	rv = 0;
+	zfs_env_index = ZFS_BE_FIRST;
+	SLIST_FOREACH_SAFE(zfs_be, &zfs_be_head, entries, zfs_be_tmp) {
+		/* Skip to the requested page number */
+		if (ctr <= ((ZFS_BE_LAST - ZFS_BE_FIRST + 1) * (page - 1))) {
+			ctr++;
+			continue;
+		}
+		
+		snprintf(envname, sizeof(envname), "bootenvmenu_caption[%d]", zfs_env_index);
+		snprintf(envval, sizeof(envval), "%s", zfs_be->name);
+		rv = setenv(envname, envval, 1);
+		if (rv != 0) {
+			break;
+		}
+
+		snprintf(envname, sizeof(envname), "bootenvansi_caption[%d]", zfs_env_index);
+		rv = setenv(envname, envval, 1);
+		if (rv != 0){
+			break;
+		}
+
+		snprintf(envname, sizeof(envname), "bootenvmenu_command[%d]", zfs_env_index);
+		rv = setenv(envname, "set_bootenv", 1);
+		if (rv != 0){
+			break;
+		}
+
+		snprintf(envname, sizeof(envname), "bootenv_root[%d]", zfs_env_index);
+		snprintf(envval, sizeof(envval), "zfs:%s/%s", beroot, zfs_be->name);
+		rv = setenv(envname, envval, 1);
+		if (rv != 0){
+			break;
+		}
+
+		zfs_env_index++;
+		if (zfs_env_index > ZFS_BE_LAST) {
+			break;
+		}
+
+	}
+	
+	for (; zfs_env_index <= ZFS_BE_LAST; zfs_env_index++) {
+		snprintf(envname, sizeof(envname), "bootenvmenu_caption[%d]", zfs_env_index);
+		(void)unsetenv(envname);
+		snprintf(envname, sizeof(envname), "bootenvansi_caption[%d]", zfs_env_index);
+		(void)unsetenv(envname);
+		snprintf(envname, sizeof(envname), "bootenvmenu_command[%d]", zfs_env_index);
+		(void)unsetenv(envname);
+		snprintf(envname, sizeof(envname), "bootenv_root[%d]", zfs_env_index);
+		(void)unsetenv(envname);
+	}
+
+	return (rv);
+}
\ No newline at end of file
diff --git a/sys/boot/zfs/zfsimpl.c b/sys/boot/zfs/zfsimpl.c
index d889047..927fbad 100644
--- a/sys/boot/zfs/zfsimpl.c
+++ b/sys/boot/zfs/zfsimpl.c
@@ -1473,7 +1473,7 @@ zap_lookup(const spa_t *spa, const dnode_phys_t *dnode, const char *name, uint64
  * the directory contents.
  */
 static int
-mzap_list(const dnode_phys_t *dnode)
+mzap_list(const dnode_phys_t *dnode, int (*callback)(const char *))
 {
 	const mzap_phys_t *mz;
 	const mzap_ent_phys_t *mze;
@@ -1492,7 +1492,7 @@ mzap_list(const dnode_phys_t *dnode)
 		mze = &mz->mz_chunk[i];
 		if (mze->mze_name[0])
 			//printf("%-32s 0x%jx\n", mze->mze_name, (uintmax_t)mze->mze_value);
-			printf("%s\n", mze->mze_name);
+			callback(mze->mze_name);
 	}
 
 	return (0);
@@ -1503,7 +1503,7 @@ mzap_list(const dnode_phys_t *dnode)
  * the directory header.
  */
 static int
-fzap_list(const spa_t *spa, const dnode_phys_t *dnode)
+fzap_list(const spa_t *spa, const dnode_phys_t *dnode, int (*callback)(const char *))
 {
 	int bsize = dnode->dn_datablkszsec << SPA_MINBLOCKSHIFT;
 	zap_phys_t zh = *(zap_phys_t *) zap_scratch;
@@ -1566,13 +1566,21 @@ fzap_list(const spa_t *spa, const dnode_phys_t *dnode)
 			value = fzap_leaf_value(&zl, zc);
 
 			//printf("%s 0x%jx\n", name, (uintmax_t)value);
-			printf("%s\n", name);
+			callback((const char *)name);
 		}
 	}
 
 	return (0);
 }
 
+static int zfs_printf(const char *name)
+{
+
+	printf("%s\n", name);
+
+	return (0);
+}
+
 /*
  * List a zap directory.
  */
@@ -1587,9 +1595,9 @@ zap_list(const spa_t *spa, const dnode_phys_t *dnode)
 
 	zap_type = *(uint64_t *) zap_scratch;
 	if (zap_type == ZBT_MICRO)
-		return mzap_list(dnode);
+		return mzap_list(dnode, zfs_printf);
 	else
-		return fzap_list(spa, dnode);
+		return fzap_list(spa, dnode, zfs_printf);
 }
 
 static int
@@ -1858,6 +1866,48 @@ zfs_list_dataset(const spa_t *spa, uint64_t objnum/*, int pos, char *entry*/)
 
 	return (zap_list(spa, &child_dir_zap) != 0);
 }
+
+int
+zfs_callback_dataset(const spa_t *spa, uint64_t objnum, int (*callback)(const char *name))
+{
+	uint64_t dir_obj, child_dir_zapobj, zap_type;
+	dnode_phys_t child_dir_zap, dir, dataset;
+	dsl_dataset_phys_t *ds;
+	dsl_dir_phys_t *dd;
+	int err;
+
+	err = objset_get_dnode(spa, &spa->spa_mos, objnum, &dataset);
+	if (err != 0) {
+		printf("ZFS: can't find dataset %ju\n", (uintmax_t)objnum);
+		return (err);
+	}
+	ds = (dsl_dataset_phys_t *) &dataset.dn_bonus;
+	dir_obj = ds->ds_dir_obj;
+
+	err = objset_get_dnode(spa, &spa->spa_mos, dir_obj, &dir);
+	if (err != 0) {
+		printf("ZFS: can't find dirobj %ju\n", (uintmax_t)dir_obj);
+		return (err);
+	}
+	dd = (dsl_dir_phys_t *)&dir.dn_bonus;
+
+	child_dir_zapobj = dd->dd_child_dir_zapobj;
+	err = objset_get_dnode(spa, &spa->spa_mos, child_dir_zapobj, &child_dir_zap);
+	if (err != 0) {
+		printf("ZFS: can't find child zap %ju\n", (uintmax_t)dir_obj);
+		return (err);
+	}
+
+	err = dnode_read(spa, &child_dir_zap, 0, zap_scratch, child_dir_zap.dn_datablkszsec * 512);
+	if (err != 0)
+		return (err);
+
+	zap_type = *(uint64_t *) zap_scratch;
+	if (zap_type == ZBT_MICRO)
+		return mzap_list(&child_dir_zap, callback);
+	else
+		return fzap_list(spa, &child_dir_zap, callback);
+}
 #endif
 
 /*