WIP: rewrite Windows code

.github/workflows/clang-format-check.yml

@@ -11,4 +11,4 @@ jobs:
  with:
  clang-format-version: '17'
  check-path: 'src'
- exclude-regex: 'src\/helper\/win|src\/app\/crashhandler.cpp'
+ exclude-regex: 'src\/app\/crashhandler.cpp'

src/app/CMakeLists.txt

@@ -90,7 +90,7 @@ if (UNIX AND NOT APPLE)
 endif()
 
 if (WIN32)
- target_link_libraries(mediawriter dbghelp)
+ target_link_libraries(mediawriter dbghelp libwmi)
 endif()
 
 if (APPLE)

src/app/main.cpp

@@ -1,5 +1,6 @@
 /*
  * Fedora Media Writer
+ * Copyright (C) 2024 Jan Grulich <jgrulichredhat.com>
  * Copyright (C) 2016 Martin Bříza <mbriza@redhat.com>
  *
  * This program is free software; you can redistribute it and/or

src/app/windrivemanager.cpp

@@ -1,6 +1,6 @@
 /*
  * Fedora Media Writer
- * Copyright (C) 2022 Jan Grulich <jgrulich@redhat.com>
+ * Copyright (C) 2022-2024 Jan Grulich <jgrulich@redhat.com>
  * Copyright (C) 2011-2022 Pete Batard <pete@akeo.ie>
  * Copyright (C) 2016 Martin Bříza <mbriza@redhat.com>
  *
@@ -22,262 +22,99 @@
 #include "windrivemanager.h"
 #include "notifications.h"
 
-#include <QDebug>
 #include <QStringList>
 #include <QTimer>
 
-#include <windows.h>
-#define INITGUID
-#include <guiddef.h>
+#include <dbt.h>
 
-#include <cmath>
-#include <cstring>
-
-const int maxPartitionCount = 16;
-
-DEFINE_GUID(PARTITION_MICROSOFT_DATA, 0xEBD0A0A2, 0xB9E5, 0x4433, 0x87, 0xC0, 0x68, 0xB6, 0xB7, 0x26, 0x99, 0xC7);
+#pragma comment(lib, "wbemuuid.lib")
 
 WinDriveProvider::WinDriveProvider(DriveManager *parent)
  : DriveProvider(parent)
+ , m_wmi(std::make_unique<LibWMI>(this))
 {
  mDebug() << this->metaObject()->className() << "construction";
+ qApp->installNativeEventFilter(this);
  QTimer::singleShot(0, this, &WinDriveProvider::checkDrives);
 }
 
-void WinDriveProvider::checkDrives()
+bool WinDriveProvider::nativeEventFilter(const QByteArray &eventType, void *message, qintptr *result)
 {
- static bool firstRun = true;
-
- if (firstRun)
- mDebug() << this->metaObject()->className() << "Looking for the drives for the first time";
-
- for (int i = 0; i < 64; i++) {
- bool present = describeDrive(i, firstRun);
- if (!present && m_drives.contains(i)) {
- emit driveRemoved(m_drives[i]);
- m_drives[i]->deleteLater();
- m_drives.remove(i);
+ Q_UNUSED(eventType);
+
+ MSG *msg = static_cast<MSG *>(message);
+ if ((msg->message == WM_DEVICECHANGE) && ((msg->wParam == DBT_DEVICEARRIVAL) || (msg->wParam == DBT_DEVICEREMOVECOMPLETE))) {
+ mDebug() << "Recieved device change event";
+ *result = TRUE;
+ for (const WinDrive *drive : m_drives) {
+ // Ignore device change events when we are restoring or writting and schedule
+ // re-check once we are done
+ if (drive->busy()) {
+ return true;
+ }
  }
+ QTimer::singleShot(0, this, &WinDriveProvider::checkDrives);
+ return true;
  }
-
- if (firstRun)
- mDebug() << this->metaObject()->className() << "Finished looking for the drives for the first time";
- firstRun = false;
- QTimer::singleShot(2500, this, &WinDriveProvider::checkDrives);
-}
-
-QString getPhysicalName(int driveNumber)
-{
- return QString("\\\\.\\PhysicalDrive%0").arg(driveNumber);
-}
-
-HANDLE getPhysicalHandle(int driveNumber)
-{
- HANDLE physicalHandle = INVALID_HANDLE_VALUE;
- QString physicalPath = getPhysicalName(driveNumber);
- physicalHandle = CreateFileA(physicalPath.toStdString().c_str(), GENERIC_READ, FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
- return physicalHandle;
+ return false;
 }
 
-bool WinDriveProvider::isMountable(int driveNumber)
+void WinDriveProvider::checkDrives()
 {
- mDebug() << this->metaObject()->className() << "Checking whether " << getPhysicalName(driveNumber) << " is mountable";
-
- HANDLE physicalHandle = getPhysicalHandle(driveNumber);
- if (physicalHandle == INVALID_HANDLE_VALUE) {
- mDebug() << this->metaObject()->className() << "Could not get physical handle for drive " << getPhysicalName(driveNumber);
- return false;
- }
-
- DWORD size;
- BYTE geometry[256];
- bool ret = DeviceIoControl(physicalHandle, IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, NULL, 0, geometry, sizeof(geometry), &size, NULL);
- if (!ret || size <= 0) {
- mDebug() << this->metaObject()->className() << "Could not get geometry for drive " << getPhysicalName(driveNumber);
- CloseHandle(physicalHandle);
- return false;
- }
-
- PDISK_GEOMETRY_EX diskGeometry = (PDISK_GEOMETRY_EX)(void *)geometry;
- // Drive info
- LONGLONG diskSize;
- DWORD sectorSize;
- DWORD sectorsPerTrack;
- DWORD firstDataSector;
- MEDIA_TYPE mediaType;
-
- diskSize = diskGeometry->DiskSize.QuadPart;
- sectorSize = diskGeometry->Geometry.BytesPerSector;
- firstDataSector = MAXDWORD;
- if (sectorSize < 512) {
- mDebug() << this->metaObject()->className() << "Warning: Drive " << getPhysicalName(driveNumber) << " reports a sector size of " << sectorSize << " - Correcting to 512 bytes.";
- sectorSize = 512;
- }
- sectorsPerTrack = diskGeometry->Geometry.SectorsPerTrack;
- mediaType = diskGeometry->Geometry.MediaType;
-
- BYTE layout[4096] = {0};
- ret = DeviceIoControl(physicalHandle, IOCTL_DISK_GET_DRIVE_LAYOUT_EX, NULL, 0, layout, sizeof(layout), &size, NULL);
- if (!ret || size <= 0) {
- mDebug() << this->metaObject()->className() << "Could not get layout for drive " << getPhysicalName(driveNumber);
- CloseHandle(physicalHandle);
- return false;
- }
-
- PDRIVE_LAYOUT_INFORMATION_EX driveLayout = (PDRIVE_LAYOUT_INFORMATION_EX)(void *)layout;
-
- switch (driveLayout->PartitionStyle) {
- case PARTITION_STYLE_MBR:
- mDebug() << this->metaObject()->className() << "MBR partition style";
- for (int i = 0; i < driveLayout->PartitionCount; i++) {
- if (driveLayout->PartitionEntry[i].Mbr.PartitionType != PARTITION_ENTRY_UNUSED) {
- QVector<uint8_t> mbrMountable = {0x01, 0x04, 0x06, 0x07, 0x0b, 0x0c, 0x0e};
- BYTE partType = driveLayout->PartitionEntry[i].Mbr.PartitionType;
- mDebug() << this->metaObject()->className() << "Partition type: " << partType;
- if (!mbrMountable.contains(partType)) {
- CloseHandle(physicalHandle);
- mDebug() << this->metaObject()->className() << getPhysicalName(driveNumber) << " is not mountable";
- return false;
- }
- }
- }
- break;
- case PARTITION_STYLE_GPT:
- mDebug() << this->metaObject()->className() << "GPT partition style";
- for (int i = 0; i < driveLayout->PartitionCount; i++) {
- if (memcmp(&driveLayout->PartitionEntry[i].Gpt.PartitionType, &PARTITION_MICROSOFT_DATA, sizeof(GUID)) != 0) {
- CloseHandle(physicalHandle);
- mDebug() << this->metaObject()->className() << getPhysicalName(driveNumber) << " is not mountable";
- return false;
+ mDebug() << this->metaObject()->className() << "Looking for the drives";
+
+ QMap<int, WinDrive *> drives;
+ auto usbDeviceList = m_wmi->getUSBDeviceList();
+ for (auto it = usbDeviceList.cbegin(); it != usbDeviceList.cend(); it++) {
+ bool noLogicalDevice = false;
+ QStringList partitionList = m_wmi->getDevicePartitions(it.key());
+ for (const QString &partition : partitionList) {
+ if (m_wmi->getLogicalDisks(partition).isEmpty()) {
+ noLogicalDevice = true;
+ break;
  }
  }
- break;
- default:
- mDebug() << this->metaObject()->className() << "Partition type: RAW";
- break;
- }
-
- mDebug() << this->metaObject()->className() << getPhysicalName(driveNumber) << " is mountable";
 
- CloseHandle(physicalHandle);
- return true;
-}
-
-bool WinDriveProvider::describeDrive(int nDriveNumber, bool verbose)
-{
- BOOL removable;
- QString productVendor;
- QString productId;
- QString serialNumber;
- uint64_t deviceBytes;
- STORAGE_BUS_TYPE storageBus;
-
- BOOL bResult = FALSE; // results flag
- // DWORD dwRet = NO_ERROR;
-
- // Format physical drive path (may be '\\.\PhysicalDrive0', '\\.\PhysicalDrive1' and so on).
- QString strDrivePath = getPhysicalName(nDriveNumber);
-
- // Get a handle to physical drive
- HANDLE hDevice = ::CreateFile(strDrivePath.toStdWString().c_str(), 0, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
-
- if (hDevice == INVALID_HANDLE_VALUE)
- return false; //::GetLastError();
-
- if (verbose)
- mDebug() << this->metaObject()->className() << strDrivePath << "is present";
-
- // Set the input data structure
- STORAGE_PROPERTY_QUERY storagePropertyQuery;
- ZeroMemory(&storagePropertyQuery, sizeof(STORAGE_PROPERTY_QUERY));
- storagePropertyQuery.PropertyId = StorageDeviceProperty;
- storagePropertyQuery.QueryType = PropertyStandardQuery;
-
- // Get the necessary output buffer size
- STORAGE_DESCRIPTOR_HEADER storageDescriptorHeader;
- ZeroMemory(&storageDescriptorHeader, sizeof(STORAGE_DESCRIPTOR_HEADER));
- DWORD dwBytesReturned = 0;
- if (verbose)
- mDebug() << this->metaObject()->className() << strDrivePath << "IOCTL_STORAGE_QUERY_PROPERTY";
- if (!::DeviceIoControl(hDevice, IOCTL_STORAGE_QUERY_PROPERTY, &storagePropertyQuery, sizeof(STORAGE_PROPERTY_QUERY), &storageDescriptorHeader, sizeof(STORAGE_DESCRIPTOR_HEADER), &dwBytesReturned, NULL)) {
- // dwRet = ::GetLastError();
- ::CloseHandle(hDevice);
- return false; // dwRet;
- }
+ auto diskDrive = m_wmi->getDiskDriveInformation(it.key(), it.value());
+ if (diskDrive->model().isEmpty() || !diskDrive->size() || diskDrive->serialNumber().isEmpty()) {
+ continue;
+ }
 
- // Alloc the output buffer
- const DWORD dwOutBufferSize = storageDescriptorHeader.Size;
- BYTE *pOutBuffer = new BYTE[dwOutBufferSize];
- ZeroMemory(pOutBuffer, dwOutBufferSize);
-
- if (verbose)
- mDebug() << this->metaObject()->className() << strDrivePath << "IOCTL_STORAGE_QUERY_PROPERTY with a bigger buffer";
- // Get the storage device descriptor
- if (!(bResult = ::DeviceIoControl(hDevice, IOCTL_STORAGE_QUERY_PROPERTY, &storagePropertyQuery, sizeof(STORAGE_PROPERTY_QUERY), pOutBuffer, dwOutBufferSize, &dwBytesReturned, NULL))) {
- // dwRet = ::GetLastError();
- delete[] pOutBuffer;
- ::CloseHandle(hDevice);
- return false; // dwRet;
+ WinDrive *currentDrive = new WinDrive(this, diskDrive->model(), diskDrive->size(), noLogicalDevice, diskDrive->index(), diskDrive->serialNumber());
+ drives[diskDrive->index()] = currentDrive;
  }
 
- // Now, the output buffer points to a STORAGE_DEVICE_DESCRIPTOR structure
- // followed by additional info like vendor ID, product ID, serial number, and so on.
- STORAGE_DEVICE_DESCRIPTOR *pDeviceDescriptor = (STORAGE_DEVICE_DESCRIPTOR *)pOutBuffer;
- removable = pDeviceDescriptor->RemovableMedia;
- if (pDeviceDescriptor->ProductIdOffset != 0)
- productId = QString((char *)pOutBuffer + pDeviceDescriptor->ProductIdOffset).trimmed();
- if (pDeviceDescriptor->VendorIdOffset != 0)
- productVendor = QString((char *)pOutBuffer + pDeviceDescriptor->VendorIdOffset).trimmed();
- if (pDeviceDescriptor->SerialNumberOffset != 0)
- serialNumber = QString((char *)pOutBuffer + pDeviceDescriptor->SerialNumberOffset).trimmed();
- storageBus = pDeviceDescriptor->BusType;
-
- if (verbose)
- mDebug() << this->metaObject()->className() << strDrivePath << "detected:" << productVendor << productId << (removable ? ", removable" : ", nonremovable") << (storageBus == BusTypeUsb ? "USB" : "notUSB");
-
- if (!removable && storageBus != BusTypeUsb)
- return false;
-
- DISK_GEOMETRY pdg;
- DWORD junk = 0; // discard results
-
- if (verbose)
- mDebug() << this->metaObject()->className() << strDrivePath << "IOCTL_DISK_GET_DRIVE_GEOMETRY";
- bResult = DeviceIoControl(hDevice, // device to be queried
- IOCTL_DISK_GET_DRIVE_GEOMETRY, // operation to perform
- NULL,
- 0, // no input buffer
- &pdg,
- sizeof(pdg), // output buffer
- &junk, // # bytes returned
- (LPOVERLAPPED)NULL); // synchronous I/O
-
- if (!bResult || pdg.MediaType == Unknown)
- return false;
+ // Update our list of drives and notify about added and removed drives
+ QList<int> driveIndexes = m_drives.keys();
+ for (auto it = drives.constBegin(); it != drives.constEnd(); it++) {
+ if (m_drives.contains(it.key()) && *m_drives[it.key()] == *it.value()) {
+ mDebug() << "Drive " << it.key() << " already exists";
+ it.value()->deleteLater();
+ driveIndexes.removeAll(it.key());
+ continue;
+ }
 
- deviceBytes = pdg.Cylinders.QuadPart * pdg.TracksPerCylinder * pdg.SectorsPerTrack * pdg.BytesPerSector;
+ if (m_drives.contains(it.key())) {
+ mDebug() << "Replacing old drive in the list on index " << it.key();
+ emit driveRemoved(m_drives[it.key()]);
+ m_drives[it.key()]->deleteLater();
+ m_drives.remove(it.key());
+ }
 
- // Do cleanup and return
- if (verbose)
- mDebug() << this->metaObject()->className() << strDrivePath << "cleanup, adding to the list";
- delete[] pOutBuffer;
- ::CloseHandle(hDevice);
+ mDebug() << "Adding new drive to the list with index " << it.key();
+ m_drives[it.key()] = it.value();
+ emit driveConnected(it.value());
 
- WinDrive *currentDrive = new WinDrive(this, productVendor + " " + productId, deviceBytes, !isMountable(nDriveNumber), nDriveNumber, serialNumber);
- if (m_drives.contains(nDriveNumber) && *m_drives[nDriveNumber] == *currentDrive) {
- currentDrive->deleteLater();
- return true;
+ driveIndexes.removeAll(it.key());
  }
 
- if (m_drives.contains(nDriveNumber)) {
- emit driveRemoved(m_drives[nDriveNumber]);
- m_drives[nDriveNumber]->deleteLater();
+ // Remove our previously stored drives that were not present in the last check
+ for (int index : driveIndexes) {
+ mDebug() << "Removing old drive with index" << index;
+ emit driveRemoved(m_drives[index]);
+ m_drives[index]->deleteLater();
+ m_drives.remove(index);
  }
-
- m_drives[nDriveNumber] = currentDrive;
- emit driveConnected(currentDrive);
-
- return true;
 }
 
 WinDrive::WinDrive(WinDriveProvider *parent, const QString &name, uint64_t size, bool containsLive, int device, const QString &serialNumber)
@@ -379,6 +216,11 @@ void WinDrive::restore()
  m_child->start(QIODevice::ReadOnly);
 }
 
+bool WinDrive::busy() const
+{
+ return (m_child && m_child->state() == QProcess::Running);
+}
+
 QString WinDrive::serialNumber() const
 {
  return m_serialNo;
@@ -455,7 +297,7 @@ void WinDrive::onReadyRead()
  Notifications::notify(tr("Finished!"), tr("Writing %1 was successful").arg(m_image->fullName()));
  } else {
  bool ok;
- qreal bytes = line.toLongLong(&ok);
+ qint64 bytes = line.toLongLong(&ok);
  if (ok) {
  if (bytes < 0)
  m_progress->setValue(NAN);