From a0631c21e675918424ee79cd8331867e18f6404b Mon Sep 17 00:00:00 2001
From: Nav <nav@navnav.me>
Date: Sun, 29 Dec 2024 04:41:36 +0000
Subject: [PATCH] Tidying WCH-LinkE IAP exit code

---
 src/DebugToolDrivers/USB/UsbDevice.cpp   |  5 ++-
 src/DebugToolDrivers/USB/UsbDevice.hpp   |  1 +
 src/DebugToolDrivers/WCH/WchLinkBase.cpp | 40 ++++++++----------------
 3 files changed, 18 insertions(+), 28 deletions(-)

diff --git a/src/DebugToolDrivers/USB/UsbDevice.cpp b/src/DebugToolDrivers/USB/UsbDevice.cpp
index 5a96e5d6..010c2214 100644
--- a/src/DebugToolDrivers/USB/UsbDevice.cpp
+++ b/src/DebugToolDrivers/USB/UsbDevice.cpp
@@ -49,7 +49,6 @@ namespace Usb
             };
         }
 
-        // For now, just use the first device found.
         this->libusbDevice.swap(devices.front());
         ::libusb_device_handle* deviceHandle = nullptr;
 
@@ -120,6 +119,10 @@ namespace Usb
         return device;
     }
 
+    bool UsbDevice::devicePresent(std::uint16_t vendorId, std::uint16_t productId) {
+        return !UsbDevice::findMatchingDevices(vendorId, productId).empty();
+    }
+
     bool UsbDevice::waitForDevice(std::uint16_t vendorId, std::uint16_t productId, std::chrono::milliseconds timeout) {
         static constexpr auto DELAY = std::chrono::milliseconds{50};
         for (auto i = 0; (i * DELAY) <= timeout; ++i){
diff --git a/src/DebugToolDrivers/USB/UsbDevice.hpp b/src/DebugToolDrivers/USB/UsbDevice.hpp
index 9267b5e4..7f819292 100644
--- a/src/DebugToolDrivers/USB/UsbDevice.hpp
+++ b/src/DebugToolDrivers/USB/UsbDevice.hpp
@@ -69,6 +69,7 @@ namespace Usb
         virtual void setConfiguration(std::uint8_t configurationIndex);
 
         static std::optional<UsbDevice> tryDevice(std::uint16_t vendorId, std::uint16_t productId);
+        static bool devicePresent(std::uint16_t vendorId, std::uint16_t productId);
         static bool waitForDevice(std::uint16_t vendorId, std::uint16_t productId, std::chrono::milliseconds timeout);
 
     protected:
diff --git a/src/DebugToolDrivers/WCH/WchLinkBase.cpp b/src/DebugToolDrivers/WCH/WchLinkBase.cpp
index b76f5767..19b439ad 100644
--- a/src/DebugToolDrivers/WCH/WchLinkBase.cpp
+++ b/src/DebugToolDrivers/WCH/WchLinkBase.cpp
@@ -4,7 +4,6 @@
 
 #include "Protocols/WchLink/WchLinkInterface.hpp"
 
-#include "src/TargetController/Exceptions/DeviceNotFound.hpp"
 #include "src/TargetController/Exceptions/DeviceInitializationFailure.hpp"
 
 #include "src/Logger/Logger.hpp"
@@ -31,37 +30,24 @@ namespace DebugToolDrivers::Wch
     {}
 
     void WchLinkBase::init() {
-        using Exceptions::DeviceNotFound;
-
-        try {
-            UsbDevice::init();
-
-        } catch (const DeviceNotFound& exception) {
+        if (this->toolConfig.exitIapMode && !UsbDevice::devicePresent(this->vendorId, this->productId)) {
             auto iapDevice = Usb::UsbDevice::tryDevice(this->iapVendorId, this->iapProductId);
 
-            if (!iapDevice.has_value()) {
-                throw exception;
+            if (iapDevice.has_value()) {
+                Logger::warning("Found device in IAP mode - attempting exit operation");
+                this->exitIapMode(*iapDevice);
+
+                Logger::info("Waiting for device to re-enumerate...");
+                if (!Usb::UsbDevice::waitForDevice(this->vendorId, this->productId, std::chrono::seconds{8})) {
+                    throw DeviceInitializationFailure{"Timeout exceeded whilst waiting for device to re-enumerate"};
+                }
+
+                Logger::info("Re-enumerated device found - IAP exit operation was successful");
             }
-
-            if (!this->toolConfig.exitIapMode) {
-                throw DeviceInitializationFailure{
-                    "Device found in IAP mode - Bloom can have the device exit this mode - see the 'exit_iap_mode' "
-                        "tool config parameter, for more"
-                };
-            }
-
-            Logger::warning("Found device in IAP mode - attempting exit operation");
-            this->exitIapMode(*iapDevice);
-
-            Logger::info("Waiting for device to re-enumerate...");
-            if (!Usb::UsbDevice::waitForDevice(this->vendorId, this->productId, std::chrono::seconds{8})) {
-                throw DeviceInitializationFailure{"Timeout exceeded whilst waiting for device to re-enumerate"};
-            }
-
-            Logger::info("Re-enumerated device found - IAP exit operation was successful");
-            UsbDevice::init();
         }
 
+        UsbDevice::init();
+
         this->detachKernelDriverFromInterface(this->wchLinkUsbInterfaceNumber);
 
         this->wchLinkUsbInterface = std::make_unique<Usb::UsbInterface>(