From 096df29c2fb41b4f7c53e693c47ad9454b377f60 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 | 342 +++----- src/app/windrivemanager.h | 12 +- src/helper/win/CMakeLists.txt | 1 + src/helper/win/main.cpp | 24 +- src/helper/win/restorejob.cpp | 168 +++- src/helper/win/restorejob.h | 13 +- src/helper/win/writejob.cpp | 538 ++++++------ src/helper/win/writejob.h | 44 +- src/lib/CMakeLists.txt | 3 + src/lib/libwindisk/CMakeLists.txt | 9 + src/lib/libwindisk/windisk.cpp | 1014 ++++++++++++++++++++++ src/lib/libwindisk/windisk.h | 117 +++ 16 files changed, 1739 insertions(+), 553 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..c553b1f8 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,262 +22,127 @@ #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"; + qApp->installNativeEventFilter(this); QTimer::singleShot(0, this, &WinDriveProvider::checkDrives); } -void WinDriveProvider::checkDrives() +bool WinDriveProvider::nativeEventFilter(const QByteArray &eventType, void *message, qintptr *result) { - static bool firstRun = true; - - if (firstRun) - mDebug() << this->metaObject()->className() << "Looking for the drives for the first time"; - - for (int i = 0; i < 64; i++) { - bool present = describeDrive(i, firstRun); - if (!present && m_drives.contains(i)) { - emit driveRemoved(m_drives[i]); - m_drives[i]->deleteLater(); - m_drives.remove(i); + Q_UNUSED(eventType); + + MSG *msg = static_cast(message); + if ((msg->message == WM_DEVICECHANGE) && ((msg->wParam == DBT_DEVICEARRIVAL) || (msg->wParam == DBT_DEVICEREMOVECOMPLETE))) { + mDebug() << "Recieved device change event"; + *result = TRUE; + for (const WinDrive *drive : m_drives) { + // Ignore device change events when we are restoring or writting and schedule + // re-check once we are done + if (drive->busy()) { + return true; + } } + // REMOVE + QTimer::singleShot(0, this, &WinDriveProvider::checkDrives); + return true; } - - if (firstRun) - mDebug() << this->metaObject()->className() << "Finished looking for the drives for the first time"; - firstRun = false; - QTimer::singleShot(2500, this, &WinDriveProvider::checkDrives); + return false; } -QString getPhysicalName(int driveNumber) +bool zeroDrive(const std::wstring &drivePath) { - 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; + HANDLE hDrive = CreateFileW(drivePath.c_str(), // Drive path (e.g., \\.\PhysicalDrive1) + GENERIC_WRITE, // Write access + FILE_SHARE_READ | FILE_SHARE_WRITE, // Share mode + NULL, // Default security + OPEN_EXISTING, // Open the existing drive + FILE_FLAG_NO_BUFFERING | FILE_FLAG_WRITE_THROUGH, // Unbuffered write, bypass cache + NULL // No template file + ); + + if (hDrive == INVALID_HANDLE_VALUE) { + mDebug() << "Failed to open drive: " << GetLastError(); } - 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; - } + BYTE buffer[4096] = {0}; // Buffer filled with zeroes + DWORD bytesWritten = 0; - 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; - } - } - } - 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; + // Write zeroes to the entire drive + while (WriteFile(hDrive, buffer, sizeof(buffer), &bytesWritten, NULL) && bytesWritten > 0) { + // Keep writing zeroes until the entire disk is zeroed out } - mDebug() << this->metaObject()->className() << getPhysicalName(driveNumber) << " is mountable"; - - CloseHandle(physicalHandle); + CloseHandle(hDrive); + mDebug() << "Drive zeroed successfully."; return true; } -bool WinDriveProvider::describeDrive(int nDriveNumber, bool verbose) +void WinDriveProvider::checkDrives() { - 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; - } + mDebug() << this->metaObject()->className() << "Looking for the drives"; + + 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()); + for (auto it = partitionList.constBegin(); it != partitionList.constEnd(); it++) { + if (!it.value()) { + mountable = false; + break; + } + } - // 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; } WinDrive::WinDrive(WinDriveProvider *parent, const QString &name, uint64_t size, bool containsLive, int device, const QString &serialNumber) @@ -290,7 +155,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 +204,7 @@ void WinDrive::cancel() { Drive::cancel(); if (m_child) { - m_child->kill(); + m_child->terminate(); m_child->deleteLater(); m_child = nullptr; } @@ -379,6 +244,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 +272,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 +286,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 +306,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 +332,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..be87e48e 100644 --- a/src/app/windrivemanager.h +++ b/src/app/windrivemanager.h @@ -21,27 +21,28 @@ #define WINDRIVEMANAGER_H #include "drivemanager.h" +#include "libwindisk/windisk.h" +#include #include 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 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 +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; diff --git a/src/helper/win/CMakeLists.txt b/src/helper/win/CMakeLists.txt index 60695802..724d0898 100644 --- a/src/helper/win/CMakeLists.txt +++ b/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} ) 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..72f1c3f2 100644 --- a/src/helper/win/restorejob.cpp +++ b/src/helper/win/restorejob.cpp @@ -18,42 +18,160 @@ */ #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()); + + /* + * 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_diskManagement->logMessage(QtCriticalMsg, "Couldn't unmount volumes on the drive"); + 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_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); + } + + /* + * 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, "Failed to lock the drive"); + 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_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); + } - 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_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); + } - 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_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); } - 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_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); + } + + CloseHandle(drive); + + 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..f7061f1a 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,444 @@ #include "writejob.h" #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(0); + return; } - return hVol; + qApp->exit(0); } -bool WriteJob::lockDrive(HANDLE drive) { - int attempts = 0; - DWORD status; - - while (true) { - if (!DeviceIoControl(drive, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &status, NULL)) { - attempts++; - } - else { - return true; - } - - if (attempts == 10) { - TCHAR message[256]; - FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); - - err << tr("Couldn't lock the drive") << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; - err.flush(); - break; - } +void WriteJob::onFileChanged(const QString &path) +{ + if (QFile::exists(path)) + return; + QRegularExpression reg("[.]part$"); + m_image = m_image.replace(reg, ""); - QThread::sleep(2); - } + m_out << "WRITE\n"; + m_out.flush(); - return false; + work(); } -bool WriteJob::removeMountPoints(uint diskNumber) { - DWORD drives = ::GetLogicalDrives(); - - for (char i = 0; i < 26; i++) { - if (drives & (1 << i)) { - char currentDrive = 'A' + i; - QString drivePath = QString("\\\\.\\%1:").arg(currentDrive); - - HANDLE hDevice = ::CreateFile(drivePath.toStdWString().c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL); +bool WriteJob::write() +{ + HANDLE drive; + const QString drivePath = QString("\\\\.\\PhysicalDrive%0").arg(m_disk->index()); + + m_diskManagement->logMessage(QtDebugMsg, "Preparing device for image writing"); + /* + * Device preparation part + */ + + if (!m_diskManagement->unmountVolumes(m_disk->index())) { + m_diskManagement->logMessage(QtCriticalMsg, "Couldn't unmount volumes on the drive"); + m_err << tr("Couldn't unmount volumes on the drive") << "\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->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(); + return false; + } - if (bResult) { - for (uint j = 0; j < vde.NumberOfDiskExtents; j++) { - if (vde.Extents[j].DiskNumber == diskNumber) { - QString volumePath = QString("%1:\\").arg(currentDrive); + 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(); + return false; + } - CloseHandle(hDevice); - hDevice = nullptr; + if (!m_diskManagement->lockDrive(drive, 10)) { + m_diskManagement->logMessage(QtCriticalMsg, "Failed to lock the drive"); + m_err << tr("Failed to lock the drive") << "\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; - } + m_diskManagement->refreshPartitionLayout(drive); - break; - } - } - } - if (hDevice) - CloseHandle(hDevice); - } + 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(); + return false; } - return true; -} - -bool WriteJob::cleanDrive(uint driveNumber) { - QProcess diskpart; - diskpart.setProgram("diskpart.exe"); - diskpart.setProcessChannelMode(QProcess::ForwardedChannels); + 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(); + return false; + } - diskpart.start(QIODevice::ReadWrite); + QThread::sleep(5); - 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"); + CloseHandle(drive); - diskpart.waitForFinished(); + /* + * Writing part + */ - if (diskpart.exitCode() == 0) { - // as advised in the diskpart documentation - QThread::sleep(15); + m_diskManagement->logMessage(QtDebugMsg, "Starting to write image file"); - return true; + drive = CreateFile(drivePath.toStdWString().c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_FLAG_WRITE_THROUGH | FILE_FLAG_NO_BUFFERING, NULL); + if (drive == INVALID_HANDLE_VALUE) { + m_diskManagement->logMessage(QtCriticalMsg, "Failed to open the drive for writing"); + m_err << tr("Failed to open the drive for writing") << "\n"; + m_err.flush(); + return false; + } + auto driveloseGuard = qScopeGuard([&drive] { + CloseHandle(drive); + }); + + if (!m_diskManagement->lockDrive(drive, 10)) { + m_diskManagement->logMessage(QtCriticalMsg, "Failed to lock the drive"); + m_err << tr("Failed to lock the drive") << "\n"; + m_err.flush(); + return false; } - return false; -} - -bool WriteJob::writeBlock(HANDLE drive, OVERLAPPED *overlap, char *data, uint size) { - DWORD bytesWritten; + bool result; + if (m_image.endsWith(".xz")) { + result = writeCompressed(drive); + } else { + result = writePlain(drive); + } - 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 (result) { + m_diskManagement->refreshPartitionLayout(drive); } - if (bytesWritten != size) { - err << tr("Destination drive is not writable") << "\n"; - err.flush(); + if (!result) { + qApp->exit(1); return false; } return true; } +bool WriteJob::writeCompressed(HANDLE driveHandle) +{ + const qint64 blockSize = m_disk->sectorSize() * 512; -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(); - } -} - -void WriteJob::work() { - if (!write()) { - out << "0\n"; - out.flush(); - QThread::sleep(5); - if (!write()) - return; + QFile isoFile(m_image); + isoFile.open(QIODevice::ReadOnly); + if (!isoFile.isOpen()) { + m_err << tr("Source image is not readable"); + m_err.flush(); + return false; } - - if (!check()) - return; - - qApp->exit(0); -} - -void WriteJob::onFileChanged(const QString &path) { - if (QFile::exists(path)) - return; - QRegularExpression reg("[.]part$"); - what = what.replace(reg, ""); - - out << "WRITE\n"; - out.flush(); - - work(); -} - -bool WriteJob::write() { - removeMountPoints(where); - cleanDrive(where); - - HANDLE drive = openDrive(where); - if (!lockDrive(drive)) { + 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 the drive for writing") << "\n"; + m_err.flush(); qApp->exit(1); 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(static_cast(outBuffer), readBytes); + if (writtenBytes < 0) { + m_diskManagement->logMessage(QtCriticalMsg, QStringLiteral("Failed to write to the device: %1").arg(getLastError())); + m_err << tr("Failed to write to the device: ") << drive.errorString() << "\n"; + m_err << tr("Failed to write to the device: ") << 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(static_cast(outBuffer), sectorSize); + if (writtenBytes < 0) { + m_diskManagement->logMessage(QtCriticalMsg, QStringLiteral("Failed to write to the device: %1").arg(getLastError())); + m_err << tr("Failed to write to the device: ") << drive.errorString() << "\n"; + m_err << tr("Failed to write to the device: ") << 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; +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(); + m_err << tr("Source image is not readable") << "\n"; + 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 the drive for writing") << "\n"; + m_err.flush(); qApp->exit(1); return false; } - + auto driveCleanup = qScopeGuard([&drive] { + drive.close(); + }); + + void *buffer = NULL; + buffer = VirtualAlloc(NULL, blockSize, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE); + if (buffer == NULL) { + m_err << tr("Failed to allocate the buffer") << "\n"; + m_err.flush(); + return false; + } + auto bufferCleanup = qScopeGuard([buffer, blockSize] { + VirtualFree(buffer, blockSize, MEM_DECOMMIT | MEM_RELEASE); + }); + + 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(static_cast(buffer), blockSize)) <= 0) { + break; + } + + readBytes = ((readBytes + sectorSize - 1) / sectorSize) * sectorSize; + writtenBytes = drive.write(static_cast(buffer), readBytes); + if (writtenBytes < 0) { + m_diskManagement->logMessage(QtCriticalMsg, QStringLiteral("Failed to write to the device: %1").arg(getLastError())); + m_err << tr("Failed to write to the device: ") << drive.errorString() << "\n"; + m_err << tr("Failed to write to the device: ") << getLastError() << "\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(); + if (writtenBytes != readBytes) { + m_err << tr("The last block was not fully written") << "\n"; + m_err.flush(); + return false; + } - if (buffer.size() != BLOCK_SIZE || isoFile.atEnd()) + totalBytes += readBytes; + m_out << totalBytes << "\n"; + m_out.flush(); + + 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_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_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_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..eac76a85 --- /dev/null +++ b/src/lib/libwindisk/windisk.cpp @@ -0,0 +1,1014 @@ +/* + * 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("Calling %1()").arg(__func__)); + + 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 + res = m_IWbemServices->ExecQuery(_bstr_t(L"WQL"), _bstr_t(L"SELECT * FROM MSFT_Disk WHERE BusType = 7"), 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("Calling %1(%2)").arg(__func__).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("Calling %1(%2)").arg(__func__).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("Calling %1(%2)").arg(__func__).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("Calling %1(%2)").arg(__func__).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("Calling %1()").arg(__func__)); + + 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("Calling %1()").arg(__func__)); + + int attempts = 0; + DWORD status; + + while (true) { + if (!DeviceIoControl(driveHandle, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &status, NULL)) { + attempts++; + } else { + return true; + } + + if (attempts == numRetries) { + logMessage(QtWarningMsg, QStringLiteral("Couldn't lock the drive: %1").arg(getLastError())); + break; + } + + QThread::sleep(2); + } + + return false; +} + +bool WinDiskManagement::unmountVolumes(quint32 index) +{ + logMessage(QtDebugMsg, QStringLiteral("Calling %1(%2)").arg(__func__).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(QtWarningMsg, 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 unmounted volume: %1").arg(currentDrive)); + break; + } + } + } + } + } + + return true; +} + +QChar WinDiskManagement::mountVolume(const QString &logicalName) +{ + logMessage(QtDebugMsg, QStringLiteral("Calling %1()").arg(__func__)); + + 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 %1 as %2").arg(logicalName).arg(driveLetter)); + return QChar(driveLetter); +} + +bool WinDiskManagement::writeFileWithRetry(HANDLE driveHandle, LPCVOID buffer, WORD 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")); + 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("Calling %1()").arg(__func__)); + + 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, static_cast(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, static_cast(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("Calling %1()").arg(__func__)); + + 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("Calling %1()").arg(__func__)); + + 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) +{ + logMessage(QtDebugMsg, QStringLiteral("Calling %1(%2)").arg(__func__).arg(index)); + + 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(QtWarningMsg, 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; + } + + volumeName[len - 1] = '\\'; + result = QString(volumeName); + break; + } + + return result; + ; +} + +bool WinDiskManagement::refreshPartitionLayout(HANDLE driveHandle) +{ + logMessage(QtDebugMsg, QStringLiteral("Calling %1()").arg(__func__)); + + 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..2c003e10 --- /dev/null +++ b/src/lib/libwindisk/windisk.h @@ -0,0 +1,117 @@ +/* + * 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); + // Unmounts all logical volumes + bool unmountVolumes(quint32 index); + // 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); + // Refreshes the partition layout + bool refreshPartitionLayout(HANDLE driveHandle); + +private: + bool writeFileWithRetry(HANDLE driveHandle, LPCVOID buffer, WORD numberOfBytesToWrite, int numberOfRetries); + + 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