1 /* basicmbr.cc -- Functions for loading, saving, and manipulating legacy MBR partition
2    data. */
3 
4 /* Initial coding by Rod Smith, January to February, 2009 */
5 
6 /* This program is copyright (c) 2009-2013 by Roderick W. Smith. It is distributed
7   under the terms of the GNU GPL version 2, as detailed in the COPYING file. */
8 
9 #define __STDC_LIMIT_MACROS
10 #ifndef __STDC_CONSTANT_MACROS
11 #define __STDC_CONSTANT_MACROS
12 #endif
13 
14 #include <stdio.h>
15 #include <stdlib.h>
16 #include <stdint.h>
17 #include <fcntl.h>
18 #include <string.h>
19 #include <time.h>
20 #include <sys/stat.h>
21 #include <errno.h>
22 #include <iostream>
23 #include <algorithm>
24 #include "mbr.h"
25 #include "support.h"
26 
27 using namespace std;
28 
29 /****************************************
30  *                                      *
31  * MBRData class and related structures *
32  *                                      *
33  ****************************************/
34 
BasicMBRData(void)35 BasicMBRData::BasicMBRData(void) {
36    blockSize = SECTOR_SIZE;
37    diskSize = 0;
38    device = "";
39    state = invalid;
40    numHeads = MAX_HEADS;
41    numSecspTrack = MAX_SECSPERTRACK;
42    myDisk = NULL;
43    canDeleteMyDisk = 0;
44 //   memset(&EbrLocations, 0, MAX_MBR_PARTS * sizeof(uint32_t));
45    EmptyMBR();
46 } // BasicMBRData default constructor
47 
BasicMBRData(const BasicMBRData & orig)48 BasicMBRData::BasicMBRData(const BasicMBRData & orig) {
49    int i;
50 
51    if (&orig != this) {
52       memcpy(code, orig.code, 440);
53       diskSignature = orig.diskSignature;
54       nulls = orig.nulls;
55       MBRSignature = orig.MBRSignature;
56       blockSize = orig.blockSize;
57       diskSize = orig.diskSize;
58       numHeads = orig.numHeads;
59       numSecspTrack = orig.numSecspTrack;
60       canDeleteMyDisk = orig.canDeleteMyDisk;
61       device = orig.device;
62       state = orig.state;
63 
64       myDisk = new DiskIO;
65       if (myDisk == NULL) {
66          cerr << "Unable to allocate memory in BasicMBRData copy constructor! Terminating!\n";
67          exit(1);
68       } // if
69       if (orig.myDisk != NULL)
70          myDisk->OpenForRead(orig.myDisk->GetName());
71 
72       for (i = 0; i < MAX_MBR_PARTS; i++) {
73          partitions[i] = orig.partitions[i];
74       } // for
75    } // if
76 } // BasicMBRData copy constructor
77 
BasicMBRData(string filename)78 BasicMBRData::BasicMBRData(string filename) {
79    blockSize = SECTOR_SIZE;
80    diskSize = 0;
81    device = filename;
82    state = invalid;
83    numHeads = MAX_HEADS;
84    numSecspTrack = MAX_SECSPERTRACK;
85    myDisk = NULL;
86    canDeleteMyDisk = 0;
87 //   memset(&EbrLocations, 0, MAX_MBR_PARTS * sizeof(uint32_t));
88 
89    // Try to read the specified partition table, but if it fails....
90    if (!ReadMBRData(filename)) {
91       EmptyMBR();
92       device = "";
93    } // if
94 } // BasicMBRData(string filename) constructor
95 
96 // Free space used by myDisk only if that's OK -- sometimes it will be
97 // copied from an outside source, in which case that source should handle
98 // it!
~BasicMBRData(void)99 BasicMBRData::~BasicMBRData(void) {
100    if (canDeleteMyDisk)
101       delete myDisk;
102 } // BasicMBRData destructor
103 
104 // Assignment operator -- copy entire set of MBR data.
operator =(const BasicMBRData & orig)105 BasicMBRData & BasicMBRData::operator=(const BasicMBRData & orig) {
106    int i;
107 
108    if (&orig != this) {
109       memcpy(code, orig.code, 440);
110       diskSignature = orig.diskSignature;
111       nulls = orig.nulls;
112       MBRSignature = orig.MBRSignature;
113       blockSize = orig.blockSize;
114       diskSize = orig.diskSize;
115       numHeads = orig.numHeads;
116       numSecspTrack = orig.numSecspTrack;
117       canDeleteMyDisk = orig.canDeleteMyDisk;
118       device = orig.device;
119       state = orig.state;
120 
121       myDisk = new DiskIO;
122       if (myDisk == NULL) {
123          cerr << "Unable to allocate memory in BasicMBRData::operator=()! Terminating!\n";
124          exit(1);
125       } // if
126       if (orig.myDisk != NULL)
127          myDisk->OpenForRead(orig.myDisk->GetName());
128 
129       for (i = 0; i < MAX_MBR_PARTS; i++) {
130          partitions[i] = orig.partitions[i];
131       } // for
132    } // if
133    return *this;
134 } // BasicMBRData::operator=()
135 
136 /**********************
137  *                    *
138  * Disk I/O functions *
139  *                    *
140  **********************/
141 
142 // Read data from MBR. Returns 1 if read was successful (even if the
143 // data isn't a valid MBR), 0 if the read failed.
ReadMBRData(const string & deviceFilename)144 int BasicMBRData::ReadMBRData(const string & deviceFilename) {
145    int allOK = 1;
146 
147    if (myDisk == NULL) {
148       myDisk = new DiskIO;
149       if (myDisk == NULL) {
150          cerr << "Unable to allocate memory in BasicMBRData::ReadMBRData()! Terminating!\n";
151          exit(1);
152       } // if
153       canDeleteMyDisk = 1;
154    } // if
155    if (myDisk->OpenForRead(deviceFilename)) {
156       allOK = ReadMBRData(myDisk);
157    } else {
158       allOK = 0;
159    } // if
160 
161    if (allOK)
162       device = deviceFilename;
163 
164    return allOK;
165 } // BasicMBRData::ReadMBRData(const string & deviceFilename)
166 
167 // Read data from MBR. If checkBlockSize == 1 (the default), the block
168 // size is checked; otherwise it's set to the default (512 bytes).
169 // Note that any extended partition(s) present will be omitted from
170 // in the partitions[] array; these partitions must be re-created when
171 // the partition table is saved in MBR format.
ReadMBRData(DiskIO * theDisk,int checkBlockSize)172 int BasicMBRData::ReadMBRData(DiskIO * theDisk, int checkBlockSize) {
173    int allOK = 1, i, logicalNum = 3;
174    int err = 1;
175    TempMBR tempMBR;
176 
177    if ((myDisk != NULL) && (myDisk != theDisk) && (canDeleteMyDisk)) {
178       delete myDisk;
179       canDeleteMyDisk = 0;
180    } // if
181 
182    myDisk = theDisk;
183 
184    // Empty existing MBR data, including the logical partitions...
185    EmptyMBR(0);
186 
187    if (myDisk->Seek(0))
188      if (myDisk->Read(&tempMBR, 512))
189         err = 0;
190    if (err) {
191       cerr << "Problem reading disk in BasicMBRData::ReadMBRData()!\n";
192    } else {
193       for (i = 0; i < 440; i++)
194          code[i] = tempMBR.code[i];
195       diskSignature = tempMBR.diskSignature;
196       nulls = tempMBR.nulls;
197       for (i = 0; i < 4; i++) {
198          partitions[i] = tempMBR.partitions[i];
199          if (partitions[i].GetLengthLBA() > 0)
200             partitions[i].SetInclusion(PRIMARY);
201       } // for i... (reading all four partitions)
202       MBRSignature = tempMBR.MBRSignature;
203       ReadCHSGeom();
204 
205       // Reverse the byte order, if necessary
206       if (IsLittleEndian() == 0) {
207          ReverseBytes(&diskSignature, 4);
208          ReverseBytes(&nulls, 2);
209          ReverseBytes(&MBRSignature, 2);
210          for (i = 0; i < 4; i++) {
211             partitions[i].ReverseByteOrder();
212          } // for
213       } // if
214 
215       if (MBRSignature != MBR_SIGNATURE) {
216          allOK = 0;
217          state = invalid;
218       } // if
219 
220       // Find disk size
221       diskSize = myDisk->DiskSize(&err);
222 
223       // Find block size
224       if (checkBlockSize) {
225          blockSize = myDisk->GetBlockSize();
226       } // if (checkBlockSize)
227 
228       // Load logical partition data, if any is found....
229       if (allOK) {
230          for (i = 0; i < 4; i++) {
231             if ((partitions[i].GetType() == 0x05) || (partitions[i].GetType() == 0x0f)
232                 || (partitions[i].GetType() == 0x85)) {
233                // Found it, so call a function to load everything from them....
234                logicalNum = ReadLogicalParts(partitions[i].GetStartLBA(), abs(logicalNum) + 1);
235                if (logicalNum < 0) {
236                   cerr << "Error reading logical partitions! List may be truncated!\n";
237                } // if maxLogicals valid
238                DeletePartition(i);
239             } // if primary partition is extended
240          } // for primary partition loop
241          if (allOK) { // Loaded logicals OK
242             state = mbr;
243          } else {
244             state = invalid;
245          } // if
246       } // if
247 
248       // Check to see if it's in GPT format....
249       if (allOK) {
250          for (i = 0; i < 4; i++) {
251             if (partitions[i].GetType() == UINT8_C(0xEE)) {
252                state = gpt;
253             } // if
254          } // for
255       } // if
256 
257       // If there's an EFI GPT partition, look for other partition types,
258       // to flag as hybrid
259       if (state == gpt) {
260          for (i = 0 ; i < 4; i++) {
261             if ((partitions[i].GetType() != UINT8_C(0xEE)) &&
262                 (partitions[i].GetType() != UINT8_C(0x00)))
263                state = hybrid;
264             if (logicalNum != 3)
265                cerr << "Warning! MBR Logical partitions found on a hybrid MBR disk! This is an\n"
266                     << "EXTREMELY dangerous configuration!\n\a";
267          } // for
268       } // if (hybrid detection code)
269    } // no initial error
270    return allOK;
271 } // BasicMBRData::ReadMBRData(DiskIO * theDisk, int checkBlockSize)
272 
273 // This is a function to read all the logical partitions, following the
274 // logical partition linked list from the disk and storing the basic data in the
275 // partitions[] array. Returns last index to partitions[] used, or -1 times the
276 // that index if there was a problem. (Some problems can leave valid logical
277 // partition data.)
278 // Parameters:
279 // extendedStart = LBA of the start of the extended partition
280 // partNum = number of first partition in extended partition (normally 4).
ReadLogicalParts(uint64_t extendedStart,int partNum)281 int BasicMBRData::ReadLogicalParts(uint64_t extendedStart, int partNum) {
282    struct TempMBR ebr;
283    int i, another = 1, allOK = 1;
284    uint8_t ebrType;
285    uint64_t offset;
286    uint64_t EbrLocations[MAX_MBR_PARTS];
287 
288    offset = extendedStart;
289    memset(&EbrLocations, 0, MAX_MBR_PARTS * sizeof(uint64_t));
290    while (another && (partNum < MAX_MBR_PARTS) && (partNum >= 0) && (allOK > 0)) {
291       for (i = 0; i < MAX_MBR_PARTS; i++) {
292          if (EbrLocations[i] == offset) { // already read this one; infinite logical partition loop!
293             cerr << "Logical partition infinite loop detected! This is being corrected.\n";
294             allOK = -1;
295             if(partNum > 0) //don't go negative
296                 partNum -= 1;
297          } // if
298       } // for
299       EbrLocations[partNum] = offset;
300       if (myDisk->Seek(offset) == 0) { // seek to EBR record
301          cerr << "Unable to seek to " << offset << "! Aborting!\n";
302          allOK = -1;
303       }
304       if (myDisk->Read(&ebr, 512) != 512) { // Load the data....
305          cerr << "Error seeking to or reading logical partition data from " << offset
306               << "!\nSome logical partitions may be missing!\n";
307          allOK = -1;
308       } else if (IsLittleEndian() != 1) { // Reverse byte ordering of some data....
309          ReverseBytes(&ebr.MBRSignature, 2);
310          ReverseBytes(&ebr.partitions[0].firstLBA, 4);
311          ReverseBytes(&ebr.partitions[0].lengthLBA, 4);
312          ReverseBytes(&ebr.partitions[1].firstLBA, 4);
313          ReverseBytes(&ebr.partitions[1].lengthLBA, 4);
314       } // if/else/if
315 
316       if (ebr.MBRSignature != MBR_SIGNATURE) {
317          allOK = -1;
318          cerr << "EBR signature for logical partition invalid; read 0x";
319          cerr.fill('0');
320          cerr.width(4);
321          cerr.setf(ios::uppercase);
322          cerr << hex << ebr.MBRSignature << ", but should be 0x";
323          cerr.width(4);
324          cerr << MBR_SIGNATURE << dec << "\n";
325          cerr.fill(' ');
326       } // if
327 
328       if ((partNum >= 0) && (partNum < MAX_MBR_PARTS) && (allOK > 0)) {
329          // Sometimes an EBR points directly to another EBR, rather than defining
330          // a logical partition and then pointing to another EBR. Thus, we skip
331          // the logical partition when this is the case....
332          ebrType = ebr.partitions[0].partitionType;
333          if ((ebrType == 0x05) || (ebrType == 0x0f) || (ebrType == 0x85)) {
334             cout << "EBR points to an EBR!\n";
335             offset = extendedStart + ebr.partitions[0].firstLBA;
336          } else {
337             // Copy over the basic data....
338             partitions[partNum] = ebr.partitions[0];
339             // Adjust the start LBA, since it's encoded strangely....
340             partitions[partNum].SetStartLBA(ebr.partitions[0].firstLBA + offset);
341             partitions[partNum].SetInclusion(LOGICAL);
342 
343             // Find the next partition (if there is one)
344             if ((ebr.partitions[1].firstLBA != UINT32_C(0)) && (partNum < (MAX_MBR_PARTS - 1))) {
345                offset = extendedStart + ebr.partitions[1].firstLBA;
346                partNum++;
347             } else {
348                another = 0;
349             } // if another partition
350          } // if/else
351       } // if
352    } // while()
353    return (partNum * allOK);
354 } // BasicMBRData::ReadLogicalPart()
355 
356 // Write the MBR data to the default defined device. This writes both the
357 // MBR itself and any defined logical partitions, provided there's an
358 // MBR extended partition.
WriteMBRData(void)359 int BasicMBRData::WriteMBRData(void) {
360    int allOK = 1;
361 
362    if (myDisk != NULL) {
363       if (myDisk->OpenForWrite() != 0) {
364          allOK = WriteMBRData(myDisk);
365          cout << "Done writing data!\n";
366       } else {
367          allOK = 0;
368       } // if/else
369       myDisk->Close();
370    } else allOK = 0;
371    return allOK;
372 } // BasicMBRData::WriteMBRData(void)
373 
374 // Save the MBR data to a file. This writes both the
375 // MBR itself and any defined logical partitions.
WriteMBRData(DiskIO * theDisk)376 int BasicMBRData::WriteMBRData(DiskIO *theDisk) {
377    int i, j, partNum, next, allOK = 1, moreLogicals = 0;
378    uint64_t extFirstLBA = 0;
379    uint64_t writeEbrTo; // 64-bit because we support extended in 2-4TiB range
380    TempMBR tempMBR;
381 
382    allOK = CreateExtended();
383    if (allOK) {
384       // First write the main MBR data structure....
385       memcpy(tempMBR.code, code, 440);
386       tempMBR.diskSignature = diskSignature;
387       tempMBR.nulls = nulls;
388       tempMBR.MBRSignature = MBRSignature;
389       for (i = 0; i < 4; i++) {
390          partitions[i].StoreInStruct(&tempMBR.partitions[i]);
391          if (partitions[i].GetType() == 0x0f) {
392             extFirstLBA = partitions[i].GetStartLBA();
393             moreLogicals = 1;
394          } // if
395       } // for i...
396    } // if
397    allOK = allOK && WriteMBRData(tempMBR, theDisk, 0);
398 
399    // Set up tempMBR with some constant data for logical partitions...
400    tempMBR.diskSignature = 0;
401    for (i = 2; i < 4; i++) {
402       tempMBR.partitions[i].firstLBA = tempMBR.partitions[i].lengthLBA = 0;
403       tempMBR.partitions[i].partitionType = 0x00;
404       for (j = 0; j < 3; j++) {
405          tempMBR.partitions[i].firstSector[j] = 0;
406          tempMBR.partitions[i].lastSector[j] = 0;
407       } // for j
408    } // for i
409 
410    partNum = FindNextInUse(4);
411    writeEbrTo = (uint64_t) extFirstLBA;
412    // Write logicals...
413    while (allOK && moreLogicals && (partNum < MAX_MBR_PARTS) && (partNum >= 0)) {
414       partitions[partNum].StoreInStruct(&tempMBR.partitions[0]);
415       tempMBR.partitions[0].firstLBA = 1;
416       // tempMBR.partitions[1] points to next EBR or terminates EBR linked list...
417       next = FindNextInUse(partNum + 1);
418       if ((next < MAX_MBR_PARTS) && (next > 0) && (partitions[next].GetStartLBA() > 0)) {
419          tempMBR.partitions[1].partitionType = 0x0f;
420          tempMBR.partitions[1].firstLBA = (uint32_t) (partitions[next].GetStartLBA() - extFirstLBA - 1);
421          tempMBR.partitions[1].lengthLBA = (uint32_t) (partitions[next].GetLengthLBA() + 1);
422          LBAtoCHS((uint64_t) tempMBR.partitions[1].firstLBA,
423                   (uint8_t *) &tempMBR.partitions[1].firstSector);
424          LBAtoCHS(tempMBR.partitions[1].lengthLBA - extFirstLBA,
425                   (uint8_t *) &tempMBR.partitions[1].lastSector);
426       } else {
427          tempMBR.partitions[1].partitionType = 0x00;
428          tempMBR.partitions[1].firstLBA = 0;
429          tempMBR.partitions[1].lengthLBA = 0;
430          moreLogicals = 0;
431       } // if/else
432       allOK = WriteMBRData(tempMBR, theDisk, writeEbrTo);
433       writeEbrTo = (uint64_t) tempMBR.partitions[1].firstLBA + (uint64_t) extFirstLBA;
434       partNum = next;
435    } // while
436    DeleteExtendedParts();
437    return allOK;
438 } // BasicMBRData::WriteMBRData(DiskIO *theDisk)
439 
WriteMBRData(const string & deviceFilename)440 int BasicMBRData::WriteMBRData(const string & deviceFilename) {
441    device = deviceFilename;
442    return WriteMBRData();
443 } // BasicMBRData::WriteMBRData(const string & deviceFilename)
444 
445 // Write a single MBR record to the specified sector. Used by the like-named
446 // function to write both the MBR and multiple EBR (for logical partition)
447 // records.
448 // Returns 1 on success, 0 on failure
WriteMBRData(struct TempMBR & mbr,DiskIO * theDisk,uint64_t sector)449 int BasicMBRData::WriteMBRData(struct TempMBR & mbr, DiskIO *theDisk, uint64_t sector) {
450    int i, allOK;
451 
452    // Reverse the byte order, if necessary
453    if (IsLittleEndian() == 0) {
454       ReverseBytes(&mbr.diskSignature, 4);
455       ReverseBytes(&mbr.nulls, 2);
456       ReverseBytes(&mbr.MBRSignature, 2);
457       for (i = 0; i < 4; i++) {
458          ReverseBytes(&mbr.partitions[i].firstLBA, 4);
459          ReverseBytes(&mbr.partitions[i].lengthLBA, 4);
460       } // for
461    } // if
462 
463    // Now write the data structure...
464    allOK = theDisk->OpenForWrite();
465    if (allOK && theDisk->Seek(sector)) {
466       if (theDisk->Write(&mbr, 512) != 512) {
467          allOK = 0;
468          cerr << "Error " << errno << " when saving MBR!\n";
469       } // if
470    } else {
471       allOK = 0;
472       cerr << "Error " << errno << " when seeking to MBR to write it!\n";
473    } // if/else
474    theDisk->Close();
475 
476    // Reverse the byte order back, if necessary
477    if (IsLittleEndian() == 0) {
478       ReverseBytes(&mbr.diskSignature, 4);
479       ReverseBytes(&mbr.nulls, 2);
480       ReverseBytes(&mbr.MBRSignature, 2);
481       for (i = 0; i < 4; i++) {
482          ReverseBytes(&mbr.partitions[i].firstLBA, 4);
483          ReverseBytes(&mbr.partitions[i].lengthLBA, 4);
484       } // for
485    }// if
486    return allOK;
487 } // BasicMBRData::WriteMBRData(uint64_t sector)
488 
489 // Set a new disk device; used in copying one disk's partition
490 // table to another disk.
SetDisk(DiskIO * theDisk)491 void BasicMBRData::SetDisk(DiskIO *theDisk) {
492    int err;
493 
494    myDisk = theDisk;
495    diskSize = theDisk->DiskSize(&err);
496    canDeleteMyDisk = 0;
497    ReadCHSGeom();
498 } // BasicMBRData::SetDisk()
499 
500 /********************************************
501  *                                          *
502  * Functions that display data for the user *
503  *                                          *
504  ********************************************/
505 
506 // Show the MBR data to the user, up to the specified maximum number
507 // of partitions....
DisplayMBRData(void)508 void BasicMBRData::DisplayMBRData(void) {
509    int i;
510 
511    cout << "\nDisk size is " << diskSize << " sectors ("
512         << BytesToIeee(diskSize, blockSize) << ")\n";
513    cout << "MBR disk identifier: 0x";
514    cout.width(8);
515    cout.fill('0');
516    cout.setf(ios::uppercase);
517    cout << hex << diskSignature << dec << "\n";
518    cout << "MBR partitions:\n\n";
519    if ((state == gpt) || (state == hybrid)) {
520       cout << "Number  Boot  Start Sector   End Sector   Status      Code\n";
521    } else {
522       cout << "                                                   Can Be   Can Be\n";
523       cout << "Number  Boot  Start Sector   End Sector   Status   Logical  Primary   Code\n";
524       UpdateCanBeLogical();
525    } //
526    for (i = 0; i < MAX_MBR_PARTS; i++) {
527       if (partitions[i].GetLengthLBA() != 0) {
528          cout.fill(' ');
529          cout.width(4);
530          cout << i + 1 << "      ";
531          partitions[i].ShowData((state == gpt) || (state == hybrid));
532       } // if
533       cout.fill(' ');
534    } // for
535 } // BasicMBRData::DisplayMBRData()
536 
537 // Displays the state, as a word, on stdout. Used for debugging & to
538 // tell the user about the MBR state when the program launches....
ShowState(void)539 void BasicMBRData::ShowState(void) {
540    switch (state) {
541       case invalid:
542          cout << "  MBR: not present\n";
543          break;
544       case gpt:
545          cout << "  MBR: protective\n";
546          break;
547       case hybrid:
548          cout << "  MBR: hybrid\n";
549          break;
550       case mbr:
551          cout << "  MBR: MBR only\n";
552          break;
553       default:
554          cout << "\a  MBR: unknown -- bug!\n";
555          break;
556    } // switch
557 } // BasicMBRData::ShowState()
558 
559 /************************
560  *                      *
561  * GPT Checks and fixes *
562  *                      *
563  ************************/
564 
565 // Perform a very rudimentary check for GPT data on the disk; searches for
566 // the GPT signature in the main and backup metadata areas.
567 // Returns 0 if GPT data not found, 1 if main data only is found, 2 if
568 // backup only is found, 3 if both main and backup data are found, and
569 // -1 if a disk error occurred.
CheckForGPT(void)570 int BasicMBRData::CheckForGPT(void) {
571    int retval = 0, err;
572    char signature1[9], signature2[9];
573 
574    if (myDisk != NULL) {
575       if (myDisk->OpenForRead() != 0) {
576          if (myDisk->Seek(1)) {
577             myDisk->Read(signature1, 8);
578             signature1[8] = '\0';
579          } else retval = -1;
580          if (myDisk->Seek(myDisk->DiskSize(&err) - 1)) {
581             myDisk->Read(signature2, 8);
582             signature2[8] = '\0';
583          } else retval = -1;
584          if ((retval >= 0) && (strcmp(signature1, "EFI PART") == 0))
585             retval += 1;
586          if ((retval >= 0) && (strcmp(signature2, "EFI PART") == 0))
587             retval += 2;
588       } else {
589          retval = -1;
590       } // if/else
591       myDisk->Close();
592    } else retval = -1;
593    return retval;
594 } // BasicMBRData::CheckForGPT()
595 
596 // Blanks the 2nd (sector #1, numbered from 0) and last sectors of the disk,
597 // but only if GPT data are verified on the disk, and only for the sector(s)
598 // with GPT signatures.
599 // Returns 1 if operation completes successfully, 0 if not (returns 1 if
600 // no GPT data are found on the disk).
BlankGPTData(void)601 int BasicMBRData::BlankGPTData(void) {
602    int allOK = 1, err;
603    uint8_t blank[512];
604 
605    memset(blank, 0, 512);
606    switch (CheckForGPT()) {
607       case -1:
608          allOK = 0;
609          break;
610       case 0:
611          break;
612       case 1:
613          if ((myDisk != NULL) && (myDisk->OpenForWrite())) {
614             if (!((myDisk->Seek(1)) && (myDisk->Write(blank, 512) == 512)))
615                allOK = 0;
616             myDisk->Close();
617          } else allOK = 0;
618          break;
619       case 2:
620          if ((myDisk != NULL) && (myDisk->OpenForWrite())) {
621             if (!((myDisk->Seek(myDisk->DiskSize(&err) - 1)) &&
622                (myDisk->Write(blank, 512) == 512)))
623                allOK = 0;
624             myDisk->Close();
625          } else allOK = 0;
626          break;
627       case 3:
628          if ((myDisk != NULL) && (myDisk->OpenForWrite())) {
629             if (!((myDisk->Seek(1)) && (myDisk->Write(blank, 512) == 512)))
630                allOK = 0;
631             if (!((myDisk->Seek(myDisk->DiskSize(&err) - 1)) &&
632                 (myDisk->Write(blank, 512) == 512)))
633                 allOK = 0;
634             myDisk->Close();
635          } else allOK = 0;
636          break;
637       default:
638          break;
639    } // switch()
640    return allOK;
641 } // BasicMBRData::BlankGPTData
642 
643 /*********************************************************************
644  *                                                                   *
645  * Functions that set or get disk metadata (CHS geometry, disk size, *
646  * etc.)                                                             *
647  *                                                                   *
648  *********************************************************************/
649 
650 // Read the CHS geometry using OS calls, or if that fails, set to
651 // the most common value for big disks (255 heads, 63 sectors per
652 // track, & however many cylinders that computes to).
ReadCHSGeom(void)653 void BasicMBRData::ReadCHSGeom(void) {
654    int err;
655 
656    numHeads = myDisk->GetNumHeads();
657    numSecspTrack = myDisk->GetNumSecsPerTrack();
658    diskSize = myDisk->DiskSize(&err);
659    blockSize = myDisk->GetBlockSize();
660    partitions[0].SetGeometry(numHeads, numSecspTrack, diskSize, blockSize);
661 } // BasicMBRData::ReadCHSGeom()
662 
663 // Find the low and high used partition numbers (numbered from 0).
664 // Return value is the number of partitions found. Note that the
665 // *low and *high values are both set to 0 when no partitions
666 // are found, as well as when a single partition in the first
667 // position exists. Thus, the return value is the only way to
668 // tell when no partitions exist.
GetPartRange(uint32_t * low,uint32_t * high)669 int BasicMBRData::GetPartRange(uint32_t *low, uint32_t *high) {
670    uint32_t i;
671    int numFound = 0;
672 
673    *low = MAX_MBR_PARTS + 1; // code for "not found"
674    *high = 0;
675    for (i = 0; i < MAX_MBR_PARTS; i++) {
676       if (partitions[i].GetStartLBA() != UINT32_C(0)) { // it exists
677          *high = i; // since we're counting up, set the high value
678          // Set the low value only if it's not yet found...
679          if (*low == (MAX_MBR_PARTS + 1))
680             *low = i;
681          numFound++;
682       } // if
683    } // for
684 
685    // Above will leave *low pointing to its "not found" value if no partitions
686    // are defined, so reset to 0 if this is the case....
687    if (*low == (MAX_MBR_PARTS + 1))
688       *low = 0;
689    return numFound;
690 } // GPTData::GetPartRange()
691 
692 // Converts 64-bit LBA value to MBR-style CHS value. Returns 1 if conversion
693 // was within the range that can be expressed by CHS (including 0, for an
694 // empty partition), 0 if the value is outside that range, and -1 if chs is
695 // invalid.
LBAtoCHS(uint64_t lba,uint8_t * chs)696 int BasicMBRData::LBAtoCHS(uint64_t lba, uint8_t * chs) {
697    uint64_t cylinder, head, sector; // all numbered from 0
698    uint64_t remainder;
699    int retval = 1;
700    int done = 0;
701 
702    if (chs != NULL) {
703       // Special case: In case of 0 LBA value, zero out CHS values....
704       if (lba == 0) {
705          chs[0] = chs[1] = chs[2] = UINT8_C(0);
706          done = 1;
707       } // if
708       // If LBA value is too large for CHS, max out CHS values....
709       if ((!done) && (lba >= ((uint64_t) numHeads * numSecspTrack * MAX_CYLINDERS))) {
710          chs[0] = 254;
711          chs[1] = chs[2] = 255;
712          done = 1;
713          retval = 0;
714       } // if
715       // If neither of the above applies, compute CHS values....
716       if (!done) {
717          cylinder = lba / (uint64_t) (numHeads * numSecspTrack);
718          remainder = lba - (cylinder * numHeads * numSecspTrack);
719          head = remainder / numSecspTrack;
720          remainder -= head * numSecspTrack;
721          sector = remainder;
722          if (head < numHeads)
723             chs[0] = (uint8_t) head;
724          else
725             retval = 0;
726          if (sector < numSecspTrack) {
727             chs[1] = (uint8_t) ((sector + 1) + (cylinder >> 8) * 64);
728             chs[2] = (uint8_t) (cylinder & UINT64_C(0xFF));
729          } else {
730             retval = 0;
731          } // if/else
732       } // if value is expressible and non-0
733    } else { // Invalid (NULL) chs pointer
734       retval = -1;
735    } // if CHS pointer valid
736    return (retval);
737 } // BasicMBRData::LBAtoCHS()
738 
739 // Look for overlapping partitions. Also looks for a couple of non-error
740 // conditions that the user should be told about.
741 // Returns the number of problems found
FindOverlaps(void)742 int BasicMBRData::FindOverlaps(void) {
743    int i, j, numProbs = 0, numEE = 0, ProtectiveOnOne = 0;
744 
745    for (i = 0; i < MAX_MBR_PARTS; i++) {
746       for (j = i + 1; j < MAX_MBR_PARTS; j++) {
747          if ((partitions[i].GetInclusion() != NONE) && (partitions[j].GetInclusion() != NONE) &&
748              (partitions[i].DoTheyOverlap(partitions[j]))) {
749             numProbs++;
750             cout << "\nProblem: MBR partitions " << i + 1 << " and " << j + 1
751                  << " overlap!\n";
752          } // if
753       } // for (j...)
754       if (partitions[i].GetType() == 0xEE) {
755          numEE++;
756          if (partitions[i].GetStartLBA() == 1)
757             ProtectiveOnOne = 1;
758       } // if
759    } // for (i...)
760 
761    if (numEE > 1)
762       cout << "\nCaution: More than one 0xEE MBR partition found. This can cause problems\n"
763            << "in some OSes.\n";
764    if (!ProtectiveOnOne && (numEE > 0))
765       cout << "\nWarning: 0xEE partition doesn't start on sector 1. This can cause "
766            << "problems\nin some OSes.\n";
767 
768    return numProbs;
769 } // BasicMBRData::FindOverlaps()
770 
771 // Returns the number of primary partitions, including the extended partition
772 // required to hold any logical partitions found.
NumPrimaries(void)773 int BasicMBRData::NumPrimaries(void) {
774    int i, numPrimaries = 0, logicalsFound = 0;
775 
776    for (i = 0; i < MAX_MBR_PARTS; i++) {
777       if (partitions[i].GetLengthLBA() > 0) {
778          if (partitions[i].GetInclusion() == PRIMARY)
779             numPrimaries++;
780          if (partitions[i].GetInclusion() == LOGICAL)
781             logicalsFound = 1;
782       } // if
783    } // for
784    return (numPrimaries + logicalsFound);
785 } // BasicMBRData::NumPrimaries()
786 
787 // Returns the number of logical partitions.
NumLogicals(void)788 int BasicMBRData::NumLogicals(void) {
789    int i, numLogicals = 0;
790 
791    for (i = 0; i < MAX_MBR_PARTS; i++) {
792       if (partitions[i].GetInclusion() == LOGICAL)
793          numLogicals++;
794    } // for
795    return numLogicals;
796 } // BasicMBRData::NumLogicals()
797 
798 // Returns the number of partitions (primaries plus logicals), NOT including
799 // the extended partition required to house the logicals.
CountParts(void)800 int BasicMBRData::CountParts(void) {
801    int i, num = 0;
802 
803    for (i = 0; i < MAX_MBR_PARTS; i++) {
804       if ((partitions[i].GetInclusion() == LOGICAL) ||
805           (partitions[i].GetInclusion() == PRIMARY))
806          num++;
807    } // for
808    return num;
809 } // BasicMBRData::CountParts()
810 
811 // Updates the canBeLogical and canBePrimary flags for all the partitions.
UpdateCanBeLogical(void)812 void BasicMBRData::UpdateCanBeLogical(void) {
813    int i, j, sectorBefore, numPrimaries, numLogicals, usedAsEBR;
814    uint64_t firstLogical, lastLogical, lStart, pStart;
815 
816    numPrimaries = NumPrimaries();
817    numLogicals = NumLogicals();
818    firstLogical = FirstLogicalLBA() - 1;
819    lastLogical = LastLogicalLBA();
820    for (i = 0; i < MAX_MBR_PARTS; i++) {
821       usedAsEBR = (SectorUsedAs(partitions[i].GetLastLBA()) == EBR);
822       if (usedAsEBR) {
823          partitions[i].SetCanBeLogical(0);
824          partitions[i].SetCanBePrimary(0);
825       } else if (partitions[i].GetLengthLBA() > 0) {
826          // First determine if it can be logical....
827          sectorBefore = SectorUsedAs(partitions[i].GetStartLBA() - 1);
828          lStart = partitions[i].GetStartLBA(); // start of potential logical part.
829          if ((lastLogical > 0) &&
830              ((sectorBefore == EBR) || (sectorBefore == NONE))) {
831             // Assume it can be logical, then search for primaries that make it
832             // not work and, if found, flag appropriately.
833             partitions[i].SetCanBeLogical(1);
834             for (j = 0; j < MAX_MBR_PARTS; j++) {
835                if ((i != j) && (partitions[j].GetInclusion() == PRIMARY)) {
836                   pStart = partitions[j].GetStartLBA();
837                   if (((pStart < lStart) && (firstLogical < pStart)) ||
838                       ((pStart > lStart) && (firstLogical > pStart))) {
839                      partitions[i].SetCanBeLogical(0);
840                   } // if/else
841                } // if
842             } // for
843          } else {
844             if ((sectorBefore != EBR) && (sectorBefore != NONE))
845                partitions[i].SetCanBeLogical(0);
846             else
847                partitions[i].SetCanBeLogical(lastLogical == 0); // can be logical only if no logicals already
848          } // if/else
849          // Now determine if it can be primary. Start by assuming it can be...
850          partitions[i].SetCanBePrimary(1);
851          if ((numPrimaries >= 4) && (partitions[i].GetInclusion() != PRIMARY)) {
852             partitions[i].SetCanBePrimary(0);
853             if ((partitions[i].GetInclusion() == LOGICAL) && (numLogicals == 1) &&
854                 (numPrimaries == 4))
855                partitions[i].SetCanBePrimary(1);
856          } // if
857          if ((partitions[i].GetStartLBA() > (firstLogical + 1)) &&
858              (partitions[i].GetLastLBA() < lastLogical))
859             partitions[i].SetCanBePrimary(0);
860       } // else if
861    } // for
862 } // BasicMBRData::UpdateCanBeLogical()
863 
864 // Returns the first sector occupied by any logical partition. Note that
865 // this does NOT include the logical partition's EBR! Returns UINT32_MAX
866 // if there are no logical partitions defined.
FirstLogicalLBA(void)867 uint64_t BasicMBRData::FirstLogicalLBA(void) {
868    int i;
869    uint64_t firstFound = UINT32_MAX;
870 
871    for (i = 0; i < MAX_MBR_PARTS; i++) {
872       if ((partitions[i].GetInclusion() == LOGICAL) &&
873           (partitions[i].GetStartLBA() < firstFound)) {
874          firstFound = partitions[i].GetStartLBA();
875       } // if
876    } // for
877    return firstFound;
878 } // BasicMBRData::FirstLogicalLBA()
879 
880 // Returns the last sector occupied by any logical partition, or 0 if
881 // there are no logical partitions defined.
LastLogicalLBA(void)882 uint64_t BasicMBRData::LastLogicalLBA(void) {
883    int i;
884    uint64_t lastFound = 0;
885 
886    for (i = 0; i < MAX_MBR_PARTS; i++) {
887       if ((partitions[i].GetInclusion() == LOGICAL) &&
888           (partitions[i].GetLastLBA() > lastFound))
889          lastFound = partitions[i].GetLastLBA();
890    } // for
891    return lastFound;
892 } // BasicMBRData::LastLogicalLBA()
893 
894 // Returns 1 if logical partitions are contiguous (have no primaries
895 // in their midst), or 0 if one or more primaries exist between
896 // logicals.
AreLogicalsContiguous(void)897 int BasicMBRData::AreLogicalsContiguous(void) {
898    int allOK = 1, i = 0;
899    uint64_t firstLogical, lastLogical;
900 
901    firstLogical = FirstLogicalLBA() - 1; // subtract 1 for EBR
902    lastLogical = LastLogicalLBA();
903    if (lastLogical > 0) {
904       do {
905          if ((partitions[i].GetInclusion() == PRIMARY) &&
906              (partitions[i].GetStartLBA() >= firstLogical) &&
907              (partitions[i].GetStartLBA() <= lastLogical)) {
908             allOK = 0;
909          } // if
910          i++;
911       } while ((i < MAX_MBR_PARTS) && allOK);
912    } // if
913    return allOK;
914 } // BasicMBRData::AreLogicalsContiguous()
915 
916 // Returns 1 if all partitions fit on the disk, given its size; 0 if any
917 // partition is too big.
DoTheyFit(void)918 int BasicMBRData::DoTheyFit(void) {
919    int i, allOK = 1;
920 
921    for (i = 0; i < MAX_MBR_PARTS; i++) {
922       if ((partitions[i].GetStartLBA() > diskSize) || (partitions[i].GetLastLBA() > diskSize)) {
923          allOK = 0;
924       } // if
925    } // for
926    return allOK;
927 } // BasicMBRData::DoTheyFit(void)
928 
929 // Returns 1 if there's at least one free sector immediately preceding
930 // all partitions flagged as logical; 0 if any logical partition lacks
931 // this space.
SpaceBeforeAllLogicals(void)932 int BasicMBRData::SpaceBeforeAllLogicals(void) {
933    int i = 0, allOK = 1;
934 
935    do {
936       if ((partitions[i].GetStartLBA() > 0) && (partitions[i].GetInclusion() == LOGICAL)) {
937          allOK = allOK && (SectorUsedAs(partitions[i].GetStartLBA() - 1) == EBR);
938       } // if
939       i++;
940    } while (allOK && (i < MAX_MBR_PARTS));
941    return allOK;
942 } // BasicMBRData::SpaceBeforeAllLogicals()
943 
944 // Returns 1 if the partitions describe a legal layout -- all logicals
945 // are contiguous and have at least one preceding empty sector,
946 // the number of primaries is under 4 (or under 3 if there are any
947 // logicals), there are no overlapping partitions, etc.
948 // Does NOT assume that primaries are numbered 1-4; uses the
949 // IsItPrimary() function of the MBRPart class to determine
950 // primary status. Also does NOT consider partition order; there
951 // can be gaps and it will still be considered legal.
IsLegal(void)952 int BasicMBRData::IsLegal(void) {
953    int allOK = 1;
954 
955    allOK = (FindOverlaps() == 0);
956    allOK = (allOK && (NumPrimaries() <= 4));
957    allOK = (allOK && AreLogicalsContiguous());
958    allOK = (allOK && DoTheyFit());
959    allOK = (allOK && SpaceBeforeAllLogicals());
960    return allOK;
961 } // BasicMBRData::IsLegal()
962 
963 // Returns 1 if the 0xEE partition in the protective/hybrid MBR is marked as
964 // active/bootable.
IsEEActive(void)965 int BasicMBRData::IsEEActive(void) {
966    int i, IsActive = 0;
967 
968    for (i = 0; i < MAX_MBR_PARTS; i++) {
969       if ((partitions[i].GetStatus() & 0x80) && (partitions[i].GetType() == 0xEE))
970          IsActive = 1;
971    }
972    return IsActive;
973 } // BasicMBRData::IsEEActive()
974 
975 // Finds the next in-use partition, starting with start (will return start
976 // if it's in use). Returns -1 if no subsequent partition is in use.
FindNextInUse(int start)977 int BasicMBRData::FindNextInUse(int start) {
978    if (start >= MAX_MBR_PARTS)
979       start = -1;
980    while ((start < MAX_MBR_PARTS) && (start >= 0) && (partitions[start].GetInclusion() == NONE))
981       start++;
982    if ((start < 0) || (start >= MAX_MBR_PARTS))
983       start = -1;
984    return start;
985 } // BasicMBRData::FindFirstLogical();
986 
987 /*****************************************************
988  *                                                   *
989  * Functions to create, delete, or change partitions *
990  *                                                   *
991  *****************************************************/
992 
993 // Empty all data. Meant mainly for calling by constructors, but it's also
994 // used by the hybrid MBR functions in the GPTData class.
EmptyMBR(int clearBootloader)995 void BasicMBRData::EmptyMBR(int clearBootloader) {
996    int i;
997 
998    // Zero out the boot loader section, the disk signature, and the
999    // 2-byte nulls area only if requested to do so. (This is the
1000    // default.)
1001    if (clearBootloader == 1) {
1002       EmptyBootloader();
1003    } // if
1004 
1005    // Blank out the partitions
1006    for (i = 0; i < MAX_MBR_PARTS; i++) {
1007       partitions[i].Empty();
1008    } // for
1009    MBRSignature = MBR_SIGNATURE;
1010    state = mbr;
1011 } // BasicMBRData::EmptyMBR()
1012 
1013 // Blank out the boot loader area. Done with the initial MBR-to-GPT
1014 // conversion, since MBR boot loaders don't understand GPT, and so
1015 // need to be replaced....
EmptyBootloader(void)1016 void BasicMBRData::EmptyBootloader(void) {
1017    int i;
1018 
1019    for (i = 0; i < 440; i++)
1020       code[i] = 0;
1021    nulls = 0;
1022 } // BasicMBRData::EmptyBootloader
1023 
1024 // Create a partition of the specified number based on the passed
1025 // partition. This function does *NO* error checking, so it's possible
1026 // to seriously screw up a partition table using this function!
1027 // Note: This function should NOT be used to create the 0xEE partition
1028 // in a conventional GPT configuration, since that partition has
1029 // specific size requirements that this function won't handle. It may
1030 // be used for creating the 0xEE partition(s) in a hybrid MBR, though,
1031 // since those toss the rulebook away anyhow....
AddPart(int num,const MBRPart & newPart)1032 void BasicMBRData::AddPart(int num, const MBRPart& newPart) {
1033    partitions[num] = newPart;
1034 } // BasicMBRData::AddPart()
1035 
1036 // Create a partition of the specified number, starting LBA, and
1037 // length. This function does almost no error checking, so it's possible
1038 // to seriously screw up a partition table using this function!
1039 // Note: This function should NOT be used to create the 0xEE partition
1040 // in a conventional GPT configuration, since that partition has
1041 // specific size requirements that this function won't handle. It may
1042 // be used for creating the 0xEE partition(s) in a hybrid MBR, though,
1043 // since those toss the rulebook away anyhow....
MakePart(int num,uint64_t start,uint64_t length,int type,int bootable)1044 void BasicMBRData::MakePart(int num, uint64_t start, uint64_t length, int type, int bootable) {
1045    if ((num >= 0) && (num < MAX_MBR_PARTS) && (start <= UINT32_MAX) && (length <= UINT32_MAX)) {
1046       partitions[num].Empty();
1047       partitions[num].SetType(type);
1048       partitions[num].SetLocation(start, length);
1049       if (num < 4)
1050          partitions[num].SetInclusion(PRIMARY);
1051       else
1052          partitions[num].SetInclusion(LOGICAL);
1053       SetPartBootable(num, bootable);
1054    } // if valid partition number & size
1055 } // BasicMBRData::MakePart()
1056 
1057 // Set the partition's type code.
1058 // Returns 1 if successful, 0 if not (invalid partition number)
SetPartType(int num,int type)1059 int BasicMBRData::SetPartType(int num, int type) {
1060    int allOK = 1;
1061 
1062    if ((num >= 0) && (num < MAX_MBR_PARTS)) {
1063       if (partitions[num].GetLengthLBA() != UINT32_C(0)) {
1064          allOK = partitions[num].SetType(type);
1065       } else allOK = 0;
1066    } else allOK = 0;
1067    return allOK;
1068 } // BasicMBRData::SetPartType()
1069 
1070 // Set (or remove) the partition's bootable flag. Setting it is the
1071 // default; pass 0 as bootable to remove the flag.
1072 // Returns 1 if successful, 0 if not (invalid partition number)
SetPartBootable(int num,int bootable)1073 int BasicMBRData::SetPartBootable(int num, int bootable) {
1074    int allOK = 1;
1075 
1076    if ((num >= 0) && (num < MAX_MBR_PARTS)) {
1077       if (partitions[num].GetLengthLBA() != UINT32_C(0)) {
1078          if (bootable == 0)
1079             partitions[num].SetStatus(UINT8_C(0x00));
1080          else
1081             partitions[num].SetStatus(UINT8_C(0x80));
1082       } else allOK = 0;
1083    } else allOK = 0;
1084    return allOK;
1085 } // BasicMBRData::SetPartBootable()
1086 
1087 // Create a partition that fills the most available space. Returns
1088 // 1 if partition was created, 0 otherwise. Intended for use in
1089 // creating hybrid MBRs.
MakeBiggestPart(int i,int type)1090 int BasicMBRData::MakeBiggestPart(int i, int type) {
1091    uint64_t start = UINT64_C(1); // starting point for each search
1092    uint64_t firstBlock; // first block in a segment
1093    uint64_t lastBlock; // last block in a segment
1094    uint64_t segmentSize; // size of segment in blocks
1095    uint64_t selectedSegment = UINT64_C(0); // location of largest segment
1096    uint64_t selectedSize = UINT64_C(0); // size of largest segment in blocks
1097    int found = 0;
1098    string anything;
1099 
1100    do {
1101       firstBlock = FindFirstAvailable(start);
1102       if (firstBlock > UINT64_C(0)) { // something's free...
1103          lastBlock = FindLastInFree(firstBlock);
1104          segmentSize = lastBlock - firstBlock + UINT64_C(1);
1105          if (segmentSize > selectedSize) {
1106             selectedSize = segmentSize;
1107             selectedSegment = firstBlock;
1108          } // if
1109          start = lastBlock + 1;
1110       } // if
1111    } while (firstBlock != 0);
1112    if ((selectedSize > UINT64_C(0)) && (selectedSize < diskSize)) {
1113       found = 1;
1114       MakePart(i, selectedSegment, selectedSize, type, 0);
1115    } else {
1116       found = 0;
1117    } // if/else
1118    return found;
1119 } // BasicMBRData::MakeBiggestPart(int i)
1120 
1121 // Delete partition #i
DeletePartition(int i)1122 void BasicMBRData::DeletePartition(int i) {
1123    partitions[i].Empty();
1124 } // BasicMBRData::DeletePartition()
1125 
1126 // Set the inclusion status (PRIMARY, LOGICAL, or NONE) with some sanity
1127 // checks to ensure the table remains legal.
1128 // Returns 1 on success, 0 on failure.
SetInclusionwChecks(int num,int inclStatus)1129 int BasicMBRData::SetInclusionwChecks(int num, int inclStatus) {
1130    int allOK = 1, origValue;
1131 
1132    if (IsLegal()) {
1133       if ((inclStatus == PRIMARY) || (inclStatus == LOGICAL) || (inclStatus == NONE)) {
1134          origValue = partitions[num].GetInclusion();
1135          partitions[num].SetInclusion(inclStatus);
1136          if (!IsLegal()) {
1137             partitions[num].SetInclusion(origValue);
1138             cerr << "Specified change is not legal! Aborting change!\n";
1139          } // if
1140       } else {
1141          cerr << "Invalid partition inclusion code in BasicMBRData::SetInclusionwChecks()!\n";
1142       } // if/else
1143    } else {
1144       cerr << "Partition table is not currently in a valid state. Aborting change!\n";
1145       allOK = 0;
1146    } // if/else
1147    return allOK;
1148 } // BasicMBRData::SetInclusionwChecks()
1149 
1150 // Recomputes the CHS values for the specified partition and adjusts the value.
1151 // Note that this will create a technically incorrect CHS value for EFI GPT (0xEE)
1152 // protective partitions, but this is required by some buggy BIOSes, so I'm
1153 // providing a function to do this deliberately at the user's command.
1154 // This function does nothing if the partition's length is 0.
RecomputeCHS(int partNum)1155 void BasicMBRData::RecomputeCHS(int partNum) {
1156    partitions[partNum].RecomputeCHS();
1157 } // BasicMBRData::RecomputeCHS()
1158 
1159 // Sorts the partitions starting with partition #start. This function
1160 // does NOT pay attention to primary/logical assignment, which is
1161 // critical when writing the partitions.
SortMBR(int start)1162 void BasicMBRData::SortMBR(int start) {
1163    if ((start < MAX_MBR_PARTS) && (start >= 0))
1164       sort(partitions + start, partitions + MAX_MBR_PARTS);
1165 } // BasicMBRData::SortMBR()
1166 
1167 // Delete any partitions that are too big to fit on the disk
1168 // or that are too big for MBR (32-bit limits).
1169 // This deletes the partitions by setting values to 0, not just
1170 // by setting them as being omitted.
1171 // Returns the number of partitions deleted in this way.
DeleteOversizedParts()1172 int BasicMBRData::DeleteOversizedParts() {
1173    int num = 0, i;
1174 
1175    for (i = 0; i < MAX_MBR_PARTS; i++) {
1176       if ((partitions[i].GetStartLBA() > diskSize) || (partitions[i].GetLastLBA() > diskSize) ||
1177           (partitions[i].GetStartLBA() > UINT32_MAX) || (partitions[i].GetLengthLBA() > UINT32_MAX)) {
1178          cerr << "\aWarning: Deleting oversized partition #" << i + 1 << "! Start = "
1179               << partitions[i].GetStartLBA() << ", length = " << partitions[i].GetLengthLBA() << "\n";
1180          partitions[i].Empty();
1181          num++;
1182       } // if
1183    } // for
1184    return num;
1185 } // BasicMBRData::DeleteOversizedParts()
1186 
1187 // Search for and delete extended partitions.
1188 // Returns the number of partitions deleted.
DeleteExtendedParts()1189 int BasicMBRData::DeleteExtendedParts() {
1190    int i, numDeleted = 0;
1191    uint8_t type;
1192 
1193    for (i = 0; i < MAX_MBR_PARTS; i++) {
1194       type = partitions[i].GetType();
1195       if (((type == 0x05) || (type == 0x0f) || (type == (0x85))) &&
1196           (partitions[i].GetLengthLBA() > 0)) {
1197          partitions[i].Empty();
1198          numDeleted++;
1199       } // if
1200    } // for
1201    return numDeleted;
1202 } // BasicMBRData::DeleteExtendedParts()
1203 
1204 // Finds any overlapping partitions and omits the smaller of the two.
OmitOverlaps()1205 void BasicMBRData::OmitOverlaps() {
1206    int i, j;
1207 
1208    for (i = 0; i < MAX_MBR_PARTS; i++) {
1209       for (j = i + 1; j < MAX_MBR_PARTS; j++) {
1210          if ((partitions[i].GetInclusion() != NONE) &&
1211              partitions[i].DoTheyOverlap(partitions[j])) {
1212             if (partitions[i].GetLengthLBA() < partitions[j].GetLengthLBA())
1213                partitions[i].SetInclusion(NONE);
1214             else
1215                partitions[j].SetInclusion(NONE);
1216          } // if
1217       } // for (j...)
1218    } // for (i...)
1219 } // BasicMBRData::OmitOverlaps()
1220 
1221 // Convert as many partitions into logicals as possible, except for
1222 // the first partition, if possible.
MaximizeLogicals()1223 void BasicMBRData::MaximizeLogicals() {
1224    int earliestPart = 0, earliestPartWas = NONE, i;
1225 
1226    for (i = MAX_MBR_PARTS - 1; i >= 0; i--) {
1227       UpdateCanBeLogical();
1228       earliestPart = i;
1229       if (partitions[i].CanBeLogical()) {
1230          partitions[i].SetInclusion(LOGICAL);
1231       } else if (partitions[i].CanBePrimary()) {
1232          partitions[i].SetInclusion(PRIMARY);
1233       } else {
1234          partitions[i].SetInclusion(NONE);
1235       } // if/elseif/else
1236    } // for
1237    // If we have spare primaries, convert back the earliest partition to
1238    // its original state....
1239    if ((NumPrimaries() < 4) && (partitions[earliestPart].GetInclusion() == LOGICAL))
1240       partitions[earliestPart].SetInclusion(earliestPartWas);
1241 } // BasicMBRData::MaximizeLogicals()
1242 
1243 // Add primaries up to the maximum allowed, from the omitted category.
MaximizePrimaries()1244 void BasicMBRData::MaximizePrimaries() {
1245    int num, i = 0;
1246 
1247    num = NumPrimaries();
1248    while ((num < 4) && (i < MAX_MBR_PARTS)) {
1249       if ((partitions[i].GetInclusion() == NONE) && (partitions[i].CanBePrimary())) {
1250          partitions[i].SetInclusion(PRIMARY);
1251          num++;
1252          UpdateCanBeLogical();
1253       } // if
1254       i++;
1255    } // while
1256 } // BasicMBRData::MaximizePrimaries()
1257 
1258 // Remove primary partitions in excess of 4, starting with the later ones,
1259 // in terms of the array location....
TrimPrimaries(void)1260 void BasicMBRData::TrimPrimaries(void) {
1261    int numToDelete, i = MAX_MBR_PARTS - 1;
1262 
1263    numToDelete = NumPrimaries() - 4;
1264    while ((numToDelete > 0) && (i >= 0)) {
1265       if (partitions[i].GetInclusion() == PRIMARY) {
1266          partitions[i].SetInclusion(NONE);
1267          numToDelete--;
1268       } // if
1269       i--;
1270    } // while (numToDelete > 0)
1271 } // BasicMBRData::TrimPrimaries()
1272 
1273 // Locates primary partitions located between logical partitions and
1274 // either converts the primaries into logicals (if possible) or omits
1275 // them.
MakeLogicalsContiguous(void)1276 void BasicMBRData::MakeLogicalsContiguous(void) {
1277    uint64_t firstLogicalLBA, lastLogicalLBA;
1278    int i;
1279 
1280    firstLogicalLBA = FirstLogicalLBA();
1281    lastLogicalLBA = LastLogicalLBA();
1282    for (i = 0; i < MAX_MBR_PARTS; i++) {
1283       if ((partitions[i].GetInclusion() == PRIMARY) &&
1284           (partitions[i].GetStartLBA() >= firstLogicalLBA) &&
1285           (partitions[i].GetLastLBA() <= lastLogicalLBA)) {
1286          if (SectorUsedAs(partitions[i].GetStartLBA() - 1) == NONE)
1287             partitions[i].SetInclusion(LOGICAL);
1288          else
1289             partitions[i].SetInclusion(NONE);
1290       } // if
1291    } // for
1292 } // BasicMBRData::MakeLogicalsContiguous()
1293 
1294 // If MBR data aren't legal, adjust primary/logical assignments and,
1295 // if necessary, drop partitions, to make the data legal.
MakeItLegal(void)1296 void BasicMBRData::MakeItLegal(void) {
1297    if (!IsLegal()) {
1298       DeleteOversizedParts();
1299       MaximizeLogicals();
1300       MaximizePrimaries();
1301       if (!AreLogicalsContiguous())
1302          MakeLogicalsContiguous();
1303       if (NumPrimaries() > 4)
1304          TrimPrimaries();
1305       OmitOverlaps();
1306    } // if
1307 } // BasicMBRData::MakeItLegal()
1308 
1309 // Removes logical partitions and deactivated partitions from first four
1310 // entries (primary space).
1311 // Returns the number of partitions moved.
RemoveLogicalsFromFirstFour(void)1312 int BasicMBRData::RemoveLogicalsFromFirstFour(void) {
1313    int i, j = 4, numMoved = 0, swapped = 0;
1314    MBRPart temp;
1315 
1316    for (i = 0; i < 4; i++) {
1317       if ((partitions[i].GetInclusion() != PRIMARY) && (partitions[i].GetLengthLBA() > 0)) {
1318          j = 4;
1319          swapped = 0;
1320          do {
1321             if ((partitions[j].GetInclusion() == NONE) && (partitions[j].GetLengthLBA() == 0)) {
1322                temp = partitions[j];
1323                partitions[j] = partitions[i];
1324                partitions[i] = temp;
1325                swapped = 1;
1326                numMoved++;
1327             } // if
1328             j++;
1329          } while ((j < MAX_MBR_PARTS) && !swapped);
1330          if (j >= MAX_MBR_PARTS)
1331             cerr << "Warning! Too many partitions in BasicMBRData::RemoveLogicalsFromFirstFour()!\n";
1332       } // if
1333    } // for i...
1334    return numMoved;
1335 } // BasicMBRData::RemoveLogicalsFromFirstFour()
1336 
1337 // Move all primaries into the first four partition spaces
1338 // Returns the number of partitions moved.
MovePrimariesToFirstFour(void)1339 int BasicMBRData::MovePrimariesToFirstFour(void) {
1340    int i, j = 0, numMoved = 0, swapped = 0;
1341    MBRPart temp;
1342 
1343    for (i = 4; i < MAX_MBR_PARTS; i++) {
1344       if (partitions[i].GetInclusion() == PRIMARY) {
1345          j = 0;
1346          swapped = 0;
1347          do {
1348             if (partitions[j].GetInclusion() != PRIMARY) {
1349                temp = partitions[j];
1350                partitions[j] = partitions[i];
1351                partitions[i] = temp;
1352                swapped = 1;
1353                numMoved++;
1354             } // if
1355             j++;
1356          } while ((j < 4) && !swapped);
1357       } // if
1358    } // for
1359    return numMoved;
1360 } // BasicMBRData::MovePrimariesToFirstFour()
1361 
1362 // Create an extended partition, if necessary, to hold the logical partitions.
1363 // This function also sorts the primaries into the first four positions of
1364 // the table.
1365 // Returns 1 on success, 0 on failure.
CreateExtended(void)1366 int BasicMBRData::CreateExtended(void) {
1367    int allOK = 1, i = 0, swapped = 0;
1368    MBRPart temp;
1369 
1370    if (IsLegal()) {
1371       // Move logicals out of primary space...
1372       RemoveLogicalsFromFirstFour();
1373       // Move primaries out of logical space...
1374       MovePrimariesToFirstFour();
1375 
1376       // Create the extended partition
1377       if (NumLogicals() > 0) {
1378          SortMBR(4); // sort starting from 4 -- that is, logicals only
1379          temp.Empty();
1380          temp.SetStartLBA(FirstLogicalLBA() - 1);
1381          temp.SetLengthLBA(LastLogicalLBA() - FirstLogicalLBA() + 2);
1382          temp.SetType(0x0f, 1);
1383          temp.SetInclusion(PRIMARY);
1384          do {
1385             if ((partitions[i].GetInclusion() == NONE) || (partitions[i].GetLengthLBA() == 0)) {
1386                partitions[i] = temp;
1387                swapped = 1;
1388             } // if
1389             i++;
1390          } while ((i < 4) && !swapped);
1391          if (!swapped) {
1392             cerr << "Could not create extended partition; no room in primary table!\n";
1393             allOK = 0;
1394          } // if
1395       } // if (NumLogicals() > 0)
1396    } else allOK = 0;
1397    // Do a final check for EFI GPT (0xEE) partitions & flag as a problem if found
1398    // along with an extended partition
1399    for (i = 0; i < MAX_MBR_PARTS; i++)
1400       if (swapped && partitions[i].GetType() == 0xEE)
1401          allOK = 0;
1402    return allOK;
1403 } // BasicMBRData::CreateExtended()
1404 
1405 /****************************************
1406  *                                      *
1407  * Functions to find data on free space *
1408  *                                      *
1409  ****************************************/
1410 
1411 // Finds the first free space on the disk from start onward; returns 0
1412 // if none available....
FindFirstAvailable(uint64_t start)1413 uint64_t BasicMBRData::FindFirstAvailable(uint64_t start) {
1414    uint64_t first;
1415    uint64_t i;
1416    int firstMoved;
1417 
1418    if ((start >= (UINT32_MAX - 1)) || (start >= (diskSize - 1)))
1419       return 0;
1420 
1421    first = start;
1422 
1423    // ...now search through all partitions; if first is within an
1424    // existing partition, move it to the next sector after that
1425    // partition and repeat. If first was moved, set firstMoved
1426    // flag; repeat until firstMoved is not set, so as to catch
1427    // cases where partitions are out of sequential order....
1428    do {
1429       firstMoved = 0;
1430       for (i = 0; i < 4; i++) {
1431          // Check if it's in the existing partition
1432          if ((first >= partitions[i].GetStartLBA()) &&
1433              (first < (partitions[i].GetStartLBA() + partitions[i].GetLengthLBA()))) {
1434             first = partitions[i].GetStartLBA() + partitions[i].GetLengthLBA();
1435             firstMoved = 1;
1436          } // if
1437       } // for
1438    } while (firstMoved == 1);
1439    if ((first >= diskSize) || (first > UINT32_MAX))
1440       first = 0;
1441    return (first);
1442 } // BasicMBRData::FindFirstAvailable()
1443 
1444 // Finds the last free sector on the disk from start forward.
FindLastInFree(uint64_t start)1445 uint64_t BasicMBRData::FindLastInFree(uint64_t start) {
1446    uint64_t nearestStart;
1447    uint64_t i;
1448 
1449    if ((diskSize <= UINT32_MAX) && (diskSize > 0))
1450       nearestStart = diskSize - 1;
1451    else
1452       nearestStart = UINT32_MAX - 1;
1453 
1454    for (i = 0; i < 4; i++) {
1455       if ((nearestStart > partitions[i].GetStartLBA()) &&
1456           (partitions[i].GetStartLBA() > start)) {
1457          nearestStart = partitions[i].GetStartLBA() - 1;
1458       } // if
1459    } // for
1460    return (nearestStart);
1461 } // BasicMBRData::FindLastInFree()
1462 
1463 // Finds the first free sector on the disk from start backward.
FindFirstInFree(uint64_t start)1464 uint64_t BasicMBRData::FindFirstInFree(uint64_t start) {
1465    uint64_t bestLastLBA, thisLastLBA;
1466    int i;
1467 
1468    bestLastLBA = 1;
1469    for (i = 0; i < 4; i++) {
1470       thisLastLBA = partitions[i].GetLastLBA() + 1;
1471       if (thisLastLBA > 0)
1472          thisLastLBA--;
1473       if ((thisLastLBA > bestLastLBA) && (thisLastLBA < start))
1474          bestLastLBA = thisLastLBA + 1;
1475    } // for
1476    return (bestLastLBA);
1477 } // BasicMBRData::FindFirstInFree()
1478 
1479 // Returns NONE (unused), PRIMARY, LOGICAL, EBR (for EBR or MBR), or INVALID.
1480 // Note: If the sector immediately before a logical partition is in use by
1481 // another partition, this function returns PRIMARY or LOGICAL for that
1482 // sector, rather than EBR.
SectorUsedAs(uint64_t sector,int topPartNum)1483 int BasicMBRData::SectorUsedAs(uint64_t sector, int topPartNum) {
1484    int i = 0, usedAs = NONE;
1485 
1486    do {
1487       if ((partitions[i].GetStartLBA() <= sector) && (partitions[i].GetLastLBA() >= sector))
1488          usedAs = partitions[i].GetInclusion();
1489       if ((partitions[i].GetStartLBA() == (sector + 1)) && (partitions[i].GetInclusion() == LOGICAL))
1490          usedAs = EBR;
1491       if (sector == 0)
1492          usedAs = EBR;
1493       if (sector >= diskSize)
1494          usedAs = INVALID;
1495       i++;
1496    } while ((i < topPartNum) && ((usedAs == NONE) || (usedAs == EBR)));
1497    return usedAs;
1498 } // BasicMBRData::SectorUsedAs()
1499 
1500 /******************************************************
1501  *                                                    *
1502  * Functions that extract data on specific partitions *
1503  *                                                    *
1504  ******************************************************/
1505 
GetStatus(int i)1506 uint8_t BasicMBRData::GetStatus(int i) {
1507    MBRPart* thePart;
1508    uint8_t retval;
1509 
1510    thePart = GetPartition(i);
1511    if (thePart != NULL)
1512       retval = thePart->GetStatus();
1513    else
1514       retval = UINT8_C(0);
1515    return retval;
1516 } // BasicMBRData::GetStatus()
1517 
GetType(int i)1518 uint8_t BasicMBRData::GetType(int i) {
1519    MBRPart* thePart;
1520    uint8_t retval;
1521 
1522    thePart = GetPartition(i);
1523    if (thePart != NULL)
1524       retval = thePart->GetType();
1525    else
1526       retval = UINT8_C(0);
1527    return retval;
1528 } // BasicMBRData::GetType()
1529 
GetFirstSector(int i)1530 uint64_t BasicMBRData::GetFirstSector(int i) {
1531    MBRPart* thePart;
1532    uint64_t retval;
1533 
1534    thePart = GetPartition(i);
1535    if (thePart != NULL)
1536       retval = thePart->GetStartLBA();
1537    else
1538       retval = UINT32_C(0);
1539    return retval;
1540 } // BasicMBRData::GetFirstSector()
1541 
GetLength(int i)1542 uint64_t BasicMBRData::GetLength(int i) {
1543    MBRPart* thePart;
1544    uint64_t retval;
1545 
1546    thePart = GetPartition(i);
1547    if (thePart != NULL)
1548       retval = thePart->GetLengthLBA();
1549    else
1550       retval = UINT64_C(0);
1551    return retval;
1552 } // BasicMBRData::GetLength()
1553 
1554 /***********************
1555  *                     *
1556  * Protected functions *
1557  *                     *
1558  ***********************/
1559 
1560 // Return a pointer to a primary or logical partition, or NULL if
1561 // the partition is out of range....
GetPartition(int i)1562 MBRPart* BasicMBRData::GetPartition(int i) {
1563    MBRPart* thePart = NULL;
1564 
1565    if ((i >= 0) && (i < MAX_MBR_PARTS))
1566       thePart = &partitions[i];
1567    return thePart;
1568 } // GetPartition()
1569 
1570 /*******************************************
1571  *                                         *
1572  * Functions that involve user interaction *
1573  *                                         *
1574  *******************************************/
1575 
1576 // Present the MBR operations menu. Note that the 'w' option does not
1577 // immediately write data; that's handled by the calling function.
1578 // Returns the number of partitions defined on exit, or -1 if the
1579 // user selected the 'q' option. (Thus, the caller should save data
1580 // if the return value is >0, or possibly >=0 depending on intentions.)
DoMenu(const string & prompt)1581 int BasicMBRData::DoMenu(const string& prompt) {
1582    int goOn = 1, quitting = 0, retval, num, haveShownInfo = 0;
1583    unsigned int hexCode;
1584    string tempStr;
1585 
1586    do {
1587       cout << prompt;
1588       switch (ReadString()[0]) {
1589          case '\0':
1590             goOn = cin.good();
1591             break;
1592          case 'a': case 'A':
1593             num = GetNumber(1, MAX_MBR_PARTS, 1, "Toggle active flag for partition: ") - 1;
1594             if (partitions[num].GetInclusion() != NONE)
1595                partitions[num].SetStatus(partitions[num].GetStatus() ^ 0x80);
1596             break;
1597          case 'c': case 'C':
1598             for (num = 0; num < MAX_MBR_PARTS; num++)
1599                RecomputeCHS(num);
1600             break;
1601          case 'l': case 'L':
1602             num = GetNumber(1, MAX_MBR_PARTS, 1, "Partition to set as logical: ") - 1;
1603             SetInclusionwChecks(num, LOGICAL);
1604             break;
1605          case 'o': case 'O':
1606             num = GetNumber(1, MAX_MBR_PARTS, 1, "Partition to omit: ") - 1;
1607             SetInclusionwChecks(num, NONE);
1608             break;
1609          case 'p': case 'P':
1610             if (!haveShownInfo) {
1611                cout << "\n** NOTE: Partition numbers do NOT indicate final primary/logical "
1612                     << "status,\n** unlike in most MBR partitioning tools!\n\a";
1613                cout << "\n** Extended partitions are not displayed, but will be generated "
1614                     << "as required.\n";
1615                haveShownInfo = 1;
1616             } // if
1617             DisplayMBRData();
1618             break;
1619          case 'q': case 'Q':
1620             cout << "This will abandon your changes. Are you sure? ";
1621             if (GetYN() == 'Y') {
1622                goOn = 0;
1623                quitting = 1;
1624             } // if
1625             break;
1626          case 'r': case 'R':
1627             num = GetNumber(1, MAX_MBR_PARTS, 1, "Partition to set as primary: ") - 1;
1628             SetInclusionwChecks(num, PRIMARY);
1629             break;
1630          case 's': case 'S':
1631             SortMBR();
1632             break;
1633          case 't': case 'T':
1634             num = GetNumber(1, MAX_MBR_PARTS, 1, "Partition to change type code: ") - 1;
1635             hexCode = 0x00;
1636             if (partitions[num].GetLengthLBA() > 0) {
1637                while ((hexCode <= 0) || (hexCode > 255)) {
1638                   cout << "Enter an MBR hex code: ";
1639                   tempStr = ReadString();
1640                   if (IsHex(tempStr))
1641                      sscanf(tempStr.c_str(), "%x", &hexCode);
1642                } // while
1643                partitions[num].SetType(hexCode);
1644             } // if
1645             break;
1646          case 'w': case 'W':
1647             goOn = 0;
1648             break;
1649          default:
1650             ShowCommands();
1651             break;
1652       } // switch
1653    } while (goOn);
1654    if (quitting)
1655       retval = -1;
1656    else
1657       retval = CountParts();
1658    return (retval);
1659 } // BasicMBRData::DoMenu()
1660 
ShowCommands(void)1661 void BasicMBRData::ShowCommands(void) {
1662    cout << "a\ttoggle the active/boot flag\n";
1663    cout << "c\trecompute all CHS values\n";
1664    cout << "l\tset partition as logical\n";
1665    cout << "o\tomit partition\n";
1666    cout << "p\tprint the MBR partition table\n";
1667    cout << "q\tquit without saving changes\n";
1668    cout << "r\tset partition as primary\n";
1669    cout << "s\tsort MBR partitions\n";
1670    cout << "t\tchange partition type code\n";
1671    cout << "w\twrite the MBR partition table to disk and exit\n";
1672 } // BasicMBRData::ShowCommands()
1673