From 924ddf3296739a0cc5e643424a27eb3e205ca689 Mon Sep 17 00:00:00 2001
From: Daniel Schaefer <dhs@frame.work>
Date: Tue, 17 Jan 2023 11:29:52 +0800
Subject: [PATCH] lotus: Clean up serial number code

More close to upstreaming ready.

Signed-off-by: Daniel Schaefer <dhs@frame.work>
---
 keyboards/lotus/config.h           |  5 ++---
 keyboards/lotus/factory.c          |  1 -
 keyboards/lotus/lotus.c            | 36 +++++++++++++++---------------
 keyboards/lotus/lotus.h            |  6 +++++
 tmk_core/protocol/usb_descriptor.c | 12 +---------
 5 files changed, 27 insertions(+), 33 deletions(-)

diff --git a/keyboards/lotus/config.h b/keyboards/lotus/config.h
index b818a878f2..3a3f0ea11a 100644
--- a/keyboards/lotus/config.h
+++ b/keyboards/lotus/config.h
@@ -13,9 +13,8 @@
 // SDB is mapped to GP22
 #define PICO_LOTUS
 
-// 18-digit serial number. Must be programmed at factory
-// TODO: Figure out how to do that. Probably a script to compile qmk
-#define SERIAL_NUMBER "FRALDLvvccxxxxxxxx"
+// Enable serial number
+#define SERIAL_NUMBER
 
 #ifdef KEYBOARD_lotus_ansi
     #define LED_CAPS_LOCK_PIN GP24
