From c4bbeccac0495685a5e6cca54e28cddd14071be5 Mon Sep 17 00:00:00 2001
From: Torrie Fischer <torrie@ripple.com>
Date: Sat, 23 Dec 2023 11:05:45 +0100
Subject: [PATCH] platform: arduino: ota: prevent watchdog hang on ota update

---
 src/platform/arduino/OTA.cpp | 23 ++++++++++++++++++++++-
 src/platform/arduino/OTA.h   |  3 +++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/src/platform/arduino/OTA.cpp b/src/platform/arduino/OTA.cpp
index 7f205a0..44feacd 100644
--- a/src/platform/arduino/OTA.cpp
+++ b/src/platform/arduino/OTA.cpp
@@ -4,10 +4,13 @@
 ArduinoOTAUpdater::ArduinoOTAUpdater() : BufferedInputSource("ArduinoOTA") {
   ArduinoOTA.onStart(&ArduinoOTAUpdater::s_onStart);
   ArduinoOTA.onProgress(&ArduinoOTAUpdater::s_onProgress);
+  ArduinoOTA.onError(&ArduinoOTAUpdater::s_onError);
+  ArduinoOTA.onEnd(&ArduinoOTAUpdater::s_onFinished);
+  ArduinoOTA.setRebootOnSuccess(false);
 }
 
 void ArduinoOTAUpdater::loop() {
-    if (m_online) {
+    if (m_online && !m_updating) {
         ArduinoOTA.handle();
     }
     BufferedInputSource::loop();
@@ -20,14 +23,32 @@ void ArduinoOTAUpdater::handleEvent(const InputEvent& evt) {
         ArduinoOTA.begin();
     }
 }
+
+void ArduinoOTAUpdater::s_onFinished()
+{
+  Log.notice("OTA complete!");
+  Static<ArduinoOTAUpdater>::instance()->m_updating = false;
+  Platform::restart();
+}
+
 void ArduinoOTAUpdater::s_onStart() {
     Log.notice("OTA Start!");
+    Static<ArduinoOTAUpdater>::instance()->m_updating = true;
     Static<ArduinoOTAUpdater>::instance()->setEvent(InputEvent::FirmwareUpdate);
 }
 
+void ArduinoOTAUpdater::s_onError(ota_error_t err)
+{
+  Log.notice("OTA failure: %d", (int)err);
+  Static<ArduinoOTAUpdater>::instance()->m_updating = false;
+}
+
 void ArduinoOTAUpdater::s_onProgress(unsigned int progress, unsigned int total) {
     Log.notice("OTA Progress! %d / %d", progress, total);
     Static<ArduinoOTAUpdater>::instance()->setEvent(InputEvent{InputEvent::FirmwareUpdate, progress});
+    // Try to run the main loop in case we're not in a threaded updater
+    // environment. This keeps rendering running, and pets the watchdog.
+    MainLoop::instance()->loop();
 }
 
 STATIC_ALLOC(ArduinoOTAUpdater);
diff --git a/src/platform/arduino/OTA.h b/src/platform/arduino/OTA.h
index f840666..6300e6b 100644
--- a/src/platform/arduino/OTA.h
+++ b/src/platform/arduino/OTA.h
@@ -9,6 +9,9 @@ class ArduinoOTAUpdater : public BufferedInputSource {
 
   private:
     bool m_online = false;
+    bool m_updating = false;
     static void s_onStart();
     static void s_onProgress(unsigned int progress, unsigned int total);
+    static void s_onError(ota_error_t err);
+    static void s_onFinished();
 };