From 45406054ccdbba56e736f1d64bc159dc4a56b928 Mon Sep 17 00:00:00 2001 From: Pete Batard Date: Wed, 23 Nov 2011 12:27:51 +0000 Subject: [PATCH] [partition] initial partitioning code * also added bWriteAccess parameter to GetDriveHandle() * also added locking (FSCTL_LOCK_VOLUME) to GetDriveHandle() * also added extra security checks to GetDriveHandle() * also added status output * also added ReadSectors/WriteSectors --- .msvc/rufus_2010.vcxproj | 4 + rufus.c | 191 +++++++++++++++++++++++++++++++++++---- rufus.h | 25 +---- 3 files changed, 179 insertions(+), 41 deletions(-) diff --git a/.msvc/rufus_2010.vcxproj b/.msvc/rufus_2010.vcxproj index 95b8c38e..8ab4c692 100644 --- a/.msvc/rufus_2010.vcxproj +++ b/.msvc/rufus_2010.vcxproj @@ -69,6 +69,10 @@ $(SolutionDir)x86_64\$(Configuration)\$(ProjectName)\ $(SolutionDir)x86_64\$(Configuration)\ $(SolutionDir)x86_64\$(Configuration)\$(ProjectName)\ + false + false + false + false diff --git a/rufus.c b/rufus.c index 31f9501e..e2a678eb 100644 --- a/rufus.c +++ b/rufus.c @@ -40,6 +40,9 @@ //#include // http://git.kernel.org/?p=fs/ext2/e2fsprogs.git;a=blob;f=misc/badblocks.c +// http://ms-sys.sourceforge.net/ +// http://thestarman.pcministry.com/asm/mbr/MSWIN41.htm + #include "msapi_utf8.h" #include "resource.h" #include "rufus.h" @@ -85,6 +88,26 @@ void _uprintf(const char *format, ...) #endif +void StatusPrintf(const char *format, ...) +{ + char buf[256], *p = buf; + va_list args; + int n; + + va_start(args, format); + n = safe_vsnprintf(p, sizeof(buf)-1, format, args); // room for NUL + va_end(args); + + p += (n < 0)?sizeof(buf)-1:n; + + while((p>buf) && (isspace(p[-1]))) + *--p = '\0'; + + *p = '\0'; + + SetDlgItemTextU(hMainDialog, IDC_STATUS, buf); +} + /* * Convert a partition type to its human readable form * http://www.win.tue.nl/~aeb/partitions/partition_types-1.html @@ -170,17 +193,29 @@ static const char* GetPartitionType(BYTE Type) } /* - * Open a drive - return both the handle and the drive letter + * Open a drive - return a drive HANDLE and the drive letter + * This call is quite risky (left unchecked, inadvertently passing 0 as index would + * return a handle to C:, which we might then proceed to unknowingly repartition!), + * so we apply the following mitigation factors: + * - Valid indexes must belong to a specific range [DRIVE_INDEX_MIN; DRIVE_INDEX_MAX] + * - When opening for write access, we lock the volume. If that fails, which would + * typically be the case on C:\ or any other drive in use, we fail the call + * - We report the full path of any drive that was successfully opened for write acces */ -static BOOL GetDriveHandle(DWORD num, HANDLE* hDrive, char* DriveLetter) +static BOOL GetDriveHandle(DWORD DriveIndex, HANDLE* hDrive, char* DriveLetter, BOOL bWriteAccess) { BOOL r; DWORD size; STORAGE_DEVICE_NUMBER_REDEF device_number = {0}; - static char drives[26*4]; /* "D:\", "E:\", etc. */ + char drives[26*4]; /* "D:\", "E:\", etc. */ char *drive = drives; char drive_name[] = "\\\\.\\#:"; + if ((DriveIndex < DRIVE_INDEX_MIN) || (DriveIndex > DRIVE_INDEX_MAX)) { + uprintf("WARNING: Bad index value. Please check your code!\n"); + } + DriveIndex -= DRIVE_INDEX_MIN; + size = GetLogicalDriveStringsA(sizeof(drives), drives); if (size == 0) { uprintf("GetLogicalDriveStrings failed: %s\n", WindowsErrorString()); @@ -197,12 +232,13 @@ static BOOL GetDriveHandle(DWORD num, HANDLE* hDrive, char* DriveLetter) continue; } safe_sprintf(drive_name, sizeof(drive_name), "\\\\.\\%c:", drive[0]); - *hDrive = CreateFileA(drive_name, GENERIC_READ, FILE_SHARE_READ|FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, 0); + *hDrive = CreateFileA(drive_name, GENERIC_READ|(bWriteAccess?GENERIC_WRITE:0), + FILE_SHARE_READ|FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, 0); if (hDrive == INVALID_HANDLE_VALUE) { - uprintf("Could not open drive %c: %s\n", WindowsErrorString()); + uprintf("Could not open drive %c: %s\n", drive[0], WindowsErrorString()); continue; } - + r = DeviceIoControl(*hDrive, IOCTL_STORAGE_GET_DEVICE_NUMBER, NULL, 0, &device_number, sizeof(device_number), &size, NULL); if ((!r) || (size <= 0)) { @@ -210,12 +246,22 @@ static BOOL GetDriveHandle(DWORD num, HANDLE* hDrive, char* DriveLetter) safe_closehandle(*hDrive); break; } - if (device_number.DeviceNumber == num) + if (device_number.DeviceNumber == DriveIndex) { + if (bWriteAccess) { + if (!DeviceIoControl(*hDrive, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &size, NULL)) { + uprintf("Could not get exclusive access to %s: %s\n", drive_name, WindowsErrorString()); + safe_closehandle(*hDrive); + return FALSE; + } + uprintf("Warning: Opening %s drive for write access\n", drive_name); + } break; + } + safe_closehandle(*hDrive); } if (DriveLetter != NULL) { - *DriveLetter = *drive; + *DriveLetter = *drive?*drive:' '; // TODO: handle NUL char upstream } return (*hDrive != INVALID_HANDLE_VALUE); @@ -232,15 +278,13 @@ static BOOL GetDriveLabel(DWORD num, char* letter, char** label) *label = "NO_LABEL"; - if (!GetDriveHandle(num, &hDrive, DrivePath)) + if (!GetDriveHandle(num, &hDrive, DrivePath, FALSE)) return FALSE; safe_closehandle(hDrive); *letter = DrivePath[0]; if (GetVolumeInformationA(DrivePath, volume_label, sizeof(volume_label), NULL, NULL, NULL, NULL, 0)) { *label = volume_label; - } else { - uprintf("GetVolumeInformation (Label) failed: %s\n", WindowsErrorString()); } return TRUE; @@ -264,7 +308,7 @@ static BOOL GetDriveInfo(DWORD num, LONGLONG* DriveSize, char* FSType, DWORD FST *DriveSize = 0; - if (!GetDriveHandle(num, &hDrive, DrivePath)) + if (!GetDriveHandle(num, &hDrive, DrivePath, FALSE)) return FALSE; r = DeviceIoControl(hDrive, IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, @@ -275,6 +319,9 @@ static BOOL GetDriveInfo(DWORD num, LONGLONG* DriveSize, char* FSType, DWORD FST return FALSE; } *DriveSize = DiskGeometry->DiskSize.QuadPart; + uprintf("Cylinders: %lld, TracksPerCylinder: %d, SectorsPerTrack: %d, BytesPerSector: %d\n", + DiskGeometry->Geometry.Cylinders, DiskGeometry->Geometry.TracksPerCylinder, + DiskGeometry->Geometry.SectorsPerTrack, DiskGeometry->Geometry.BytesPerSector); r = DeviceIoControl(hDrive, IOCTL_DISK_GET_DRIVE_LAYOUT_EX, NULL, 0, layout, sizeof(layout), &size, NULL ); @@ -309,7 +356,7 @@ static BOOL GetDriveInfo(DWORD num, LONGLONG* DriveSize, char* FSType, DWORD FST if (!GetVolumeInformationA(DrivePath, NULL, 0, NULL, NULL, NULL, FSType, FSTypeSize)) { safe_sprintf(FSType, FSTypeSize, "Non Windows (Please Select)"); - uprintf("GetVolumeInformation (Properties) failed: %s\n", WindowsErrorString()); +// uprintf("GetVolumeInformation (Properties) failed: %s\n", WindowsErrorString()); } return TRUE; @@ -352,6 +399,109 @@ static BOOL PopulateProperties(int index) return TRUE; } +BOOL WriteSectors(HANDLE hDrive, size_t SectorSize, size_t StartSector, size_t nSectors, void* Buf, size_t BufSize) +{ + LARGE_INTEGER ptr; + DWORD Size; + + if (SectorSize * nSectors > BufSize) { + uprintf("WriteSectors: Buffer is too small\n"); + return FALSE; + } + + ptr.QuadPart = StartSector*SectorSize; + if (!SetFilePointerEx(hDrive, ptr, NULL, FILE_BEGIN)) { + uprintf("WriteSectors: Could not access sector %lld - %s\n", StartSector, WindowsErrorString()); + return FALSE; + } + + if ((!WriteFile(hDrive, Buf, BufSize, &Size, NULL)) || (Size != BufSize)) { + uprintf("WriteSectors: Write error - %s\n", WindowsErrorString()); + return FALSE; + } + + return TRUE; +} + +BOOL ReadSectors(HANDLE hDrive, size_t SectorSize, size_t StartSector, size_t nSectors, void* Buf, size_t BufSize) +{ + LARGE_INTEGER ptr; + DWORD size; + + if (SectorSize * nSectors > BufSize) { + uprintf("ReadSectors: Buffer is too small\n"); + return FALSE; + } + + ptr.QuadPart = StartSector*SectorSize; + if (!SetFilePointerEx(hDrive, ptr, NULL, FILE_BEGIN)) { + uprintf("ReadSectors: Could not access sector %lld - %s\n", StartSector, WindowsErrorString()); + return FALSE; + } + + if ((!ReadFile(hDrive, Buf, BufSize, &size, NULL)) || (size != BufSize)) { + uprintf("ReadSectors: Write error - %s\n", WindowsErrorString()); + return FALSE; + } + + return TRUE; +} + +/* + * Create a partition table + */ +BOOL CreatePartition(HANDLE hDrive) +{ + BYTE layout[sizeof(DRIVE_LAYOUT_INFORMATION_EX) + 3*sizeof(PARTITION_INFORMATION_EX)] = {0}; + PDRIVE_LAYOUT_INFORMATION_EX DriveLayoutEx = (PDRIVE_LAYOUT_INFORMATION_EX)layout; + BOOL r; + DWORD size; + int nbHidden = 63; + + DriveLayoutEx->PartitionStyle = PARTITION_STYLE_MBR; + DriveLayoutEx->PartitionCount = 4; // Must be multiple of 4 for MBR + DriveLayoutEx->Mbr.Signature = GetTickCount(); + DriveLayoutEx->PartitionEntry[0].PartitionStyle = PARTITION_STYLE_MBR; + DriveLayoutEx->PartitionEntry[0].StartingOffset.QuadPart = nbHidden*512; // TODO + DriveLayoutEx->PartitionEntry[0].PartitionLength.QuadPart = 63*2*512; // TODO + DriveLayoutEx->PartitionEntry[0].PartitionNumber = 1; + DriveLayoutEx->PartitionEntry[0].RewritePartition = TRUE; + DriveLayoutEx->PartitionEntry[0].Mbr.PartitionType = 0x83; // TODO + DriveLayoutEx->PartitionEntry[0].Mbr.HiddenSectors = nbHidden; // TODO + + // For the remaining partitions, PartitionType has already been zeroed (= set to unused) + DriveLayoutEx->PartitionEntry[1].PartitionStyle = PARTITION_STYLE_MBR; + DriveLayoutEx->PartitionEntry[2].PartitionStyle = PARTITION_STYLE_MBR; + DriveLayoutEx->PartitionEntry[3].PartitionStyle = PARTITION_STYLE_MBR; + + r = DeviceIoControl(hDrive, IOCTL_DISK_SET_DRIVE_LAYOUT_EX, + layout, sizeof(layout), NULL, 0, &size, NULL ); + if (!r) { + uprintf("IOCTL_DISK_SET_DRIVE_LAYOUT_EX failed: %s\n", WindowsErrorString()); + safe_closehandle(hDrive); + return FALSE; + } + StatusPrintf("Successfully Created Partition"); + + return TRUE; +} + +BOOL FormatDrive(DWORD num) +{ + HANDLE hDrive; + BOOL r; + + if (!GetDriveHandle(num, &hDrive, NULL, TRUE)) { + // TODO: report an error for exclusive access + return FALSE; + } + + r = CreatePartition(hDrive); + + CloseHandle(hDrive); + return r; +} + /* * Refresh the list of USB devices */ @@ -395,7 +545,7 @@ static BOOL GetUSBDevices(void) uprintf("SetupDiGetDeviceRegistryProperty (Friendly Name) failed: %d\n", WindowsErrorString()); continue; } - uprintf("found drive '%s'\n", buffer); + uprintf("Found drive '%s'\n", buffer); StrArrayAdd(&DriveID, buffer); devint_data.cbSize = sizeof(devint_data); @@ -444,9 +594,10 @@ static BOOL GetUSBDevices(void) continue; } - if (GetDriveLabel(device_number.DeviceNumber, &drive_letter, &label)) { + if (GetDriveLabel(device_number.DeviceNumber + DRIVE_INDEX_MIN, &drive_letter, &label)) { safe_sprintf(entry, sizeof(entry), "%s (%c:)\n", label, drive_letter); - IGNORE_RETVAL(ComboBox_SetItemData(hDeviceList, ComboBox_AddStringU(hDeviceList, entry), device_number.DeviceNumber)); + IGNORE_RETVAL(ComboBox_SetItemData(hDeviceList, ComboBox_AddStringU(hDeviceList, entry), + device_number.DeviceNumber + DRIVE_INDEX_MIN)); } } } @@ -456,8 +607,6 @@ static BOOL GetUSBDevices(void) return TRUE; } -// TODO: the device is currently in use by another application (find application a la TGit installer?) - /* * Main dialog callback */ @@ -528,6 +677,12 @@ static INT_PTR CALLBACK MainCallback(HWND hDlg, UINT message, WPARAM wParam, LPA break; } break; + case IDC_START: + nDeviceIndex = ComboBox_GetCurSel(hDeviceList); + if (nDeviceIndex != CB_ERR) { + FormatDrive(ComboBox_GetItemData(hDeviceList, nDeviceIndex)); + } + break; default: return (INT_PTR)FALSE; } diff --git a/rufus.h b/rufus.h index e80a0ef5..8175e636 100644 --- a/rufus.h +++ b/rufus.h @@ -21,6 +21,8 @@ #define RUFUS_DEBUG #define APP_VERSION "Rufus v1.0.0.1" +#define DRIVE_INDEX_MIN 0x80 +#define DRIVE_INDEX_MAX 0xC0 #define MAX_TOOLTIPS 16 #define WHITE RGB(255,255,255) #define SEPARATOR_GREY RGB(223,223,223) @@ -103,26 +105,3 @@ typedef struct { ULONG DeviceNumber; ULONG PartitionNumber; } STORAGE_DEVICE_NUMBER_REDEF; - -typedef struct _SCSI_PASS_THROUGH { - USHORT Length; - UCHAR ScsiStatus; - UCHAR PathId; - UCHAR TargetId; - UCHAR Lun; - UCHAR CdbLength; - UCHAR SenseInfoLength; - UCHAR DataIn; - ULONG DataTransferLength; - ULONG TimeOutValue; - ULONG_PTR DataBufferOffset; - ULONG SenseInfoOffset; - UCHAR Cdb[16]; -} SCSI_PASS_THROUGH,*PSCSI_PASS_THROUGH; - -typedef struct _SCSI_PASS_THROUGH_WITH_BUFFERS { - SCSI_PASS_THROUGH Spt; - ULONG Filler; - UCHAR SenseBuf[32]; - UCHAR DataBuf[512]; -} SCSI_PASS_THROUGH_WITH_BUFFERS, *PSCSI_PASS_THROUGH_WITH_BUFFERS;