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/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..a0d7196c 100644 --- a/src/helper/win/restorejob.cpp +++ b/src/helper/win/restorejob.cpp @@ -18,42 +18,150 @@ */ #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_err << tr("Couldn't unmount all volumes on the drive") << "\n"; + m_err.flush(); + return; + } + + /* + * 2) Remove all the existing partitions + * This uses "DeleteObject" on MSFT_Partition using QMI query + */ + if (!m_diskManagement->clearPartitions(m_disk->index())) { + m_err << tr("Failed to clear the drive") << "\n"; + m_err.flush(); + qApp->exit(1); + return; + } + + drive = CreateFile(drivePath.toStdWString().c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL); + if (drive == INVALID_HANDLE_VALUE) { + m_err << tr("Failed to open the drive for GPT/MBR removal") << "\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_out << tr("Failed to lock the drive") << "\n"; + m_out.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_out << tr("Failed to clear the partition table on the drive.") << "\n"; + m_out.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_out << tr("Failed to clear the drive: %1") << "\n"; + m_out.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_out << tr("Failed to create GPT partition on the drive") << "\n"; + m_out.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_out << tr("Failed to get GUID volume path on the drive: %1").arg(m_disk->index()) << "\n"; + m_out.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_out << tr("Failed to mount the partition") << "\n"; + m_out.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_out << tr("Failed to format the partition to exFAT") << "\n"; + m_out.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..83fd4b4f 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,422 @@ #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); - - 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; +void WriteJob::work() +{ + if (!write()) { + qApp->exit(1); + return; } - return hVol; -} - -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; - } - - QThread::sleep(2); + if (!check()) { + qApp->exit(0); + return; } - return false; + qApp->exit(0); } -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); - - 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 (bResult) { - for (uint j = 0; j < vde.NumberOfDiskExtents; j++) { - if (vde.Extents[j].DiskNumber == diskNumber) { - QString volumePath = QString("%1:\\").arg(currentDrive); - - CloseHandle(hDevice); - hDevice = nullptr; - - 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; - } +void WriteJob::onFileChanged(const QString &path) +{ + if (QFile::exists(path)) + return; + QRegularExpression reg("[.]part$"); + m_image = m_image.replace(reg, ""); - break; - } - } - } - if (hDevice) - CloseHandle(hDevice); - } - } + m_out << "WRITE\n"; + m_out.flush(); - return true; + work(); } -bool WriteJob::cleanDrive(uint driveNumber) { - QProcess diskpart; - diskpart.setProgram("diskpart.exe"); - diskpart.setProcessChannelMode(QProcess::ForwardedChannels); - - diskpart.start(QIODevice::ReadWrite); - - 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"); - - diskpart.waitForFinished(); +bool WriteJob::write() +{ + HANDLE drive; + const QString drivePath = QString("\\\\.\\PhysicalDrive%0").arg(m_disk->index()); - if (diskpart.exitCode() == 0) { - // as advised in the diskpart documentation - QThread::sleep(15); + /* + * Device preparation part + */ - return true; + if (!m_diskManagement->unmountVolumes(m_disk->index())) { + m_out << tr("Failed to unmount volumes on the drive before writing"); + m_out.flush(); + return false; } - return false; -} + if (!m_diskManagement->clearPartitions(m_disk->index())) { + m_out << tr("Failed to clear the drive before writing"); + m_out.flush(); + return false; + } -bool WriteJob::writeBlock(HANDLE drive, OVERLAPPED *overlap, char *data, uint size) { - DWORD bytesWritten; + 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("Couldn't open the drive for GPT/MBR removal for %1:").arg(drivePath) << " (" << getLastError() << ")\n"; + m_err.flush(); + return false; + } - 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 (!m_diskManagement->lockDrive(drive, 10)) { + m_err << tr("Failed to lock the drive"); + m_err.flush(); + return false; } - if (bytesWritten != size) { - err << tr("Destination drive is not writable") << "\n"; - err.flush(); + m_diskManagement->refreshPartitionLayout(drive); + + if (!m_diskManagement->clearDiskDrive(drive)) { + m_err << tr("Failed to clear the drive before writing"); + m_err.flush(); return false; } - return true; -} + CloseHandle(drive); + /* + * Writing part + */ -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(); + 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_err << tr("Couldn't open the drive for partition removal for %1:").arg(drivePath) << " (" << getLastError() << ")\n"; + m_err.flush(); + return false; } -} + auto driveloseGuard = qScopeGuard([&drive] { + CloseHandle(drive); + }); -void WriteJob::work() { - if (!write()) { - out << "0\n"; - out.flush(); - QThread::sleep(5); - if (!write()) - return; + if (!m_diskManagement->lockDrive(drive, 10)) { + m_out << tr("Failed to lock the drive"); + m_out.flush(); + return false; } - if (!check()) - return; + QThread::sleep(5); - qApp->exit(0); -} + bool result; + if (m_image.endsWith(".xz")) { + result = writeCompressed(drive); + } else { + result = writePlain(drive); + } -void WriteJob::onFileChanged(const QString &path) { - if (QFile::exists(path)) - return; - QRegularExpression reg("[.]part$"); - what = what.replace(reg, ""); + m_diskManagement->refreshPartitionLayout(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)) { + 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 the device for writing"); + 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"); + 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."); 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_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"); + 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."); 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."); break; case LZMA_OPTIONS_ERROR: - err << tr("Unsupported compression options."); + m_err << tr("Unsupported compression options."); break; default: - err << tr("Unknown decompression error."); + m_err << tr("Unknown decompression error."); 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_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"); + 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"); + 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 device for writing"); + 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"); + 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_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"); + 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(); + 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..4c192487 --- /dev/null +++ b/src/lib/libwindisk/windisk.cpp @@ -0,0 +1,988 @@ +/* + * 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() +{ + QMap result; + if (!m_wmiInitialized) { + logMessage(QtWarningMsg, 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(QtWarningMsg, 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; + 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; + logMessage(QtDebugMsg, QStringLiteral("Disk Index: %1").arg(index)); + } else if (var.vt == VT_UI4) { + index = var.uintVal; + logMessage(QtDebugMsg, QStringLiteral("Disk Index: %1").arg(index)); + } + 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); + logMessage(QtDebugMsg, QStringLiteral("Device ID%1").arg(deviceID)); + } + VariantClear(&var); + } + pDiskObject->Release(); + + result.insert(index, deviceID); + } + pEnumDiskObjects->Release(); + + return result; +} + +QMap WinDiskManagement::getDevicePartitions(quint32 index) +{ + QMap result; + if (!m_wmiInitialized) { + logMessage(QtWarningMsg, 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(QtWarningMsg, 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; + logMessage(QtDebugMsg, QStringLiteral("Partition Index: %1").arg(partitionIndex)); + } else if (var.vt == VT_UI4) { + partitionIndex = var.uintVal; + logMessage(QtDebugMsg, QStringLiteral("Partition Index: %1").arg(partitionIndex)); + } + 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); + logMessage(QtDebugMsg, QStringLiteral("Partition path: %1").arg(partitionPath)); + } + VariantClear(&var); + } + + if ((pPartitionObject->Get(_bstr_t(L"DriveLetter"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_I2) { + wchar_t driveLetter = static_cast(var.iVal); + mountable = driveLetter != L'\0'; + if (mountable) { + logMessage(QtDebugMsg, QStringLiteral("DriveLetter: %1. Partition is mountable").arg(QString::fromWCharArray(&driveLetter, 1))); + } else { + logMessage(QtDebugMsg, QStringLiteral("No assigned letter. Partition is not mountable")); + } + } + VariantClear(&var); + } + + pPartitionObject->Release(); + + if (partitionIndex != -1) { + result.insert(static_cast(partitionIndex), mountable); + } + } + pPartitionObjects->Release(); + + return result; +} + +std::unique_ptr WinDiskManagement::getDiskDriveInformation(quint32 index, const QString &diskPath) +{ + std::unique_ptr result = std::make_unique(index, diskPath); + if (!m_wmiInitialized) { + logMessage(QtWarningMsg, 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(QtWarningMsg, 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) +{ + 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(QtWarningMsg, 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(QtWarningMsg, 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(QtWarningMsg, 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) +{ + if (!m_wmiInitialized) { + logMessage(QtWarningMsg, 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(QtWarningMsg, 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(QtWarningMsg, 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(QtWarningMsg, 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(QtWarningMsg, 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(QtWarningMsg, 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(QtWarningMsg, 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(QtWarningMsg, 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(QtWarningMsg, 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) +{ + if (!m_wmiInitialized) { + logMessage(QtWarningMsg, 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(QtWarningMsg, 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) +{ + 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) +{ + 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(QtWarningMsg, 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) +{ + 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(QtWarningMsg, 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(QtWarningMsg, 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(QtWarningMsg, 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(QtWarningMsg, QStringLiteral("Could not set file pointer")); + break; + } + if (WriteFile(driveHandle, buffer, numberOfBytesToWrite, &writtenBytes, NULL)) { + if (numberOfBytesToWrite == writtenBytes) { + return true; + } + } else { + logMessage(QtWarningMsg, 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) +{ + quint64 sectorsToClear = 128; + LARGE_INTEGER filePointer; + + uint8_t *zeroBuffer = static_cast(calloc(sectorSize, sectorsToClear)); + if (!zeroBuffer) { + logMessage(QtWarningMsg, QStringLiteral("Couldn't allocate zero buffer")); + return false; + } + + filePointer.QuadPart = 0ULL; + if (!SetFilePointerEx(driveHandle, filePointer, &filePointer, FILE_BEGIN) || (filePointer.QuadPart != 0ULL)) { + logMessage(QtWarningMsg, QStringLiteral("Couldn't reset disk position: %1").arg(getLastError())); + } + + if (!writeFileWithRetry(driveHandle, zeroBuffer, static_cast(sectorSize * sectorsToClear), 4)) { + logMessage(QtWarningMsg, 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(QtWarningMsg, 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) +{ + 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(QtWarningMsg, 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) +{ + 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(QtWarningMsg, 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(QtWarningMsg, 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(QtWarningMsg, 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) +{ + 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) +{ + BOOL ret = DeviceIoControl(driveHandle, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, NULL, NULL); + if (!ret) { + logMessage(QtWarningMsg, 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