ota_update.cpp

#include "ota_update.h"
#include "wled.h"

#ifdef ESP32
#include <esp_app_format.h>
#include <esp_ota_ops.h>
#include <esp_flash.h>
#include <mbedtls/sha256.h>

// Platform-specific metadata locations
constexpr size_t METADATA_OFFSET = 256;          // ESP32: metadata appears after Espressif metadata
#define UPDATE_ERROR errorString

// Bootloader is at fixed offset 0x1000 (4KB), 0x0000 (0KB), or 0x2000 (8KB), and is typically 32KB
// Bootloader offsets for different MCUs => see https://github.com/wled/WLED/issues/5064
#if defined(CONFIG_IDF_TARGET_ESP32S3) || defined(CONFIG_IDF_TARGET_ESP32C3) || defined(CONFIG_IDF_TARGET_ESP32C6)
constexpr size_t BOOTLOADER_OFFSET = 0x0000; // esp32-S3, esp32-C3 and (future support) esp32-c6
constexpr size_t BOOTLOADER_SIZE   = 0x8000; // 32KB, typical bootloader size
#define BOOTLOADER_OTA_UNSUPPORTED    // still needs validation on these platforms.
#elif defined(CONFIG_IDF_TARGET_ESP32P4) || defined(CONFIG_IDF_TARGET_ESP32C5)
constexpr size_t BOOTLOADER_OFFSET = 0x2000; // (future support) esp32-P4 and esp32-C5
constexpr size_t BOOTLOADER_SIZE   = 0x8000; // 32KB, typical bootloader size
#define BOOTLOADER_OTA_UNSUPPORTED    // still needs testing on these platforms
#else
constexpr size_t BOOTLOADER_OFFSET = 0x1000; // esp32 and esp32-s2
constexpr size_t BOOTLOADER_SIZE   = 0x8000; // 32KB, typical bootloader size
#endif

#elif defined(ESP8266)
constexpr size_t METADATA_OFFSET = 0x1000;     // ESP8266: metadata appears at 4KB offset
#define UPDATE_ERROR getErrorString
#define SUPPORT_GZIPPED_OTA
#endif

#ifdef SUPPORT_GZIPPED_OTA
#include <uzlib.h>  // ArduinoUZlib library: gzipped firmware uploads are decoded by eboot at boot,
                    // so the upload stream itself can arrive compressed; we decompress just enough
                    // to find the metadata for OTA validation.
#endif

constexpr size_t METADATA_SEARCH_RANGE = 512;  // bytes

// State structure for update process
namespace {
  struct UpdateContext {
    // State flags
    // FUTURE: the flags could be replaced by a state machine
    bool replySent = false;
    bool needsRestart = false;
    bool updateStarted = false;
    bool uploadComplete = false;
    bool releaseCheckPassed = false;
    String errorMessage;

    // Buffer to hold block data across posts, if needed
    std::vector<uint8_t> releaseMetadataBuffer;

    #ifdef SUPPORT_GZIPPED_OTA
    bool gzipDetected         = false;  // upload is a gzip stream (eboot will decompress at boot)
    #endif
  };
}

/**
 * Check if OTA should be allowed based on release compatibility using custom description
 * @param binaryData Pointer to binary file data (not modified)
 * @param dataSize Size of binary data in bytes
 * @param errorMessage Buffer to store error message if validation fails
 * @param errorMessageLen Maximum length of error message buffer
 * @return true if OTA should proceed, false if it should be blocked
 */

static bool validateOTA(const uint8_t* binaryData, size_t dataSize, char* errorMessage, size_t errorMessageLen) {
  // Clear error message
  if (errorMessage && errorMessageLen > 0) {
    errorMessage[0] = '\0';
  }

  // Try to extract WLED structure directly from binary data
  wled_metadata_t extractedDesc;
  bool hasDesc = findWledMetadata(binaryData, dataSize, &extractedDesc);

  if (hasDesc) {
    return shouldAllowOTA(extractedDesc, errorMessage, errorMessageLen);
  } else {
    // No custom description - this could be a legacy binary
    if (errorMessage && errorMessageLen > 0) {
      strncpy_P(errorMessage, PSTR("This firmware file is missing compatibility metadata."), errorMessageLen - 1);
      errorMessage[errorMessageLen - 1] = '\0';
    }
    return false;
  }
}

#ifdef SUPPORT_GZIPPED_OTA
constexpr size_t GZIP_DECOMPRESSED_SIZE = METADATA_OFFSET + METADATA_SEARCH_RANGE;
// We use the JSON buffer to hold the decompressed metadata, so we need to ensure it's large enough
static_assert(JSON_BUFFER_SIZE >= GZIP_DECOMPRESSED_SIZE, "JSON_BUFFER_SIZE must be at least GZIP_DECOMPRESSED_SIZE to support gzipped OTA updates");

