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

[core] factorize the img/zero disk writing code

Browse source 
* Also fix a WDK warning
This commit is contained in:
Pete Batard 2016-03-04 16:34:50 +00:00
parent d385a097c0
commit 71520baf31
3 changed files with 109 additions and 138 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

234
src/format.c
View file

@ -1264,7 +1264,7 @@ out:
// incompatibilities from one version of Windows to the next.
// Maybe when we use wimlib we'll review this, but for now just turn it off.
//#define SET_INTERNAL_DRIVES_OFFLINE
BOOL SetupWinToGo(const char* drive_name, BOOL use_ms_efi)
static BOOL SetupWinToGo(const char* drive_name, BOOL use_ms_efi)
{
#ifdef SET_INTERNAL_DRIVES_OFFLINE
static char san_policy_path[] = "?:\\san_policy.xml";
@ -1432,7 +1432,7 @@ static BOOL CALLBACK FormatPromptCallback(HWND hWnd, LPARAM lParam)
* disk in drive X: before you can use it'. You will also get that popup if you start a
* bad blocks check and cancel it before it completes. We have to close that popup manually.
*/
DWORD WINAPI CloseFormatPromptThread(LPVOID param) {
static DWORD WINAPI CloseFormatPromptThread(LPVOID param) {
HWND hFormatPrompt;
while(format_op_in_progress) {
@ -1447,7 +1447,7 @@ DWORD WINAPI CloseFormatPromptThread(LPVOID param) {
ExitThread(0);
}
void update_progress(const uint64_t processed_bytes)
static void update_progress(const uint64_t processed_bytes)
{
if (_GetTickCount64() > LastRefresh + 25) {
LastRefresh = _GetTickCount64();
@ -1457,6 +1457,99 @@ void update_progress(const uint64_t processed_bytes)
}
}
/* Write an image file or zero a drive */
static BOOL WriteDrive(HANDLE hPhysicalDrive, HANDLE hSourceImage)
{
const DWORD SectorSize = SelectedDrive.Geometry.BytesPerSector;
BOOL s, ret = FALSE;
LARGE_INTEGER li;
DWORD rSize, wSize, BufSize;
uint64_t wb, target_size = hSourceImage?img_report.projected_size:SelectedDrive.DiskSize;
uint8_t *buffer = NULL, *aligned_buffer;
int i;
// We poked the MBR and other stuff, so we need to rewind
li.QuadPart = 0;
if (!SetFilePointerEx(hPhysicalDrive, li, NULL, FILE_BEGIN))
uprintf("Warning: Unable to rewind image position - wrong data might be copied!");
LastRefresh = 0;
if (img_report.compression_type != BLED_COMPRESSION_NONE) {
uprintf("Writing Compressed Image...");
bled_init(_uprintf, update_progress, &FormatStatus);
bled_uncompress_with_handles(hSourceImage, hPhysicalDrive, img_report.compression_type);
bled_exit();
} else {
uprintf(hSourceImage?"Writing Image...":"Zeroing drive...");
// Our buffer size must be a multiple of the sector size
BufSize = ((DD_BUFFER_SIZE + SectorSize - 1) / SectorSize) * SectorSize;
buffer = (uint8_t*)calloc(BufSize + SectorSize, 1); // +1 sector for align
if (buffer == NULL) {
FormatStatus = ERROR_SEVERITY_ERROR | FAC(FACILITY_STORAGE) | ERROR_NOT_ENOUGH_MEMORY;
uprintf("could not allocate disk write buffer");
goto out;
}
// http://msdn.microsoft.com/en-us/library/windows/desktop/aa365747.aspx does buffer sector alignment
aligned_buffer = ((void *)((((uintptr_t)(buffer)) + (SectorSize)-1) & (~(((uintptr_t)(SectorSize)) - 1))));
// Don't bother trying for something clever, using double buffering overlapped and whatnot:
// With Windows' default optimizations, sync read + sync write for sequential operations
// will be as fast, if not faster, than whatever async scheme you can come up with.
rSize = BufSize;
for (wb = 0, wSize = 0; wb < (uint64_t)SelectedDrive.DiskSize; wb += wSize) {
if (_GetTickCount64() > LastRefresh + 25) {
LastRefresh = _GetTickCount64();
format_percent = (100.0f*wb) / (1.0f*target_size);
PrintInfo(0, hSourceImage?MSG_261:MSG_286, format_percent);
UpdateProgress(OP_FORMAT, format_percent);
}
if (hSourceImage != NULL) {
s = ReadFile(hSourceImage, aligned_buffer, BufSize, &rSize, NULL);
if (!s) {
FormatStatus = ERROR_SEVERITY_ERROR | FAC(FACILITY_STORAGE) | ERROR_READ_FAULT;
uprintf("read error: %s", WindowsErrorString());
goto out;
}
if (rSize == 0)
break;
}
// Don't overflow our projected size (mostly for VHDs)
if (wb + rSize > target_size) {
rSize = (DWORD)(target_size - wb);
}
// WriteFile fails unless the size is a multiple of sector size
if (rSize % SectorSize != 0)
rSize = ((rSize + SectorSize - 1) / SectorSize) * SectorSize;
for (i = 0; i < WRITE_RETRIES; i++) {
CHECK_FOR_USER_CANCEL;
s = WriteFile(hPhysicalDrive, aligned_buffer, rSize, &wSize, NULL);
if ((s) && (wSize == rSize))
break;
if (s)
uprintf("write error: Wrote %d bytes, expected %d bytes\n", wSize, rSize);
else
uprintf("write error: %s", WindowsErrorString());
if (i < WRITE_RETRIES - 1) {
li.QuadPart = wb;
SetFilePointerEx(hPhysicalDrive, li, NULL, FILE_BEGIN);
uprintf(" RETRYING...\n");
} else {
FormatStatus = ERROR_SEVERITY_ERROR | FAC(FACILITY_STORAGE) | ERROR_WRITE_FAULT;
goto out;
}
}
if (i >= WRITE_RETRIES) goto out;
}
}
RefreshDriveLayout(hPhysicalDrive);
ret = TRUE;
out:
safe_free(buffer);
return ret;
}
/*
* Standalone thread for the formatting operation
* According to http://msdn.microsoft.com/en-us/library/windows/desktop/aa364562.aspx
@ -1471,20 +1564,18 @@ void update_progress(const uint64_t processed_bytes)
DWORD WINAPI FormatThread(void* param)
{
int i, r, pt, tt, fs, bt;
BOOL s, ret, use_large_fat32, windows_to_go;
BOOL ret, use_large_fat32, windows_to_go;
const DWORD SectorSize = SelectedDrive.Geometry.BytesPerSector;
DWORD rSize, wSize, BufSize, DriveIndex = (DWORD)(uintptr_t)param;
DWORD DriveIndex = (DWORD)(uintptr_t)param;
HANDLE hPhysicalDrive = INVALID_HANDLE_VALUE;
HANDLE hLogicalVolume = INVALID_HANDLE_VALUE;
HANDLE hSourceImage = INVALID_HANDLE_VALUE;
SYSTEMTIME lt;
FILE* log_fd;
LARGE_INTEGER li;
uint64_t wb;
uint8_t *buffer = NULL, *aligned_buffer, extra_partitions = 0;
uint8_t *buffer = NULL, extra_partitions = 0;
char *bb_msg, *guid_volume = NULL;
char drive_name[] = "?:\\";
char drive_letters[27];
char drive_letters[27], fs_type[32];
char logfile[MAX_PATH], *userdir;
char efi_dst[] = "?:\\efi\\boot\\bootx64.efi";
char kolibri_dst[] = "?:\\MTLD_F32";
@ -1586,57 +1677,8 @@ DWORD WINAPI FormatThread(void* param)
}
CreateThread(NULL, 0, CloseFormatPromptThread, NULL, 0, NULL);
// TODO: factorize this with DD write?
if (zero_drive) {
li.QuadPart = 0;
SetFilePointerEx(hPhysicalDrive, li, NULL, FILE_BEGIN);
uprintf("Zeroing drive...");
// Our buffer size must be a multiple of the sector size
BufSize = ((DD_BUFFER_SIZE + SectorSize - 1) / SectorSize) * SectorSize;
buffer = (uint8_t*)calloc(BufSize + SectorSize, 1); // +1 sector for align
if (buffer == NULL) {
FormatStatus = ERROR_SEVERITY_ERROR | FAC(FACILITY_STORAGE) | ERROR_NOT_ENOUGH_MEMORY;
uprintf("could not allocate zeroing buffer");
goto out;
}
// http://msdn.microsoft.com/en-us/library/windows/desktop/aa365747.aspx does buffer sector alignment
aligned_buffer = ((void *)((((uintptr_t)(buffer)) + (SectorSize)-1) & (~(((uintptr_t)(SectorSize)) - 1))));
for (wb = 0, wSize = 0; wb < (uint64_t)SelectedDrive.DiskSize; wb += wSize) {
if (_GetTickCount64() > LastRefresh + 25) {
LastRefresh = _GetTickCount64();
format_percent = (100.0f*wb) / (1.0f*SelectedDrive.DiskSize);
PrintInfo(0, MSG_286, format_percent);
UpdateProgress(OP_FORMAT, format_percent);
}
// Don't overflow our projected size
if (wb + BufSize > (uint64_t)SelectedDrive.DiskSize) {
BufSize = (DWORD)(SelectedDrive.DiskSize - wb);
}
// WriteFile fails unless the size is a multiple of sector size
if (BufSize % SectorSize != 0)
BufSize = ((BufSize + SectorSize - 1) / SectorSize) * SectorSize;
for (i = 0; i < WRITE_RETRIES; i++) {
CHECK_FOR_USER_CANCEL;
s = WriteFile(hPhysicalDrive, aligned_buffer, BufSize, &wSize, NULL);
if ((s) && (wSize == BufSize))
break;
if (s)
uprintf("write error: Wrote %d bytes, expected %d bytes\n", wSize, BufSize);
else
uprintf("write error: %s", WindowsErrorString());
if (i < WRITE_RETRIES - 1) {
li.QuadPart = wb;
SetFilePointerEx(hPhysicalDrive, li, NULL, FILE_BEGIN);
uprintf(" RETRYING...\n");
}
else {
FormatStatus = ERROR_SEVERITY_ERROR | FAC(FACILITY_STORAGE) | ERROR_WRITE_FAULT;
goto out;
}
}
if (i >= WRITE_RETRIES) goto out;
}
RefreshDriveLayout(hPhysicalDrive);
WriteDrive(hPhysicalDrive, NULL);
goto out;
}
@ -1707,11 +1749,6 @@ DWORD WINAPI FormatThread(void* param)
// Write an image file
if (IsChecked(IDC_BOOT) && (bt == BT_IMG)) {
char fs_type[32];
// We poked the MBR and other stuff, so we need to rewind
li.QuadPart = 0;
if (!SetFilePointerEx(hPhysicalDrive, li, NULL, FILE_BEGIN))
uprintf("Warning: Unable to rewind image position - wrong data might be copied!");
hSourceImage = CreateFileU(image_path, GENERIC_READ, FILE_SHARE_READ, NULL,
OPEN_EXISTING, FILE_FLAG_SEQUENTIAL_SCAN, NULL);
if (hSourceImage == INVALID_HANDLE_VALUE) {
@ -1719,75 +1756,10 @@ DWORD WINAPI FormatThread(void* param)
FormatStatus = ERROR_SEVERITY_ERROR|FAC(FACILITY_STORAGE)|ERROR_OPEN_FAILED;
goto out;
}
LastRefresh = 0;
if (img_report.compression_type != BLED_COMPRESSION_NONE) {
uprintf("Writing Compressed Image...");
bled_init(_uprintf, update_progress, &FormatStatus);
bled_uncompress_with_handles(hSourceImage, hPhysicalDrive, img_report.compression_type);
bled_exit();
} else {
uprintf("Writing Image...");
// Our buffer size must be a multiple of the sector size
BufSize = ((DD_BUFFER_SIZE + SectorSize - 1) / SectorSize) * SectorSize;
buffer = (uint8_t*)malloc(BufSize + SectorSize); // +1 sector for align
if (buffer == NULL) {
FormatStatus = ERROR_SEVERITY_ERROR|FAC(FACILITY_STORAGE)|ERROR_NOT_ENOUGH_MEMORY;
uprintf("could not allocate DD buffer");
goto out;
}
// http://msdn.microsoft.com/en-us/library/windows/desktop/aa365747.aspx does buffer sector alignment
aligned_buffer = ((void *) ((((uintptr_t)(buffer)) + (SectorSize) - 1) & (~(((uintptr_t)(SectorSize)) - 1))));
// Don't bother trying for something clever, using double buffering overlapped and whatnot:
// With Windows' default optimizations, sync read + sync write for sequential operations
// will be as fast, if not faster, than whatever async scheme you can come up with.
for (wb = 0, wSize = 0; ; wb += wSize) {
s = ReadFile(hSourceImage, aligned_buffer, BufSize, &rSize, NULL);
if (!s) {
FormatStatus = ERROR_SEVERITY_ERROR|FAC(FACILITY_STORAGE)|ERROR_READ_FAULT;
uprintf("read error: %s", WindowsErrorString());
goto out;
}
if (rSize == 0)
break;
if (_GetTickCount64() > LastRefresh + 25) {
LastRefresh = _GetTickCount64();
format_percent = (100.0f*wb)/(1.0f*img_report.projected_size);
PrintInfo(0, MSG_261, format_percent);
UpdateProgress(OP_FORMAT, format_percent);
}
// Don't overflow our projected size (mostly for VHDs)
if (wb + rSize > img_report.projected_size) {
rSize = (DWORD)(img_report.projected_size - wb);
}
// WriteFile fails unless the size is a multiple of sector size
if (rSize % SectorSize != 0)
rSize = ((rSize + SectorSize -1) / SectorSize) * SectorSize;
for (i=0; i < WRITE_RETRIES; i++) {
CHECK_FOR_USER_CANCEL;
s = WriteFile(hPhysicalDrive, aligned_buffer, rSize, &wSize, NULL);
if ((s) && (wSize == rSize))
break;
if (s)
uprintf("write error: Wrote %d bytes, expected %d bytes\n", wSize, rSize);
else
uprintf("write error: %s", WindowsErrorString());
if (i < WRITE_RETRIES-1) {
li.QuadPart = wb;
SetFilePointerEx(hPhysicalDrive, li, NULL, FILE_BEGIN);
uprintf(" RETRYING...\n");
} else {
FormatStatus = ERROR_SEVERITY_ERROR|FAC(FACILITY_STORAGE)|ERROR_WRITE_FAULT;
goto out;
}
}
if (i >= WRITE_RETRIES) goto out;
}
}
WriteDrive(hPhysicalDrive, hSourceImage);
// If the image contains a partition we might be able to access, try to re-mount it
RefreshDriveLayout(hPhysicalDrive);
safe_unlockclose(hPhysicalDrive);
safe_unlockclose(hLogicalVolume);
Sleep(200);
@ -1797,8 +1769,6 @@ DWORD WINAPI FormatThread(void* param)
if ((guid_volume != NULL) && (MountVolume(drive_name, guid_volume)))
uprintf("Remounted %s on %s\n", guid_volume, drive_name);
}
uprintf("Done");
goto out;
}

10
src/rufus.rc
View file

@ -33,7 +33,7 @@ LANGUAGE LANG_NEUTRAL, SUBLANG_NEUTRAL
IDD_DIALOG DIALOGEX 12, 12, 242, 376
STYLE DS_SETFONT | DS_MODALFRAME | DS_CENTER | WS_MINIMIZEBOX | WS_POPUP | WS_CAPTION | WS_SYSMENU
EXSTYLE WS_EX_ACCEPTFILES
CAPTION "Rufus 2.8.877"
CAPTION "Rufus 2.8.878"
FONT 8, "Segoe UI Symbol", 400, 0, 0x0
BEGIN
LTEXT "Device",IDS_DEVICE_TXT,9,6,200,8
@ -320,8 +320,8 @@ END
//
VS_VERSION_INFO VERSIONINFO
FILEVERSION 2,8,877,0
PRODUCTVERSION 2,8,877,0
FILEVERSION 2,8,878,0
PRODUCTVERSION 2,8,878,0
FILEFLAGSMASK 0x3fL
#ifdef _DEBUG
FILEFLAGS 0x1L
@ -338,13 +338,13 @@ BEGIN
BEGIN
VALUE "CompanyName", "Akeo Consulting (http://akeo.ie)"
VALUE "FileDescription", "Rufus"
VALUE "FileVersion", "2.8.877"
VALUE "FileVersion", "2.8.878"
VALUE "InternalName", "Rufus"
VALUE "LegalCopyright", "© 2011-2016 Pete Batard (GPL v3)"
VALUE "LegalTrademarks", "http://www.gnu.org/copyleft/gpl.html"
VALUE "OriginalFilename", "rufus.exe"
VALUE "ProductName", "Rufus"
VALUE "ProductVersion", "2.8.877"
VALUE "ProductVersion", "2.8.878"
END
END
BLOCK "VarFileInfo"

3
src/stdfn.c
View file

@ -820,7 +820,8 @@ BOOL SetLGP(BOOL bRestore, BOOL* bExistingKey, const char* szPath, const char* s
*/
BOOL SetThreadAffinity(DWORD_PTR* thread_affinity, size_t num_threads)
{
int i, j, pc;
size_t i, j, pc;
DWORD_PTR affinity, dummy;
memset(thread_affinity, 0, num_threads * sizeof(DWORD_PTR));