From 181d211d5fd8f0b1b2762f6a136f7e61cb39d2aa Mon Sep 17 00:00:00 2001
From: Nodir Temirkhodjaev <nodir.temir@gmail.com>
Date: Tue, 9 Jan 2024 14:33:36 +0300
Subject: [PATCH] UI: DriveListManager: Filter mounted drives

---
 src/tests/LogReaderTest/tst_logreader.h |  2 +-
 src/ui/driver/driverworker.cpp          |  2 +-
 src/ui/manager/drivelistmanager.cpp     | 10 +++++-
 src/ui/manager/drivelistmanager.h       |  2 ++
 src/ui/util/device.cpp                  |  2 +-
 src/ui/util/device.h                    |  2 +-
 src/ui/util/fileutil.cpp                | 41 +++++++++++++++++++++++++
 src/ui/util/fileutil.h                  |  2 ++
 8 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/src/tests/LogReaderTest/tst_logreader.h b/src/tests/LogReaderTest/tst_logreader.h
index 38820bdd..df87ebdc 100644
--- a/src/tests/LogReaderTest/tst_logreader.h
+++ b/src/tests/LogReaderTest/tst_logreader.h
@@ -114,7 +114,7 @@ TEST_F(LogReaderTest, logRead)
     LogBuffer buf(DriverCommon::bufferSize());
 
     for (;;) {
-        int nr;
+        qsizetype nr;
         QByteArray &array = buf.array();
         ASSERT_TRUE(device.ioctl(
                 DriverCommon::ioctlGetLog(), nullptr, 0, array.data(), array.size(), &nr));
diff --git a/src/ui/driver/driverworker.cpp b/src/ui/driver/driverworker.cpp
index ff1baf21..256b5ac4 100644
--- a/src/ui/driver/driverworker.cpp
+++ b/src/ui/driver/driverworker.cpp
@@ -110,7 +110,7 @@ void DriverWorker::readLog()
         return;
 
     QByteArray &array = m_logBuffer->array();
-    int nr;
+    qsizetype nr;
 
     const bool success = m_device->ioctl(
             DriverCommon::ioctlGetLog(), nullptr, 0, array.data(), array.size(), &nr);
diff --git a/src/ui/manager/drivelistmanager.cpp b/src/ui/manager/drivelistmanager.cpp
index 6abbb83d..66a05d4c 100644
--- a/src/ui/manager/drivelistmanager.cpp
+++ b/src/ui/manager/drivelistmanager.cpp
@@ -18,7 +18,10 @@ void DriveListManager::initialize()
 
 void DriveListManager::onDriveListChanged()
 {
-    const quint32 driveMask = FileUtil::driveMask();
+    quint32 driveMask = FileUtil::driveMask();
+    if (m_checkMounted) {
+        driveMask = FileUtil::mountedDriveMask(driveMask);
+    }
 
     if (m_driveMask == driveMask)
         return;
@@ -35,6 +38,11 @@ void DriveListManager::onDriveListChanged()
 
 void DriveListManager::startPolling()
 {
+    if (m_checkMounted)
+        return;
+
+    m_checkMounted = true;
+
     setupPollingTimer();
 
     m_pollingTimer->start();
diff --git a/src/ui/manager/drivelistmanager.h b/src/ui/manager/drivelistmanager.h
index c76156e1..d9831249 100644
--- a/src/ui/manager/drivelistmanager.h
+++ b/src/ui/manager/drivelistmanager.h
@@ -28,6 +28,8 @@ private:
     void setupPollingTimer();
 
 private:
+    bool m_checkMounted = false;
+
     quint32 m_driveMask = 0;
 
     QTimer *m_pollingTimer = nullptr;
diff --git a/src/ui/util/device.cpp b/src/ui/util/device.cpp
index 442e74f3..1bea72bc 100644
--- a/src/ui/util/device.cpp
+++ b/src/ui/util/device.cpp
@@ -49,7 +49,7 @@ bool Device::cancelIo()
     return CancelIoEx(m_handle, nullptr);
 }
 
-bool Device::ioctl(quint32 code, char *in, int inSize, char *out, int outSize, int *retSize)
+bool Device::ioctl(quint32 code, char *in, int inSize, char *out, int outSize, qsizetype *retSize)
 {
     LPOVERLAPPED overlapped = isOverlapped() ? (LPOVERLAPPED) m_buffer.data() : nullptr;
 
diff --git a/src/ui/util/device.h b/src/ui/util/device.h
index cfc70614..75a2856b 100644
--- a/src/ui/util/device.h
+++ b/src/ui/util/device.h
@@ -33,7 +33,7 @@ public slots:
     bool cancelIo();
 
     bool ioctl(quint32 code, char *in = nullptr, int inSize = 0, char *out = nullptr,
-            int outSize = 0, int *retSize = nullptr);
+            int outSize = 0, qsizetype *retSize = nullptr);
 
     void initOverlapped(void *eventHandle = nullptr);
 
diff --git a/src/ui/util/fileutil.cpp b/src/ui/util/fileutil.cpp
index d33e92cd..c0772596 100644
--- a/src/ui/util/fileutil.cpp
+++ b/src/ui/util/fileutil.cpp
@@ -8,6 +8,7 @@
 
 #define WIN32_LEAN_AND_MEAN
 #include <qt_windows.h>
+#include <winioctl.h>
 
 namespace FileUtil {
 
@@ -53,6 +54,46 @@ quint32 driveMask()
     return GetLogicalDrives();
 }
 
+static bool isDriveMounted(WCHAR drive)
+{
+    const WCHAR volume[] = { L'\\', L'\\', L'.', L'\\', drive, L':', L'\0' };
+
+    const DWORD shareMode = FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE;
+    const HANDLE volumeHandle =
+            CreateFileW(volume, GENERIC_READ, shareMode, nullptr, OPEN_EXISTING, 0, nullptr);
+
+    if (volumeHandle == INVALID_HANDLE_VALUE)
+        return false;
+
+    DWORD nr;
+    const bool ok = DeviceIoControl(
+            volumeHandle, FSCTL_IS_VOLUME_MOUNTED, nullptr, 0, nullptr, 0, &nr, nullptr);
+
+    CloseHandle(volumeHandle);
+
+    return ok;
+}
+
+quint32 mountedDriveMask(quint32 driveMask)
+{
+    quint32 mask = driveMask;
+    while (mask != 0) {
+        unsigned long index;
+        if (!_BitScanForward(&index, mask))
+            break;
+
+        const quint32 bit = (1u << index);
+
+        if (!isDriveMounted(L'A' + index)) {
+            driveMask ^= bit;
+        }
+
+        mask ^= bit;
+    }
+
+    return driveMask;
+}
+
 quint32 driveMaskByPath(const QString &path)
 {
     if (!isDriveFilePath(path))
diff --git a/src/ui/util/fileutil.h b/src/ui/util/fileutil.h
index 9b947e6d..3b2a1c07 100644
--- a/src/ui/util/fileutil.h
+++ b/src/ui/util/fileutil.h
@@ -17,6 +17,8 @@ bool isDriveFilePath(const QString &path);
 
 quint32 driveMask();
 
+quint32 mountedDriveMask(quint32 driveMask);
+
 quint32 driveMaskByPath(const QString &path);
 
 // Convert DOS device name to drive letter (A: .. Z:)