add bad usb

fw/Core/Hitcon/App/BadUsbApp.cc

@@ -0,0 +1,23 @@
+#include <App/BadUsbApp.h>
+#include <Logic/BadgeController.h>
+#include <Logic/UsbLogic.h>
+
+namespace hitcon {
+namespace usb {
+BadUsbApp bad_usb_app;
+
+BadUsbApp::BadUsbApp() {}
+
+void BadUsbApp::OnEntry() {
+  g_usb_logic.RunScript((callback_t)&BadUsbApp::OnScriptFinished, this);
+}
+
+void BadUsbApp::OnExit() { g_usb_logic.StopScript(); }
+
+void BadUsbApp::OnButton(button_t button) {}
+
+void BadUsbApp::OnScriptFinished(void *unsed) {
+  badge_controller.OnAppEnd(this);
+}
+}  // namespace usb
+}  // namespace hitcon

fw/Core/Hitcon/App/BadUsbApp.h

@@ -0,0 +1,26 @@
+#ifndef BAD_USB_APP_H
+#define BAD_USB_APP_H
+
+#include <App/app.h>
+
+namespace hitcon {
+namespace usb {
+
+class BadUsbApp : public App {
+ private:
+  void OnScriptFinished(void* unused);
+
+ public:
+  BadUsbApp();
+  virtual ~BadUsbApp() = default;
+
+  void OnEntry() override;
+  void OnExit() override;
+  void OnButton(button_t button) override;
+};
+
+extern BadUsbApp bad_usb_app;
+}  // namespace usb
+}  // namespace hitcon
+
+#endif

fw/Core/Hitcon/App/MainMenuApp.h

@@ -1,3 +1,4 @@
+#include <App/BadUsbApp.h>
 #include <App/DinoApp.h>
 #include <App/ShowNameApp.h>
 #include <App/SnakeApp.h>
