lywsvip/reactos
1
0
Fork 0
mirror of https://github.com/reactos/reactos.git synced 2026-06-03 01:41:13 +08:00
Code Issues Packages Projects Releases Wiki Activity

** wip ** [SETUPLIB] Add support for volume drive letters assignment

Browse Source
This commit is contained in:
Hermès Bélusca-Maïto
2024-06-05 22:44:02 +02:00
parent b390614b7e
commit fa97f68041
5 changed files with 714 additions and 144 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
base/setup/lib/setuplib.c
View File

@@ -1372,7 +1372,8 @@ DoUpdate:
/* Update the mounted devices list */
// FIXME: This should technically be done by mountmgr (if AutoMount is enabled)!
SetMountedDeviceValues(PartitionList);
// SetMountedDeviceValues(PartitionList);
ExportMountedDevices();
}
#ifdef __REACTOS__

591
base/setup/lib/utils/partlist.c
View File

@@ -15,8 +15,6 @@
#include "fsrec.h" // For FileSystemToMBRPartitionType()
#include "devutils.h"
#include "registry.h"
#define NDEBUG
#include <debug.h>
@@ -29,15 +27,6 @@
(PartEntry->SectorCount.QuadPart != 0LL))
#include <pshpack1.h>
typedef struct _REG_DISK_MOUNT_INFO
{
ULONG Signature;
ULONGLONG StartingOffset;
} REG_DISK_MOUNT_INFO, *PREG_DISK_MOUNT_INFO;
#include <poppack.h>
/* FUNCTIONS *****************************************************************/
VOID
@@ -92,7 +81,7 @@ DumpMountedDevices(VOID)
return;
}
Buffer = RtlAllocateHeap(RtlGetProcessHeap(), 0, BufferSize);
Buffer = RtlAllocateHeap(ProcessHeap, 0, BufferSize);
if (!Buffer)
{
DPRINT1("RtlAllocateHeap() failed\n");
@@ -111,9 +100,9 @@ DumpMountedDevices(VOID)
&RequiredSize);
if (Status == STATUS_BUFFER_OVERFLOW || Status == STATUS_BUFFER_TOO_SMALL)
{
RtlFreeHeap(RtlGetProcessHeap(), 0, Buffer);
RtlFreeHeap(ProcessHeap, 0, Buffer);
BufferSize = RequiredSize;
Buffer = RtlAllocateHeap(RtlGetProcessHeap(), 0, BufferSize);
Buffer = RtlAllocateHeap(ProcessHeap, 0, BufferSize);
if (!Buffer)
{
DPRINT1("RtlAllocateHeap() failed\n");
@@ -155,7 +144,7 @@ DumpMountedDevices(VOID)
}
DbgPrint("**** End Dumping ****\n");
RtlFreeHeap(RtlGetProcessHeap(), 0, Buffer);
RtlFreeHeap(ProcessHeap, 0, Buffer);
NtClose(hKey);
}
@@ -261,20 +250,103 @@ GetDriverName(
}
}
#define SET_MAP_LETTER(map, letter) \
do { \
ULONG _i = towupper(letter) - L'A'; \
if (0 <= _i && _i <= 'Z'-'A') \
(map) |= (1 << _i); \
} while (0)
#define REMOVE_MAP_LETTER(map, letter) \
do { \
ULONG _i = towupper(letter) - L'A'; \
if (0 <= _i && _i <= 'Z'-'A') \
(map) &= ~(1 << _i); \
} while (0)
static
WCHAR
AssignNextDriveLetter(
_In_ PPARTLIST List,
_Inout_ PVOLINFO Volume,
_In_ BOOLEAN TempAssign)
{
WCHAR Letter = UNICODE_NULL;
if (TempAssign)
{
ULONG i;
/* If the volume already has a drive letter, just return it */
if (Volume->DriveLetter)
return Volume->DriveLetter;
/* Scan the drive letters map and find the first free letter.
* Start with the C drive, except on NEC PC-98. */
for (i = (IsNEC_98 ? 0 : 2); i <= 'Z'-'A'; ++i)
{
if (!(List->DriveMap & (1 << i)))
{
/* Return the letter */
Letter = L'A' + i;
break;
}
}
}
else
{
NTSTATUS Status;
UNICODE_STRING Name;
/* Remove the temporary drive letter from the map */
REMOVE_MAP_LETTER(List->DriveMap, Volume->DriveLetter);
/* Re-assign the drive letter */
RtlInitUnicodeString(&Name, Volume->DeviceName);
Status = GetOrAssignNextVolumeDriveLetter(&Name, &Letter);
DBG_UNREFERENCED_PARAMETER(Status);
}
/* Reserve the letter and set it */
SET_MAP_LETTER(List->DriveMap, Letter);
Volume->DriveLetter = Letter;
return Letter;
}
/**
* @brief
* Assign drive letters to created volumes.
*
* @param[in] List
* Disks/Partitions/Volumes list.
*
* @param[in] TempAssign
* Whether to do a temporary assignment (TRUE: letters are maintained by us
* and are not reported to the system), or to do a permanent assignment
* (FALSE: letters are queried from the MountMgr and assigned to the volumes).
*
* @note
* For the moment, we do it ourselves, by assigning drives to partitions
* that are *only on MBR disks*. We first assign letters to each active
* partition on each disk, then assign letters to each logical partition,
* and finish by assigning letters to the remaining primary partitions.
* (This algorithm is the one that can be observed in the Windows Setup.)
**/
static
VOID
AssignDriveLetters(
IN PPARTLIST List)
_In_ PPARTLIST List,
_In_ BOOLEAN TempAssign)
{
PDISKENTRY DiskEntry;
PPARTENTRY PartEntry;
PLIST_ENTRY Entry1;
PLIST_ENTRY Entry2;
WCHAR Letter;
PLIST_ENTRY Entry1, Entry2;
Letter = L'C';
__debugbreak();
/* Assign drive letters to primary partitions */
/* Assign drive letters to volumes on primary partitions */
for (Entry1 = List->DiskListHead.Flink;
Entry1 != &List->DiskListHead;
Entry1 = Entry1->Flink)
@@ -286,23 +358,20 @@ AssignDriveLetters(
Entry2 = Entry2->Flink)
{
PartEntry = CONTAINING_RECORD(Entry2, PARTENTRY, ListEntry);
if (!PartEntry->Volume)
continue;
PartEntry->Volume->Info.DriveLetter = UNICODE_NULL;
if (PartEntry->IsPartitioned &&
!IsContainerPartition(PartEntry->PartitionType) &&
(IsRecognizedPartition(PartEntry->PartitionType) ||
PartEntry->SectorCount.QuadPart != 0LL))
ASSERT_IS_VOLUME_PARTITION_VALID(PartEntry);
if (!IsRecognizedPartition(PartEntry->PartitionType))
{
if (Letter <= L'Z')
PartEntry->Volume->Info.DriveLetter = Letter++;
ASSERT(PartEntry->Volume->Info.DriveLetter == UNICODE_NULL);
continue;
}
AssignNextDriveLetter(List, &PartEntry->Volume->Info, TempAssign);
}
}
/* Assign drive letters to logical drives */
/* Assign drive letters to volumes on logical drives */
for (Entry1 = List->DiskListHead.Flink;
Entry1 != &List->DiskListHead;
Entry1 = Entry1->Flink)
@@ -314,22 +383,66 @@ AssignDriveLetters(
Entry2 = Entry2->Flink)
{
PartEntry = CONTAINING_RECORD(Entry2, PARTENTRY, ListEntry);
if (!PartEntry->Volume)
continue;
PartEntry->Volume->Info.DriveLetter = UNICODE_NULL;
if (PartEntry->IsPartitioned &&
(IsRecognizedPartition(PartEntry->PartitionType) ||
PartEntry->SectorCount.QuadPart != 0LL))
ASSERT_IS_VOLUME_PARTITION_VALID(PartEntry);
if (!IsRecognizedPartition(PartEntry->PartitionType))
{
if (Letter <= L'Z')
PartEntry->Volume->Info.DriveLetter = Letter++;
ASSERT(PartEntry->Volume->Info.DriveLetter == UNICODE_NULL);
continue;
}
AssignNextDriveLetter(List, &PartEntry->Volume->Info, TempAssign);
}
}
}
static
VOID
InitDriveLettersMap(
_Out_ PULONG DriveMap)
{
PROCESS_DEVICEMAP_INFORMATION DeviceMap;
NTSTATUS Status;
ULONG i;
/* Get the system drive letters map (seen by this process) */
Status = NtQueryInformationProcess(NtCurrentProcess(),
ProcessDeviceMap,
&DeviceMap.Query,
sizeof(DeviceMap.Query),
NULL);
/* Zero the map if we failed */
if (!NT_SUCCESS(Status))
RtlZeroMemory(&DeviceMap, sizeof(DeviceMap));
/* Keep only the letters used by other resources (remote drives, CD-ROM,
* RamDisks, ...) that either we or the MountMgr do not maintain */
for (i = 0; i <= 'Z'-'A'; ++i)
{
/* Skip the drive if it's not in the map */
if (!(DeviceMap.Query.DriveMap & (1 << i)))
continue;
/* Disable unknown drives */
if (!((DeviceMap.Query.DriveType[i] >= DOSDEVICE_DRIVE_REMOVABLE) &&
(DeviceMap.Query.DriveType[i] <= DOSDEVICE_DRIVE_RAMDISK)))
{
DeviceMap.Query.DriveMap &= ~(1 << i);
}
#if 0
/* Disable letters on fixed drives, these should be handled by MountMgr */
if (DeviceMap.Query.DriveType[i] == DOSDEVICE_DRIVE_FIXED)
{
DeviceMap.Query.DriveMap &= ~(1 << i);
}
#endif
}
*DriveMap = DeviceMap.Query.DriveMap;
}
static NTSTATUS
NTAPI
DiskIdentifierQueryRoutine(
@@ -2233,6 +2346,9 @@ __debugbreak();
else // !IsUnknown && !IsUnformatted == IsFormatted
Volume->FormatState = Formatted;
/* Set the drive letter in the map */
SET_MAP_LETTER(List->DriveMap, Volume->Info.DriveLetter);
return STATUS_SUCCESS;
}
@@ -2361,6 +2477,9 @@ CreatePartitionList(VOID)
InitializeListHead(&List->VolumesList);
InitializeListHead(&List->PendingUnmountVolumesList);
/* Get the system drive letters map (seen by this process) */
InitDriveLettersMap(&List->DriveMap);
/*
* Enumerate the disks seen by the BIOS; this will be used later
* to map drives seen by NTOS with their corresponding BIOS names.
@@ -2387,7 +2506,6 @@ CreatePartitionList(VOID)
UpdateDiskSignatures(List);
UpdateHwDiskNumbers(List);
AssignDriveLetters(List);
/*
* Retrieve the system partition: the active partition on the system
@@ -3296,10 +3414,9 @@ CreatePartition(
* associated with this partition has to be created. */
PartEntry->Volume = InitVolume(DiskEntry->PartList, PartEntry);
ASSERT(PartEntry->Volume);
AssignNextDriveLetter(List, &PartEntry->Volume->Info, TRUE);
}
AssignDriveLetters(List);
return TRUE;
}
@@ -3319,6 +3436,9 @@ DismountPartition(
ASSERT(Volume->PartEntry == PartEntry);
ASSERT_IS_VOLUME_PARTITION_VALID(PartEntry);
/* Remove the drive letter from the map */
REMOVE_MAP_LETTER(List->DriveMap, Volume->Info.DriveLetter);
/* Dismount the basic volume: unlink the volume from the volumes list */
PartEntry->Volume = NULL;
Volume->PartEntry = NULL;
@@ -3478,7 +3598,7 @@ DeletePartition(
}
UpdateDiskLayout(DiskEntry);
AssignDriveLetters(List);
AssignDriveLetters(List, TRUE);
return TRUE;
}
@@ -4165,110 +4285,15 @@ WritePartitionsToDisk(
InitVolumeDeviceName(Volume, NULL);
}
/* Assign persistent drive letters to volumes */
AssignDriveLetters(List, FALSE);
//** Re-dump the list of MountedDevices **//
DumpMountedDevices();
return TRUE;
}
/**
* @brief
* Assign a "\DosDevices\#:" mount point drive letter to a disk partition or
* volume, specified by a given disk signature and starting partition offset.
**/
static BOOLEAN
SetMountedDeviceValue(
_In_ PVOLENTRY Volume)
{
PPARTENTRY PartEntry = Volume->PartEntry;
WCHAR Letter = Volume->Info.DriveLetter;
NTSTATUS Status;
OBJECT_ATTRIBUTES ObjectAttributes;
UNICODE_STRING KeyName = RTL_CONSTANT_STRING(L"SYSTEM\\MountedDevices");
UNICODE_STRING ValueName;
WCHAR Buffer[16];
HANDLE KeyHandle;
REG_DISK_MOUNT_INFO MountInfo;
/* Ignore no letter */
if (!Letter)
return TRUE;
RtlStringCchPrintfW(Buffer, _countof(Buffer),
L"\\DosDevices\\%c:", Letter);
RtlInitUnicodeString(&ValueName, Buffer);
InitializeObjectAttributes(&ObjectAttributes,
&KeyName,
OBJ_CASE_INSENSITIVE,
GetRootKeyByPredefKey(HKEY_LOCAL_MACHINE, NULL),
NULL);
Status = NtOpenKey(&KeyHandle,
KEY_ALL_ACCESS,
&ObjectAttributes);
if (!NT_SUCCESS(Status))
{
Status = NtCreateKey(&KeyHandle,
KEY_ALL_ACCESS,
&ObjectAttributes,
0,
NULL,
REG_OPTION_NON_VOLATILE,
NULL);
}
if (!NT_SUCCESS(Status))
{
DPRINT1("NtCreateKey() failed (Status %lx)\n", Status);
return FALSE;
}
MountInfo.Signature = PartEntry->DiskEntry->LayoutBuffer->Signature;
MountInfo.StartingOffset = GetPartEntryOffsetInBytes(PartEntry);
Status = NtSetValueKey(KeyHandle,
&ValueName,
0,
REG_BINARY,
(PVOID)&MountInfo,
sizeof(MountInfo));
NtClose(KeyHandle);
if (!NT_SUCCESS(Status))
{
DPRINT1("NtSetValueKey() failed (Status %lx)\n", Status);
return FALSE;
}
return TRUE;
}
BOOLEAN
SetMountedDeviceValues(
_In_ PPARTLIST List)
{
PLIST_ENTRY Entry;
PVOLENTRY Volume;
if (!List)
return FALSE;
//** Last time dumping of the list of MountedDevices **//
DumpMountedDevices();
for (Entry = List->VolumesList.Flink;
Entry != &List->VolumesList;
Entry = Entry->Flink)
{
Volume = CONTAINING_RECORD(Entry, VOLENTRY, ListEntry);
/* Assign a "\DosDevices\#:" mount point to this volume */
if (!SetMountedDeviceValue(Volume))
return FALSE;
}
return TRUE;
}
VOID
SetMBRPartitionType(
IN PPARTENTRY PartEntry,
@@ -4292,4 +4317,292 @@ SetMBRPartitionType(
DiskEntry->LayoutBuffer->PartitionEntry[PartEntry->PartitionIndex].RewritePartition = TRUE;
}
/*****************************************************************************\
* SETUP-specific support functions
*/
#include "registry.h" // For GetRootKeyByPredefKey()
#include <pshpack1.h>
typedef struct _REG_DISK_MOUNT_INFO
{
ULONG Signature;
ULONGLONG StartingOffset;
} REG_DISK_MOUNT_INFO, *PREG_DISK_MOUNT_INFO;
#include <poppack.h>
/**
* @brief
* Assign a "\DosDevices\#:" mount point drive letter to a disk partition or
* volume, specified by a given disk signature and starting partition offset.
*
* @note
* The association is stored in the registry of the **TARGET**
* NT installation, not of the current running one.
*
* We use it to update the mounted devices list
* // FIXME: This should technically be done by mountmgr (if AutoMount is enabled)!
*
* TODO:
* - Make the function more generic for MBR and GPT.
*
* @note
* The stored data actually corresponds to a "unique ID" the partition manager
* gives to the mount manager. In this function below, the format is actually
* the one used for partitions on MBR disks only.
*
* @see
* https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/supporting-mount-manager-requests-in-a-storage-class-driver
* https://winreg-kb.readthedocs.io/en/latest/sources/system-keys/Mounted-devices.html
**/
static NTSTATUS
SetMountedDeviceValue(
_In_ PVOLENTRY Volume)
{
PPARTENTRY PartEntry = Volume->PartEntry;
WCHAR Letter = Volume->Info.DriveLetter;
NTSTATUS Status;
OBJECT_ATTRIBUTES ObjectAttributes;
UNICODE_STRING KeyName = RTL_CONSTANT_STRING(L"SYSTEM\\MountedDevices");
UNICODE_STRING ValueName;
WCHAR Buffer[16];
HANDLE KeyHandle;
REG_DISK_MOUNT_INFO MountInfo;
/* Ignore no letter */
if (!Letter)
return STATUS_SUCCESS;
RtlStringCchPrintfW(Buffer, _countof(Buffer),
L"\\DosDevices\\%c:", Letter);
RtlInitUnicodeString(&ValueName, Buffer);
InitializeObjectAttributes(&ObjectAttributes,
&KeyName,
OBJ_CASE_INSENSITIVE,
GetRootKeyByPredefKey(HKEY_LOCAL_MACHINE, NULL),
NULL);
Status = NtOpenKey(&KeyHandle,
KEY_ALL_ACCESS,
&ObjectAttributes);
if (!NT_SUCCESS(Status))
{
Status = NtCreateKey(&KeyHandle,
KEY_ALL_ACCESS,
&ObjectAttributes,
0,
NULL,
REG_OPTION_NON_VOLATILE,
NULL);
}
if (!NT_SUCCESS(Status))
{
DPRINT1("NtCreateKey() failed (Status %lx)\n", Status);
return Status;
}
//
// NOTE: Alternatively, just ask the MountMgr (if we have it)
// for the volume unique ID, via IOCTL_MOUNTDEV_QUERY_UNIQUE_ID !!
//
// _In_ ULONG Signature, // DiskSignature
MountInfo.Signature = PartEntry->DiskEntry->LayoutBuffer->Signature;
MountInfo.StartingOffset = GetPartEntryOffsetInBytes(PartEntry);
Status = NtSetValueKey(KeyHandle,
&ValueName,
0,
REG_BINARY,
(PVOID)&MountInfo,
sizeof(MountInfo));
NtClose(KeyHandle);
if (!NT_SUCCESS(Status))
{
DPRINT1("NtSetValueKey() failed (Status %lx)\n", Status);
return Status;
}
return STATUS_SUCCESS;
}
/*
* TODO:
* - It may actually make sense to just migrate the volume letters database
* from the current active system to the installation TARGET one.
*/
BOOLEAN
SetMountedDeviceValues(
_In_ PPARTLIST List)
{
PLIST_ENTRY Entry;
PVOLENTRY Volume;
NTSTATUS Status;
if (!List)
return FALSE;
//** Last time dumping of the list of MountedDevices **//
DumpMountedDevices();
for (Entry = List->VolumesList.Flink;
Entry != &List->VolumesList;
Entry = Entry->Flink)
{
Volume = CONTAINING_RECORD(Entry, VOLENTRY, ListEntry);
/* Assign a "\DosDevices\#:" mount point to this volume */
Status = SetMountedDeviceValue(Volume);
if (!NT_SUCCESS(Status))
return FALSE;
}
return TRUE;
}
/**
* @brief
* Exports the \Registry\Machine\SYSTEM\MountedDevices database of the
* current running ReactOS installation to the target ReactOS installation.
*
* Used when both installations are for the same machine.
**/
NTSTATUS
ExportMountedDevices(VOID)
{
NTSTATUS Status;
UNICODE_STRING SrcPath
= RTL_CONSTANT_STRING(L"\\Registry\\Machine\\SYSTEM\\MountedDevices");
UNICODE_STRING DstPath
= RTL_CONSTANT_STRING(L"SYSTEM\\MountedDevices");
OBJECT_ATTRIBUTES ObjectAttributes;
HANDLE SrcKey = NULL;
HANDLE DstKey = NULL;
PKEY_VALUE_FULL_INFORMATION Buffer;
ULONG BufferSize = sizeof(KEY_VALUE_FULL_INFORMATION) + MAX_PATH * sizeof(WCHAR);
ULONG RequiredSize;
ULONG i = 0;
UNICODE_STRING Name;
/* Open the source key */
InitializeObjectAttributes(&ObjectAttributes,
&SrcPath,
OBJ_CASE_INSENSITIVE,
NULL, NULL);
Status = NtOpenKey(&SrcKey, KEY_QUERY_VALUE, &ObjectAttributes);
if (!NT_SUCCESS(Status))
{
DPRINT1("NtOpenKey() failed with status 0x%08lx\n", Status);
return Status;
}
/* Open the target key */
InitializeObjectAttributes(&ObjectAttributes,
&DstPath,
OBJ_CASE_INSENSITIVE,
GetRootKeyByPredefKey(HKEY_LOCAL_MACHINE, NULL),
NULL);
Status = NtOpenKey(&DstKey, KEY_ALL_ACCESS, &ObjectAttributes);
if (!NT_SUCCESS(Status))
{
Status = NtCreateKey(&DstKey,
KEY_ALL_ACCESS,
&ObjectAttributes,
0,
NULL,
REG_OPTION_NON_VOLATILE,
NULL);
}
if (!NT_SUCCESS(Status))
{
DPRINT1("NtCreateKey() failed (Status %lx)\n", Status);
goto Quit;
}
/*
* Enumerate each value in the source key and copy it into the target key.
*/
Buffer = RtlAllocateHeap(ProcessHeap, 0, BufferSize);
if (!Buffer)
{
DPRINT1("RtlAllocateHeap() failed\n");
Status = STATUS_NO_MEMORY;
goto Quit;
}
DbgPrint("\n**** Importing HKLM\\SYSTEM\\MountedDevices ****\n");
while (TRUE)
{
Status = NtEnumerateValueKey(SrcKey,
i,
KeyValueFullInformation,
Buffer,
BufferSize,
&RequiredSize);
if (Status == STATUS_BUFFER_OVERFLOW || Status == STATUS_BUFFER_TOO_SMALL)
{
RtlFreeHeap(ProcessHeap, 0, Buffer);
BufferSize = RequiredSize;
Buffer = RtlAllocateHeap(ProcessHeap, 0, BufferSize);
if (!Buffer)
{
DPRINT1("RtlAllocateHeap() failed\n");
// Status = STATUS_NO_MEMORY;
break;
// continue;
}
Status = NtEnumerateValueKey(SrcKey,
i,
KeyValueFullInformation,
Buffer,
BufferSize,
&RequiredSize);
}
if (!NT_SUCCESS(Status))
{
DPRINT("NtEnumerateKey() failed with status 0x%08lx\n", Status);
break;
}
i++;
if (Buffer->Type != REG_BINARY)
{
DPRINT1("Wrong registry type: got 0x%lx, expected 0x%lx (REG_BINARY)\n",
Buffer->Type, REG_BINARY);
}
Name.Length = Name.MaximumLength = Buffer->NameLength;
Name.Buffer = Buffer->Name;
DbgPrint(" '%wZ' =>\n", &Name);
{
PVOID Data = (PVOID)((ULONG_PTR)Buffer + Buffer->DataOffset);
DebugDumpBuffer(Data, Buffer->DataLength);
Status = NtSetValueKey(DstKey,
&Name,
0,
Buffer->Type,
Data,
Buffer->DataLength);
}
DbgPrint("\n");
}
DbgPrint("**** End Importing ****\n");
RtlFreeHeap(ProcessHeap, 0, Buffer);
Quit:
if (DstKey)
NtClose(DstKey);
if (SrcKey)
NtClose(SrcKey);
return Status;
}
/*
* End of SETUP-specific support functions
*****************************************************************************/
/* EOF */

13
base/setup/lib/utils/partlist.h
View File

@@ -186,6 +186,7 @@ typedef struct _PARTLIST
/* (Basic) Volumes management */
LIST_ENTRY VolumesList; ///< List of active volumes
LIST_ENTRY PendingUnmountVolumesList; ///< List of volumes to unmount
ULONG DriveMap; ///< Drive letters map, used by AssignNextDriveLetter()
} PARTLIST, *PPARTLIST;
@@ -386,13 +387,17 @@ BOOLEAN
WritePartitionsToDisk(
IN PPARTLIST List);
BOOLEAN
SetMountedDeviceValues(
_In_ PPARTLIST List);
VOID
SetMBRPartitionType(
IN PPARTENTRY PartEntry,
IN UCHAR PartitionType);
BOOLEAN
SetMountedDeviceValues(
_In_ PPARTLIST List);
NTSTATUS
ExportMountedDevices(VOID);
/* EOF */

