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