Cleanup, clarify and generally polish the MBR/GPT initialization

code.

'-g' is promoted to be independant of '-i'. This makes it clearer
that there are four mutually exclusive initialization options.

'-i' puts the default MBR on disk, '-g' puts the default GPT on
disk, '-u' updates the MBR boot code on disk and '-A' puts a new
set of GPT partitions on disk without overwriting 'protected'
partitions.

The last initialization option specified is the one executed, so
existing '-i -g' finger memory, etc. continue to work as before.

man page/usage feedback/tweaks from jmc@

7 files changed, 132 insertions, 118 deletions

diff --git a/sbin/fdisk/cmd.c b/sbin/fdisk/cmd.c
index ebd9270d394..3cb58bbffdc 100644
--- a/sbin/fdisk/cmd.c
+++ b/sbin/fdisk/cmd.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: cmd.c,v 1.135 2021/07/21 20:26:30 krw Exp $	*/
+/*	$OpenBSD: cmd.c,v 1.136 2021/08/06 10:41:31 krw Exp $	*/
 
 /*
  * Copyright (c) 1997 Tobias Weingartner
@@ -51,7 +51,6 @@ extern const int		manpage_sz;
 int
 Xreinit(char *args, struct mbr *mbr)
 {
-	struct dos_mbr		dos_mbr;
 	int			dogpt;
 
 	dogpt = 0;
@@ -65,9 +64,6 @@ Xreinit(char *args, struct mbr *mbr)
 		return CMD_CONT;
 	}
 
-	MBR_make(&initial_mbr, &dos_mbr);
-	MBR_parse(&dos_mbr, mbr->mbr_lba_self, mbr->mbr_lba_firstembr, mbr);
-
 	if (dogpt) {
 		GPT_init(GHANDGP);
 		GPT_print("s", TERSE);
@@ -511,7 +507,7 @@ Xhelp(char *args, struct mbr *mbr)
 int
 Xupdate(char *args, struct mbr *mbr)
 {
-	memcpy(mbr->mbr_code, initial_mbr.mbr_code, sizeof(mbr->mbr_code));
+	memcpy(mbr->mbr_code, default_dmbr.dmbr_boot, sizeof(mbr->mbr_code));
 	mbr->mbr_signature = DOSMBR_SIGNATURE;
 	printf("Machine code updated.\n");
 	return CMD_DIRTY;
diff --git a/sbin/fdisk/fdisk.8 b/sbin/fdisk/fdisk.8
index c83ddb670c9..08e4b5fab21 100644
--- a/sbin/fdisk/fdisk.8
+++ b/sbin/fdisk/fdisk.8
@@ -1,4 +1,4 @@
-.\"	$OpenBSD: fdisk.8,v 1.104 2021/06/23 13:07:13 krw Exp $
+.\"	$OpenBSD: fdisk.8,v 1.105 2021/08/06 10:41:31 krw Exp $
 .\"
 .\"
 .\" Copyright (c) 1997 Tobias Weingartner
@@ -15,7 +15,7 @@
 .\" ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 .\" OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 .\"
-.Dd $Mdocdate: June 23 2021 $
+.Dd $Mdocdate: August 6 2021 $
 .Dt FDISK 8
 .Os
 .Sh NAME
@@ -23,8 +23,8 @@
 .Nd partition table maintenance program
 .Sh SYNOPSIS
 .Nm fdisk
-.Op Fl ey
-.Op Fl i Oo Fl g Oc | Fl u | Fl A | Fl v
+.Op Fl evy
+.Op Fl A  | g | i | u
 .Op Fl b Ar blocks Ns Op @ Ns Ar offset Ns Op : Ns Ar type
 .Op Fl l Ar blocks | Fl c Ar cylinders Fl h Ar heads Fl s Ar sectors
 .Op Fl f Ar mbrfile
@@ -54,7 +54,7 @@ at install time.
 The options are as follows:
 .Bl -tag -width Ds
 .It Fl A
-Allocates the largest chunks of free space in an existing GPT to the
+Modifies existing GPT partitions to allocate all possible space to an
 .Ox
 partition
 and the optional
@@ -82,12 +82,18 @@ The
 partition will follow the boot partition and use the remaining
 available space.
 .Pp
-Can only be used when initializing a disk with
-.Fl i Op Fl g
+Can only be used when initializing
+.Ar disk
+with
+.Fl A ,
+.Fl g ,
+or
+.Fl i .
+If
+.Fl A
 or
-.Fl A .
-If a GPT is being initialized
-only the
+.Fl g
+is specified, only the
 .Ar blocks
 value is used.
 .It Xo
@@ -118,40 +124,33 @@ Specifies an alternate MBR template file.
 The default file is
 .Pa /usr/mdec/mbr .
 .It Fl g
-A protective MBR and a default GPT will be written to disk.
-The default GPT will have a single
+A default GPT, including a protective MBR, will be written to
+.Ar disk .
+The default GPT has a single
 .Ox
-partition containing all available space not allocated by a
-.Fl b
-specification.
+partition containing all available space not allocated by
+.Fl b .
 .It Fl i
-Requests that the partition table data be re-initialized.
-In this mode,
-.Nm
-will completely overwrite the primary MBR bootcode and MBR partition table.
-.Pp
-If
-.Fl g
-is not specified then the MBR will be overwritten by
-the default MBR template
+A default MBR will be written to
+.Ar disk .
+The default MBR will be read from
+the file specified by
+.Fl f ,
+the default file
 .Pa /usr/mdec/mbr
-(or the one optionally specified by the
-.Fl f
-flag).
-In the default template, MBR partition number 3 will be configured as an
+or the built-in default.
+Normally the default MBR has a single
 .Ox
-MBR partition containing all available space not allocated by a
-.Fl b
-specification.
+partition containing all available space not allocated by
+.Fl b .
 .It Fl l Ar blocks
-Treat the disk as if it has the specified number of blocks.
+Treat
+.Ar disk
+as if it has the specified number of blocks.
 .It Fl u
 Update MBR bootcode, preserving existing MBR partition table.
 The MBR bootcode extends from offset 0x000 to the start of the MBR partition
 table at offset 0x1BE.
-It is similar to the
-.Fl i
-flag, except the existing MBR partition table is preserved.
 This is useful for writing new MBR bootcode onto an existing drive, and is
 equivalent to the DOS command
 .Dq FDISK /MBR .
@@ -159,10 +158,6 @@ Note that this option will overwrite the NT disk signature, if present.
 .It Fl v
 Print the contents of the MBR, the Primary GPT and the Secondary GPT.
 Also print more detailed GPT header and partition entry information.
-.Fl i ,
-.Fl u ,
-or
-.Fl e .
 .It Fl y
 Avoid asking yes/no questions when not desirable.
 .It Ar disk
@@ -184,10 +179,12 @@ or
 .Sh TYPICAL LAYOUT
 When called with no special flags,
 .Nm
-prints the partition table of the specified disk.
+prints the partition table of the
+.Ar disk .
 .Pp
-If the disk does not contain a valid GPT the contents
-of the MBR are displayed.
+If
+.Ar disk
+does not contain a valid GPT the contents of the MBR are displayed.
 For example:
 .Bd -literal -offset 1n
 # fdisk sd0
@@ -239,7 +236,9 @@ for the system to be able to boot and use the drive correctly.
 These values must be kept correctly synchronized or a variety of
 problems develop which are very difficult to diagnose.
 .Pp
-If the disk contains a protective MBR and a valid GPT, the contents of the GPT
+If
+.Ar disk
+contains a protective MBR and a valid GPT, the contents of the GPT
 are displayed.
 For example:
 .Bd -literal -offset 1n
@@ -255,7 +254,7 @@ partitions that span the whole disk.
 The first partition is a 960 sector EFI Sys partition;
 the second is a 7764929 sector
 .Ox
-partition using the remainder of the disk.
+partition using the remaing space.
 The fields of the output are:
 .Bl -tag -width "type"
 .It Em "#"
diff --git a/sbin/fdisk/fdisk.c b/sbin/fdisk/fdisk.c
index e4af09328f3..1ba3cfbdeeb 100644
--- a/sbin/fdisk/fdisk.c
+++ b/sbin/fdisk/fdisk.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: fdisk.c,v 1.129 2021/07/22 18:54:17 krw Exp $	*/
+/*	$OpenBSD: fdisk.c,v 1.130 2021/08/06 10:41:31 krw Exp $	*/
 
 /*
  * Copyright (c) 1997 Tobias Weingartner
@@ -37,15 +37,20 @@
 #include "user.h"
 #include "gpt.h"
 
+#define	INIT_GPT		1
+#define	INIT_GPTPARTITIONS	2
+#define	INIT_MBR		3
+#define	INIT_MBRBOOTCODE	4
+
 #define	_PATH_MBR		_PATH_BOOTDIR "mbr"
 static unsigned char		builtin_mbr[] = {
 #include "mbrcode.h"
 };
 
-int			A_flag, y_flag;
+int			y_flag;
 
 void			parse_bootprt(const char *);
-void			get_default_mbr(const char *, struct mbr *);
+void			get_default_dmbr(const char *, struct dos_mbr *);
 
 static void
 usage(void)
@@ -53,7 +58,7 @@ usage(void)
 	extern char		* __progname;
 
 	fprintf(stderr, "usage: %s "
-	    "[-evy] [-i [-g] | -u | -A ] [-b blocks[@offset[:type]]]\n"
+	    "[-evy] [-A | -g | -i | -u] [-b blocks[@offset[:type]]]\n"
 	    "\t[-l blocks | -c cylinders -h heads -s sectors] [-f mbrfile] disk\n",
 	    __progname);
 	exit(1);
@@ -69,30 +74,28 @@ main(int argc, char *argv[])
 	const char		*mbrfile = NULL;
 #endif
 	int			 ch, error;
-	int			 e_flag = 0, g_flag = 0, i_flag = 0, u_flag = 0;
+	int			 e_flag = 0, init = 0;
 	int			 verbosity = TERSE;
 	int			 oflags = O_RDONLY;
-	char			*query;
 
 	while ((ch = getopt(argc, argv, "Aiegpuvf:c:h:s:l:b:y")) != -1) {
 		const char *errstr;
 
 		switch(ch) {
 		case 'A':
-			A_flag = 1;
-			oflags = O_RDWR;
+			init = INIT_GPTPARTITIONS;
 			break;
 		case 'i':
-			i_flag = 1;
-			oflags = O_RDWR;
+			init = INIT_MBR;
+			break;
+		case 'g':
+			init = INIT_GPT;
 			break;
 		case 'u':
-			u_flag = 1;
-			oflags = O_RDWR;
+			init = INIT_MBRBOOTCODE;
 			break;
 		case 'e':
 			e_flag = 1;
-			oflags = O_RDWR;
 			break;
 		case 'f':
 			mbrfile = optarg;
@@ -116,9 +119,6 @@ main(int argc, char *argv[])
 				errx(1, "Sector argument %s [1..63].", errstr);
 			disk.dk_size = 0;
 			break;
-		case 'g':
-			g_flag = 1;
-			break;
 		case 'b':
 			parse_bootprt(optarg);
 			break;
@@ -142,14 +142,16 @@ main(int argc, char *argv[])
 	argc -= optind;
 	argv += optind;
 
-	if (argc != 1 || (i_flag && u_flag) ||
-	    (i_flag == 0 && g_flag))
+	if (argc != 1)
 		usage();
 
 	if ((disk.dk_cylinders || disk.dk_heads || disk.dk_sectors) &&
 	    (disk.dk_cylinders * disk.dk_heads * disk.dk_sectors == 0))
 		usage();
 
+	if (init || e_flag)
+		oflags = O_RDWR;
+
 	DISK_open(argv[0], oflags);
 	if (oflags == O_RDONLY) {
 		if (pledge("stdio", NULL) == -1)
@@ -162,35 +164,39 @@ main(int argc, char *argv[])
 	if (pledge("stdio rpath wpath disklabel proc exec", NULL) == -1)
 		err(1, "pledge");
 
-	get_default_mbr(mbrfile, &initial_mbr);
+	get_default_dmbr(mbrfile, &default_dmbr);
 
-	query = NULL;
-	if (A_flag) {
+	switch (init) {
+	case INIT_GPT:
+		GPT_init(GHANDGP);
+		if (ask_yn("Do you wish to write new GPT?"))
+			GPT_write();
+		break;
+	case INIT_GPTPARTITIONS:
 		if (GPT_read(ANYGPT))
 			errx(1, "-A requires a valid GPT");
-		else {
-			GPT_init(GPONLY);
-			query = "Do you wish to write new GPT?";
-		}
-	} else if (i_flag) {
-		if (g_flag) {
-			GPT_init(GHANDGP);
-			query = "Do you wish to write new GPT?";
-		} else {
-			MBR_init(&initial_mbr);
-			query = "Do you wish to write new MBR and "
-			    "partition table?";
-		}
-	} else if (u_flag) {
+		GPT_init(GPONLY);
+		if (ask_yn("Do you wish to write new GPT?"))
+			GPT_write();
+		break;
+	case INIT_MBR:
+		mbr.mbr_lba_self = mbr.mbr_lba_firstembr = 0;
+		MBR_init(&mbr);
+		if (ask_yn("Do you wish to write new MBR?"))
+			MBR_write(&mbr);
+		break;
+	case INIT_MBRBOOTCODE:
 		error = MBR_read(0, 0, &mbr);
 		if (error)
 			errx(1, "Can't read MBR!");
-		memcpy(initial_mbr.mbr_prt, mbr.mbr_prt,
-		    sizeof(initial_mbr.mbr_prt));
-		query = "Do you wish to write new MBR?";
+		memcpy(mbr.mbr_code, default_dmbr.dmbr_boot,
+		    sizeof(mbr.mbr_code));
+		if (ask_yn("Do you wish to write new MBR?"))
+			MBR_write(&mbr);
+		break;
+	default:
+		break;
 	}
-	if (query && ask_yn(query))
-		Xwrite(NULL, &initial_mbr);
 
 	if (e_flag)
 		USER_edit(0, 0);
@@ -252,30 +258,27 @@ parse_bootprt(const char *arg)
 }
 
 void
-get_default_mbr(const char *mbrfile, struct mbr *mbr)
+get_default_dmbr(const char *mbrfile, struct dos_mbr *dmbr)
 {
-	struct dos_mbr		dos_mbr;
 	ssize_t			len;
 	int			fd;
 
 	if (mbrfile == NULL) {
-		memcpy(&dos_mbr, builtin_mbr, sizeof(dos_mbr));
+		memcpy(dmbr, builtin_mbr, sizeof(*dmbr));
 	} else {
 		fd = open(mbrfile, O_RDONLY);
 		if (fd == -1) {
 			warn("%s", mbrfile);
 			warnx("using builtin MBR");
-			memcpy(&dos_mbr, builtin_mbr, sizeof(dos_mbr));
+			memcpy(dmbr, builtin_mbr, sizeof(*dmbr));
 		} else {
-			len = read(fd, &dos_mbr, sizeof(dos_mbr));
+			len = read(fd, dmbr, sizeof(*dmbr));
 			close(fd);
 			if (len == -1)
 				err(1, "Unable to read MBR from '%s'", mbrfile);
-			else if (len != sizeof(dos_mbr))
+			else if (len != sizeof(*dmbr))
 				errx(1, "Unable to read complete MBR from '%s'",
 				    mbrfile);
 		}
 	}
-
-	MBR_parse(&dos_mbr, 0, 0, mbr);
 }
diff --git a/sbin/fdisk/gpt.c b/sbin/fdisk/gpt.c
index b964e647e78..b4f3722e52d 100644
--- a/sbin/fdisk/gpt.c
+++ b/sbin/fdisk/gpt.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: gpt.c,v 1.48 2021/07/26 13:05:14 krw Exp $	*/
+/*	$OpenBSD: gpt.c,v 1.49 2021/08/06 10:41:31 krw Exp $	*/
 /*
  * Copyright (c) 2015 Markus Muller <mmu@grummel.net>
  * Copyright (c) 2015 Kenneth R Westerback <krw@openbsd.org>
@@ -447,7 +447,7 @@ init_gh(void)
 	memset(&gmbr, 0, sizeof(gmbr));
 
 	/* XXX Do we need the boot code? UEFI spec & Apple says no. */