245
base/setup/lib/utils/volutil.c
View File

@@ -8,6 +8,7 @@
/* INCLUDES ******************************************************************/
#include "precomp.h"
#include <mountdev.h>
#include "volutil.h"
#include "fsrec.h"
@@ -19,12 +20,249 @@
/* FUNCTIONS *****************************************************************/
/**
* @brief
* Retrieves a handle to the MountMgr controlling device.
* The handle should be closed with NtClose() once it is no longer in use.
**/
NTSTATUS
GetMountMgrHandle(
_Out_ PHANDLE MountMgrHandle,
_In_ ACCESS_MASK DesiredAccess)
{
NTSTATUS Status;
UNICODE_STRING MountMgrDevice;
OBJECT_ATTRIBUTES ObjectAttributes;
IO_STATUS_BLOCK IoStatusBlock;
*MountMgrHandle = NULL;
RtlInitUnicodeString(&MountMgrDevice, MOUNTMGR_DEVICE_NAME);
InitializeObjectAttributes(&ObjectAttributes,
&MountMgrDevice,
OBJ_CASE_INSENSITIVE,
NULL,
NULL);
Status = NtOpenFile(MountMgrHandle,
DesiredAccess | SYNCHRONIZE,
&ObjectAttributes,
&IoStatusBlock,
FILE_SHARE_READ | FILE_SHARE_WRITE,
FILE_SYNCHRONOUS_IO_NONALERT);
if (!NT_SUCCESS(Status))
{
DPRINT1("NtOpenFile(%wZ) failed, Status 0x%08lx\n",
&MountMgrDevice, Status);
}
return Status;
}
/**
* @brief
* Requests the MountMgr to retrieve the drive letter,
* if any, associated to the given volume.
*
* @param[in] VolumeName
* Device name of the volume.
*
* @param[out] DriveLetter
* Pointer to a WCHAR buffer receiving the corresponding drive letter,
* or UNICODE_NULL if none has been assigned.
*
* @return A status code.
**/
NTSTATUS
GetVolumeDriveLetter(
_In_ PCUNICODE_STRING VolumeName,
_Out_ PWCHAR DriveLetter)
{
/* Differs from MOUNTMGR_IS_DRIVE_LETTER(): no '\DosDevices\' accounted for */
#define IS_DRIVE_LETTER(s) \
((s)->Length == 2*sizeof(WCHAR) && (s)->Buffer[0] >= 'A' && \
(s)->Buffer[0] <= 'Z' && (s)->Buffer[1] == ':')
NTSTATUS Status;
IO_STATUS_BLOCK IoStatusBlock;
HANDLE MountMgrHandle;
ULONG Length;
MOUNTMGR_VOLUME_PATHS VolumePath;
PMOUNTMGR_VOLUME_PATHS VolumePathPtr;
UNICODE_STRING DosPath;
ULONG DeviceNameLength;
/*
* This variable is used to store the device name.
* for the input buffer to IOCTL_MOUNTMGR_QUERY_DOS_VOLUME_PATH.
* It's based on MOUNTMGR_TARGET_NAME (mountmgr.h).
* Doing it this way prevents memory allocation.
* The device name won't be longer.
*/
struct
{
USHORT NameLength;
WCHAR DeviceName[256];
} DeviceName;
/* Default to no letter */
*DriveLetter = UNICODE_NULL;
/* First, build the corresponding device name */
DeviceName.NameLength = VolumeName->Length;
RtlCopyMemory(&DeviceName.DeviceName, VolumeName->Buffer, VolumeName->Length);
DeviceNameLength = FIELD_OFFSET(MOUNTMGR_TARGET_NAME, DeviceName) + DeviceName.NameLength;
/* Now, query the MountMgr for the DOS path */
Status = GetMountMgrHandle(&MountMgrHandle, FILE_READ_ATTRIBUTES);
if (!NT_SUCCESS(Status))
{
DPRINT1("MountMgr unavailable: Status 0x%08lx\n", Status);
return Status;
}
VolumePathPtr = NULL;
Status = NtDeviceIoControlFile(MountMgrHandle,
NULL, NULL, NULL,
&IoStatusBlock,
IOCTL_MOUNTMGR_QUERY_DOS_VOLUME_PATH,
&DeviceName, DeviceNameLength,
&VolumePath, sizeof(VolumePath));
/* The only tolerated failure here is buffer too small, which is expected */
if (!NT_SUCCESS(Status) && (Status != STATUS_BUFFER_OVERFLOW))
{
goto Quit;
}
/* Compute the needed size to store the DOS path */
Length = FIELD_OFFSET(MOUNTMGR_VOLUME_PATHS, MultiSz) + VolumePath.MultiSzLength;
if (Length > MAXUSHORT)
{
Status = STATUS_INVALID_BUFFER_SIZE;
goto Quit;
}
/* Reallocate the buffer */
VolumePathPtr = RtlAllocateHeap(ProcessHeap, HEAP_ZERO_MEMORY, Length);
if (!VolumePathPtr)
{
Status = STATUS_INSUFFICIENT_RESOURCES;
goto Quit;
}
/* Re-query the DOS path with the proper size */
Status = NtDeviceIoControlFile(MountMgrHandle,
NULL, NULL, NULL,
&IoStatusBlock,
IOCTL_MOUNTMGR_QUERY_DOS_VOLUME_PATH,
&DeviceName, DeviceNameLength,
VolumePathPtr, Length);
if (!NT_SUCCESS(Status))
goto Quit;
/* Retrieve the drive letter */
DosPath.Length = DosPath.MaximumLength = (USHORT)VolumePathPtr->MultiSzLength - 2 * sizeof(UNICODE_NULL);
DosPath.Buffer = VolumePathPtr->MultiSz;
if (IS_DRIVE_LETTER(&DosPath))
{
/* Return the drive letter, ensuring it's uppercased */
*DriveLetter = towupper(DosPath.Buffer[0]);
}
/* We are done, return success */
Status = STATUS_SUCCESS;
Quit:
if (VolumePathPtr)
RtlFreeHeap(ProcessHeap, 0, VolumePathPtr);
NtClose(MountMgrHandle);
return Status;
}
/**
* @brief
* Requests the MountMgr to either assign the next available drive letter
* to the given volume, if none already exists, or to retrieve its existing
* associated drive letter.
*
* @param[in] VolumeName
* Device name of the volume.
*
* @param[out] DriveLetter
* Pointer to a WCHAR buffer receiving the corresponding drive letter,
* or UNICODE_NULL if none has been assigned.
*
* @return A status code.
**/
NTSTATUS
GetOrAssignNextVolumeDriveLetter(
_In_ PCUNICODE_STRING VolumeName,
_Out_ PWCHAR DriveLetter)
{
NTSTATUS Status;
IO_STATUS_BLOCK IoStatusBlock;
HANDLE MountMgrHandle;
MOUNTMGR_DRIVE_LETTER_INFORMATION LetterInfo;
/*
* This variable is used to store the device name
* for the input buffer to IOCTL_MOUNTMGR_NEXT_DRIVE_LETTER.
* It's based on MOUNTMGR_DRIVE_LETTER_TARGET (mountmgr.h).
* Doing it this way prevents memory allocation.
* The device name won't be longer.
*/
struct
{
USHORT DeviceNameLength;
WCHAR DeviceName[256];
} DeviceName;
/* Default to no letter */
*DriveLetter = UNICODE_NULL;
/* First, build the corresponding device name */
DeviceName.DeviceNameLength = VolumeName->Length;
RtlCopyMemory(&DeviceName.DeviceName, VolumeName->Buffer, VolumeName->Length);
/* Now, query the MountMgr for the drive letter */
Status = GetMountMgrHandle(&MountMgrHandle, FILE_READ_ACCESS | FILE_WRITE_ACCESS);
if (!NT_SUCCESS(Status))
{
DPRINT1("MountMgr unavailable: Status 0x%08lx\n", Status);
return Status;
}
Status = NtDeviceIoControlFile(MountMgrHandle,
NULL, NULL, NULL,
&IoStatusBlock,
IOCTL_MOUNTMGR_NEXT_DRIVE_LETTER,
&DeviceName, sizeof(DeviceName),
&LetterInfo, sizeof(LetterInfo));
if (!NT_SUCCESS(Status))
goto Quit;
DPRINT1("DriveLetterWasAssigned = %s, CurrentDriveLetter = %c\n",
LetterInfo.DriveLetterWasAssigned ? "TRUE" : "FALSE",
LetterInfo.CurrentDriveLetter ? LetterInfo.CurrentDriveLetter : '-');
/* Return the drive letter the MountMgr potentially assigned,
* ensuring it's uppercased */
*DriveLetter = towupper(LetterInfo.CurrentDriveLetter);
Quit:
NtClose(MountMgrHandle);
return Status;
}
/**
* @brief
* Attempts to mount the designated volume, and retrieve and cache
* some of its properties (file system, volume label, ...).
**/
NTSTATUS
MountVolume(
_Inout_ PVOLINFO Volume,
_In_opt_ UCHAR MbrPartitionType)
{
NTSTATUS Status;
UNICODE_STRING Name;
HANDLE VolumeHandle;
/* If the volume is already mounted, just return success */
@@ -128,6 +366,12 @@ MountVolume(
}
}
/* Get the volume drive letter */
__debugbreak();
RtlInitUnicodeString(&Name, Volume->DeviceName);
Status = GetVolumeDriveLetter(&Name, &Volume->DriveLetter);
UNREFERENCED_PARAMETER(Status);
/* Close the volume */
if (VolumeHandle)
NtClose(VolumeHandle);
@@ -215,6 +459,7 @@ DismountVolume(
/* Reset some data only if dismount succeeded */
if (NT_SUCCESS(Status))
{
// TODO: Should we notify MountMgr to delete the drive letter?
Volume->DriveLetter = UNICODE_NULL;
Volume->VolumeLabel[0] = UNICODE_NULL;
Volume->FileSystem[0] = UNICODE_NULL;

6
base/setup/lib/utils/volutil.h
View File

@@ -40,6 +40,12 @@ typedef struct _VOLINFO
(!IsUnknown(VolInfo) && !IsUnformatted(VolInfo))
NTSTATUS
GetOrAssignNextVolumeDriveLetter(
_In_ PCUNICODE_STRING VolumeName,
_Out_ PWCHAR DriveLetter);
// DetectFileSystem()
NTSTATUS
MountVolume(
_Inout_ PVOLINFO Volume,