From 9ba61043bbb3c1c5354cd29f6ffba92b19fffc3e Mon Sep 17 00:00:00 2001
From: krw <krw@openbsd.org>
Date: Mon, 25 Jul 2022 17:45:16 +0000
Subject: [PATCH] Only MBR and GPT structures can be edited by fdisk.

If neither are found, restrict user actions to printing basic
information on the disk, reading the man page, initializing an
MBR or GPT, or terminating fdisk without changing anything.

Feedback on earlier attempt by miod@ and brynet@.
---
 sbin/fdisk/cmd.c  | 11 ++++++++---
 sbin/fdisk/mbr.c  | 22 +++++++++++++++++++++-
 sbin/fdisk/mbr.h  |  3 ++-
 sbin/fdisk/user.c | 42 ++++++++++++++++++++++++++++--------------
 sbin/fdisk/user.h |  4 ++--
 5 files changed, 61 insertions(+), 21 deletions(-)

diff --git a/sbin/fdisk/cmd.c b/sbin/fdisk/cmd.c
index 3ced91da012..dc8698901b9 100644
--- a/sbin/fdisk/cmd.c
+++ b/sbin/fdisk/cmd.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: cmd.c,v 1.163 2022/07/10 20:34:31 krw Exp $	*/
+/*	$OpenBSD: cmd.c,v 1.164 2022/07/25 17:45:16 krw Exp $	*/
 
 /*
  * Copyright (c) 1997 Tobias Weingartner
@@ -371,8 +371,13 @@ Xprint(const char *args, struct mbr *mbr)
 {
 	if (gh.gh_sig == GPTSIGNATURE)
 		GPT_print(args, VERBOSE);
-	else
+	else if (MBR_valid_prt(mbr))
 		MBR_print(mbr, args);
+	else {
+		DISK_printgeometry("s");
+		printf("Offset: %d\tSignature: 0x%X.\tNo MBR or GPT.\n",
+		    DOSBBSECTOR, (int)mbr->mbr_signature);
+	}
 
 	return CMD_CONT;
 }
@@ -430,7 +435,7 @@ Xexit(const char *args, struct mbr *mbr)
 int
 Xhelp(const char *args, struct mbr *mbr)
 {
-	USER_help();
+	USER_help(mbr);
 
 	return CMD_CONT;
 }
diff --git a/sbin/fdisk/mbr.c b/sbin/fdisk/mbr.c
index 11308d30445..2ef6783110f 100644
--- a/sbin/fdisk/mbr.c
+++ b/sbin/fdisk/mbr.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: mbr.c,v 1.119 2022/04/25 17:10:09 krw Exp $	*/
+/*	$OpenBSD: mbr.c,v 1.120 2022/07/25 17:45:16 krw Exp $	*/
 
 /*
  * Copyright (c) 1997 Tobias Weingartner
@@ -180,3 +180,23 @@ MBR_write(const struct mbr *mbr)
 
 	return 0;
 }
+
+int
+MBR_valid_prt(const struct mbr *mbr)
+{
+	uint64_t		bs, ns;
+	unsigned int		i, nprt;
+	unsigned char		id;
+
+	nprt = 0;
+	for (i = 0; i < NDOSPART; i++) {
+		bs = mbr->mbr_prt[i].prt_bs;
+		ns = mbr->mbr_prt[i].prt_ns;
+		id = mbr->mbr_prt[i].prt_id;
+		if ((bs == 0 && ns == 0 && id == 0) ||
+		    (bs < DL_GETDSIZE(&dl) && ns > 0 && ns <= DL_GETDSIZE(&dl)))
+			nprt++;
+	}
+
+	return nprt > 0 && mbr->mbr_signature == DOSMBR_SIGNATURE;
+}
diff --git a/sbin/fdisk/mbr.h b/sbin/fdisk/mbr.h
index ad449080f5a..0c56a3db4e0 100644
--- a/sbin/fdisk/mbr.h
+++ b/sbin/fdisk/mbr.h
@@ -1,4 +1,4 @@
-/*	$OpenBSD: mbr.h,v 1.42 2021/10/25 13:51:25 krw Exp $	*/
+/*	$OpenBSD: mbr.h,v 1.43 2022/07/25 17:45:16 krw Exp $	*/
 
 /*
  * Copyright (c) 1997 Tobias Weingartner
@@ -30,3 +30,4 @@ void		MBR_print(const struct mbr *, const char *);
 void		MBR_init(struct mbr *);
 int		MBR_read(const uint64_t, const uint64_t, struct mbr *);
 int		MBR_write(const struct mbr *);
+int		MBR_valid_prt(const struct mbr *);
diff --git a/sbin/fdisk/user.c b/sbin/fdisk/user.c
index 50b1d4b9abd..f4389ab18e6 100644
--- a/sbin/fdisk/user.c
+++ b/sbin/fdisk/user.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: user.c,v 1.83 2022/07/18 15:06:22 krw Exp $	*/
+/*	$OpenBSD: user.c,v 1.84 2022/07/25 17:45:16 krw Exp $	*/
 
 /*
  * Copyright (c) 1997 Tobias Weingartner
@@ -33,31 +33,34 @@
 
 struct cmd {
 	char	*cmd_name;
-	int	 cmd_gpt;
+	int	 cmd_valid;
 	int	(*cmd_fcn)(const char *, struct mbr *);
 	char	*cmd_help;
 };
 
 const struct cmd		cmd_table[] = {
-	{"help",   1, Xhelp,   "Display summary of available commands"},
-	{"manual", 1, Xmanual, "Display fdisk man page"},
-	{"reinit", 1, Xreinit, "Initialize the partition table"},
+	{"help",   2, Xhelp,   "Display summary of available commands"},
+	{"manual", 2, Xmanual, "Display fdisk man page"},
+	{"reinit", 2, Xreinit, "Initialize the partition table"},
 	{"setpid", 1, Xsetpid, "Set identifier of table entry"},
 	{"edit",   1, Xedit,   "Edit table entry"},
 	{"flag",   1, Xflag,   "Set flag value of table entry"},
 	{"update", 0, Xupdate, "Update MBR bootcode"},
 	{"select", 0, Xselect, "Select MBR extended table entry"},
 	{"swap",   1, Xswap,   "Swap two table entries"},
-	{"print",  1, Xprint,  "Print partition table"},
+	{"print",  2, Xprint,  "Print partition table"},
 	{"write",  1, Xwrite,  "Write partition table to disk"},
 	{"exit",   1, Xexit,   "Discard changes and exit edit level"},
 	{"quit",   1, Xquit,   "Save changes and exit edit level"},
-	{"abort",  1, Xabort,  "Discard changes and terminate fdisk"},
+	{"abort",  2, Xabort,  "Discard changes and terminate fdisk"},
 };
 
+#define	ANY_CMD(_i)	(cmd_table[(_i)].cmd_valid > 1)
+#define	GPT_CMD(_i)	(cmd_table[(_i)].cmd_valid > 0)
+
 int			modified;
 
-int			ask_cmd(char **);
+int			ask_cmd(const struct mbr *, char **);
 
 void
 USER_edit(const uint64_t lba_self, const uint64_t lba_firstembr)
@@ -84,7 +87,7 @@ USER_edit(const uint64_t lba_self, const uint64_t lba_firstembr)
 		printf("%s%s: %d> ", disk.dk_name, modified ? "*" : "",
 		    editlevel);
 		fflush(stdout);
-		i = ask_cmd(&args);
+		i = ask_cmd(&mbr, &args);
 		if (i == -1)
 			continue;
 
@@ -121,7 +124,14 @@ USER_print_disk(const int verbosity)
 	do {
 		if (MBR_read(lba_self, lba_firstembr, &mbr))
 			break;
-		if (lba_self == 0) {
+		if (lba_self == DOSBBSECTOR) {
+			if (MBR_valid_prt(&mbr) == 0) {
+				DISK_printgeometry("s");
+				printf("Offset: %d\tSignature: 0x%X.\t"
+				    "No MBR or GPT.\n", DOSBBSECTOR,
+				    (int)mbr.mbr_signature);
+				return;
+			}
 			if (GPT_read(ANYGPT)) {
 				if (verbosity == VERBOSE) {
 					printf("Primary GPT:\nNot Found\n");
@@ -161,20 +171,22 @@ USER_print_disk(const int verbosity)
 }
 
 void
-USER_help(void)
+USER_help(const struct mbr *mbr)
 {
 	unsigned int		i;
 
 	for (i = 0; i < nitems(cmd_table); i++) {
-		if (gh.gh_sig == GPTSIGNATURE && cmd_table[i].cmd_gpt == 0)
+		if (gh.gh_sig == GPTSIGNATURE && GPT_CMD(i) == 0)
 				continue;
+		if (MBR_valid_prt(mbr) == 0 && ANY_CMD(i) == 0)
+			continue;
 		printf("\t%s\t\t%s\n", cmd_table[i].cmd_name,
 		    cmd_table[i].cmd_help);
 	}
 }
 
 int
-ask_cmd(char **arg)
+ask_cmd(const struct mbr *mbr, char **arg)
 {
 	static char		 lbuf[LINEBUFSZ];
 	char			*cmd;
@@ -194,7 +206,9 @@ ask_cmd(char **arg)
 		cmd = "help";
 
 	for (i = 0; i < nitems(cmd_table); i++) {
-		if (gh.gh_sig == GPTSIGNATURE && cmd_table[i].cmd_gpt == 0)
+		if (gh.gh_sig == GPTSIGNATURE && GPT_CMD(i) == 0)
+			continue;
+		if (MBR_valid_prt(mbr) == 0 && ANY_CMD(i) == 0)
 			continue;
 		if (strstr(cmd_table[i].cmd_name, cmd) == cmd_table[i].cmd_name)
 			return i;
diff --git a/sbin/fdisk/user.h b/sbin/fdisk/user.h
index 11010a1aac2..e358f81f5b1 100644
--- a/sbin/fdisk/user.h
+++ b/sbin/fdisk/user.h
@@ -1,4 +1,4 @@
-/*	$OpenBSD: user.h,v 1.23 2022/07/10 17:46:03 krw Exp $	*/
+/*	$OpenBSD: user.h,v 1.24 2022/07/25 17:45:16 krw Exp $	*/
 
 /*
  * Copyright (c) 1997 Tobias Weingartner
@@ -21,4 +21,4 @@
 
 void		USER_edit(const uint64_t, const uint64_t);
 void		USER_print_disk(const int);
-void		USER_help(void);
+void		USER_help(const struct mbr *);
-- 
2.20.1