diff --git a/keyboards/lotus/factory.c b/keyboards/lotus/factory.c
index 4d21c75949..b65e2f6734 100644
--- a/keyboards/lotus/factory.c
+++ b/keyboards/lotus/factory.c
@@ -29,7 +29,6 @@ void emulate_rgb_keycode_press(uint16_t target_keycode) {
 #endif
 
 extern char ascii_serialnum[SERIALNUM_LEN+1];
-extern char utf16_serialnum[(SERIALNUM_LEN+1) * 2];
 
 void handle_factory_command(uint8_t *data) {
     uint8_t factory_command_id = data[0];
diff --git a/keyboards/lotus/lotus.c b/keyboards/lotus/lotus.c
index c91f48b7c4..c6f07c6d65 100644
--- a/keyboards/lotus/lotus.c
+++ b/keyboards/lotus/lotus.c
@@ -9,43 +9,43 @@
 #define USBSTR(s) USBCONCAT(L, s)
 USB_Descriptor_String_t PROGMEM SerialNumberString = {
     .Header = {
-        .Size                   = sizeof(USBSTR(SERIAL_NUMBER)),
+        .Size                   = sizeof(FALLBACK_SERIAL_NUMBER),
         .Type                   = DTYPE_String
     },
-    .UnicodeString              = USBSTR(SERIAL_NUMBER)
+    .UnicodeString              = FALLBACK_SERIAL_NUMBER
 };
 
 char ascii_serialnum[SERIALNUM_LEN+1];
-char utf16_serialnum[(SERIALNUM_LEN+1) * 2];
 
 void *lotus_serial_number_string(void) {
-  if (utf16_serialnum[0] != '\0') {
-    return utf16_serialnum;
+  // Exit early, if it was previously read and converted
+  if (ascii_serialnum[0] != '\0' || ascii_serialnum[0] == 0xFF) {
+      return &SerialNumberString;
   }
-  ascii_serialnum[SERIALNUM_LEN] = '\0';
-  utf16_serialnum[SERIALNUM_LEN*2] = '\0';
-  utf16_serialnum[SERIALNUM_LEN*2 + 1] = '\0';
-  char *serialnum_ptr = (char*) (FLASH_OFFSET + LAST_4K_BLOCK);
 
+  // Read ASCII serial number from memory-mapped flash
+  char *serialnum_ptr = (char*) (FLASH_OFFSET + LAST_4K_BLOCK);
   memcpy(ascii_serialnum, serialnum_ptr, SERIALNUM_LEN);
 
-  // Replace with dummy serialnumber if the flash is erased
+  // Just keep fallback serialnumber if the flash is erased
   if (ascii_serialnum[0] == 0xFF) {
-      memcpy(ascii_serialnum, SERIAL_NUMBER, sizeof(SERIAL_NUMBER));
+      return &SerialNumberString;
   }
 
-  for (int c = 0; c < SERIALNUM_LEN; c++) {
-      utf16_serialnum[c*2] = ascii_serialnum[c];
+  // Convert to UCS-2, which is equivalent to UTF-16, if the input is ASCII
+  for (int i = 0; i < SERIALNUM_LEN; i++) {
+      if (ascii_serialnum[i] > 128) {
+          dprintf("Serial number character %d is not valid ASCII.", ascii_serialnum[i]);
+          SerialNumberString.Header.Size = i * 2;
+          break;
+      }
+      SerialNumberString.UnicodeString[i] = ascii_serialnum[i];
   }
 
-  //SerialNumberString.UnicodeString = &utf16_serialnum[0];
-  memcpy(SerialNumberString.UnicodeString, utf16_serialnum, SERIALNUM_LEN*2);
-  SerialNumberString.Header.Size = SERIALNUM_LEN * 2;
-
   return &SerialNumberString;
 }
 uint16_t lotus_serial_number_string_len(void) {
-  return SERIALNUM_LEN*2;
+  return pgm_read_byte(&SerialNumberString.Header.Size);
 }
 
 void keyboard_post_init_kb(void) {
diff --git a/keyboards/lotus/lotus.h b/keyboards/lotus/lotus.h
index b3c19dd4ae..134461f2fc 100644
--- a/keyboards/lotus/lotus.h
+++ b/keyboards/lotus/lotus.h
@@ -38,3 +38,9 @@ uint16_t lotus_serial_number_string_len(void);
 #define FLASH_OFFSET 0x10000000
 #define LAST_4K_BLOCK 0xff000
 #define SERIALNUM_LEN 18
+
+// 18-digit serial number. Must be programmed at factory
+// TODO: Figure out how to do that. Probably a script to compile qmk
+// Must be wide-string (UTF-16)
+#define FALLBACK_SERIAL_NUMBER L"FRALDLvvccxxxxxxxx"
+_Static_assert(sizeof(FALLBACK_SERIAL_NUMBER) == (SERIALNUM_LEN+1)*2, "Fallback serial number must be same length");
diff --git a/tmk_core/protocol/usb_descriptor.c b/tmk_core/protocol/usb_descriptor.c
index 96a942f97c..02f1c089d3 100644
--- a/tmk_core/protocol/usb_descriptor.c
+++ b/tmk_core/protocol/usb_descriptor.c
@@ -1078,17 +1078,6 @@ const USB_Descriptor_String_t PROGMEM ProductString = {
     .UnicodeString              = USBSTR(PRODUCT)
 };
 
-#if defined(SERIAL_NUMBER)
-extern const USB_Descriptor_String_t PROGMEM SerialNumberString;
-//const USB_Descriptor_String_t PROGMEM SerialNumberString = {
-//    .Header = {
-//        .Size                   = sizeof(USBSTR(SERIAL_NUMBER)),
-//        .Type                   = DTYPE_String
-//    },
-//    .UnicodeString              = USBSTR(SERIAL_NUMBER)
-//};
-#endif
-
 // clang-format on
 
 /**
@@ -1134,6 +1123,7 @@ uint16_t get_usb_descriptor(const uint16_t wValue, const uint16_t wIndex, const
                     break;
 #if defined(SERIAL_NUMBER)
                 case 0x03:
+                    // TODO: Give these functions a generic name and let anyone override it
                     Address = lotus_serial_number_string();
                     Size    = lotus_serial_number_string_len();