File indexing completed on 2024-04-14 05:41:26

 /*
     SPDX-FileCopyrightText: 2016 ROSA, 2023 Jonathan Esk-Riddell <jr@jriddell.org>
     SPDX-License-Identifier: GPL-3.0-or-later
 */
 
 ////////////////////////////////////////////////////////////////////////////////
 // Implementation of ImageWriter
 
 #include "imagewriter.h"
 
 #include <KLocalizedString>
 
 #include <QtDBus/QtDBus>
 #include <QFile>
 #include <KCompressionDevice>
 
 #include <fcntl.h>
 
 typedef QHash<QString, QVariant> Properties;
 typedef QHash<QString, Properties> InterfacesAndProperties;
 typedef QHash<QDBusObjectPath, InterfacesAndProperties> DBusIntrospection;
 Q_DECLARE_METATYPE(Properties)
 Q_DECLARE_METATYPE(InterfacesAndProperties)
 Q_DECLARE_METATYPE(DBusIntrospection)
 
 #include "common.h"
 #include "physicaldevice.h"
 
 ImageWriter::ImageWriter(const QString& ImageFile, UsbDevice* Device, QObject *parent) :
     QObject(parent),
     m_Device(Device),
     m_ImageFile(ImageFile),
     m_CancelWriting(false)
 {
 }
 
 // The main method that writes the image
 void ImageWriter::writeImage()
 {
     const qint64 TRANSFER_BLOCK_SIZE = 1024 * 1024;
     void* buffer = NULL;
 
     bool isError = false;
     bool cancelRequested = false;
     bool zeroing = (m_ImageFile == "");
 
     // Using try-catch for processing errors
     // Invalid values are used for indication non-initialized objects;
     // after the try-catch block all the initialized objects are freed
     try
     {
 #if defined(Q_OS_WIN32)
         // Using VirtualAlloc so that the buffer was properly aligned (required for
         // direct access to devices and for unbuffered reading/writing)
         buffer = VirtualAlloc(NULL, TRANSFER_BLOCK_SIZE, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
         if (buffer == NULL)
             throw formatErrorMessageFromCode(i18n("Failed to allocate memory for buffer:"));
 #elif defined(Q_OS_LINUX) || defined(Q_OS_MAC) || defined(Q_OS_FREEBSD)
         buffer = malloc(TRANSFER_BLOCK_SIZE);
         if (buffer == NULL)
             throw i18n("Failed to allocate memory for buffer.");
 #endif
 
         QFile imageFile;
         if (zeroing)
         {
             // Prepare zero-filled buffer
             memset(buffer, 0, TRANSFER_BLOCK_SIZE);
         }
         else
         {
             // Open the source image file for reading
             imageFile.setFileName(m_ImageFile);
             if (!imageFile.open(QIODevice::ReadOnly))
                 throw i18n("Failed to open the image file: %1", imageFile.errorString());
         }
 
         QIODevice* device;
         if (imageFile.fileName().endsWith(".gz")) {
             device = new KCompressionDevice(&imageFile, true, KCompressionDevice::GZip);
             if (!device->open(QIODevice::ReadOnly)) {
                 throw i18n("Failed to open compression device: %1", device->errorString());
             }
         } else if (imageFile.fileName().endsWith(".xz")) {
             device = new KCompressionDevice(&imageFile, true, KCompressionDevice::Xz);
             if (!device->open(QIODevice::ReadOnly)) {
                 throw i18n("Failed to open compression device: %1", device->errorString());
             }
         } else if (imageFile.fileName().endsWith(".zstd")) {
             device = new KCompressionDevice(&imageFile, true, KCompressionDevice::Zstd);
             if (!device->open(QIODevice::ReadOnly)) {
                 throw i18n("Failed to open compression device: %1", device->errorString());
             }
         } else {
             device = &imageFile;
         }
 
         // Unmount volumes that belong to the selected target device
         // TODO: Check first if they are used and show warning
         // (problem: have to show request in the GUI thread and return reply back here)
         QStringList errMessages;
 
 #if defined(Q_OS_WIN32)
         for (int i = 0; i < m_Device->m_Volumes.size(); ++i)
         {
             DWORD bret;
             HANDLE volume = CreateFile(
                 reinterpret_cast<const wchar_t*>(("\\\\.\\" + m_Device->m_Volumes[i]).utf16()),
                 GENERIC_READ | GENERIC_WRITE,
                 FILE_SHARE_READ | FILE_SHARE_WRITE,
                 NULL,
                 OPEN_EXISTING,
                 0,
                 NULL
             );
             if (volume == INVALID_HANDLE_VALUE)
             {
                 errMessages << formatErrorMessageFromCode(i18n("Failed to open the drive %1", m_Device->m_Volumes[i]));
                 continue;
             }
             // Trying to lock the volume but ignore if we failed (such call seems to be required for
             // dismounting the volume on WinXP)
             DeviceIoControl(volume, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &bret, NULL);
             if (!DeviceIoControl(volume, FSCTL_DISMOUNT_VOLUME, NULL, 0, NULL, 0, &bret, NULL))
                 errMessages << formatErrorMessageFromCode(i18n("Failed to unmount the drive %1", m_Device->m_Volumes[i]));
             CloseHandle(volume);
             volume = INVALID_HANDLE_VALUE;
         }
 #elif defined(Q_OS_MAC)
         struct statfs* mntEntries = NULL;
         int mntEntriesNum = getmntinfo(&mntEntries, MNT_WAIT);
         for (int i = 0; i < mntEntriesNum; ++i)
         {
             for (int j = 0; j < m_Device->m_Volumes.size(); ++j)
             {
                 // Check that the mount point is either our target device itself or a partition on it
                 if ((mntEntries[i].f_mntfromname == m_Device->m_Volumes[j]) ||
                     QString(mntEntries[i].f_mntfromname).startsWith(m_Device->m_Volumes[j] + 's'))
                 {
                     // Mount point is the selected device or one of its partitions - try to unmount it
                     if (unmount(mntEntries[i].f_mntonname, MNT_FORCE) != 0)
                         errMessages << i18n("Failed to unmount the volume %1\n%2", m_Device->m_Volumes[i], strerror(errno));
                 }
             }
         }
 #endif
         if (errMessages.size() > 0)
             throw errMessages.join("\n\n");
 
 #if defined(Q_OS_WIN32) || defined(Q_OS_MAC)
         // Open the target USB device for writing and lock it
         PhysicalDevice deviceFile(m_Device->m_PhysicalDevice);
         if (!deviceFile.open())
             throw i18n("Failed to open the target device:\n%1", deviceFile.errorString());
 #endif
 #if defined(Q_OS_LINUX) || defined(Q_OS_FREEBSD)
         QDBusInterface deviceDBus("org.freedesktop.UDisks2", m_Device->m_PhysicalDevice, "org.freedesktop.UDisks2.Block", QDBusConnection::systemBus(), this);
         QDBusReply<QDBusUnixFileDescriptor> reply = deviceDBus.call(QDBus::Block, "OpenDevice", "rw", Properties{{"flags", O_EXCL | O_SYNC | O_CLOEXEC}} );
         QDBusUnixFileDescriptor fd = reply.value();
         QFile deviceFile;
         deviceFile.open(fd.fileDescriptor(), QIODevice::WriteOnly);
 #endif
 
         qint64 readBytes;
         qint64 writtenBytes;
         // Start reading/writing cycle
         for (;;)
         {
             if (zeroing)
             {
                 readBytes = TRANSFER_BLOCK_SIZE;
             }
             else
             {
                 if ((readBytes = device->read(static_cast<char*>(buffer), TRANSFER_BLOCK_SIZE)) <= 0) {
                     break;
                 }
             }
             // Align the number of bytes to the sector size
             readBytes = alignNumber(readBytes, (qint64)m_Device->m_SectorSize);
             writtenBytes = deviceFile.write(static_cast<char*>(buffer), readBytes);
             if (writtenBytes < 0) {
                 qDebug() << "write writtenBytes: " << writtenBytes;
                 //throw i18n("Failed to write to the device:\n%1"); //, "ook"); //deviceFile.errorString());
             }
             if (writtenBytes != readBytes)
                 throw i18n("The last block was not fully written (%1 of %2 bytes)!\nAborting.", writtenBytes, readBytes);
 #if defined(Q_OS_LINUX) || defined(Q_OS_MAC) || defined(Q_OS_FREEBSD)
             // In Linux/MacOS the USB device is opened with buffering. Using forced sync to validate progress bar.
             // For unknown reason, deviceFile.flush() does not work as intended here.
             fsync(deviceFile.handle());
 #endif
             const int percent = (100 * imageFile.pos()) / imageFile.size();
             // Inform the GUI thread that next block was written
             // TODO: Make sure that when TRANSFER_BLOCK_SIZE is not a multiple of DEFAULT_UNIT
             // this still works or at least fails compilation
             emit progressChanged(percent);
 
             // Check for the cancel request (using temporary variable to avoid multiple unlock calls in the code)
             m_Mutex.lock();
             cancelRequested = m_CancelWriting;
             m_Mutex.unlock();
 
             if (cancelRequested)
             {
                 // The cancel request was issued
                 emit cancelled();
                 break;
             }
             if (zeroing)
             {
                 // In zeroing mode only write 1 block - 1 MB is enough to clear both MBR and GPT
                 break;
             }
         }
         if (!zeroing)
         {
             if (readBytes < 0) {
                 throw i18n("Failed to read the image file:\n%1", device->errorString());
             }
             imageFile.close();
         }
         deviceFile.close();
     }
     catch (QString msg)
     {
         // Something went wrong :-(
         emit error(msg);
         isError = true;
     }
 
     if (buffer != NULL)
 #if defined(Q_OS_WIN32)
         VirtualFree(buffer, TRANSFER_BLOCK_SIZE, MEM_DECOMMIT | MEM_RELEASE);
 #elif defined(Q_OS_LINUX) || defined(Q_OS_MAC)
         free(buffer);
 #endif
 
     // If no errors occurred and user did not stop the operation, it means everything went fine
     if (!isError && !cancelRequested) {
         QString message = i18n("The operation completed successfully.") +
             "<br><br>" +
             (zeroing ? i18n("Now you need to format your device.") : i18n("To be able to store data on this device again, please, use the button \"Wipe USB Disk\"."));
         emit success(message);
     }
 
     // In any case the operation is finished
     emit finished();
 }
 
 // Implements reaction to the cancel request from user
 void ImageWriter::cancelWriting()
 {
     m_Mutex.lock();
     m_CancelWriting = true;
     m_Mutex.unlock();
 }
 
 #include "moc_imagewriter.cpp"