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

4 files changed, 142 insertions, 0 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