From 44dd485e179b110924b2e600f45cd56b6cc41dce Mon Sep 17 00:00:00 2001
From: Manuel Weiser <manuel.weiser@me.com>
Date: Thu, 20 Feb 2025 20:04:58 +0100
Subject: [PATCH] fix: improve OTA update handling and logging for better error
 reporting

---
 platformio.ini |  3 ++-
 src/ota.cpp    | 43 +++++++++++++++++++++++++++++++++----------
 2 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/platformio.ini b/platformio.ini
index da18ffb..de6da6e 100644
--- a/platformio.ini
+++ b/platformio.ini
@@ -45,7 +45,7 @@ build_flags =
     -mtext-section-literals
     '-D VERSION="${common.version}"'
     -DASYNCWEBSERVER_REGEX
-    -DCORE_DEBUG_LEVEL=1
+    -DCORE_DEBUG_LEVEL=3
     -DCONFIG_ARDUHAL_LOG_COLORS=1
     -DOTA_DEBUG=1
     -DARDUINO_RUNNING_CORE=1
@@ -66,6 +66,7 @@ build_flags =
     -DCONFIG_PARTITION_TABLE_OFFSET=0x8000
     -DCONFIG_PARTITION_TABLE_MD5=y
     -DBOOT_APP_PARTITION_OTA_0=1
+    -DCONFIG_LOG_DEFAULT_LEVEL=3
     
 extra_scripts = 
     scripts/extra_script.py
diff --git a/src/ota.cpp b/src/ota.cpp
index dc79953..24468b5 100644
--- a/src/ota.cpp
+++ b/src/ota.cpp
@@ -22,7 +22,7 @@ void stopAllTasks() {
 
 void handleOTAUpload(AsyncWebServerRequest *request, String filename, size_t index, uint8_t *data, size_t len, bool final) {
     if (!index) {
-        bool isFullImage = true;
+        bool isFullImage = filename.endsWith(".bin");
         Serial.printf("Update Start: %s (type: %s)\n", filename.c_str(), isFullImage ? "full" : "OTA");
         
         if (request->contentLength() == 0) {
@@ -30,46 +30,69 @@ void handleOTAUpload(AsyncWebServerRequest *request, String filename, size_t ind
             return;
         }
 
+        // Berechne verfügbaren Speicherplatz
+        size_t updateSize = request->contentLength();
+        size_t maxSketchSpace = (ESP.getFreeSketchSpace() - 0x1000) & 0xFFFFF000;
+        
+        Serial.printf("Update size: %u bytes\n", updateSize);
+        Serial.printf("Available space: %u bytes\n", maxSketchSpace);
+        
+        if (updateSize > maxSketchSpace) {
+            Serial.printf("Error: Not enough space. Need %u bytes but only have %u bytes available\n", 
+                        updateSize, maxSketchSpace);
+            request->send(400, "application/json", 
+                        "{\"status\":\"error\",\"message\":\"Not enough space for update\"}");
+            return;
+        }
+
         if (!tasksAreStopped && (RfidReaderTask || BambuMqttTask || ScaleTask)) {
             stopAllTasks();
             tasksAreStopped = true;
         }
 
+        // Ensure SPIFFS is ended before update
+        if (SPIFFS.begin()) {
+            SPIFFS.end();
+        }
+
         bool success;
         if (isFullImage) {
-            // Full image update ohne Magic Byte Check, aber mit U_FLASH
-            success = Update.begin(UPDATE_SIZE_UNKNOWN, U_FLASH);
+            success = Update.begin(updateSize, U_FLASH);
         } else {
-            // Normales OTA update mit Magic Byte Check
             if (data[0] != 0xE9) {
                 Serial.printf("Wrong magic byte: 0x%02X (expected 0xE9)\n", data[0]);
-                request->send(400, "application/json", "{\"status\":\"error\",\"message\":\"Invalid firmware format\"}");
+                request->send(400, "application/json", 
+                            "{\"status\":\"error\",\"message\":\"Invalid firmware format\"}");
                 return;
             }
-            success = Update.begin(request->contentLength());
+            success = Update.begin(updateSize);
         }
 
         if (!success) {
             Update.printError(Serial);
-            request->send(400, "application/json", "{\"status\":\"error\",\"message\":\"Update start failed\"}");
+            request->send(400, "application/json", 
+                        "{\"status\":\"error\",\"message\":\"Update initialization failed\"}");
             return;
         }
     }
 
     if (Update.write(data, len) != len) {
         Update.printError(Serial);
-        request->send(400, "application/json", "{\"status\":\"error\",\"message\":\"Write failed\"}");
+        request->send(400, "application/json", 
+                    "{\"status\":\"error\",\"message\":\"Write failed\"}");
         return;
     }
 
     if (final) {
         if (!Update.end(true)) {
             Update.printError(Serial);
-            request->send(400, "application/json", "{\"status\":\"error\",\"message\":\"Update failed\"}");
+            request->send(400, "application/json", 
+                        "{\"status\":\"error\",\"message\":\"Update failed\"}");
             return;
         }
         Serial.println("Update successful, restarting...");
-        request->send(200, "application/json", "{\"status\":\"success\",\"message\":\"Update successful! Device will restart...\",\"restart\":true}");
+        request->send(200, "application/json", 
+                    "{\"status\":\"success\",\"message\":\"Update successful! Device will restart...\",\"restart\":true}");
         delay(500);
         ESP.restart();
     }