From 9f814dad1b43153b0665b5284fc71e7c61b0ead8 Mon Sep 17 00:00:00 2001 From: Jan Grulich Date: Fri, 4 Oct 2024 16:12:21 +0200 Subject: [PATCH] WIP: rewrite Windows code --- .github/workflows/clang-format-check.yml | 2 +- src/app/CMakeLists.txt | 2 +- src/app/main.cpp | 1 + src/app/windrivemanager.cpp | 298 +++---------- src/app/windrivemanager.h | 12 +- src/helper/win/CMakeLists.txt | 1 + src/helper/win/main.cpp | 24 +- src/helper/win/restorejob.cpp | 10 +- src/helper/win/restorejob.h | 6 +- src/helper/win/writejob.cpp | 534 +++++++++++++---------- src/helper/win/writejob.h | 45 +- src/lib/CMakeLists.txt | 3 + src/lib/libwmi/CMakeLists.txt | 9 + src/lib/libwmi/libwmi.cpp | 361 +++++++++++++++ src/lib/libwmi/libwmi.h | 81 ++++ 15 files changed, 880 insertions(+), 509 deletions(-) create mode 100644 src/lib/libwmi/CMakeLists.txt create mode 100644 src/lib/libwmi/libwmi.cpp create mode 100644 src/lib/libwmi/libwmi.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 9449894a..06613cf7 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 libwmi) 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..ee1d7ecb 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,99 @@ #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_wmi(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; + } } + QTimer::singleShot(0, this, &WinDriveProvider::checkDrives); + return true; } - - if (firstRun) - mDebug() << this->metaObject()->className() << "Finished looking for the drives for the first time"; - firstRun = false; - QTimer::singleShot(2500, this, &WinDriveProvider::checkDrives); -} - -QString getPhysicalName(int driveNumber) -{ - return QString("\\\\.\\PhysicalDrive%0").arg(driveNumber); -} - -HANDLE getPhysicalHandle(int driveNumber) -{ - HANDLE physicalHandle = INVALID_HANDLE_VALUE; - QString physicalPath = getPhysicalName(driveNumber); - physicalHandle = CreateFileA(physicalPath.toStdString().c_str(), GENERIC_READ, FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); - return physicalHandle; + return false; } -bool WinDriveProvider::isMountable(int driveNumber) +void WinDriveProvider::checkDrives() { - mDebug() << this->metaObject()->className() << "Checking whether " << getPhysicalName(driveNumber) << " is mountable"; - - HANDLE physicalHandle = getPhysicalHandle(driveNumber); - if (physicalHandle == INVALID_HANDLE_VALUE) { - mDebug() << this->metaObject()->className() << "Could not get physical handle for drive " << getPhysicalName(driveNumber); - return false; - } - - DWORD size; - BYTE geometry[256]; - bool ret = DeviceIoControl(physicalHandle, IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, NULL, 0, geometry, sizeof(geometry), &size, NULL); - if (!ret || size <= 0) { - mDebug() << this->metaObject()->className() << "Could not get geometry for drive " << getPhysicalName(driveNumber); - CloseHandle(physicalHandle); - return false; - } - - PDISK_GEOMETRY_EX diskGeometry = (PDISK_GEOMETRY_EX)(void *)geometry; - // Drive info - LONGLONG diskSize; - DWORD sectorSize; - DWORD sectorsPerTrack; - DWORD firstDataSector; - MEDIA_TYPE mediaType; - - diskSize = diskGeometry->DiskSize.QuadPart; - sectorSize = diskGeometry->Geometry.BytesPerSector; - firstDataSector = MAXDWORD; - if (sectorSize < 512) { - mDebug() << this->metaObject()->className() << "Warning: Drive " << getPhysicalName(driveNumber) << " reports a sector size of " << sectorSize << " - Correcting to 512 bytes."; - sectorSize = 512; - } - sectorsPerTrack = diskGeometry->Geometry.SectorsPerTrack; - mediaType = diskGeometry->Geometry.MediaType; - - BYTE layout[4096] = {0}; - ret = DeviceIoControl(physicalHandle, IOCTL_DISK_GET_DRIVE_LAYOUT_EX, NULL, 0, layout, sizeof(layout), &size, NULL); - if (!ret || size <= 0) { - mDebug() << this->metaObject()->className() << "Could not get layout for drive " << getPhysicalName(driveNumber); - CloseHandle(physicalHandle); - return false; - } - - PDRIVE_LAYOUT_INFORMATION_EX driveLayout = (PDRIVE_LAYOUT_INFORMATION_EX)(void *)layout; - - switch (driveLayout->PartitionStyle) { - case PARTITION_STYLE_MBR: - mDebug() << this->metaObject()->className() << "MBR partition style"; - for (int i = 0; i < driveLayout->PartitionCount; i++) { - if (driveLayout->PartitionEntry[i].Mbr.PartitionType != PARTITION_ENTRY_UNUSED) { - QVector mbrMountable = {0x01, 0x04, 0x06, 0x07, 0x0b, 0x0c, 0x0e}; - BYTE partType = driveLayout->PartitionEntry[i].Mbr.PartitionType; - mDebug() << this->metaObject()->className() << "Partition type: " << partType; - if (!mbrMountable.contains(partType)) { - CloseHandle(physicalHandle); - mDebug() << this->metaObject()->className() << getPhysicalName(driveNumber) << " is not mountable"; - return false; - } - } - } - break; - case PARTITION_STYLE_GPT: - mDebug() << this->metaObject()->className() << "GPT partition style"; - for (int i = 0; i < driveLayout->PartitionCount; i++) { - if (memcmp(&driveLayout->PartitionEntry[i].Gpt.PartitionType, &PARTITION_MICROSOFT_DATA, sizeof(GUID)) != 0) { - CloseHandle(physicalHandle); - mDebug() << this->metaObject()->className() << getPhysicalName(driveNumber) << " is not mountable"; - return false; + mDebug() << this->metaObject()->className() << "Looking for the drives"; + + QMap drives; + auto usbDeviceList = m_wmi->getUSBDeviceList(); + for (auto it = usbDeviceList.cbegin(); it != usbDeviceList.cend(); it++) { + bool noLogicalDevice = false; + QStringList partitionList = m_wmi->getDevicePartitions(it.key()); + for (const QString &partition : partitionList) { + if (m_wmi->getLogicalDisks(partition).isEmpty()) { + noLogicalDevice = true; + break; } } - break; - default: - mDebug() << this->metaObject()->className() << "Partition type: RAW"; - break; - } - - mDebug() << this->metaObject()->className() << getPhysicalName(driveNumber) << " is mountable"; - CloseHandle(physicalHandle); - return true; -} - -bool WinDriveProvider::describeDrive(int nDriveNumber, bool verbose) -{ - BOOL removable; - QString productVendor; - QString productId; - QString serialNumber; - uint64_t deviceBytes; - STORAGE_BUS_TYPE storageBus; - - BOOL bResult = FALSE; // results flag - // DWORD dwRet = NO_ERROR; - - // Format physical drive path (may be '\\.\PhysicalDrive0', '\\.\PhysicalDrive1' and so on). - QString strDrivePath = getPhysicalName(nDriveNumber); - - // Get a handle to physical drive - HANDLE hDevice = ::CreateFile(strDrivePath.toStdWString().c_str(), 0, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL); - - if (hDevice == INVALID_HANDLE_VALUE) - return false; //::GetLastError(); - - if (verbose) - mDebug() << this->metaObject()->className() << strDrivePath << "is present"; - - // Set the input data structure - STORAGE_PROPERTY_QUERY storagePropertyQuery; - ZeroMemory(&storagePropertyQuery, sizeof(STORAGE_PROPERTY_QUERY)); - storagePropertyQuery.PropertyId = StorageDeviceProperty; - storagePropertyQuery.QueryType = PropertyStandardQuery; - - // Get the necessary output buffer size - STORAGE_DESCRIPTOR_HEADER storageDescriptorHeader; - ZeroMemory(&storageDescriptorHeader, sizeof(STORAGE_DESCRIPTOR_HEADER)); - DWORD dwBytesReturned = 0; - if (verbose) - mDebug() << this->metaObject()->className() << strDrivePath << "IOCTL_STORAGE_QUERY_PROPERTY"; - if (!::DeviceIoControl(hDevice, IOCTL_STORAGE_QUERY_PROPERTY, &storagePropertyQuery, sizeof(STORAGE_PROPERTY_QUERY), &storageDescriptorHeader, sizeof(STORAGE_DESCRIPTOR_HEADER), &dwBytesReturned, NULL)) { - // dwRet = ::GetLastError(); - ::CloseHandle(hDevice); - return false; // dwRet; - } + auto diskDrive = m_wmi->getDiskDriveInformation(it.key(), it.value()); + if (diskDrive->model().isEmpty() || !diskDrive->size() || diskDrive->serialNumber().isEmpty()) { + continue; + } - // Alloc the output buffer - const DWORD dwOutBufferSize = storageDescriptorHeader.Size; - BYTE *pOutBuffer = new BYTE[dwOutBufferSize]; - ZeroMemory(pOutBuffer, dwOutBufferSize); - - if (verbose) - mDebug() << this->metaObject()->className() << strDrivePath << "IOCTL_STORAGE_QUERY_PROPERTY with a bigger buffer"; - // Get the storage device descriptor - if (!(bResult = ::DeviceIoControl(hDevice, IOCTL_STORAGE_QUERY_PROPERTY, &storagePropertyQuery, sizeof(STORAGE_PROPERTY_QUERY), pOutBuffer, dwOutBufferSize, &dwBytesReturned, NULL))) { - // dwRet = ::GetLastError(); - delete[] pOutBuffer; - ::CloseHandle(hDevice); - return false; // dwRet; + WinDrive *currentDrive = new WinDrive(this, diskDrive->model(), diskDrive->size(), noLogicalDevice, diskDrive->index(), diskDrive->serialNumber()); + drives[diskDrive->index()] = currentDrive; } - // Now, the output buffer points to a STORAGE_DEVICE_DESCRIPTOR structure - // followed by additional info like vendor ID, product ID, serial number, and so on. - STORAGE_DEVICE_DESCRIPTOR *pDeviceDescriptor = (STORAGE_DEVICE_DESCRIPTOR *)pOutBuffer; - removable = pDeviceDescriptor->RemovableMedia; - if (pDeviceDescriptor->ProductIdOffset != 0) - productId = QString((char *)pOutBuffer + pDeviceDescriptor->ProductIdOffset).trimmed(); - if (pDeviceDescriptor->VendorIdOffset != 0) - productVendor = QString((char *)pOutBuffer + pDeviceDescriptor->VendorIdOffset).trimmed(); - if (pDeviceDescriptor->SerialNumberOffset != 0) - serialNumber = QString((char *)pOutBuffer + pDeviceDescriptor->SerialNumberOffset).trimmed(); - storageBus = pDeviceDescriptor->BusType; - - if (verbose) - mDebug() << this->metaObject()->className() << strDrivePath << "detected:" << productVendor << productId << (removable ? ", removable" : ", nonremovable") << (storageBus == BusTypeUsb ? "USB" : "notUSB"); - - if (!removable && storageBus != BusTypeUsb) - return false; - - DISK_GEOMETRY pdg; - DWORD junk = 0; // discard results - - if (verbose) - mDebug() << this->metaObject()->className() << strDrivePath << "IOCTL_DISK_GET_DRIVE_GEOMETRY"; - bResult = DeviceIoControl(hDevice, // device to be queried - IOCTL_DISK_GET_DRIVE_GEOMETRY, // operation to perform - NULL, - 0, // no input buffer - &pdg, - sizeof(pdg), // output buffer - &junk, // # bytes returned - (LPOVERLAPPED)NULL); // synchronous I/O - - if (!bResult || pdg.MediaType == Unknown) - return false; + // Update our list of drives and notify about added and removed drives + QList driveIndexes = m_drives.keys(); + for (auto it = drives.constBegin(); it != drives.constEnd(); it++) { + if (m_drives.contains(it.key()) && *m_drives[it.key()] == *it.value()) { + mDebug() << "Drive " << it.key() << " already exists"; + it.value()->deleteLater(); + driveIndexes.removeAll(it.key()); + continue; + } - deviceBytes = pdg.Cylinders.QuadPart * pdg.TracksPerCylinder * pdg.SectorsPerTrack * pdg.BytesPerSector; + if (m_drives.contains(it.key())) { + mDebug() << "Replacing old drive in the list on index " << it.key(); + emit driveRemoved(m_drives[it.key()]); + m_drives[it.key()]->deleteLater(); + m_drives.remove(it.key()); + } - // Do cleanup and return - if (verbose) - mDebug() << this->metaObject()->className() << strDrivePath << "cleanup, adding to the list"; - delete[] pOutBuffer; - ::CloseHandle(hDevice); + mDebug() << "Adding new drive to the list with index " << it.key(); + m_drives[it.key()] = it.value(); + emit driveConnected(it.value()); - WinDrive *currentDrive = new WinDrive(this, productVendor + " " + productId, deviceBytes, !isMountable(nDriveNumber), nDriveNumber, serialNumber); - if (m_drives.contains(nDriveNumber) && *m_drives[nDriveNumber] == *currentDrive) { - currentDrive->deleteLater(); - return true; + driveIndexes.removeAll(it.key()); } - if (m_drives.contains(nDriveNumber)) { - emit driveRemoved(m_drives[nDriveNumber]); - m_drives[nDriveNumber]->deleteLater(); + // Remove our previously stored drives that were not present in the last check + for (int index : driveIndexes) { + mDebug() << "Removing old drive with index" << index; + emit driveRemoved(m_drives[index]); + m_drives[index]->deleteLater(); + m_drives.remove(index); } - - m_drives[nDriveNumber] = currentDrive; - emit driveConnected(currentDrive); - - return true; } WinDrive::WinDrive(WinDriveProvider *parent, const QString &name, uint64_t size, bool containsLive, int device, const QString &serialNumber) @@ -379,6 +216,11 @@ void WinDrive::restore() m_child->start(QIODevice::ReadOnly); } +bool WinDrive::busy() const +{ + return (m_child && m_child->state() == QProcess::Running); +} + QString WinDrive::serialNumber() const { return m_serialNo; @@ -455,7 +297,7 @@ void WinDrive::onReadyRead() Notifications::notify(tr("Finished!"), tr("Writing %1 was successful").arg(m_image->fullName())); } else { bool ok; - qreal bytes = line.toLongLong(&ok); + qint64 bytes = line.toLongLong(&ok); if (ok) { if (bytes < 0) m_progress->setValue(NAN); diff --git a/src/app/windrivemanager.h b/src/app/windrivemanager.h index 3d1ac43d..c99b18d9 100644 --- a/src/app/windrivemanager.h +++ b/src/app/windrivemanager.h @@ -21,27 +21,28 @@ #define WINDRIVEMANAGER_H #include "drivemanager.h" +#include "libwmi/libwmi.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_wmi; }; 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 50e8d5f7..26e8c4fb 100644 --- a/src/helper/win/CMakeLists.txt +++ b/src/helper/win/CMakeLists.txt @@ -26,6 +26,7 @@ remove_definitions(-std=c++17) target_link_libraries(helper Qt6::Core isomd5 + libwmi ${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..40e96204 100644 --- a/src/helper/win/restorejob.cpp +++ b/src/helper/win/restorejob.cpp @@ -22,8 +22,8 @@ #include #include -RestoreJob::RestoreJob(const QString &where) - : QObject(nullptr) +RestoreJob::RestoreJob(const QString &where, QObject *parent) + : QObject(parent) { bool ok = false; m_where = where.toInt(&ok); @@ -33,7 +33,8 @@ RestoreJob::RestoreJob(const QString &where) QTimer::singleShot(0, this, &RestoreJob::work); } -void RestoreJob::work() { +void RestoreJob::work() +{ m_diskpart.setProgram("diskpart.exe"); m_diskpart.setProcessChannelMode(QProcess::ForwardedChannels); @@ -50,8 +51,7 @@ void RestoreJob::work() { if (m_diskpart.waitForFinished()) { qApp->exit(0); - } - else { + } else { err << m_diskpart.readAllStandardError(); err.flush(); qApp->exit(1); diff --git a/src/helper/win/restorejob.h b/src/helper/win/restorejob.h index a46ada9b..75fe074c 100644 --- a/src/helper/win/restorejob.h +++ b/src/helper/win/restorejob.h @@ -28,7 +28,7 @@ class RestoreJob : public QObject { Q_OBJECT public: - explicit RestoreJob(const QString &where); + explicit RestoreJob(const QString &where, QObject *parent); signals: @@ -36,8 +36,8 @@ private slots: void work(); private: - QTextStream out { stdout }; - QTextStream err { stderr }; + QTextStream out{stdout}; + QTextStream err{stderr}; QProcess m_diskpart; int m_where; diff --git a/src/helper/win/writejob.cpp b/src/helper/win/writejob.cpp index 30947f30..ee39c3c3 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,83 +21,65 @@ #include "writejob.h" #include -#include -#include -#include #include -#include +#include #include -#include +#include +#include #include -#include #include #include "isomd5/libcheckisomd5.h" - -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_wmi = std::make_unique(this); + m_wmiDiskDrive = m_wmi->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; - } - - return hVol; -} - -bool WriteJob::lockDrive(HANDLE drive) { +bool WriteJob::lockDrive(HANDLE driveHandle) +{ int attempts = 0; DWORD status; while (true) { - if (!DeviceIoControl(drive, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &status, NULL)) { + if (!DeviceIoControl(driveHandle, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &status, NULL)) { attempts++; - } - else { + } 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(); + m_err << tr("Couldn't lock the drive") << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; + m_err.flush(); break; } @@ -106,310 +89,405 @@ bool WriteJob::lockDrive(HANDLE drive) { return false; } -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::work() +{ + if (!write()) { + return; + } - break; - } - } - } - if (hDevice) - CloseHandle(hDevice); - } + if (!check()) { + return; } - return true; + qApp->exit(0); } -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"); +void WriteJob::onFileChanged(const QString &path) +{ + if (QFile::exists(path)) + return; + QRegularExpression reg("[.]part$"); + m_image = m_image.replace(reg, ""); - diskpart.waitForFinished(); + m_out << "WRITE\n"; + m_out.flush(); - if (diskpart.exitCode() == 0) { - // as advised in the diskpart documentation - QThread::sleep(15); + work(); +} - return true; +bool WriteJob::write() +{ + HANDLE drive = CreateFile(m_wmiDiskDrive->deviceID().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) { + TCHAR message[256]; + FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); + m_err << tr("Couldn't open the drive for partition removal for %1:").arg(m_wmiDiskDrive->deviceID()) << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; + m_err.flush(); + qApp->exit(1); + return false; } + auto driveloseGuard = qScopeGuard([&drive] { + CloseHandle(drive); + }); + + QList volumeHandles; + auto volumeCloseGuard = qScopeGuard([&volumeHandles] { + for (HANDLE volume : volumeHandles) { + CloseHandle(volume); + } + }); + + const QStringList partitions = m_wmi->getDevicePartitions(m_wmiDiskDrive->index()); + for (const QString &partition : partitions) { + QStringList volumes = m_wmi->getLogicalDisks(partition); + for (const QString &volumeID : volumes) { + const QString volumePath = QString("\\\\.\\%1").arg(volumeID); + m_err << "Volume path " << volumePath << "\n"; + m_out << "Unmounting: " << volumePath; + m_out.flush(); + + DWORD ret; + HANDLE volume = CreateFile(volumePath.toStdWString().c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL); + if (volume == INVALID_HANDLE_VALUE) { + TCHAR message[256]; + FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); + m_err << tr("Couldn't open the volume %1:").arg(volumeID) << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; + m_err.flush(); + qApp->exit(1); + return false; + } + volumeHandles << volume; - return false; -} - -bool WriteJob::writeBlock(HANDLE drive, OVERLAPPED *overlap, char *data, uint size) { - DWORD bytesWritten; + if (!DeviceIoControl(volume, FSCTL_DISMOUNT_VOLUME, NULL, 0, NULL, 0, &ret, NULL)) { + TCHAR message[256]; + FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); + m_err << tr("Couldn't unmount the volume %1:").arg(volumeID) << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; + m_err.flush(); + qApp->exit(1); + 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 (!DeviceIoControl(volume, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &ret, NULL)) { + TCHAR message[256]; + FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); + m_err << tr("Couldn't lock the volume %1:").arg(volumeID) << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; + m_err.flush(); + } } } - if (bytesWritten != size) { - err << tr("Destination drive is not writable") << "\n"; - err.flush(); + if (!lockDrive(drive)) { + m_out << "Failed to lock the drive"; + m_out.flush(); + qApp->exit(1); return false; } - return true; -} + // Remove partitions + DWORD ret; + DRIVE_LAYOUT_INFORMATION_EX *driveLayout = (DRIVE_LAYOUT_INFORMATION_EX *)malloc(sizeof(DRIVE_LAYOUT_INFORMATION_EX)); + ZeroMemory(driveLayout, sizeof(DRIVE_LAYOUT_INFORMATION_EX)); + driveLayout->PartitionCount = 0; + driveLayout->PartitionStyle = PARTITION_STYLE_RAW; + driveLayout->Mbr.Signature = 0; - -void WriteJob::unlockDrive(HANDLE drive) { - DWORD status; - if (!DeviceIoControl(drive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &status, NULL)) { + if (!DeviceIoControl(drive, IOCTL_DISK_SET_DRIVE_LAYOUT_EX, driveLayout, sizeof(DRIVE_LAYOUT_INFORMATION_EX), NULL, 0, &ret, 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(); + m_err << tr("Couldn't remove the partition table for %1:").arg(m_wmiDiskDrive->deviceID()) << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; + m_err.flush(); + qApp->exit(1); + free(driveLayout); + return false; } -} - -void WriteJob::work() { - if (!write()) { - out << "0\n"; - out.flush(); - QThread::sleep(5); - if (!write()) - return; + free(driveLayout); + + // FIXME: not sure if this is needed and it has been working + // also without it, but previously we did that for diskpart + // so it might maybe help in some situations and 5 seconds + // delay is not a big deal. + QThread::sleep(5); + + bool result; + if (m_image.endsWith(".xz")) { + result = writeCompressed(drive); + } else { + result = writePlain(drive); } - 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(); + 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_wmiDiskDrive->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_wmiDiskDrive->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"; + TCHAR message[256]; + FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); + m_err << tr("Failed to write to the device:") << " (" << QString::fromWCharArray(message).trimmed() << ")\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"; + TCHAR message[256]; + FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); + m_err << tr("Failed to write to the device:") << " (" << QString::fromWCharArray(message).trimmed() << ")\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_wmiDiskDrive->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_wmiDiskDrive->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"; + TCHAR message[256]; + FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); + m_err << tr("Failed to write to the device:") << " (" << QString::fromWCharArray(message).trimmed() << ")\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(); +bool WriteJob::check() +{ + m_out << "CHECK\n"; + m_out.flush(); - HANDLE drive = openDrive(where); + HANDLE drive = CreateFile(m_wmiDiskDrive->deviceID().toStdWString().c_str(), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL); + if (drive == INVALID_HANDLE_VALUE) { + TCHAR message[256]; + FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); + m_err << tr("Couldn't open the drive for data verification for %1:").arg(m_wmiDiskDrive->deviceID()) << " (" << QString::fromWCharArray(message).trimmed() << ")\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(); + m_out << "DONE\n"; + m_out.flush(); + m_err << "OK\n"; + m_err.flush(); qApp->exit(0); break; case ISOMD5SUM_CHECK_FAILED: - err << tr("Your drive is probably damaged.") << "\n"; - err.flush(); + m_err << tr("Your drive is probably damaged.") << "\n"; + m_err.flush(); qApp->exit(1); return false; default: - err << tr("Unexpected error occurred during media check.") << "\n"; - err.flush(); + m_err << tr("Unexpected error occurred during media check.") << "\n"; + m_err.flush(); qApp->exit(1); return false; } diff --git a/src/helper/win/writejob.h b/src/helper/win/writejob.h index a5aa222b..f8533d02 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,48 @@ #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 lockDrive(HANDLE driveHandle); + 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_wmi; + std::unique_ptr m_wmiDiskDrive; - 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..2142e16d 100644 --- a/src/lib/CMakeLists.txt +++ b/src/lib/CMakeLists.txt @@ -1 +1,4 @@ add_subdirectory(isomd5) +if (WIN32) + add_subdirectory(libwmi) +endif() diff --git a/src/lib/libwmi/CMakeLists.txt b/src/lib/libwmi/CMakeLists.txt new file mode 100644 index 00000000..221bff5c --- /dev/null +++ b/src/lib/libwmi/CMakeLists.txt @@ -0,0 +1,9 @@ +set(LIBWMI_SRCS + libwmi.cpp +) + +add_library(libwmi STATIC ${LIBWMI_SRCS}) + +target_link_libraries(libwmi + Qt6::Core +) diff --git a/src/lib/libwmi/libwmi.cpp b/src/lib/libwmi/libwmi.cpp new file mode 100644 index 00000000..92674395 --- /dev/null +++ b/src/lib/libwmi/libwmi.cpp @@ -0,0 +1,361 @@ +/* + * 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 "libwmi.h" + +#include + +#include +#include + +#pragma comment(lib, "wbemuuid.lib") + +LibWMI::LibWMI(QObject *parent) + : QObject(parent) +{ + HRESULT res = S_OK; + // This needs to be initialized here before any RPC communication occurs + 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); + qWarning() << "Failed to initialize security. Error = " << err.ErrorMessage(); + return; + } + + res = CoCreateInstance(CLSID_WbemAdministrativeLocator, NULL, CLSCTX_INPROC_SERVER, IID_IWbemLocator, reinterpret_cast(&m_IWbemLocator)); + if (FAILED(res)) { + _com_error err(res); + qWarning() << "Failed to create IWbemLocator object. Error = " << QString::fromWCharArray(err.ErrorMessage()); + return; + } + + res = m_IWbemLocator->ConnectServer(_bstr_t(L"root\\cimv2"), NULL, NULL, NULL, 0, NULL, NULL, &m_IWbemServices); + if (FAILED(res)) { + _com_error err(res); + qWarning() << "Could not connect to WMI. Error = " << QString::fromWCharArray(err.ErrorMessage()); + return; + } + + initialized = true; +} + +LibWMI::~LibWMI() +{ + if (m_IWbemLocator) { + m_IWbemLocator->Release(); + } + if (m_IWbemServices) { + m_IWbemServices->Release(); + } + CoUninitialize(); +} + +QMap LibWMI::getUSBDeviceList() +{ + QMap result; + if (!initialized) { + return result; + } + + HRESULT res = S_OK; + IEnumWbemClassObject *pEnumDiskObjects = NULL; + + res = m_IWbemServices->ExecQuery(_bstr_t(L"WQL"), _bstr_t(L"SELECT * FROM Win32_DiskDrive WHERE InterfaceType = 'USB'"), WBEM_FLAG_RETURN_IMMEDIATELY, NULL, &pEnumDiskObjects); + if (FAILED(res)) { + _com_error err(res); + qWarning() << "WMI query failed. Error = " << 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; + } + + // Fetch disk information + quint32 index; + QString deviceID; + VARIANT var; + + if ((pDiskObject->Get(_bstr_t(L"Index"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_I4) { + index = var.intVal; + qDebug() << "Disk Index: " << index; + } else if (var.vt == VT_UI4) { + index = var.uintVal; + qDebug() << "Disk Index: " << index; + } + VariantClear(&var); + } + + if ((pDiskObject->Get(_bstr_t(L"DeviceID"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_BSTR) { + deviceID = QString::fromWCharArray(var.bstrVal); + qDebug() << "Device ID" << deviceID; + } + VariantClear(&var); + } + pDiskObject->Release(); + + result.insert(index, deviceID); + } + pEnumDiskObjects->Release(); + + return result; +} + +QStringList LibWMI::getDevicePartitions(quint32 index) +{ + QStringList result; + if (!initialized) { + return result; + } + + HRESULT res = S_OK; + IEnumWbemClassObject *pPartitionObjects = NULL; + std::wstring partitionQuery = L"SELECT * FROM Win32_DiskPartition WHERE DiskIndex = " + 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); + qWarning() << "Query for disk partitions failed. Error = " << 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; + } + + QString deviceID; + VARIANT partitionDeviceID; + if ((pPartitionObject->Get(_bstr_t(L"DeviceID"), 0, &partitionDeviceID, 0, 0)) == WBEM_S_NO_ERROR) { + if (partitionDeviceID.vt == VT_BSTR) { + deviceID = QString::fromWCharArray(partitionDeviceID.bstrVal); + qDebug() << "Partition device ID " << deviceID; + } + VariantClear(&partitionDeviceID); + } + pPartitionObject->Release(); + if (!deviceID.isEmpty()) { + result << deviceID; + } + } + pPartitionObjects->Release(); + + return result; +} + +QStringList LibWMI::getLogicalDisks(const QString &partitionID) +{ + QStringList result; + if (!initialized) { + return result; + } + + HRESULT res = S_OK; + std::wstring partitionToLogicalQuery = L"ASSOCIATORS OF {Win32_DiskPartition.DeviceID='" + partitionID.toStdWString() + L"'} WHERE AssocClass = Win32_LogicalDiskToPartition"; + IEnumWbemClassObject *pLogicalDiskObjects = NULL; + res = m_IWbemServices->ExecQuery(_bstr_t(L"WQL"), _bstr_t(partitionToLogicalQuery.c_str()), WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY, NULL, &pLogicalDiskObjects); + if (FAILED(res)) { + _com_error err(res); + qWarning() << "Query for logical disks failed. Error = " << QString::fromWCharArray(err.ErrorMessage()); + return result; + } + + while (true) { + IWbemClassObject *pLogicalDiskObject = NULL; + ULONG uReturnLogicalDisk = 0; + pLogicalDiskObjects->Next(WBEM_INFINITE, 1, &pLogicalDiskObject, &uReturnLogicalDisk); + if (uReturnLogicalDisk == 0) { + break; + } + + QString deviceID; + VARIANT logicalDiskDeviceID; + if ((pLogicalDiskObject->Get(_bstr_t(L"DeviceID"), 0, &logicalDiskDeviceID, 0, 0)) == WBEM_S_NO_ERROR) { + if (logicalDiskDeviceID.vt == VT_BSTR) { + deviceID = QString::fromWCharArray(logicalDiskDeviceID.bstrVal); + qDebug() << "Logical disk device ID " << deviceID; + } + VariantClear(&logicalDiskDeviceID); + } + pLogicalDiskObject->Release(); + if (!deviceID.isEmpty()) { + result << deviceID; + } + } + pLogicalDiskObjects->Release(); + + return result; +} + +std::unique_ptr LibWMI::getDiskDriveInformation(quint32 index, const QString &deviceID) +{ + std::unique_ptr result = std::make_unique(index, deviceID); + if (!initialized) { + return result; + } + + HRESULT res = S_OK; + IEnumWbemClassObject *pEnumDiskObjects = NULL; + + std::wstring deviceQuery = L"SELECT * FROM Win32_DiskDrive WHERE Index = " + 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); + qWarning() << "WMI query failed. Error = " << 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->deviceID().isEmpty()) { + if ((pDiskObject->Get(_bstr_t(L"DeviceID"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_BSTR) { + result->setDeviceID(QString::fromWCharArray(var.bstrVal)); + qDebug() << "DeviceID " << result->deviceID(); + } + VariantClear(&var); + } + } + + if ((pDiskObject->Get(_bstr_t(L"Model"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_BSTR) { + result->setModel(QString::fromWCharArray(var.bstrVal)); + qDebug() << "Disk model: " << result->model(); + } + 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()); + qDebug() << "Size " << result->size(); + } else if (var.vt == VT_I4) { + result->setSize(var.intVal); + qDebug() << "Size " << result->size(); + } else if (var.vt == VT_UI4) { + result->setSize(var.uintVal); + qDebug() << "Size " << result->size(); + } + VariantClear(&var); + } + + if ((pDiskObject->Get(_bstr_t(L"BytesPerSector"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) { + if (var.vt == VT_BSTR) { + result->setSectorSize(QString::fromWCharArray(var.bstrVal).toULongLong()); + qDebug() << "Sector size " << result->sectorSize(); + } else if (var.vt == VT_I4) { + result->setSectorSize(var.intVal); + qDebug() << "Sector size " << result->sectorSize(); + } else if (var.vt == VT_UI4) { + result->setSectorSize(var.uintVal); + qDebug() << "Sector size " << 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)); + qDebug() << "Serial number " << result->serialNumber(); + } + VariantClear(&var); + } + pDiskObject->Release(); + } + pEnumDiskObjects->Release(); + + return result; +} + +LibWMIDiskDrive::LibWMIDiskDrive(quint32 index, const QString &deviceID) + : m_index(index) + , m_deviceID(deviceID) +{ +} + +quint32 LibWMIDiskDrive::index() const +{ + return m_index; +} + +QString LibWMIDiskDrive::deviceID() const +{ + return m_deviceID; +} + +void LibWMIDiskDrive::setDeviceID(const QString &deviceID) +{ + m_deviceID = deviceID; +} + +QString LibWMIDiskDrive::model() const +{ + return m_model; +} + +void LibWMIDiskDrive::setModel(const QString &model) +{ + m_model = model; +} + +quint64 LibWMIDiskDrive::size() const +{ + return m_size; +} + +void LibWMIDiskDrive::setSize(quint64 size) +{ + m_size = size; +} + +QString LibWMIDiskDrive::serialNumber() const +{ + return m_serialNumber; +} + +void LibWMIDiskDrive::setSerialNumber(const QString &serialNumber) +{ + m_serialNumber = serialNumber; +} + +quint32 LibWMIDiskDrive::sectorSize() +{ + return m_sectorSize; +} + +void LibWMIDiskDrive::setSectorSize(quint32 sectorSize) +{ + m_sectorSize = sectorSize; +} diff --git a/src/lib/libwmi/libwmi.h b/src/lib/libwmi/libwmi.h new file mode 100644 index 00000000..fa6ef485 --- /dev/null +++ b/src/lib/libwmi/libwmi.h @@ -0,0 +1,81 @@ +/* + * 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 LIBWMI_H +#define LIBWMI_H + +#include +#include + +class LibWMIDiskDrive; + +class LibWMI : public QObject +{ + Q_OBJECT +public: + LibWMI(QObject *parent); + ~LibWMI(); + + // Returns a map with list of devices + QMap getUSBDeviceList(); + // Returns a list of partition IDs for device with given index + QStringList getDevicePartitions(quint32 index); + // Returns a list of logical disks IDs for given partition + QStringList getLogicalDisks(const QString &partitionID); + // Returns information about disk drive + std::unique_ptr getDiskDriveInformation(quint32 index, const QString &deviceID = QString()); + +private: + bool initialized = false; + IWbemLocator *m_IWbemLocator = NULL; + IWbemServices *m_IWbemServices = NULL; +}; + +class LibWMIDiskDrive +{ +public: + LibWMIDiskDrive(quint32 index, const QString &deviceID = QString()); + + quint32 index() const; + + QString deviceID() const; + void setDeviceID(const QString &deviceID); + + QString model() const; + void setModel(const QString &model); + + quint64 size() const; + void setSize(quint64 size); + + QString serialNumber() const; + void setSerialNumber(const QString &serialNumber); + + quint32 sectorSize(); + void setSectorSize(quint32 sectorSize); + +private: + quint32 m_index = 0; + quint32 m_sectorSize = 0; + quint64 m_size = 0; + QString m_deviceID; + QString m_model; + QString m_serialNumber; +}; + +#endif // LIBWMI_H