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