// File-static flash source context for the uzlib source_read_cb.
// The 256-byte read buffer must be 4-byte aligned for ESP.flashRead().
static struct {
  uint32_t flashBase;
  uint32_t flashOffset;
  uint8_t  buf[256] __attribute__((aligned(4)));
} s_gzipFlashSrc;

static int gzip_flash_read_cb(struct uzlib_uncomp* d)
{
  if (!ESP.flashRead(s_gzipFlashSrc.flashBase + s_gzipFlashSrc.flashOffset,
                     reinterpret_cast<uint32_t*>(s_gzipFlashSrc.buf),
                     sizeof(s_gzipFlashSrc.buf))) {
    return -1;
  }
  s_gzipFlashSrc.flashOffset += sizeof(s_gzipFlashSrc.buf);
  d->source       = s_gzipFlashSrc.buf;
  d->source_limit = s_gzipFlashSrc.buf + sizeof(s_gzipFlashSrc.buf);
  return *d->source++;
}

static bool unzipFlash(uint32_t flashBase, uint8_t* outBuf, char* errorMessage, size_t errorMessageLen)
{
  uzlib_uncomp uncomp_ctx;  // large, but we can get away with this since we're on the system stack

  s_gzipFlashSrc.flashBase   = flashBase;
  s_gzipFlashSrc.flashOffset = 0;
  uncomp_ctx.source         = s_gzipFlashSrc.buf;
  uncomp_ctx.source_limit   = s_gzipFlashSrc.buf;   // empty: first byte triggers callback
  uncomp_ctx.source_read_cb = gzip_flash_read_cb;
  uzlib_uncompress_init(&uncomp_ctx, nullptr, 0);

  bool ok = false;
  if (uzlib_gzip_parse_header(&uncomp_ctx) != TINF_OK) {
    strncpy_P(errorMessage, PSTR("Invalid gzip header in OTA firmware"), errorMessageLen - 1);
    errorMessage[errorMessageLen - 1] = '\0';
  } else {
    uncomp_ctx.dest_start = uncomp_ctx.dest = outBuf;
    uncomp_ctx.dest_limit = outBuf + GZIP_DECOMPRESSED_SIZE;
    int res = TINF_OK;
    while (res == TINF_OK && uncomp_ctx.dest < uncomp_ctx.dest_limit) res = uzlib_uncompress(&uncomp_ctx);
    if (uncomp_ctx.dest >= uncomp_ctx.dest_limit) {
      ok = true;
    } else {
      strncpy_P(errorMessage, PSTR("Not enough decompressed data for validation"), errorMessageLen - 1);
      errorMessage[errorMessageLen - 1] = '\0';
    }
  }

  return ok;
}

static void validateGzippedOTA(UpdateContext* context) {
  // Compute the flash address to read from based on the expected update size and the location of the FS partition, and then decompress just enough of the firmware to find the metadata for validation.
  extern uint32_t _FS_start;
  const size_t updateSize = Update.size();
  // Logic borrowed from ESP8266 Updater.cpp
  const size_t roundedSize = (updateSize + FLASH_SECTOR_SIZE - 1) & ~(FLASH_SECTOR_SIZE - 1);
  const uintptr_t fsPhysAddr = (uintptr_t)&_FS_start - 0x40200000;
  const uint32_t flashBase = (fsPhysAddr > roundedSize) ? (fsPhysAddr - roundedSize) : 0;
  uint8_t* compressionBuf = static_cast<uint8_t*>(pDoc->memoryPool().buffer()); // borrow JSON buffer

  char errorMessage[128];
  bool ok = unzipFlash(flashBase, compressionBuf, errorMessage, sizeof(errorMessage));
  if (!ok) {
    DEBUG_PRINTF_P(PSTR("OTA unzip failed: %s\n"), errorMessage);
    context->errorMessage = errorMessage;
    return;
  }

  ok = validateOTA(compressionBuf, GZIP_DECOMPRESSED_SIZE, errorMessage, sizeof(errorMessage));
  if (!ok) {
    DEBUG_PRINTF_P(PSTR("OTA declined: %s\n"), errorMessage);
    context->errorMessage = errorMessage;
    context->errorMessage += F(" Enable 'Ignore firmware validation' to proceed anyway.");
    return;
  }

  DEBUG_PRINTLN(F("OTA allowed: Release compatibility check passed (gzipped)"));
  context->releaseCheckPassed = true;

  // Write the final block
  if (!Update.hasError()) {
    const auto& buf = context->releaseMetadataBuffer;
    if (Update.write(const_cast<uint8_t*>(buf.data()), buf.size()) != buf.size()) {
      DEBUG_PRINTF_P(PSTR("OTA write failed on final chunk: %s\n"), Update.UPDATE_ERROR());
    }
  }

  DEBUG_PRINTLN(F("OTA Update End")); // for symmetry with the non-gzip path
  context->uploadComplete = true;   // Final block was passed to Update.write()
}
#endif

