Plug leaks, fix unaligned writes, refactor some, slightly better test

5 files changed, 144 insertions, 51 deletions

diff --git a/Makefile b/Makefile
index 707b325..fe356f0 100644
--- a/Makefile
+++ b/Makefile
@@ -1,4 +1,4 @@
-CFLAGS+=-Wall -Wextra -Wpedantic -std=c11 -D_DEFAULT_SOURCE
+CFLAGS+=-Wall -Wextra -Wpedantic -std=c11 -D_DEFAULT_SOURCE #-D_FORTIFY_SOURCE=2
 ALL:=mkgpt
 PREFIX:=/usr/local
 
diff --git a/README.md b/README.md
index 1e05497..32bbf0d 100644
--- a/README.md
+++ b/README.md
@@ -60,7 +60,7 @@ Optionally, the string `random` can be used to generate a random GUID.
 - no way to create swap partitions (borked GUID parsing)
 - some overly specific code (Linux, Windoze) without a strong need for it
 - some strange code (static info in oversized, dynamically populated array?)
-- some rather broken code (required calls in asserts, misaligned pointers)
+- some rather broken code (required calls in asserts, misaligned pointers, memory leaks)
 
 ## References
 
@@ -71,6 +71,7 @@ Not enough people seem to know about these things, so here you go:
 - https://www.embedded.com/reduce-c-language-coding-errors-with-x-macros-part-2/
 - https://www.embedded.com/reduce-c-language-coding-errors-with-x-macros-part-3/
 - http://www.catb.org/esr/structure-packing/
+- https://nullprogram.com/blog/2016/11/22/
 
 Here's some background on the whole GPT mess:
 
@@ -79,6 +80,9 @@ Here's some background on the whole GPT mess:
   (I misunderstood this at first, read if you come from MBR)
 - Apple's [Secrets of the GPT](https://developer.apple.com/library/archive/technotes/tn2166/_index.html)
   has some useful background (once you ignore the Apple-rific stuff)
+- https://thestarman.pcministry.com/asm/mbr/GPT.htm
+  (points out disagreements between Windoze protective MBRs and the UEFI
+  specification)
 
 Want to scare yourself?
 
diff --git a/mkgpt.c b/mkgpt.c
index 70f267b..3f8ad36 100644
--- a/mkgpt.c
+++ b/mkgpt.c
@@ -25,6 +25,7 @@
 #include "guid.h"
 #include "part.h"
 #include "part_ids.h"
+#include "unaligned.h"
 
 #include <assert.h>
 #include <errno.h>
@@ -34,7 +35,9 @@
 #include <stdlib.h>
 #include <string.h>
 
-#define MAX_PART_NAME (36)
+#define MAX_PART_NAME (36U)
+#define MIN_SECTOR_SIZE (512U)
+#define MAX_SECTOR_SIZE (4096U)
 
 static void
 dump_help(char *fname);
@@ -51,7 +54,7 @@ min(const int a, const int b)
 	return a < b ? a : b;
 }
 
-static size_t sect_size = 512;
+static size_t sect_size = MIN_SECTOR_SIZE;
 static long image_sects = 0;
 static long min_image_sects = 2048;
 static struct partition *first_part = NULL;
@@ -144,13 +147,15 @@ parse_opts(int argc, char **argv)
 
 			sect_size = atoi(argv[i]);
 
-			if (sect_size < 512 || sect_size > 4096 ||
-				sect_size % 512) {
+			if (sect_size < MIN_SECTOR_SIZE ||
+				sect_size > MAX_SECTOR_SIZE ||
+				sect_size % MIN_SECTOR_SIZE) {
 				fprintf(stderr,
 					"invalid sector size (%zu) - must be "
-					">= 512 and <= 4096 and "
-					"a multiple of 512",
-					sect_size);
+					">= %u and <= %u and "
+					"a multiple of %u",
+					sect_size, MIN_SECTOR_SIZE,
+					MAX_SECTOR_SIZE, MIN_SECTOR_SIZE);
 				return -1;
 			}
 			i++;
@@ -456,48 +461,52 @@ panic(const char *msg)
 	exit(EXIT_FAILURE);
 }
 
+/*
+ * Write the "protective MBR" to FILE *output.
+ */
 static void
