From 63af2727012b8f34439a6ff09edb005d7e5af054 Mon Sep 17 00:00:00 2001
From: Keeley Hoek <keeley@hoek.io>
Date: Fri, 5 Jul 2024 10:08:03 -0400
Subject: [PATCH] fix(hid): Eliminate data race in USB pathway causing dropped
 keys

Closes #2253
---
 app/Kconfig           |  9 +++++++++
 app/include/zmk/hid.h | 24 ++++++++++++++++++++++++
 app/src/usb_hid.c     | 16 +++++++++++++++-
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/app/Kconfig b/app/Kconfig
index a45f2dc23f0..931834eaa5e 100644
--- a/app/Kconfig
+++ b/app/Kconfig
@@ -648,6 +648,15 @@ config ZMK_SETTINGS_SAVE_DEBOUNCE
 #SETTINGS
 endif
 
+config ZMK_WORKAROUND_COPY_USB_TX_DATA_BUFFER
+    bool "Enable workaround for USB driver not copying data before transmit"
+    help
+      Work around some hardware USB drivers not copying the passed USB TX buffer
+      contents before beginning a transaction, causing corruption of the USB
+      message, by copying the TX data into a temporary buffer first. For
+      example, the NRFx family of Zephyr drivers suffer this issue.
+    default y if SOC_SERIES_NRF52X
+
 config ZMK_BATTERY_REPORT_INTERVAL
     depends on ZMK_BATTERY_REPORTING
     int "Battery level report interval in seconds"
diff --git a/app/include/zmk/hid.h b/app/include/zmk/hid.h
index 41f559b5189..2a027d0a6d2 100644
--- a/app/include/zmk/hid.h
+++ b/app/include/zmk/hid.h
@@ -254,6 +254,30 @@ struct zmk_hid_mouse_report {
 
 #endif // IS_ENABLED(CONFIG_ZMK_MOUSE)
 
+struct zmk_hid_report_body {
+    union {
+#if IS_ENABLED(CONFIG_ZMK_MOUSE)
+        struct zmk_hid_mouse_report_body mouse;
+#endif
+
+        struct zmk_hid_keyboard_report_body keyboard;
+        struct zmk_hid_consumer_report_body consumer;
+    } __packed;
+} __packed;
+
+struct zmk_hid_report {
+    union {
+#if IS_ENABLED(CONFIG_ZMK_USB_BOOT)
+        zmk_hid_boot_report_t boot;
+#endif
+
+        struct {
+            uint8_t report_id;
+            struct zmk_hid_report_body body;
+        } __packed;
+    } __packed;
+} __packed;
+
 zmk_mod_flags_t zmk_hid_get_explicit_mods(void);
 int zmk_hid_register_mod(zmk_mod_t modifier);
 int zmk_hid_unregister_mod(zmk_mod_t modifier);
diff --git a/app/src/usb_hid.c b/app/src/usb_hid.c
index 9db10952c95..c9e7f2b2a96 100644
--- a/app/src/usb_hid.c
+++ b/app/src/usb_hid.c
@@ -4,6 +4,8 @@
  * SPDX-License-Identifier: MIT
  */
 
+#include <zephyr/sys/util.h>
+
 #include <zephyr/device.h>
 #include <zephyr/init.h>
 
@@ -24,6 +26,17 @@ static const struct device *hid_dev;
 
 static K_SEM_DEFINE(hid_sem, 1, 1);
 
+#if IS_ENABLED(CONFIG_ZMK_WORKAROUND_COPY_USB_TX_DATA_BUFFER)
+static uint8_t hid_ep_write_buf[sizeof(struct zmk_hid_report)];
+
+static const uint8_t *prepare_report_buf(const uint8_t *report, size_t len) {
+    memcpy(hid_ep_write_buf, report, len);
+    return hid_ep_write_buf;
+}
+#else
+static const uint8_t *prepare_report_buf(const uint8_t *report, size_t len) { return report; }
+#endif
+
 static void in_ready_cb(const struct device *dev) { k_sem_give(&hid_sem); }
 
 #define HID_GET_REPORT_TYPE_MASK 0xff00
@@ -137,7 +150,8 @@ static int zmk_usb_hid_send_report(const uint8_t *report, size_t len) {
         return -ENODEV;
     default:
         k_sem_take(&hid_sem, K_MSEC(30));
-        int err = hid_int_ep_write(hid_dev, report, len, NULL);
+        const uint8_t *buf = prepare_report_buf(report, len);
+        int err = hid_int_ep_write(hid_dev, buf, len, NULL);
 
         if (err) {
             k_sem_give(&hid_sem);