static void endOTA(AsyncWebServerRequest *request) {
  UpdateContext* context = reinterpret_cast<UpdateContext*>(request->_tempObject);
  request->_tempObject = nullptr;

  DEBUG_PRINTF_P(PSTR("EndOTA %x --> %x (%d)\n"), (uintptr_t)request,(uintptr_t) context, context ? context->uploadComplete : 0);
  if (context) {
    if (context->updateStarted) {  // We initialized the update
      // We use Update.end() because not all forms of Update() support an abort.
      // If the upload is incomplete, Update.end(false) should error out.
      if (Update.end(context->uploadComplete)) {
        // Update successful!
        #ifndef ESP8266
        bootloopCheckOTA(); // let the bootloop-checker know there was an OTA update
        #endif
        doReboot = true;
        context->needsRestart = false;
      }
    }

    if (context->needsRestart) {
      strip.resume();
      UsermodManager::onUpdateBegin(false);
      #if WLED_WATCHDOG_TIMEOUT > 0
      WLED::instance().enableWatchdog();
      #endif
    }
    delete context;
  }
};

static bool beginOTA(AsyncWebServerRequest *request, UpdateContext* context)
{
  #ifdef ESP8266
  Update.runAsync(true);
  #endif

  if (Update.isRunning()) {
      request->send(503);
      setOTAReplied(request);
      return false;
  }

  #if WLED_WATCHDOG_TIMEOUT > 0
  WLED::instance().disableWatchdog();
  #endif
  UsermodManager::onUpdateBegin(true); // notify usermods that update is about to begin (some may require task de-init)

  strip.suspend();
  backupConfig(); // backup current config in case the update ends badly
  strip.resetSegments();  // free as much memory as you can
  context->needsRestart = true;

  DEBUG_PRINTF_P(PSTR("OTA Update Start, %x --> %x\n"), (uintptr_t)request,(uintptr_t) context);

  auto skipValidationParam = request->getParam("skipValidation", true);
  if (skipValidationParam && (skipValidationParam->value() == "1")) {
    context->releaseCheckPassed = true;
    DEBUG_PRINTLN(F("OTA validation skipped by user"));
  }

  // Begin update with the firmware size from content length
  size_t updateSize = request->contentLength() > 0 ? request->contentLength() : ((ESP.getFreeSketchSpace() - 0x1000) & 0xFFFFF000);
  if (!Update.begin(updateSize)) {
    context->errorMessage = Update.UPDATE_ERROR();
    DEBUG_PRINTF_P(PSTR("OTA Failed to begin: %s\n"), context->errorMessage.c_str());
    return false;
  }

  context->updateStarted = true;
  return true;
}

// Create an OTA context object on an AsyncWebServerRequest
// Returns true if successful, false on failure.
bool initOTA(AsyncWebServerRequest *request) {
  // Allocate update context
  UpdateContext* context = new (std::nothrow) UpdateContext {};
  if (context) {
    request->_tempObject = context;
    request->onDisconnect([=]() { endOTA(request); });  // ensures we restart on failure
  };

  DEBUG_PRINTF_P(PSTR("OTA Update init, %x --> %x\n"), (uintptr_t)request,(uintptr_t) context);
  return (context != nullptr);
}

void setOTAReplied(AsyncWebServerRequest *request) {
  UpdateContext* context = reinterpret_cast<UpdateContext*>(request->_tempObject);
  if (!context) return;
  context->replySent = true;
};

// Returns the OTA outcome as an OTAResultStatus plus an optional error string.
// TryAgain: JSON lock busy; caller must call deferResponse() and retry.
// Replied:  response already sent; no action needed.
// Ready:    result available; empty string = success (200), non-empty = failure (500).
std::pair<OTAResultStatus, String> getOTAResult(AsyncWebServerRequest* request) {
  UpdateContext* context = reinterpret_cast<UpdateContext*>(request->_tempObject);
  if (!context) return { OTAResultStatus::Ready, F("OTA context unexpectedly missing") };
  if (context->replySent) return { OTAResultStatus::Replied, {} };

  #ifdef SUPPORT_GZIPPED_OTA
  if (context->gzipDetected && !context->errorMessage.length()) {
    JSONBufferGuard jsonGuard(JSON_LOCK_OTA);
    if (!jsonGuard) return { OTAResultStatus::TryAgain, {} };  // busy — caller will deferResponse()
    validateGzippedOTA(context); // Stores error in context->errorMessage if there's a problem
    context->gzipDetected = false;  // all done
  }
  #endif

  if (context->errorMessage.length()) return { OTAResultStatus::Ready, context->errorMessage };

  if (context->updateStarted) {
    // Release the OTA context now.
    endOTA(request);
    if (Update.hasError()) {
      return { OTAResultStatus::Ready, Update.UPDATE_ERROR() };
    } else {
      return { OTAResultStatus::Ready, {} };
    }
  }

  // Should never happen
  return { OTAResultStatus::Ready, F("Internal software failure") };
}




