1 files changed, 12 insertions, 15 deletions

diff --git a/tmk_core/protocol/chibios/usb_main.c b/tmk_core/protocol/chibios/usb_main.c
index 8fbe877db..d71e7ce44 100644
--- a/tmk_core/protocol/chibios/usb_main.c
+++ b/tmk_core/protocol/chibios/usb_main.c
@@ -62,7 +62,7 @@ extern keymap_config_t keymap_config;
 
 uint8_t                keyboard_idle __attribute__((aligned(2)))      = 0;
 uint8_t                keyboard_protocol __attribute__((aligned(2)))  = 1;
-uint16_t               keyboard_led_stats __attribute__((aligned(2))) = 0;
+uint8_t                keyboard_led_stats                             = 0;
 volatile uint16_t      keyboard_idle_count                            = 0;
 static virtual_timer_t keyboard_idle_timer;
 static void            keyboard_idle_timer_cb(void *arg);
@@ -381,14 +381,17 @@ static uint16_t get_hword(uint8_t *p) {
  * Other Device    Required    Optional    Optional    Optional    Optional    Optional
  */
 
-#if defined(SHARED_EP_ENABLE) && !defined(KEYBOARD_SHARED_EP)
 static uint8_t set_report_buf[2] __attribute__((aligned(2)));
 static void    set_led_transfer_cb(USBDriver *usbp) {
-    if ((set_report_buf[0] == REPORT_ID_KEYBOARD) || (set_report_buf[0] == REPORT_ID_NKRO)) {
-        keyboard_led_stats = set_report_buf[1];
+    if (usbp->setup[6] == 2) { /* LSB(wLength) */
+        uint8_t report_id = set_report_buf[0];
+        if ((report_id == REPORT_ID_KEYBOARD) || (report_id == REPORT_ID_NKRO)) {
+            keyboard_led_stats = set_report_buf[1];
+        }
+    } else {
+        keyboard_led_stats = set_report_buf[0];
     }
 }
-#endif
 
 /* Callback for SETUP request on the endpoint 0 (control) */
 static bool usb_request_hook_cb(USBDriver *usbp) {
@@ -444,18 +447,12 @@ static bool usb_request_hook_cb(USBDriver *usbp) {
             case USB_RTYPE_DIR_HOST2DEV:
                 switch (usbp->setup[1]) { /* bRequest */
                     case HID_SET_REPORT:
-                        switch (usbp->setup[4]) { /* LSB(wIndex) (check MSB==0 and wLength==1?) */
+                        switch (usbp->setup[4]) { /* LSB(wIndex) (check MSB==0?) */
+                            case KEYBOARD_INTERFACE:
 #if defined(SHARED_EP_ENABLE) && !defined(KEYBOARD_SHARED_EP)
                             case SHARED_INTERFACE:
-                                usbSetupTransfer(usbp, set_report_buf, sizeof(set_report_buf), set_led_transfer_cb);
-                                return TRUE;
-                                break;
 #endif
-
-                            case KEYBOARD_INTERFACE:
-                                /* keyboard_led_stats = <read byte from next OUT report>
-                                 * keyboard_led_stats needs be word (or dword), otherwise we get an exception on F0 */
-                                usbSetupTransfer(usbp, (uint8_t *)&keyboard_led_stats, 1, NULL);
+                                usbSetupTransfer(usbp, set_report_buf, sizeof(set_report_buf), set_led_transfer_cb);
                                 return TRUE;
                                 break;
                         }
@@ -613,7 +610,7 @@ static void keyboard_idle_timer_cb(void *arg) {
 }
 
 /* LED status */
-uint8_t keyboard_leds(void) { return (uint8_t)(keyboard_led_stats & 0xFF); }
+uint8_t keyboard_leds(void) { return keyboard_led_stats; }
 
 /* prepare and start sending a report IN
  * not callable from ISR or locked state */