From 0c432bc854837c483f2a2f5766552a9ecc0ebee1 Mon Sep 17 00:00:00 2001
From: Romain Vimont <rom@rom1v.com>
Date: Fri, 6 Sep 2024 23:08:08 +0200
Subject: [PATCH] Unregister all AOA devices automatically on exit

Pushing a close event from the keyboard_aoa or mouse_aoa implementation
was racy, because the AOA thread might be stopped before these events
were processed.

Instead, keep the list of open AOA devices to close them automatically
from the AOA thread before exiting.
---
 app/src/usb/aoa_hid.c      | 43 ++++++++++++++++++++++++++++++++++----
 app/src/usb/keyboard_aoa.c |  9 ++------
 app/src/usb/mouse_aoa.c    |  9 ++------
 3 files changed, 43 insertions(+), 18 deletions(-)

diff --git a/app/src/usb/aoa_hid.c b/app/src/usb/aoa_hid.c
index 621dc289..fd2e7ae3 100644
--- a/app/src/usb/aoa_hid.c
+++ b/app/src/usb/aoa_hid.c
@@ -7,6 +7,7 @@
 #include "aoa_hid.h"
 #include "util/log.h"
 #include "util/str.h"
+#include "util/vector.h"
 
 // See <https://source.android.com/devices/accessories/aoa2#hid-support>.
 #define ACCESSORY_REGISTER_HID 54
@@ -19,6 +20,8 @@
 // Drop droppable events above this limit
 #define SC_AOA_EVENT_QUEUE_LIMIT 60
 
+struct sc_vec_hid_ids SC_VECTOR(uint16_t);
+
 static void
 sc_hid_input_log(const struct sc_hid_input *hid_input) {
     // HID input: [00] FF FF FF FF...
@@ -305,7 +308,8 @@ sc_aoa_push_close(struct sc_aoa *aoa, const struct sc_hid_close *hid_close) {
 }
 
 static bool
-sc_aoa_process_event(struct sc_aoa *aoa, struct sc_aoa_event *event) {
+sc_aoa_process_event(struct sc_aoa *aoa, struct sc_aoa_event *event,
+                     struct sc_vec_hid_ids *vec_open) {
     switch (event->type) {
         case SC_AOA_EVENT_TYPE_INPUT: {
             uint64_t ack_to_wait = event->input.ack_to_wait;
@@ -346,7 +350,16 @@ sc_aoa_process_event(struct sc_aoa *aoa, struct sc_aoa_event *event) {
             bool ok = sc_aoa_setup_hid(aoa, hid_open->hid_id,
                                        hid_open->report_desc,
                                        hid_open->report_desc_size);
-            if (!ok) {
+            if (ok) {
+                // The device is now open, add it to the list of devices to
+                // close automatically on exit
+                bool pushed = sc_vector_push(vec_open, hid_open->hid_id);
+                if (!pushed) {
+                    LOG_OOM();
+                    // this is not fatal, the HID device will just not be
+                    // explicitly unregistered
+                }
+            } else {
                 LOGW("Could not open AOA device: %" PRIu16, hid_open->hid_id);
             }
 
@@ -355,7 +368,14 @@ sc_aoa_process_event(struct sc_aoa *aoa, struct sc_aoa_event *event) {
         case SC_AOA_EVENT_TYPE_CLOSE: {
             struct sc_hid_close *hid_close = &event->close.hid;
             bool ok = sc_aoa_unregister_hid(aoa, hid_close->hid_id);
-            if (!ok) {
+            if (ok) {
+                // The device is not open anymore, remove it from the list of
+                // devices to close automatically on exit
+                ssize_t idx = sc_vector_index_of(vec_open, hid_close->hid_id);
+                if (idx >= 0) {
+                    sc_vector_remove(vec_open, idx);
+                }
+            } else {
                 LOGW("Could not close AOA device: %" PRIu16, hid_close->hid_id);
             }
 
@@ -371,6 +391,9 @@ static int
 run_aoa_thread(void *data) {
     struct sc_aoa *aoa = data;
 
+    // Store the HID ids of opened devices to unregister them all before exiting
+    struct sc_vec_hid_ids vec_open = SC_VECTOR_INITIALIZER;
+
     for (;;) {
         sc_mutex_lock(&aoa->mutex);
         while (!aoa->stopped && sc_vecdeque_is_empty(&aoa->queue)) {
@@ -386,12 +409,24 @@ run_aoa_thread(void *data) {
         struct sc_aoa_event event = sc_vecdeque_pop(&aoa->queue);
         sc_mutex_unlock(&aoa->mutex);
 
-        bool cont = sc_aoa_process_event(aoa, &event);
+        bool cont = sc_aoa_process_event(aoa, &event, &vec_open);
         if (!cont) {
             // stopped
             break;
         }
     }
+
+    // Explicitly unregister all registered HID ids before exiting
+    for (size_t i = 0; i < vec_open.size; ++i) {
+        uint16_t hid_id = vec_open.data[i];
+        LOGD("Unregistering AOA device %" PRIu16 "...", hid_id);
+        bool ok = sc_aoa_unregister_hid(aoa, hid_id);
+        if (!ok) {
+            LOGW("Could not close AOA device: %" PRIu16, hid_id);
+        }
+    }
+    sc_vector_destroy(&vec_open);
+
     return 0;
 }
 
diff --git a/app/src/usb/keyboard_aoa.c b/app/src/usb/keyboard_aoa.c
index 738f6875..b7834b0f 100644
--- a/app/src/usb/keyboard_aoa.c
+++ b/app/src/usb/keyboard_aoa.c
@@ -98,11 +98,6 @@ sc_keyboard_aoa_init(struct sc_keyboard_aoa *kb, struct sc_aoa *aoa) {
 
 void
 sc_keyboard_aoa_destroy(struct sc_keyboard_aoa *kb) {
-    struct sc_hid_close hid_close;
-    sc_hid_keyboard_generate_close(&hid_close);
-
-    bool ok = sc_aoa_push_close(kb->aoa, &hid_close);
-    if (!ok) {
-        LOGW("Could not push AOA HID close (keyboard)");
-    }
+    (void) kb;
+    // Do nothing, kb->aoa will automatically unregister all devices
 }
diff --git a/app/src/usb/mouse_aoa.c b/app/src/usb/mouse_aoa.c
index b4eb4eb8..33b777c4 100644
--- a/app/src/usb/mouse_aoa.c
+++ b/app/src/usb/mouse_aoa.c
@@ -78,11 +78,6 @@ sc_mouse_aoa_init(struct sc_mouse_aoa *mouse, struct sc_aoa *aoa) {
 
 void
 sc_mouse_aoa_destroy(struct sc_mouse_aoa *mouse) {
-    struct sc_hid_close hid_close;
-    sc_hid_mouse_generate_close(&hid_close);
-
-    bool ok = sc_aoa_push_close(mouse->aoa, &hid_close);
-    if (!ok) {
-        LOGW("Could not push AOA HID close (mouse)");
-    }
+    (void) mouse;
+    // Do nothing, mouse->aoa will automatically unregister all devices
 }