nerd
/
rufus
mirror of https://github.com/pbatard/rufus.git
1
1
Fork 0
Code Issues 1 Projects Releases Wiki Activity
rufus/src/rufus.c

1095 lines
36 KiB
C
Raw Blame History

/*
* Rufus: The Reliable USB Formatting Utility
* Copyright (c) 2011 Pete Batard <pete@akeo.ie>
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
/* Memory leaks detection - define _CRTDBG_MAP_ALLOC as preprocessor macro */
#ifdef _CRTDBG_MAP_ALLOC
#include <stdlib.h>
#include <crtdbg.h>
#endif
#include <windows.h>
#include <windowsx.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <math.h>
#include <commctrl.h>
#include <setupapi.h>
#include <winioctl.h>
#include <process.h>
#include <dbt.h>
#include "msapi_utf8.h"
#include "resource.h"
#include "rufus.h"
#include "sys_types.h"
static const char* FileSystemLabel[FS_MAX] = { "FAT", "FAT32", "NTFS", "exFAT" };
// Don't ask me - just following the MS standard here
static const char* ClusterSizeLabel[] = { "512 bytes", "1024 bytes","2048 bytes","4096 bytes","8192 bytes",
"16 kilobytes", "32 kilobytes", "64 kilobytes", "128 kilobytes", "256 kilobytes", "512 kilobytes",
"1024 kilobytes","2048 kilobytes","4096 kilobytes","8192 kilobytes","16 megabytes","32 megabytes" };
/*
* Globals
*/
HINSTANCE hMainInstance;
HWND hMainDialog;
char szFolderPath[MAX_PATH];
float fScale = 1.0f;
int default_fs;
HWND hDeviceList, hCapacity, hFileSystem, hClusterSize, hLabel;
static HWND hDeviceTooltip = NULL, hFSTooltip = NULL, hProgress = NULL;
static StrArray DriveID, DriveLabel;
static char szTimer[10] = "00:00:00";
static unsigned int timer;
/*
* The following is used to allocate slots within the progress bar
* 0 means unused (no operation or no progress allocated to it)
* +n means allocate exactly n bars (n percents of the progress bar)
* -n means allocate a weighted slot of n from all remaining
* bars. Eg if 80 slots remain and the sum of all negative entries
* is 10, -4 will allocate 4/10*80 = 32 bars (32%) for OP progress
*/
static int nb_slots[OP_MAX];
static float slot_end[OP_MAX+1]; // shifted +1 so that we can substract 1 to OP indexes
static float previous_end = 0.0f;
/*
* Convert a partition type to its human readable form using
* (slightly modified) entries from GNU fdisk
*/
static const char* GetPartitionType(BYTE Type)
{
int i;
for (i=0; i<ARRAYSIZE(msdos_systypes); i++) {
if (msdos_systypes[i].type == Type)
return msdos_systypes[i].name;
}
return "Unknown";
}
#define KB 1024LL
#define MB 1048576LL
#define GB 1073741824LL
#define TB 1099511627776LL
/*
* Set cluster size values according to http://support.microsoft.com/kb/140365
* this call will return FALSE if we can't find a supportable FS for the drive
*/
static BOOL DefineClusterSizes(void)
{
LONGLONG i;
int fs;
BOOL r = FALSE;
char tmp[64] = "";
default_fs = FS_UNKNOWN;
memset(&SelectedDrive.ClusterSize, 0, sizeof(SelectedDrive.ClusterSize));
if (SelectedDrive.DiskSize < 8*MB) {
// TODO: muck with FAT12 and Small FAT16 like Microsoft does
uprintf("This application does not support volumes smaller than 8 MB yet\n");
goto out;
}
// FAT 16
if (SelectedDrive.DiskSize < 4*GB) {
// TODO: Refine the following according to size
SelectedDrive.ClusterSize[FS_FAT16].Allowed = 0x0001FE00;
for (i=32; i<=4096; i<<=1) { // 8 MB -> 4 GB
if (SelectedDrive.DiskSize < i*MB) {
SelectedDrive.ClusterSize[FS_FAT16].Default = 16*(ULONG)i;
break;
}
}
}
// FAT 32
if (SelectedDrive.DiskSize < 256*MB) {
// TODO: Refine the following according to size
SelectedDrive.ClusterSize[FS_FAT32].Allowed = 0x0001FE00;
for (i=64; i<=256; i<<=1) { // 8 MB -> 256 MB
if (SelectedDrive.DiskSize < i*MB) {
SelectedDrive.ClusterSize[FS_FAT32].Default = 8*(ULONG)i;
break;
}
}
} else if (SelectedDrive.DiskSize < 32*GB) {
SelectedDrive.ClusterSize[FS_FAT32].Allowed = 0x0001FE00;
for (i=8; i<=32; i<<=1) { // 256 MB -> 32 GB
if (SelectedDrive.DiskSize < i*GB) {
SelectedDrive.ClusterSize[FS_FAT32].Default = ((ULONG)i/2)*1024;
break;
}
}
}
// NTFS
if (SelectedDrive.DiskSize < 256*TB) {
SelectedDrive.ClusterSize[FS_NTFS].Allowed = 0x0001FE00;
for (i=16; i<=256; i<<=1) { // 7 MB -> 256 TB
if (SelectedDrive.DiskSize < i*TB) {
SelectedDrive.ClusterSize[FS_NTFS].Default = ((ULONG)i/4)*1024;
break;
}
}
}
// exFAT
if (SelectedDrive.DiskSize < 256*TB) {
SelectedDrive.ClusterSize[FS_EXFAT].Allowed = 0x03FFFE00;
if (SelectedDrive.DiskSize < 256*MB) // < 256 MB
SelectedDrive.ClusterSize[FS_EXFAT].Default = 4*1024;
else if (SelectedDrive.DiskSize < 32*GB) // < 32 GB
SelectedDrive.ClusterSize[FS_EXFAT].Default = 32*1024;
else
SelectedDrive.ClusterSize[FS_EXFAT].Default = 28*1024;
}
out:
// Only add the filesystems we can service
for (fs=0; fs<FS_MAX; fs++) {
if (SelectedDrive.ClusterSize[fs].Allowed != 0) {
safe_sprintf(tmp, sizeof(tmp), FileSystemLabel[fs]);
if (default_fs == FS_UNKNOWN) {
safe_strcat(tmp, sizeof(tmp), " (Default)");
default_fs = fs;
}
IGNORE_RETVAL(ComboBox_SetItemData(hFileSystem,
ComboBox_AddStringU(hFileSystem, tmp), fs));
r = TRUE;
}
}
return r;
}
#undef KB
#undef MB
#undef GB
#undef TB
/*
* Populate the Allocation unit size field
*/
static BOOL SetClusterSizes(int FSType)
{
char szClustSize[64];
int i, default_index = 0;
ULONG j;
IGNORE_RETVAL(ComboBox_ResetContent(hClusterSize));
if ((FSType < 0) || (FSType >= FS_MAX)) {
uprintf("Invalid FS value passed to SetClusterSizes\n");
return FALSE;
}
if ( (SelectedDrive.ClusterSize[FSType].Allowed == 0)
|| (SelectedDrive.ClusterSize[FSType].Default == 0) ) {
uprintf("The drive is incompatible with FS type #%d\n", FSType);
return FALSE;
}
for(i=0,j=0x200;j<0x10000000;i++,j<<=1) {
if (j & SelectedDrive.ClusterSize[FSType].Allowed) {
safe_sprintf(szClustSize, sizeof(szClustSize), "%s", ClusterSizeLabel[i]);
if (j == SelectedDrive.ClusterSize[FSType].Default) {
safe_strcat(szClustSize, sizeof(szClustSize), " (Default)");
default_index = i;
}
IGNORE_RETVAL(ComboBox_SetItemData(hClusterSize, ComboBox_AddStringU(hClusterSize, szClustSize), j));
}
}
IGNORE_RETVAL(ComboBox_SetCurSel(hClusterSize, default_index));
return TRUE;
}
/*
* Fill the drive properties (size, FS, etc)
*/
static BOOL GetDriveInfo(void)
{
BOOL r;
HANDLE hDrive;
DWORD size;
BYTE geometry[128], layout[1024];
void* disk_geometry = (void*)geometry;
void* drive_layout = (void*)layout;
PDISK_GEOMETRY_EX DiskGeometry = (PDISK_GEOMETRY_EX)disk_geometry;
PDRIVE_LAYOUT_INFORMATION_EX DriveLayout = (PDRIVE_LAYOUT_INFORMATION_EX)drive_layout;
char DrivePath[] = "#:\\", tmp[128], fs_type[32];
DWORD i, nb_partitions = 0;
SelectedDrive.DiskSize = 0;
hDrive = GetDriveHandle(SelectedDrive.DeviceNumber, DrivePath, FALSE, FALSE);
if (hDrive == INVALID_HANDLE_VALUE)
return FALSE;
r = DeviceIoControl(hDrive, IOCTL_DISK_GET_DRIVE_GEOMETRY_EX,
NULL, 0, geometry, sizeof(geometry), &size, NULL );
if (!r || size <= 0) {
uprintf("IOCTL_DISK_GET_DRIVE_GEOMETRY_EX failed: %s\n", WindowsErrorString());
safe_closehandle(hDrive);
return FALSE;
}
SelectedDrive.DiskSize = DiskGeometry->DiskSize.QuadPart;
memcpy(&SelectedDrive.Geometry, &DiskGeometry->Geometry, sizeof(DISK_GEOMETRY));
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 );
if (!r || size <= 0) {
uprintf("IOCTL_DISK_GET_DRIVE_LAYOUT_EX failed: %s\n", WindowsErrorString());
} else {
DestroyTooltip(hFSTooltip);
hFSTooltip = NULL;
switch (DriveLayout->PartitionStyle) {
case PARTITION_STYLE_MBR:
for (i=0; i<DriveLayout->PartitionCount; i++) {
if (DriveLayout->PartitionEntry[i].Mbr.PartitionType != PARTITION_ENTRY_UNUSED) {
uprintf("Partition #%d:\n", ++nb_partitions);
if (hFSTooltip == NULL) {
// TODO: provide all partitions FS on tooltip, not just the one
safe_sprintf(tmp, sizeof(tmp), "Current file system: %s (0x%02x)",
GetPartitionType(DriveLayout->PartitionEntry[i].Mbr.PartitionType),
DriveLayout->PartitionEntry[i].Mbr.PartitionType);
hFSTooltip = CreateTooltip(hFileSystem, tmp, -1);
}
uprintf(" Type: %s (0x%02x)\n Boot: %s\n Recognized: %s\n Hidden Sectors: %d\n",
GetPartitionType(DriveLayout->PartitionEntry[i].Mbr.PartitionType),
DriveLayout->PartitionEntry[i].Mbr.PartitionType,
DriveLayout->PartitionEntry[i].Mbr.BootIndicator?"Yes":"No",
DriveLayout->PartitionEntry[i].Mbr.RecognizedPartition?"Yes":"No",
DriveLayout->PartitionEntry[i].Mbr.HiddenSectors);
}
}
uprintf("Partition type: MBR, NB Partitions: %d\n", nb_partitions);
break;
case PARTITION_STYLE_GPT:
uprintf("Partition type: GPT, NB Partitions: %d\n", DriveLayout->PartitionCount);
break;
default:
uprintf("Partition type: RAW\n");
break;
}
}
safe_closehandle(hDrive);
if (!DefineClusterSizes()) {
uprintf("no file system is selectable for this drive\n");
return FALSE;
}
// re-select existing FS if it's one we know
if (GetVolumeInformationA(DrivePath, NULL, 0, NULL, NULL, NULL,
fs_type, sizeof(fs_type))) {
for (SelectedDrive.FSType=FS_MAX-1; SelectedDrive.FSType>=0; SelectedDrive.FSType--) {
if (safe_strcmp(fs_type, FileSystemLabel[SelectedDrive.FSType]) == 0) {
break;
}
}
} else {
SelectedDrive.FSType = FS_UNKNOWN;
}
for (i=0; i<ComboBox_GetCount(hFileSystem); i++) {
if (ComboBox_GetItemData(hFileSystem, i) == SelectedDrive.FSType) {
IGNORE_RETVAL(ComboBox_SetCurSel(hFileSystem, i));
break;
}
}
if (i == ComboBox_GetCount(hFileSystem)) {
// failed to reselect => pick default
for (i=0; i<ComboBox_GetCount(hFileSystem); i++) {
if (ComboBox_GetItemData(hFileSystem, i) == default_fs) {
IGNORE_RETVAL(ComboBox_SetCurSel(hFileSystem, i));
break;
}
}
}
// At least one filesystem is go => enable formatting
EnableWindow(GetDlgItem(hMainDialog, IDC_START), TRUE);
return SetClusterSizes((int)ComboBox_GetItemData(hFileSystem, ComboBox_GetCurSel(hFileSystem)));
}
/*
* Populate the UI properties
*/
static BOOL PopulateProperties(int ComboIndex)
{
double HumanReadableSize;
char capacity[64];
static char *suffix[] = { "KB", "MB", "GB", "TB", "PB"};
char proposed_label[16], no_label[] = STR_NO_LABEL;
int i;
IGNORE_RETVAL(ComboBox_ResetContent(hCapacity));
IGNORE_RETVAL(ComboBox_ResetContent(hFileSystem));
IGNORE_RETVAL(ComboBox_ResetContent(hClusterSize));
EnableWindow(GetDlgItem(hMainDialog, IDC_START), FALSE);
SetWindowTextA(hLabel, "");
DestroyTooltip(hDeviceTooltip);
DestroyTooltip(hFSTooltip);
hDeviceTooltip = NULL;
hFSTooltip = NULL;
memset(&SelectedDrive, 0, sizeof(SelectedDrive));
if (ComboIndex < 0) {
return TRUE;
}
SelectedDrive.DeviceNumber = (DWORD)ComboBox_GetItemData(hDeviceList, ComboIndex);
if (!GetDriveInfo())
return FALSE;
HumanReadableSize = (double)SelectedDrive.DiskSize;
for (i=0; i<ARRAYSIZE(suffix); i++) {
HumanReadableSize /= 1024.0;
if (HumanReadableSize < 512.0) {
safe_sprintf(capacity, sizeof(capacity), "%0.2f %s", HumanReadableSize, suffix[i]);
break;
}
}
IGNORE_RETVAL(ComboBox_AddStringU(hCapacity, capacity));
IGNORE_RETVAL(ComboBox_SetCurSel(hCapacity, 0));
hDeviceTooltip = CreateTooltip(hDeviceList, DriveID.Table[ComboIndex], -1);
// If no existing label is available, propose one according to the size (eg: "256MB", "8GB")
if (safe_strcmp(no_label, DriveLabel.Table[ComboIndex]) == 0) {
if (HumanReadableSize < 1.0) {
HumanReadableSize *= 1024.0;
i--;
}
// If we're beneath the tolerance, round proposed label to an integer, if not, show one decimal point
if (fabs(HumanReadableSize / ceil(HumanReadableSize) - 1.0) < PROPOSEDLABEL_TOLERANCE) {
safe_sprintf(proposed_label, sizeof(proposed_label), "%0.0f%s", ceil(HumanReadableSize), suffix[i]);
} else {
safe_sprintf(proposed_label, sizeof(proposed_label), "%0.1f%s", HumanReadableSize, suffix[i]);
}
SetWindowTextA(hLabel, proposed_label);
} else {
SetWindowTextA(hLabel, DriveLabel.Table[ComboIndex]);
}
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;
PrintStatus(0, "Partitioning...");
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 =
SelectedDrive.Geometry.BytesPerSector * SelectedDrive.Geometry.SectorsPerTrack;
DriveLayoutEx->PartitionEntry[0].PartitionLength.QuadPart = SelectedDrive.DiskSize -
DriveLayoutEx->PartitionEntry[0].StartingOffset.QuadPart;
DriveLayoutEx->PartitionEntry[0].PartitionNumber = 1;
DriveLayoutEx->PartitionEntry[0].RewritePartition = TRUE;
DriveLayoutEx->PartitionEntry[0].Mbr.HiddenSectors = SelectedDrive.Geometry.SectorsPerTrack;
switch (ComboBox_GetItemData(hFileSystem, ComboBox_GetCurSel(hFileSystem))) {
case FS_FAT16:
DriveLayoutEx->PartitionEntry[0].Mbr.PartitionType = 0x0e; // FAT16 LBA
break;
case FS_NTFS:
case FS_EXFAT:
// TODO: but how do we set this thing up afterwards?
DriveLayoutEx->PartitionEntry[0].Mbr.PartitionType = 0x07; // NTFS
break;
default:
DriveLayoutEx->PartitionEntry[0].Mbr.PartitionType = 0x0c; // FAT32 LBA
break;
}
// For the remaining partitions, PartitionStyle & PartitionType have already
// been zeroed => set to MBR/unused
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;
}
return TRUE;
}
/*
* Refresh the list of USB devices
*/
static BOOL GetUSBDevices(void)
{
BOOL r;
HDEVINFO dev_info = NULL;
SP_DEVINFO_DATA dev_info_data;
SP_DEVICE_INTERFACE_DATA devint_data;
PSP_DEVICE_INTERFACE_DETAIL_DATA_A devint_detail_data;
STORAGE_DEVICE_NUMBER_REDEF device_number;
DWORD size, i, j, datatype;
HANDLE hDrive;
char drive_letter;
char *label, entry[MAX_PATH], buffer[MAX_PATH];
const char* usbstor_name = "USBSTOR";
GUID _GUID_DEVINTERFACE_DISK = // only known to some...
{ 0x53f56307L, 0xb6bf, 0x11d0, {0x94, 0xf2, 0x00, 0xa0, 0xc9, 0x1e, 0xfb, 0x8b} };
IGNORE_RETVAL(ComboBox_ResetContent(hDeviceList));
StrArrayClear(&DriveID);
StrArrayClear(&DriveLabel);
dev_info = SetupDiGetClassDevsA(&_GUID_DEVINTERFACE_DISK, NULL, NULL, DIGCF_PRESENT|DIGCF_DEVICEINTERFACE);
if (dev_info == INVALID_HANDLE_VALUE) {
uprintf("SetupDiGetClassDevs (Interface) failed: %d\n", WindowsErrorString());
return FALSE;
}
dev_info_data.cbSize = sizeof(dev_info_data);
for (i=0; SetupDiEnumDeviceInfo(dev_info, i, &dev_info_data); i++) {
memset(buffer, 0, sizeof(buffer));
if (!SetupDiGetDeviceRegistryPropertyA(dev_info, &dev_info_data, SPDRP_ENUMERATOR_NAME,
&datatype, (LPBYTE)buffer, sizeof(buffer), &size)) {
uprintf("SetupDiGetDeviceRegistryProperty (Enumerator Name) failed: %d\n", WindowsErrorString());
continue;
}
if (safe_strcmp(buffer, usbstor_name) != 0)
continue;
memset(buffer, 0, sizeof(buffer));
if (!SetupDiGetDeviceRegistryPropertyA(dev_info, &dev_info_data, SPDRP_FRIENDLYNAME,
&datatype, (LPBYTE)buffer, sizeof(buffer), &size)) {
uprintf("SetupDiGetDeviceRegistryProperty (Friendly Name) failed: %d\n", WindowsErrorString());
continue;
}
uprintf("Found drive '%s'\n", buffer);
devint_data.cbSize = sizeof(devint_data);
hDrive = INVALID_HANDLE_VALUE;
devint_detail_data = NULL;
for (j=0; ;j++) {
safe_closehandle(hDrive);
safe_free(devint_detail_data);
if (!SetupDiEnumDeviceInterfaces(dev_info, &dev_info_data, &_GUID_DEVINTERFACE_DISK, j, &devint_data)) {
if(GetLastError() != ERROR_NO_MORE_ITEMS) {
uprintf("SetupDiEnumDeviceInterfaces failed: %s\n", WindowsErrorString());
}
break;
}
if (!SetupDiGetDeviceInterfaceDetailA(dev_info, &devint_data, NULL, 0, &size, NULL)) {
if(GetLastError() == ERROR_INSUFFICIENT_BUFFER) {
devint_detail_data = (PSP_DEVICE_INTERFACE_DETAIL_DATA_A)calloc(1, size);
if (devint_detail_data == NULL) {
uprintf("unable to allocate data for SP_DEVICE_INTERFACE_DETAIL_DATA\n");
return FALSE;
}
devint_detail_data->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA_A);
} else {
uprintf("SetupDiGetDeviceInterfaceDetail (dummy) failed: %s\n", WindowsErrorString());
continue;
}
}
if(!SetupDiGetDeviceInterfaceDetailA(dev_info, &devint_data, devint_detail_data, size, &size, NULL)) {
uprintf("SetupDiGetDeviceInterfaceDetail (actual) failed: %s\n", WindowsErrorString());
continue;
}
hDrive = CreateFileA(devint_detail_data->DevicePath, GENERIC_READ|GENERIC_WRITE, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
if(hDrive == INVALID_HANDLE_VALUE) {
uprintf("could not open '%s': %s\n", devint_detail_data->DevicePath, WindowsErrorString());
continue;
}
memset(&device_number, 0, sizeof(device_number));
r = DeviceIoControl(hDrive, IOCTL_STORAGE_GET_DEVICE_NUMBER,
NULL, 0, &device_number, sizeof(device_number), &size, NULL );
if (!r || size <= 0) {
uprintf("IOCTL_STORAGE_GET_DEVICE_NUMBER (GetUSBDevices) failed: %s\n", WindowsErrorString());
continue;
}
if (GetDriveLabel(device_number.DeviceNumber + DRIVE_INDEX_MIN, &drive_letter, &label)) {
// Must ensure that the combo box is UNSORTED for indexes to be the same
StrArrayAdd(&DriveID, buffer);
StrArrayAdd(&DriveLabel, label);
safe_sprintf(entry, sizeof(entry), "%s (%c:)", label, drive_letter);
IGNORE_RETVAL(ComboBox_SetItemData(hDeviceList, ComboBox_AddStringU(hDeviceList, entry),
device_number.DeviceNumber + DRIVE_INDEX_MIN));
safe_closehandle(hDrive);
safe_free(devint_detail_data);
break;
}
}
}
SetupDiDestroyDeviceInfoList(dev_info);
IGNORE_RETVAL(ComboBox_SetCurSel(hDeviceList, 0));
SendMessage(hMainDialog, WM_COMMAND, (CBN_SELCHANGE<<16) | IDC_DEVICE, 0);
SendMessage(hMainDialog, WM_COMMAND, (CBN_SELCHANGE<<16) | IDC_FILESYSTEM,
ComboBox_GetCurSel(hFileSystem));
return TRUE;
}
/*
* Set up progress bar real estate allocation
*/
static void InitProgress(void)
{
int i;
float last_end = 0.0f, slots_discrete = 0.0f, slots_analog = 0.0f;
memset(&nb_slots, 0, sizeof(nb_slots));
memset(&slot_end, 0, sizeof(slot_end));
previous_end = 0.0f;
nb_slots[OP_ZERO_MBR] = 1;
if (IsChecked(IDC_BADBLOCKS)) {
nb_slots[OP_BADBLOCKS] = -1;
}
if (IsChecked(IDC_DOS)) {
// TODO: this should reflect the number of files to copy +1 for PBR writing
nb_slots[OP_DOS] = 3+1;
}
nb_slots[OP_PARTITION] = 1;
nb_slots[OP_FIX_MBR] = 1;
nb_slots[OP_CREATE_FS] =
nb_steps[ComboBox_GetItemData(hFileSystem, ComboBox_GetCurSel(hFileSystem))];
if (!IsChecked(IDC_QUICKFORMAT)) {
nb_slots[OP_FORMAT] = -1;
}
for (i=0; i<OP_MAX; i++) {
if (nb_slots[i] > 0) {
slots_discrete += nb_slots[i]*1.0f;
}
if (nb_slots[i] < 0) {
slots_analog += nb_slots[i]*1.0f;
}
}
for (i=0; i<OP_MAX; i++) {
if (nb_slots[i] == 0) {
slot_end[i+1] = last_end;
} else if (nb_slots[i] > 0) {
slot_end[i+1] = last_end + (1.0f * nb_slots[i]);
} else if (nb_slots[i] < 0) {
slot_end[i+1] = last_end + (( (100.0f-slots_discrete) * nb_slots[i]) / slots_analog);
}
last_end = slot_end[i+1];
}
/* Is there's no analog, adjust our discrete ends to fill the whole bar */
if (slots_analog == 0.0f) {
for (i=0; i<OP_MAX; i++) {
slot_end[i+1] *= 100.0f / slots_discrete;
}
}
}
/*
* Position the progress bar within each operation range
*/
void UpdateProgress(int op, float percent)
{
int pos;
if ((op < 0) || (op > OP_MAX)) {
uprintf("UpdateProgress: invalid op %d\n", op);
return;
}
if (percent > 100.1f) {
uprintf("UpdateProgress(%d): invalid percentage %0.2f\n", op, percent);
return;
}
if ((percent < 0.0f) && (nb_slots[op] <= 0)) {
uprintf("UpdateProgress(%d): error negative percentage sent for negative slot value\n", op);
return;
}
if (nb_slots[op] == 0)
return;
if (previous_end < slot_end[op]) {
previous_end = slot_end[op];
}
if (percent < 0.0f) {
// Negative means advance one slot (1.0%) - requires a positive slot allocation
previous_end += (slot_end[op+1] - slot_end[op]) / (1.0f * nb_slots[op]);
pos = (int)(previous_end / 100.0f * MAX_PROGRESS);
} else {
pos = (int)((previous_end + ((slot_end[op+1] - previous_end) * (percent / 100.0f))) / 100.0f * MAX_PROGRESS);
}
if (pos > MAX_PROGRESS) {
uprintf("UpdateProgress(%d): rounding error - pos %d is greater than %d\n", op, pos, MAX_PROGRESS);
pos = MAX_PROGRESS;
}
SendMessage(hProgress, PBM_SETPOS, (WPARAM)pos, 0);
}
/*
* Set or restore a Local Group Policy DWORD key indexed by szPath/SzPolicy
*/
typedef enum _GROUP_POLICY_OBJECT_TYPE {
GPOTypeLocal = 0,GPOTypeRemote,GPOTypeDS
} GROUP_POLICY_OBJECT_TYPE,*PGROUP_POLICY_OBJECT_TYPE;
#define REGISTRY_EXTENSION_GUID { 0x35378EAC,0x683F,0x11D2, {0xA8,0x9A,0x00,0xC0,0x4F,0xBB,0xCF,0xA2} }
#define GPO_OPEN_LOAD_REGISTRY 0x00000001
#define GPO_SECTION_MACHINE 2
#undef INTERFACE
#define INTERFACE IGroupPolicyObject
DECLARE_INTERFACE_(IGroupPolicyObject,IUnknown) {
STDMETHOD(QueryInterface) (THIS_ REFIID riid,LPVOID *ppvObj) PURE;
STDMETHOD_(ULONG,AddRef) (THIS) PURE;
STDMETHOD_(ULONG,Release) (THIS) PURE;
STDMETHOD(New) (THIS_ LPOLESTR pszDomainName,LPOLESTR pszDisplayName,DWORD dwFlags) PURE;
STDMETHOD(OpenDSGPO) (THIS_ LPOLESTR pszPath,DWORD dwFlags) PURE;
STDMETHOD(OpenLocalMachineGPO) (THIS_ DWORD dwFlags) PURE;
STDMETHOD(OpenRemoteMachineGPO) (THIS_ LPOLESTR pszComputerName,DWORD dwFlags) PURE;
STDMETHOD(Save) (THIS_ BOOL bMachine, BOOL bAdd,GUID *pGuidExtension,GUID *pGuid) PURE;
STDMETHOD(Delete) (THIS) PURE;
STDMETHOD(GetName) (THIS_ LPOLESTR pszName,int cchMaxLength) PURE;
STDMETHOD(GetDisplayName) (THIS_ LPOLESTR pszName,int cchMaxLength) PURE;
STDMETHOD(SetDisplayName) (THIS_ LPOLESTR pszName) PURE;
STDMETHOD(GetPath) (THIS_ LPOLESTR pszPath,int cchMaxPath) PURE;
STDMETHOD(GetDSPath) (THIS_ DWORD dwSection,LPOLESTR pszPath,int cchMaxPath) PURE;
STDMETHOD(GetFileSysPath) (THIS_ DWORD dwSection,LPOLESTR pszPath,int cchMaxPath) PURE;
STDMETHOD(GetRegistryKey) (THIS_ DWORD dwSection,HKEY *hKey) PURE;
STDMETHOD(GetOptions) (THIS_ DWORD *dwOptions) PURE;
STDMETHOD(SetOptions) (THIS_ DWORD dwOptions,DWORD dwMask) PURE;
STDMETHOD(GetType) (THIS_ GROUP_POLICY_OBJECT_TYPE *gpoType) PURE;
STDMETHOD(GetMachineName) (THIS_ LPOLESTR pszName,int cchMaxLength) PURE;
STDMETHOD(GetPropertySheetPages) (THIS_ HPROPSHEETPAGE **hPages,UINT *uPageCount) PURE;
};
typedef IGroupPolicyObject *LPGROUPPOLICYOBJECT;
BOOL SetLGP(BOOL bRestore, const char* szPath, const char* szPolicy, DWORD dwValue)
{
OSVERSIONINFO os_version;
LONG r;
DWORD disp, regtype, val, val_size=sizeof(DWORD);
HRESULT hr;
IGroupPolicyObject* pLGPO;
// These statuc values are used to restore initial state
static BOOL key_was_present = FALSE;
static DWORD original_val;
HKEY path_key = NULL, policy_key = NULL;
// MSVC is finicky about these ones => redefine them
const IID my_IID_IGroupPolicyObject =
{ 0xea502723, 0xa23d, 0x11d1, { 0xa7, 0xd3, 0x0, 0x0, 0xf8, 0x75, 0x71, 0xe3 } };
const IID my_CLSID_GroupPolicyObject =
{ 0xea502722, 0xa23d, 0x11d1, { 0xa7, 0xd3, 0x0, 0x0, 0xf8, 0x75, 0x71, 0xe3 } };
GUID ext_guid = REGISTRY_EXTENSION_GUID;
// Can be anything really
GUID snap_guid = { 0x3D271CFC, 0x2BC6, 0x4AC2, {0xB6, 0x33, 0x3B, 0xDF, 0xF5, 0xBD, 0xAB, 0x2A} };
// Vista or later
// TODO: should we do this on XP too?
memset(&os_version, 0, sizeof(OSVERSIONINFO));
os_version.dwOSVersionInfoSize = sizeof(OSVERSIONINFO);
if ((GetVersionEx(&os_version) != 0) && (os_version.dwPlatformId == VER_PLATFORM_WIN32_NT)) {
if (os_version.dwMajorVersion < 6) {
return TRUE;
}
}
// We need an IGroupPolicyObject instance to set a Local Group Policy
hr = CoCreateInstance(&my_CLSID_GroupPolicyObject, NULL, CLSCTX_INPROC_SERVER, &my_IID_IGroupPolicyObject, (LPVOID*)&pLGPO);
if (FAILED(hr)) {
uprintf("SetLGP: CoCreateInstance failed; hr = %x\n", hr);
goto error;
}
hr = pLGPO->lpVtbl->OpenLocalMachineGPO(pLGPO, GPO_OPEN_LOAD_REGISTRY);
if (FAILED(hr)) {
uprintf("SetLGP: OpenLocalMachineGPO failed - error %x\n", hr);
goto error;
}
hr = pLGPO->lpVtbl->GetRegistryKey(pLGPO, GPO_SECTION_MACHINE, &path_key);
if (FAILED(hr)) {
uprintf("SetLGP: GetRegistryKey failed - error %x\n", hr);
goto error;
}
// The DisableSystemRestore is set in Software\Policies\Microsoft\Windows\DeviceInstall\Settings
r = RegCreateKeyExA(path_key, szPath, 0, NULL, 0, KEY_SET_VALUE | KEY_QUERY_VALUE,
NULL, &policy_key, &disp);
if (r != ERROR_SUCCESS) {
uprintf("SetLGP: Failed to open LGPO path %s - error %x\n", szPath, hr);
goto error;
}
if ((disp == REG_OPENED_EXISTING_KEY) && (!bRestore) && (!key_was_present)) {
// backup existing value for restore
key_was_present = TRUE;
regtype = REG_DWORD;
r = RegQueryValueExA(policy_key, szPolicy, NULL, &regtype, (LPBYTE)&original_val, &val_size);
if (r == ERROR_FILE_NOT_FOUND) {
// The Key exists but not its value, which is OK
key_was_present = FALSE;
} else if (r != ERROR_SUCCESS) {
uprintf("SetLGP: Failed to read original %s policy value - error %x\n", szPolicy, r);
}
}
if ((!bRestore) || (key_was_present)) {
val = (bRestore)?original_val:dwValue;
r = RegSetValueExA(policy_key, szPolicy, 0, REG_DWORD, (BYTE*)&val, sizeof(val));
} else {
r = RegDeleteValueA(policy_key, szPolicy);
}
if (r != ERROR_SUCCESS) {
uprintf("SetLGP: RegSetValueEx / RegDeleteValue failed - error %x\n", r);
}
RegCloseKey(policy_key);
policy_key = NULL;
// Apply policy
hr = pLGPO->lpVtbl->Save(pLGPO, TRUE, (bRestore)?FALSE:TRUE, &ext_guid, &snap_guid);
if (r != S_OK) {
uprintf("SetLGP: Unable to apply %s policy - error %x\n", szPolicy, hr);
goto error;
} else {
uprintf("SetLGP: Successfully %s %s policy\n", (bRestore)?"restored":"disabled", szPolicy);
}
RegCloseKey(path_key);
pLGPO->lpVtbl->Release(pLGPO);
return TRUE;
error:
if (path_key != NULL) RegCloseKey(path_key);
if (policy_key != NULL) RegCloseKey(policy_key);
if (pLGPO != NULL) pLGPO->lpVtbl->Release(pLGPO);
return FALSE;
}
/*
* Toggle controls according to operation
*/
static void EnableControls(BOOL bEnable)
{
int fs;
EnableWindow(GetDlgItem(hMainDialog, IDC_DEVICE), bEnable);
EnableWindow(GetDlgItem(hMainDialog, IDC_CAPACITY), bEnable);
EnableWindow(GetDlgItem(hMainDialog, IDC_FILESYSTEM), bEnable);
EnableWindow(GetDlgItem(hMainDialog, IDC_CLUSTERSIZE), bEnable);
EnableWindow(GetDlgItem(hMainDialog, IDC_LABEL), bEnable);
EnableWindow(GetDlgItem(hMainDialog, IDC_QUICKFORMAT), bEnable);
if (bEnable) {
fs = (int)ComboBox_GetItemData(hFileSystem, ComboBox_GetCurSel(hFileSystem));
EnableWindow(GetDlgItem(hMainDialog, IDC_DOS), (fs == FS_FAT16) || (fs == FS_FAT32));
} else {
EnableWindow(GetDlgItem(hMainDialog, IDC_DOS), FALSE);
}
EnableWindow(GetDlgItem(hMainDialog, IDC_BADBLOCKS), bEnable);
EnableWindow(GetDlgItem(hMainDialog, IDC_ABOUT), bEnable);
EnableWindow(GetDlgItem(hMainDialog, IDC_START), bEnable);
SetDlgItemTextA(hMainDialog, IDCANCEL, bEnable?"Close":"Cancel");
}
/*
* Timer in the right part of the status area
*/
static void CALLBACK ClockTimer(HWND hWnd, UINT uMsg, UINT_PTR idEvent, DWORD dwTime)
{
timer++;
safe_sprintf(szTimer, sizeof(szTimer), "%02d:%02d:%02d",
timer/3600, (timer%3600)/60, timer%60);
SendMessageA(GetDlgItem(hWnd, IDC_STATUS), SB_SETTEXTA, SBT_OWNERDRAW | 1, (LPARAM)szTimer);
}
/*
* Main dialog callback
*/
static INT_PTR CALLBACK MainCallback(HWND hDlg, UINT message, WPARAM wParam, LPARAM lParam)
{
HDC hDC;
HICON hSmallIcon, hBigIcon;
DRAWITEMSTRUCT* pDI;
int nDeviceIndex, fs;
DWORD DeviceNum;
char str[MAX_PATH], tmp[128];
static uintptr_t format_thid = -1L;
static HWND hDOS;
static UINT uDOSChecked = BST_CHECKED;
switch (message) {
case WM_DEVICECHANGE:
if ( (format_thid == -1L) &&
((wParam == DBT_DEVICEARRIVAL) || (wParam == DBT_DEVICEREMOVECOMPLETE)) ) {
GetUSBDevices();
return (INT_PTR)TRUE;
}
break;
case WM_INITDIALOG:
hMainDialog = hDlg;
hDeviceList = GetDlgItem(hDlg, IDC_DEVICE);
hCapacity = GetDlgItem(hDlg, IDC_CAPACITY);
hFileSystem = GetDlgItem(hDlg, IDC_FILESYSTEM);
hClusterSize = GetDlgItem(hDlg, IDC_CLUSTERSIZE);
hLabel = GetDlgItem(hDlg, IDC_LABEL);
hProgress = GetDlgItem(hDlg, IDC_PROGRESS);
hDOS = GetDlgItem(hDlg, IDC_DOS);
// High DPI scaling
hDC = GetDC(hDlg);
fScale = GetDeviceCaps(hDC, LOGPIXELSX) / 96.0f;
ReleaseDC(hDlg, hDC);
// Create the title bar icon
hSmallIcon = (HICON)LoadImage(hMainInstance, MAKEINTRESOURCE(IDI_ICON), IMAGE_ICON, 16, 16, 0);
SendMessage (hDlg, WM_SETICON, ICON_SMALL, (LPARAM)hSmallIcon);
hBigIcon = (HICON)LoadImage(hMainInstance, MAKEINTRESOURCE(IDI_ICON), IMAGE_ICON, 32, 32, 0);
SendMessage (hDlg, WM_SETICON, ICON_BIG, (LPARAM)hBigIcon);
// Create the status line
CreateStatusBar();
// Use maximum granularity for the progress bar
SendMessage(hProgress, PBM_SETRANGE, 0, MAX_PROGRESS<<16);
// Create the string array
StrArrayCreate(&DriveID, MAX_DRIVES);
StrArrayCreate(&DriveLabel, MAX_DRIVES);
// Set the quick format & create DOS disk checkboxes
CheckDlgButton(hDlg, IDC_QUICKFORMAT, BST_CHECKED);
CheckDlgButton(hDlg, IDC_DOS, BST_CHECKED);
GetUSBDevices();
return (INT_PTR)TRUE;
// Change the colour of the version text in the status bar
case WM_DRAWITEM:
if (wParam == IDC_STATUS) {
pDI = (DRAWITEMSTRUCT*)lParam;
SetBkMode(pDI->hDC, TRANSPARENT);
SetTextColor(pDI->hDC, GetSysColor(COLOR_3DSHADOW));
pDI->rcItem.top += (int)(2.0f * fScale);
pDI->rcItem.left += (int)(4.0f * fScale);
DrawTextExA(pDI->hDC, szTimer, -1, &pDI->rcItem, DT_LEFT, NULL);
return (INT_PTR)TRUE;
}
break;
case WM_COMMAND:
switch(LOWORD(wParam)) {
case IDOK: // close application
case IDCANCEL:
if (format_thid != -1L) {
if (MessageBoxA(hMainDialog, "Cancelling may leave the device in an UNUSABLE state.\r\n"
"If you are sure you want to cancel, click YES. Otherwise, click NO.",
RUFUS_CANCELBOX_TITLE, MB_YESNO|MB_ICONWARNING) == IDYES) {
// Operation may have completed in the meantime
if (format_thid != -1L) {
FormatStatus = ERROR_SEVERITY_ERROR|FAC(FACILITY_STORAGE)|ERROR_CANCELLED;
PrintStatus(0, "Cancelling - please wait...");
}
}
return (INT_PTR)TRUE;
}
PostQuitMessage(0);
StrArrayDestroy(&DriveID);
StrArrayDestroy(&DriveLabel);
DestroyAllTooltips();
EndDialog(hDlg, 0);
break;
case IDC_ABOUT:
CreateAboutBox();
break;
case IDC_DEVICE:
switch (HIWORD(wParam)) {
case CBN_SELCHANGE:
PrintStatus(0, "%d device%s found.", ComboBox_GetCount(hDeviceList),
(ComboBox_GetCount(hDeviceList)!=1)?"s":"");
PopulateProperties(ComboBox_GetCurSel(hDeviceList));
break;
}
break;
case IDC_FILESYSTEM:
switch (HIWORD(wParam)) {
case CBN_SELCHANGE:
fs = (int)ComboBox_GetItemData(hFileSystem, ComboBox_GetCurSel(hFileSystem));
SetClusterSizes(fs);
// Disable/Restore the DOS checkbox according to FS
if ((fs == FS_FAT16) || (fs == FS_FAT32)) {
if (!IsWindowEnabled(hDOS)) {
EnableWindow(hDOS, TRUE);
CheckDlgButton(hDlg, IDC_DOS, uDOSChecked);
}
} else {
if (IsWindowEnabled(hDOS)) {
uDOSChecked = IsDlgButtonChecked(hMainDialog, IDC_DOS);
CheckDlgButton(hDlg, IDC_DOS, BST_UNCHECKED);
EnableWindow(hDOS, FALSE);
}
}
break;
}
break;
case IDC_START:
if (format_thid != -1L) {
return (INT_PTR)TRUE;
}
nDeviceIndex = ComboBox_GetCurSel(hDeviceList);
if (nDeviceIndex != CB_ERR) {
GetWindowTextA(hDeviceList, tmp, sizeof(tmp));
safe_sprintf(str, sizeof(str), "WARNING: ALL DATA ON DEVICE %s\r\nWILL BE DESTROYED.\r\n"
"To continue with this operation, click OK. To quit click CANCEL.", tmp);
if (MessageBoxA(hMainDialog, str, "Rufus", MB_OKCANCEL|MB_ICONWARNING) == IDOK) {
// Disable all controls except cancel
EnableControls(FALSE);
DeviceNum = (DWORD)ComboBox_GetItemData(hDeviceList, nDeviceIndex);
FormatStatus = 0;
InitProgress();
format_thid = _beginthread(FormatThread, 0, (void*)(uintptr_t)DeviceNum);
if (format_thid == -1L) {
uprintf("Unable to start formatting thread");
FormatStatus = ERROR_SEVERITY_ERROR|FAC(FACILITY_STORAGE)|APPERR(ERROR_CANT_START_THREAD);
PostMessage(hMainDialog, UM_FORMAT_COMPLETED, 0, 0);
}
timer = 0;
safe_sprintf(szTimer, sizeof(szTimer), "00:00:00");
SendMessageA(GetDlgItem(hMainDialog, IDC_STATUS), SB_SETTEXTA,
SBT_OWNERDRAW | 1, (LPARAM)szTimer);
SetTimer(hMainDialog, TID_APP_TIMER, 1000, ClockTimer);
}
}
break;
default:
return (INT_PTR)FALSE;
}
return (INT_PTR)TRUE;
case WM_CLOSE:
if (format_thid != -1L) {
return (INT_PTR)TRUE;
}
PostQuitMessage(0);
break;
case UM_FORMAT_COMPLETED:
format_thid = -1L;
// Stop the timer
KillTimer(hMainDialog, TID_APP_TIMER);
// Close the cancel MessageBox if active
SendMessage(FindWindowA(MAKEINTRESOURCEA(32770), RUFUS_CANCELBOX_TITLE), WM_COMMAND, IDNO, 0);
EnableControls(TRUE);
GetUSBDevices();
if (!IS_ERROR(FormatStatus)) {
PrintStatus(0, "DONE");
} else if (SCODE_CODE(FormatStatus) == ERROR_CANCELLED) {
PrintStatus(0, "Cancelled");
Notification(MSG_INFO, "Cancelled", "Operation cancelled by the user.");
} else {
PrintStatus(0, "FAILED");
Notification(MSG_ERROR, "Error", "Error: %s", StrError(FormatStatus));
}
if (FormatStatus) {
SendMessage(hProgress, PBM_SETPOS, 0, 0);
} else {
// This is the only way to achieve instantenous progress transition
SendMessage(hProgress, PBM_SETRANGE, 0, (MAX_PROGRESS+1)<<16);
SendMessage(hProgress, PBM_SETPOS, (MAX_PROGRESS+1), 0);
SendMessage(hProgress, PBM_SETRANGE, 0, MAX_PROGRESS<<16);
}
return (INT_PTR)TRUE;
}
return (INT_PTR)FALSE;
}
/*
* Application Entrypoint
*/
int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLine, int nCmdShow)
{
HANDLE mutex = NULL;
HWND hDlg = NULL;
MSG msg;
uprintf("*** RUFUS INIT ***\n");
// Prevent 2 applications from running at the same time
mutex = CreateMutexA(NULL, TRUE, "Global/RUFUS");
if ((mutex == NULL) || (GetLastError() == ERROR_ALREADY_EXISTS))
{
MessageBoxA(NULL, "Another Rufus application is running.\n"
"Please close the first application before running another one.",
"Other instance detected", MB_ICONSTOP);
return 0;
}
// Save instance of the application for further reference
hMainInstance = hInstance;
// Initialize COM for folder selection
CoInitializeEx(NULL, COINIT_APARTMENTTHREADED);
#ifdef DISABLE_AUTORUN
// We use local group policies rather than direct registry manipulation
// 0x9e disables removable and fixed drive notifications
SetLGP(FALSE, "Software\\Microsoft\\Windows\\CurrentVersion\\Policies\\Explorer", "NoDriveTypeAutorun", 0x9e);
#endif
// Create the main Window
if ( (hDlg = CreateDialogA(hInstance, MAKEINTRESOURCEA(IDD_DIALOG), NULL, MainCallback)) == NULL ) {
MessageBoxA(NULL, "Could not create Window", "DialogBox failure", MB_ICONSTOP);
goto out;
}
CenterDialog(hDlg);
ShowWindow(hDlg, SW_SHOWNORMAL);
UpdateWindow(hDlg);
// Do our own event processing
while(GetMessage(&msg, NULL, 0, 0)) {
TranslateMessage(&msg);
DispatchMessage(&msg);
}
out:
#ifdef DISABLE_AUTORUN
SetLGP(TRUE, "Software\\Microsoft\\Windows\\CurrentVersion\\Policies\\Explorer", "NoDriveTypeAutorun", 0);
#endif
CloseHandle(mutex);
uprintf("*** RUFUS EXIT ***\n");
#ifdef _CRTDBG_MAP_ALLOC
_CrtDumpMemoryLeaks();
#endif
return 0;
}
Reference in New Issue View Git Blame Copy Permalink