[misc] fix small issues

* SetAutoMount()/GetAutoMount() should check for INVALID_HANDLE_VALUE and not NULL.
  Also we don't actually need to open MOUNTMGR_DOS_DEVICE_NAME rw to issue an IOCTL.
* ToggleEsp() failed to exit properly when an ESP offset was specified.
* Introduce PI_MAX to explicitly set the size of the partition_information table.
* write_sectors() has write retry, so there's no need to perform one on top of it.
* When we exit FormatThread(), GetLogicalName() should attempt to look for the the
  main partition and be silent.
This commit is contained in:
Pete Batard 2020-10-25 12:31:30 +00:00
parent a17ca005ce
commit 8085a2846d
No known key found for this signature in database
GPG key ID: 38E0CF5E69EDD671
4 changed files with 27 additions and 40 deletions

View file

@ -74,7 +74,7 @@ PF_TYPE_DECL(NTAPI, NTSTATUS, NtQueryVolumeInformationFile, (HANDLE, PIO_STATUS_
RUFUS_DRIVE_INFO SelectedDrive;
extern BOOL installed_uefi_ntfs, write_as_esp;
extern int nWindowsVersion, nWindowsBuildNumber;
uint64_t partition_offset[3];
uint64_t partition_offset[PI_MAX];
uint64_t persistence_size = 0;
/*
@ -93,9 +93,8 @@ BOOL SetAutoMount(BOOL enable)
DWORD size;
BOOL ret = FALSE;
hMountMgr = CreateFileA(MOUNTMGR_DOS_DEVICE_NAME, GENERIC_READ|GENERIC_WRITE,
FILE_SHARE_READ|FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
if (hMountMgr == NULL)
hMountMgr = CreateFileA(MOUNTMGR_DOS_DEVICE_NAME, 0, FILE_SHARE_READ|FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
if (hMountMgr == INVALID_HANDLE_VALUE)
return FALSE;
ret = DeviceIoControl(hMountMgr, IOCTL_MOUNTMGR_SET_AUTO_MOUNT, &enable, sizeof(enable), NULL, 0, &size, NULL);
CloseHandle(hMountMgr);
@ -110,9 +109,8 @@ BOOL GetAutoMount(BOOL* enabled)
if (enabled == NULL)
return FALSE;
hMountMgr = CreateFileA(MOUNTMGR_DOS_DEVICE_NAME, GENERIC_READ|GENERIC_WRITE,
FILE_SHARE_READ|FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
if (hMountMgr == NULL)
hMountMgr = CreateFileA(MOUNTMGR_DOS_DEVICE_NAME, 0, FILE_SHARE_READ|FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
if (hMountMgr == INVALID_HANDLE_VALUE)
return FALSE;
ret = DeviceIoControl(hMountMgr, IOCTL_MOUNTMGR_QUERY_AUTO_MOUNT, NULL, 0, enabled, sizeof(*enabled), &size, NULL);
CloseHandle(hMountMgr);
@ -1292,8 +1290,10 @@ BOOL ToggleEsp(DWORD DriveIndex, uint64_t PartitionOffset)
}
} else {
for (i = 0, j = 0; i < DriveLayout->PartitionCount; i++) {
if (DriveLayout->PartitionEntry[i].StartingOffset.QuadPart == PartitionOffset)
if (DriveLayout->PartitionEntry[i].StartingOffset.QuadPart == PartitionOffset) {
DriveLayout->PartitionEntry[i].Gpt.PartitionType = PARTITION_GENERIC_ESP;
break;
}
}
}
if (i >= DriveLayout->PartitionCount) {
@ -1853,6 +1853,8 @@ BOOL CreatePartition(HANDLE hDrive, int partition_style, int file_system, BOOL m
} else {
assert(FALSE);
}
// NB: Because we already subtracted the backup GPT size from the main partition size and
// this extra partition is indexed on main size, it does not overflow into the backup GPT.
main_part_size_in_sectors = ((main_part_size_in_sectors / SelectedDrive.SectorsPerTrack) -
extra_part_size_in_tracks) * SelectedDrive.SectorsPerTrack;
}