[core] make sure the system disk is never listed

* When using VHD or VMDK drives, it was possible for the system disk to be listed.
* Closes #541

src/format.c

@@ -1276,7 +1276,8 @@ BOOL SetupWinToGo(const char* drive_name, BOOL use_ms_efi)
 	static char san_policy_path[] = "?:\\san_policy.xml";
 #endif
 	static char unattend_path[] = "?:\\Windows\\System32\\sysprep\\unattend.xml";
-	char *mounted_iso, *ms_efi = NULL, image[128], cmd[MAX_PATH], system_root[128];
+	char *mounted_iso, *ms_efi = NULL, image[128], cmd[MAX_PATH];
+	char usb_system_dir[] = "?:\\Windows\\System32";
 	unsigned char *buffer;
 	int i;
 	wchar_t wVolumeName[] = L"?:";
@@ -1287,6 +1288,7 @@ BOOL SetupWinToGo(const char* drive_name, BOOL use_ms_efi)
 	PF_INIT(FormatEx, Fmifs);
 
 	uprintf("Windows To Go mode selected");
+	usb_system_dir[0] = drive_name[0];
 	// Additional sanity checks
 	if ( ((use_ms_efi) && (SelectedDrive.Geometry.MediaType != FixedMedia)) ||
 		 ((nWindowsVersion < WINDOWS_8) || ((WimExtractCheck() & 4) == 0)) ) {
@@ -1353,16 +1355,12 @@ BOOL SetupWinToGo(const char* drive_name, BOOL use_ms_efi)
 	// Try the 'bcdboot' command, first using the one from the target drive and, if that doesn't work, the system's
 	uprintf("Enabling boot...");
 	for (i = 0; i < 2; i++) {
-		if (i == 0)
-			static_sprintf(system_root, "%s\\Windows\\System32", drive_name);
-		else
-			GetSystemDirectoryA(system_root, sizeof(system_root));
 		static_sprintf(cmd, "%s\\bcdboot.exe %s\\Windows /f ALL /s %s",
-			system_root, drive_name, (use_ms_efi)?ms_efi:drive_name);
+			(i==0)?usb_system_dir:system_dir, drive_name, (use_ms_efi)?ms_efi:drive_name);
 		if (RunCommand(cmd, NULL, TRUE) == 0)
 			break;
 	}
-	if (i >= 2)	{
+	if (i >= 2) {
 		// Fatal, as the UFD is unlikely to boot then
 		uprintf("Failed to enable boot - aborting", cmd);
 		FormatStatus = ERROR_SEVERITY_ERROR|FAC(FACILITY_STORAGE)|APPERR(ERROR_ISO_EXTRACT);

src/msapi_utf8.h

@@ -407,6 +407,20 @@ static __inline DWORD GetCurrentDirectoryU(DWORD nBufferLength, char* lpBuffer)
 	return ret;
 }
 
+static __inline UINT GetSystemDirectoryU(char* lpBuffer, UINT uSize)
+{
+	UINT ret = 0, err = ERROR_INVALID_DATA;
+	walloc(lpBuffer, uSize);
+	ret = GetSystemDirectoryW(wlpBuffer, uSize);
+	err = GetLastError();
+	if ((ret != 0) && ((ret = wchar_to_utf8_no_alloc(wlpBuffer, lpBuffer, uSize)) == 0)) {
+		err = GetLastError();
+	}
+	wfree(lpBuffer);
+	SetLastError(err);
+	return ret;
+}
+
 static __inline DWORD GetTempPathU(DWORD nBufferLength, char* lpBuffer)
 {
 	DWORD ret = 0, err = ERROR_INVALID_DATA;

src/rufus.c

@@ -125,7 +125,7 @@ char lost_translators[][6] = LOST_TRANSLATORS;
 OPENED_LIBRARIES_VARS;
 HINSTANCE hMainInstance;
 HWND hMainDialog, hLangToolbar = NULL, hUpdatesDlg = NULL;
-char szFolderPath[MAX_PATH], app_dir[MAX_PATH];
+char szFolderPath[MAX_PATH], app_dir[MAX_PATH], system_dir[MAX_PATH];
 char* image_path = NULL;
 float fScale = 1.0f;
 int default_fs;
@@ -2677,11 +2677,18 @@ int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLine
 				goto out;
 		}
 	} else {
-		uprintf("unable to access UTF-16 args");
+		uprintf("Could not access UTF-16 args");
 	}
 
-	// Retrieve the current application directory
+	// Retrieve the current application directory as well as the system dir
-	GetCurrentDirectoryU(MAX_PATH, app_dir);
+	if (GetCurrentDirectoryU(sizeof(app_dir), app_dir) == 0) {
+		uprintf("Could not get current directory: %s", WindowsErrorString());
+		app_dir[0] = 0;
+	}
+	if (GetSystemDirectoryU(system_dir, sizeof(system_dir)) == 0) {
+		uprintf("Could not get system directory: %s", WindowsErrorString());
+		safe_strcpy(system_dir, sizeof(system_dir), "C:\\Windows\\System32");
+	}
 
 	// Look for a .ini file in the current app directory
 	static_sprintf(ini_path, "%s\\rufus.ini", app_dir);