void handleOTAData(AsyncWebServerRequest *request, size_t index, uint8_t *data, size_t len, bool isFinal)
{
  UpdateContext* context = reinterpret_cast<UpdateContext*>(request->_tempObject);
  if (!context) return;

  //DEBUG_PRINTF_P(PSTR("HandleOTAData: %d %d %d\n"), index, len, isFinal);

  if (context->replySent || (context->errorMessage.length())) return;

  if (index == 0) {
    if (!beginOTA(request, context)) return;

    #ifdef SUPPORT_GZIPPED_OTA
    if (len >= 2 && data[0] == 0x1F && data[1] == 0x8B) {
      DEBUG_PRINTLN(F("OTA: gzipped firmware detected"));
      context->gzipDetected = true;
    }
    #endif
  }

  // Perform plain-firmware validation as data arrives.
  // Gzip validation is deferred to the end of the upload; skip this path for gzip.
  if (!context->releaseCheckPassed
      #ifdef SUPPORT_GZIPPED_OTA
      && !context->gzipDetected
      #endif
  ) {
    if ((index+len) > METADATA_OFFSET) {
      // Current chunk contains the metadata offset
      size_t availableDataAfterOffset = (index + len) - METADATA_OFFSET;

      DEBUG_PRINTF_P(PSTR("OTA metadata check: %d in buffer, %d received, %d available\n"), context->releaseMetadataBuffer.size(), len, availableDataAfterOffset);

      if (availableDataAfterOffset >= METADATA_SEARCH_RANGE) {
        // We have enough data to validate, one way or another
        const uint8_t* search_data = data;
        size_t search_len = len;

        // If we have saved data, use that instead
        if (context->releaseMetadataBuffer.size()) {
          // Add this data
          context->releaseMetadataBuffer.insert(context->releaseMetadataBuffer.end(), data, data+len);
          search_data = context->releaseMetadataBuffer.data();
          search_len = context->releaseMetadataBuffer.size();
        }

        // Do the checking
        char errorMessage[128];
        bool OTA_ok = validateOTA(search_data, search_len, errorMessage, sizeof(errorMessage));

        // Release buffer if there was one
        context->releaseMetadataBuffer = decltype(context->releaseMetadataBuffer){};

        if (!OTA_ok) {
          DEBUG_PRINTF_P(PSTR("OTA declined: %s\n"), errorMessage);
          context->errorMessage = errorMessage;
          context->errorMessage += F(" Enable 'Ignore firmware validation' to proceed anyway.");
          return;
        } else {
          DEBUG_PRINTLN(F("OTA allowed: Release compatibility check passed"));
          context->releaseCheckPassed = true;
        }
      } else {
        // Store the data we just got for next pass
        context->releaseMetadataBuffer.insert(context->releaseMetadataBuffer.end(), data, data+len);
      }
    }
  }

#ifdef SUPPORT_GZIPPED_OTA
  // Gzip: stash the final chunk and defer validation to getOTAResult(), where
  // deferResponse() can be used if the JSON lock is not yet available.
  if (isFinal && context->gzipDetected && !context->releaseCheckPassed) {
    context->releaseMetadataBuffer.assign(data, data + len);
    return;  // getOTAResult() will validate, write this chunk, and complete the update
  }
#endif

  // Check if validation was still pending (shouldn't happen normally)
  // This is done before writing the last chunk, so endOTA can abort
  if (isFinal && !context->releaseCheckPassed) {
    DEBUG_PRINTLN(F("OTA failed: Validation never completed"));
    context->errorMessage = F("Release check data never arrived?");
    return;
  }

  if (!Update.hasError()) {
    if (Update.write(data, len) != len) {
      DEBUG_PRINTF_P(PSTR("OTA write failed on chunk %zu: %s\n"), index, Update.UPDATE_ERROR());
    }
  }

  if (isFinal) {
    DEBUG_PRINTLN(F("OTA Update End"));
    context->uploadComplete = true;
  }
}

void markOTAvalid() {
  #ifndef ESP8266
  const esp_partition_t* running = esp_ota_get_running_partition();
  esp_ota_img_states_t ota_state;
  if (esp_ota_get_state_partition(running, &ota_state) == ESP_OK) {
    if (ota_state == ESP_OTA_IMG_PENDING_VERIFY) {
      esp_ota_mark_app_valid_cancel_rollback(); // only needs to be called once, it marks the ota_state as ESP_OTA_IMG_VALID
      DEBUG_PRINTLN(F("Current firmware validated"));
    }
  }
  #endif
}

