From f72e57a6b3b5d6704733e29a8572940e00272b33 Mon Sep 17 00:00:00 2001
From: Edward Hesketh <headblockhead@gmail.com>
Date: Wed, 4 Sep 2024 23:32:15 +0100
Subject: [PATCH] fix: multikey test

---
 include/squirrel.h      |  4 +++-
 include/squirrel_init.h |  2 --
 src/squirrel_init.c     | 21 ++++++++-------------
 src/squirrel_quantum.c  |  1 +
 tests/multikey.c        | 35 +++++++++++++++++++++--------------
 5 files changed, 33 insertions(+), 30 deletions(-)

diff --git a/include/squirrel.h b/include/squirrel.h
index bcb67ec..8944a2a 100644
--- a/include/squirrel.h
+++ b/include/squirrel.h
@@ -6,7 +6,9 @@
 enum squirrel_error {
   ERR_NONE = 0,
   ERR_KEY_FUNC_WRONG_ARGUMENT_COUNT,
-  ERR_PASSTHROUGH_ON_BOTTOM_LAYER
+  ERR_PASSTHROUGH_ON_BOTTOM_LAYER,
+  ERR_OUT_OF_MEMORY_KEYS,
+  ERR_OUT_OF_MEMORY_KEY_STATES,
 };
 
 #endif
diff --git a/include/squirrel_init.h b/include/squirrel_init.h
index 58d4bcf..f33daac 100644
--- a/include/squirrel_init.h
+++ b/include/squirrel_init.h
@@ -2,6 +2,4 @@
 #define SQUIRREL_INIT_H
 enum squirrel_error squirrel_init(
     int total_keys); // Initialize the keyboard with the total number of keys.
-enum squirrel_error
-squirrel_deinit(); // Deinitialize the keyboard. Frees memory.
 #endif
diff --git a/src/squirrel_init.c b/src/squirrel_init.c
index 8bf2545..665ed51 100644
--- a/src/squirrel_init.c
+++ b/src/squirrel_init.c
@@ -11,28 +11,23 @@ enum squirrel_error squirrel_init(int total_keys) {
   struct key passthrough_key;
   passthrough_key.pressed = quantum_passthrough_press;
   passthrough_key.pressed_argument_count = 0;
+  passthrough_key.pressed_arguments = NULL;
   passthrough_key.released = quantum_passthrough_release;
   passthrough_key.released_argument_count = 0;
+  passthrough_key.released_arguments = NULL;
   for (uint8_t j = 16; j != 255; j--) {
     layers[j].keys = (struct key *)malloc(total_keys * sizeof(struct key));
+    if (layers[j].keys == NULL) {
+      return ERR_OUT_OF_MEMORY_KEYS;
+    }
     for (int i = 0; i < total_keys; i++) {
       copy_key(&passthrough_key, &(layers[j].keys[i]));
     }
   }
   key_states = (bool *)malloc(total_keys * sizeof(bool));
-  layers[16].active = true;
-  return ERR_NONE;
-};
-
-enum squirrel_error squirrel_deinit() {
-  for (uint8_t j = 16; j != 255; j--) {
-    free(layers[j].keys);
-  }
-  free(key_states);
-  // NULL the pointers to prevent use after free
-  for (uint8_t j = 16; j != 255; j--) {
-    layers[j].keys = NULL;
+  if (key_states == NULL) {
+    return ERR_OUT_OF_MEMORY_KEY_STATES;
   }
-  key_states = NULL;
+  layers[16].active = true;
   return ERR_NONE;
 };
diff --git a/src/squirrel_quantum.c b/src/squirrel_quantum.c
index b84f2ca..d7f495d 100644
--- a/src/squirrel_quantum.c
+++ b/src/squirrel_quantum.c
@@ -4,6 +4,7 @@
 #include "squirrel_key.h"
 #include "squirrel_keyboard.h"
 #include <stdint.h>
+#include <stdlib.h>
 
 struct layer layers[17];
 
diff --git a/tests/multikey.c b/tests/multikey.c
index 356e53e..373f35d 100644
--- a/tests/multikey.c
+++ b/tests/multikey.c
@@ -40,15 +40,11 @@ enum squirrel_error test_release(struct key *key, uint8_t layer,
   return ERR_NONE;
 }
 
-// test: press_key, release_key + check_key - in squirrel_key.c,
-// but with multiple keys in a keyboard.
-int main() {
-  squirrel_init(1);
+uint8_t code1 = 0x0F;
+uint8_t code2 = 0xF0;
 
+void make_testkey(int index) {
   // press_key + release_key
-  uint8_t code1 = 0x0F;
-  uint8_t code2 = 0xF0;
-
   struct key testkey;
   testkey.pressed = test_press;
   testkey.pressed_argument_count = 2;
@@ -62,11 +58,18 @@ int main() {
   testkey.released_arguments[0] = &code2;
   testkey.released_arguments[1] = &code1;
 
-  layers[0].keys[0] = testkey;
-  layers[0].keys[1] = testkey;
-  layers[0].keys[2] = testkey;
-  layers[0].keys[3] = testkey;
-  layers[0].keys[4] = testkey;
+  layers[0].keys[index] = testkey;
+};
+
+// test: press_key, release_key + check_key - in squirrel_key.c,
+// but with multiple keys in a keyboard.
+int main() {
+  squirrel_init(5);
+  make_testkey(0);
+  make_testkey(1);
+  make_testkey(2);
+  make_testkey(3);
+  make_testkey(4);
   layers[0].active = true;
 
   for (int i = 0; i < 5; i++) {
@@ -75,6 +78,7 @@ int main() {
     if (err != ERR_NONE) {
       return 2;
     }
+    return 0;
     if (test_result != 0) {
       return 3;
     }
@@ -157,8 +161,11 @@ int main() {
     }
   };
 
-  free(testkey.pressed_arguments);
-  free(testkey.released_arguments);
+  // cleanup
+  for (int i = 0; i < 5; i++) {
+    free(layers[0].keys[i].pressed_arguments);
+    free(layers[0].keys[i].released_arguments);
+  }
 
   return 0;
 }