-write_output()
+write_mbr(void)
 {
-	int i;
-	uint8_t *mbr, *gpt, *gpt2, *parts, *image_buf;
-	struct partition *cur_part;
-
-	/* Write MBR */
-	mbr = calloc(1, sect_size);
-	if (mbr == NULL) {
-		panic("calloc failed");
+	uint8_t mbr[MAX_SECTOR_SIZE] = {0};
+	assert(sect_size <= sizeof(mbr));
+
+	/* entry for "Partition 1" starts here */
+	uint8_t *p1 = mbr + 446;
+
+	/* boot indicator = 0, start CHS = 0x000200 */
+	set_u32(p1 + 0, 0x00020000);
+	/* OSType 0xee = GPT Protective, EndingCHS = 0xffffff */
+	set_u32(p1 + 4, 0xffffffee);
+	/* StartingLBA = 1 */
+	set_u32(p1 + 8, 0x00000001);
+	/* number of sectors in partition */
+	if (image_sects > 0xffffffff) {
+		set_u32(p1 + 12, 0xffffffff);
+	} else {
+		assert(image_sects > 1); /* 0 sectors is not allowed */
+		set_u32(p1 + 12, image_sects - 1);
 	}
-
-	*(uint32_t *)&mbr[446] =
-		0x00020000; /* boot indicator = 0, start CHS = 0x000200 */
-	mbr[446 + 4] = 0xee; /* OSType = GPT Protective */
-	mbr[446 + 5] = 0xff;
-	mbr[446 + 6] = 0xff;
-	mbr[446 + 7] = 0xff; /* EndingCHS = 0xffffff */
-	*(uint32_t *)&mbr[446 + 8] = 0x1; /* StartingLBA = 1 */
-
-	if (image_sects > 0xffffffff)
-		*(uint32_t *)&mbr[446 + 12] = 0xffffffff;
-	else
-		*(uint32_t *)&mbr[446 + 12] = (uint32_t)image_sects - 1;
-
-	mbr[510] = 0x55;
-	mbr[511] = 0xaa; /* Signature */
+	/* Signature */
+	set_u16(mbr + 510, 0xaa55);
 
 	if (fwrite(mbr, 1, sect_size, output) != sect_size) {
 		panic("fwrite failed");
 	}
+}
+
+static void
+write_output(void)
+{
+	uint8_t *parts;
+	struct partition *cur_part;
+
+	write_mbr();
 
 	/* Define GPT headers */
-	gpt = calloc(1, sect_size);
-	if (gpt == NULL) {
-		panic("calloc failed");
-	}
-	gpt2 = calloc(1, sect_size);
-	if (gpt2 == NULL) {
-		panic("calloc failed");
-	}
+	uint8_t gpt[MAX_SECTOR_SIZE] = {0};
+	assert(sect_size <= sizeof(gpt));
+	uint8_t gpt2[MAX_SECTOR_SIZE] = {0};
+	assert(sect_size <= sizeof(gpt2));
 
 	*(uint64_t *)&gpt[0] = 0x5452415020494645ULL; /* Signature */
 	*(uint32_t *)&gpt[8] = 0x00010000UL; /* Revision */
@@ -521,7 +530,7 @@ write_output()
 	}
 
 	cur_part = first_part;
-	i = 0;
+	int i = 0;
 	while (cur_part) {
 		int char_id;
 
@@ -540,6 +549,9 @@ write_output()
 		/*
 		 * TODO settle missing UTF-16LE conversion issue somehow,
 		 * possibly by simply limiting the tool to ASCII here?
+		 * TODO used to be "&& char_id < 35" but MAX_PART_NAME is 36
+		 * now so which is correct? do we need a "double zero" at the
+		 * end or not? what does the spec say?
 		 */
 		int len = min(strlen(cur_part->name), MAX_PART_NAME);
 		for (char_id = 0; char_id < len; char_id++) {
@@ -573,7 +585,8 @@ write_output()
 
 	/* Write partitions */
 	cur_part = first_part;
-	image_buf = malloc(sect_size);
+	uint8_t image_buf[MAX_SECTOR_SIZE] = {0};
+	assert(sect_size <= sizeof(image_buf));
 	while (cur_part) {
 		size_t bytes_read;
 		size_t bytes_written = 0;
@@ -611,4 +624,6 @@ write_output()
 	if (fwrite(gpt2, 1, sect_size, output) != sect_size) {
 		panic("fwrite failed");
 	}
+
+	free(parts);
 }
diff --git a/test-mkgpt.sh b/test-mkgpt.sh
index 7c3c54d..d952705 100755
--- a/test-mkgpt.sh
+++ b/test-mkgpt.sh
@@ -15,14 +15,22 @@ for name in a.img b.img c.img d.img e.img; do
 done
 
 # TODO something goes wrong with --sector-size 1024 or 4096
-./mkgpt -o ${tmpdir}/bla.img -s 131072 \
-	--part ${tmpdir}/a.img --type system --name part_system_a \
-	--part ${tmpdir}/b.img --type fat32 --name part_fat32_b \
-	--part ${tmpdir}/c.img --type linux --name X \
-	--part ${tmpdir}/d.img --type 0x82 --name 123456789012345678901234567890123456 \
-	--part ${tmpdir}/e.img --type 21686148-6449-6E6F-744E-656564454649
+# TODO --uuid "all zero" is taken to mean "random uuid"?
+./mkgpt -o ${tmpdir}/bla.img -s 131072 --disk-guid 1ABC2ABC-1111-2222-3333-1ABC2ABC3ABC \
+	--part ${tmpdir}/a.img --type system --name part_system_a --uuid 33333333-3333-3333-3333-333333333333 \
+	--part ${tmpdir}/b.img --type fat32 --name part_fat32_b --uuid 11111111-1111-1111-1111-111111111111 \
+	--part ${tmpdir}/c.img --type linux --name X --uuid 01234567-89AB-CDEF-0123-456789ABCDEF \
+	--part ${tmpdir}/d.img --type 0x82 --name 123456789012345678901234567890123456 --uuid 22222222-2222-2222-2222-222222222222 \
+	--part ${tmpdir}/e.img --type 21686148-6449-6E6F-744E-656564454649 --uuid 44444444-4444-4444-4444-444444444444
 
 fdisk -l ${tmpdir}/bla.img
-#xxd ${tmpdir}/bla.img | less
+sfdisk --verify ${tmpdir}/bla.img
+head -c 512 ${tmpdir}/bla.img | xxd -s 446
+
+checksum=$(md5sum ${tmpdir}/bla.img | cut -c1-32)
+if [ ! "$checksum" = "a4abc130196a29288ab9af1c4489fe24" ]; then
+	echo "checksum didn't match, regression!"
+	exit 1
+fi
 
 rm -rfv ${tmpdir}
diff --git a/unaligned.h b/unaligned.h
new file mode 100644
index 0000000..02cff81
--- /dev/null
+++ b/unaligned.h
@@ -0,0 +1,66 @@
+#pragma once
+
+/* SPDX-License-Identifier: MIT */
+
+#ifndef UNALIGNED_H
+#define UNALIGNED_H
+
+/*
+ * Unaligned little-endian memory access. See Chris Wellons' excellent post at
+ * https://nullprogram.com/blog/2016/11/22/ if you're confused by this.
+ */
+
+#include <stdint.h>
+
+static inline uint16_t
+get_u16(const uint8_t *buf)
+{
+	return (uint16_t)buf[1] << 8 | (uint16_t)buf[0] << 0;
+}
+
+static inline uint32_t
+get_u32(const uint8_t *buf)
+{
+	return (uint32_t)buf[3] << 24 | (uint32_t)buf[2] << 16 |
+	       (uint32_t)buf[1] << 8 | (uint32_t)buf[0] << 0;
+}
+
+static inline uint64_t
+get_u64(const uint8_t *buf)
+{
+	return (uint64_t)buf[7] << 56 | (uint64_t)buf[6] << 48 |
+	       (uint64_t)buf[5] << 40 | (uint64_t)buf[4] << 32 |
+	       (uint64_t)buf[3] << 24 | (uint64_t)buf[2] << 16 |
+	       (uint64_t)buf[1] << 8 | (uint64_t)buf[0] << 0;
+}
+
+static inline void
+set_u16(uint8_t *buf, const uint16_t val)
+{
+	buf[0] = val >> 0;
+	buf[1] = val >> 8;
+}
+
+static inline void
+set_u32(uint8_t *buf, const uint32_t val)
+{
+	buf[0] = val >> 0;
+	buf[1] = val >> 8;
+	buf[2] = val >> 16;
+	buf[3] = val >> 24;
+}
+
+static inline void
+set_u64(uint8_t *buf, const uint64_t val)
+{
+	buf[0] = val >> 0;
+	buf[1] = val >> 8;
+	buf[2] = val >> 16;
+	buf[3] = val >> 24;
+	buf[4] = val >> 32;
+	buf[5] = val >> 40;
+	buf[6] = val >> 48;
+	buf[7] = val >> 56;
+}
+
+#endif