From a2c2373e17db7c27399a485ae91dd70644945353 Mon Sep 17 00:00:00 2001 From: Jan Grulich Date: Fri, 4 Oct 2024 16:12:21 +0200 Subject: [PATCH] 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 | 2 +- src/app/CMakeLists.txt | 2 +- src/app/main.cpp | 1 + src/app/utilities.h | 2 +- src/app/windrivemanager.cpp | 319 ++----- src/app/windrivemanager.h | 7 +- src/helper/win/CMakeLists.txt | 2 + src/helper/win/main.cpp | 24 +- src/helper/win/restorejob.cpp | 173 +++- src/helper/win/restorejob.h | 13 +- src/helper/win/writejob.cpp | 579 +++++++----- src/helper/win/writejob.h | 44 +- src/lib/CMakeLists.txt | 3 + src/lib/libwindisk/CMakeLists.txt | 9 + src/lib/libwindisk/windisk.cpp | 1061 ++++++++++++++++++++++ src/lib/libwindisk/windisk.h | 122 +++ 16 files changed, 1812 insertions(+), 551 deletions(-) create mode 100644 src/lib/libwindisk/CMakeLists.txt create mode 100644 src/lib/libwindisk/windisk.cpp create mode 100644 src/lib/libwindisk/windisk.h diff --git a/.github/workflows/clang-format-check.yml b/.github/workflows/clang-format-check.yml index 5d45c7e1..9c856529 100644 --- a/.github/workflows/clang-format-check.yml +++ b/.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' diff --git a/src/app/CMakeLists.txt b/src/app/CMakeLists.txt index 75fdf759..528ca8f8 100644 --- a/src/app/CMakeLists.txt +++ b/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) diff --git a/src/app/main.cpp b/src/app/main.cpp index e7512088..eb09fd7a 100644 --- a/src/app/main.cpp +++ b/src/app/main.cpp @@ -1,5 +1,6 @@ /* * Fedora Media Writer + * Copyright (C) 2024 Jan Grulich * Copyright (C) 2016 Martin Bříza * * This program is free software; you can redistribute it and/or diff --git a/src/app/utilities.h b/src/app/utilities.h index 424c477b..6c714110 100644 --- a/src/app/utilities.h +++ b/src/app/utilities.h @@ -78,7 +78,7 @@ class Options bool testing{false}; bool verbose{false}; -#ifdef QT_NO_DEBUG +#if defined(QT_NO_DEBUG) && !defined(_WIN32) bool logging{false}; #else bool logging{true}; diff --git a/src/app/windrivemanager.cpp b/src/app/windrivemanager.cpp index 7fb1035c..38ba6a76 100644 --- a/src/app/windrivemanager.cpp +++ b/src/app/windrivemanager.cpp @@ -1,6 +1,6 @@ /* * Fedora Media Writer - * Copyright (C) 2022 Jan Grulich + * Copyright (C) 2022-2024 Jan Grulich * Copyright (C) 2011-2022 Pete Batard * Copyright (C) 2016 Martin Bříza * @@ -22,23 +22,16 @@ #include "windrivemanager.h" #include "notifications.h" -#include #include #include -#include -#define INITGUID -#include +#include -#include -#include - -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_diskManagement(std::make_unique(this)) { mDebug() << this->metaObject()->className() << "construction"; QTimer::singleShot(0, this, &WinDriveProvider::checkDrives); @@ -46,238 +39,75 @@ WinDriveProvider::WinDriveProvider(DriveManager *parent) void WinDriveProvider::checkDrives() { - 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); + mDebug() << this->metaObject()->className() << "Looking for the drives"; + + 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()) { + // Skip this round + QTimer::singleShot(2500, this, &WinDriveProvider::checkDrives); + return; } } - 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; -} - -bool WinDriveProvider::isMountable(int driveNumber) -{ - 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 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; + QMap drives; + auto usbDeviceList = m_diskManagement->getUSBDeviceList(); + for (auto it = usbDeviceList.cbegin(); it != usbDeviceList.cend(); it++) { + bool mountable = true; + auto partitionList = m_diskManagement->getDevicePartitions(it.key()); + if (partitionList.isEmpty()) { + mountable = false; + } else { + for (auto it = partitionList.constBegin(); it != partitionList.constEnd(); it++) { + if (!it.value()) { + mountable = false; + break; } } } - 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; - } - } - 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; - } - - // 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; + auto diskDrive = m_diskManagement->getDiskDriveInformation(it.key(), it.value()); + if (diskDrive->name().isEmpty() || !diskDrive->size() || diskDrive->serialNumber().isEmpty()) { + continue; + } + WinDrive *currentDrive = new WinDrive(this, diskDrive->name(), diskDrive->size(), !mountable, 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 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; + QTimer::singleShot(2500, this, &WinDriveProvider::checkDrives); } WinDrive::WinDrive(WinDriveProvider *parent, const QString &name, uint64_t size, bool containsLive, int device, const QString &serialNumber) @@ -290,7 +120,7 @@ WinDrive::WinDrive(WinDriveProvider *parent, const QString &name, uint64_t size, WinDrive::~WinDrive() { if (m_child) - m_child->kill(); + m_child->terminate(); } bool WinDrive::write(ReleaseVariant *data) @@ -339,7 +169,7 @@ void WinDrive::cancel() { Drive::cancel(); if (m_child) { - m_child->kill(); + m_child->terminate(); m_child->deleteLater(); m_child = nullptr; } @@ -379,6 +209,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; @@ -402,9 +237,12 @@ void WinDrive::onFinished(int exitCode, QProcess::ExitStatus exitStatus) if (exitCode == 0) { m_image->setStatus(ReleaseVariant::FINISHED); Notifications::notify(tr("Finished!"), tr("Writing %1 was successful").arg(m_image->fullName())); - } else { + } else if (exitCode == 1) { m_image->setErrorString(m_child->readAllStandardError().trimmed()); m_image->setStatus(ReleaseVariant::FAILED); + } else if (exitCode == 2) { + m_image->setErrorString(tr("Writing has been cancelled")); + m_image->setStatus(ReleaseVariant::FAILED); } m_child->deleteLater(); @@ -413,16 +251,18 @@ void WinDrive::onFinished(int exitCode, QProcess::ExitStatus exitStatus) void WinDrive::onRestoreFinished(int exitCode, QProcess::ExitStatus exitStatus) { - if (!m_child) + if (!m_child) { return; + } mCritical() << "Process finished" << exitCode << exitStatus; mCritical() << m_child->readAllStandardError(); - if (exitCode == 0) + if (exitCode == 0) { m_restoreStatus = RESTORED; - else + } else { m_restoreStatus = RESTORE_ERROR; + } emit restoreStatusChanged(); m_child->deleteLater(); @@ -431,14 +271,16 @@ void WinDrive::onRestoreFinished(int exitCode, QProcess::ExitStatus exitStatus) void WinDrive::onReadyRead() { - if (!m_child) + if (!m_child) { return; + } m_progress->setTo(m_image->size()); m_progress->setValue(NAN); - if (m_image->status() != ReleaseVariant::WRITE_VERIFYING && m_image->status() != ReleaseVariant::WRITING) + if (m_image->status() != ReleaseVariant::WRITE_VERIFYING && m_image->status() != ReleaseVariant::WRITING) { m_image->setStatus(ReleaseVariant::WRITING); + } while (m_child->bytesAvailable() > 0) { QString line = m_child->readLine().trimmed(); @@ -455,12 +297,13 @@ 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) + if (bytes < 0) { m_progress->setValue(NAN); - else + } else { m_progress->setValue(bytes); + } } } } diff --git a/src/app/windrivemanager.h b/src/app/windrivemanager.h index 3d1ac43d..54e9d094 100644 --- a/src/app/windrivemanager.h +++ b/src/app/windrivemanager.h @@ -21,6 +21,7 @@ #define WINDRIVEMANAGER_H #include "drivemanager.h" +#include "libwindisk/windisk.h" #include @@ -37,11 +38,8 @@ public slots: void checkDrives(); private: - QSet findPhysicalDrive(char driveLetter); - bool describeDrive(int driveNumber, bool verbose); - bool isMountable(int driveNumber); - QMap m_drives; + std::unique_ptr m_diskManagement; }; class WinDrive : public Drive @@ -55,6 +53,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; diff --git a/src/helper/win/CMakeLists.txt b/src/helper/win/CMakeLists.txt index 60695802..21cb76df 100644 --- a/src/helper/win/CMakeLists.txt +++ b/src/helper/win/CMakeLists.txt @@ -25,7 +25,9 @@ target_sources(helper PRIVATE helper.exe.rc) target_link_libraries(helper Qt6::Core + Qt6::Gui isomd5 + libwindisk ${LIBLZMA_LIBRARIES} ) diff --git a/src/helper/win/main.cpp b/src/helper/win/main.cpp index 4152eec7..0053533f 100644 --- a/src/helper/win/main.cpp +++ b/src/helper/win/main.cpp @@ -1,5 +1,6 @@ /* * Fedora Media Writer + * Copyright (C) 2024 Jan Grulich * Copyright (C) 2016 Martin Bříza * * This program is free software; you can redistribute it and/or @@ -18,27 +19,28 @@ */ #include +#include #include #include -#include #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; diff --git a/src/helper/win/restorejob.cpp b/src/helper/win/restorejob.cpp index b4633756..2a9c6301 100644 --- a/src/helper/win/restorejob.cpp +++ b/src/helper/win/restorejob.cpp @@ -1,5 +1,6 @@ /* * Fedora Media Writer + * Copyright (C) 2024 Jan Grulich * Copyright (C) 2016 Martin Bříza * * This program is free software; you can redistribute it and/or @@ -18,42 +19,164 @@ */ #include "restorejob.h" + #include #include +#include #include -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(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()); - m_diskpart.start(QIODevice::ReadWrite); + /* + * Formatting has to be apparently done in this order to be successful + */ - 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"); + /* + * 0) Refresh information about partitions + * Uses WMI query + */ + m_diskManagement->refreshDiskDrive(m_disk->path()); - if (m_diskpart.waitForFinished()) { - qApp->exit(0); + /* + * 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->removeDriveLetters(m_disk->index())) { + m_diskManagement->logMessage(QtCriticalMsg, "Couldn't remove drive mountpoints"); + m_err << tr("Couldn't remove drive mountpoints") << "\n"; + m_err.flush(); + return; } - else { - err << m_diskpart.readAllStandardError(); - err.flush(); + + /* + * 2) Remove all the existing partitions + * This uses "DeleteObject" on MSFT_Partition using QMI query + */ + if (!m_diskManagement->clearPartitions(m_disk->index())) { + m_diskManagement->logMessage(QtCriticalMsg, "Failed to remove partitions from the drive"); + 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_diskManagement->logMessage(QtCriticalMsg, "Failed to open the drive for formatting"); + m_err << tr("Failed to open the drive for formatting") << "\n"; + m_err.flush(); + qApp->exit(1); + } + auto cleanup = qScopeGuard([=] { + m_diskManagement->unlockDrive(drive); + CloseHandle(drive); + }); + + /* + * 3) Lock the drive for the rest of the process + * Uses DeviceIoControl(FSCTL_LOCK_VOLUME) from WinAPI + */ + if (!m_diskManagement->lockDrive(drive, 10)) { + m_diskManagement->logMessage(QtCriticalMsg, "Couldn't lock the drive"); + m_err << tr("Couldn't 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); + + /* + * 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_diskManagement->logMessage(QtCriticalMsg, "Failed to clear the partition table on the drive"); + m_err << tr("Failed to clear the partition table on the drive") << "\n"; + m_err.flush(); + qApp->exit(1); + } + + /* + * 6) Sets the drive to the RAW state + * Uses DeviceIoControl(IOCTL_DISK_CREATE_DISK) with PARTITION_STYLE_RAW + */ + if (!m_diskManagement->clearDiskDrive(drive)) { + m_diskManagement->logMessage(QtCriticalMsg, "Failed to set the drive to RAW partition style"); + m_err << tr("Failed to set the drive to RAW partition style") << "\n"; + m_err.flush(); + qApp->exit(1); + } + + /* + * 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_diskManagement->logMessage(QtCriticalMsg, "Failed to create a GPT partition on the drive"); + m_err << tr("Failed to create a GPT partition on the drive") << "\n"; + m_err.flush(); + qApp->exit(1); + } + + // FIXME: isn't this too much? We used to have this even before as + // apparently it was suggested after diskpart operations + QThread::sleep(15); + + /* + * 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_diskManagement->logMessage(QtCriticalMsg, "Failed to get GUID volume path on the drive"); + 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_diskManagement->logMessage(QtCriticalMsg, "Failed to remove partitions from the drive"); + 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_diskManagement->logMessage(QtCriticalMsg, "Failed to format the partition to exFAT"); + m_err << tr("Failed to format the partition to exFAT") << "\n"; + m_err.flush(); + qApp->exit(1); + } + + qApp->exit(0); } diff --git a/src/helper/win/restorejob.h b/src/helper/win/restorejob.h index a46ada9b..414226a7 100644 --- a/src/helper/win/restorejob.h +++ b/src/helper/win/restorejob.h @@ -1,5 +1,6 @@ /* * Fedora Media Writer + * Copyright (C) 2024 Jan Grulich * Copyright (C) 2016 Martin Bříza * * This program is free software; you can redistribute it and/or @@ -20,6 +21,8 @@ #ifndef RESTOREJOB_H #define RESTOREJOB_H +#include + #include #include #include @@ -28,7 +31,7 @@ class RestoreJob : public QObject { Q_OBJECT public: - explicit RestoreJob(const QString &where); + explicit RestoreJob(const QString &where, QObject *parent); signals: @@ -36,11 +39,11 @@ 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 m_diskManagement; + std::unique_ptr m_disk; }; #endif // RESTOREJOB_H diff --git a/src/helper/win/writejob.cpp b/src/helper/win/writejob.cpp index 30947f30..8f2c1b7c 100644 --- a/src/helper/win/writejob.cpp +++ b/src/helper/win/writejob.cpp @@ -1,5 +1,6 @@ /* * Fedora Media Writer + * Copyright (C) 2024 Jan Grulich * Copyright (C) 2016 Martin Bříza * * This program is free software; you can redistribute it and/or @@ -20,397 +21,499 @@ #include "writejob.h" #include -#include -#include -#include #include -#include +#include +#include +#include #include -#include +#include +#include #include -#include #include #include "isomd5/libcheckisomd5.h" +static QString getLastError() +{ + TCHAR message[256]; + FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); + return QString::fromWCharArray(message).trimmed(); +} -WriteJob::WriteJob(const QString &what, const QString &where) - : QObject(nullptr), what(what) +WriteJob::WriteJob(const QString &image, const QString &driveNumber, QObject *parent) + : QObject(parent) + , m_image(image) { - bool ok = false; - this->where = where.toInt(&ok); + const int wmiDriveNumber = driveNumber.toInt(); - if (what.endsWith(".part")) { - connect(&watcher, &QFileSystemWatcher::fileChanged, this, &WriteJob::onFileChanged); - watcher.addPath(what); - } - else { + m_diskManagement = std::make_unique(this, true); + m_disk = m_diskManagement->getDiskDriveInformation(wmiDriveNumber); + + if (m_image.endsWith(".part")) { + connect(&m_watcher, &QFileSystemWatcher::fileChanged, this, &WriteJob::onFileChanged); + m_watcher.addPath(m_image); + } else { QTimer::singleShot(0, this, &WriteJob::work); } } -int WriteJob::staticOnMediaCheckAdvanced(void *data, long long offset, long long total) { - return ((WriteJob*)data)->onMediaCheckAdvanced(offset, total); +int WriteJob::staticOnMediaCheckAdvanced(void *data, long long offset, long long total) +{ + return ((WriteJob *)data)->onMediaCheckAdvanced(offset, total); } -int WriteJob::onMediaCheckAdvanced(long long offset, long long total) { +int WriteJob::onMediaCheckAdvanced(long long offset, long long total) +{ Q_UNUSED(total); - out << offset << "\n"; - out.flush(); + m_out << offset << "\n"; + m_out.flush(); return 0; } -HANDLE WriteJob::openDrive(int physicalDriveNumber) { - HANDLE hVol; - QString drivePath = QString("\\\\.\\PhysicalDrive%0").arg(physicalDriveNumber); - - hVol = CreateFile(drivePath.toStdWString().c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, 0, OPEN_EXISTING, 0, 0); +void WriteJob::work() +{ + if (!write()) { + qApp->exit(1); + return; + } - if( hVol == INVALID_HANDLE_VALUE ) { - TCHAR message[256]; - FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); - err << tr("Couldn't open the drive for writing") << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; - err.flush(); - return hVol; + if (!check()) { + qApp->exit(1); + return; } - return hVol; + qApp->exit(0); } -bool WriteJob::lockDrive(HANDLE drive) { - int attempts = 0; - DWORD status; +void WriteJob::onFileChanged(const QString &path) +{ + if (QFile::exists(path)) + return; + QRegularExpression reg("[.]part$"); + m_image = m_image.replace(reg, ""); - while (true) { - if (!DeviceIoControl(drive, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &status, NULL)) { - attempts++; - } - else { - return true; - } + work(); +} - if (attempts == 10) { - TCHAR message[256]; - FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); +bool WriteJob::write() +{ + m_out << "WRITE\n"; + m_out.flush(); - err << tr("Couldn't lock the drive") << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; - err.flush(); - break; - } + m_diskManagement->logMessage(QtDebugMsg, "Preparing device for image writing"); - QThread::sleep(2); - } + /* + * Device preparation part + */ - return false; -} + if (!m_diskManagement->removeDriveLetters(m_disk->index())) { + m_err << tr("Couldn't remove drive mountpoints") << "\n"; + m_err.flush(); + return false; + } -bool WriteJob::removeMountPoints(uint diskNumber) { - DWORD drives = ::GetLogicalDrives(); + // This doesn't need to be fatal and we can try to continue + if (!m_diskManagement->clearPartitions(m_disk->index())) { + m_err << tr("Failed to remove partitions from the drive") << "\n"; + m_err.flush(); + } - for (char i = 0; i < 26; i++) { - if (drives & (1 << i)) { - char currentDrive = 'A' + i; - QString drivePath = QString("\\\\.\\%1:").arg(currentDrive); + HANDLE drive; + const QString drivePath = QString("\\\\.\\PhysicalDrive%0").arg(m_disk->index()); - HANDLE hDevice = ::CreateFile(drivePath.toStdWString().c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL); + 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(); + return false; + } - DWORD bytesReturned; - VOLUME_DISK_EXTENTS vde; // TODO FIXME: handle ERROR_MORE_DATA (this is an extending structure) - BOOL bResult = DeviceIoControl(hDevice, IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS, NULL, 0, &vde, sizeof(vde), &bytesReturned, NULL); + if (!m_diskManagement->lockDrive(drive, 10)) { + m_err << tr("Couldn't lock the drive") << "\n"; + m_err.flush(); + return false; + } - if (bResult) { - for (uint j = 0; j < vde.NumberOfDiskExtents; j++) { - if (vde.Extents[j].DiskNumber == diskNumber) { - QString volumePath = QString("%1:\\").arg(currentDrive); + m_diskManagement->refreshPartitionLayout(drive); - CloseHandle(hDevice); - hDevice = nullptr; + HANDLE logicalHandle = NULL; + const QString logicalVolumePath = m_diskManagement->getLogicalName(m_disk->index(), false); + if (!logicalVolumePath.isEmpty()) { + m_diskManagement->logMessage(QtDebugMsg, "Trying to lock and unmount logical volume"); + logicalHandle = CreateFile(logicalVolumePath.toStdWString().c_str(), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); + if (logicalHandle == INVALID_HANDLE_VALUE) { + m_err << tr("Couldn't open the logical handle") << "\n"; + m_err.flush(); + return false; + } - if (!DeleteVolumeMountPointA(volumePath.toStdString().c_str())) { - TCHAR message[256]; - FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); - err << tr("Couldn't remove the drive %1:").arg(currentDrive) << " (" << QString::fromWCharArray(message).trimmed() << "\n"; - err.flush(); - return false; - } + if (!m_diskManagement->lockDrive(logicalHandle, 10)) { + m_err << tr("Couldn't lock the logical drive") << "\n"; + m_err.flush(); + return false; + } - break; - } - } - } - if (hDevice) - CloseHandle(hDevice); + // This doesn't need to be fatal and we can try to continue + if (!m_diskManagement->unmountVolume(logicalHandle)) { + m_err << tr("Couldn't unmount drive") << "\n"; + m_err.flush(); } } - return true; -} - -bool WriteJob::cleanDrive(uint driveNumber) { - QProcess diskpart; - diskpart.setProgram("diskpart.exe"); - diskpart.setProcessChannelMode(QProcess::ForwardedChannels); + 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(); + return false; + } - diskpart.start(QIODevice::ReadWrite); + if (!m_diskManagement->clearDiskDrive(drive)) { + m_err << tr("Failed to set the drive to RAW partition style") << "\n"; + m_err.flush(); + return false; + } - diskpart.write(qPrintable(QString("select disk %0\r\n").arg(driveNumber))); - diskpart.write("clean\r\n"); - // for some reason this works (tm) - diskpart.write("create part pri\r\n"); - diskpart.write("clean\r\n"); - diskpart.write("exit\r\n"); + m_diskManagement->unlockDrive(drive); + CloseHandle(drive); - diskpart.waitForFinished(); + // FIXME: isn't this too much? We used to have this even before as + // apparently it was suggested after diskpart operations + QThread::sleep(15); - if (diskpart.exitCode() == 0) { - // as advised in the diskpart documentation - QThread::sleep(15); + /* + * Writing part + */ - return true; - } + m_diskManagement->logMessage(QtDebugMsg, "Starting to write image file"); - return false; -} + // drive = CreateFile(drivePath.toStdWString().c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_FLAG_NO_BUFFERING | FILE_FLAG_WRITE_THROUGH, NULL); -bool WriteJob::writeBlock(HANDLE drive, OVERLAPPED *overlap, char *data, uint size) { - DWORD bytesWritten; - - if (!WriteFile(drive, data, size, &bytesWritten, overlap)) { - DWORD Errorcode = GetLastError(); - if (Errorcode == ERROR_IO_PENDING) { - WaitForSingleObject(overlap->hEvent, INFINITE); - } - else { - TCHAR message[256]; - FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); - err << tr("Destination drive is not writable") << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; - err.flush(); - return false; - } - } - - if (bytesWritten != size) { - err << tr("Destination drive is not writable") << "\n"; - err.flush(); + drive = CreateFile(drivePath.toStdWString().c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); + if (drive == INVALID_HANDLE_VALUE) { + m_err << tr("Couldn't open the drive for writing") << "\n"; + m_err.flush(); return false; } - return true; -} - + // m_diskManagement->disableIOBoundaryChecks(drive); -void WriteJob::unlockDrive(HANDLE drive) { - DWORD status; - if (!DeviceIoControl(drive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &status, NULL)) { - TCHAR message[256]; - FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); - err << tr("Couldn't unlock the drive") << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; - err.flush(); + if (!m_diskManagement->lockDrive(drive, 10)) { + m_err << tr("Couldn't lock the drive") << "\n"; + m_err.flush(); + return false; } -} -void WriteJob::work() { - if (!write()) { - out << "0\n"; - out.flush(); - QThread::sleep(5); - if (!write()) - return; + bool result; + if (m_image.endsWith(".xz")) { + result = writeCompressed(drive); + } else { + result = writePlain(drive); } - if (!check()) - return; - - qApp->exit(0); -} + if (result) { + m_diskManagement->refreshPartitionLayout(drive); + } -void WriteJob::onFileChanged(const QString &path) { - if (QFile::exists(path)) - return; - QRegularExpression reg("[.]part$"); - what = what.replace(reg, ""); + if (logicalHandle && logicalHandle != INVALID_HANDLE_VALUE) { + m_diskManagement->unlockDrive(logicalHandle); + CloseHandle(logicalHandle); + } + CloseHandle(drive); - out << "WRITE\n"; - out.flush(); + if (!result) { + qApp->exit(1); + return false; + } - work(); + return true; } -bool WriteJob::write() { - removeMountPoints(where); - cleanDrive(where); +bool WriteJob::writeCompressed(HANDLE driveHandle) +{ + const qint64 blockSize = m_disk->sectorSize() * 512; - HANDLE drive = openDrive(where); - if (!lockDrive(drive)) { - qApp->exit(1); + QFile isoFile(m_image); + isoFile.open(QIODevice::ReadOnly); + if (!isoFile.isOpen()) { + m_err << tr("Source image is not readable"); + m_err.flush(); return false; } + auto isoCleanup = qScopeGuard([&isoFile] { + isoFile.close(); + }); + + QFile drive; + drive.open(_open_osfhandle(reinterpret_cast(driveHandle), 0), QIODevice::WriteOnly /*| QIODevice::Unbuffered*/, QFile::AutoCloseHandle); + if (!drive.isOpen()) { + m_err << tr("Failed to open device for writing") << "\n"; + m_err.flush(); + return false; + } + auto driveCleanup = qScopeGuard([&drive] { + drive.close(); + }); + + void *outBuffer = NULL; + outBuffer = VirtualAlloc(NULL, blockSize, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE); + if (outBuffer == NULL) { + m_err << tr("Failed to allocate the buffer") << "\n"; + m_err.flush(); + return false; + } + auto outBufferCleanup = qScopeGuard([outBuffer, blockSize] { + VirtualFree(outBuffer, blockSize, MEM_DECOMMIT | MEM_RELEASE); + }); - if (what.endsWith(".xz")) - return writeCompressed(drive); - else - return writePlain(drive); -} - -bool WriteJob::writeCompressed(HANDLE drive) { qint64 totalRead = 0; lzma_stream strm = LZMA_STREAM_INIT; lzma_ret ret; - uint8_t *inBuffer = new uint8_t[BLOCK_SIZE]; - uint8_t *outBuffer = new uint8_t[BLOCK_SIZE]; + uint8_t *inBuffer = new uint8_t[blockSize]; + auto inBufferCleanup = qScopeGuard([inBuffer] { + delete[] inBuffer; + }); - QFile file(what); + QFile file(m_image); file.open(QIODevice::ReadOnly); ret = lzma_stream_decoder(&strm, MEDIAWRITER_LZMA_LIMIT, LZMA_CONCATENATED); if (ret != LZMA_OK) { - err << tr("Failed to start decompressing."); + m_err << tr("Failed to start decompressing.") << "\n"; return false; } strm.next_in = inBuffer; strm.avail_in = 0; - strm.next_out = outBuffer; - strm.avail_out = BLOCK_SIZE; - - OVERLAPPED osWrite; - memset(&osWrite, 0, sizeof(osWrite)); - osWrite.hEvent = 0; + strm.next_out = static_cast(outBuffer); + strm.avail_out = blockSize; + const qint64 sectorSize = m_disk->sectorSize(); while (true) { if (strm.avail_in == 0) { - qint64 len = file.read((char*) inBuffer, BLOCK_SIZE); + qint64 len = file.read((char *)inBuffer, blockSize); totalRead += len; strm.next_in = inBuffer; strm.avail_in = len; - out << totalRead << "\n"; - out.flush(); + m_out << totalRead << "\n"; + m_out.flush(); } ret = lzma_code(&strm, strm.avail_in == 0 ? LZMA_FINISH : LZMA_RUN); if (ret == LZMA_STREAM_END) { - if (!writeBlock(drive, &osWrite, (char *) outBuffer, BLOCK_SIZE - strm.avail_out)) { + qint64 writtenBytes = 0; + qint64 readBytes = blockSize - strm.avail_out; + readBytes = ((readBytes + sectorSize - 1) / sectorSize) * sectorSize; + writtenBytes = drive.write(reinterpret_cast(outBuffer), readBytes); + + if (writtenBytes <= 0) { + m_diskManagement->logMessage(QtCriticalMsg, QStringLiteral("Destination drive is not writable: %1").arg(getLastError())); + m_err << tr("Destination drive is not writable") << ": " << getLastError() << "\n"; + m_err.flush(); qApp->exit(1); - CloseHandle(drive); return false; } - if (osWrite.Offset + BLOCK_SIZE < osWrite.Offset) - osWrite.OffsetHigh++; - osWrite.Offset += BLOCK_SIZE; - - CloseHandle(drive); + if (writtenBytes != readBytes) { + m_err << tr("The last block was not fully written") << "\n"; + m_err.flush(); + return false; + } return true; } + if (ret != LZMA_OK) { switch (ret) { case LZMA_MEM_ERROR: - err << tr("There is not enough memory to decompress the file."); + m_err << tr("There is not enough memory to decompress the file.") << "\n"; break; case LZMA_FORMAT_ERROR: case LZMA_DATA_ERROR: case LZMA_BUF_ERROR: - err << tr("The downloaded compressed file is corrupted."); + m_err << tr("The downloaded compressed file is corrupted.") << "\n"; break; case LZMA_OPTIONS_ERROR: - err << tr("Unsupported compression options."); + m_err << tr("Unsupported compression options.") << "\n"; break; default: - err << tr("Unknown decompression error."); + m_err << tr("Unknown decompression error.") << "\n"; break; } qApp->exit(4); - CloseHandle(drive); return false; } if (strm.avail_out == 0) { - if (!writeBlock(drive, &osWrite, (char *) outBuffer, BLOCK_SIZE - strm.avail_out)) { + qint64 writtenBytes = 0; + writtenBytes = drive.write(reinterpret_cast(outBuffer), sectorSize); + if (writtenBytes <= 0) { + m_diskManagement->logMessage(QtCriticalMsg, QStringLiteral("Destination drive is not writable: %1").arg(getLastError())); + m_err << tr("Destination drive is not writable") << ": " << getLastError() << "\n"; + m_err.flush(); qApp->exit(1); - CloseHandle(drive); return false; } - if (osWrite.Offset + BLOCK_SIZE < osWrite.Offset) - osWrite.OffsetHigh++; - osWrite.Offset += BLOCK_SIZE; + if (writtenBytes != sectorSize) { + m_err << tr("The last block was not fully written") << "\n"; + m_err.flush(); + qApp->exit(1); + return false; + } - strm.next_out = outBuffer; - strm.avail_out = BLOCK_SIZE; + strm.next_out = static_cast(outBuffer); + strm.avail_out = blockSize; } } + return false; } -bool WriteJob::writePlain(HANDLE drive) { - OVERLAPPED osWrite; - memset(&osWrite, 0, sizeof(osWrite)); - osWrite.hEvent = 0; +static QString calculcateSizeString(quint64 bytes) +{ + const quint64 kb = 1024; + const quint64 mb = 1024 * kb; + const quint64 gb = 1024 * mb; + const quint64 tb = 1024 * gb; + + if (bytes >= tb) { + return QFileSystemModel::tr("%1 TB").arg(QLocale().toString(qreal(bytes) / tb, 'f', 3)); + } else if (bytes >= gb) { + return QFileSystemModel::tr("%1 GB").arg(QLocale().toString(qreal(bytes) / gb, 'f', 2)); + } else if (bytes >= mb) { + return QFileSystemModel::tr("%1 MB").arg(QLocale().toString(qreal(bytes) / mb, 'f', 1)); + } else if (bytes >= kb) { + return QFileSystemModel::tr("%1 KB").arg(QLocale().toString(bytes / kb)); + } + + return QFileSystemModel::tr("%1 byte(s)").arg(QLocale().toString(bytes)); +} + +bool WriteJob::writePlain(HANDLE driveHandle) +{ + const qint64 blockSize = m_disk->sectorSize() * 512; - uint64_t cnt = 0; - QByteArray buffer; - QFile isoFile(what); + QFile isoFile(m_image); isoFile.open(QIODevice::ReadOnly); if (!isoFile.isOpen()) { - err << tr("Source image is not readable"); - err.flush(); - qApp->exit(1); + m_err << tr("Source image is not readable") << "\n"; + m_err.flush(); + return false; + } + const qint64 imageSize = isoFile.size(); + auto isoCleanup = qScopeGuard([&isoFile] { + isoFile.close(); + }); + + QFile drive; + drive.open(_open_osfhandle(reinterpret_cast(driveHandle), 0), QIODevice::WriteOnly /*| QIODevice::Unbuffered*/, QFile::AutoCloseHandle); + if (!drive.isOpen()) { + m_err << tr("Failed to open device for writing") << "\n"; + m_err.flush(); + return false; + } + auto driveCleanup = qScopeGuard([&drive] { + drive.close(); + }); + + uint8_t *buffer = NULL; + buffer = static_cast(_mm_malloc(blockSize, m_disk->sectorSize())); + if (!buffer) { + m_err << tr("Failed to allocate the buffer") << "\n"; + m_err.flush(); return false; } + auto bufferCleanup = qScopeGuard([buffer] { + if (buffer) { + _mm_free(buffer); + } + }); + qint64 sectorSize = m_disk->sectorSize(); + qint64 totalBytes = 0; + qint64 readBytes; + qint64 writtenBytes; while (true) { - buffer = isoFile.read(BLOCK_SIZE); - if (!writeBlock(drive, &osWrite, buffer.data(), buffer.size())) { - qApp->exit(1); + if ((readBytes = isoFile.read(reinterpret_cast(buffer), blockSize)) <= 0) { + break; + } + + readBytes = ((readBytes + sectorSize - 1) / sectorSize) * sectorSize; + writtenBytes = drive.write(reinterpret_cast(buffer), readBytes); + if (writtenBytes <= 0) { + m_diskManagement->logMessage(QtCriticalMsg, QStringLiteral("Destination drive is not writable: %1").arg(getLastError())); + m_diskManagement->logMessage(QtCriticalMsg, QStringLiteral("Wrote %1 out of %2").arg(calculcateSizeString(totalBytes)).arg(calculcateSizeString(imageSize))); + m_err << tr("Destination drive is not writable") << ": " << getLastError() << "\n"; + m_err << tr("Wrote %1 out of %2").arg(calculcateSizeString(totalBytes)).arg(calculcateSizeString(imageSize)) << "\n"; + + m_err.flush(); + return false; + } + + if (writtenBytes != readBytes) { + m_diskManagement->logMessage(QtCriticalMsg, QStringLiteral("The last block was not fully written: %1").arg(getLastError())); + m_err << tr("The last block was not fully written") << "\n"; + m_err.flush(); return false; } - if (osWrite.Offset + BLOCK_SIZE < osWrite.Offset) - osWrite.OffsetHigh++; - osWrite.Offset += BLOCK_SIZE; - cnt += buffer.size(); - out << cnt << "\n"; - out.flush(); + totalBytes += readBytes; + m_out << totalBytes << "\n"; + m_out.flush(); - if (buffer.size() != BLOCK_SIZE || isoFile.atEnd()) + if (readBytes != blockSize || isoFile.atEnd()) { break; + } } - CloseHandle(drive); + if (readBytes < 0) { + m_err << tr("Failed to read the image file: ") << isoFile.errorString() << "\n"; + m_err.flush(); + return false; + } return true; } -bool WriteJob::check() { - out << "CHECK\n"; - out.flush(); - - HANDLE drive = openDrive(where); +bool WriteJob::check() +{ + m_out << "CHECK\n"; + m_out.flush(); + + const QString drivePath = QString("\\\\.\\PhysicalDrive%0").arg(m_disk->index()); + HANDLE drive = CreateFile(drivePath.toStdWString().c_str(), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL); + if (drive == INVALID_HANDLE_VALUE) { + m_err << tr("Couldn't open the drive for data verification for %1:").arg(drivePath) << " (" << getLastError() << ")\n"; + m_err.flush(); + qApp->exit(1); + return false; + } + auto driveloseGuard = qScopeGuard([&drive] { + CloseHandle(drive); + }); switch (mediaCheckFD(_open_osfhandle(reinterpret_cast(drive), 0), &WriteJob::staticOnMediaCheckAdvanced, this)) { case ISOMD5SUM_CHECK_NOT_FOUND: case ISOMD5SUM_CHECK_PASSED: - out << "DONE\n"; - out.flush(); - err << "OK\n"; - err.flush(); - qApp->exit(0); - break; + m_diskManagement->logMessage(QtCriticalMsg, QStringLiteral("Check passed")); + m_out << "DONE\n"; + m_out.flush(); + m_err << "OK\n"; + m_err.flush(); + return true; case ISOMD5SUM_CHECK_FAILED: - err << tr("Your drive is probably damaged.") << "\n"; - err.flush(); - qApp->exit(1); + m_diskManagement->logMessage(QtCriticalMsg, QStringLiteral("Check failed")); + m_err << tr("Your drive is probably damaged.") << "\n"; + m_err.flush(); return false; default: - err << tr("Unexpected error occurred during media check.") << "\n"; - err.flush(); - qApp->exit(1); + m_diskManagement->logMessage(QtCriticalMsg, QStringLiteral("Check unexpected error")); + m_err << tr("Unexpected error occurred during media check.") << "\n"; + m_err.flush(); return false; } diff --git a/src/helper/win/writejob.h b/src/helper/win/writejob.h index a5aa222b..6be99701 100644 --- a/src/helper/win/writejob.h +++ b/src/helper/win/writejob.h @@ -1,5 +1,6 @@ /* * Fedora Media Writer + * Copyright (C) 2024 Jan Grulich * Copyright (C) 2016 Martin Bříza * * This program is free software; you can redistribute it and/or @@ -20,58 +21,47 @@ #ifndef WRITEJOB_H #define WRITEJOB_H -#include -#include -#include +#include #include +#include +#include #include #ifndef MEDIAWRITER_LZMA_LIMIT // 256MB memory limit for the decompressor -# define MEDIAWRITER_LZMA_LIMIT (1024*1024*256) +#define MEDIAWRITER_LZMA_LIMIT (1024 * 1024 * 256) #endif class WriteJob : public QObject { Q_OBJECT public: - explicit WriteJob(const QString &what, const QString &where); + explicit WriteJob(const QString &image, const QString &driveNumber, QObject *parent); static int staticOnMediaCheckAdvanced(void *data, long long offset, long long total); int onMediaCheckAdvanced(long long offset, long long total); private: - HANDLE openDrive(int physicalDriveNumber); - bool lockDrive(HANDLE drive); - bool removeMountPoints(uint diskNumber); - // bool dismountDrive(HANDLE drive, int diskNumber); - bool cleanDrive(uint diskNumber); - - bool writeBlock(HANDLE drive, OVERLAPPED *overlap, char *data, uint size); - - void unlockDrive(HANDLE drive); - + bool check(); + bool write(); + bool writeCompressed(HANDLE driveHandle); + bool writePlain(HANDLE driveHandle); private slots: - void work(); void onFileChanged(const QString &path); + void work(); - bool write(); - bool writeCompressed(HANDLE drive); - bool writePlain(HANDLE drive); - bool check(); private: - QString what; - uint where; - - QTextStream out { stdout }; - QTextStream err { stderr }; + QString m_image; + std::unique_ptr m_diskManagement; + std::unique_ptr m_disk; - QFileSystemWatcher watcher { }; + QTextStream m_out{stdout}; + QTextStream m_err{stderr}; - const int BLOCK_SIZE { 512 * 128 }; + QFileSystemWatcher m_watcher; }; #endif // WRITEJOB_H diff --git a/src/lib/CMakeLists.txt b/src/lib/CMakeLists.txt index 3a352a7a..ec6fca63 100644 --- a/src/lib/CMakeLists.txt +++ b/src/lib/CMakeLists.txt @@ -1 +1,4 @@ add_subdirectory(isomd5) +if (WIN32) + add_subdirectory(libwindisk) +endif() diff --git a/src/lib/libwindisk/CMakeLists.txt b/src/lib/libwindisk/CMakeLists.txt new file mode 100644 index 00000000..dd4e7efd --- /dev/null +++ b/src/lib/libwindisk/CMakeLists.txt @@ -0,0 +1,9 @@ +set(LIBWINDISK_SRCS + windisk.cpp +) + +add_library(libwindisk STATIC ${LIBWINDISK_SRCS}) + +target_link_libraries(libwindisk + Qt6::Core +) diff --git a/src/lib/libwindisk/windisk.cpp b/src/lib/libwindisk/windisk.cpp new file mode 100644 index 00000000..09caec91 --- /dev/null +++ b/src/lib/libwindisk/windisk.cpp @@ -0,0 +1,1061 @@ +/* + * Fedora Media Writer + * Copyright (C) 2024 Jan Grulich + * + * This program is free software; you can redistribute it and/or + * modify it under the terms of the GNU General Public License + * as published by the Free Software Foundation; either version 2 + * of the License, or (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU General Public License for more details. + * + * You should have received a copy of the GNU General Public License + * along with this program; if not, write to the Free Software + * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. + */ + +#include "windisk.h" + +#include + +#include +#include + +#include +#include + +#define INITGUID +#include +// Define the partition GUID for a basic data partition (for GPT) +DEFINE_GUID(PARTITION_BASIC_DATA_GUID, 0xebd0a0a2, 0xb9e5, 0x4433, 0x87, 0xc0, 0x68, 0xb6, 0xb7, 0x26, 0x99, 0xc7); + +#pragma comment(lib, "wbemuuid.lib") + +WinDiskManagement::WinDiskManagement(QObject *parent, bool isHelper) + : QObject(parent) +{ + if (isHelper) { + QString debugFileName = QStandardPaths::writableLocation(QStandardPaths::DocumentsLocation) + "/FedoraMediaWriter-helper.log"; + fopen_s(&m_debugFile, debugFileName.toStdString().c_str(), "w"); + } + + HRESULT res = S_OK; + // This needs to be initialized before any RPC communication occurs + // Currently when used in WinDriveManager we are good. + res = CoInitializeSecurity(NULL, -1, NULL, NULL, RPC_C_AUTHN_LEVEL_DEFAULT, RPC_C_IMP_LEVEL_IMPERSONATE, NULL, EOAC_NONE, 0); + if (FAILED(res)) { + _com_error err(res); + logMessage(QtWarningMsg, QStringLiteral("Failed to initialize security. Error = %1").arg(err.ErrorMessage())); + return; + } + + res = CoCreateInstance(CLSID_WbemAdministrativeLocator, NULL, CLSCTX_INPROC_SERVER, IID_IWbemLocator, reinterpret_cast(&m_IWbemLocator)); + if (FAILED(res)) { + _com_error err(res); + logMessage(QtWarningMsg, QStringLiteral("Failed to create IWbemLocator object. Error = %1").arg(QString::fromWCharArray(err.ErrorMessage()))); + return; + } + + res = m_IWbemLocator->ConnectServer(_bstr_t(L"ROOT\\Microsoft\\Windows\\Storage"), NULL, NULL, NULL, 0, NULL, NULL, &m_IWbemServices); + if (FAILED(res)) { + _com_error err(res); + logMessage(QtWarningMsg, QStringLiteral("Could not connect to WMI. Error = %1").arg(QString::fromWCharArray(err.ErrorMessage()))); + return; + } + + m_wmiInitialized = true; +} + +WinDiskManagement::~WinDiskManagement() +{ + if (m_IWbemLocator) { + m_IWbemLocator->Release(); + } + if (m_IWbemServices) { + m_IWbemServices->Release(); + } + CoUninitialize(); + + if (m_debugFile) { + fclose(m_debugFile); + } +} + +void WinDiskManagement::logMessage(QtMsgType type, const QString &msg) +{ + if (m_debugFile) { + QString txt; + switch (type) { + case QtDebugMsg: + txt = QString("WinDiskManagement[D]: %1").arg(msg); + break; + case QtInfoMsg: + txt = QString("WinDiskManagement[I]: %1").arg(msg); + break; + case QtWarningMsg: + txt = QString("WinDiskManagement[W]: %1").arg(msg); + break; + case QtCriticalMsg: + txt = QString("WinDiskManagement[C]: %1").arg(msg); + break; + case QtFatalMsg: + txt = QString("WinDiskManagement[F]: %1").arg(msg); + break; + } + fprintf(m_debugFile, "%s\n", txt.toStdString().c_str()); + fflush(m_debugFile); + return; + } + + switch (type) { + case QtDebugMsg: + qDebug() << "WinDiskManagement[D]: " << msg; + break; + case QtInfoMsg: + qInfo() << "WinDiskManagement[I]: " << msg; + break; + case QtWarningMsg: + qWarning() << "WinDiskManagement[W]: " << msg; + break; + case QtCriticalMsg: + qCritical() << "WinDiskManagement[C]: " << msg; + break; + case QtFatalMsg: + qFatal() << "WinDiskManagement[F]: " << msg; + break; + } +} + +QMap WinDiskManagement::getUSBDeviceList() +{ + logMessage(QtDebugMsg, QStringLiteral("Enumerating USB devices")); + + QMap result; + if (!m_wmiInitialized) { + logMessage(QtCriticalMsg, QStringLiteral("WMI interface is not initialized")); + return result; + } + + HRESULT res = S_OK; + IEnumWbemClassObject *pEnumDiskObjects = NULL; + + // BusType = 7 = USB | BusType = 12 = SD + res = m_IWbemServices->ExecQuery(_bstr_t(L"WQL"), _bstr_t(L"SELECT * FROM MSFT_Disk WHERE BusType = 7 OR BusType = 12"), WBEM_FLAG_RETURN_IMMEDIATELY, NULL, &pEnumDiskObjects); + if (FAILED(res)) { + _com_error err(res); + logMessage(QtCriticalMsg, QStringLiteral("WMI query failed. Error = %1").arg(QString::fromWCharArray(err.ErrorMessage()))); + return result; + } + + while (true) { + IWbemClassObject *pDiskObject = NULL; + ULONG uReturn = 0; + pEnumDiskObjects->Next(WBEM_INFINITE, 1, &pDiskObject, &uReturn); + if (uReturn == 0) { + break; + } + + quint32 index = 0; + QString deviceID; + VARIANT var; + + if ((pDiskObject->Get(_bstr_t(L"Number"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_I4) { + index = var.intVal; + } else if (var.vt == VT_UI4) { + index = var.uintVal; + } + VariantClear(&var); + } + + if ((pDiskObject->Get(_bstr_t(L"__PATH"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_BSTR) { + deviceID = QString::fromWCharArray(var.bstrVal); + } + VariantClear(&var); + } + pDiskObject->Release(); + + if (index > 0) { + logMessage(QtDebugMsg, QStringLiteral("Found device with index %1").arg(index)); + result.insert(index, deviceID); + } + } + pEnumDiskObjects->Release(); + + return result; +} + +QMap WinDiskManagement::getDevicePartitions(quint32 index) +{ + logMessage(QtDebugMsg, QStringLiteral("Enumerating partitions for device with index %1").arg(index)); + + QMap result; + if (!m_wmiInitialized) { + logMessage(QtCriticalMsg, QStringLiteral("WMI interface is not initialized")); + return result; + } + + HRESULT res = S_OK; + IEnumWbemClassObject *pPartitionObjects = NULL; + std::wstring partitionQuery = L"SELECT * FROM MSFT_Partition WHERE DiskNumber = " + std::to_wstring(index); + + res = m_IWbemServices->ExecQuery(_bstr_t(L"WQL"), _bstr_t(partitionQuery.c_str()), WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY, NULL, &pPartitionObjects); + if (FAILED(res)) { + _com_error err(res); + logMessage(QtCriticalMsg, QStringLiteral("Query for disk partitions failed. Error = %1").arg(QString::fromWCharArray(err.ErrorMessage()))); + return result; + } + + while (true) { + IWbemClassObject *pPartitionObject = NULL; + ULONG uReturnPartition = 0; + pPartitionObjects->Next(WBEM_INFINITE, 1, &pPartitionObject, &uReturnPartition); + if (uReturnPartition == 0) { + break; + } + + VARIANT var; + bool mountable = false; + qint32 partitionIndex = -1; + QString partitionPath; + + if ((pPartitionObject->Get(_bstr_t(L"PartitionNumber"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_I4) { + partitionIndex = var.intVal; + } else if (var.vt == VT_UI4) { + partitionIndex = var.uintVal; + } + VariantClear(&var); + } + + if ((pPartitionObject->Get(_bstr_t(L"__PATH"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_BSTR) { + partitionPath = QString::fromWCharArray(var.bstrVal); + } + VariantClear(&var); + } + + wchar_t driveLetter; + if ((pPartitionObject->Get(_bstr_t(L"DriveLetter"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_I2) { + driveLetter = static_cast(var.iVal); + mountable = driveLetter != L'\0'; + } + VariantClear(&var); + } + + pPartitionObject->Release(); + + if (partitionIndex != -1) { + if (mountable) { + logMessage(QtDebugMsg, QStringLiteral("Found partition with index %1 mounted to %2").arg(partitionIndex).arg(QString::fromWCharArray(&driveLetter, 1))); + } else { + logMessage(QtDebugMsg, QStringLiteral("Found unmounted partition with index %1").arg(partitionIndex)); + } + result.insert(static_cast(partitionIndex), mountable); + } + } + pPartitionObjects->Release(); + + return result; +} + +std::unique_ptr WinDiskManagement::getDiskDriveInformation(quint32 index, const QString &diskPath) +{ + logMessage(QtDebugMsg, QStringLiteral("Obtaining disk drive information for disk with index %1").arg(index)); + + std::unique_ptr result = std::make_unique(index, diskPath); + if (!m_wmiInitialized) { + logMessage(QtCriticalMsg, QStringLiteral("WMI interface is not initialized")); + return result; + } + + HRESULT res = S_OK; + IEnumWbemClassObject *pEnumDiskObjects = NULL; + + std::wstring deviceQuery = L"SELECT * FROM MSFT_Disk WHERE Number = " + std::to_wstring(index); + res = m_IWbemServices->ExecQuery(_bstr_t(L"WQL"), _bstr_t(deviceQuery.c_str()), WBEM_FLAG_RETURN_IMMEDIATELY, NULL, &pEnumDiskObjects); + if (FAILED(res)) { + _com_error err(res); + logMessage(QtCriticalMsg, QStringLiteral("WMI query failed. Error = %1").arg(QString::fromWCharArray(err.ErrorMessage()))); + return result; + } + + while (true) { + IWbemClassObject *pDiskObject = NULL; + ULONG uReturn = 0; + pEnumDiskObjects->Next(WBEM_INFINITE, 1, &pDiskObject, &uReturn); + if (uReturn == 0) { + break; + } + + VARIANT var; + if (result->path().isEmpty()) { + if ((pDiskObject->Get(_bstr_t(L"__PATH"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_BSTR) { + result->setPath(QString::fromWCharArray(var.bstrVal)); + logMessage(QtDebugMsg, QStringLiteral("DeviceID %1").arg(result->path())); + } + VariantClear(&var); + } + } + + if ((pDiskObject->Get(_bstr_t(L"IsOffline"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_BOOL) { + result->setIsOffline(var.boolVal); + logMessage(QtDebugMsg, QStringLiteral("Disk is offline: %1").arg(result->isOffline())); + } + VariantClear(&var); + } + + if ((pDiskObject->Get(_bstr_t(L"FriendlyName"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_BSTR) { + result->setName(QString::fromWCharArray(var.bstrVal)); + logMessage(QtDebugMsg, QStringLiteral("Disk name: %1").arg(result->name())); + } + VariantClear(&var); + } + + if ((pDiskObject->Get(_bstr_t(L"Size"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_BSTR) { + result->setSize(QString::fromWCharArray(var.bstrVal).toULongLong()); + logMessage(QtDebugMsg, QStringLiteral("Size %1").arg(result->size())); + } else if (var.vt == VT_I4) { + result->setSize(var.intVal); + logMessage(QtDebugMsg, QStringLiteral("Size %1").arg(result->size())); + } else if (var.vt == VT_UI4) { + result->setSize(var.uintVal); + logMessage(QtDebugMsg, QStringLiteral("Size %1").arg(result->size())); + } + VariantClear(&var); + } + + if ((pDiskObject->Get(_bstr_t(L"PhysicalSectorSize"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_BSTR) { + result->setSectorSize(QString::fromWCharArray(var.bstrVal).toULongLong()); + logMessage(QtDebugMsg, QStringLiteral("Sector size %1").arg(result->sectorSize())); + } else if (var.vt == VT_I4) { + result->setSectorSize(var.intVal); + logMessage(QtDebugMsg, QStringLiteral("Sector size %1").arg(result->sectorSize())); + } else if (var.vt == VT_UI4) { + result->setSectorSize(var.uintVal); + logMessage(QtDebugMsg, QStringLiteral("Sector size %1").arg(result->sectorSize())); + } + VariantClear(&var); + } + + if ((pDiskObject->Get(_bstr_t(L"SerialNumber"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_BSTR) { + result->setSerialNumber(QString::fromWCharArray(var.bstrVal)); + logMessage(QtDebugMsg, QStringLiteral("Serial number %1").arg(result->serialNumber())); + } + VariantClear(&var); + } + pDiskObject->Release(); + } + pEnumDiskObjects->Release(); + + return result; +} + +bool WinDiskManagement::clearPartitions(qint32 index) +{ + logMessage(QtDebugMsg, QStringLiteral("Removing partitions on disk with index %1").arg(index)); + + if (!m_wmiInitialized) { + logMessage(QtWarningMsg, QStringLiteral("WMI interface is not initialized")); + return false; + } + + HRESULT res = S_OK; + + IEnumWbemClassObject *pPartitionObjects = NULL; + std::wstring partitionQuery = L"SELECT * FROM MSFT_Partition WHERE DiskNumber = " + std::to_wstring(index); + + res = m_IWbemServices->ExecQuery(_bstr_t(L"WQL"), _bstr_t(partitionQuery.c_str()), WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY, NULL, &pPartitionObjects); + if (FAILED(res)) { + _com_error err(res); + logMessage(QtCriticalMsg, QStringLiteral("Query for disk partitions failed. Error = %1").arg(QString::fromWCharArray(err.ErrorMessage()))); + return false; + } + + while (true) { + IWbemClassObject *pPartitionObject = NULL; + ULONG uReturnPartition = 0; + pPartitionObjects->Next(WBEM_INFINITE, 1, &pPartitionObject, &uReturnPartition); + if (uReturnPartition == 0) { + break; + } + VARIANT var; + res = pPartitionObject->Get(L"__PATH", 0, &var, NULL, NULL); + if (SUCCEEDED(res)) { + QString partitionPath = QString::fromWCharArray(var.bstrVal); + IWbemClassObject *pOutParams = NULL; + res = m_IWbemServices->ExecMethod(_bstr_t(partitionPath.toStdWString().c_str()), _bstr_t(L"DeleteObject"), 0, NULL, NULL, &pOutParams, NULL); + if (FAILED(res)) { + _com_error err(res); + logMessage(QtCriticalMsg, QStringLiteral("Failed to delete partition. Error = %1").arg(QString::fromWCharArray(err.ErrorMessage()))); + return false; + } + + if (pOutParams) { + VARIANT returnValueVar; + VariantInit(&returnValueVar); + pOutParams->Get(L"ReturnValue", 0, &returnValueVar, NULL, NULL); + if (returnValueVar.vt == VT_I4 && returnValueVar.intVal != 0) { + logMessage(QtCriticalMsg, QStringLiteral("Failed to delete partition. Error code: %1").arg(QString::number(returnValueVar.intVal, 16))); + VariantClear(&returnValueVar); + return false; + } + pOutParams->Release(); + } + } + pPartitionObject->Release(); + } + pPartitionObjects->Release(); + + logMessage(QtDebugMsg, QStringLiteral("Partitions deleted successfully.")); + + return true; +} + +bool WinDiskManagement::formatPartition(const QChar &driveLetter) +{ + logMessage(QtDebugMsg, QStringLiteral("Formatting partition mounted to drive letter %1:").arg(driveLetter)); + + if (!m_wmiInitialized) { + logMessage(QtCriticalMsg, QStringLiteral("WMI interface is not initialized")); + return false; + } + + HRESULT res = S_OK; + IEnumWbemClassObject *pVolumeObjects = NULL; + std::wstring query = L"SELECT * FROM MSFT_Volume WHERE DriveLetter='"; + query.push_back(driveLetter.toUpper().unicode()); + query.push_back(L'\''); + + res = m_IWbemServices->ExecQuery(_bstr_t(L"WQL"), _bstr_t(query.c_str()), WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY, NULL, &pVolumeObjects); + if (FAILED(res)) { + _com_error err(res); + logMessage(QtCriticalMsg, QStringLiteral("Query for MSFT_Volume failed. Error = %1").arg(QString::fromWCharArray(err.ErrorMessage()))); + return false; + } + + bool volumeFound = false; + while (true) { + IWbemClassObject *pVolumeObject = NULL; + ULONG uReturnVolume = 0; + pVolumeObjects->Next(WBEM_INFINITE, 1, &pVolumeObject, &uReturnVolume); + if (uReturnVolume == 0) { + logMessage(QtWarningMsg, QStringLiteral("No volume object")); + break; + } + + QString volumePath; + VARIANT volumePathVar; + VariantInit(&volumePathVar); + res = pVolumeObject->Get(L"__PATH", 0, &volumePathVar, NULL, NULL); + if (SUCCEEDED(res)) { + volumePath = QString::fromWCharArray(volumePathVar.bstrVal); + VariantClear(&volumePathVar); + } + pVolumeObject->Release(); + + IWbemClassObject *pClass = NULL; + IWbemClassObject *pInParamsDefinition = NULL; + IWbemClassObject *pInParams = NULL; + IWbemClassObject *pOutParams = NULL; + + auto cleanup = qScopeGuard([=] { + if (pOutParams) { + pOutParams->Release(); + } + if (pInParams) { + pInParams->Release(); + } + if (pInParamsDefinition) { + pInParamsDefinition->Release(); + } + if (pClass) { + pClass->Release(); + } + }); + + res = m_IWbemServices->GetObject(_bstr_t(L"MSFT_Volume"), 0, NULL, &pClass, NULL); + if (FAILED(res)) { + _com_error err(res); + logMessage(QtCriticalMsg, QStringLiteral("WMI query to get MSFT_Volume object failed. Error = %1").arg(QString::fromWCharArray(err.ErrorMessage()))); + return false; + } + + res = pClass->GetMethod(_bstr_t(L"Format"), 0, &pInParamsDefinition, NULL); + if (FAILED(res)) { + _com_error err(res); + logMessage(QtCriticalMsg, QStringLiteral("WMI query to get 'Format' method on MSFT_Volume object failed. Error = %1").arg(QString::fromWCharArray(err.ErrorMessage()))); + return false; + } + + res = pInParamsDefinition->SpawnInstance(0, &pInParams); + if (FAILED(res)) { + _com_error err(res); + logMessage(QtCriticalMsg, QStringLiteral("WMI spawn instance failed. Error = %1").arg(QString::fromWCharArray(err.ErrorMessage()))); + return false; + } + + VARIANT var; + VariantInit(&var); + var.vt = VT_BSTR; + var.bstrVal = _bstr_t(L"exFAT"); + res = pInParams->Put(L"FileSystem", 0, &var, 0); + VariantClear(&var); + + if (FAILED(res)) { + _com_error err(res); + logMessage(QtCriticalMsg, QStringLiteral("Failed to set 'FileSystem' parameter. Error = %1").arg(QString::fromWCharArray(err.ErrorMessage()))); + return false; + } + + var.vt = VT_BOOL; + var.boolVal = VARIANT_FALSE; // Quick format + res = pInParams->Put(L"Full", 0, &var, 0); + VariantClear(&var); + + if (FAILED(res)) { + logMessage(QtCriticalMsg, QStringLiteral("Failed to set 'Full' parameter.")); + return false; + } + + volumeFound = true; + res = m_IWbemServices->ExecMethod(_bstr_t(volumePath.toStdString().c_str()), _bstr_t(L"Format"), 0, NULL, pInParams, &pOutParams, NULL); + if (FAILED(res)) { + _com_error err(res); + logMessage(QtCriticalMsg, QStringLiteral("Failed to format the volume. Error = %1").arg(QString::fromWCharArray(err.ErrorMessage()))); + } + + if (pOutParams) { + VARIANT returnValueVar; + VariantInit(&returnValueVar); + pOutParams->Get(L"ReturnValue", 0, &returnValueVar, NULL, NULL); + if (returnValueVar.vt == VT_I4 && returnValueVar.intVal != 0) { + logMessage(QtCriticalMsg, QStringLiteral("Failed to format the volume. Error code: %1").arg(QString::number(returnValueVar.intVal, 16))); + VariantClear(&returnValueVar); + return false; + } else { + logMessage(QtDebugMsg, QStringLiteral("Volume successfully formatted to exFat.")); + volumeFound = true; + } + pOutParams->Release(); + } + + if (volumeFound) { + break; + } + } + pVolumeObjects->Release(); + + if (!volumeFound) { + logMessage(QtWarningMsg, QStringLiteral("No volumes found for the partition.")); + return false; + } + + return true; +} + +bool WinDiskManagement::refreshDiskDrive(const QString &diskPath) +{ + logMessage(QtDebugMsg, QStringLiteral("Refreshing disk drive information")); + + if (!m_wmiInitialized) { + logMessage(QtCriticalMsg, QStringLiteral("WMI interface is not initialized")); + return false; + } + + HRESULT res = S_OK; + IWbemClassObject *pOutParams = NULL; + + res = m_IWbemServices->ExecMethod(_bstr_t(diskPath.toStdWString().c_str()), _bstr_t(L"Refresh"), 0, NULL, NULL, &pOutParams, NULL); + + if (pOutParams) { + pOutParams->Release(); + } + + if (FAILED(res)) { + _com_error err(res); + logMessage(QtCriticalMsg, QStringLiteral("Failed to refresh the disk. Error = %1").arg(QString::fromWCharArray(err.ErrorMessage()))); + return false; + } + + logMessage(QtDebugMsg, QStringLiteral("Successfully refreshed the disk.")); + return true; +} + +static QString getLastError() +{ + TCHAR message[256]; + FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); + return QString::fromWCharArray(message).trimmed(); +} + +bool WinDiskManagement::lockDrive(HANDLE driveHandle, int numRetries) +{ + logMessage(QtDebugMsg, QStringLiteral("Trying to lock the drive")); + + int attempts = 0; + DWORD status; + + while (true) { + if (!DeviceIoControl(driveHandle, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &status, NULL)) { + attempts++; + } else { + logMessage(QtDebugMsg, QStringLiteral("Successfully locked the drive")); + return true; + } + + if (attempts == numRetries) { + logMessage(QtWarningMsg, QStringLiteral("Couldn't lock the drive: %1").arg(getLastError())); + break; + } + + QThread::sleep(2); + } + + return false; +} + +bool WinDiskManagement::unlockDrive(HANDLE driveHandle) +{ + logMessage(QtDebugMsg, QStringLiteral("Trying to unlock the drive")); + + bool ret = DeviceIoControl(driveHandle, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, NULL, NULL); + if (!ret) { + logMessage(QtFatalMsg, QStringLiteral("Failed to unlock the drive")); + return false; + } + + logMessage(QtDebugMsg, QStringLiteral("Successfully unlocked the drive")); + return true; +} + +bool WinDiskManagement::disableIOBoundaryChecks(HANDLE driveHandle) +{ + logMessage(QtDebugMsg, QStringLiteral("Trying to disable I/O boundary checks")); + + bool ret = DeviceIoControl(driveHandle, FSCTL_ALLOW_EXTENDED_DASD_IO, NULL, 0, NULL, 0, NULL, NULL); + if (!ret) { + logMessage(QtFatalMsg, QStringLiteral("Failed to disable I/O boundary checks")); + return false; + } + + logMessage(QtDebugMsg, QStringLiteral("Successfully disabled I/O boundary checks")); + return true; +} + +bool WinDiskManagement::removeDriveLetters(quint32 index) +{ + logMessage(QtDebugMsg, QStringLiteral("Removing assigned drive letters for device with index %1").arg(index)); + + DWORD drives = ::GetLogicalDrives(); + + for (char i = 0; i < 26; i++) { + if (drives & (1 << i)) { + char currentDrive = 'A' + i; + QString drivePath = QString("\\\\.\\%1:").arg(currentDrive); + logMessage(QtDebugMsg, QStringLiteral("Checking drive: %1").arg(drivePath)); + HANDLE device = ::CreateFile(drivePath.toStdWString().c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL); + if (device == INVALID_HANDLE_VALUE) { + logMessage(QtWarningMsg, QStringLiteral("Failed to open logical drive: %1").arg(currentDrive)); + continue; + } + + auto cleanup = qScopeGuard([device] { + CloseHandle(device); + }); + + DWORD bytesReturned; + VOLUME_DISK_EXTENTS vde; + if (DeviceIoControl(device, IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS, NULL, 0, &vde, sizeof(vde), &bytesReturned, NULL)) { + for (uint j = 0; j < vde.NumberOfDiskExtents; j++) { + if (vde.Extents[j].DiskNumber == index) { + QString volumePath = QString("%1:\\").arg(currentDrive); + logMessage(QtWarningMsg, QStringLiteral("Checking volume: %1").arg(volumePath)); + if (!DeleteVolumeMountPointA(volumePath.toStdString().c_str())) { + logMessage(QtCriticalMsg, QStringLiteral("Couldn't remove the drive: %1").arg(getLastError())); + return false; + } + logMessage(QtDebugMsg, QStringLiteral("Successfully removed mountpoints for volume: %1").arg(currentDrive)); + break; + } + } + } + } + } + + return true; +} + +bool WinDiskManagement::unmountVolume(HANDLE logicalHandle) +{ + logMessage(QtDebugMsg, QStringLiteral("Unmounting logical volume")); + + if (!DeviceIoControl(logicalHandle, FSCTL_DISMOUNT_VOLUME, NULL, 0, NULL, 0, NULL, NULL)) { + logMessage(QtCriticalMsg, QStringLiteral("Couldn't unmount drive: %1").arg(getLastError())); + return false; + } + + logMessage(QtDebugMsg, QStringLiteral("Successfully unmounted logical volume")); + + return true; +} + +QChar WinDiskManagement::mountVolume(const QString &logicalName) +{ + logMessage(QtDebugMsg, QStringLiteral("Mounting logical volume")); + + char mountedLetter[27] = {0}; + DWORD size; + if (::GetVolumePathNamesForVolumeNameA(logicalName.toStdString().c_str(), mountedLetter, sizeof(mountedLetter), &size) && (size > 1)) { + logMessage(QtDebugMsg, QStringLiteral("Volume is already mounted")); + return QChar(mountedLetter[0]); + } + + char drives[256]; + DWORD driveSize = GetLogicalDriveStringsA(sizeof(drives), drives); + if (!driveSize || driveSize > sizeof(drives)) { + logMessage(QtCriticalMsg, QStringLiteral("Failed to get drive letter mountpoint: %1").arg(getLastError())); + return QChar(); + } + + char driveLetter = 0; + for (char letter = 'C'; letter <= 'Z'; letter++) { + bool isDriveUsed = std::any_of(drives, drives + driveSize, [letter](char drive) { + return toupper(drive) == letter; + }); + + if (!isDriveUsed) { + driveLetter = letter; + break; + } + } + + if (driveLetter == 0) { + logMessage(QtWarningMsg, QStringLiteral("Couldn't find available drive letter for mountpoint.")); + return QChar(); + } + + std::string drivePath = std::string(1, driveLetter) + ":\\"; + if (!::SetVolumeMountPointA(drivePath.c_str(), logicalName.toStdString().c_str())) { + logMessage(QtCriticalMsg, QStringLiteral("Couldn't mount %1 as %2: %3").arg(logicalName).arg(driveLetter).arg(getLastError())); + return QChar(); + } + + logMessage(QtDebugMsg, QStringLiteral("Successfuly mounted logical volume %1 as %2").arg(logicalName).arg(driveLetter)); + return QChar(driveLetter); +} + +bool WinDiskManagement::writeFileWithRetry(HANDLE driveHandle, void *buffer, qint64 numberOfBytesToWrite, int numberOfRetries) +{ + bool success = false; + bool readFilePointer = false; + LARGE_INTEGER filePointer; + LARGE_INTEGER filePointerZero = {{0, 0}}; + DWORD writtenBytes; + + readFilePointer = SetFilePointerEx(driveHandle, filePointerZero, &filePointer, FILE_CURRENT); + if (!readFilePointer) { + logMessage(QtCriticalMsg, QStringLiteral("Could not read file pointer: %1").arg(getLastError())); + } + + for (int attempts = 1; attempts <= numberOfRetries; attempts++) { + if ((attempts > 1) && (!SetFilePointerEx(driveHandle, filePointer, NULL, FILE_BEGIN))) { + logMessage(QtCriticalMsg, QStringLiteral("Could not set file pointer")); + break; + } + if (WriteFile(driveHandle, buffer, numberOfBytesToWrite, &writtenBytes, NULL)) { + if (numberOfBytesToWrite == writtenBytes) { + return true; + } + } else { + logMessage(QtCriticalMsg, QStringLiteral("Failed to write data to the drive: %1").arg(getLastError())); + success = false; + } + + if (!readFilePointer) { + break; + } + + if (attempts < numberOfRetries) { + QThread::sleep(5); + } + } + + return success; +} + +bool WinDiskManagement::clearPartitionTable(HANDLE driveHandle, quint64 driveSize, quint32 sectorSize) +{ + logMessage(QtDebugMsg, QStringLiteral("Clearing partition table information")); + + quint64 sectorsToClear = 128; + LARGE_INTEGER filePointer; + + uint8_t *zeroBuffer = static_cast(calloc(sectorSize, sectorsToClear)); + if (!zeroBuffer) { + logMessage(QtCriticalMsg, QStringLiteral("Couldn't allocate zero buffer")); + return false; + } + + filePointer.QuadPart = 0ULL; + if (!SetFilePointerEx(driveHandle, filePointer, &filePointer, FILE_BEGIN) || (filePointer.QuadPart != 0ULL)) { + logMessage(QtCriticalMsg, QStringLiteral("Couldn't reset disk position: %1").arg(getLastError())); + } + + if (!writeFileWithRetry(driveHandle, zeroBuffer, sectorSize * sectorsToClear, 4)) { + logMessage(QtCriticalMsg, QStringLiteral("Couldn't write zero data to the drive: %1").arg(getLastError())); + } + + filePointer.QuadPart = driveSize - (LONGLONG)sectorSize * sectorsToClear; + if (SetFilePointerEx(driveHandle, filePointer, &filePointer, FILE_BEGIN)) { + if (!writeFileWithRetry(driveHandle, zeroBuffer, sectorSize * sectorsToClear, 4)) { + logMessage(QtCriticalMsg, QStringLiteral("Couldn't write zero data to the drive: %1").arg(getLastError())); + } + } + + free(zeroBuffer); + + if (!refreshPartitionLayout(driveHandle)) { + logMessage(QtWarningMsg, QStringLiteral("Couldn't update drive properties")); + } + + logMessage(QtDebugMsg, QStringLiteral("Successfully cleared the partition table")); + return true; +} + +bool WinDiskManagement::clearDiskDrive(HANDLE driveHandle) +{ + logMessage(QtDebugMsg, QStringLiteral("Clearing disk drive")); + + BOOL ret; + CREATE_DISK createDisk = {PARTITION_STYLE_RAW, {{0}}}; + + ret = DeviceIoControl(driveHandle, IOCTL_DISK_CREATE_DISK, &createDisk, sizeof(createDisk), NULL, 0, NULL, NULL); + if (!ret) { + logMessage(QtCriticalMsg, QStringLiteral("Couldn't delete drive layout")); + return false; + } + + if (!refreshPartitionLayout(driveHandle)) { + logMessage(QtWarningMsg, QStringLiteral("Couldn't update drive properties")); + } + + logMessage(QtDebugMsg, QStringLiteral("Successfully cleared the disk drive")); + return true; +} + +bool WinDiskManagement::createGPTPartition(HANDLE driveHandle, quint64 diskSize, quint32 sectorSize) +{ + logMessage(QtDebugMsg, QStringLiteral("Creating GPT partition table")); + + BOOL ret; + + CREATE_DISK createDisk = {PARTITION_STYLE_GPT, {{0}}}; + CoCreateGuid(&createDisk.Gpt.DiskId); + + ret = DeviceIoControl(driveHandle, IOCTL_DISK_CREATE_DISK, &createDisk, sizeof(createDisk), NULL, 0, NULL, NULL); + if (!ret) { + logMessage(QtCriticalMsg, QStringLiteral("Failed to create GPT partition table on the disk.")); + return false; + } + + ret = DeviceIoControl(driveHandle, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, NULL, NULL); + if (!ret) { + logMessage(QtCriticalMsg, QStringLiteral("Couldn't update disk properties.")); + return false; + } + + DRIVE_LAYOUT_INFORMATION_EX driveLayout = {0}; + driveLayout.PartitionStyle = PARTITION_STYLE_GPT; + driveLayout.PartitionCount = 1; + + PARTITION_INFORMATION_EX &partitionInfo = driveLayout.PartitionEntry[0]; + partitionInfo.PartitionStyle = PARTITION_STYLE_GPT; + // GPT starts at sector 34 (after the GPT header) + partitionInfo.StartingOffset.QuadPart = 34 * sectorSize; + // Disk size - GPT header/footer sectors ; + partitionInfo.PartitionLength.QuadPart = diskSize - ((34 + 33) * sectorSize); + partitionInfo.Gpt.PartitionType = PARTITION_BASIC_DATA_GUID; + partitionInfo.Gpt.PartitionId = createDisk.Gpt.DiskId; + + ret = DeviceIoControl(driveHandle, IOCTL_DISK_SET_DRIVE_LAYOUT_EX, &driveLayout, sizeof(driveLayout), NULL, 0, NULL, NULL); + if (!ret) { + logMessage(QtCriticalMsg, QStringLiteral("Failed to set GPT partition layout on the disk.")); + return false; + } + + if (!refreshPartitionLayout(driveHandle)) { + logMessage(QtWarningMsg, QStringLiteral("Couldn't update drive properties")); + } + + logMessage(QtDebugMsg, QStringLiteral("Successfully created GPT partition over the whole disk.")); + return true; +} + +QString WinDiskManagement::getLogicalName(quint32 index, bool keepTrailingBackslash) +{ + QString result; + static const char *volumeStart = "\\\\?\\"; + char volumeName[2048]; + char path[2048]; + HANDLE drive = INVALID_HANDLE_VALUE; + HANDLE volume = INVALID_HANDLE_VALUE; + + for (int i = 0; drive == INVALID_HANDLE_VALUE; i++) { + if (i == 0) { + volume = FindFirstVolumeA(volumeName, sizeof(volumeName)); + if (volume == INVALID_HANDLE_VALUE) { + logMessage(QtWarningMsg, QStringLiteral("Couldn't access GUID volume: %1").arg(getLastError())); + continue; + } + } else { + if (!FindNextVolumeA(volume, volumeName, sizeof(volumeName))) { + logMessage(QtDebugMsg, QStringLiteral("Couldn't access next GUID volume: %1").arg(getLastError())); + break; + } + } + + size_t len = strnlen_s(volumeName, 2048); + if (len <= 4 || _strnicmp(volumeName, volumeStart, 4) != 0 || volumeName[len - 1] != '\\') { + logMessage(QtWarningMsg, QStringLiteral("Obtained wrong volume name: %1").arg(volumeName)); + continue; + } + + volumeName[len - 1] = 0; + if (QueryDosDeviceA(&volumeName[4], path, sizeof(path)) == 0) { + logMessage(QtWarningMsg, QStringLiteral("Failed to get device path for GUID volume %1: %2").arg(volumeName).arg(getLastError())); + continue; + } + drive = ::CreateFileA(volumeName, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); + if (drive == INVALID_HANDLE_VALUE) { + logMessage(QtWarningMsg, QStringLiteral("Couldn't open GUID volume %1: %2").arg(volumeName).arg(getLastError())); + continue; + } + + DWORD size = 0; + VOLUME_DISK_EXTENTS diskExtents; + BOOL ret = DeviceIoControl(drive, IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS, NULL, 0, &diskExtents, sizeof(diskExtents), &size, NULL); + CloseHandle(drive); + drive = INVALID_HANDLE_VALUE; + if (!ret || size == 0) { + logMessage(QtWarningMsg, QStringLiteral("Couldn't open GUID volume %1: %2").arg(volumeName).arg(getLastError())); + continue; + } + + if (diskExtents.NumberOfDiskExtents == 0 || diskExtents.NumberOfDiskExtents != 1) { + logMessage(QtWarningMsg, QStringLiteral("Wrong number of disk extents.")); + continue; + } + + if (diskExtents.Extents[0].DiskNumber != index) { + continue; + } + + if (keepTrailingBackslash) { + volumeName[len - 1] = '\\'; + } + result = QString(volumeName); + break; + } + + if (result.isEmpty()) { + logMessage(QtDebugMsg, QStringLiteral("No logical volume found. Device doesn't have any partition")); + } + + return result; +} + +bool WinDiskManagement::refreshPartitionLayout(HANDLE driveHandle) +{ + logMessage(QtDebugMsg, QStringLiteral("Refreshing information about partition layout")); + + BOOL ret = DeviceIoControl(driveHandle, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, NULL, NULL); + if (!ret) { + logMessage(QtCriticalMsg, QStringLiteral("Couldn't update disk properties.")); + return false; + } + + return true; +} + +WinDisk::WinDisk(quint32 index, const QString &path) + : m_index(index) + , m_path(path) +{ +} + +quint32 WinDisk::index() const +{ + return m_index; +} + +bool WinDisk::isOffline() const +{ + return m_isOffline; +} + +void WinDisk::setIsOffline(bool isOffline) +{ + m_isOffline = isOffline; +} + +QString WinDisk::path() const +{ + return m_path; +} + +void WinDisk::setPath(const QString &path) +{ + m_path = path; +} + +QString WinDisk::name() const +{ + return m_name; +} + +void WinDisk::setName(const QString &name) +{ + m_name = name; +} + +quint64 WinDisk::size() const +{ + return m_size; +} + +void WinDisk::setSize(quint64 size) +{ + m_size = size; +} + +QString WinDisk::serialNumber() const +{ + return m_serialNumber; +} + +void WinDisk::setSerialNumber(const QString &serialNumber) +{ + m_serialNumber = serialNumber; +} + +quint32 WinDisk::sectorSize() +{ + return m_sectorSize; +} + +void WinDisk::setSectorSize(quint32 sectorSize) +{ + m_sectorSize = sectorSize; +} diff --git a/src/lib/libwindisk/windisk.h b/src/lib/libwindisk/windisk.h new file mode 100644 index 00000000..9eb2776a --- /dev/null +++ b/src/lib/libwindisk/windisk.h @@ -0,0 +1,122 @@ +/* + * Fedora Media Writer + * Copyright (C) 2024 Jan Grulich + * + * This program is free software; you can redistribute it and/or + * modify it under the terms of the GNU General Public License + * as published by the Free Software Foundation; either version 2 + * of the License, or (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU General Public License for more details. + * + * You should have received a copy of the GNU General Public License + * along with this program; if not, write to the Free Software + * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. + */ + +#ifndef WINDISKMANAGEMENT_H +#define WINDISKMANAGEMENT_H + +#include +#include + +class WinDisk; + +class WinDiskManagement : public QObject +{ + Q_OBJECT +public: + WinDiskManagement(QObject *parent, bool isHelper = false); + ~WinDiskManagement(); + + void logMessage(QtMsgType type, const QString &msg); + + /* + * WMI - Windows Management Instrumentation + */ + // Returns a map with list of devices + QMap getUSBDeviceList(); + // Returns a map of partitions for device with given @index + QMap getDevicePartitions(quint32 index); + // Returns information about disk drive on given @index + std::unique_ptr getDiskDriveInformation(quint32 index, const QString &diskPath = QString()); + // Remove all partitions on device on given @index + bool clearPartitions(qint32 index); + // Formats partition to exFAT on given @partitionPath + bool formatPartition(const QChar &driveLetter); + // Refreshes disk drive on given @diskPath + bool refreshDiskDrive(const QString &diskPath); + + /* + * WinAPI + */ + // Locks the drive and try @numRetries attempts if we fail + bool lockDrive(HANDLE driveHandle, int numRetries = 1); + // Unlocks the drive + bool unlockDrive(HANDLE driveHandle); + // Try to disable I/O boundary checks + bool disableIOBoundaryChecks(HANDLE driveHandle); + // Remove all assigned drive letters + bool removeDriveLetters(quint32 index); + // Unmount volume provided by logical handle + bool unmountVolume(HANDLE volumeHandle); + // Mount volume provided by GUID path and return drive letter it's mounted to + QChar mountVolume(const QString &volume); + // Clears GPT/MBR records after writing ISO image + bool clearPartitionTable(HANDLE driveHandle, quint64 driveSize, quint32 sectorSize); + // Clears the drive and sets it to RAW state + bool clearDiskDrive(HANDLE driveHandle); + // Creates a GPT partition table on the drive provided by @driveHandle + bool createGPTPartition(HANDLE driveHandle, quint64 diskSize, quint32 sectorSize); + // Returns the GUID volume name + QString getLogicalName(quint32 index, bool keepTrailingBackslash = true); + // Refreshes the partition layout + bool refreshPartitionLayout(HANDLE driveHandle); + bool writeFileWithRetry(HANDLE driveHandle, void *buffer, qint64 numberOfBytesToWrite, int numberOfRetries = 1); + +private: + bool m_wmiInitialized = false; + IWbemLocator *m_IWbemLocator = NULL; + IWbemServices *m_IWbemServices = NULL; + FILE *m_debugFile = nullptr; +}; + +class WinDisk +{ +public: + WinDisk(quint32 index, const QString &path = QString()); + + quint32 index() const; + + bool isOffline() const; + void setIsOffline(bool offline); + + QString path() const; + void setPath(const QString &path); + + QString name() const; + void setName(const QString &name); + + quint64 size() const; + void setSize(quint64 size); + + QString serialNumber() const; + void setSerialNumber(const QString &serialNumber); + + quint32 sectorSize(); + void setSectorSize(quint32 sectorSize); + +private: + bool m_isOffline = false; + quint32 m_index = 0; + quint32 m_sectorSize = 0; + quint64 m_size = 0; + QString m_path; + QString m_name; + QString m_serialNumber; +}; + +#endif // WINDISKMANAGEMENT_H