-	memcpy(gmbr.mbr_code, initial_mbr.mbr_code, sizeof(gmbr.mbr_code));
+	memcpy(gmbr.mbr_code, default_dmbr.dmbr_boot, sizeof(gmbr.mbr_code));
 	gmbr.mbr_prt[0].prt_id = DOSPTYP_EFI;
 	gmbr.mbr_prt[0].prt_bs = 1;
 	gmbr.mbr_prt[0].prt_ns = UINT32_MAX;
diff --git a/sbin/fdisk/mbr.c b/sbin/fdisk/mbr.c
index eeabdb3dba6..abd93c040dc 100644
--- a/sbin/fdisk/mbr.c
+++ b/sbin/fdisk/mbr.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: mbr.c,v 1.95 2021/07/26 13:05:14 krw Exp $	*/
+/*	$OpenBSD: mbr.c,v 1.96 2021/08/06 10:41:31 krw Exp $	*/
 
 /*
  * Copyright (c) 1997 Tobias Weingartner
@@ -33,7 +33,11 @@
 #include "mbr.h"
 #include "gpt.h"
 
-struct mbr		initial_mbr;
+struct dos_mbr		default_dmbr;
+
+void		mbr_to_dos_mbr(const struct mbr *, struct dos_mbr *);
+void		dos_mbr_to_mbr(const struct dos_mbr *, const uint64_t,
+    const uint64_t, struct mbr *);
 
 void
 MBR_init(struct mbr *mbr)
@@ -45,11 +49,22 @@ MBR_init(struct mbr *mbr)
 	memset(&gh, 0, sizeof(gh));
 	memset(&gp, 0, sizeof(gp));
 
+	if (mbr->mbr_lba_self != 0) {
+		/* Extended MBR - save lba's, set sig, zap everything else. */
+		memset(mbr->mbr_code, 0, sizeof(mbr->mbr_code));
+		memset(mbr->mbr_prt, 0, sizeof(mbr->mbr_prt));
+		mbr->mbr_signature = DOSMBR_SIGNATURE;
+		return;
+	}
+
+	dos_mbr_to_mbr(&default_dmbr, 0, 0, mbr);
+
 	/*
-	 * XXX Do *NOT* zap all MBR parts! Some archs still read initmbr
-	 * from disk!! Just mark them inactive until -b goodness spreads
+	 * XXX Do *NOT* zap all MBR parts! Some archs still read default mmbr
+	 * from disk! Just mark them inactive until -b goodness spreads
 	 * further.
 	 */
