From b66401b3559dd2e69c1c8cbfe9b996d1a612f3ba Mon Sep 17 00:00:00 2001
From: Pete Batard <pete@akeo.ie>
Date: Wed, 4 Nov 2015 13:58:19 +0000
Subject: [PATCH] [core] add a cheat mode to zero the USB device

* Alt-Z to zap!
* Also set rufus-next to 2.6
---
 configure                  | 20 +++++------
 configure.ac               |  2 +-
 res/localization/rufus.loc |  1 +
 src/format.c               | 73 +++++++++++++++++++++++++++++++++-----
 src/rufus.c                |  9 ++++-
 src/rufus.rc               | 10 +++---
 6 files changed, 89 insertions(+), 26 deletions(-)

diff --git a/configure b/configure
index d23c588c..594e7162 100644
--- a/configure
+++ b/configure
@@ -1,6 +1,6 @@
 #! /bin/sh
 # Guess values for system-dependent variables and create Makefiles.
-# Generated by GNU Autoconf 2.69 for rufus 2.5.
+# Generated by GNU Autoconf 2.69 for rufus 2.6.
 #
 # Report bugs to <https://github.com/pbatard/rufus/issues>.
 #
@@ -580,8 +580,8 @@ MAKEFLAGS=
 # Identity of this package.
 PACKAGE_NAME='rufus'
 PACKAGE_TARNAME='rufus'
-PACKAGE_VERSION='2.5'
-PACKAGE_STRING='rufus 2.5'
+PACKAGE_VERSION='2.6'
+PACKAGE_STRING='rufus 2.6'
 PACKAGE_BUGREPORT='https://github.com/pbatard/rufus/issues'
 PACKAGE_URL='http://rufus.akeo.ie'
 
@@ -1228,7 +1228,7 @@ if test "$ac_init_help" = "long"; then
   # Omit some internal or obsolete options to make the list less imposing.
   # This message is too long to be a string in the A/UX 3.1 sh.
   cat <<_ACEOF
