File indexing completed on 2024-04-28 09:43:51
0001 /* 0002 SPDX-FileCopyrightText: 2016 ROSA, 2023 Jonathan Esk-Riddell <jr@jriddell.org> 0003 SPDX-License-Identifier: GPL-3.0-or-later 0004 */ 0005 0006 //////////////////////////////////////////////////////////////////////////////// 0007 // Implementation of ImageWriter 0008 0009 #include "imagewriter.h" 0010 0011 #include <KLocalizedString> 0012 0013 #include <QtDBus/QtDBus> 0014 #include <QFile> 0015 #include <KCompressionDevice> 0016 0017 #include <fcntl.h> 0018 0019 typedef QHash<QString, QVariant> Properties; 0020 typedef QHash<QString, Properties> InterfacesAndProperties; 0021 typedef QHash<QDBusObjectPath, InterfacesAndProperties> DBusIntrospection; 0022 Q_DECLARE_METATYPE(Properties) 0023 Q_DECLARE_METATYPE(InterfacesAndProperties) 0024 Q_DECLARE_METATYPE(DBusIntrospection) 0025 0026 #include "common.h" 0027 #include "physicaldevice.h" 0028 0029 ImageWriter::ImageWriter(const QString& ImageFile, UsbDevice* Device, QObject *parent) : 0030 QObject(parent), 0031 m_Device(Device), 0032 m_ImageFile(ImageFile), 0033 m_CancelWriting(false) 0034 { 0035 } 0036 0037 // The main method that writes the image 0038 void ImageWriter::writeImage() 0039 { 0040 const qint64 TRANSFER_BLOCK_SIZE = 1024 * 1024; 0041 void* buffer = NULL; 0042 0043 bool isError = false; 0044 bool cancelRequested = false; 0045 bool zeroing = (m_ImageFile == ""); 0046 0047 // Using try-catch for processing errors 0048 // Invalid values are used for indication non-initialized objects; 0049 // after the try-catch block all the initialized objects are freed 0050 try 0051 { 0052 #if defined(Q_OS_WIN32) 0053 // Using VirtualAlloc so that the buffer was properly aligned (required for 0054 // direct access to devices and for unbuffered reading/writing) 0055 buffer = VirtualAlloc(NULL, TRANSFER_BLOCK_SIZE, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE); 0056 if (buffer == NULL) 0057 throw formatErrorMessageFromCode(i18n("Failed to allocate memory for buffer:")); 0058 #elif defined(Q_OS_LINUX) || defined(Q_OS_MAC) || defined(Q_OS_FREEBSD) 0059 buffer = malloc(TRANSFER_BLOCK_SIZE); 0060 if (buffer == NULL) 0061 throw i18n("Failed to allocate memory for buffer."); 0062 #endif 0063 0064 QFile imageFile; 0065 if (zeroing) 0066 { 0067 // Prepare zero-filled buffer 0068 memset(buffer, 0, TRANSFER_BLOCK_SIZE); 0069 } 0070 else 0071 { 0072 // Open the source image file for reading 0073 imageFile.setFileName(m_ImageFile); 0074 if (!imageFile.open(QIODevice::ReadOnly)) 0075 throw i18n("Failed to open the image file: %1", imageFile.errorString()); 0076 } 0077 0078 QIODevice* device; 0079 if (imageFile.fileName().endsWith(".gz")) { 0080 device = new KCompressionDevice(&imageFile, true, KCompressionDevice::GZip); 0081 if (!device->open(QIODevice::ReadOnly)) { 0082 throw i18n("Failed to open compression device: %1", device->errorString()); 0083 } 0084 } else if (imageFile.fileName().endsWith(".xz")) { 0085 device = new KCompressionDevice(&imageFile, true, KCompressionDevice::Xz); 0086 if (!device->open(QIODevice::ReadOnly)) { 0087 throw i18n("Failed to open compression device: %1", device->errorString()); 0088 } 0089 } else if (imageFile.fileName().endsWith(".zstd")) { 0090 device = new KCompressionDevice(&imageFile, true, KCompressionDevice::Zstd); 0091 if (!device->open(QIODevice::ReadOnly)) { 0092 throw i18n("Failed to open compression device: %1", device->errorString()); 0093 } 0094 } else { 0095 device = &imageFile; 0096 } 0097 0098 // Unmount volumes that belong to the selected target device 0099 // TODO: Check first if they are used and show warning 0100 // (problem: have to show request in the GUI thread and return reply back here) 0101 QStringList errMessages; 0102 0103 #if defined(Q_OS_WIN32) 0104 for (int i = 0; i < m_Device->m_Volumes.size(); ++i) 0105 { 0106 DWORD bret; 0107 HANDLE volume = CreateFile( 0108 reinterpret_cast<const wchar_t*>(("\\\\.\\" + m_Device->m_Volumes[i]).utf16()), 0109 GENERIC_READ | GENERIC_WRITE, 0110 FILE_SHARE_READ | FILE_SHARE_WRITE, 0111 NULL, 0112 OPEN_EXISTING, 0113 0, 0114 NULL 0115 ); 0116 if (volume == INVALID_HANDLE_VALUE) 0117 { 0118 errMessages << formatErrorMessageFromCode(i18n("Failed to open the drive %1", m_Device->m_Volumes[i])); 0119 continue; 0120 } 0121 // Trying to lock the volume but ignore if we failed (such call seems to be required for 0122 // dismounting the volume on WinXP) 0123 DeviceIoControl(volume, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &bret, NULL); 0124 if (!DeviceIoControl(volume, FSCTL_DISMOUNT_VOLUME, NULL, 0, NULL, 0, &bret, NULL)) 0125 errMessages << formatErrorMessageFromCode(i18n("Failed to unmount the drive %1", m_Device->m_Volumes[i])); 0126 CloseHandle(volume); 0127 volume = INVALID_HANDLE_VALUE; 0128 } 0129 #elif defined(Q_OS_MAC) 0130 struct statfs* mntEntries = NULL; 0131 int mntEntriesNum = getmntinfo(&mntEntries, MNT_WAIT); 0132 for (int i = 0; i < mntEntriesNum; ++i) 0133 { 0134 for (int j = 0; j < m_Device->m_Volumes.size(); ++j) 0135 { 0136 // Check that the mount point is either our target device itself or a partition on it 0137 if ((mntEntries[i].f_mntfromname == m_Device->m_Volumes[j]) || 0138 QString(mntEntries[i].f_mntfromname).startsWith(m_Device->m_Volumes[j] + 's')) 0139 { 0140 // Mount point is the selected device or one of its partitions - try to unmount it 0141 if (unmount(mntEntries[i].f_mntonname, MNT_FORCE) != 0) 0142 errMessages << i18n("Failed to unmount the volume %1\n%2", m_Device->m_Volumes[i], strerror(errno)); 0143 } 0144 } 0145 } 0146 #endif 0147 if (errMessages.size() > 0) 0148 throw errMessages.join("\n\n"); 0149 0150 #if defined(Q_OS_WIN32) || defined(Q_OS_MAC) 0151 // Open the target USB device for writing and lock it 0152 PhysicalDevice deviceFile(m_Device->m_PhysicalDevice); 0153 if (!deviceFile.open()) 0154 throw i18n("Failed to open the target device:\n%1", deviceFile.errorString()); 0155 #endif 0156 #if defined(Q_OS_LINUX) || defined(Q_OS_FREEBSD) 0157 QDBusInterface deviceDBus("org.freedesktop.UDisks2", m_Device->m_PhysicalDevice, "org.freedesktop.UDisks2.Block", QDBusConnection::systemBus(), this); 0158 QDBusReply<QDBusUnixFileDescriptor> reply = deviceDBus.call(QDBus::Block, "OpenDevice", "rw", Properties{{"flags", O_EXCL | O_SYNC | O_CLOEXEC}} ); 0159 QDBusUnixFileDescriptor fd = reply.value(); 0160 QFile deviceFile; 0161 deviceFile.open(fd.fileDescriptor(), QIODevice::WriteOnly); 0162 #endif 0163 0164 qint64 readBytes; 0165 qint64 writtenBytes; 0166 // Start reading/writing cycle 0167 for (;;) 0168 { 0169 if (zeroing) 0170 { 0171 readBytes = TRANSFER_BLOCK_SIZE; 0172 } 0173 else 0174 { 0175 if ((readBytes = device->read(static_cast<char*>(buffer), TRANSFER_BLOCK_SIZE)) <= 0) { 0176 break; 0177 } 0178 } 0179 // Align the number of bytes to the sector size 0180 readBytes = alignNumber(readBytes, (qint64)m_Device->m_SectorSize); 0181 writtenBytes = deviceFile.write(static_cast<char*>(buffer), readBytes); 0182 if (writtenBytes < 0) { 0183 qDebug() << "write writtenBytes: " << writtenBytes; 0184 //throw i18n("Failed to write to the device:\n%1"); //, "ook"); //deviceFile.errorString()); 0185 } 0186 if (writtenBytes != readBytes) 0187 throw i18n("The last block was not fully written (%1 of %2 bytes)!\nAborting.", writtenBytes, readBytes); 0188 #if defined(Q_OS_LINUX) || defined(Q_OS_MAC) || defined(Q_OS_FREEBSD) 0189 // In Linux/MacOS the USB device is opened with buffering. Using forced sync to validate progress bar. 0190 // For unknown reason, deviceFile.flush() does not work as intended here. 0191 fsync(deviceFile.handle()); 0192 #endif 0193 const int percent = (100 * imageFile.pos()) / imageFile.size(); 0194 // Inform the GUI thread that next block was written 0195 // TODO: Make sure that when TRANSFER_BLOCK_SIZE is not a multiple of DEFAULT_UNIT 0196 // this still works or at least fails compilation 0197 emit progressChanged(percent); 0198 0199 // Check for the cancel request (using temporary variable to avoid multiple unlock calls in the code) 0200 m_Mutex.lock(); 0201 cancelRequested = m_CancelWriting; 0202 m_Mutex.unlock(); 0203 0204 if (cancelRequested) 0205 { 0206 // The cancel request was issued 0207 emit cancelled(); 0208 break; 0209 } 0210 if (zeroing) 0211 { 0212 // In zeroing mode only write 1 block - 1 MB is enough to clear both MBR and GPT 0213 break; 0214 } 0215 } 0216 if (!zeroing) 0217 { 0218 if (readBytes < 0) { 0219 throw i18n("Failed to read the image file:\n%1", device->errorString()); 0220 } 0221 imageFile.close(); 0222 } 0223 deviceFile.close(); 0224 } 0225 catch (QString msg) 0226 { 0227 // Something went wrong :-( 0228 emit error(msg); 0229 isError = true; 0230 } 0231 0232 if (buffer != NULL) 0233 #if defined(Q_OS_WIN32) 0234 VirtualFree(buffer, TRANSFER_BLOCK_SIZE, MEM_DECOMMIT | MEM_RELEASE); 0235 #elif defined(Q_OS_LINUX) || defined(Q_OS_MAC) 0236 free(buffer); 0237 #endif 0238 0239 // If no errors occurred and user did not stop the operation, it means everything went fine 0240 if (!isError && !cancelRequested) { 0241 QString message = i18n("The operation completed successfully.") + 0242 "<br><br>" + 0243 (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\".")); 0244 emit success(message); 0245 } 0246 0247 // In any case the operation is finished 0248 emit finished(); 0249 } 0250 0251 // Implements reaction to the cancel request from user 0252 void ImageWriter::cancelWriting() 0253 { 0254 m_Mutex.lock(); 0255 m_CancelWriting = true; 0256 m_Mutex.unlock(); 0257 } 0258 0259 #include "moc_imagewriter.cpp"