src/rufus.h

@@ -350,7 +350,7 @@ extern HWND hMainDialog, hLogDlg, hStatus, hDeviceList, hCapacity;
 extern HWND hPartitionScheme, hFileSystem, hClusterSize, hLabel, hBootType, hNBPasses, hLog;
 extern HWND hInfo, hProgress, hDiskID;
 extern float fScale;
-extern char szFolderPath[MAX_PATH], app_dir[MAX_PATH];
+extern char szFolderPath[MAX_PATH], app_dir[MAX_PATH], system_dir[MAX_PATH];
 extern char* image_path;
 extern DWORD FormatStatus, DownloadStatus;
 extern BOOL PromptOnError;

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.3.683"
+CAPTION "Rufus 2.3.684"
 FONT 8, "Segoe UI", 400, 0, 0x1
 BEGIN
     LTEXT           "Device",IDS_DEVICE_TXT,9,6,200,8
@@ -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.3.683"
+CAPTION "Rufus 2.3.684"
 FONT 8, "MS Shell Dlg", 400, 0, 0x1
 BEGIN
     LTEXT           "Device",IDS_DEVICE_TXT,9,6,200,8
@@ -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.3.683"
+CAPTION "Rufus 2.3.684"
 FONT 8, "Segoe UI", 400, 0, 0x1
 BEGIN
     LTEXT           "Device",IDS_DEVICE_TXT,9,6,200,8
@@ -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.3.683"
+CAPTION "Rufus 2.3.684"
 FONT 8, "MS Shell Dlg", 400, 0, 0x1
 BEGIN
     LTEXT           "Device",IDS_DEVICE_TXT,9,6,200,8
@@ -671,8 +671,8 @@ END
 //
 
 VS_VERSION_INFO VERSIONINFO
- FILEVERSION 2,3,683,0
+ FILEVERSION 2,3,684,0
- PRODUCTVERSION 2,3,683,0
+ PRODUCTVERSION 2,3,684,0
  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.3.683"
+            VALUE "FileVersion", "2.3.684"
             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.3.683"
+            VALUE "ProductVersion", "2.3.684"
         END
     END
     BLOCK "VarFileInfo"

src/usb.c

@@ -147,7 +147,7 @@ BOOL GetUSBDevices(DWORD devnum)
 	ULONG list_size[ARRAYSIZE(storage_name)] = { 0 }, list_start[ARRAYSIZE(storage_name)] = { 0 }, full_list_size, ulFlags;
 	HANDLE hDrive;
 	LONG maxwidth = 0;
-	int s, score, drive_number;
+	int s, score, drive_number, remove_drive;
 	char drive_letters[27], *device_id, *devid_list = NULL, entry_msg[128];
 	char *label, *entry, buffer[MAX_PATH], str[MAX_PATH], *method_str;
 	usb_device_props props;
@@ -448,17 +448,21 @@ BOOL GetUSBDevices(DWORD devnum)
 					// If that is the case, use "Multiple Volumes" instead of the label
 					safe_strcpy(entry_msg, sizeof(entry_msg), ((drive_letters[0] != 0) && (drive_letters[1] != 0))?
 						lmprintf(MSG_047):label);
-					for (k=0; drive_letters[k]; k++) {
+					for (k=0, remove_drive=0; drive_letters[k] && (!remove_drive); k++) {
 						// Append all the drive letters we detected
 						letter_name[2] = drive_letters[k];
 						if (right_to_left_mode)
 							safe_strcat(entry_msg, sizeof(entry_msg), RIGHT_TO_LEFT_MARK);
 						safe_strcat(entry_msg, sizeof(entry_msg), letter_name);
-						if (drive_letters[k] == (PathGetDriveNumberU(app_dir) + 'A')) break;
+						if (drive_letters[k] == (PathGetDriveNumberU(app_dir) + 'A'))
+							remove_drive = 1;
+						if (drive_letters[k] == (PathGetDriveNumberU(system_dir) + 'A'))
+							remove_drive = 2;
 					}
-					// Repeat as we need to break the outside loop
+					// Make sure that we don't list any drive that should not be listed
-					if (drive_letters[k] == (PathGetDriveNumberU(app_dir) + 'A')) {
+					if (remove_drive) {
-						uprintf("Removing %c: from the list: This is the disk from which " APPLICATION_NAME " is running!\n", app_dir[0]);
+						uprintf("Removing %C: from the list: This is the %s!", drive_letters[--k],
+							(remove_drive==1)?"disk from which " APPLICATION_NAME " is running":"system disk");
 						safe_closehandle(hDrive);
 						safe_free(devint_detail_data);
 						break;