#if defined(ARDUINO_ARCH_ESP32) && !defined(WLED_DISABLE_OTA)

// Class for computing the expected bootloader data size given a stream of the data.
// If the image includes an SHA256 appended after the data stream, we do not consider it here.
class BootloaderImageSizer {
public:

  bool feed(const uint8_t* data, size_t len) {
    if (error) return false;

    //DEBUG_PRINTF("Feed %d\n", len);

    if (imageSize == 0) {
      // Parse header first
      if (len < sizeof(esp_image_header_t)) {
        error = true;
        return false;
      }

      esp_image_header_t header;
      memcpy(&header, data, sizeof(esp_image_header_t));

      if (header.segment_count == 0) {
        error = true;
        return false;
      }

      imageSize = sizeof(esp_image_header_t);
      segmentsLeft = header.segment_count;
      data += sizeof(esp_image_header_t);
      len -= sizeof(esp_image_header_t);
      //DEBUG_PRINTF("BLS parsed image header, segment count %d, is %d\n", segmentsLeft, imageSize);
    }

    while (len && segmentsLeft) {
      if (segmentHeaderBytes < sizeof(esp_image_segment_header_t)) {
        size_t headerBytes = std::min(len, sizeof(esp_image_segment_header_t) - segmentHeaderBytes);
        memcpy(reinterpret_cast<uint8_t*>(&segmentHeader) + segmentHeaderBytes, data, headerBytes);
        segmentHeaderBytes += headerBytes;
        if (segmentHeaderBytes < sizeof(esp_image_segment_header_t)) {
          return true;  // needs more bytes for the header
        }

        //DEBUG_PRINTF("BLS parsed segment [%08X %08X=%d], segment count %d, is %d\n", segmentHeader.load_addr, segmentHeader.data_len, segmentHeader.data_len, segmentsLeft, imageSize);

        // Validate segment size
        if (segmentHeader.data_len > BOOTLOADER_SIZE) {
          error = true;
          return false;
        }

        data += headerBytes;
        len -= headerBytes;
        imageSize += sizeof(esp_image_segment_header_t) + segmentHeader.data_len;
        --segmentsLeft;
        if (segmentsLeft == 0) {
          // all done, actually; we don't need to read any more

          // Round up to nearest 16 bytes.
          // Always add 1 to account for the checksum byte.
          imageSize = ((imageSize/ 16) + 1) * 16;

          //DEBUG_PRINTF("BLS complete, is %d\n", imageSize);
          return false;
        }
      }

      // If we don't have enough bytes ...
      if (len < segmentHeader.data_len) {
        //DEBUG_PRINTF("Needs more bytes\n");
        segmentHeader.data_len -= len;
        return true;  // still in this segment
      }

      // Segment complete
      len -= segmentHeader.data_len;
      data += segmentHeader.data_len;
      segmentHeaderBytes = 0;
      //DEBUG_PRINTF("Segment complete: len %d\n", len);
    }

    return !error;
  }

  bool hasError() const { return error; }
  bool isSizeKnown() const { return !error && imageSize != 0 && segmentsLeft == 0; }
  size_t totalSize() const {
    if (!isSizeKnown()) return 0;
    return imageSize;
  }

private:
  size_t imageSize = 0;
  size_t segmentsLeft = 0;
  esp_image_segment_header_t segmentHeader;
  size_t segmentHeaderBytes = 0;
  bool error = false;
};

static bool bootloaderSHA256CacheValid = false;
static uint8_t bootloaderSHA256Cache[32];

/**
 * Calculate and cache the bootloader SHA256 digest
 * Reads the bootloader from flash and computes SHA256 hash
 *
 * Strictly speaking, most bootloader images already contain a hash at the end of the image;
 * we could in theory just read it.  The trouble is that we have to parse the structure anyways
 * to find the actual endpoint, so we might as well always calculate it ourselves rather than
 * handle a special case if the hash isn't stored.
 *
 */
static void calculateBootloaderSHA256() {
  // Calculate SHA256
  mbedtls_sha256_context ctx;
  mbedtls_sha256_init(&ctx);
  mbedtls_sha256_starts(&ctx, 0); // 0 = SHA256 (not SHA224)

  const size_t chunkSize = 256;
  alignas(esp_image_header_t) uint8_t buffer[chunkSize];
  size_t bootloaderSize = BOOTLOADER_SIZE;
  BootloaderImageSizer sizer;
  size_t totalHashLen = 0;

  for (uint32_t offset = 0; offset < bootloaderSize; offset += chunkSize) {
    size_t readSize = min((size_t)(bootloaderSize - offset), chunkSize);
    if (esp_flash_read(NULL, buffer, BOOTLOADER_OFFSET + offset, readSize) == ESP_OK) {
      sizer.feed(buffer, readSize);

      size_t hashLen = readSize;
      if (sizer.isSizeKnown()) {
        size_t totalSize = sizer.totalSize();
        if (totalSize > 0 && totalSize <= BOOTLOADER_SIZE) {
          bootloaderSize = totalSize;
          if (offset + readSize > totalSize) {
            hashLen = (totalSize > offset) ? (totalSize - offset) : 0;
          }
        }
      }

      if (hashLen > 0) {
        totalHashLen += hashLen;
        mbedtls_sha256_update(&ctx, buffer, hashLen);
      }
    }
  }

  mbedtls_sha256_finish(&ctx, bootloaderSHA256Cache);
  mbedtls_sha256_free(&ctx);

  bootloaderSHA256CacheValid = true;
}