-\`configure' configures rufus 2.5 to adapt to many kinds of systems.
+\`configure' configures rufus 2.6 to adapt to many kinds of systems.
 
 Usage: $0 [OPTION]... [VAR=VALUE]...
 
@@ -1294,7 +1294,7 @@ fi
 
 if test -n "$ac_init_help"; then
   case $ac_init_help in
-     short | recursive ) echo "Configuration of rufus 2.5:";;
+     short | recursive ) echo "Configuration of rufus 2.6:";;
    esac
   cat <<\_ACEOF
 
@@ -1385,7 +1385,7 @@ fi
 test -n "$ac_init_help" && exit $ac_status
 if $ac_init_version; then
   cat <<\_ACEOF
-rufus configure 2.5
+rufus configure 2.6
 generated by GNU Autoconf 2.69
 
 Copyright (C) 2012 Free Software Foundation, Inc.
@@ -1440,7 +1440,7 @@ cat >config.log <<_ACEOF
 This file contains any messages produced by compilers while
 running configure, to aid debugging if configure makes a mistake.
 
-It was created by rufus $as_me 2.5, which was
+It was created by rufus $as_me 2.6, which was
 generated by GNU Autoconf 2.69.  Invocation command line was
 
   $ $0 $@
@@ -2303,7 +2303,7 @@ fi
 
 # Define the identity of the package.
  PACKAGE='rufus'
- VERSION='2.5'
+ VERSION='2.6'
 
 
 cat >>confdefs.h <<_ACEOF
@@ -4480,7 +4480,7 @@ cat >>$CONFIG_STATUS <<\_ACEOF || ac_write_fail=1
 # report actual input values of CONFIG_FILES etc. instead of their
 # values after options handling.
 ac_log="
-This file was extended by rufus $as_me 2.5, which was
+This file was extended by rufus $as_me 2.6, which was
 generated by GNU Autoconf 2.69.  Invocation command line was
 
   CONFIG_FILES    = $CONFIG_FILES
@@ -4534,7 +4534,7 @@ _ACEOF
 cat >>$CONFIG_STATUS <<_ACEOF || ac_write_fail=1
 ac_cs_config="`$as_echo "$ac_configure_args" | sed 's/^ //; s/[\\""\`\$]/\\\\&/g'`"
 ac_cs_version="\\
-rufus config.status 2.5
+rufus config.status 2.6
 configured by $0, generated by GNU Autoconf 2.69,
   with options \\"\$ac_cs_config\\"
 
diff --git a/configure.ac b/configure.ac
index 57437056..7f20422a 100644
--- a/configure.ac
+++ b/configure.ac
@@ -1,4 +1,4 @@
-AC_INIT([rufus], [2.5], [https://github.com/pbatard/rufus/issues], [rufus], [http://rufus.akeo.ie])
+AC_INIT([rufus], [2.6], [https://github.com/pbatard/rufus/issues], [rufus], [http://rufus.akeo.ie])
 AM_INIT_AUTOMAKE([-Wno-portability foreign no-dist no-dependencies])
 AC_CONFIG_SRCDIR([src/rufus.c])
 AC_CONFIG_MACRO_DIR([m4])
diff --git a/res/localization/rufus.loc b/res/localization/rufus.loc
index 97db183c..687723af 100644
--- a/res/localization/rufus.loc
+++ b/res/localization/rufus.loc
@@ -545,6 +545,7 @@ t MSG_283 "Invalid signature"
 t MSG_284 "The downloaded executable is missing a digital signature."
 t MSG_285 "The downloaded executable is signed by '%s'.\nThis is not a signature we recognize and could "
 	"indicate some form of malicious activity...\nAre you sure you want to run this file?"
+t MSG_286 "Zeroing drive: %0.1f%% completed"
 
 ################################################################################
 ############################# TRANSLATOR END COPY ##############################
diff --git a/src/format.c b/src/format.c
index a2e5f4bf..d2218ee7 100644
--- a/src/format.c
+++ b/src/format.c
@@ -59,7 +59,7 @@ static int task_number = 0;
 extern const int nb_steps[FS_MAX];
 extern uint32_t dur_mins, dur_secs;
 static int fs_index = 0;
-extern BOOL force_large_fat32, enable_ntfs_compression, lock_drive;
+extern BOOL force_large_fat32, enable_ntfs_compression, lock_drive, zero_drive;
 uint8_t *grub2_buf = NULL;
 long grub2_len;
 static BOOL WritePBR(HANDLE hLogicalDrive);
@@ -1577,12 +1577,14 @@ DWORD WINAPI FormatThread(void* param)
 	}
 	CHECK_FOR_USER_CANCEL;
 
-	PrintInfoDebug(0, MSG_226);
-	AnalyzeMBR(hPhysicalDrive, "Drive");
-	if ((hLogicalVolume != NULL) && (hLogicalVolume != INVALID_HANDLE_VALUE)) {
-		AnalyzePBR(hLogicalVolume);
+	if (!zero_drive) {
+		PrintInfoDebug(0, MSG_226);
+		AnalyzeMBR(hPhysicalDrive, "Drive");
+		if ((hLogicalVolume != NULL) && (hLogicalVolume != INVALID_HANDLE_VALUE)) {
+			AnalyzePBR(hLogicalVolume);
+		}
+		UpdateProgress(OP_ANALYZE_MBR, -1.0f);
 	}
-	UpdateProgress(OP_ANALYZE_MBR, -1.0f);
 
 	// Zap any existing partitions. This helps prevent access errors.
 	// As this creates issues with FAT16 formatted MS drives, only do this for other filesystems
@@ -1591,9 +1593,61 @@ DWORD WINAPI FormatThread(void* param)
 		FormatStatus = ERROR_SEVERITY_ERROR|FAC(FACILITY_STORAGE)|ERROR_PARTITION_FAILURE;
 		goto out;
 	}
-
 	CreateThread(NULL, 0, CloseFormatPromptThread, NULL, 0, NULL);
-	if (IsChecked(IDC_BADBLOCKS)) {
+
+	// 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 (GetTickCount() > LastRefresh + 25) {
+				LastRefresh = GetTickCount();
+				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);
+		goto out;
+	} else if (IsChecked(IDC_BADBLOCKS)) {
 		do {
 			// create a log file for bad blocks report. Since %USERPROFILE% may
 			// have localized characters, we use the UTF-8 API.
@@ -1717,7 +1771,7 @@ DWORD WINAPI FormatThread(void* param)
 				// 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++) {
+				for (i=0; i < WRITE_RETRIES; i++) {
 					CHECK_FOR_USER_CANCEL;
 					s = WriteFile(hPhysicalDrive, aligned_buffer, rSize, &wSize, NULL);
 					if ((s) && (wSize == rSize))
@@ -1954,6 +2008,7 @@ DWORD WINAPI FormatThread(void* param)
 	}
 
 out:
+	zero_drive = FALSE;
 	safe_free(guid_volume);
 	safe_free(buffer);
 	safe_closehandle(hSourceImage);
diff --git a/src/rufus.c b/src/rufus.c
index 4e64dd60..06e10680 100644
--- a/src/rufus.c
+++ b/src/rufus.c
@@ -141,6 +141,7 @@ BOOL use_own_c32[NB_OLD_C32] = {FALSE, FALSE}, mbr_selected_by_user = FALSE, tog
 BOOL iso_op_in_progress = FALSE, format_op_in_progress = FALSE, right_to_left_mode = FALSE;
 BOOL enable_HDDs = FALSE, force_update = FALSE, enable_ntfs_compression = FALSE, no_confirmation_on_cancel = FALSE, lock_drive = TRUE;
 BOOL advanced_mode, allow_dual_uefi_bios, detect_fakes, enable_vmdk, force_large_fat32, usb_debug, use_fake_units, preserve_timestamps;
+BOOL zero_drive = FALSE;
 int dialog_showing = 0, lang_button_id = 0;
 uint16_t rufus_version[3], embedded_sl_version[2];
 char embedded_sl_version_str[2][12] = { "?.??", "?.??" };
@@ -2514,7 +2515,7 @@ static INT_PTR CALLBACK MainCallback(HWND hDlg, UINT message, WPARAM wParam, LPA
 				// Disable all controls except cancel
 				EnableControls(FALSE);
 				DeviceNum = (DWORD)ComboBox_GetItemData(hDeviceList, nDeviceIndex);
-				InitProgress(FALSE);
+				InitProgress(zero_drive);
 				format_thid = CreateThread(NULL, 0, FormatThread, (LPVOID)(uintptr_t)DeviceNum, 0, NULL);
 				if (format_thid == NULL) {
 					uprintf("Unable to start formatting thread");
@@ -3173,6 +3174,12 @@ relaunch:
 			SHDeleteDirectoryExU(NULL, tmp_path, FOF_SILENT | FOF_NOERRORUI | FOF_NOCONFIRMATION);
 			continue;
 		}
+		// Alt-Z => Zero the drive
+		if ((msg.message == WM_SYSKEYDOWN) && (msg.wParam == 'Z')) {
+			zero_drive = TRUE;
+			// Simulate a button click for Start
+			PostMessage(hDlg, WM_COMMAND, (WPARAM)IDC_START, 0);
+		}
 
 		// Let the system handle dialog messages (e.g. those from the tab key)
 		if (!IsDialogMessage(hDlg, &msg) && !IsDialogMessage(hLogDlg, &msg)) {
diff --git a/src/rufus.rc b/src/rufus.rc
index 9d9c6bfc..c5952d28 100644
--- a/src/rufus.rc
+++ b/src/rufus.rc
@@ -32,7 +32,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
-CAPTION "Rufus 2.5.801"
+CAPTION "Rufus 2.6.802"
 FONT 8, "Segoe UI Symbol", 400, 0, 0x0
 BEGIN
     LTEXT           "Device",IDS_DEVICE_TXT,9,6,200,8
@@ -319,8 +319,8 @@ END
 //
 
 VS_VERSION_INFO VERSIONINFO
- FILEVERSION 2,5,801,0
- PRODUCTVERSION 2,5,801,0
+ FILEVERSION 2,6,802,0
+ PRODUCTVERSION 2,6,802,0
  FILEFLAGSMASK 0x3fL
 #ifdef _DEBUG
  FILEFLAGS 0x1L
@@ -337,13 +337,13 @@ BEGIN
         BEGIN
             VALUE "CompanyName", "Akeo Consulting (http://akeo.ie)"
             VALUE "FileDescription", "Rufus"
-            VALUE "FileVersion", "2.5.801"
+            VALUE "FileVersion", "2.6.802"
             VALUE "InternalName", "Rufus"
             VALUE "LegalCopyright", "© 2011-2015 Pete Batard (GPL v3)"
             VALUE "LegalTrademarks", "http://www.gnu.org/copyleft/gpl.html"
             VALUE "OriginalFilename", "rufus.exe"
             VALUE "ProductName", "Rufus"
-            VALUE "ProductVersion", "2.5.801"
+            VALUE "ProductVersion", "2.6.802"
         END
     END
     BLOCK "VarFileInfo"