From 5c5cd20b9de38568fc48801f9259bf211c94ce37 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adam=20S=C5=82abo=C5=84?= <asaillen@protonmail.com>
Date: Sun, 3 Mar 2024 00:27:49 +0100
Subject: [PATCH] [FMIFS] Improve FormatEx and Chkdsk

- Remove the hack for missing IOCTL_MOUNTDEV_QUERY_DEVICE_NAME
- Handle DriveRoot without a trailing backslash
---
 dll/win32/fmifs/chkdsk.c | 45 ++++++++++++++++++++++++----------------
 dll/win32/fmifs/format.c | 45 ++++++++++++++++++++++++----------------
 2 files changed, 54 insertions(+), 36 deletions(-)

diff --git a/dll/win32/fmifs/chkdsk.c b/dll/win32/fmifs/chkdsk.c
index c3be682eeea..5a24a3f8645 100644
--- a/dll/win32/fmifs/chkdsk.c
+++ b/dll/win32/fmifs/chkdsk.c
@@ -8,6 +8,7 @@
  */
 
 #include "precomp.h"
+#include <ntstrsafe.h>
 
 #define NDEBUG
 #include <debug.h>
@@ -30,31 +31,38 @@ Chkdsk(
     UNICODE_STRING usDriveRoot;
     NTSTATUS Status;
     BOOLEAN Success = FALSE;
+    WCHAR DriveName[MAX_PATH];
     WCHAR VolumeName[MAX_PATH];
-    //CURDIR CurDir;
 
     Provider = GetProvider(Format);
     if (!Provider)
     {
         /* Unknown file system */
-        Callback(DONE, 0, &Success);
-        return;
+        goto Quit;
     }
 
-#if 1
-    DPRINT1("Warning: use GetVolumeNameForVolumeMountPointW() instead!\n");
-    swprintf(VolumeName, L"\\??\\%c:", towupper(DriveRoot[0]));
-    RtlCreateUnicodeString(&usDriveRoot, VolumeName);
-    /* Code disabled as long as our storage stack doesn't understand IOCTL_MOUNTDEV_QUERY_DEVICE_NAME */
-#else
-    if (!GetVolumeNameForVolumeMountPointW(DriveRoot, VolumeName, RTL_NUMBER_OF(VolumeName)) ||
-        !RtlDosPathNameToNtPathName_U(VolumeName, &usDriveRoot, NULL, &CurDir))
+    if (!NT_SUCCESS(RtlStringCchCopyW(DriveName, ARRAYSIZE(DriveName), DriveRoot)))
+        goto Quit;
+
+    if (DriveName[wcslen(DriveName) - 1] != L'\\')
     {
-        /* Report an error */
-        Callback(DONE, 0, &Success);
-        return;
+        /* Append the trailing backslash for GetVolumeNameForVolumeMountPointW */
+        if (!NT_SUCCESS(RtlStringCchCatW(DriveName, ARRAYSIZE(DriveName), L"\\")))
+            goto Quit;
     }
-#endif
+
+    if (!GetVolumeNameForVolumeMountPointW(DriveName, VolumeName, ARRAYSIZE(VolumeName)))
+    {
+        /* Couldn't get a volume GUID path, try checking using a parameter provided path */
+        DPRINT1("Couldn't get a volume GUID path for drive %S\n", DriveName);
+        wcscpy(VolumeName, DriveName);
+    }
+
+    if (!RtlDosPathNameToNtPathName_U(VolumeName, &usDriveRoot, NULL, NULL))
+        goto Quit;
+
+    /* Trim the trailing backslash since we will work with a device object */
+    usDriveRoot.Length -= sizeof(WCHAR);
 
     DPRINT("Chkdsk() - %S\n", Format);
     Status = STATUS_SUCCESS;
@@ -72,10 +80,11 @@ Chkdsk(
     if (!Success)
         DPRINT1("Chkdsk() failed with Status 0x%lx\n", Status);
 
-    /* Report success */
-    Callback(DONE, 0, &Success);
-
     RtlFreeUnicodeString(&usDriveRoot);
+
+Quit:
+    /* Report result */
+    Callback(DONE, 0, &Success);
 }
 
 /* EOF */
diff --git a/dll/win32/fmifs/format.c b/dll/win32/fmifs/format.c
index 82fa203b1e4..ecab033f98e 100644
--- a/dll/win32/fmifs/format.c
+++ b/dll/win32/fmifs/format.c
@@ -9,6 +9,7 @@
  */
 
 #include "precomp.h"
+#include <ntstrsafe.h>
 
 #define NDEBUG
 #include <debug.h>
@@ -50,8 +51,8 @@ FormatEx(
     BOOLEAN Success = FALSE;
     BOOLEAN BackwardCompatible = FALSE; // Default to latest FS versions.
     MEDIA_TYPE MediaType;
+    WCHAR DriveName[MAX_PATH];
     WCHAR VolumeName[MAX_PATH];
-    //CURDIR CurDir;
 
 //
 // TODO: Convert filesystem Format into ULIB format string.
@@ -61,24 +62,31 @@ FormatEx(
     if (!Provider)
     {
         /* Unknown file system */
-        Callback(DONE, 0, &Success);
-        return;
+        goto Quit;
     }
 
-#if 1
-    DPRINT1("Warning: use GetVolumeNameForVolumeMountPointW() instead!\n");
-    swprintf(VolumeName, L"\\??\\%c:", towupper(DriveRoot[0]));
-    RtlCreateUnicodeString(&usDriveRoot, VolumeName);
-    /* Code disabled as long as our storage stack doesn't understand IOCTL_MOUNTDEV_QUERY_DEVICE_NAME */
-#else
-    if (!GetVolumeNameForVolumeMountPointW(DriveRoot, VolumeName, RTL_NUMBER_OF(VolumeName)) ||
-        !RtlDosPathNameToNtPathName_U(VolumeName, &usDriveRoot, NULL, &CurDir))
+    if (!NT_SUCCESS(RtlStringCchCopyW(DriveName, ARRAYSIZE(DriveName), DriveRoot)))
+        goto Quit;
+
+    if (DriveName[wcslen(DriveName) - 1] != L'\\')
     {
-        /* Report an error */
-        Callback(DONE, 0, &Success);
-        return;
+        /* Append the trailing backslash for GetVolumeNameForVolumeMountPointW */
+        if (!NT_SUCCESS(RtlStringCchCatW(DriveName, ARRAYSIZE(DriveName), L"\\")))
+            goto Quit;
     }
-#endif
+
+    if (!GetVolumeNameForVolumeMountPointW(DriveName, VolumeName, ARRAYSIZE(VolumeName)))
+    {
+        /* Couldn't get a volume GUID path, try formatting using a parameter provided path */
+        DPRINT1("Couldn't get a volume GUID path for drive %S\n", DriveName);
+        wcscpy(VolumeName, DriveName);
+    }
+
+    if (!RtlDosPathNameToNtPathName_U(VolumeName, &usDriveRoot, NULL, NULL))
+        goto Quit;
+
+    /* Trim the trailing backslash since we will work with a device object */
+    usDriveRoot.Length -= sizeof(WCHAR);
 
     RtlInitUnicodeString(&usLabel, Label);
 
@@ -119,10 +127,11 @@ FormatEx(
     if (!Success)
         DPRINT1("Format() failed\n");
 
-    /* Report success */
-    Callback(DONE, 0, &Success);
-
     RtlFreeUnicodeString(&usDriveRoot);
+
+Quit:
+    /* Report result */
+    Callback(DONE, 0, &Success);
 }
 
 /* EOF */