6 files changed, 120 insertions, 2 deletions

diff --git a/common_features.mk b/common_features.mk
index 619359717..a71bcace0 100644
--- a/common_features.mk
+++ b/common_features.mk
@@ -29,6 +29,7 @@ QUANTUM_SRC += \
     $(QUANTUM_DIR)/keyboard.c \
     $(QUANTUM_DIR)/keymap_common.c \
     $(QUANTUM_DIR)/keycode_config.c \
+    $(QUANTUM_DIR)/usb_device_state.c \
     $(QUANTUM_DIR)/logging/debug.c \
     $(QUANTUM_DIR)/logging/sendchar.c \
 
diff --git a/quantum/usb_device_state.c b/quantum/usb_device_state.c
new file mode 100644
index 000000000..5ccd309ec
--- /dev/null
+++ b/quantum/usb_device_state.c
@@ -0,0 +1,51 @@
+/*
+ * Copyright 2021 Andrei Purdea <andrei@purdea.ro>
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#include "usb_device_state.h"
+
+enum usb_device_state usb_device_state = USB_DEVICE_STATE_NO_INIT;
+
+__attribute__((weak)) void notify_usb_device_state_change_kb(enum usb_device_state usb_device_state) { notify_usb_device_state_change_user(usb_device_state); }
+
+__attribute__((weak)) void notify_usb_device_state_change_user(enum usb_device_state usb_device_state) {}
+
+static void notify_usb_device_state_change(enum usb_device_state usb_device_state) { notify_usb_device_state_change_kb(usb_device_state); }
+
+void usb_device_state_set_configuration(bool isConfigured, uint8_t configurationNumber) {
+    usb_device_state = isConfigured ? USB_DEVICE_STATE_CONFIGURED : USB_DEVICE_STATE_INIT;
+    notify_usb_device_state_change(usb_device_state);
+}
+
+void usb_device_state_set_suspend(bool isConfigured, uint8_t configurationNumber) {
+    usb_device_state = USB_DEVICE_STATE_SUSPEND;
+    notify_usb_device_state_change(usb_device_state);
+}
+
+void usb_device_state_set_resume(bool isConfigured, uint8_t configurationNumber) {
+    usb_device_state = isConfigured ? USB_DEVICE_STATE_CONFIGURED : USB_DEVICE_STATE_INIT;
+    notify_usb_device_state_change(usb_device_state);
+}
+
+void usb_device_state_set_reset(void) {
+    usb_device_state = USB_DEVICE_STATE_INIT;
+    notify_usb_device_state_change(usb_device_state);
+}
+
+void usb_device_state_init(void) {
+    usb_device_state = USB_DEVICE_STATE_INIT;
+    notify_usb_device_state_change(usb_device_state);
+}
diff --git a/quantum/usb_device_state.h b/quantum/usb_device_state.h
new file mode 100644
index 000000000..c229311d4
--- /dev/null
+++ b/quantum/usb_device_state.h
@@ -0,0 +1,39 @@
+/*
+ * Copyright 2021 Andrei Purdea <andrei@purdea.ro>
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#pragma once
+
+#include <stdbool.h>
+#include <stdint.h>
+
+void usb_device_state_set_configuration(bool isConfigured, uint8_t configurationNumber);
+void usb_device_state_set_suspend(bool isConfigured, uint8_t configurationNumber);
+void usb_device_state_set_resume(bool isConfigured, uint8_t configurationNumber);
+void usb_device_state_set_reset(void);
+void usb_device_state_init(void);
+
+enum usb_device_state {
+    USB_DEVICE_STATE_NO_INIT    = 0,  // We're in this state before calling usb_device_state_init()
+    USB_DEVICE_STATE_INIT       = 1,  // Can consume up to 100mA
+    USB_DEVICE_STATE_CONFIGURED = 2,  // Can consume up to what is specified in configuration descriptor, typically 500mA
+    USB_DEVICE_STATE_SUSPEND    = 3   // Can consume only suspend current
+};
+
+extern enum usb_device_state usb_device_state;
+
+void notify_usb_device_state_change_kb(enum usb_device_state usb_device_state);
+void notify_usb_device_state_change_user(enum usb_device_state usb_device_state);
diff --git a/tmk_core/protocol/chibios/chibios.c b/tmk_core/protocol/chibios/chibios.c
index 78a2e3fcb..e5d3a4c54 100644
--- a/tmk_core/protocol/chibios/chibios.c
+++ b/tmk_core/protocol/chibios/chibios.c
@@ -27,6 +27,7 @@
 #include "keyboard.h"
 #include "action.h"
 #include "action_util.h"
+#include "usb_device_state.h"
 #include "mousekey.h"
 #include "led.h"
 #include "sendchar.h"
@@ -139,6 +140,8 @@ void boardInit(void) {
 }
 
 void protocol_setup(void) {
+    usb_device_state_init();
+
     // TESTING
     // chThdCreateStatic(waThread1, sizeof(waThread1), NORMALPRIO, Thread1, NULL);
 
diff --git a/tmk_core/protocol/chibios/usb_main.c b/tmk_core/protocol/chibios/usb_main.c
index cc282e6a9..9b139b399 100644
--- a/tmk_core/protocol/chibios/usb_main.c
+++ b/tmk_core/protocol/chibios/usb_main.c
@@ -39,6 +39,7 @@
 #    include "led.h"
 #endif
 #include "wait.h"
+#include "usb_device_state.h"
 #include "usb_descriptor.h"
 #include "usb_driver.h"
 
@@ -412,6 +413,7 @@ static inline bool usb_event_queue_dequeue(usbevent_t *event) {
 }
 
 static inline void usb_event_suspend_handler(void) {
+    usb_device_state_set_suspend(USB_DRIVER.configuration != 0, USB_DRIVER.configuration);
 #ifdef SLEEP_LED_ENABLE
     sleep_led_enable();
 #endif /* SLEEP_LED_ENABLE */
