EdkCompatibilityPkg/Sample/Tools/Source/GenBootsector/genbootsector.c

   1 /*++
   2
   3 Copyright 2006 - 2007, Intel Corporation
   4 All rights reserved. This program and the accompanying materials
   5 are licensed and made available under the terms and conditions of the BSD License
   6 which accompanies this distribution.  The full text of the license may be found at
   7 http://opensource.org/licenses/bsd-license.php
   8
   9 THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
  10 WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
  11
  12 Module Name:
  13
  14   genbootsector.c
  15
  16 Abstract:
  17   Reading/writing MBR/DBR.
  18   NOTE:
  19     If we write MBR to disk, we just update the MBR code and the partition table wouldn't be over written.
  20     If we process DBR, we will patch MBR to set first partition active if no active partition exists.
  21
  22 --*/
  23
  24 #include <windows.h>
  25 #include <stdio.h>
  26 #include <string.h>
  27
  28 #define MAX_DRIVE                             26
  29 #define PARTITION_TABLE_OFFSET                0x1BE
  30
  31 #define SIZE_OF_PARTITION_ENTRY               0x10
  32
  33 #define PARTITION_ENTRY_STARTLBA_OFFSET       8
  34
  35 #define PARTITION_ENTRY_NUM                   4
  36
  37 INT
  38 GetDrvNumOffset (
  39   IN VOID *BootSector
  40   );
  41
  42 typedef enum {
  43   PatchTypeUnknown,
  44   PatchTypeFloppy,
  45   PatchTypeIde,
  46   PatchTypeUsb,
  47 } PATCH_TYPE;
  48
  49 typedef enum {
  50   ErrorSuccess,
  51   ErrorFileCreate,
  52   ErrorFileReadWrite,
  53   ErrorNoMbr,
  54   ErrorFatType
  55 } ERROR_STATUS;
  56
  57 CHAR *ErrorStatusDesc[] = {
  58   "Success",
  59   "Failed to create files",
  60   "Failed to read/write files",
  61   "No MBR exists",
  62   "Failed to detect Fat type"
  63 };
  64
  65 typedef struct _DRIVE_TYPE_DESC {
  66   UINT  Type;
  67   CHAR  *Description;
  68 } DRIVE_TYPE_DESC;
  69
  70 #define DRIVE_TYPE_ITEM(x) {x, #x}
  71 DRIVE_TYPE_DESC DriveTypeDesc[] = {
  72   DRIVE_TYPE_ITEM (DRIVE_UNKNOWN),
  73   DRIVE_TYPE_ITEM (DRIVE_NO_ROOT_DIR),
  74   DRIVE_TYPE_ITEM (DRIVE_REMOVABLE),
  75   DRIVE_TYPE_ITEM (DRIVE_FIXED),
  76   DRIVE_TYPE_ITEM (DRIVE_REMOTE),
  77   DRIVE_TYPE_ITEM (DRIVE_CDROM),
  78   DRIVE_TYPE_ITEM (DRIVE_RAMDISK),
  79   (UINT) -1, NULL
  80 };
  81
  82 typedef struct _DRIVE_INFO {
  83   CHAR              VolumeLetter;
  84   DRIVE_TYPE_DESC   *DriveType;
  85   UINT              DiskNumber;
  86 } DRIVE_INFO;
  87
  88 #define BOOT_SECTOR_LBA_OFFSET 0x1FA
  89
  90 #define IsLetter(x) (((x) >= 'a' && (x) <= 'z') || ((x) >= 'A' && (x) <= 'Z'))
  91
  92 BOOL
  93 GetDriveInfo (
  94   CHAR       VolumeLetter,
  95   DRIVE_INFO *DriveInfo
  96   )
  97 /*++
  98 Routine Description:
  99   Get drive information including disk number and drive type,
 100   where disknumber is useful for reading/writing disk raw data.
 101   NOTE: Floppy disk doesn't have disk number but it doesn't matter because
 102         we can reading/writing floppy disk without disk number.
 103
 104 Arguments:
 105   VolumeLetter : volume letter, e.g.: C for C:, A for A:
 106   DriveInfo    : pointer to DRIVE_INFO structure receiving drive information.
 107
 108 Return:
 109   TRUE  : successful
 110   FALSE : failed
 111 --*/
 112 {
 113   HANDLE                  VolumeHandle;
 114   STORAGE_DEVICE_NUMBER   StorageDeviceNumber;
 115   DWORD                   BytesReturned;
 116   BOOL                    Success;
 117   UINT                    DriveType;
 118   UINT                    Index;
 119
 120   CHAR RootPath[]         = "X:\\";       // "X:\"  -> for GetDriveType
 121   CHAR VolumeAccessPath[] = "\\\\.\\X:";  // "\\.\X:"  -> to open the volume
 122
 123   RootPath[0] = VolumeAccessPath[4] = VolumeLetter;
 124   DriveType = GetDriveType(RootPath);
 125   if (DriveType != DRIVE_REMOVABLE && DriveType != DRIVE_FIXED) {
 126     return FALSE;
 127   }
 128
 129   DriveInfo->VolumeLetter = VolumeLetter;
 130   VolumeHandle = CreateFile (
 131                    VolumeAccessPath,
 132                    0,
 133                    FILE_SHARE_READ | FILE_SHARE_WRITE,
 134                    NULL,
 135                    OPEN_EXISTING,
 136                    0,
 137                    NULL
 138                    );
 139   if (VolumeHandle == INVALID_HANDLE_VALUE) {
 140     fprintf (
 141       stderr,
 142       "ERROR: CreateFile failed: Volume = %s, LastError = 0x%x\n",
 143       VolumeAccessPath,
 144       GetLastError ()
 145       );
 146     return FALSE;
 147   }
 148
 149   //
 150   // Get Disk Number. It should fail when operating on floppy. That's ok
 151   //  because Disk Number is only needed when operating on Hard or USB disk.
 152   //
 153   // To direct write to disk:
 154   //   for USB and HD: use path = \\.\PHYSICALDRIVEx, where x is Disk Number
 155   //   for floppy:     use path = \\.\X:, where X can be A or B
 156   //
 157   Success = DeviceIoControl(
 158               VolumeHandle,
 159               IOCTL_STORAGE_GET_DEVICE_NUMBER,
 160               NULL,
 161               0,
 162               &StorageDeviceNumber,
 163               sizeof(StorageDeviceNumber),
 164               &BytesReturned,
 165               NULL
 166               );
 167   //
 168   // DeviceIoControl should fail if Volume is floppy or network drive.
 169   //
 170   if (!Success) {
 171     DriveInfo->DiskNumber = (UINT) -1;
 172   } else if (StorageDeviceNumber.DeviceType != FILE_DEVICE_DISK) {
 173     //
 174     // Only care about the disk.
 175     //
 176     return FALSE;
 177   } else{
 178     DriveInfo->DiskNumber = StorageDeviceNumber.DeviceNumber;
 179   }
 180   CloseHandle(VolumeHandle);
 181
 182   //
 183   // Fill in the type string
 184   //
 185   DriveInfo->DriveType = NULL;
 186   for (Index = 0; DriveTypeDesc[Index].Description != NULL; Index ++) {
 187     if (DriveType == DriveTypeDesc[Index].Type) {
 188       DriveInfo->DriveType = &DriveTypeDesc[Index];
 189       break;
 190     }
 191   }
 192
 193   if (DriveInfo->DriveType == NULL) {
 194     //
 195     // Should have a type.
 196     //
 197     fprintf (stderr, "ERROR: fetal error!!!\n");
 198     return FALSE;
 199   }
 200   return TRUE;
 201 }
 202
 203 VOID
 204 ListDrive (
 205   VOID
 206   )
 207 /*++
 208 Routine Description:
 209   List every drive in current system and their information.
 210
 211 --*/
 212 {
 213   UINT       Index;
 214   DRIVE_INFO DriveInfo;
 215
 216   UINT Mask =  GetLogicalDrives();
 217
 218   for (Index = 0; Index < MAX_DRIVE; Index++) {
 219     if (((Mask >> Index) & 0x1) == 1) {
 220       if (GetDriveInfo ('A' + (CHAR) Index, &DriveInfo)) {
 221         if (Index < 2) {
 222           // Floppy will occupy 'A' and 'B'
 223           fprintf (
 224             stdout,
 225             "%c: - Type: %s\n",
 226             DriveInfo.VolumeLetter,
 227             DriveInfo.DriveType->Description
 228             );
 229         }
 230         else {
 231           fprintf (
 232             stdout,
 233             "%c: - DiskNum: %d, Type: %s\n",
 234             DriveInfo.VolumeLetter,
 235             DriveInfo.DiskNumber,
 236             DriveInfo.DriveType->Description
 237             );
 238         }
 239       }
 240     }
 241   }
 242
 243 }
 244
 245 INT
 246 GetBootSectorOffset (
 247   HANDLE     DiskHandle,
 248   BOOL       WriteToDisk,
 249   PATCH_TYPE PatchType
 250   )
 251 /*++
 252 Description:
 253   Get the offset of boot sector.
 254   For non-MBR disk, offset is just 0
 255   for disk with MBR, offset needs to be caculated by parsing MBR
 256
 257   NOTE: if no one is active, we will patch MBR to select first partition as active.
 258
 259 Arguments:
 260   DiskHandle  : HANDLE of disk
 261   WriteToDisk : TRUE indicates writing
 262   PatchType   : PatchTypeFloppy, PatchTypeIde, PatchTypeUsb
 263
 264 Return:
 265   -1   : failed
 266   o.w. : Offset to boot sector
 267 --*/
 268 {
 269   BYTE    DiskPartition[0x200];
 270   DWORD   BytesReturn;
 271   DWORD   DbrOffset;
 272   DWORD   Index;
 273   BOOL    HasMbr;
 274
 275   DbrOffset = 0;
 276   HasMbr    = FALSE;
 277
 278   SetFilePointer(DiskHandle, 0, NULL, FILE_BEGIN);
 279   if (!ReadFile (DiskHandle, DiskPartition, 0x200, &BytesReturn, NULL)) {
 280     return -1;
 281   }
 282
 283   //
 284   // Check Signature, Jmp, and Boot Indicator.
 285   // if all pass, we assume MBR found.
 286   //
 287
 288   // Check Signature: 55AA
 289   if ((DiskPartition[0x1FE] == 0x55) && (DiskPartition[0x1FF] == 0xAA)) {
 290     // Check Jmp: (EB ?? 90) or (E9 ?? ??)
 291     if (((DiskPartition[0] != 0xEB) || (DiskPartition[2] != 0x90)) &&
 292         (DiskPartition[0] != 0xE9)) {
 293       // Check Boot Indicator: 0x00 or 0x80
 294       // Boot Indicator is the first byte of Partition Entry
 295       HasMbr = TRUE;
 296       for (Index = 0; Index < PARTITION_ENTRY_NUM; ++Index) {
 297         if ((DiskPartition[PARTITION_TABLE_OFFSET + Index * SIZE_OF_PARTITION_ENTRY] & 0x7F) != 0) {
 298           HasMbr = FALSE;
 299           break;
 300         }
 301       }
 302     }
 303   }
 304
 305   if (HasMbr) {
 306     //
 307     // Skip MBR
 308     //
 309     for (Index = 0; Index < PARTITION_ENTRY_NUM; Index++) {
 310       //
 311       // Found Boot Indicator.
 312       //
 313       if (DiskPartition[PARTITION_TABLE_OFFSET + (Index * SIZE_OF_PARTITION_ENTRY)] == 0x80) {
 314         DbrOffset = *(DWORD *)&DiskPartition[PARTITION_TABLE_OFFSET + (Index * SIZE_OF_PARTITION_ENTRY) + PARTITION_ENTRY_STARTLBA_OFFSET];
 315         break;
 316       }
 317     }
 318     //
 319     // If no boot indicator, we manually select 1st partition, and patch MBR.
 320     //
 321     if (Index == PARTITION_ENTRY_NUM) {
 322       DbrOffset = *(DWORD *)&DiskPartition[PARTITION_TABLE_OFFSET + PARTITION_ENTRY_STARTLBA_OFFSET];
 323       if (WriteToDisk && (PatchType == PatchTypeUsb)) {
 324         SetFilePointer(DiskHandle, 0, NULL, FILE_BEGIN);
 325         DiskPartition[PARTITION_TABLE_OFFSET] = 0x80;
 326         WriteFile (DiskHandle, DiskPartition, 0x200, &BytesReturn, NULL);
 327       }
 328     }
 329   }
 330
 331   return DbrOffset;
 332 }
 333
 334 ERROR_STATUS
 335 ProcessBsOrMbr (
 336   CHAR        *DiskName,
 337   CHAR        *FileName,
 338   BOOL        WriteToDisk,
 339   PATCH_TYPE  PatchType,
 340   BOOL        ProcessMbr
 341   )
 342 /*++
 343 Routine Description:
 344   Writing or reading boot sector or MBR according to the argument.
 345
 346 Arguments:
 347   DiskName    : Win32 API recognized string name of disk
 348   FileName    : file name
 349   WriteToDisk : TRUE is to write content of file to disk, otherwise, reading content of disk to file
 350   PatchType   : PatchTypeFloppy, PatchTypeIde, PatchTypeUsb
 351   ProcessMbr  : TRUE is to process MBR, otherwise, processing boot sector
 352
 353 Return:
 354   ErrorSuccess
 355   ErrorFileCreate
 356   ErrorFileReadWrite
 357   ErrorNoMbr
 358   ErrorFatType
 359 --*/
 360 {
 361   BYTE    DiskPartition[0x200];
 362   BYTE    DiskPartitionBackup[0x200];
 363   HANDLE  DiskHandle;
 364   HANDLE  FileHandle;
 365   DWORD   BytesReturn;
 366   DWORD   DbrOffset;
 367   INT     DrvNumOffset;
 368
 369   DiskHandle = CreateFile (
 370                  DiskName,
 371                  GENERIC_READ | GENERIC_WRITE,
 372                  FILE_SHARE_READ,
 373                  NULL,
 374                  OPEN_EXISTING,
 375                  FILE_ATTRIBUTE_NORMAL,
 376                  NULL
 377                  );
 378   if (DiskHandle == INVALID_HANDLE_VALUE) {
 379     return ErrorFileCreate;
 380   }
 381
 382   FileHandle = CreateFile (
 383                  FileName,
 384                  GENERIC_READ | GENERIC_WRITE,
 385                  0,
 386                  NULL,
 387                  OPEN_ALWAYS,
 388                  FILE_ATTRIBUTE_NORMAL,
 389                  NULL
 390                  );
 391   if (FileHandle == INVALID_HANDLE_VALUE) {
 392     return ErrorFileCreate;
 393   }
 394
 395   DbrOffset = 0;
 396   //
 397   // Skip potential MBR for Ide & USB disk
 398   //
 399   if ((PatchType == PatchTypeIde) || (PatchType == PatchTypeUsb)) {
 400     //
 401     // Even user just wants to process MBR, we get offset of boot sector here to validate the disk
 402     //  if disk have MBR, DbrOffset should be greater than 0
 403     //
 404     DbrOffset = GetBootSectorOffset (DiskHandle, WriteToDisk, PatchType);
 405
 406     if (!ProcessMbr) {
 407       //
 408       // 1. Process boot sector, set file pointer to the beginning of boot sector
 409       //
 410       SetFilePointer (DiskHandle, DbrOffset * 0x200, NULL, FILE_BEGIN);
 411     } else if(DbrOffset == 0) {
 412       //
 413       // If user want to process Mbr, but no Mbr exists, simply return FALSE
 414       //
 415       return ErrorNoMbr;
 416     } else {
 417       //
 418       // 2. Process MBR, set file pointer to 0
 419       //
 420       SetFilePointer (DiskHandle, 0, NULL, FILE_BEGIN);
 421     }
 422   }
 423
 424   //
 425   // [File Pointer is pointed to beginning of Mbr or Dbr]
 426   //
 427   if (WriteToDisk) {
 428     //
 429     // Write
 430     //
 431     if (!ReadFile (FileHandle, DiskPartition, 0x200, &BytesReturn, NULL)) {
 432       return ErrorFileReadWrite;
 433     }
 434     if (ProcessMbr) {
 435       //
 436       // Use original partition table
 437       //
 438       if (!ReadFile (DiskHandle, DiskPartitionBackup, 0x200, &BytesReturn, NULL)) {
 439         return ErrorFileReadWrite;
 440       }
 441       memcpy (DiskPartition + 0x1BE, DiskPartitionBackup + 0x1BE, 0x40);
 442       SetFilePointer (DiskHandle, 0, NULL, FILE_BEGIN);
 443     }
 444
 445     if (!WriteFile (DiskHandle, DiskPartition, 0x200, &BytesReturn, NULL)) {
 446       return ErrorFileReadWrite;
 447     }
 448
 449   } else {
 450     //
 451     // Read
 452     //
 453     if (!ReadFile (DiskHandle, DiskPartition, 0x200, &BytesReturn, NULL)) {
 454       return ErrorFileReadWrite;
 455     }
 456
 457     if (PatchType == PatchTypeUsb) {
 458       // Manually set BS_DrvNum to 0x80 as window's format.exe has a bug which will clear this field discarding USB disk's MBR.
 459       // offset of BS_DrvNum is 0x24 for FAT12/16
 460       //                        0x40 for FAT32
 461       //
 462       DrvNumOffset = GetDrvNumOffset (DiskPartition);
 463       if (DrvNumOffset == -1) {
 464         return ErrorFatType;
 465       }
 466       //
 467       // Some legacy BIOS require 0x80 discarding MBR.
 468       // Question left here: is it needed to check Mbr before set 0x80?
 469       //
 470       DiskPartition[DrvNumOffset] = ((DbrOffset > 0) ? 0x80 : 0);
 471   }
 472
 473
 474     if (PatchType == PatchTypeIde) {
 475       //
 476       // Patch LBAOffsetForBootSector
 477       //
 478       *(DWORD *)&DiskPartition [BOOT_SECTOR_LBA_OFFSET] = DbrOffset;
 479     }
 480     if (!WriteFile (FileHandle, DiskPartition, 0x200, &BytesReturn, NULL)) {
 481       return ErrorFileReadWrite;
 482     }
 483   }
 484   CloseHandle (FileHandle);
 485   CloseHandle (DiskHandle);
 486   return ErrorSuccess;
 487 }
 488
 489 VOID
 490 PrintUsage (
 491   CHAR* AppName
 492   )
 493 {
 494   fprintf (
 495     stdout,
 496     "Usage: %s [OPTIONS]...\n"
 497     "Copy file content from/to bootsector.\n"
 498     "\n"
 499     "  -l        list disks\n"
 500     "  -if=FILE  specified an input, can be files or disks\n"
 501     "  -of=FILE  specified an output, can be files or disks\n"
 502     "  -mbr      process MBR also\n"
 503     "  -h        print this message\n"
 504     "\n"
 505     "FILE providing a volume plus a colon (X:), indicates a disk\n"
 506     "FILE providing other format, indicates a file\n",
 507     AppName
 508     );
 509 }
 510
 511 INT
 512 main (
 513   INT  argc,
 514   CHAR *argv[]
 515   )
 516 {
 517   CHAR          *AppName;
 518   INT           Index;
 519   BOOL          ProcessMbr;
 520   CHAR          VolumeLetter;
 521   CHAR          *FilePath;
 522   BOOL          WriteToDisk;
 523   DRIVE_INFO    DriveInfo;
 524   PATCH_TYPE    PatchType;
 525   ERROR_STATUS  Status;
 526
 527   CHAR        FloppyPathTemplate[] = "\\\\.\\%c:";
 528   CHAR        DiskPathTemplate[]   = "\\\\.\\PHYSICALDRIVE%u";
 529   CHAR        DiskPath[MAX_PATH];
 530
 531   AppName = *argv;
 532   argv ++;
 533   argc --;
 534
 535   ProcessMbr    = FALSE;
 536   WriteToDisk   = TRUE;
 537   FilePath      = NULL;
 538   VolumeLetter  = 0;
 539
 540   //
 541   // Parse command line
 542   //
 543   for (Index = 0; Index < argc; Index ++) {
 544     if (_stricmp (argv[Index], "-l") == 0) {
 545       ListDrive ();
 546       return 0;
 547     }
 548     else if (_stricmp (argv[Index], "-mbr") == 0) {
 549       ProcessMbr = TRUE;
 550     }
 551     else if ((_strnicmp (argv[Index], "-if=", 4) == 0) ||
 552              (_strnicmp (argv[Index], "-of=", 4) == 0)
 553              ) {
 554       if (argv[Index][6] == '\0' && argv[Index][5] == ':' && IsLetter (argv[Index][4])) {
 555         VolumeLetter = argv[Index][4];
 556         if (_strnicmp (argv[Index], "-if=", 4) == 0) {
 557           WriteToDisk = FALSE;
 558         }
 559       }
 560       else {
 561         FilePath = &argv[Index][4];
 562       }
 563     }
 564     else {
 565       PrintUsage (AppName);
 566       return 1;
 567     }
 568   }
 569
 570   //
 571   // Check parameter
 572   //
 573   if (VolumeLetter == 0) {
 574     fprintf (stderr, "ERROR: Volume isn't provided!\n");
 575     PrintUsage (AppName);
 576     return 1;
 577   }
 578
 579   if (FilePath == NULL) {
 580     fprintf (stderr, "ERROR: File isn't pvovided!\n");
 581     PrintUsage (AppName);
 582     return 1;
 583   }
 584
 585   PatchType = PatchTypeUnknown;
 586
 587   if ((VolumeLetter == 'A') || (VolumeLetter == 'a') ||
 588       (VolumeLetter == 'B') || (VolumeLetter == 'b')
 589       ) {
 590     //
 591     // Floppy
 592     //
 593     sprintf (DiskPath, FloppyPathTemplate, VolumeLetter);
 594     PatchType = PatchTypeFloppy;
 595   }
 596   else {
 597     //
 598     // Hard/USB disk
 599     //
 600     if (!GetDriveInfo (VolumeLetter, &DriveInfo)) {
 601       fprintf (stderr, "ERROR: GetDriveInfo - 0x%x\n", GetLastError ());
 602       return 1;
 603     }
 604
 605     //
 606     // Shouldn't patch my own hard disk, but can read it.
 607     // very safe then:)
 608     //
 609     if (DriveInfo.DriveType->Type == DRIVE_FIXED && WriteToDisk) {
 610       fprintf (stderr, "ERROR: Write to local harddisk - permission denied!\n");
 611       return 1;
 612     }
 613
 614     sprintf (DiskPath, DiskPathTemplate, DriveInfo.DiskNumber);
 615     if (DriveInfo.DriveType->Type == DRIVE_REMOVABLE) {
 616       PatchType = PatchTypeUsb;
 617     }
 618     else if (DriveInfo.DriveType->Type == DRIVE_FIXED) {
 619       PatchType = PatchTypeIde;
 620     }
 621   }
 622
 623   if (PatchType == PatchTypeUnknown) {
 624     fprintf (stderr, "ERROR: PatchType unknown!\n");
 625     return 1;
 626   }
 627
 628   //
 629   // Process DBR (Patch or Read)
 630   //
 631   Status = ProcessBsOrMbr (DiskPath, FilePath, WriteToDisk, PatchType, ProcessMbr);
 632   if (Status == ErrorSuccess) {
 633     fprintf (
 634       stdout,
 635       "%s %s: successfully!\n",
 636       WriteToDisk ? "Write" : "Read",
 637       ProcessMbr ? "MBR" : "DBR"
 638       );
 639     return 0;
 640   } else {
 641     fprintf (
 642       stderr,
 643       "%s: %s %s: failed - %s (LastError: 0x%x)!\n",
 644       (Status == ErrorNoMbr) ? "WARNING" : "ERROR",
 645       WriteToDisk ? "Write" : "Read",
 646       ProcessMbr ? "MBR" : "DBR",
 647       ErrorStatusDesc[Status],
 648       GetLastError ()
 649       );
 650     return 1;
 651   }
 652 }