5 files changed, 27 insertions, 35 deletions

diff --git a/quantum/keyboard.c b/quantum/keyboard.c
index f2a0889c1..f8e36994d 100644
--- a/quantum/keyboard.c
+++ b/quantum/keyboard.c
@@ -335,6 +335,12 @@ void keyboard_init(void) {
 #ifdef DIP_SWITCH_ENABLE
     dip_switch_init();
 #endif
+#ifdef SLEEP_LED_ENABLE
+    sleep_led_init();
+#endif
+#ifdef VIRTSER_ENABLE
+    virtser_init();
+#endif
 
 #if defined(DEBUG_MATRIX_SCAN_RATE) && defined(CONSOLE_ENABLE)
     debug_enable = true;
diff --git a/quantum/main.c b/quantum/main.c
index 2cbcd73d8..a896a67c6 100644
--- a/quantum/main.c
+++ b/quantum/main.c
@@ -20,7 +20,18 @@ void platform_setup(void);
 
 void protocol_setup(void);
 void protocol_init(void);
-void protocol_task(void);
+void protocol_pre_task(void);
+void protocol_post_task(void);
+
+// Bodge as refactoring vusb sucks....
+void protocol_task(void) __attribute__((weak));
+void protocol_task(void) {
+    protocol_pre_task();
+
+    keyboard_task();
+
+    protocol_post_task();
+}
 
 /** \brief Main
  *
@@ -30,8 +41,10 @@ int main(void) __attribute__((weak));
 int main(void) {
     platform_setup();
     protocol_setup();
+    keyboard_setup();
 
     protocol_init();
+    keyboard_init();
 
     /* Main loop */
     while (true) {
diff --git a/tmk_core/protocol/chibios/chibios.c b/tmk_core/protocol/chibios/chibios.c
index c860328c8..98b330508 100644
--- a/tmk_core/protocol/chibios/chibios.c
+++ b/tmk_core/protocol/chibios/chibios.c
@@ -138,8 +138,6 @@ void protocol_setup(void) {
 
     // TESTING
     // chThdCreateStatic(waThread1, sizeof(waThread1), NORMALPRIO, Thread1, NULL);
-
-    keyboard_setup();
 }
 
 void protocol_init(void) {
@@ -176,18 +174,10 @@ void protocol_init(void) {
 
     print("USB configured.\n");
 
-    /* init TMK modules */
-    keyboard_init();
     host_set_driver(driver);
-
-#ifdef SLEEP_LED_ENABLE
-    sleep_led_init();
-#endif
-
-    print("Keyboard start.\n");
 }
 
-void protocol_task(void) {
+void protocol_pre_task(void) {
     usb_event_queue_task();
 
 #if !defined(NO_USB_STARTUP_CHECK)
@@ -210,8 +200,9 @@ void protocol_task(void) {
 #    endif /* MOUSEKEY_ENABLE */
     }
 #endif
+}
 
-    keyboard_task();
+void protocol_post_task(void) {
 #ifdef CONSOLE_ENABLE
     console_task();
 #endif
diff --git a/tmk_core/protocol/lufa/lufa.c b/tmk_core/protocol/lufa/lufa.c
index 753762358..80781d2f3 100644
--- a/tmk_core/protocol/lufa/lufa.c
+++ b/tmk_core/protocol/lufa/lufa.c
@@ -1070,7 +1070,6 @@ void protocol_setup(void) {
 
     setup_mcu();
     usb_device_state_init();
-    keyboard_setup();
 }
 
 void protocol_init(void) {
@@ -1095,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]");
@@ -1133,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
diff --git a/tmk_core/protocol/vusb/protocol.c b/tmk_core/protocol/vusb/protocol.c
index 89dc795b2..947c3383f 100644
--- a/tmk_core/protocol/vusb/protocol.c
+++ b/tmk_core/protocol/vusb/protocol.c
@@ -111,22 +111,15 @@ void protocol_setup(void) {
     // clock prescaler
     clock_prescale_set(clock_div_1);
 #endif
-    keyboard_setup();
 }
 
 void protocol_init(void) {
     setup_usb();
     sei();
 
-    keyboard_init();
-
     host_set_driver(vusb_driver());
 
     wait_ms(50);
-
-#ifdef SLEEP_LED_ENABLE
-    sleep_led_init();
-#endif
 }
 
 void protocol_task(void) {