// Get bootloader SHA256 as hex string
String getBootloaderSHA256Hex() {
  if (!bootloaderSHA256CacheValid) {
    calculateBootloaderSHA256();
  }

  // Convert to hex string
  String result;
  result.reserve(65);
  for (int i = 0; i < 32; i++) {
    char b1 = bootloaderSHA256Cache[i];
    char b2 = b1 >> 4;
    b1 &= 0x0F; b2 &= 0x0F;
    b1 += '0'; b2 += '0';
    if (b1 > '9') b1 += 39;
    if (b2 > '9') b2 += 39;
    result.concat(b2);
    result.concat(b1);
  }
  return result;
}

/**
 * Invalidate cached bootloader SHA256 (call after bootloader update)
 * Forces recalculation on next call to calculateBootloaderSHA256 or getBootloaderSHA256Hex
 */
static void invalidateBootloaderSHA256Cache() {
  bootloaderSHA256CacheValid = false;
}

/**
 * Verify complete buffered bootloader using ESP-IDF validation approach
 * This matches the key validation steps from esp_image_verify() in ESP-IDF
 * @param buffer Reference to pointer to bootloader binary data (will be adjusted if offset detected)
 * @param len Reference to length of bootloader data (will be adjusted to actual size)
 * @param bootloaderErrorMsg Pointer to String to store error message (must not be null)
 * @return true if validation passed, false otherwise
 */
static bool verifyBootloaderImage(const uint8_t* &buffer, size_t &len, String& bootloaderErrorMsg) {
  const size_t MIN_IMAGE_HEADER_SIZE = sizeof(esp_image_header_t);

  // 1. Validate minimum size for header
  if (len < MIN_IMAGE_HEADER_SIZE) {
    bootloaderErrorMsg = "Too small";
    return false;
  }

  // Check if the bootloader starts at offset 0x1000 (common in partition table dumps)
  // This happens when someone uploads a complete flash dump instead of just the bootloader
  if (len > BOOTLOADER_OFFSET + MIN_IMAGE_HEADER_SIZE &&
      buffer[BOOTLOADER_OFFSET] == ESP_IMAGE_HEADER_MAGIC &&
      buffer[0] != ESP_IMAGE_HEADER_MAGIC) {
    DEBUG_PRINTF_P(PSTR("Bootloader detected at offset\n"));
    // Adjust buffer pointer to start at the actual bootloader
    buffer = buffer + BOOTLOADER_OFFSET;
    len = len - BOOTLOADER_OFFSET;

    // Re-validate size after adjustment
    if (len < MIN_IMAGE_HEADER_SIZE) {
      bootloaderErrorMsg = "Too small";
      return false;
    }
  }

  size_t availableLen = len;
  esp_image_header_t imageHeader{};
  memcpy(&imageHeader, buffer, sizeof(imageHeader));

  // 2. Basic header sanity checks (matches early esp_image_verify checks)
  if (imageHeader.magic != ESP_IMAGE_HEADER_MAGIC ||
      imageHeader.segment_count == 0 || imageHeader.segment_count > 16 ||
      imageHeader.spi_mode > 3 ||
      imageHeader.entry_addr < 0x40000000 || imageHeader.entry_addr > 0x50000000) {
    bootloaderErrorMsg = "Invalid header";
    return false;
  }

  // 3. Chip ID validation (matches esp_image_verify step 3)
  if (imageHeader.chip_id != CONFIG_IDF_FIRMWARE_CHIP_ID) {
    bootloaderErrorMsg = "Chip ID mismatch";
    return false;
  }

  // 4. Validate image size
  BootloaderImageSizer sizer;
  sizer.feed(buffer, availableLen);
  if (!sizer.isSizeKnown()) {
    bootloaderErrorMsg = "Invalid image";
    return false;
  }
  size_t actualBootloaderSize = sizer.totalSize();

  // 5. SHA256 checksum (optional)
  if (imageHeader.hash_appended == 1) {
    actualBootloaderSize += 32;
  }

  if (actualBootloaderSize > len) {
    // Same as above
    bootloaderErrorMsg = "Too small";
    return false;
  }

  DEBUG_PRINTF_P(PSTR("Bootloader validation: %d segments, actual size %d bytes (buffer size %d bytes, hash_appended=%d)\n"),
                 imageHeader.segment_count, actualBootloaderSize, len, imageHeader.hash_appended);

  // Update len to reflect actual bootloader size (including hash and checksum, with alignment)
  // This is critical - we must write the complete image including checksums
  len = actualBootloaderSize;

  return true;
}

