Use nitems() of source and destination partition arrays

rather than assuming they are the same size. Zero
destination entries when no source partition available.

No intentional functional change.

2 files changed, 21 insertions, 13 deletions

diff --git a/sbin/fdisk/gpt.c b/sbin/fdisk/gpt.c
index 7f82960ead5..bff1007e09e 100644
--- a/sbin/fdisk/gpt.c
+++ b/sbin/fdisk/gpt.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: gpt.c,v 1.85 2023/03/04 23:09:15 krw Exp $	*/
+/*	$OpenBSD: gpt.c,v 1.86 2023/03/06 13:24:40 krw Exp $	*/
 /*
  * Copyright (c) 2015 Markus Muller <mmu@grummel.net>
  * Copyright (c) 2015 Kenneth R Westerback <krw@openbsd.org>
@@ -124,14 +124,16 @@ int
 protective_mbr(const struct mbr *mbr)
 {
 	struct dos_partition	dp[NDOSPART], dos_partition;
-	int			i;
+	unsigned int		i;
 
 	if (mbr->mbr_lba_self != 0)
 		return -1;
 
-	for (i = 0; i < NDOSPART; i++) {
-		PRT_make(&mbr->mbr_prt[i], mbr->mbr_lba_self,
-		    mbr->mbr_lba_firstembr, &dos_partition);
+	for (i = 0; i < nitems(dp); i++) {
+		memset(&dos_partition, 0, sizeof(dos_partition));
+		if (i < nitems(mbr->mbr_prt))
+			PRT_make(&mbr->mbr_prt[i], mbr->mbr_lba_self,
+			    mbr->mbr_lba_firstembr, &dos_partition);
 		memcpy(&dp[i], &dos_partition, sizeof(dp[i]));
 	}
 
diff --git a/sbin/fdisk/mbr.c b/sbin/fdisk/mbr.c
index ef88ed4cb4a..b7aeea00d9e 100644
--- a/sbin/fdisk/mbr.c
+++ b/sbin/fdisk/mbr.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: mbr.c,v 1.122 2023/03/04 21:22:51 krw Exp $	*/
+/*	$OpenBSD: mbr.c,v 1.123 2023/03/06 13:24:40 krw Exp $	*/
 
 /*
  * Copyright (c) 1997 Tobias Weingartner
@@ -120,23 +120,29 @@ dos_mbr_to_mbr(const struct dos_mbr *dmbr, const uint64_t lba_self,
 
 	memcpy(dos_parts, dmbr->dmbr_parts, sizeof(dos_parts));
 
-	for (i = 0; i < NDOSPART; i++)
-		PRT_parse(&dos_parts[i], lba_self, lba_firstembr,
-		    &mbr->mbr_prt[i]);
+	for (i = 0; i < nitems(mbr->mbr_prt); i++) {
+		memset(&mbr->mbr_prt[i], 0, sizeof(mbr->mbr_prt[i]));
+		if (i < nitems(dmbr->dmbr_parts))
+			PRT_parse(&dos_parts[i], lba_self, lba_firstembr,
+			    &mbr->mbr_prt[i]);
+	}
 }
 
 void
 mbr_to_dos_mbr(const struct mbr *mbr, struct dos_mbr *dos_mbr)
 {
 	struct dos_partition	dos_partition;
-	int			i;
+	unsigned int		i;
 
 	memcpy(dos_mbr->dmbr_boot, mbr->mbr_code, sizeof(dos_mbr->dmbr_boot));
 	dos_mbr->dmbr_sign = htole16(DOSMBR_SIGNATURE);
 
-	for (i = 0; i < NDOSPART; i++) {
-		PRT_make(&mbr->mbr_prt[i], mbr->mbr_lba_self,
-		    mbr->mbr_lba_firstembr, &dos_partition);
+	for (i = 0; i < nitems(dos_mbr->dmbr_parts); i++) {
+		memset(&dos_partition, 0, sizeof(dos_partition));
+		if (i < nitems(mbr->mbr_prt)) {
+			PRT_make(&mbr->mbr_prt[i], mbr->mbr_lba_self,
+			    mbr->mbr_lba_firstembr, &dos_partition);
+		}
 		memcpy(&dos_mbr->dmbr_parts[i], &dos_partition,
 		    sizeof(dos_mbr->dmbr_parts[i]));
 	}