@@ -14,7 +15,7 @@ using hitcon::app::tetris::tetris_app;
 
 constexpr menu_entry_t main_menu_entries[] = {
     // TODO : change app
-    {"BadUSB", nullptr, nullptr},
+    {"BadUSB", &hitcon::usb::bad_usb_app, nullptr},
     {"Snake", &snake_app, &hitcon::app::snake::SetSingleplayer},
     {"Dino", &dino_app, nullptr},
     {"Tetris", &tetris_app, &hitcon::app::tetris::SetSingleplayer},

fw/Core/Hitcon/Hitcon.cpp

@@ -77,8 +77,8 @@ void hitcon_run() {
   hitcon::ir::irController.Init();
   hitcon::app::snake::snake_app.Init();
   hitcon::app::dino::dino_app.Init();
-  g_usb_logic.Init();
   g_xboard_game_controller.Init();
+  hitcon::usb::g_usb_logic.Init();
 
   // run hardware test mode if MODE/SETTINGS Button is pressed during
   // initializing

fw/Core/Hitcon/Logic/NvStorage.cc

@@ -23,7 +23,7 @@ NvStorage::NvStorage()
 void NvStorage::Init() {
   int32_t newest_version = -1;
   my_assert(!g_flash_service.IsBusy());
-  for (size_t i = 0; i < FLASH_PAGE_COUNT; i++) {
+  for (size_t i = 1; i < FLASH_PAGE_COUNT; i++) {
     uint8_t* page_data =
         reinterpret_cast<uint8_t*>(g_flash_service.GetPagePointer(i));
     nv_storage_content* page_content =
@@ -64,6 +64,7 @@ void NvStorage::ForceFlushInternal() {
     storage_dirty_ = false;
     next_available_page = (next_available_page + 1) %
                           FLASH_PAGE_COUNT;  // Increment for the next write
+    if (next_available_page == 0) next_available_page = 1;
     content_.version++;
     last_flush_cycle = current_cycle;  // Record the current cycle
   }

fw/Core/Hitcon/Logic/UsbLogic.cc

@@ -1,106 +1,155 @@
 #include <App/ShowNameApp.h>
 #include <Logic/NvStorage.h>
 #include <Logic/UsbLogic.h>
+#include <Service/FlashService.h>
 #include <Service/Sched/Scheduler.h>
-#include <main.h>
 #include <usb_device.h>
 #include <usbd_custom_hid_if.h>
 
 using namespace hitcon::service::sched;
-// TODO: check name length
+
+/* TODO:
+ * 1. check name length
+ * 2. add usb on connect
+ * 3. add run script progress (XX%)
+ * 4. add hack api
+ */
 
 namespace hitcon {
+namespace usb {
+
 UsbLogic g_usb_logic;
 
 UsbLogic::UsbLogic()
-    : _routine_task(810, (task_callback_t)&UsbLogic::Routine, (void*)this,
-                    DELAY_INTERVAL) {}
+    : on_recv_task(823, (task_callback_t)&UsbLogic::OnDataRecv, this),
+      _routine_task(810, (task_callback_t)&UsbLogic::Routine, (void*)this,
+                    DELAY_INTERVAL),
+      _write_routine_task(810, (task_callback_t)&UsbLogic::WriteRoutine,
+                          (void*)this, WAIT_INTERVAL) {}
 
 void UsbLogic::Init() {
   _state = USB_STATE_HEADER;
   _index = 0;
   scheduler.Queue(&_routine_task, nullptr);
+  scheduler.Queue(&_write_routine_task, nullptr);
 }
-/* TODO:
- * 1. check recv have buffer? or set
- * 2. get recv data length
- * 3. implement write script on flash
- * 4. test run script
- * 5. fix set_name
- */
-void UsbLogic::OnDataRecv(uint8_t* data) {
-  // check length?, call this function at outevent
+
+void UsbLogic::OnDataRecv(void* arg) {
+  uint8_t* data = reinterpret_cast<uint8_t*>(arg);
+  static size_t counter = 0;
+  counter++;
+  if (_state == USB_STATE_ERASE) return;  // erase not finish
+  for (uint8_t i = 0; i < RECV_BUF_LEN; i++) {
+    _temp[i] = data[i];
+  }
+
   if (_state == USB_STATE_HEADER) {
     if (data[0]) {
       _state = static_cast<usb_state_t>(data[0]);
       _index = 0;
     }
   }
   nv_storage_content& content = g_nv_storage.GetCurrentStorage();
-
-  FLASH_EraseInitTypeDef erase_struct = {
-      .TypeErase = FLASH_TYPEERASE_PAGES,
-      .PageAddress = SCRIPT_BEGIN_ADDR,
-      .NbPages = MY_FLASH_PAGE_SIZE / FLASH_PAGE_SIZE,
-  };
-  uint32_t error;
   switch (_state) {
     case USB_STATE_SET_NAME:
-      if (_index == NAME_LEN) {
-        //        g_nv_storage.MarkDirty();
-        show_name_app.SetName(reinterpret_cast<const char*>(_name));
-        _state = USB_STATE_HEADER;
-      }
       for (size_t i = (_index == 0); i < RECV_BUF_LEN; i++, _index++) {
-        //        content.name[_index] = data[i];
-        _name[_index] = data[i];
-        // TODO: maybe need a temp buffer?
+        content.name[_index] = data[i];
+        if (_index - 1 == hitcon::ShowNameApp::NAME_LEN) {
+          show_name_app.SetName(reinterpret_cast<const char*>(content.name));
+          _state = USB_STATE_HEADER;
+          break;
+        }
       }
       break;
-    case USB_STATE_CLEAR:
-      HAL_FLASH_Unlock();
-      HAL_FLASHEx_Erase(&erase_struct, &error);
-      HAL_FLASH_Lock();
-      if (error != 0xFFFFFFFF) {
-        // TODO: error erase flash
-        return;
-      }
-      _state = USB_STATE_HEADER;
+    case USB_STATE_ERASE:
+      g_flash_service.ErasePage(SCRIPT_FLASH_INDEX);  // TODO: check is busy
+      scheduler.EnablePeriodic(&_write_routine_task);
       break;
     case USB_STATE_START_WRITE:
-      // TODO: use flash service to program page
       // TODO: add checksum
-
-      break;
-    case USB_STATE_STOP_WRITE:
-      _state = USB_STATE_HEADER;
+      _script_len = (data[1] << 8) | data[2];
+      memcpy(_temp, data, RECV_BUF_LEN);
+      _state = USB_STATE_WRITING;
+      _new_data = true;
+      scheduler.EnablePeriodic(&_write_routine_task);
+    case USB_STATE_WRITING:
+      // it seems like first byte of data cannot be 0 so use 0xFC instead
+      if (data[0] == 0xFC) data[0] = 0;
+      memcpy(_temp, data, RECV_BUF_LEN);
+      _new_data = true;
       break;
     default:
       break;
   }
 }
 
-void UsbLogic::RunScript() {
+void UsbLogic::WriteRoutine(void* unused) {
+  if (_state == USB_STATE_ERASE) {
+    if (!g_flash_service.IsBusy()) {
+      _state = USB_STATE_HEADER;
+      scheduler.DisablePeriodic(&_write_routine_task);
+      keyboard_report = {0, 0, 0xFF, 0, 0, 0, 0, 0};
+      USBD_CUSTOM_HID_SendReport(
+          &hUsbDeviceFS, reinterpret_cast<uint8_t*>(&keyboard_report), 8);
+    }
+  } else if (_state == USB_STATE_WRITING) {
+    if (!g_flash_service.IsBusy() && _new_data) {
+      keyboard_report = {0, 0, 0xFF, 0, 0, 0, 0, 0};
+      USBD_CUSTOM_HID_SendReport(
+          &hUsbDeviceFS, reinterpret_cast<uint8_t*>(&keyboard_report), 8);
+      g_flash_service.ProgramOnly(SCRIPT_FLASH_INDEX, _index,
+                                  reinterpret_cast<uint32_t*>(_temp), 8);
+      _index += 8;
+      _new_data = false;
+      if (_index >= _script_len + 8) {
+        _state = USB_STATE_HEADER;
+        scheduler.DisablePeriodic(&_write_routine_task);
+      }
+    }
+  }
+}
+
+void UsbLogic::RunScript(callback_t cb, void* arg1) {
+  _on_finish_cb = cb;
+  _on_finish_arg1 = arg1;
   _index = -1;
   keyboard_report = {0, 0, 0, 0, 0, 0, 0, 0};
+  empty_report = {0, 0, 0, 0, 0, 0, 0, 0};
   scheduler.EnablePeriodic(&_routine_task);
+  flag = false;
 }
 
-// run every 10ms
+void UsbLogic::StopScript() { scheduler.DisablePeriodic(&_routine_task); }
+
+// run every 20ms
 void UsbLogic::Routine(void* unused) {
   static uint8_t delay_count = 0;
+  uint16_t len = (*(uint8_t*)(SCRIPT_BEGIN_ADDR - 2) << 8) |
+                 *(uint8_t*)(SCRIPT_BEGIN_ADDR - 1);
+  if (flag) {
+    flag = false;
+
+    if (keyboard_report.modifier == 0)
+      USBD_CUSTOM_HID_SendReport(&hUsbDeviceFS,
+                                 reinterpret_cast<uint8_t*>(&empty_report), 8);
+    return;
+  }
   if (delay_count != 0) {
     delay_count--;
     return;
-  } else if (_index == *reinterpret_cast<uint16_t*>(SCRIPT_BEGIN_ADDR -
-                                                    2)) {  // finish script
+  } else if (_index == len) {
     scheduler.DisablePeriodic(&_routine_task);
+    _on_finish_cb(_on_finish_arg1, nullptr);
     return;
   }
-
   if (_index == -1) {  // new script begin
     delay_count = 0;
   } else {
+    uint8_t progress = _index * 100 / len;
+    char str[4] = "XX%";
+    str[0] = progress / 10 + '0';
+    str[1] = progress % 10 + '0';
+    display_set_mode_text((str));
     uint8_t* addr = reinterpret_cast<uint8_t*>(SCRIPT_BEGIN_ADDR + _index);
     switch (*addr) {
       case CODE_DELAY:
@@ -121,13 +170,15 @@ void UsbLogic::Routine(void* unused) {
     }
     USBD_CUSTOM_HID_SendReport(&hUsbDeviceFS,
                                reinterpret_cast<uint8_t*>(&keyboard_report), 8);
+    flag = true;
+    delay_count++;
   }
   _index++;
 }
-
+}  // namespace usb
 }  // namespace hitcon
 
 void UsbServiceOnDataReceived(uint8_t* data) {
-  hitcon::g_usb_logic.OnDataRecv(data);
-  // TODO: queue task
+  scheduler.Queue(&hitcon::usb::g_usb_logic.on_recv_task,
+                  reinterpret_cast<uint8_t*>(data));
 }