// Bootloader OTA context structure
struct BootloaderUpdateContext {
  // State flags
  bool replySent = false;
  bool uploadComplete = false;
  String errorMessage;

  // Buffer to hold bootloader data
  uint8_t* buffer = nullptr;
  size_t bytesBuffered = 0;
  const uint32_t bootloaderOffset = 0x1000;
  const uint32_t maxBootloaderSize = 0x10000; // 64KB buffer size
};

// Cleanup bootloader OTA context
static void endBootloaderOTA(AsyncWebServerRequest *request) {
  BootloaderUpdateContext* context = reinterpret_cast<BootloaderUpdateContext*>(request->_tempObject);
  request->_tempObject = nullptr;

  DEBUG_PRINTF_P(PSTR("EndBootloaderOTA %x --> %x\n"), (uintptr_t)request, (uintptr_t)context);
  if (context) {
    if (context->buffer) {
      free(context->buffer);
      context->buffer = nullptr;
    }

    // If update failed, restore system state
    if (!context->uploadComplete || !context->errorMessage.isEmpty()) {
      strip.resume();
      #if WLED_WATCHDOG_TIMEOUT > 0
      WLED::instance().enableWatchdog();
      #endif
    }

    delete context;
  }
}

// Initialize bootloader OTA context
bool initBootloaderOTA(AsyncWebServerRequest *request) {
  if (request->_tempObject) {
    return true; // Already initialized
  }

  BootloaderUpdateContext* context = new BootloaderUpdateContext();
  if (!context) {
    DEBUG_PRINTLN(F("Failed to allocate bootloader OTA context"));
    return false;
  }
  request->_tempObject = context;
  request->onDisconnect([=]() { endBootloaderOTA(request); });  // ensures cleanup on disconnect

#ifdef BOOTLOADER_OTA_UNSUPPORTED
  context->errorMessage = F("Bootloader update not supported on this chip");
  return false;
#else
  DEBUG_PRINTLN(F("Bootloader Update Start - initializing buffer"));
  #if WLED_WATCHDOG_TIMEOUT > 0
  WLED::instance().disableWatchdog();
  #endif
  lastEditTime = millis(); // make sure PIN does not lock during update
  strip.suspend();
  strip.resetSegments();

  // Check available heap before attempting allocation
  DEBUG_PRINTF_P(PSTR("Free heap before bootloader buffer allocation: %d bytes (need %d bytes)\n"), getContiguousFreeHeap(), context->maxBootloaderSize);

  context->buffer = (uint8_t*)malloc(context->maxBootloaderSize);
  if (!context->buffer) {
    size_t freeHeapNow = getContiguousFreeHeap();
    DEBUG_PRINTF_P(PSTR("Failed to allocate %d byte bootloader buffer! Contiguous heap: %d bytes\n"), context->maxBootloaderSize, freeHeapNow);
    context->errorMessage = "Out of memory! Contiguous heap: " + String(freeHeapNow) + " bytes, need: " + String(context->maxBootloaderSize) + " bytes";
    strip.resume();
    #if WLED_WATCHDOG_TIMEOUT > 0
    WLED::instance().enableWatchdog();
    #endif
    return false;
  }

  context->bytesBuffered = 0;
  return true;
#endif
}

// Set bootloader OTA replied flag
void setBootloaderOTAReplied(AsyncWebServerRequest *request) {
  BootloaderUpdateContext* context = reinterpret_cast<BootloaderUpdateContext*>(request->_tempObject);
  if (context) {
    context->replySent = true;
  }
}

// Get bootloader OTA result
std::pair<bool, String> getBootloaderOTAResult(AsyncWebServerRequest *request) {
  BootloaderUpdateContext* context = reinterpret_cast<BootloaderUpdateContext*>(request->_tempObject);

  if (!context) {
    return std::make_pair(true, String(F("Internal error: No bootloader OTA context")));
  }

  bool needsReply = !context->replySent;
  String errorMsg = context->errorMessage;

  // If upload was successful, return empty string and trigger reboot
  if (context->uploadComplete && errorMsg.isEmpty()) {
    doReboot = true;
    endBootloaderOTA(request);
    return std::make_pair(needsReply, String());
  }

  // If there was an error, return it
  if (!errorMsg.isEmpty()) {
    endBootloaderOTA(request);
    return std::make_pair(needsReply, errorMsg);
  }

  // Should never happen
  return std::make_pair(true, String(F("Internal software failure")));
}

