From fdcb17aa7ae000070d4fa8613d5615f5365f1f69 Mon Sep 17 00:00:00 2001
From: Pete Batard <pete@akeo.ie>
Date: Thu, 22 Jan 2015 22:31:34 +0000
Subject: [PATCH] [misc] pocketful of enhancement and fixes - part 3

* Add Windows To Go tooltips
* Set UEFI:TOGO partition type to EFI System Partition
* Improve apply-image progress report
* Fix MBR analysis for 4k disks
* Do not select Rufus MBR for Windows To Go
---
 res/localization/rufus.loc |  4 ++-
 src/drive.c                | 18 +++++++---
 src/format.c               |  9 +++--
 src/rufus.c                | 11 ++++++-
 src/rufus.h                |  2 +-
 src/rufus.rc               | 16 ++++-----
 src/sys_types.h            |  1 +
 src/vhd.c                  | 67 +++++++++++++++++++++++++++++++++++---
 8 files changed, 105 insertions(+), 23 deletions(-)

diff --git a/res/localization/rufus.loc b/res/localization/rufus.loc
index 6ef80e6c..170f9d41 100644
--- a/res/localization/rufus.loc
+++ b/res/localization/rufus.loc
@@ -419,6 +419,8 @@ t MSG_196 "IMPORTANT: THIS DRIVE USES A NONSTANDARD SECTOR SIZE!!\n\n"
 t MSG_197 "Nonstandard sector size detected"
 t MSG_198 "Windows To Go can only be installed on a GPT partitioned drive if it has "
 	"the FIXED attribute set. The current drive was not detected as FIXED."
+t MSG_199 "Choose this if you want to INSTALL Windows, on a different disk."
+t MSG_200 "Choose this if you want to RUN Windows, from the USB."
 
 # Status messages - these messages will appear on the status bar
 t MSG_201 "Cancelling - Please wait..."
@@ -507,7 +509,7 @@ t MSG_263 "Use PROPER size units"
 t MSG_264 "Deleting directory '%s'"
 t MSG_265 "VMWare disk detection"
 t MSG_266 "Dual UEFI/BIOS mode"
-t MSG_267 "Applying Windows image: %d%% completed"
+t MSG_267 "Applying Windows image: %0.1f%% completed"
 t MSG_268 "Applying Windows image..."
 
 ################################################################################
diff --git a/src/drive.c b/src/drive.c
index 1bb15399..1e5f5800 100644
--- a/src/drive.c
+++ b/src/drive.c
@@ -589,8 +589,10 @@ BOOL AnalyzeMBR(HANDLE hPhysicalDrive, const char* TargetName)
 	int i;
 
 	fake_fd._ptr = (char*)hPhysicalDrive;