@@ -419,6 +421,7 @@ static inline void usb_event_suspend_handler(void) {
 
 static inline void usb_event_wakeup_handler(void) {
     suspend_wakeup_init();
+    usb_device_state_set_resume(USB_DRIVER.configuration != 0, USB_DRIVER.configuration);
 #ifdef SLEEP_LED_ENABLE
     sleep_led_disable();
     // NOTE: converters may not accept this
@@ -440,6 +443,15 @@ void usb_event_queue_task(void) {
                 last_suspend_state = false;
                 usb_event_wakeup_handler();
                 break;
+            case USB_EVENT_CONFIGURED:
+                usb_device_state_set_configuration(USB_DRIVER.configuration != 0, USB_DRIVER.configuration);
+                break;
+            case USB_EVENT_UNCONFIGURED:
+                usb_device_state_set_configuration(false, 0);
+                break;
+            case USB_EVENT_RESET:
+                usb_device_state_set_reset();
+                break;
             default:
                 // Nothing to do, we don't handle it.
                 break;
@@ -482,13 +494,14 @@ static void usb_event_cb(USBDriver *usbp, usbevent_t event) {
             if (last_suspend_state) {
                 usb_event_queue_enqueue(USB_EVENT_WAKEUP);
             }
+            usb_event_queue_enqueue(USB_EVENT_CONFIGURED);
             return;
         case USB_EVENT_SUSPEND:
-            usb_event_queue_enqueue(USB_EVENT_SUSPEND);
             /* Falls into.*/
         case USB_EVENT_UNCONFIGURED:
             /* Falls into.*/
         case USB_EVENT_RESET:
+            usb_event_queue_enqueue(event);
             for (int i = 0; i < NUM_USB_DRIVERS; i++) {
                 chSysLockFromISR();
                 /* Disconnection event on suspend.*/
diff --git a/tmk_core/protocol/lufa/lufa.c b/tmk_core/protocol/lufa/lufa.c
index 4a30d2257..753762358 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
@@ -414,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
  *
@@ -422,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
@@ -437,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
@@ -529,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
@@ -1059,6 +1069,7 @@ void protocol_setup(void) {
 #endif
 
     setup_mcu();
+    usb_device_state_init();
     keyboard_setup();
 }