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
print_usage(bool error)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
checkError(bool failed,const char * message)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
mbrWipe(int handle)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
mbrValid(int handle)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
createPartition(int handle,int index,bool active,uint8_t type,uint64_t offset,uint64_t size)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
main(int argc,char * argv[])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