Implement metadata-based OTA release checking system

Copilot · netmindz · Copilot · commit 42ff73ffe734 · 2025-09-14T13:11:24.000Z
Co-authored-by: netmindz &lt;442066+netmindz@users.noreply.github.com&gt;
diff --git 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
@@ -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;
-    }
+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;
   }
-  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
-    }
+  // 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;
   }
-  
-  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;
-}
+  DEBUG_PRINTLN(F("Found WLED metadata header"));
 
-bool extractReleaseNameFromBinary(const uint8_t* binaryData, size_t dataSize, char* extractedRelease) {
-  if (!binaryData || !extractedRelease || dataSize == 0) {
-    return false;
-  }
+  // 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;
   
-  // 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;
-    }
-  }
-  
-  // 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;
 }
diff --git 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 <Arduino.h>
 
+#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
diff --git a/wled00/wled_server.cpp b/wled00/wled_server.cpp
