Merge pull request #116 from endlessm/T13969

Refactor and improve CreateUSBStick

src/drive.c

@@ -889,8 +889,62 @@ BOOL MountVolume(char* drive_name, char *drive_guid)
 }
 
 /*
- * Mount partition #part_nr, residing on the same disk as drive_name to an available
- * drive letter. Returns the newly allocated drive string.
+ * Mount the volume identified by drive_guid on a drive letter, if it is not already, and return the mount point in drive_guid.
+ * drive_guid should be initialized to "?:\"; the first character will be overwritten with the mounted drive letter.
+ */
+BOOL EnsureVolumeMounted(char* drive_name, char *drive_guid)
+{
+	char mounted_guid[52];	// You need at least 51 characters on XP
+	char mounted_letter[16] = {0};
+	DWORD size;
+
+	if (drive_name[0] != '?')
+		return FALSE;
+
+	// For fixed disks, Windows may already have remounted the volume.
+	if ( (GetVolumePathNamesForVolumeNameA(drive_guid, mounted_letter, sizeof(mounted_letter), &size))
+	  && (size > 1) ) {
+		uprintf("Volume is already mounted at %c", mounted_letter[0]);
+		drive_name[0] = mounted_letter[0];
+		return TRUE;
+	}
+
+	drive_name[0] = GetUnusedDriveLetter();
+	if (drive_name[0] == 0) {
+		uprintf("Could not find an unused drive letter");
+		return FALSE;
+	}
+
+	if (SetVolumeMountPointA(drive_name, drive_guid)) {
+		uprintf("%s mounted as %s", drive_guid, drive_name);
+		return TRUE;
+	} else {
+		// If the OS mounted something on that letter in parallel, this operation can fail
+		// with ERROR_DIR_NOT_EMPTY. If that's the case, just check that mountpoints match
+		if (GetLastError() == ERROR_DIR_NOT_EMPTY) {
+			if (!GetVolumeNameForVolumeMountPointA(drive_name, mounted_guid, sizeof(mounted_guid))) {
+				uprintf("%s already mounted, but volume GUID could not be checked: %s",
+					drive_name, WindowsErrorString());
+				return FALSE;
+			}
+			if (safe_strcmp(drive_guid, mounted_guid) != 0) {
+				uprintf("%s already mounted, but volume GUID doesn't match. Hoped for %s, got %s",
+					drive_name, drive_guid, mounted_guid);
+				return FALSE;
+			}
+			uprintf("%s was already mounted as %s\n", drive_guid, drive_name);
+			return TRUE;
+		} else {
+			return FALSE;
+		}
+	}
+}
+
+/*
+ * Mount partition #part_nr, residing on the same disk as drive_name (which must be
+ * of the form "X:" with no trailing slash) to an available
+ * drive letter. Returns a borrowed pointer to the newly-allocated drive letter of
+ * the form "Y:", which must not be freed, or NULL on error.
  * We need to do this because, for instance, EFI System partitions are not assigned
  * Volume GUIDs by the OS, and we need to have a letter assigned, for when we invoke
  * bcdtool for Windows To Go. All in all, the process looks like this:

src/drive.h

@@ -93,6 +93,7 @@ BOOL AnalyzePBR(HANDLE hLogicalVolume);
 BOOL GetDrivePartitionData(DWORD DriveIndex, char* FileSystemName, DWORD FileSystemNameSize, BOOL bSilent);
 BOOL UnmountVolume(HANDLE hDrive);
 BOOL MountVolume(char* drive_name, char *drive_guid);
+BOOL EnsureVolumeMounted(char* drive_name, char *drive_guid);
 BOOL AltUnmountVolume(const char* drive_name);
 char* AltMountVolume(const char* drive_name, uint8_t part_nr);
 BOOL RemountVolume(char* drive_name);

src/endless/EndlessUsbTool.vcxproj

@@ -271,6 +271,7 @@
   <ItemGroup>
     <ClInclude Include="..\dos.h" />
     <ClInclude Include="..\drive.h" />
+    <ClInclude Include="..\format.h" />
     <ClInclude Include="..\localization.h" />
     <ClInclude Include="..\msapi_utf8.h" />
     <ClInclude Include="..\rufus.h" />
@@ -289,6 +290,7 @@
     <ClInclude Include="grub\grub.h" />
     <ClInclude Include="grub\reed_solomon.h" />
     <ClInclude Include="Images.h" />
+    <ClInclude Include="IOUtil.h" />
     <ClInclude Include="json\json-forwards.h" />
     <ClInclude Include="json\json.h" />
     <ClInclude Include="localization_data.h" />
@@ -513,6 +515,7 @@
       <PrecompiledHeader Condition="'$(Configuration)|$(Platform)'=='Release|x64'">NotUsing</PrecompiledHeader>
     </ClCompile>
     <ClCompile Include="Images.cpp" />
+    <ClCompile Include="IOUtil.cpp" />
     <ClCompile Include="StringHelperMethods.cpp" />
     <ClCompile Include="jsoncpp.cpp">
       <PrecompiledHeader Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'">NotUsing</PrecompiledHeader>

src/endless/EndlessUsbTool.vcxproj.filters

@@ -90,6 +90,9 @@
     <ClInclude Include="..\dos.h">
       <Filter>Header Files\Rufus Headers</Filter>
     </ClInclude>
+    <ClInclude Include="..\format.h">
+      <Filter>Header Files\Rufus Headers</Filter>
+    </ClInclude>
     <ClInclude Include="DownloadManager.h">
       <Filter>Header Files</Filter>
     </ClInclude>
@@ -153,6 +156,9 @@
     <ClInclude Include="UEFIBootSetup.h">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="IOUtil.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
   </ItemGroup>
   <ItemGroup>
     <ClCompile Include="EndlessUsbTool.cpp">
@@ -257,6 +263,9 @@
     <ClCompile Include="Eosldr.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="IOUtil.cpp">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
   <ItemGroup>
     <None Include="res\EndlessUsbTool.rc2">