diff --git a/pio-scripts/output_bins.py b/pio-scripts/output_bins.py index 4d1594d84..9250a3118 100644 --- a/pio-scripts/output_bins.py +++ b/pio-scripts/output_bins.py @@ -18,6 +18,37 @@ def _create_dirs(dirs=["map", "release", "firmware"]): for d in dirs: os.makedirs(os.path.join(OUTPUT_DIR, d), exist_ok=True) +def add_metadata_header(binary_file, release_name): + """Add WLED release metadata header to binary file""" + # Metadata format: "WLED_META:" + release_name + null terminator + padding to 64 bytes + original binary + header_prefix = b"WLED_META:" + release_bytes = release_name.encode('utf-8') + + # Ensure total header is exactly 64 bytes for alignment + header_size = 64 + header_data_size = len(header_prefix) + len(release_bytes) + 1 # +1 for null terminator + + if header_data_size > header_size: + print(f"Warning: Release name '{release_name}' too long, truncating") + max_release_len = header_size - len(header_prefix) - 1 + release_bytes = release_bytes[:max_release_len] + header_data_size = len(header_prefix) + len(release_bytes) + 1 + + # Create header with padding + header = header_prefix + release_bytes + b'\0' + header += b'\xFF' * (header_size - header_data_size) # Pad with 0xFF (erased flash pattern) + + # Read original binary + with open(binary_file, 'rb') as f: + original_data = f.read() + + # Write header + original binary + with open(binary_file, 'wb') as f: + f.write(header) + f.write(original_data) + + print(f"Added WLED metadata header with release name '{release_name}' to {binary_file}") + def create_release(source): release_name_def = _get_cpp_define_value(env, "WLED_RELEASE_NAME") if release_name_def: @@ -27,6 +58,10 @@ def create_release(source): release_gz_file = release_file + ".gz" print(f"Copying {source} to {release_file}") shutil.copy(source, release_file) + + # Add metadata header to the release binary + add_metadata_header(release_file, release_name) + bin_gzip(release_file, release_gz_file) else: variant = env["PIOENV"] diff --git a/wled00/ota_release_check.cpp b/wled00/ota_release_check.cpp index cae6c6e6d..756a3e831 100644 --- a/wled00/ota_release_check.cpp +++ b/wled00/ota_release_check.cpp @@ -1,136 +1,39 @@ #include "ota_release_check.h" #include "wled.h" -// Maximum size to scan in binary (we don't need to scan the entire file) -#define MAX_SCAN_SIZE 32768 - -/** - * Find a string in binary data - */ -static int findStringInBinary(const uint8_t* data, size_t dataSize, const char* pattern) { - size_t patternLen = strlen(pattern); - if (patternLen == 0 || patternLen > dataSize) return -1; - - for (size_t i = 0; i <= dataSize - patternLen; i++) { - if (memcmp(data + i, pattern, patternLen) == 0) { - return i; - } - } - return -1; -} - -/** - * Check if a string looks like a valid WLED release name - */ -static bool isValidReleaseNameFormat(const char* name, size_t len) { - if (len < 3 || len > 63) return false; - - // Should contain at least one letter and be mostly alphanumeric with underscores/dashes - bool hasLetter = false; - for (size_t i = 0; i < len; i++) { - char c = name[i]; - if (isalpha(c)) { - hasLetter = true; - } else if (!isdigit(c) && c != '_' && c != '-') { - return false; // Invalid character - } - } - - return hasLetter; // Must have at least one letter -} - -/** - * Extract release name by searching for any reasonable string that could be a release name - * This method is very permissive to handle custom builds and unknown release formats - */ -static bool extractByGenericStringSearch(const uint8_t* data, size_t dataSize, char* extractedRelease) { - // Search for null-terminated strings that could be release names - char bestCandidate[64] = ""; - int bestScore = -1; - - for (size_t i = 0; i < dataSize - 4; i++) { - // Look for potential start of a string (printable character) - if (isalpha(data[i])) { - // Find the end of this string (null terminator) - size_t j = i; - while (j < dataSize && data[j] != 0) { - j++; - } - - if (j < dataSize) { // Found null terminator - size_t len = j - i; - if (len >= 3 && len <= 63) { // reasonable length for a release name - char candidate[64]; - strncpy(candidate, (const char*)(data + i), len); - candidate[len] = '\0'; - - // Check if this looks like a valid release name format - if (isValidReleaseNameFormat(candidate, len)) { - // Score candidates to find the most likely release name - int score = 0; - - // High score for common patterns - if (strstr(candidate, "ESP") != NULL) score += 100; - if (strstr(candidate, "WLED") != NULL) score += 100; - if (strstr(candidate, "Custom") != NULL) score += 50; - if (strstr(candidate, "Build") != NULL) score += 30; - - // Medium score for reasonable structure - if (len >= 5 && len <= 32) score += 20; // reasonable length - if (strchr(candidate, '_') != NULL) score += 10; // contains underscore (common in release names) - if (strchr(candidate, '-') != NULL) score += 10; // contains dash (common in release names) - - // Basic score for any valid format - if (score == 0) score = 5; // Any valid format gets minimum score - - if (score > bestScore) { - bestScore = score; - strcpy(bestCandidate, candidate); - } - } - } - } - } - } - - if (bestScore > 0) { - strcpy(extractedRelease, bestCandidate); - DEBUG_PRINTF_P(PSTR("Found release name by generic search (score %d): %s\n"), bestScore, extractedRelease); - return true; - } - - return false; -} - -bool extractReleaseNameFromBinary(const uint8_t* binaryData, size_t dataSize, char* extractedRelease) { - if (!binaryData || !extractedRelease || dataSize == 0) { +bool extractMetadataHeader(uint8_t* binaryData, size_t dataSize, char* extractedRelease, size_t* actualBinarySize) { + if (!binaryData || !extractedRelease || !actualBinarySize || dataSize < WLED_META_HEADER_SIZE) { + *actualBinarySize = dataSize; return false; } - - // Limit scan size to avoid performance issues with large binaries - size_t scanSize = (dataSize > MAX_SCAN_SIZE) ? MAX_SCAN_SIZE : dataSize; - - // First, try to find the exact current release string in the binary - // This is the most reliable method since we know what we're looking for - int pos = findStringInBinary(binaryData, scanSize, releaseString); - if (pos >= 0) { - // Verify it's properly null-terminated - size_t releaseLen = strlen(releaseString); - if (pos + releaseLen < scanSize && binaryData[pos + releaseLen] == 0) { - strcpy(extractedRelease, releaseString); - DEBUG_PRINTF_P(PSTR("Found exact current release string in binary: %s\n"), extractedRelease); - return true; - } + + // Check if the binary starts with our metadata header + if (memcmp(binaryData, WLED_META_PREFIX, strlen(WLED_META_PREFIX)) != 0) { + // No metadata header found, this is a legacy binary + *actualBinarySize = dataSize; + DEBUG_PRINTLN(F("No WLED metadata header found - legacy binary")); + return false; } + + DEBUG_PRINTLN(F("Found WLED metadata header")); + + // Extract release name from header + const char* releaseStart = (const char*)(binaryData + strlen(WLED_META_PREFIX)); + size_t maxReleaseLen = WLED_META_HEADER_SIZE - strlen(WLED_META_PREFIX) - 1; - // Fallback: Search for any string that looks like a release name - // This handles the case where the binary has a different but valid release name - if (extractByGenericStringSearch(binaryData, scanSize, extractedRelease)) { - return true; - } - - DEBUG_PRINTLN(F("Could not extract release name from binary")); - return false; + // Copy release name (it should be null-terminated within the header) + strncpy(extractedRelease, releaseStart, maxReleaseLen); + extractedRelease[maxReleaseLen] = '\0'; // Ensure null termination + + // Remove metadata header by shifting binary data + size_t firmwareSize = dataSize - WLED_META_HEADER_SIZE; + memmove(binaryData, binaryData + WLED_META_HEADER_SIZE, firmwareSize); + *actualBinarySize = firmwareSize; + + DEBUG_PRINTF_P(PSTR("Extracted release name from metadata: '%s', firmware size: %zu bytes\n"), + extractedRelease, firmwareSize); + + return true; } bool validateReleaseCompatibility(const char* extractedRelease) { @@ -147,29 +50,41 @@ bool validateReleaseCompatibility(const char* extractedRelease) { return match; } -bool shouldAllowOTA(const uint8_t* binaryData, size_t dataSize, bool ignoreReleaseCheck, char* errorMessage) { +bool shouldAllowOTA(uint8_t* binaryData, size_t dataSize, bool ignoreReleaseCheck, char* errorMessage, size_t* actualBinarySize) { // Clear error message if (errorMessage) { errorMessage[0] = '\0'; } + // Initialize actual binary size to full size by default + if (actualBinarySize) { + *actualBinarySize = dataSize; + } + // If user chose to ignore release check, allow OTA if (ignoreReleaseCheck) { DEBUG_PRINTLN(F("OTA release check bypassed by user")); + // Still need to extract metadata header if present to get clean binary + char dummyRelease[64]; + extractMetadataHeader(binaryData, dataSize, dummyRelease, actualBinarySize); return true; } - - // Try to extract release name from binary + + // Try to extract metadata header char extractedRelease[64]; - if (!extractReleaseNameFromBinary(binaryData, dataSize, extractedRelease)) { + bool hasMetadata = extractMetadataHeader(binaryData, dataSize, extractedRelease, actualBinarySize); + + if (!hasMetadata) { + // No metadata header - this could be a legacy binary or a binary without our metadata + // We cannot determine compatibility for such binaries if (errorMessage) { - strcpy(errorMessage, "Could not determine release type of uploaded file. Check 'Ignore release name check' to proceed."); + strcpy(errorMessage, "Binary has no release compatibility metadata. Check 'Ignore release name check' to proceed."); } - DEBUG_PRINTLN(F("OTA blocked: Could not extract release name")); + DEBUG_PRINTLN(F("OTA blocked: No metadata header found")); return false; } - - // Validate compatibility + + // Validate compatibility using extracted release name if (!validateReleaseCompatibility(extractedRelease)) { if (errorMessage) { snprintf(errorMessage, 127, "Release mismatch: current='%s', uploaded='%s'. Check 'Ignore release name check' to proceed.", @@ -179,7 +94,7 @@ bool shouldAllowOTA(const uint8_t* binaryData, size_t dataSize, bool ignoreRelea releaseString, extractedRelease); return false; } - + DEBUG_PRINTLN(F("OTA allowed: Release names match")); return true; } \ No newline at end of file diff --git a/wled00/ota_release_check.h b/wled00/ota_release_check.h index 0411ee84e..6ef300558 100644 --- a/wled00/ota_release_check.h +++ b/wled00/ota_release_check.h @@ -2,20 +2,24 @@ #define WLED_OTA_RELEASE_CHECK_H /* - * OTA Release Compatibility Checking - * Functions to extract and validate release names from uploaded binary files + * OTA Release Compatibility Checking with Metadata Headers + * Functions to extract and validate release names from uploaded binary files with metadata headers */ #include +#define WLED_META_HEADER_SIZE 64 +#define WLED_META_PREFIX "WLED_META:" + /** - * Extract release name from ESP32/ESP8266 binary data - * @param binaryData Pointer to binary file data + * Extract and remove metadata header from binary data + * @param binaryData Pointer to binary file data (will be modified) * @param dataSize Size of binary data in bytes * @param extractedRelease Buffer to store extracted release name (should be at least 64 bytes) - * @return true if release name was extracted successfully, false otherwise + * @param actualBinarySize Pointer to store the size of actual firmware binary (without header) + * @return true if metadata header was found and extracted, false if no metadata header present */ -bool extractReleaseNameFromBinary(const uint8_t* binaryData, size_t dataSize, char* extractedRelease); +bool extractMetadataHeader(uint8_t* binaryData, size_t dataSize, char* extractedRelease, size_t* actualBinarySize); /** * Validate if extracted release name matches current release @@ -25,13 +29,14 @@ bool extractReleaseNameFromBinary(const uint8_t* binaryData, size_t dataSize, ch bool validateReleaseCompatibility(const char* extractedRelease); /** - * Check if OTA should be allowed based on release compatibility - * @param binaryData Pointer to binary file data + * Check if OTA should be allowed based on release compatibility using metadata headers + * @param binaryData Pointer to binary file data (will be modified if metadata header present) * @param dataSize Size of binary data in bytes * @param ignoreReleaseCheck If true, skip release validation * @param errorMessage Buffer to store error message if validation fails (should be at least 128 bytes) + * @param actualBinarySize Pointer to store the size of actual firmware binary (without header) * @return true if OTA should proceed, false if it should be blocked */ -bool shouldAllowOTA(const uint8_t* binaryData, size_t dataSize, bool ignoreReleaseCheck, char* errorMessage); +bool shouldAllowOTA(uint8_t* binaryData, size_t dataSize, bool ignoreReleaseCheck, char* errorMessage, size_t* actualBinarySize); #endif // WLED_OTA_RELEASE_CHECK_H \ No newline at end of file diff --git a/wled00/wled_server.cpp b/wled00/wled_server.cpp index 9c3f3e033..cf918b0f8 100644 --- a/wled00/wled_server.cpp +++ b/wled00/wled_server.cpp @@ -434,28 +434,26 @@ void initServer() } if (!correctPIN || otaLock) return; - // Static variables to track release check state across chunks - static bool releaseCheckDone = false; + // Static variables to track metadata processing state across chunks + static bool metadataProcessed = false; static bool releaseCheckPassed = false; - static uint8_t* releaseCheckBuffer = nullptr; - static size_t bufferPos = 0; - static const size_t RELEASE_CHECK_BUFFER_SIZE = 8192; // Check first 8KB for release name + static uint8_t* firstChunkBuffer = nullptr; + static size_t expectedFirmwareSize = 0; + static size_t totalBytesWritten = 0; if(!index){ DEBUG_PRINTLN(F("OTA Update Start")); - // Reset release check state - releaseCheckDone = false; + // Reset metadata processing state + metadataProcessed = false; releaseCheckPassed = false; - bufferPos = 0; + expectedFirmwareSize = 0; + totalBytesWritten = 0; - // Allocate buffer for release check - if (releaseCheckBuffer) free(releaseCheckBuffer); - releaseCheckBuffer = (uint8_t*)malloc(RELEASE_CHECK_BUFFER_SIZE); - if (!releaseCheckBuffer) { - DEBUG_PRINTLN(F("OTA Failed: Could not allocate buffer for release check")); - request->send(500, FPSTR(CONTENT_TYPE_PLAIN), F("OTA Failed: Memory allocation error")); - return; + // Clean up any existing buffer + if (firstChunkBuffer) { + free(firstChunkBuffer); + firstChunkBuffer = nullptr; } #if WLED_WATCHDOG_TIMEOUT > 0 @@ -465,61 +463,96 @@ void initServer() lastEditTime = millis(); // make sure PIN does not lock during update } - // Accumulate data for release check if we haven't checked yet - if (!releaseCheckDone && releaseCheckBuffer) { - size_t copyLen = min(len, RELEASE_CHECK_BUFFER_SIZE - bufferPos); - if (copyLen > 0) { - memcpy(releaseCheckBuffer + bufferPos, data, copyLen); - bufferPos += copyLen; + // Process first chunk to check metadata and start OTA + if (!metadataProcessed && index == 0) { + metadataProcessed = true; + + // Create a copy of the first chunk data that we can modify for metadata extraction + firstChunkBuffer = (uint8_t*)malloc(len); + if (!firstChunkBuffer) { + DEBUG_PRINTLN(F("OTA Failed: Could not allocate buffer for metadata processing")); + request->send(500, FPSTR(CONTENT_TYPE_PLAIN), F("OTA Failed: Memory allocation error")); + return; + } + memcpy(firstChunkBuffer, data, len); + + // Check if user wants to ignore release check + bool ignoreRelease = request->hasParam("ignoreRelease", true); + + char errorMessage[128]; + size_t actualFirmwareSize; + releaseCheckPassed = shouldAllowOTA(firstChunkBuffer, len, ignoreRelease, errorMessage, &actualFirmwareSize); + + if (!releaseCheckPassed) { + DEBUG_PRINTF_P(PSTR("OTA blocked: %s\n"), errorMessage); + free(firstChunkBuffer); + firstChunkBuffer = nullptr; + UsermodManager::onUpdateBegin(false); + #if WLED_WATCHDOG_TIMEOUT > 0 + WLED::instance().enableWatchdog(); + #endif + request->send(400, FPSTR(CONTENT_TYPE_PLAIN), errorMessage); + return; } - // Check if we have enough data or this is the final chunk - if (bufferPos >= RELEASE_CHECK_BUFFER_SIZE || isFinal) { - releaseCheckDone = true; - - // Check if user wants to ignore release check - bool ignoreRelease = request->hasParam("ignoreRelease", true); - - char errorMessage[128]; - releaseCheckPassed = shouldAllowOTA(releaseCheckBuffer, bufferPos, ignoreRelease, errorMessage); - - if (!releaseCheckPassed) { - DEBUG_PRINTF_P(PSTR("OTA blocked: %s\n"), errorMessage); - free(releaseCheckBuffer); - releaseCheckBuffer = nullptr; - strip.resume(); - UsermodManager::onUpdateBegin(false); - #if WLED_WATCHDOG_TIMEOUT > 0 - WLED::instance().enableWatchdog(); - #endif - request->send(400, FPSTR(CONTENT_TYPE_PLAIN), errorMessage); - return; - } - - DEBUG_PRINTLN(F("Release check passed, starting OTA update")); - - // Now start the actual OTA update - strip.suspend(); - backupConfig(); // backup current config in case the update ends badly - strip.resetSegments(); // free as much memory as you can - #ifdef ESP8266 - Update.runAsync(true); - #endif - Update.begin((ESP.getFreeSketchSpace() - 0x1000) & 0xFFFFF000); + DEBUG_PRINTLN(F("Release check passed, starting OTA update")); + + // Calculate expected firmware size (if metadata header was present, actualFirmwareSize will be adjusted) + // For final size calculation, we need to account for all chunks + expectedFirmwareSize = request->contentLength(); + if (actualFirmwareSize < len) { + // Metadata header was removed, adjust expected size + size_t headerSize = len - actualFirmwareSize; + expectedFirmwareSize -= headerSize; + DEBUG_PRINTF_P(PSTR("Metadata header removed, firmware size: %zu bytes\n"), expectedFirmwareSize); } + + // Start the actual OTA update + strip.suspend(); + backupConfig(); // backup current config in case the update ends badly + strip.resetSegments(); // free as much memory as you can + #ifdef ESP8266 + Update.runAsync(true); + #endif + + // Begin update with the correct firmware size + size_t updateSize = (expectedFirmwareSize > 0) ? expectedFirmwareSize : ((ESP.getFreeSketchSpace() - 0x1000) & 0xFFFFF000); + if (!Update.begin(updateSize)) { + DEBUG_PRINTF_P(PSTR("OTA Failed to begin: %s\n"), Update.getErrorString().c_str()); + free(firstChunkBuffer); + firstChunkBuffer = nullptr; + strip.resume(); + UsermodManager::onUpdateBegin(false); + #if WLED_WATCHDOG_TIMEOUT > 0 + WLED::instance().enableWatchdog(); + #endif + request->send(500, FPSTR(CONTENT_TYPE_PLAIN), String("Update.begin failed: ") + Update.getErrorString()); + return; + } + + // Write the processed first chunk (with metadata header removed if present) + if (actualFirmwareSize > 0 && !Update.hasError()) { + if (Update.write(firstChunkBuffer, actualFirmwareSize) != actualFirmwareSize) { + DEBUG_PRINTF_P(PSTR("OTA write failed on first chunk: %s\n"), Update.getErrorString().c_str()); + } else { + totalBytesWritten += actualFirmwareSize; + } + } + + free(firstChunkBuffer); + firstChunkBuffer = nullptr; } - - // Only write data if release check passed (or was bypassed) - if (releaseCheckDone && releaseCheckPassed && !Update.hasError()) { - Update.write(data, len); + // Write subsequent chunks directly to Update + else if (metadataProcessed && releaseCheckPassed && index > 0 && !Update.hasError()) { + if (Update.write(data, len) != len) { + DEBUG_PRINTF_P(PSTR("OTA write failed on chunk %zu: %s\n"), index, Update.getErrorString().c_str()); + } else { + totalBytesWritten += len; + } } if(isFinal){ - // Clean up release check buffer - if (releaseCheckBuffer) { - free(releaseCheckBuffer); - releaseCheckBuffer = nullptr; - } + DEBUG_PRINTF_P(PSTR("OTA Update End: %zu bytes written\n"), totalBytesWritten); if (releaseCheckPassed) { if(Update.end(true)){ @@ -533,6 +566,12 @@ void initServer() #endif } } + + // Clean up + if (firstChunkBuffer) { + free(firstChunkBuffer); + firstChunkBuffer = nullptr; + } } }); #else