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