From cc59ab21c5489484e0adc969b76e7b33267f85b0 Mon Sep 17 00:00:00 2001
From: Ilya Zhuravlev <whatever@xyz.is>
Date: Tue, 29 Dec 2020 15:36:12 -0500
Subject: [PATCH] via/vial: kill all rawhid handling while unlocking

---
 quantum/via.c  |  9 ++++++++-
 quantum/via.h  |  1 +
 quantum/vial.c | 12 ------------
 quantum/vial.h | 12 ++++++++++++
 4 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/quantum/via.c b/quantum/via.c
index 1afac67264..3191e965d3 100644
--- a/quantum/via.c
+++ b/quantum/via.c
@@ -218,6 +218,13 @@ __attribute__((weak)) void raw_hid_receive_kb(uint8_t *data, uint8_t length) {
 void raw_hid_receive(uint8_t *data, uint8_t length) {
     uint8_t *command_id   = &(data[0]);
     uint8_t *command_data = &(data[1]);
+
+#ifdef VIAL_ENABLE
+    /* When unlock is in progress, only command we react to is unlock_poll */
+    if (vial_unlock_in_progress && (data[0] != id_vial_prefix || data[1] != vial_unlock_poll))
+        goto skip;
+#endif
+
     switch (*command_id) {
         case id_get_protocol_version: {
             command_data[0] = VIA_PROTOCOL_VERSION >> 8;
@@ -419,7 +426,7 @@ void raw_hid_receive(uint8_t *data, uint8_t length) {
             break;
         }
 #ifdef VIAL_ENABLE
-        case 0xFE: {
+        case id_vial_prefix: {
             vial_handle_cmd(data, length);
             break;
         }
diff --git a/quantum/via.h b/quantum/via.h
index 373843f901..c17fcd515f 100644
--- a/quantum/via.h
+++ b/quantum/via.h
@@ -78,6 +78,7 @@ enum via_command_id {
     id_dynamic_keymap_get_layer_count       = 0x11,
     id_dynamic_keymap_get_buffer            = 0x12,
     id_dynamic_keymap_set_buffer            = 0x13,
+    id_vial_prefix                          = 0xFE,
     id_unhandled                            = 0xFF,
 };
 
diff --git a/quantum/vial.c b/quantum/vial.c
index dccba4240b..3512f1ecf1 100644
--- a/quantum/vial.c
+++ b/quantum/vial.c
@@ -25,18 +25,6 @@
 
 #define VIAL_UNLOCK_COUNTER_MAX 50
 
-enum {
-    vial_get_keyboard_id = 0x00,
-    vial_get_size = 0x01,
-    vial_get_def = 0x02,
-    vial_get_encoder = 0x03,
-    vial_set_encoder = 0x04,
-    vial_get_lock = 0x05,
-    vial_unlock_start = 0x06,
-    vial_unlock_poll = 0x07,
-    vial_lock = 0x08,
-};
-
 #ifdef VIAL_INSECURE
 #pragma message "Building Vial-enabled firmware in insecure mode."
 int vial_unlocked = 1;
diff --git a/quantum/vial.h b/quantum/vial.h
index 875bcb6355..514ac815d0 100644
--- a/quantum/vial.h
+++ b/quantum/vial.h
@@ -29,3 +29,15 @@ void vial_encoder_update(uint8_t index, bool clockwise);
 
 extern int vial_unlocked;
 extern int vial_unlock_in_progress;
+
+enum {
+    vial_get_keyboard_id = 0x00,
+    vial_get_size = 0x01,
+    vial_get_def = 0x02,
+    vial_get_encoder = 0x03,
+    vial_set_encoder = 0x04,
+    vial_get_lock = 0x05,
+    vial_unlock_start = 0x06,
+    vial_unlock_poll = 0x07,
+    vial_lock = 0x08,
+};