[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

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

rufus.c

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

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;