[SETUPLIB] Make InferFileSystem() compatible for both MBR and GPT disks. Simplify the FSREC interface a bit.

Instead of providing an MBR partition type to InferFileSystem(), make
it call IOCTL_DISK_GET_PARTITION_INFO(_EX) to determine whether the
partition pointed by the path/handle is MBR or GPT. Then, only if it's
MBR, we retrieve its partition type in order to "guess" an adequate file
system name, in case the latter was not recognized already via regular
ways (via GetFileSystemName() / NtQueryVolumeInformationFile()).

- Remove the GetFileSystemNameByHandle() and InferFileSystemByHandle()
  functions. Instead, make the other GetFileSystemName*() and
  InferFileSystem*() functions accept a HANDLE as an alternative to the
  already-existing partition path string. These parameters are exclusive
  to each other.

- Rename SetPartitionType() -> SetMBRPartitionType(),
  and FileSystemToPartitionType() -> FileSystemToMBRPartitionType()
  in order to really clarify what they do (since this code is meant
  for MBR partitions only, not GPT ones).

base/setup/lib/fsutil.c

@@ -679,9 +679,9 @@ PreparePartitionForFormatting(
         return FALSE;
     }
 
-    PartitionType = FileSystemToPartitionType(FileSystemName,
-                                              &PartEntry->StartSector,
-                                              &PartEntry->SectorCount);
+    PartitionType = FileSystemToMBRPartitionType(FileSystemName,
+                                                 PartEntry->StartSector.QuadPart,
+                                                 PartEntry->SectorCount.QuadPart);
     if (PartitionType == PARTITION_ENTRY_UNUSED)
     {
         /* Unknown file system */
@@ -689,7 +689,7 @@ PreparePartitionForFormatting(
         return FALSE;
     }
 
-    SetPartitionType(PartEntry, PartitionType);
+    SetMBRPartitionType(PartEntry, PartitionType);
 
     /*
      * Adjust the filesystem name in case of FAT vs. FAT32, according to

base/setup/lib/utils/fsrec.c

@@ -19,26 +19,27 @@
 /* FUNCTIONS ****************************************************************/
 
 /* NOTE: Ripped & adapted from base/system/autochk/autochk.c */
+static inline
 NTSTATUS
-GetFileSystemNameByHandle(
+GetFileSystemNameWorker(
     IN HANDLE PartitionHandle,
+    OUT PIO_STATUS_BLOCK IoStatusBlock,
     IN OUT PWSTR FileSystemName,
     IN SIZE_T FileSystemNameSize)
 {
     NTSTATUS Status;
-    IO_STATUS_BLOCK IoStatusBlock;
     UCHAR Buffer[sizeof(FILE_FS_ATTRIBUTE_INFORMATION) + MAX_PATH * sizeof(WCHAR)];
     PFILE_FS_ATTRIBUTE_INFORMATION FileFsAttribute = (PFILE_FS_ATTRIBUTE_INFORMATION)Buffer;
 
     /* Retrieve the FS attributes */
     Status = NtQueryVolumeInformationFile(PartitionHandle,
-                                          &IoStatusBlock,
+                                          IoStatusBlock,
                                           FileFsAttribute,
                                           sizeof(Buffer),
                                           FileFsAttributeInformation);
     if (!NT_SUCCESS(Status))
     {
-        DPRINT1("NtQueryVolumeInformationFile failed, Status 0x%08lx\n", Status);
+        DPRINT1("NtQueryVolumeInformationFile() failed, Status 0x%08lx\n", Status);
         return Status;
     }
 
@@ -52,69 +53,97 @@ GetFileSystemNameByHandle(
 
 NTSTATUS
 GetFileSystemName_UStr(
-    IN PUNICODE_STRING PartitionPath,
+    IN PUNICODE_STRING PartitionPath OPTIONAL,
+    IN HANDLE PartitionHandle OPTIONAL,
     IN OUT PWSTR FileSystemName,
     IN SIZE_T FileSystemNameSize)
 {
     NTSTATUS Status;
     OBJECT_ATTRIBUTES ObjectAttributes;
-    HANDLE PartitionHandle;
     IO_STATUS_BLOCK IoStatusBlock;
 
-    /* Open the partition */
-    InitializeObjectAttributes(&ObjectAttributes,
-                               PartitionPath,
-                               OBJ_CASE_INSENSITIVE,
-                               NULL,
-                               NULL);
-    Status = NtOpenFile(&PartitionHandle,
-                        FILE_GENERIC_READ /* | SYNCHRONIZE */,
-                        &ObjectAttributes,
-                        &IoStatusBlock,
-                        FILE_SHARE_READ | FILE_SHARE_WRITE,
-                        0 /* FILE_SYNCHRONOUS_IO_NONALERT */);
-    if (!NT_SUCCESS(Status))
+    if (PartitionPath && PartitionHandle)
+        return STATUS_INVALID_PARAMETER;
+
+    /* Open the partition if a path has been given;
+     * otherwise just use the provided handle. */
+    if (PartitionPath)
     {
-        DPRINT1("Failed to open partition '%wZ', Status 0x%08lx\n", PartitionPath, Status);
-        return Status;
+        InitializeObjectAttributes(&ObjectAttributes,
+                                   PartitionPath,
+                                   OBJ_CASE_INSENSITIVE,
+                                   NULL,
+                                   NULL);
+        Status = NtOpenFile(&PartitionHandle,
+                            FILE_GENERIC_READ /* | SYNCHRONIZE */,
+                            &ObjectAttributes,
+                            &IoStatusBlock,
+                            FILE_SHARE_READ | FILE_SHARE_WRITE,
+                            0 /* FILE_SYNCHRONOUS_IO_NONALERT */);
+        if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("Failed to open partition '%wZ', Status 0x%08lx\n",
+                    PartitionPath, Status);
+            return Status;
+        }
     }
 
     /* Retrieve the FS attributes */
-    Status = GetFileSystemNameByHandle(PartitionHandle, FileSystemName, FileSystemNameSize);
+    Status = GetFileSystemNameWorker(PartitionHandle,
+                                     &IoStatusBlock,
+                                     FileSystemName,
+                                     FileSystemNameSize);
     if (!NT_SUCCESS(Status))
     {
-        DPRINT1("GetFileSystemNameByHandle() failed for partition '%wZ', Status 0x%08lx\n",
-                PartitionPath, Status);
+        DPRINT1("GetFileSystemName() failed for partition '%wZ' (0x%p), Status 0x%08lx\n",
+                PartitionPath, PartitionHandle, Status);
     }
 
-    /* Close the partition */
-    NtClose(PartitionHandle);
+    if (PartitionPath)
+    {
+        /* Close the partition */
+        NtClose(PartitionHandle);
+    }
 
     return Status;
 }
 
 NTSTATUS
 GetFileSystemName(
-    IN PCWSTR Partition,
+    IN PCWSTR PartitionPath OPTIONAL,
+    IN HANDLE PartitionHandle OPTIONAL,
     IN OUT PWSTR FileSystemName,
     IN SIZE_T FileSystemNameSize)
 {
-    UNICODE_STRING PartitionPath;
+    UNICODE_STRING PartitionPathU;
 
-    RtlInitUnicodeString(&PartitionPath, Partition);
-    return GetFileSystemName_UStr(&PartitionPath,
+    if (PartitionPath && PartitionHandle)
+        return STATUS_INVALID_PARAMETER;
+
+    if (PartitionPath)
+        RtlInitUnicodeString(&PartitionPathU, PartitionPath);
+
+    return GetFileSystemName_UStr(PartitionPath ? &PartitionPathU : NULL,
+                                  PartitionPath ? NULL : PartitionHandle,
                                   FileSystemName,
                                   FileSystemNameSize);
 }
 
+static inline
 NTSTATUS
-InferFileSystemByHandle(
+InferFileSystemWorker(
     IN HANDLE PartitionHandle,
-    IN UCHAR PartitionType,
+    OUT PIO_STATUS_BLOCK IoStatusBlock,
     IN OUT PWSTR FileSystemName,
     IN SIZE_T FileSystemNameSize)
 {
-    NTSTATUS Status;
+    NTSTATUS Status, Status2;
+    union
+    {
+        PARTITION_INFORMATION_EX InfoEx;
+        PARTITION_INFORMATION Info;
+    } PartInfo;
+    UCHAR PartitionType;
 
     if (FileSystemNameSize < sizeof(WCHAR))
         return STATUS_BUFFER_TOO_SMALL;
@@ -122,25 +151,81 @@ InferFileSystemByHandle(
     *FileSystemName = L'\0';
 
     /* Try to infer a file system using NT file system recognition */
-    Status = GetFileSystemNameByHandle(PartitionHandle,
-                                       FileSystemName,
-                                       FileSystemNameSize);
+    Status = GetFileSystemName_UStr(NULL, PartitionHandle,
+                                    FileSystemName,
+                                    FileSystemNameSize);
     if (NT_SUCCESS(Status) && *FileSystemName)
-    {
         goto Quit;
+
+    /*
+     * Check whether the partition is MBR, and if so, retrieve its MBR
+     * partition type and try to infer a preferred file system for it.
+     */
+
+    // NOTE: Use Status2 in order not to clobber the original Status.
+    Status2 = NtDeviceIoControlFile(PartitionHandle,
+                                    NULL,
+                                    NULL,
+                                    NULL,
+                                    IoStatusBlock,
+                                    IOCTL_DISK_GET_PARTITION_INFO_EX,
+                                    NULL,
+                                    0,
+                                    &PartInfo.InfoEx,
+                                    sizeof(PartInfo.InfoEx));
+    if (!NT_SUCCESS(Status2))
+    {
+        DPRINT1("IOCTL_DISK_GET_PARTITION_INFO_EX failed (Status %lx)\n", Status2);
+
+        if (Status2 != STATUS_INVALID_DEVICE_REQUEST)
+            goto Quit;
+
+        /*
+         * We could have failed because the partition is on a dynamic
+         * MBR or GPT data disk, so retry with the non-EX IOCTL.
+         */
+        Status2 = NtDeviceIoControlFile(PartitionHandle,
+                                        NULL,
+                                        NULL,
+                                        NULL,
+                                        IoStatusBlock,
+                                        IOCTL_DISK_GET_PARTITION_INFO,
+                                        NULL,
+                                        0,
+                                        &PartInfo.Info,
+                                        sizeof(PartInfo.Info));
+        if (!NT_SUCCESS(Status2))
+        {
+            /* We failed again, bail out */
+            DPRINT1("IOCTL_DISK_GET_PARTITION_INFO failed (Status %lx)\n", Status2);
+            goto Quit;
+        }
+
+        /* The partition is supposed to be on an MBR disk; retrieve its type */
+        PartitionType = PartInfo.Info.PartitionType;
+    }
+    else
+    {
+        /* We succeeded; retrieve the partition type only if it is on an MBR disk */
+        if (PartInfo.InfoEx.PartitionStyle != PARTITION_STYLE_MBR)
+        {
+            /* Disk is not MBR, bail out */
+            goto Quit;
+        }
+        PartitionType = PartInfo.InfoEx.Mbr.PartitionType;
     }
 
     /*
-     * Try to infer a preferred file system for this partition, given its ID.
+     * Given an MBR partition type, try to infer a preferred file system.
      *
-     * WARNING: This is partly a hack, since partitions with the same ID can
-     * be formatted with different file systems: for example, usual Linux
+     * WARNING: This is partly a hack, since partitions with the same type
+     * can be formatted with different file systems: for example, usual Linux
      * partitions that are formatted in EXT2/3/4, ReiserFS, etc... have the
-     * same partition ID 0x83.
+     * same partition type 0x83.
      *
      * The proper fix is to make a function that detects the existing FS
-     * from a given partition (not based on the partition ID).
-     * On the contrary, for unformatted partitions with a given ID, the
+     * from a given partition (not based on the partition type).
+     * On the contrary, for unformatted partitions with a given type, the
      * following code is OK.
      */
     if ((PartitionType == PARTITION_FAT_12) ||
@@ -170,93 +255,110 @@ InferFileSystemByHandle(
     }
 
 Quit:
-    if (*FileSystemName)
+    if (*FileSystemName && wcsicmp(FileSystemName, L"NTFS") == 0)
     {
         // WARNING: We cannot write on this FS yet!
-        if (PartitionType == PARTITION_IFS)
-        {
-            DPRINT1("Recognized file system '%S' that doesn't have write support yet!\n",
-                    FileSystemName);
-        }
+        DPRINT1("Recognized file system '%S' that doesn't have write support yet!\n",
+                FileSystemName);
     }
 
-    DPRINT1("InferFileSystem -- PartitionType: 0x%02X ; FileSystem (guessed): %S\n",
-            PartitionType, *FileSystemName ? FileSystemName : L"None");
-
     return Status;
 }
 
 NTSTATUS
 InferFileSystem(
-    IN PCWSTR Partition,
-    IN UCHAR PartitionType,
+    IN PCWSTR PartitionPath OPTIONAL,
+    IN HANDLE PartitionHandle OPTIONAL,
     IN OUT PWSTR FileSystemName,
     IN SIZE_T FileSystemNameSize)
 {
     NTSTATUS Status;
-    UNICODE_STRING PartitionPath;
+    UNICODE_STRING PartitionPathU;
     OBJECT_ATTRIBUTES ObjectAttributes;
-    HANDLE PartitionHandle;
     IO_STATUS_BLOCK IoStatusBlock;
 
-    /* Open the partition */
-    RtlInitUnicodeString(&PartitionPath, Partition);
-    InitializeObjectAttributes(&ObjectAttributes,
-                               &PartitionPath,
-                               OBJ_CASE_INSENSITIVE,
-                               NULL,
-                               NULL);
-    Status = NtOpenFile(&PartitionHandle,
-                        FILE_GENERIC_READ /* | SYNCHRONIZE */,
-                        &ObjectAttributes,
-                        &IoStatusBlock,
-                        FILE_SHARE_READ | FILE_SHARE_WRITE,
-                        0 /* FILE_SYNCHRONOUS_IO_NONALERT */);
-    if (!NT_SUCCESS(Status))
+    if (PartitionPath && PartitionHandle)
+        return STATUS_INVALID_PARAMETER;
+
+    /* Open the partition if a path has been given;
+     * otherwise just use the provided handle. */
+    if (PartitionPath)
     {
-        DPRINT1("Failed to open partition '%wZ', Status 0x%08lx\n", &PartitionPath, Status);
-        return Status;
+        RtlInitUnicodeString(&PartitionPathU, PartitionPath);
+        InitializeObjectAttributes(&ObjectAttributes,
+                                   &PartitionPathU,
+                                   OBJ_CASE_INSENSITIVE,
+                                   NULL,
+                                   NULL);
+        Status = NtOpenFile(&PartitionHandle,
+                            FILE_GENERIC_READ /* | SYNCHRONIZE */,
+                            &ObjectAttributes,
+                            &IoStatusBlock,
+                            FILE_SHARE_READ | FILE_SHARE_WRITE,
+                            0 /* FILE_SYNCHRONOUS_IO_NONALERT */);
+        if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("Failed to open partition '%S', Status 0x%08lx\n",
+                    PartitionPath, Status);
+            return Status;
+        }
     }
 
     /* Retrieve the FS */
-    Status = InferFileSystemByHandle(PartitionHandle,
-                                     PartitionType,
-                                     FileSystemName,
-                                     FileSystemNameSize);
+    Status = InferFileSystemWorker(PartitionHandle,
+                                   &IoStatusBlock,
+                                   FileSystemName,
+                                   FileSystemNameSize);
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("InferFileSystem() failed for partition '%S' (0x%p), Status 0x%08lx\n",
+                PartitionPath, PartitionHandle, Status);
+    }
+    else
+    {
+        DPRINT1("InferFileSystem(): FileSystem (guessed): %S\n",
+                *FileSystemName ? FileSystemName : L"None");
+    }
 
-    /* Close the partition */
-    NtClose(PartitionHandle);
+    if (PartitionPath)
+    {
+        /* Close the partition */
+        NtClose(PartitionHandle);
+    }
 
     return Status;
 }
 
 UCHAR
-FileSystemToPartitionType(
+FileSystemToMBRPartitionType(
     IN PCWSTR FileSystem,
-    IN PULARGE_INTEGER StartSector,
-    IN PULARGE_INTEGER SectorCount)
+    IN ULONGLONG StartSector,
+    IN ULONGLONG SectorCount)
 {
-    ASSERT(FileSystem && StartSector && SectorCount);
+    ASSERT(FileSystem);
+
+    if (SectorCount == 0)
+        return PARTITION_ENTRY_UNUSED;
 
     if (wcsicmp(FileSystem, L"FAT")   == 0 ||
         wcsicmp(FileSystem, L"FAT32") == 0 ||
         wcsicmp(FileSystem, L"RAW")   == 0)
     {
-        if (SectorCount->QuadPart < 8192)
+        if (SectorCount < 8192ULL)
         {
             /* FAT12 CHS partition (disk is smaller than 4.1MB) */
             return PARTITION_FAT_12;
         }
-        else if (StartSector->QuadPart < 1450560)
+        else if (StartSector < 1450560ULL)
         {
             /* Partition starts below the 8.4GB boundary ==> CHS partition */
 
-            if (SectorCount->QuadPart < 65536)
+            if (SectorCount < 65536ULL)
             {
                 /* FAT16 CHS partition (partition size < 32MB) */
                 return PARTITION_FAT_16;
             }
-            else if (SectorCount->QuadPart < 1048576)
+            else if (SectorCount < 1048576ULL)
             {
                 /* FAT16 CHS partition (partition size < 512MB) */
                 return PARTITION_HUGE;
@@ -271,7 +373,7 @@ FileSystemToPartitionType(
         {
             /* Partition starts above the 8.4GB boundary ==> LBA partition */
 
-            if (SectorCount->QuadPart < 1048576)
+            if (SectorCount < 1048576ULL)
             {
                 /* FAT16 LBA partition (partition size < 512MB) */
                 return PARTITION_XINT13;

base/setup/lib/utils/fsrec.h

@@ -8,42 +8,31 @@
 
 #pragma once
 
-NTSTATUS
-GetFileSystemNameByHandle(
-    IN HANDLE PartitionHandle,
-    IN OUT PWSTR FileSystemName,
-    IN SIZE_T FileSystemNameSize);
-
 NTSTATUS
 GetFileSystemName_UStr(
-    IN PUNICODE_STRING PartitionPath,
+    IN PUNICODE_STRING PartitionPath OPTIONAL,
+    IN HANDLE PartitionHandle OPTIONAL,
     IN OUT PWSTR FileSystemName,
     IN SIZE_T FileSystemNameSize);
 
 NTSTATUS
 GetFileSystemName(
-    IN PCWSTR Partition,
-    IN OUT PWSTR FileSystemName,
-    IN SIZE_T FileSystemNameSize);
-
-NTSTATUS
-InferFileSystemByHandle(
-    IN HANDLE PartitionHandle,
-    IN UCHAR PartitionType,
+    IN PCWSTR PartitionPath OPTIONAL,
+    IN HANDLE PartitionHandle OPTIONAL,
     IN OUT PWSTR FileSystemName,
     IN SIZE_T FileSystemNameSize);
 
 NTSTATUS
 InferFileSystem(
-    IN PCWSTR Partition,
-    IN UCHAR PartitionType,
+    IN PCWSTR PartitionPath OPTIONAL,
+    IN HANDLE PartitionHandle OPTIONAL,
     IN OUT PWSTR FileSystemName,
     IN SIZE_T FileSystemNameSize);
 
 UCHAR
-FileSystemToPartitionType(
+FileSystemToMBRPartitionType(
     IN PCWSTR FileSystem,
-    IN PULARGE_INTEGER StartSector,
-    IN PULARGE_INTEGER SectorCount);
+    IN ULONGLONG StartSector,
+    IN ULONGLONG SectorCount);
 
 /* EOF */

base/setup/lib/utils/partlist.c

@@ -906,7 +906,10 @@ InitializePartitionEntry(
     PartEntry->New = TRUE;
     PartEntry->IsPartitioned = TRUE;
 
-    PartEntry->PartitionType = FileSystemToPartitionType(L"RAW", &PartEntry->StartSector, &PartEntry->SectorCount);
+// FIXME: Use FileSystemToMBRPartitionType() only for MBR, otherwise use PARTITION_BASIC_DATA_GUID.
+    PartEntry->PartitionType = FileSystemToMBRPartitionType(L"RAW",
+                                                            PartEntry->StartSector.QuadPart,
+                                                            PartEntry->SectorCount.QuadPart);
     ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
 
     PartEntry->FormatState = Unformatted;
@@ -1018,12 +1021,11 @@ AddPartitionToDisk(
             ASSERT(NT_SUCCESS(Status));
 
             /* We don't have a FS, try to guess one */
-            Status = InferFileSystemByHandle(PartitionHandle,
-                                             PartEntry->PartitionType,
-                                             PartEntry->FileSystem,
-                                             sizeof(PartEntry->FileSystem));
+            Status = InferFileSystem(NULL, PartitionHandle,
+                                     PartEntry->FileSystem,
+                                     sizeof(PartEntry->FileSystem));
             if (!NT_SUCCESS(Status))
-                DPRINT1("InferFileSystemByHandle() failed, Status 0x%08lx\n", Status);
+                DPRINT1("InferFileSystem() failed, Status 0x%08lx\n", Status);
         }
         if (*PartEntry->FileSystem)
         {
@@ -4105,12 +4107,14 @@ SetMountedDeviceValues(
 }
 
 VOID
-SetPartitionType(
+SetMBRPartitionType(
     IN PPARTENTRY PartEntry,
     IN UCHAR PartitionType)
 {
     PDISKENTRY DiskEntry = PartEntry->DiskEntry;
 
+    ASSERT(DiskEntry->DiskStyle == PARTITION_STYLE_MBR);
+
     PartEntry->PartitionType = PartitionType;
 
     DiskEntry->Dirty = TRUE;

base/setup/lib/utils/partlist.h

@@ -349,7 +349,7 @@ SetMountedDeviceValues(
     IN PPARTLIST List);
 
 VOID
-SetPartitionType(
+SetMBRPartitionType(
     IN PPARTENTRY PartEntry,
     IN UCHAR PartitionType);