Windows: complete rewrite of drive restoration and image writing

This is a complete rewrite of the Windows support. It introduces a new
library 'libwindisk' that combines use of WMI and WinAPI to manipulate
with devices. We use Windows Storage Management for getting information
about devices and some actions, especially since pure WinAPI does not
seem to have support for things like formatting or partition removal.
This library is used in both the app and the helper process used for
writing and formatting. We now also don't rely on diskpart since it
didn't work properly in some cases.

Fixes #626 | Fixes #575 | Fixes #574 |Fixes #555 | Fixes #96

.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 libwindisk)
 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.h

@@ -21,27 +21,28 @@
 #define WINDRIVEMANAGER_H
 
 #include "drivemanager.h"
+#include "libwindisk/windisk.h"
 
+#include <QAbstractNativeEventFilter>
 #include <QProcess>
 
 class WinDriveProvider;
 class WinDrive;
 
-class WinDriveProvider : public DriveProvider
+class WinDriveProvider : public DriveProvider, public QAbstractNativeEventFilter
 {
     Q_OBJECT
 public:
     WinDriveProvider(DriveManager *parent);
 
+    bool nativeEventFilter(const QByteArray &eventType, void *message, qintptr *result) override;
+
 public slots:
     void checkDrives();
 
 private:
-    QSet<int> findPhysicalDrive(char driveLetter);
-    bool describeDrive(int driveNumber, bool verbose);
-    bool isMountable(int driveNumber);
-
     QMap<int, WinDrive *> m_drives;
+    std::unique_ptr<WinDiskManagement> m_diskManagement;
 };
 
 class WinDrive : public Drive
@@ -55,6 +56,7 @@ class WinDrive : public Drive
     Q_INVOKABLE virtual void cancel() override;
     Q_INVOKABLE virtual void restore() override;
 
+    bool busy() const;
     QString serialNumber() const;
 
     bool operator==(const WinDrive &o) const;

src/helper/win/CMakeLists.txt

@@ -26,6 +26,7 @@ target_sources(helper PRIVATE helper.exe.rc)
 target_link_libraries(helper
     Qt6::Core
     isomd5
+    libwindisk
     ${LIBLZMA_LIBRARIES}
 )
 

src/helper/win/main.cpp

@@ -1,5 +1,6 @@
 /*
  * Fedora Media Writer
+ * Copyright (C) 2024 Jan Grulich <jgrulich@redhat.com>
  * Copyright (C) 2016 Martin Bříza <mbriza@redhat.com>
  *
  * This program is free software; you can redistribute it and/or
@@ -18,27 +19,28 @@
  */
 
 #include <QCoreApplication>
+#include <QLocale>
 #include <QTextStream>
 #include <QTranslator>
-#include <QLocale>
 
 #include "restorejob.h"
 #include "writejob.h"
 
