Fix use-after-free by keyboard commands when connection/controller is being shut down.

Author: Mysticial

SerialPrograms/Source/Controllers/KeyboardInput/KeyboardInput.cpp

@@ -27,8 +27,12 @@ void KeyboardInputController::start(){
         [this]{ thread_loop(); }
     );
 }
-void KeyboardInputController::stop(){
-    m_stop.store(true, std::memory_order_release);
+void KeyboardInputController::stop() noexcept{
+    bool expected = false;
+    if (!m_stop.compare_exchange_strong(expected, true)){
+        //  Already stopped.
+        return;
+    }
     {
         std::lock_guard<std::mutex> lg(m_sleep_lock);
         m_cv.notify_all();

SerialPrograms/Source/Controllers/KeyboardInput/KeyboardInput.h

@@ -51,8 +51,8 @@ class KeyboardInputController{
 
 
 protected:
-    void start();   //  Child class must call this at end of constructor.
-    void stop();    //  Child class must call this at start of destructor.
+    void start();           //  Child class must call this at end of constructor.
+    void stop() noexcept;   //  Child class must call this at start of destructor.
 
     virtual std::unique_ptr<ControllerState> make_state() const = 0;
     virtual void update_state(ControllerState& state, const std::set<uint32_t>& pressed_keys) = 0;

SerialPrograms/Source/Controllers/SerialPABotBase/SerialPABotBase_Connection.cpp

@@ -304,6 +304,7 @@ void SerialPABotBase_Connection::thread_body(){
             if (!m_ready.load(std::memory_order_relaxed)){
                 break;
             }
+
             m_cv.wait_for(lg, SERIAL_REFRESH_RATE);
         }
     });

SerialPrograms/Source/NintendoSwitch/Controllers/NintendoSwitch_ProController.cpp

@@ -32,7 +32,7 @@ class ProController::KeyboardManager : public KeyboardInputController{
 public:
     KeyboardManager(ProController& controller)
         : KeyboardInputController(true)
-        , m_controller(controller)
+        , m_controller(&controller)
     {
         std::vector<std::shared_ptr<EditableTableRow>> mapping =
             ConsoleSettings::instance().KEYBOARD_MAPPINGS.TABLE.current_refs();
@@ -45,6 +45,16 @@ class ProController::KeyboardManager : public KeyboardInputController{
     ~KeyboardManager(){
         stop();
     }
+    void stop() noexcept{
+        {
+            WriteSpinLock lg(m_lock);
+            if (m_controller == nullptr){
+                return;
+            }
+            m_controller = nullptr;
+        }
+        KeyboardInputController::stop();
+    }
 
     virtual std::unique_ptr<ControllerState> make_state() const override{
         return std::make_unique<SwitchControllerState>();
@@ -65,10 +75,18 @@ class ProController::KeyboardManager : public KeyboardInputController{
         deltas.to_state(static_cast<SwitchControllerState&>(state));
     }
     virtual void cancel_all_commands() override{
-        m_controller.cancel_all_commands();
+        WriteSpinLock lg(m_lock);
+        if (m_controller == nullptr){
+            return;
+        }
+        m_controller->cancel_all_commands();
     }
     virtual void replace_on_next_command() override{
-        m_controller.replace_on_next_command();
+        WriteSpinLock lg(m_lock);
+        if (m_controller == nullptr){
+            return;
+        }
+        m_controller->replace_on_next_command();
     }
     virtual void send_state(const ControllerState& state) override{
         const SwitchControllerState& switch_state = static_cast<const SwitchControllerState&>(state);
@@ -82,7 +100,11 @@ class ProController::KeyboardManager : public KeyboardInputController{
             COLOR_DARKGREEN
         );
 #endif
-        m_controller.issue_full_controller_state(
+        WriteSpinLock lg(m_lock);
+        if (m_controller == nullptr){
+            return;
+        }
+        m_controller->issue_full_controller_state(
             nullptr,
             switch_state.buttons,
             switch_state.dpad,
@@ -96,7 +118,8 @@ class ProController::KeyboardManager : public KeyboardInputController{
 
 
 private:
-    ProController& m_controller;
+    SpinLock m_lock;
+    ProController* m_controller;
     std::map<Qt::Key, ControllerDeltas> m_mapping;
 };
 
@@ -111,6 +134,9 @@ ProController::ProController(Milliseconds timing_variation)
 {
 
 }
+void ProController::stop() noexcept{
+    m_keyboard_manager->stop();
+}
 
 
 void ProController::keyboard_release_all(){

SerialPrograms/Source/NintendoSwitch/Controllers/NintendoSwitch_ProController.h

@@ -84,6 +84,9 @@ class ProController : public AbstractController{
     virtual ~ProController();
     ProController(Milliseconds timing_variation);
 
+    //  Must call before destruction begins.
+    void stop() noexcept;
+
     virtual Logger& logger() = 0;
 
     Milliseconds timing_variation() const{

SerialPrograms/Source/NintendoSwitch/Controllers/NintendoSwitch_ProController_SerialPABotBase.cpp

@@ -67,6 +67,9 @@ ProController_SerialPABotBase::ProController_SerialPABotBase(
 
     m_error_string = html_color_text("Missing Feature: " + missing_feature, COLOR_RED);
 }
+ProController_SerialPABotBase::~ProController_SerialPABotBase(){
+    stop();
+}
 
 
 

SerialPrograms/Source/NintendoSwitch/Controllers/NintendoSwitch_ProController_SerialPABotBase.h

@@ -16,7 +16,7 @@ namespace PokemonAutomation{
 namespace NintendoSwitch{
 
 
-class ProController_SerialPABotBase : public ProControllerWithScheduler{
+class ProController_SerialPABotBase final : public ProControllerWithScheduler{
 public:
     using ContextType = ProControllerContext;
 
@@ -26,6 +26,7 @@ class ProController_SerialPABotBase : public ProControllerWithScheduler{
         SerialPABotBase::SerialPABotBase_Connection& connection,
         const ControllerRequirements& requirements
     );
+    ~ProController_SerialPABotBase();
 
     virtual bool is_ready() const override{
         return m_serial && m_handle.is_ready() && m_error_string.empty();

SerialPrograms/Source/NintendoSwitch/Controllers/SysbotBase/SysbotBase_ProController.cpp

@@ -69,6 +69,7 @@ ProController_SysbotBase::ProController_SysbotBase(
 
 }
 ProController_SysbotBase::~ProController_SysbotBase(){
+    stop();
     m_stopping.store(true, std::memory_order_release);
     {
         std::lock_guard<std::mutex> lg(m_state_lock);

SerialPrograms/Source/NintendoSwitch/Controllers/SysbotBase/SysbotBase_ProController.h

@@ -19,7 +19,7 @@ namespace NintendoSwitch{
 
 
 
-class ProController_SysbotBase : public ProControllerWithScheduler{
+class ProController_SysbotBase final : public ProControllerWithScheduler{
 public:
     using ContextType = ProControllerContext;