tools/mbrtool: A simple tool to create MBR partitions

Change-Id: I0f04b257ad49d07c03d630df47c4891c1fd7a954
Reviewed-on: https://review.haiku-os.org/c/haiku/+/3570
Reviewed-by: Stephan Aßmus <[email protected]>

src/tools/Jamfile

@@ -106,6 +106,7 @@ SubInclude HAIKU_TOP src tools hvif2png ;
 SubInclude HAIKU_TOP src tools keymap ;
 SubInclude HAIKU_TOP src tools locale ;
 SubInclude HAIKU_TOP src tools makebootable ;
+SubInclude HAIKU_TOP src tools mbrtool ;
 SubInclude HAIKU_TOP src tools opd_to_package_info ;
 SubInclude HAIKU_TOP src tools package ;
 SubInclude HAIKU_TOP src tools package_repo ;

src/tools/mbrtool/Jamfile

@@ -0,0 +1,3 @@
+SubDir HAIKU_TOP src tools mbrtool ;
+
+BuildPlatformMain <build>mbrtool : mbrtool.cpp : $(HOST_LIBSUPC++) ;

src/tools/mbrtool/mbrtool.cpp

@@ -0,0 +1,212 @@
+/*
+ * Copyright 2020-2021 Haiku, Inc. All rights reserved.
+ * All rights reserved. Distributed under the terms of the MIT License.
+ *
+ * Authors:
+ * 		Alexander von Gluck IV <[email protected]>
+ */
+
+
+#include <errno.h>
+#include <getopt.h>
+#include <fcntl.h>
+#include <unistd.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/stat.h>
+
+#define EXIT_FAILURE 1
+
+
+//#define DEBUG_MBRTOOL
+#ifdef DEBUG_MBRTOOL
+#   define TRACE(x...) printf("mbrtool: " x)
+#else
+#   define TRACE(x...) ;
+#endif
+
+#define INFO(x...) printf("mbrtool: " x)
+
+
+// Disk sector size, 512 assumed!
+static const size_t kSectorSize = 512;
+
+static void
+print_usage(bool error)
+{
+	printf("\n");
+	printf("usage: mbrtool (options) <diskImage> <id> <type> <start> <len>\n");
+	printf("       <diskImage>            Disk image to operate on\n");
+	printf("       <id>                   Partition ID (0-3)\n");
+	printf("       <type>                 Partition type id (hex)\n");
+	printf("       <start>                Partition start offset (KiB)\n");
+	printf("       <len>                  Partition length (KiB)\n\n");
+	printf("  Options:\n");
+	printf("       -a, --active           Partition boot flag\n");
+	printf("\nWarning: This tool requires precision!\n");
+	printf("         Inputs are only lightly validated!\n\n");
+	exit(error ? EXIT_FAILURE : 0);
+}
+
+
+static void
+checkError(bool failed, const char *message)
+{
+	if (!failed)
+		return;
+
+	if (errno != 0)
+		INFO("Error: %s: %s\n", message, strerror(errno));
+	else
+		INFO("Error: %s\n", message);
+
+	exit(1);
+}
+
+
+static ssize_t
+mbrWipe(int handle)
+{
+	// Blow away MBR while avoiding bootloader
+	uint8_t emptyMBR[66] = {};
+	// End of MBR marker
+	emptyMBR[64] = 0x55;
+	emptyMBR[65] = 0xAA;
+	return pwrite(handle, emptyMBR, 66, 0x1BE);
+}
+
+
+static bool
+mbrValid(int handle)
+{
+	// TODO: this is a really basic check and ignores invalid
+	// partition table entries
+	uint8_t mbrBytes[66] = {};
+	ssize_t read = pread(handle, mbrBytes, 66, 0x1BE);
+	checkError(read < 0, "failed to read MBR for validation");
+	return (mbrBytes[64] == 0x55 && mbrBytes[65] == 0xAA);
+}
+
+
+static void
+createPartition(int handle, int index, bool active, uint8_t type,
+	uint64_t offset, uint64_t size)
+{
+	uint8_t bootable = active ? 0x80 : 0x0;
+	uint8_t partition[16] = {
+		bootable,		// bootable
+		0xff, 0xff, 0xff,	// CHS first block (default to LBA)
+		type,			// partition type
+		0xff, 0xff, 0xff,	// CHS last block (default to LBA)
+		0x00, 0x00, 0x00, 0x00,	// imageOffset in blocks (written below)
+		0x00, 0x00, 0x00, 0x00	// imageSize in blocks (written below)
+	};
+
+	// fill in LBA values
+	uint32_t partitionOffset = (uint32_t)(offset / kSectorSize);
+	((uint32_t *)partition)[2] = partitionOffset;
+	((uint32_t *)partition)[3] = (uint32_t)(size / kSectorSize);
+
+	TRACE("%s: #%d %c bytes: %u-%u, sectors: %u-%u\n", __func__, index,
+		active ? 'b' : '-', offset, offset + size, partitionOffset,
+		partitionOffset + uint32_t(size / kSectorSize));
+
+	ssize_t written = pwrite(handle, partition, 16, 512 - 2 - 16 * (4 - index));
+	checkError(written != 16, "failed to write partition entry");
+
+	if (active) {
+		// make it bootable
+		written = pwrite(handle, &partitionOffset, 4, offset + 512 - 2 - 4);
+		checkError(written != 4, "failed to make partition bootable");
+	}
+	return;
+}
+
+
+int
+main(int argc, char *argv[])
+{
+	const char *imageFile = NULL;
+
+	int partType = -1;
+	int partIndex = -1;
+	int64_t partStartOffset = -1;
+	int64_t partLength = -1;
+	bool partBootable = false;
+
+	while (1) {
+		int c;
+		static struct option long_options[] = {
+			{"active", no_argument, 0, 'a'},
+			{"help", no_argument, 0, 'h'},
+			{0, 0, 0, 0}
+		};
+
+		opterr = 1; /* don't print errors */
+		c = getopt_long(argc, argv, "+ha", long_options, NULL);
+
+		if (c == -1)
+			break;
+
+		switch (c) {
+			case 's':
+				print_usage(false);
+				break;
+			case 'a':
+				partBootable = true;
+				break;
+			default:
+				print_usage(true);
+		}
+	}
+
+	if ((argc - optind) != 5)
+		print_usage(true);
+
+	for (int index = optind; index < argc; index++) {
+		if (imageFile == NULL)
+			imageFile = argv[index];
+		else if (partIndex < 0)
+			partIndex = atoi(argv[index]);
+		else if (partType < 0)
+			partType = (int)strtol(argv[index], NULL, 0);
+		else if (partStartOffset < 0 ) {
+			partStartOffset = (int64_t)strtol(argv[index], NULL, 10);
+			partStartOffset *= 1024;
+		} else if (partLength < 0) {
+			partLength = (int64_t)strtol(argv[index], NULL, 10);
+			partLength *= 1024;
+		}
+	}
+
+	checkError(partIndex < 0 || partIndex > 3,
+		"invalid partition index, valid range is 0-3");
+	checkError(partType < 0,  "Incorrect Partition Type!");
+
+	checkError((partStartOffset + partLength) > 2089072000000,
+		"partitions beyond 2TiB are not accepted!");
+
+	int imageFileHandle = open(imageFile, O_RDWR);
+	checkError(imageFileHandle < 0, "failed to open disk image file");
+
+	struct stat stat;
+	int result = fstat(imageFileHandle, &stat);
+	checkError(result != 0, "failed to stat image file");
+	off_t imageSize = stat.st_size;
+
+	if (!mbrValid(imageFileHandle)) {
+		INFO("MBR of image is invalid, creating a fresh one.\n");
+		mbrWipe(imageFileHandle);
+	}
+
+	// Just a warning. This is technically valid since MBR partition
+	// definitions are purely within the first 512 bytes
+	if (partStartOffset + partLength > imageSize)
+		INFO("Warning: Partition extends beyond end of file!\n");
+
+	createPartition(imageFileHandle, partIndex, partBootable, partType,
+		partStartOffset, partLength);
+	return 0;
+}