1 files changed, 41 insertions, 26 deletions

diff --git a/tmk_core/protocol/lufa/lufa.c b/tmk_core/protocol/lufa/lufa.c
index 5b56e8a03c..80781d2f37 100644
--- a/tmk_core/protocol/lufa/lufa.c
+++ b/tmk_core/protocol/lufa/lufa.c
@@ -52,6 +52,7 @@
 #include "usb_descriptor.h"
 #include "lufa.h"
 #include "quantum.h"
+#include "usb_device_state.h"
 #include <util/atomic.h>
 
 #ifdef NKRO_ENABLE
@@ -142,7 +143,8 @@ static void    send_keyboard(report_keyboard_t *report);
 static void    send_mouse(report_mouse_t *report);
 static void    send_system(uint16_t data);
 static void    send_consumer(uint16_t data);
-host_driver_t  lufa_driver = {keyboard_leds, send_keyboard, send_mouse, send_system, send_consumer};
+static void    send_programmable_button(uint32_t data);
+host_driver_t  lufa_driver = {keyboard_leds, send_keyboard, send_mouse, send_system, send_consumer, send_programmable_button};
 
 #ifdef VIRTSER_ENABLE
 // clang-format off
@@ -413,7 +415,10 @@ void EVENT_USB_Device_Disconnect(void) {
  *
  * FIXME: Needs doc
  */
-void EVENT_USB_Device_Reset(void) { print("[R]"); }
+void EVENT_USB_Device_Reset(void) {
+    print("[R]");
+    usb_device_state_set_reset();
+}
 
 /** \brief Event USB Device Connect
  *
@@ -421,6 +426,8 @@ void EVENT_USB_Device_Reset(void) { print("[R]"); }
  */
 void EVENT_USB_Device_Suspend() {
     print("[S]");
+    usb_device_state_set_suspend(USB_Device_ConfigurationNumber != 0, USB_Device_ConfigurationNumber);
+
 #ifdef SLEEP_LED_ENABLE
     sleep_led_enable();
 #endif
@@ -436,6 +443,8 @@ void EVENT_USB_Device_WakeUp() {
     suspend_wakeup_init();
 #endif
 
+    usb_device_state_set_resume(USB_DeviceState == DEVICE_STATE_Configured, USB_Device_ConfigurationNumber);
+
 #ifdef SLEEP_LED_ENABLE
     sleep_led_disable();
     // NOTE: converters may not accept this
@@ -528,6 +537,8 @@ void EVENT_USB_Device_ConfigurationChanged(void) {
     /* Setup digitizer endpoint */
     ConfigSuccess &= Endpoint_ConfigureEndpoint((DIGITIZER_IN_EPNUM | ENDPOINT_DIR_IN), EP_TYPE_INTERRUPT, DIGITIZER_EPSIZE, 1);
 #endif
+
+    usb_device_state_set_configuration(USB_DeviceState == DEVICE_STATE_Configured, USB_Device_ConfigurationNumber);
 }
 
 /* FIXME: Expose this table in the docs somehow
@@ -760,29 +771,35 @@ static void send_mouse(report_mouse_t *report) {
 #endif
 }
 
-/** \brief Send Extra
- *
- * FIXME: Needs doc
- */
-#ifdef EXTRAKEY_ENABLE
-static void send_extra(uint8_t report_id, uint16_t data) {
+#if defined(EXTRAKEY_ENABLE) || defined(PROGRAMMABLE_BUTTON_ENABLE)
+static void send_report(void *report, size_t size) {
     uint8_t timeout = 255;
 
     if (USB_DeviceState != DEVICE_STATE_Configured) return;
 
-    static report_extra_t r;
-    r = (report_extra_t){.report_id = report_id, .usage = data};
     Endpoint_SelectEndpoint(SHARED_IN_EPNUM);
 
     /* Check if write ready for a polling interval around 10ms */
     while (timeout-- && !Endpoint_IsReadWriteAllowed()) _delay_us(40);
     if (!Endpoint_IsReadWriteAllowed()) return;
 
-    Endpoint_Write_Stream_LE(&r, sizeof(report_extra_t), NULL);
+    Endpoint_Write_Stream_LE(report, size, NULL);
     Endpoint_ClearIN();
 }
 #endif
 
+/** \brief Send Extra
+ *
+ * FIXME: Needs doc
+ */
+#ifdef EXTRAKEY_ENABLE
+static void send_extra(uint8_t report_id, uint16_t data) {
+    static report_extra_t r;
+    r = (report_extra_t){.report_id = report_id, .usage = data};
+    send_report(&r, sizeof(r));
+}
+#endif
+
 /** \brief Send System
  *
  * FIXME: Needs doc
@@ -822,6 +839,14 @@ static void send_consumer(uint16_t data) {
 #endif
 }
 
+static void send_programmable_button(uint32_t data) {
+#ifdef PROGRAMMABLE_BUTTON_ENABLE
+    static report_programmable_button_t r;
+    r = (report_programmable_button_t){.report_id = REPORT_ID_PROGRAMMABLE_BUTTON, .usage = data};
+    send_report(&r, sizeof(r));
+#endif
+}
+
 /*******************************************************************************
  * sendchar
  ******************************************************************************/
@@ -1044,7 +1069,7 @@ void protocol_setup(void) {
 #endif
 
     setup_mcu();
-    keyboard_setup();
+    usb_device_state_init();
 }
 
 void protocol_init(void) {
@@ -1069,21 +1094,11 @@ void protocol_init(void) {
 #else
     USB_USBTask();
 #endif
-    /* init modules */
-    keyboard_init();
-    host_set_driver(&lufa_driver);
-#ifdef SLEEP_LED_ENABLE
-    sleep_led_init();
-#endif
-
-#ifdef VIRTSER_ENABLE
-    virtser_init();
-#endif
 
-    print("Keyboard start.\n");
+    host_set_driver(&lufa_driver);
 }
 
-void protocol_task(void) {
+void protocol_pre_task(void) {
 #if !defined(NO_USB_STARTUP_CHECK)
     if (USB_DeviceState == DEVICE_STATE_Suspended) {
         print("[s]");
@@ -1107,9 +1122,9 @@ void protocol_task(void) {
         suspend_wakeup_init();
     }
 #endif
+}
 
-    keyboard_task();
-
+void protocol_post_task(void) {
 #ifdef MIDI_ENABLE
     MIDI_Device_USBTask(&USB_MIDI_Interface);
 #endif