-	// Must be set to 512, as we also use this method for images and we may not have a target UFD yet
-	fake_fd._bufsiz = 512;
+	fake_fd._bufsiz = SelectedDrive.Geometry.BytesPerSector;
+	// Might need correction, as we use this method for images and we may not have a target UFD yet
+	if (fake_fd._bufsiz < 512)
+		fake_fd._bufsiz = 512;
 
 	if (!is_br(&fake_fd)) {
 		uprintf("%s does not have an x86 %s\n", TargetName, mbr_name);
@@ -726,7 +728,7 @@ BOOL GetDrivePartitionData(DWORD DriveIndex, char* FileSystemName, DWORD FileSys
 		for (i=0; i<DriveLayout->PartitionCount; i++) {
 			if (DriveLayout->PartitionEntry[i].Mbr.PartitionType != PARTITION_ENTRY_UNUSED) {
 				part_type = DriveLayout->PartitionEntry[i].Mbr.PartitionType;
-				isUefiTogo = (i == 1) && (part_type == 0x01) &&
+				isUefiTogo = (i == 1) && (part_type == 0xef) &&
 					(DriveLayout->PartitionEntry[i].PartitionLength.QuadPart == uefi_togo_size);
 				suprintf("Partition %d%s:\n", i+1, isUefiTogo?" (UEFI:TOGO)":"");
 				for (j=0; j<ARRAYSIZE(mbr_mountable); j++) {
@@ -1025,8 +1027,11 @@ BOOL CreatePartition(HANDLE hDrive, int partition_style, int file_system, BOOL m
 
 	PrintInfoDebug(0, MSG_238, PartitionTypeName[partition_style]);
 
-	if ((extra_partitions & XP_UEFI_TOGO) && (uefi_togo_size == 0))
+	if ((extra_partitions & XP_UEFI_TOGO) && (uefi_togo_size == 0)) {
 		uefi_togo_size = GetResourceSize(hMainInstance, MAKEINTRESOURCEA(IDR_UEFI_TOGO), _RT_RCDATA, "uefi-togo.img");
+		if (uefi_togo_size == 0)
+			return FALSE;
+	}
 
 	// Compute the start offset of our first partition
 	if ((partition_style == PARTITION_STYLE_GPT) || (!IsChecked(IDC_EXTRA_PARTITION))) {
@@ -1064,6 +1069,9 @@ BOOL CreatePartition(HANDLE hDrive, int partition_style, int file_system, BOOL m
 	main_part_size_in_sectors = (SelectedDrive.DiskSize - DriveLayoutEx.PartitionEntry[pn].StartingOffset.QuadPart) /
 		// Need 33 sectors at the end for secondary GPT
 		SelectedDrive.Geometry.BytesPerSector - ((partition_style == PARTITION_STYLE_GPT)?33:0);
+	if (main_part_size_in_sectors <= 0)
+		return FALSE;
+
 	// Adjust the size according to extra partitions (which we always align to a track)
 	if (extra_partitions) {
 		uprintf("Adding extra partition");
@@ -1127,7 +1135,7 @@ BOOL CreatePartition(HANDLE hDrive, int partition_style, int file_system, BOOL m
 			IGNORE_RETVAL(CoCreateGuid(&DriveLayoutEx.PartitionEntry[pn].Gpt.PartitionId));
 			wcscpy(DriveLayoutEx.PartitionEntry[pn].Gpt.Name, (extra_partitions & XP_UEFI_TOGO)?L"UEFI:TOGO":L"EFI system partition");
 		} else {
-			DriveLayoutEx.PartitionEntry[pn].Mbr.PartitionType = (extra_partitions & XP_UEFI_TOGO)?0x01:RUFUS_EXTRA_PARTITION_TYPE;
+			DriveLayoutEx.PartitionEntry[pn].Mbr.PartitionType = (extra_partitions & XP_UEFI_TOGO)?0xef:RUFUS_EXTRA_PARTITION_TYPE;
 			if (extra_partitions & XP_COMPAT)
 				// Set the one track compatibility partition to be all hidden sectors
 				DriveLayoutEx.PartitionEntry[pn].Mbr.HiddenSectors = SelectedDrive.Geometry.SectorsPerTrack;
diff --git a/src/format.c b/src/format.c
index 0233f66b..0176fff2 100644
--- a/src/format.c
+++ b/src/format.c
@@ -50,8 +50,9 @@
 /*
  * Globals
  */
-DWORD FormatStatus, LastRefresh;
+DWORD FormatStatus;
 badblocks_report report;
+static DWORD LastRefresh;
 static float format_percent = 0.0f;
 static int task_number = 0;
 extern const int nb_steps[FS_MAX];
@@ -1320,7 +1321,8 @@ BOOL SetupWinToGo(const char* drive_name, BOOL use_ms_efi)
 		Sleep(200);
 		AltUnmountVolume(ms_efi);
 	}
-	UpdateProgress(OP_DOS, 99.0f);
+	PrintInfo(0, MSG_267, 99.9f);
+	UpdateProgress(OP_DOS, 99.9f);
 
 	// The following are non fatal if they fail
 	buffer = GetResource(hMainInstance, MAKEINTRESOURCEA(IDR_TOGO_SAN_POLICY_XML),
@@ -1348,6 +1350,7 @@ BOOL SetupWinToGo(const char* drive_name, BOOL use_ms_efi)
 		uprintf("Could not write '%s'\n", unattend_path);
 	}
 	fclose(fd);
+	PrintInfo(0, MSG_267, 100.0f);
 	UpdateProgress(OP_DOS, 100.0f);
 
 	return TRUE;
@@ -1452,7 +1455,7 @@ DWORD WINAPI FormatThread(void* param)
 	pt = GETPARTTYPE((int)ComboBox_GetItemData(hPartitionScheme, ComboBox_GetCurSel(hPartitionScheme)));
 	bt = GETBIOSTYPE((int)ComboBox_GetItemData(hPartitionScheme, ComboBox_GetCurSel(hPartitionScheme)));
 	use_large_fat32 = (fs == FS_FAT32) && ((SelectedDrive.DiskSize > LARGE_FAT32_SIZE) || (force_large_fat32));
-	windows_to_go = HAS_TOGO(iso_report) && (Button_GetCheck(GetDlgItem(hMainDialog, IDC_WINDOWS_TO_GO)) == BST_CHECKED);
+	windows_to_go = (togo_mode) && HAS_TOGO(iso_report) && (Button_GetCheck(GetDlgItem(hMainDialog, IDC_WINDOWS_TO_GO)) == BST_CHECKED);
 	// Find out if we need to add any extra partitions
 	if ((windows_to_go) && (bt == BT_UEFI) && (pt == PARTITION_STYLE_GPT))
 		// According to Microsoft, every GPT disk (we RUN Windows from) must have an MSR due to not having hidden sectors
diff --git a/src/rufus.c b/src/rufus.c
index a3b7749c..833890b5 100644
--- a/src/rufus.c
+++ b/src/rufus.c
@@ -498,7 +498,8 @@ static void SetMBRProps(void)
 	int dt = (int)ComboBox_GetItemData(hBootType, ComboBox_GetCurSel(hBootType));
 	BOOL needs_masquerading = (IS_WINPE(iso_report.winpe) && (!iso_report.uses_minint));
 
-	if ((!mbr_selected_by_user) && ((image_path == NULL) || (dt != DT_ISO) || (fs != FS_NTFS) || IS_GRUB(iso_report))) {
+	if ((!mbr_selected_by_user) && ((image_path == NULL) || (dt != DT_ISO) || (fs != FS_NTFS) || IS_GRUB(iso_report) ||
+		((togo_mode) && (Button_GetCheck(GetDlgItem(hMainDialog, IDC_WINDOWS_TO_GO)) == BST_CHECKED)) )) {
 		CheckDlgButton(hMainDialog, IDC_RUFUS_MBR, BST_UNCHECKED);
 		IGNORE_RETVAL(ComboBox_SetCurSel(hDiskID, 0));
 		return;
@@ -1658,6 +1659,8 @@ void InitDialog(HWND hDlg)
 	CreateTooltip(GetDlgItem(hDlg, IDC_ENABLE_FIXED_DISKS), lmprintf(MSG_170), -1);
 	CreateTooltip(GetDlgItem(hDlg, IDC_START), lmprintf(MSG_171), -1);
 	CreateTooltip(GetDlgItem(hDlg, IDC_ABOUT), lmprintf(MSG_172), -1);
+	CreateTooltip(GetDlgItem(hDlg, IDC_WINDOWS_INSTALL), lmprintf(MSG_199), -1);
+	CreateTooltip(GetDlgItem(hDlg, IDC_WINDOWS_TO_GO), lmprintf(MSG_200), -1);
 	// TODO: add new tooltips
 
 	ToggleAdvanced();	// We start in advanced mode => go to basic mode
@@ -2165,6 +2168,10 @@ static INT_PTR CALLBACK MainCallback(HWND hDlg, UINT message, WPARAM wParam, LPA
 				FormatStatus = ERROR_SEVERITY_ERROR|FAC(FACILITY_STORAGE)|APPERR(ERROR_CANT_START_THREAD);
 			}
 			break;
+		case IDC_WINDOWS_INSTALL:
+			if (Button_GetCheck(GetDlgItem(hMainDialog, IDC_WINDOWS_INSTALL)) == BST_CHECKED)
+				SetMBRProps();
+			break;
 		case IDC_WINDOWS_TO_GO:
 			if (Button_GetCheck(GetDlgItem(hMainDialog, IDC_WINDOWS_TO_GO)) == BST_CHECKED) {
 				for (i=0; i<ComboBox_GetCount(hFileSystem); i++) {
@@ -2173,6 +2180,8 @@ static INT_PTR CALLBACK MainCallback(HWND hDlg, UINT message, WPARAM wParam, LPA
 						break;
 					}
 				}
+				SetMBRProps();
+				CheckDlgButton(hMainDialog, IDC_RUFUS_MBR, BST_UNCHECKED);
 			}
 			break;
 		case IDC_RUFUS_MBR:
diff --git a/src/rufus.h b/src/rufus.h
index e4243dcc..cfe0e926 100644
--- a/src/rufus.h
+++ b/src/rufus.h
@@ -341,7 +341,7 @@ extern DWORD syslinux_ldlinux_len[2];
 extern RUFUS_DRIVE_INFO SelectedDrive;
 extern const int nb_steps[FS_MAX];
 extern BOOL use_own_c32[NB_OLD_C32], detect_fakes, iso_op_in_progress, format_op_in_progress, right_to_left_mode;
-extern BOOL allow_dual_uefi_bios;
+extern BOOL allow_dual_uefi_bios, togo_mode;
 extern RUFUS_ISO_REPORT iso_report;
 extern int64_t iso_blocking_status;
 extern uint16_t rufus_version[4], embedded_sl_version[2];
diff --git a/src/rufus.rc b/src/rufus.rc
index c2dcde34..079ba36e 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.0.0.581"
+CAPTION "Rufus 2.0.0.582"
 FONT 8, "Segoe UI", 400, 0, 0x1
 BEGIN
     DEFPUSHBUTTON   "Start",IDC_START,127,339,50,14
@@ -157,7 +157,7 @@ END
 
 IDD_DIALOG_XP DIALOGEX 12, 12, 242, 376
 STYLE DS_SETFONT | DS_MODALFRAME | DS_CENTER | WS_MINIMIZEBOX | WS_POPUP | WS_CAPTION | WS_SYSMENU
-CAPTION "Rufus 2.0.0.581"
+CAPTION "Rufus 2.0.0.582"
 FONT 8, "MS Shell Dlg", 400, 0, 0x1
 BEGIN
     DEFPUSHBUTTON   "Start",IDC_START,127,339,50,14
@@ -283,7 +283,7 @@ END
 IDD_DIALOG_RTL DIALOGEX 12, 12, 242, 376
 STYLE DS_SETFONT | DS_MODALFRAME | DS_CENTER | WS_MINIMIZEBOX | WS_POPUP | WS_CAPTION | WS_SYSMENU
 EXSTYLE WS_EX_RTLREADING | WS_EX_APPWINDOW | WS_EX_LAYOUTRTL
-CAPTION "Rufus 2.0.0.581"
+CAPTION "Rufus 2.0.0.582"
 FONT 8, "Segoe UI", 400, 0, 0x1
 BEGIN
     DEFPUSHBUTTON   "Start",IDC_START,127,339,50,14
@@ -415,7 +415,7 @@ END
 IDD_DIALOG_RTL_XP DIALOGEX 12, 12, 242, 376
 STYLE DS_SETFONT | DS_MODALFRAME | DS_CENTER | WS_MINIMIZEBOX | WS_POPUP | WS_CAPTION | WS_SYSMENU
 EXSTYLE WS_EX_RTLREADING | WS_EX_APPWINDOW | WS_EX_LAYOUTRTL
-CAPTION "Rufus 2.0.0.581"
+CAPTION "Rufus 2.0.0.582"
 FONT 8, "MS Shell Dlg", 400, 0, 0x1
 BEGIN
     DEFPUSHBUTTON   "Start",IDC_START,127,339,50,14
@@ -671,8 +671,8 @@ END
 //
 
 VS_VERSION_INFO VERSIONINFO
- FILEVERSION 2,0,0,581
- PRODUCTVERSION 2,0,0,581
+ FILEVERSION 2,0,0,582
+ PRODUCTVERSION 2,0,0,582
  FILEFLAGSMASK 0x3fL
 #ifdef _DEBUG
  FILEFLAGS 0x1L
@@ -689,13 +689,13 @@ BEGIN
         BEGIN
             VALUE "CompanyName", "Akeo Consulting (http://akeo.ie)"
             VALUE "FileDescription", "Rufus"
-            VALUE "FileVersion", "2.0.0.581"
+            VALUE "FileVersion", "2.0.0.582"
             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.0.0.581"
+            VALUE "ProductVersion", "2.0.0.582"
         END
     END
     BLOCK "VarFileInfo"
diff --git a/src/sys_types.h b/src/sys_types.h
index 7dd4044e..adf88d39 100644
--- a/src/sys_types.h
+++ b/src/sys_types.h
@@ -32,6 +32,7 @@ typedef struct {
 /*
  * File system types for MBR partition tables
  * See http://en.wikipedia.org/wiki/Partition_type
+ * Also http://www.win.tue.nl/~aeb/partitions/partition_types-1.html
  * Note: If googling APTI (Alternative Partition Table Identification)
  * doesn't return squat, then IT ISN'T A REAL THING!!
  */
diff --git a/src/vhd.c b/src/vhd.c
index 96c6ece7..5cea1a4b 100644
--- a/src/vhd.c
+++ b/src/vhd.c
@@ -109,6 +109,9 @@ PF_TYPE_DECL(RPC_ENTRY, RPC_STATUS, UuidCreate, (UUID __RPC_FAR*));
 static uint8_t wim_flags = 0;
 static char sevenzip_path[MAX_PATH];
 static const char conectix_str[] = VHD_FOOTER_COOKIE;
+static uint32_t wim_nb_files, wim_proc_files;
+static BOOL count_files;
+static DWORD LastRefresh;
 
 static BOOL Get7ZipPath(void)
 {
@@ -491,6 +494,7 @@ enum WIMMessage {
 	WIM_MSG_ALIGNMENT,	// Enables the caller to align a file resource on a particular alignment boundary.
 	WIM_MSG_RETRY,		// Sent when the file is being reapplied because of a network timeout.
 	WIM_MSG_SPLIT,		// Enables the caller to align a file resource on a particular alignment boundary.
+	WIM_MSG_FILEINFO,	// Used in conjunction with WimApplyImages()'s WIM_FLAG_FILEINFO flag to provide detailed file info.
 	WIM_MSG_INFO,		// Sent when an info message is available.
 	WIM_MSG_WARNING,	// Sent when a warning message is available.
 	WIM_MSG_CHK_PROCESS,
@@ -500,30 +504,74 @@ enum WIMMessage {
 
 #define INVALID_CALLBACK_VALUE 0xFFFFFFFF
 
+#define WIM_FLAG_RESERVED      0x00000001
+#define WIM_FLAG_VERIFY        0x00000002
+#define WIM_FLAG_INDEX         0x00000004
+#define WIM_FLAG_NO_APPLY      0x00000008
+#define WIM_FLAG_NO_DIRACL     0x00000010
+#define WIM_FLAG_NO_FILEACL    0x00000020
+#define WIM_FLAG_SHARE_WRITE   0x00000040
+#define WIM_FLAG_FILEINFO      0x00000080
+#define WIM_FLAG_NO_RP_FIX     0x00000100
+
 // Progress callback
 DWORD WINAPI WimProgressCallback(DWORD dwMsgId, WPARAM wParam, LPARAM lParam, PVOID pvIgnored)
 {
 	PBOOL pbCancel = NULL;
+	PWIN32_FIND_DATA pFileData;
 	char* str = NULL;
 	const char* level = NULL;
+	uint64_t size;
+	float apply_percent;
 
 	switch (dwMsgId) {
 	case WIM_MSG_PROGRESS:
+		// The default WIM progress is useless (freezes at 95%, which is usually when only half
+		// the files have been processed), so we don't use it
+#if 0
 		PrintInfo(0, MSG_267, (DWORD)wParam);
 		UpdateProgress(OP_DOS, 0.98f*(DWORD)wParam);
+#endif
 		break;
 	case WIM_MSG_PROCESS:
-		// The amount of files processed is a bit overwhelming, and displaying it all slows us down
+		// The amount of files processed is overwhelming (16k+ for a typical image),
+		// and trying to display it *WILL* slow us down, so we don't.
 #if 0
 		str = wchar_to_utf8((PWSTR)wParam);
-		uprintf("Applying: '%s'", str);
+		uprintf("%s", str);
 		PrintStatus(0, MSG_000, str);	// MSG_000 is "%s"
 #endif
+		if (count_files) {
+			wim_nb_files++;
+		} else {
+			wim_proc_files++;
+			if (GetTickCount() > LastRefresh + 100) {
+				// At the end of an actual apply, the WIM API re-lists a bunch of directories it
+				// already processed, so we end up with more entries than counted - ignore those.
+				if (wim_proc_files > wim_nb_files)
+					wim_proc_files = wim_nb_files;
+				LastRefresh = GetTickCount();
+				// x^3 progress, so as not to give a better idea right from the onset
+				// as to the dismal speed with which the WIM API can actually apply files...
+				apply_percent = 4.636942595f * ((float)wim_proc_files) / ((float)wim_nb_files);
+				apply_percent = apply_percent * apply_percent * apply_percent;
+				PrintInfo(0, MSG_267, apply_percent);
+				UpdateProgress(OP_DOS, apply_percent);
+			}
+		}
+		// Halt on error
 		if (IS_ERROR(FormatStatus)) {
 			pbCancel = (PBOOL)lParam;
 			*pbCancel = TRUE;
+			break;
 		}
 		break;
+	case WIM_MSG_FILEINFO:
+		str = wchar_to_utf8((PWSTR)wParam);
+		pFileData = (PWIN32_FIND_DATA)lParam;
+		size = (((uint64_t)pFileData->nFileSizeHigh) << 32) + pFileData->nFileSizeLow;
+		uprintf("'%s' (%s)", str, SizeToHumanReadable(size, FALSE, FALSE));
+		break;
 	case WIM_MSG_RETRY:
 		level = "retry";
 		// fall through
@@ -596,13 +644,24 @@ static DWORD WINAPI WimApplyImageThread(LPVOID param)
 	}
 
 	uprintf("Applying Windows image...");
+	// Run a first pass using WIM_FLAG_NO_APPLY to count the files
+	wim_nb_files = 0;
+	wim_proc_files = 0;
+	LastRefresh = 0;
+	count_files = TRUE;
+	if (!pfWIMApplyImage(hImage, wdst, WIM_FLAG_NO_APPLY)) {
+		uprintf("  Could not count the files to apply: %s", WindowsErrorString());
+		goto out;
+	}
+	count_files = FALSE;
+	// Actual apply
 	if (!pfWIMApplyImage(hImage, wdst, 0)) {
 		uprintf("  Could not apply image: %s", WindowsErrorString());
 		goto out;
 	}
-
+	PrintInfo(0, MSG_267, 99.8f);
+	UpdateProgress(OP_DOS, 99.8f);
 	r = TRUE;
-	UpdateProgress(OP_FINALIZE, -1.0f);
 
 out:
 	if ((hImage != NULL) || (hWim != NULL)) {