1 /* 2 * Copyright 2020-2021 Haiku, Inc. All rights reserved. 3 * All rights reserved. Distributed under the terms of the MIT License. 4 * 5 * Authors: 6 * Alexander von Gluck IV <kallisti5@unixzen.com> 7 */ 8 9 10 #include <errno.h> 11 #include <getopt.h> 12 #include <fcntl.h> 13 #include <unistd.h> 14 #include <stdint.h> 15 #include <stdio.h> 16 #include <stdlib.h> 17 #include <string.h> 18 #include <sys/stat.h> 19 20 #define EXIT_FAILURE 1 21 22 23 //#define DEBUG_MBRTOOL 24 #ifdef DEBUG_MBRTOOL 25 # define TRACE(x...) printf("mbrtool: " x) 26 #else 27 # define TRACE(x...) ; 28 #endif 29 30 #define INFO(x...) printf("mbrtool: " x) 31 32 33 // Disk sector size, 512 assumed! 34 static const size_t kSectorSize = 512; 35 36 static void 37 print_usage(bool error) 38 { 39 printf("\n"); 40 printf("usage: mbrtool (options) <diskImage> <id> <type> <start> <len>\n"); 41 printf(" <diskImage> Disk image to operate on\n"); 42 printf(" <id> Partition ID (0-3)\n"); 43 printf(" <type> Partition type id (hex)\n"); 44 printf(" <start> Partition start offset (KiB)\n"); 45 printf(" <len> Partition length (KiB)\n\n"); 46 printf(" Options:\n"); 47 printf(" -a, --active Partition boot flag\n"); 48 printf("\nWarning: This tool requires precision!\n"); 49 printf(" Inputs are only lightly validated!\n\n"); 50 exit(error ? EXIT_FAILURE : 0); 51 } 52 53 54 static void 55 checkError(bool failed, const char *message) 56 { 57 if (!failed) 58 return; 59 60 if (errno != 0) 61 INFO("Error: %s: %s\n", message, strerror(errno)); 62 else 63 INFO("Error: %s\n", message); 64 65 exit(1); 66 } 67 68 69 static ssize_t 70 mbrWipe(int handle) 71 { 72 // Blow away MBR while avoiding bootloader 73 uint8_t emptyMBR[66] = {}; 74 // End of MBR marker 75 emptyMBR[64] = 0x55; 76 emptyMBR[65] = 0xAA; 77 return pwrite(handle, emptyMBR, 66, 0x1BE); 78 } 79 80 81 static bool 82 mbrValid(int handle) 83 { 84 // TODO: this is a really basic check and ignores invalid 85 // partition table entries 86 uint8_t mbrBytes[66] = {}; 87 ssize_t read = pread(handle, mbrBytes, 66, 0x1BE); 88 checkError(read < 0, "failed to read MBR for validation"); 89 return (mbrBytes[64] == 0x55 && mbrBytes[65] == 0xAA); 90 } 91 92 93 static void 94 createPartition(int handle, int index, bool active, uint8_t type, 95 uint64_t offset, uint64_t size) 96 { 97 uint8_t bootable = active ? 0x80 : 0x0; 98 uint8_t partition[16] = { 99 bootable, // bootable 100 0xff, 0xff, 0xff, // CHS first block (default to LBA) 101 type, // partition type 102 0xff, 0xff, 0xff, // CHS last block (default to LBA) 103 0x00, 0x00, 0x00, 0x00, // imageOffset in blocks (written below) 104 0x00, 0x00, 0x00, 0x00 // imageSize in blocks (written below) 105 }; 106 107 // fill in LBA values 108 uint32_t partitionOffset = (uint32_t)(offset / kSectorSize); 109 ((uint32_t *)partition)[2] = partitionOffset; 110 ((uint32_t *)partition)[3] = (uint32_t)(size / kSectorSize); 111 112 TRACE("%s: #%d %c bytes: %u-%u, sectors: %u-%u\n", __func__, index, 113 active ? 'b' : '-', offset, offset + size, partitionOffset, 114 partitionOffset + uint32_t(size / kSectorSize)); 115 116 ssize_t written = pwrite(handle, partition, 16, 512 - 2 - 16 * (4 - index)); 117 checkError(written != 16, "failed to write partition entry"); 118 119 if (active) { 120 // make it bootable 121 written = pwrite(handle, &partitionOffset, 4, offset + 512 - 2 - 4); 122 checkError(written != 4, "failed to make partition bootable"); 123 } 124 return; 125 } 126 127 128 int 129 main(int argc, char *argv[]) 130 { 131 const char *imageFile = NULL; 132 133 int partType = -1; 134 int partIndex = -1; 135 int64_t partStartOffset = -1; 136 int64_t partLength = -1; 137 bool partBootable = false; 138 139 while (1) { 140 int c; 141 static struct option long_options[] = { 142 {"active", no_argument, 0, 'a'}, 143 {"help", no_argument, 0, 'h'}, 144 {0, 0, 0, 0} 145 }; 146 147 opterr = 1; /* don't print errors */ 148 c = getopt_long(argc, argv, "+ha", long_options, NULL); 149 150 if (c == -1) 151 break; 152 153 switch (c) { 154 case 's': 155 print_usage(false); 156 break; 157 case 'a': 158 partBootable = true; 159 break; 160 default: 161 print_usage(true); 162 } 163 } 164 165 if ((argc - optind) != 5) 166 print_usage(true); 167 168 for (int index = optind; index < argc; index++) { 169 if (imageFile == NULL) 170 imageFile = argv[index]; 171 else if (partIndex < 0) 172 partIndex = atoi(argv[index]); 173 else if (partType < 0) 174 partType = (int)strtol(argv[index], NULL, 0); 175 else if (partStartOffset < 0 ) { 176 partStartOffset = (int64_t)strtol(argv[index], NULL, 10); 177 partStartOffset *= 1024; 178 } else if (partLength < 0) { 179 partLength = (int64_t)strtol(argv[index], NULL, 10); 180 partLength *= 1024; 181 } 182 } 183 184 checkError(partIndex < 0 || partIndex > 3, 185 "invalid partition index, valid range is 0-3"); 186 checkError(partType < 0, "Incorrect Partition Type!"); 187 188 checkError((partStartOffset + partLength) > 2089072000000, 189 "partitions beyond 2TiB are not accepted!"); 190 191 int imageFileHandle = open(imageFile, O_RDWR); 192 checkError(imageFileHandle < 0, "failed to open disk image file"); 193 194 struct stat stat; 195 int result = fstat(imageFileHandle, &stat); 196 checkError(result != 0, "failed to stat image file"); 197 off_t imageSize = stat.st_size; 198 199 if (!mbrValid(imageFileHandle)) { 200 INFO("MBR of image is invalid, creating a fresh one.\n"); 201 mbrWipe(imageFileHandle); 202 } 203 204 // Just a warning. This is technically valid since MBR partition 205 // definitions are purely within the first 512 bytes 206 if (partStartOffset + partLength > imageSize) 207 INFO("Warning: Partition extends beyond end of file!\n"); 208 209 createPartition(imageFileHandle, partIndex, partBootable, partType, 210 partStartOffset, partLength); 211 return 0; 212 } 213