+
 	mbr->mbr_prt[0].prt_flag = 0;
 	mbr->mbr_prt[1].prt_flag = 0;
 	mbr->mbr_prt[2].prt_flag = 0;
@@ -115,18 +130,18 @@ MBR_init(struct mbr *mbr)
 }
 
 void
-MBR_parse(const struct dos_mbr *dos_mbr, const uint64_t lba_self,
+dos_mbr_to_mbr(const struct dos_mbr *dmbr, const uint64_t lba_self,
     const uint64_t lba_firstembr, struct mbr *mbr)
 {
 	struct dos_partition	dos_parts[NDOSPART];
 	int			i;
 
-	memcpy(mbr->mbr_code, dos_mbr->dmbr_boot, sizeof(mbr->mbr_code));
+	memcpy(mbr->mbr_code, dmbr->dmbr_boot, sizeof(mbr->mbr_code));
 	mbr->mbr_lba_self = lba_self;
 	mbr->mbr_lba_firstembr = lba_firstembr;
-	mbr->mbr_signature = letoh16(dos_mbr->dmbr_sign);
+	mbr->mbr_signature = letoh16(dmbr->dmbr_sign);
 
-	memcpy(dos_parts, dos_mbr->dmbr_parts, sizeof(dos_parts));
+	memcpy(dos_parts, dmbr->dmbr_parts, sizeof(dos_parts));
 
 	for (i = 0; i < NDOSPART; i++)
 		PRT_parse(&dos_parts[i], lba_self, lba_firstembr,
@@ -134,7 +149,7 @@ MBR_parse(const struct dos_mbr *dos_mbr, const uint64_t lba_self,
 }
 
 void
-MBR_make(const struct mbr *mbr, struct dos_mbr *dos_mbr)
+mbr_to_dos_mbr(const struct mbr *mbr, struct dos_mbr *dos_mbr)
 {
 	struct dos_partition	dos_partition;
 	int			i;
@@ -178,7 +193,7 @@ MBR_read(const uint64_t lba_self, const uint64_t lba_firstembr, struct mbr *mbr)
 	memcpy(&dos_mbr, secbuf, sizeof(dos_mbr));
 	free(secbuf);
 
-	MBR_parse(&dos_mbr, lba_self, lba_firstembr, mbr);
+	dos_mbr_to_mbr(&dos_mbr, lba_self, lba_firstembr, mbr);
 
 	return 0;
 }
@@ -194,7 +209,7 @@ MBR_write(const struct mbr *mbr)
 	if (secbuf == NULL)
 		return -1;
 
-	MBR_make(mbr, &dos_mbr);
+	mbr_to_dos_mbr(mbr, &dos_mbr);
 	memcpy(secbuf, &dos_mbr, sizeof(dos_mbr));
 
 	rslt = DISK_writesectors(secbuf, mbr->mbr_lba_self, 1);
diff --git a/sbin/fdisk/mbr.h b/sbin/fdisk/mbr.h
index 532969989aa..214c776d1fe 100644
--- a/sbin/fdisk/mbr.h
+++ b/sbin/fdisk/mbr.h
@@ -1,4 +1,4 @@
-/*	$OpenBSD: mbr.h,v 1.40 2021/07/21 12:22:54 krw Exp $	*/
+/*	$OpenBSD: mbr.h,v 1.41 2021/08/06 10:41:31 krw Exp $	*/
 
 /*
  * Copyright (c) 1997 Tobias Weingartner
@@ -27,12 +27,9 @@ struct mbr {
 	uint16_t	mbr_signature;
 };
 
-extern struct mbr	initial_mbr;
+extern struct dos_mbr	default_dmbr;
 
 void		MBR_print(const struct mbr *, const char *);
-void		MBR_parse(const struct dos_mbr *, const uint64_t,
-    const uint64_t, struct mbr *);
-void		MBR_make(const struct mbr *, struct dos_mbr *);
 void		MBR_init(struct mbr *);
 int		MBR_read(const uint64_t, const uint64_t, struct mbr *);
 int		MBR_write(const struct mbr *);
diff --git a/sbin/fdisk/user.c b/sbin/fdisk/user.c
index 94924c1f516..09753167c48 100644
--- a/sbin/fdisk/user.c
+++ b/sbin/fdisk/user.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: user.c,v 1.67 2021/07/21 12:22:54 krw Exp $	*/
+/*	$OpenBSD: user.c,v 1.68 2021/08/06 10:41:31 krw Exp $	*/
 
 /*
  * Copyright (c) 1997 Tobias Weingartner
@@ -79,7 +79,11 @@ USER_edit(const uint64_t lba_self, const uint64_t lba_firstembr)
 
 again:
 	do {
-		printf("%s%s: %d> ", disk.dk_name, modified ? "*" : "", editlevel);
+		if (letoh64(gh.gh_sig) == GPTSIGNATURE && editlevel > 1)
+			goto done;	/* 'reinit gpt'. Unwind recursion! */
+
+		printf("%s%s: %d> ", disk.dk_name, modified ? "*" : "",
+		    editlevel);
 		fflush(stdout);
 		ask_cmd(&cmd, &args);