check usb endpoint

4 files changed, 36 insertions, 0 deletions

diff --git a/hid/patch.py b/hid/patch.py
index 58462bc1..4892d4b2 100644
--- a/hid/patch.py
+++ b/hid/patch.py
@@ -34,6 +34,7 @@ def _patch(path: str, patch_path: str) -> None:
 
 # =====
 _patch(_get_pkg_path("framework-arduino-avr"), "patches/serial.patch")
+_patch(_get_pkg_path("framework-arduino-avr"), "patches/get-plugged-endpoint.patch")
 
 _libs = _get_libs()
 if "HID-Project" in _libs:
diff --git a/hid/patches/get-plugged-endpoint.patch b/hid/patches/get-plugged-endpoint.patch
new file mode 100644
index 00000000..c32131ff
--- /dev/null
+++ b/hid/patches/get-plugged-endpoint.patch
@@ -0,0 +1,11 @@
+--- a/cores/arduino/PluggableUSB.h	2019-05-16 15:52:01.000000000 +0300
++++ b/cores/arduino/PluggableUSB.h	2020-11-14 20:57:30.942432544 +0300
+@@ -31,6 +31,8 @@
+     numEndpoints(numEps), numInterfaces(numIfs), endpointType(epType)
+   { }
+
++  uint8_t getPluggedEndpoint() { return pluggedEndpoint; }
++
+ protected:
+   virtual bool setup(USBSetup& setup) = 0;
+   virtual int getInterface(uint8_t* interfaceCount) = 0;
diff --git a/hid/platformio.ini b/hid/platformio.ini
index a24c4aca..1f6bd069 100644
--- a/hid/platformio.ini
+++ b/hid/platformio.ini
@@ -103,6 +103,7 @@ extra_scripts =
 build_flags =
 	-DCMD_SPI
 	-DNO_SERIAL
+	-DCHECK_ENDPOINT
 upload_protocol = custom
 upload_flags =
 	-C
diff --git a/hid/src/usb/hid.h b/hid/src/usb/hid.h
index 326e8cc5..f0016931 100644
--- a/hid/src/usb/hid.h
+++ b/hid/src/usb/hid.h
@@ -22,12 +22,29 @@
 
 #pragma once
 
+#include <Arduino.h>
 #include <HID-Project.h>
 
 #include "keymap.h"
 
 
 // -----------------------------------------------------------------------------
+#ifdef CHECK_ENDPOINT
+static bool _checkEndpoint(uint8_t ep) {
+	// https://github.com/arduino/ArduinoCore-avr/blob/2f67c916f6ab6193c404eebe22efe901e0f9542d/cores/arduino/USBCore.cpp#L249
+	// https://sourceforge.net/p/arduinomidilib/svn/41/tree/branch/3.1/Teensy/teensy_core/usb_midi/usb_api.cpp#l103
+	uint8_t intr_state = SREG;
+	cli();
+	UENUM = ep & 7;
+	bool rw_allowed = UEINTX & (1 << RWAL);
+	SREG = intr_state;
+	return rw_allowed;
+}
+#	define CHECK_HID_EP(_hid) { if (!_checkEndpoint(_hid.getPluggedEndpoint())) return; }
+#else
+#	define CHECK_HID_EP(_hid)
+#endif
+
 class UsbHidKeyboard {
 	public:
 		UsbHidKeyboard() {}
@@ -41,6 +58,7 @@ class UsbHidKeyboard {
 		}
 
 		void sendKey(uint8_t code, bool state) {
+			CHECK_HID_EP(BootKeyboard);
 			KeyboardKeycode usb_code = keymapUsb(code);
 			if (usb_code != KEY_ERROR_UNDEFINED) {
 				if (state) BootKeyboard.press(usb_code);
@@ -86,17 +104,22 @@ class UsbHidMouse {
 		}
 
 		void sendMove(int x, int y) {
+			CHECK_HID_EP(SingleAbsoluteMouse);
 			SingleAbsoluteMouse.moveTo(x, y);
 		}
 
 		void sendWheel(int delta_y) {
+			CHECK_HID_EP(SingleAbsoluteMouse);
 			// delta_x is not supported by hid-project now
 			SingleAbsoluteMouse.move(0, 0, delta_y);
 		}
 
 	private:
 		void _sendButton(uint8_t button, bool state) {
+			CHECK_HID_EP(SingleAbsoluteMouse);
 			if (state) SingleAbsoluteMouse.press(button);
 			else SingleAbsoluteMouse.release(button);
 		}
 };
+
+#undef CHECK_HID_EP