From 2bbb8bfba78b2b86d0d6c6809ea43410686540c8 Mon Sep 17 00:00:00 2001
From: Floris Bos <bos@je-eigen-domein.nl>
Date: Sun, 21 Mar 2021 16:25:32 +0100
Subject: [PATCH] Windows: improve file system mounting for customization

- Check if config.txt exists on drive letter before proceeding to
  write changes. Wait up to 3 seconds if not.
- Force drive letter assignment for portable SSDs and other devices
  that pretend to be a fixed drive. (Windows only auto-assigns
  drive letters for removable media).

Ref #171
---
 CMakeLists.txt     |  4 +--
 downloadthread.cpp | 69 +++++++++++++++++++++++++++++++++++++++-------
 downloadthread.h   |  1 +
 3 files changed, 62 insertions(+), 12 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 3d7fc792..66af9a0a 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -8,8 +8,8 @@ endif()
 
 project(rpi-imager LANGUAGES CXX C)
 set(IMAGER_VERSION_MAJOR 1)
-set(IMAGER_VERSION_MINOR 6)
-set(IMAGER_VERSION_STR "${IMAGER_VERSION_MAJOR}.${IMAGER_VERSION_MINOR}")
+set(IMAGER_VERSION_MINOR 7)
+set(IMAGER_VERSION_STR "${IMAGER_VERSION_MAJOR}.${IMAGER_VERSION_MINOR}beta")
 set(IMAGER_VERSION_CSV "${IMAGER_VERSION_MAJOR},${IMAGER_VERSION_MINOR},0,0")
 add_definitions(-DIMAGER_VERSION_STR="${IMAGER_VERSION_STR}")
 add_definitions(-DIMAGER_VERSION_CSV=${IMAGER_VERSION_CSV})
diff --git a/downloadthread.cpp b/downloadthread.cpp
index 925204b8..4e51b96d 100644
--- a/downloadthread.cpp
+++ b/downloadthread.cpp
@@ -127,7 +127,7 @@ bool DownloadThread::_openAndPrepareDevice()
 
     if (std::regex_match(_filename.constData(), m, windriveregex))
     {
-        QByteArray _nr = QByteArray::fromStdString(m[1]);
+        _nr = QByteArray::fromStdString(m[1]);
 
         if (!_nr.isEmpty()) {
             qDebug() << "Removing partition table from Windows drive #" << _nr << "(" << _filename << ")";
@@ -811,6 +811,18 @@ bool DownloadThread::_customizeImage()
 
     emit preparationStatusUpdate(tr("Waiting for FAT partition to be mounted"));
 
+#ifdef Q_OS_WIN
+    qDebug() << "Running diskpart rescan";
+    QProcess proc;
+    proc.setProcessChannelMode(proc.MergedChannels);
+    proc.start("diskpart");
+    proc.waitForStarted();
+    proc.write("rescan\r\n");
+    proc.closeWriteChannel();
+    proc.waitForFinished();
+    qDebug() << proc.readAll();
+#endif
+
     /* See if OS auto-mounted the device */
     for (int tries = 0; tries < 3; tries++)
     {
@@ -826,6 +838,30 @@ bool DownloadThread::_customizeImage()
         }
     }
 
+#ifdef Q_OS_WIN
+    if (folder.isEmpty() && !_nr.isEmpty()) {
+        qDebug() << "Windows did not assign drive letter automatically. Ask diskpart to do so manually.";
+        proc.start("diskpart");
+        proc.waitForStarted();
+        proc.write("select disk "+_nr+"\r\n"
+                        "select partition 1\r\n"
+                        "assign\r\n");
+        proc.closeWriteChannel();
+        proc.waitForFinished();
+        qDebug() << proc.readAll();
+
+        auto l = Drivelist::ListStorageDevices();
+        for (auto i : l)
+        {
+            if (QByteArray::fromStdString(i.device).toLower() == devlower && i.mountpoints.size())
+            {
+                folder = QByteArray::fromStdString(i.mountpoints.front());
+                break;
+            }
+        }
+    }
+#endif
+
 #ifdef Q_OS_LINUX
     bool manualmount = false;
 
@@ -853,7 +889,7 @@ bool DownloadThread::_customizeImage()
             QTemporaryDir td;
             QStringList args;
             folder = td.path();
-            args << fatpartition << folder;
+            args << "-t" << "vfat" << fatpartition << folder;
 
             if (QProcess::execute("mount", args) != 0)
             {
@@ -883,14 +919,29 @@ bool DownloadThread::_customizeImage()
         return false;
     }
 
+    /* Some operating system take longer to complete mounting FAT32
+       wait up to 3 seconds for config.txt file to appear */
+    QString configFilename = folder+"/config.txt";
+    for (int tries = 0; tries < 3; tries++)
+    {
+        if (QFile::exists(configFilename))
+            break;
+        QThread::sleep(1);
+    }
+
+    if (!QFile::exists(configFilename))
+    {
+        emit error(tr("Unable to customize. File '%1' does not exist.").arg(configFilename));
+        return false;
+    }
+
     emit preparationStatusUpdate(tr("Customizing image"));
 
     if (!_firstrun.isEmpty())
     {
         QFile f(folder+"/firstrun.sh");
-        if (f.open(f.WriteOnly))
+        if (f.open(f.WriteOnly) && f.write(_firstrun) == _firstrun.length())
         {
-            f.write(_firstrun);
             f.close();
         }
         else
@@ -906,8 +957,8 @@ bool DownloadThread::_customizeImage()
         configItems.removeAll("");
         QByteArray config;
 
-        QFile f(folder+"/config.txt");
-        if (f.exists() && f.open(f.ReadOnly))
+        QFile f(configFilename);
+        if (f.open(f.ReadOnly))
         {
             config = f.readAll();
             f.close();
@@ -929,9 +980,8 @@ bool DownloadThread::_customizeImage()
             }
         }
 
-        if (f.open(f.WriteOnly))
+        if (f.open(f.WriteOnly) && f.write(config) == config.length())
         {
-            f.write(config);
             f.close();
         }
         else
@@ -953,9 +1003,8 @@ bool DownloadThread::_customizeImage()
         }
 
         cmdline += _cmdline;
-        if (f.open(f.WriteOnly))
+        if (f.open(f.WriteOnly) && f.write(cmdline) == cmdline.length())
         {
-            f.write(cmdline);
             f.close();
         }
         else
diff --git a/downloadthread.h b/downloadthread.h
index 103d3e47..c4100148 100644
--- a/downloadthread.h
+++ b/downloadthread.h
@@ -176,6 +176,7 @@ class DownloadThread : public QThread
 
 #ifdef Q_OS_WIN
     WinFile _file, _volumeFile;
+    QByteArray _nr;
 #elif defined(Q_OS_DARWIN)
     MacFile _file;
 #else