nerd/rufus
1
1
Fork 0
mirror of https://github.com/pbatard/rufus.git synced 2024-08-14 23:57:05 +00:00
Code Issues 1 Projects Releases Wiki Activity

[partition] initial partitioning code

Browse source 
* 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
This commit is contained in:
Pete Batard 2011-11-23 12:27:51 +00:00
parent af0fef1c8c
commit 45406054cc
3 changed files with 179 additions and 41 deletions
Download patch file Download diff file Expand all files Collapse all files

4
.msvc/rufus_2010.vcxproj
View file

@ -69,6 +69,10 @@
<IntDir Condition="'$(Configuration)|$(Platform)'=='Debug|x64'">$(SolutionDir)x86_64\$(Configuration)\$(ProjectName)\</IntDir>
<OutDir Condition="'$(Configuration)|$(Platform)'=='Release|x64'">$(SolutionDir)x86_64\$(Configuration)\</OutDir>
<IntDir Condition="'$(Configuration)|$(Platform)'=='Release|x64'">$(SolutionDir)x86_64\$(Configuration)\$(ProjectName)\</IntDir>
<LinkIncremental Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'">false</LinkIncremental>
<LinkIncremental Condition="'$(Configuration)|$(Platform)'=='Release|Win32'">false</LinkIncremental>
<LinkIncremental Condition="'$(Configuration)|$(Platform)'=='Debug|x64'">false</LinkIncremental>
<LinkIncremental Condition="'$(Configuration)|$(Platform)'=='Release|x64'">false</LinkIncremental>
</PropertyGroup>
<ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'">
<BuildLog />

191
rufus.c
View file

@ -40,6 +40,9 @@
//#include <fmifs.h>
// 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;
}

25
rufus.h
View file

@ -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;