-int main(int argc, char *argv[]) {
+int main(int argc, char *argv[])
+{
     QCoreApplication app(argc, argv);
 
     QTranslator translator;
-    if (translator.load(QLocale(QLocale().language(), QLocale().country()), QLatin1String(), QLatin1String(), ":/translations"))
-        app.installTranslator(&translator);   
-
-    if (app.arguments().count() == 3 && app.arguments()[1] == "restore") {
-        new RestoreJob(app.arguments()[2]);
+    if (translator.load(QLocale(), QLatin1String(), QLatin1String(), ":/translations")) {
+        app.installTranslator(&translator);
     }
-    else if (app.arguments().count() == 4 && app.arguments()[1] == "write") {
-        new WriteJob(app.arguments()[2], app.arguments()[3]);
-    }
-    else {
+
+    const QStringList args = app.arguments();
+    if (args.count() == 3 && args[1] == "restore") {
+        new RestoreJob(args[2], &app);
+    } else if (args.count() == 4 && args[1] == "write") {
+        new WriteJob(args[2], args[3], &app);
+    } else {
         QTextStream err(stderr);
         err << "Helper: Wrong arguments entered\n";
         return 1;

src/helper/win/restorejob.cpp

@@ -18,42 +18,150 @@
  */
 
 #include "restorejob.h"
+
 #include <QCoreApplication>
 #include <QTextStream>
+#include <QThread>
 #include <QTimer>
 
-RestoreJob::RestoreJob(const QString &where)
-    : QObject(nullptr)
+RestoreJob::RestoreJob(const QString &driveNumber, QObject *parent)
+    : QObject(parent)
 {
-    bool ok = false;
-    m_where = where.toInt(&ok);
-    if (!ok)
-        qApp->exit(1);
-    else
-        QTimer::singleShot(0, this, &RestoreJob::work);
+    auto index = driveNumber.toInt();
+    m_diskManagement = std::make_unique<WinDiskManagement>(this, true);
+    m_disk = m_diskManagement->getDiskDriveInformation(index);
+
+    QTimer::singleShot(0, this, &RestoreJob::work);
 }
 
-void RestoreJob::work() {
-    m_diskpart.setProgram("diskpart.exe");
-    m_diskpart.setProcessChannelMode(QProcess::ForwardedChannels);
+void RestoreJob::work()
+{
+    HANDLE drive;
+    const QString drivePath = QString("\\\\.\\PhysicalDrive%0").arg(m_disk->index());
+
+    /*
+     * Formatting has to be apparently done in this order to be successful
+     */
+
+    /*
+     * 0) Refresh information about partitions
+     *    Uses WMI query
+     */
+    m_diskManagement->refreshDiskDrive(m_disk->path());
+
+    /*
+     * 1) Unmount all currently mounted volumes
+     *    Uses DeleteVolumeMountPointA call from WinAPI
+     *    We probably don't need to fail on this step and can try to continue
+     */
+    if (!m_diskManagement->unmountVolumes(m_disk->index())) {
+        m_err << tr("Couldn't unmount volumes on the drive") << "\n";
+        m_err.flush();
+        return;
+    }
+
+    /*
+     * 2) Remove all the existing partitions
+     *    This uses "DeleteObject" on MSFT_Partition using QMI query
+     */
+    if (!m_diskManagement->clearPartitions(m_disk->index())) {
+        m_err << tr("Failed to remove partitions from the drive") << "\n";
+        m_err.flush();
+        qApp->exit(1);
+        return;
+    }
+
+    drive = CreateFile(drivePath.toStdWString().c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
+    if (drive == INVALID_HANDLE_VALUE) {
+        m_err << tr("Failed to open the drive for formatting") << "\n";
+        m_err.flush();
+        qApp->exit(1);
+    }
+
+    /*
+     * 3) Lock the drive for the rest of the process
+     *    Uses DeviceIoControl(FSCTL_LOCK_VOLUME) from WinAPI
+     */
+    if (!m_diskManagement->lockDrive(drive, 10)) {
+        m_err << tr("Failed to lock the drive") << "\n";
+        m_err.flush();
+        qApp->exit(1);
+    }
+
+    /*
+     * 4) Refresh information about partition layout
+     *    Uses DeviceIoControl(IOCTL_DISK_UPDATE_PROPERTIES) from WinAPI
+     */
+    m_diskManagement->refreshPartitionLayout(drive);
 
-    m_diskpart.start(QIODevice::ReadWrite);
+    /*
+     * 5) Removes GPT/MBR records at the beginning and the end of the drive
+     *    Writes zeroes to the beginning and the end of the drive
+     */
+    if (!m_diskManagement->clearPartitionTable(drive, m_disk->size(), m_disk->sectorSize())) {
+        m_err << tr("Failed to clear the partition table on the drive") << "\n";
+        m_err.flush();
+        qApp->exit(1);
+    }
 
-    m_diskpart.write(qPrintable(QString("select disk %0\r\n").arg(m_where)));
-    m_diskpart.write("clean\r\n");
-    m_diskpart.write("convert gpt\r\n");
-    m_diskpart.write("convert mbr\r\n");
-    m_diskpart.write("create part pri\r\n");
-    m_diskpart.write("format fs=exFAT quick\r\n");
-    m_diskpart.write("assign\r\n");
-    m_diskpart.write("exit\r\n");
+    /*
+     * 6) Sets the drive to the RAW state
+     *    Uses DeviceIoControl(IOCTL_DISK_CREATE_DISK) with PARTITION_STYLE_RAW
+     */
+    if (!m_diskManagement->clearDiskDrive(drive)) {
+        m_err << tr("Failed to set the drive to RAW partition style") << "\n";
+        m_err.flush();
+        qApp->exit(1);
+    }
 
-    if (m_diskpart.waitForFinished()) {
-        qApp->exit(0);
+    /*
+     * 7) Created a new GPT partition on the drive
+     *    Uses DeviceIoControl(IOCTL_DISK_CREATE_DISK) with DeviceIoControl(IOCTL_DISK_SET_DRIVE_LAYOUT_EX) from WinAPI
+     */
+    if (!m_diskManagement->createGPTPartition(drive, m_disk->size(), m_disk->sectorSize())) {
+        m_err << tr("Failed to create a GPT partition on the drive") << "\n";
+        m_err.flush();
+        qApp->exit(1);
     }
-    else {
-        err << m_diskpart.readAllStandardError();
-        err.flush();
+
+    QThread::sleep(5);
+
+    /*
+     * 8) Get GUID name of the partition
+     *    Uses WinAPI to go through volumes and to get the GUID name
+     */
+    QString logicalName = m_diskManagement->getLogicalName(m_disk->index());
+    if (logicalName.isEmpty()) {
+        m_err << tr("Failed to get GUID volume path on the drive") << "\n";
+        m_err.flush();
         qApp->exit(1);
     }
+
+    m_diskManagement->refreshDiskDrive(m_disk->path());
+
+    /*
+     * 9) Attempt to mount a volume using the GUID path we get above
+     *    Uses GetVolumePathNamesForVolumeNameA() to check whether the volume is already mounted, or
+     *    SetVolumeMountPointA() to mount the partition. Returns assigned drive letter.
+     */
+    QChar driveLetter = m_diskManagement->mountVolume(logicalName);
+    if (!driveLetter.isLetter()) {
+        m_err << tr("Failed to mount the new partition") << "\n";
+        m_err.flush();
+        qApp->exit(1);
+    }
+
+    /*
+     * 10) Format the partition to exFAT
+     *     Uses "Format" method on the MSFT_Volume object using WMI query.
+     */
+    if (!m_diskManagement->formatPartition(driveLetter)) {
+        m_err << tr("Failed to format the partition to exFAT") << "\n";
+        m_err.flush();
+        qApp->exit(1);
+    }
+
+    CloseHandle(drive);
+
+    qApp->exit(0);
 }

src/helper/win/restorejob.h

@@ -1,5 +1,6 @@
 /*
  * Fedora Media Writer
+ * Copyright (C) 2024 Jan Grulich <jgrulich@redhat.com>
  * Copyright (C) 2016 Martin Bříza <mbriza@redhat.com>
  *
  * This program is free software; you can redistribute it and/or
@@ -20,6 +21,8 @@
 #ifndef RESTOREJOB_H
 #define RESTOREJOB_H
 
+#include <libwindisk/windisk.h>
+
 #include <QObject>
 #include <QProcess>
 #include <QTextStream>
@@ -28,19 +31,19 @@ class RestoreJob : public QObject
 {
     Q_OBJECT
 public:
-    explicit RestoreJob(const QString &where);
+    explicit RestoreJob(const QString &where, QObject *parent);
 
 signals:
 
 private slots:
     void work();
 
 private:
-    QTextStream out { stdout };
-    QTextStream err { stderr };
+    QTextStream m_out{stdout};
+    QTextStream m_err{stderr};
 
-    QProcess m_diskpart;
-    int m_where;
+    std::unique_ptr<WinDiskManagement> m_diskManagement;
+    std::unique_ptr<WinDisk> m_disk;
 };
 
 #endif // RESTOREJOB_H