// Handle bootloader OTA data
void handleBootloaderOTAData(AsyncWebServerRequest *request, size_t index, uint8_t *data, size_t len, bool isFinal) {
  BootloaderUpdateContext* context = reinterpret_cast<BootloaderUpdateContext*>(request->_tempObject);

  if (!context) {
    DEBUG_PRINTLN(F("No bootloader OTA context - ignoring data"));
    return;
  }

  if (!context->errorMessage.isEmpty()) {
    return;
  }

  // Buffer the incoming data
  if (context->buffer && context->bytesBuffered + len <= context->maxBootloaderSize) {
    memcpy(context->buffer + context->bytesBuffered, data, len);
    context->bytesBuffered += len;
    DEBUG_PRINTF_P(PSTR("Bootloader buffer progress: %d / %d bytes\n"), context->bytesBuffered, context->maxBootloaderSize);
  } else if (!context->buffer) {
    DEBUG_PRINTLN(F("Bootloader buffer not allocated!"));
    context->errorMessage = "Internal error: Bootloader buffer not allocated";
    return;
  } else {
    size_t totalSize = context->bytesBuffered + len;
    DEBUG_PRINTLN(F("Bootloader size exceeds maximum!"));
    context->errorMessage = "Bootloader file too large: " + String(totalSize) + " bytes (max: " + String(context->maxBootloaderSize) + " bytes)";
    return;
  }

  // Only write to flash when upload is complete
  if (isFinal) {
    DEBUG_PRINTLN(F("Bootloader Upload Complete - validating and flashing"));

    if (context->buffer && context->bytesBuffered > 0) {
      // Prepare pointers for verification (may be adjusted if bootloader at offset)
      const uint8_t* bootloaderData = context->buffer;
      size_t bootloaderSize = context->bytesBuffered;

      // Verify the complete bootloader image before flashing
      // Note: verifyBootloaderImage may adjust bootloaderData pointer and bootloaderSize
      // for validation purposes only
      if (!verifyBootloaderImage(bootloaderData, bootloaderSize, context->errorMessage)) {
        DEBUG_PRINTLN(F("Bootloader validation failed!"));
        // Error message already set by verifyBootloaderImage
      } else {
        // Calculate offset to write to flash
        // If bootloaderData was adjusted (partition table detected), we need to skip it in flash too
        size_t flashOffset = context->bootloaderOffset;
        const uint8_t* dataToWrite = context->buffer;
        size_t bytesToWrite = context->bytesBuffered;

        // If validation adjusted the pointer, it means we have a partition table at the start
        // In this case, we should skip writing the partition table and write bootloader at 0x1000
        if (bootloaderData != context->buffer) {
          // bootloaderData was adjusted - skip partition table in our data
          size_t partitionTableSize = bootloaderData - context->buffer;
          dataToWrite = bootloaderData;
          bytesToWrite = bootloaderSize;
          DEBUG_PRINTF_P(PSTR("Skipping %d bytes of partition table data\n"), partitionTableSize);
        }

        DEBUG_PRINTF_P(PSTR("Bootloader validation passed - writing %d bytes to flash at 0x%04X\n"),
                       bytesToWrite, flashOffset);

        // Calculate erase size (must be multiple of 4KB)
        size_t eraseSize = ((bytesToWrite + 0xFFF) / 0x1000) * 0x1000;
        if (eraseSize > context->maxBootloaderSize) {
          eraseSize = context->maxBootloaderSize;
        }

        // Erase bootloader region
        DEBUG_PRINTF_P(PSTR("Erasing %d bytes at 0x%04X...\n"), eraseSize, flashOffset);
        esp_err_t err = esp_flash_erase_region(NULL, flashOffset, eraseSize);
        if (err != ESP_OK) {
          DEBUG_PRINTF_P(PSTR("Bootloader erase error: %d\n"), err);
          context->errorMessage = "Flash erase failed (error code: " + String(err) + ")";
        } else {
          // Write the validated bootloader data to flash
          err = esp_flash_write(NULL, dataToWrite, flashOffset, bytesToWrite);
          if (err != ESP_OK) {
            DEBUG_PRINTF_P(PSTR("Bootloader flash write error: %d\n"), err);
            context->errorMessage = "Flash write failed (error code: " + String(err) + ")";
          } else {
            DEBUG_PRINTF_P(PSTR("Bootloader Update Success - %d bytes written to 0x%04X\n"),
                           bytesToWrite, flashOffset);
            // Invalidate cached bootloader hash
            invalidateBootloaderSHA256Cache();
            context->uploadComplete = true;
          }
        }
      }
    } else if (context->bytesBuffered == 0) {
      context->errorMessage = "No bootloader data received";
    }
  }
}
#endif