1+ #include " ota_update.h"
2+ #include " wled.h"
3+
4+ #ifdef ESP32
5+ #include < esp_app_format.h>
6+ #include < esp_ota_ops.h>
7+ #endif
8+
9+ // Platform-specific metadata locations
10+ #ifdef ESP32
11+ constexpr size_t METADATA_OFFSET = 256 ; // ESP32: metadata appears after Espressif metadata
12+ #define UPDATE_ERROR errorString
13+ #elif defined(ESP8266)
14+ constexpr size_t METADATA_OFFSET = 0x1000 ; // ESP8266: metadata appears at 4KB offset
15+ #define UPDATE_ERROR getErrorString
16+ #endif
17+ constexpr size_t METADATA_SEARCH_RANGE = 512 ; // bytes
18+
19+
20+ /* *
21+ * Check if OTA should be allowed based on release compatibility using custom description
22+ * @param binaryData Pointer to binary file data (not modified)
23+ * @param dataSize Size of binary data in bytes
24+ * @param errorMessage Buffer to store error message if validation fails
25+ * @param errorMessageLen Maximum length of error message buffer
26+ * @return true if OTA should proceed, false if it should be blocked
27+ */
28+
29+ static bool validateOTA (const uint8_t * binaryData, size_t dataSize, char * errorMessage, size_t errorMessageLen) {
30+ // Clear error message
31+ if (errorMessage && errorMessageLen > 0 ) {
32+ errorMessage[0 ] = ' \0 '
33+ }
34+
35+ // Try to extract WLED structure directly from binary data
36+ wled_custom_desc_t extractedDesc;
37+ bool hasDesc = findWledMetadata (binaryData, dataSize, &extractedDesc);
38+
39+ if (hasDesc) {
40+ return shouldAllowOTA (extractedDesc, errorMessage, errorMessageLen);
41+ } else {
42+ // No custom description - this could be a legacy binary
43+ if (errorMessage && errorMessageLen > 0 ) {
44+ strncpy_P (errorMessage, PSTR (" This firmware file is missing compatibility metadata." 1 );
45+ errorMessage[errorMessageLen - 1 ] = ' \0 '
46+ }
47+ return false ;
48+ }
49+ }
50+
51+ struct UpdateContext {
52+ // State flags
53+ // FUTURE: the flags could be replaced by a state machine
54+ bool replySent = false ;
55+ bool needsRestart = false ;
56+ bool updateStarted = false ;
57+ bool uploadComplete = false ;
58+ bool releaseCheckPassed = false ;
59+ String errorMessage;
60+
61+ // Buffer to hold block data across posts, if needed
62+ std::vector<uint8_t > releaseMetadataBuffer;
63+ };
64+
65+
66+ static void endOTA (AsyncWebServerRequest *request) {
67+ UpdateContext* context = reinterpret_cast <UpdateContext*>(request->_tempObject );
68+ request->_tempObject = nullptr ;
69+
70+ DEBUG_PRINTF_P (PSTR (" EndOTA %x --> %x (%d)\n " uintptr_t )request,(uintptr_t ) context, context ? context->uploadComplete : 0 );
71+ if (context) {
72+ if (context->updateStarted ) { // We initialized the update
73+ // We use Update.end() because not all forms of Update() support an abort.
74+ // If the upload is incomplete, Update.end(false) should error out.
75+ if (Update.end (context->uploadComplete )) {
76+ // Update successful!
77+ #ifndef ESP8266
78+ bootloopCheckOTA (); // let the bootloop-checker know there was an OTA update
79+ #endif
80+ doReboot = true ;
81+ context->needsRestart = false ;
82+ }
83+ }
84+
85+ if (context->needsRestart ) {
86+ strip.resume ();
87+ UsermodManager::onUpdateBegin (false );
88+ #if WLED_WATCHDOG_TIMEOUT > 0
89+ WLED::instance ().enableWatchdog ();
90+ #endif
91+ }
92+ delete context;
93+ }
94+ };
95+
96+ static bool beginOTA (AsyncWebServerRequest *request, UpdateContext* context)
97+ {
98+ #ifdef ESP8266
99+ Update.runAsync (true );
100+ #endif
101+
102+ if (Update.isRunning ()) {
103+ request->send (503 );
104+ setOTAReplied (request);
105+ return false ;
106+ }
107+
108+ #if WLED_WATCHDOG_TIMEOUT > 0
109+ WLED::instance ().disableWatchdog ();
110+ #endif
111+ UsermodManager::onUpdateBegin (true ); // notify usermods that update is about to begin (some may require task de-init)
112+
113+ strip.suspend ();
114+ backupConfig (); // backup current config in case the update ends badly
115+ strip.resetSegments (); // free as much memory as you can
116+ context->needsRestart = true ;
117+
118+ DEBUG_PRINTF_P (PSTR (" OTA Update Start, %x --> %x\n " uintptr_t )request,(uintptr_t ) context);
119+
120+ auto skipValidationParam = request->getParam (" skipValidation" true );
121+ if (skipValidationParam && (skipValidationParam->value () == " 1"
122+ context->releaseCheckPassed = true ;
123+ DEBUG_PRINTLN (F (" OTA validation skipped by user"
124+ }
125+
126+ // Begin update with the firmware size from content length
127+ size_t updateSize = request->contentLength () > 0 ? request->contentLength () : ((ESP.getFreeSketchSpace () - 0x1000 ) & 0xFFFFF000 );
128+ if (!Update.begin (updateSize)) {
129+ context->errorMessage = Update.UPDATE_ERROR ();
130+ DEBUG_PRINTF_P (PSTR (" OTA Failed to begin: %s\n " errorMessage .c_str ());
131+ return false ;
132+ }
133+
134+ context->updateStarted = true ;
135+ return true ;
136+ }
137+
138+ // Create an OTA context object on an AsyncWebServerRequest
139+ // Returns true if successful, false on failure.
140+ bool initOTA (AsyncWebServerRequest *request) {
141+ // Allocate update context
142+ UpdateContext* context = new (std::nothrow) UpdateContext {};
143+ if (context) {
144+ request->_tempObject = context;
145+ request->onDisconnect ([=]() { endOTA (request); }); // ensures we restart on failure
146+ };
147+
148+ DEBUG_PRINTF_P (PSTR (" OTA Update init, %x --> %x\n " uintptr_t )request,(uintptr_t ) context);
149+ return (context != nullptr );
150+ }
151+
152+ void setOTAReplied (AsyncWebServerRequest *request) {
153+ UpdateContext* context = reinterpret_cast <UpdateContext*>(request->_tempObject );
154+ if (!context) return ;
155+ context->replySent = true ;
156+ };
157+
158+ // Returns pointer to error message, or nullptr if OTA was successful.
159+ std::pair<bool , String> getOTAResult (AsyncWebServerRequest* request) {
160+ UpdateContext* context = reinterpret_cast <UpdateContext*>(request->_tempObject );
161+ if (!context) return { true , F (" OTA context unexpectedly missing"
162+ if (context->replySent ) return { false , {} };
163+ if (context->errorMessage .length ()) return { true , context->errorMessage };
164+
165+ if (context->updateStarted ) {
166+ // Release the OTA context now.
167+ endOTA (request);
168+ if (Update.hasError ()) {
169+ return { true , Update.UPDATE_ERROR () };
170+ } else {
171+ return { true , {} };
172+ }
173+ }
174+
175+ // Should never happen
176+ return { true , F (" Internal software failure"
177+ }
178+
179+
180+
181+ void handleOTAData (AsyncWebServerRequest *request, size_t index, uint8_t *data, size_t len, bool isFinal)
182+ {
183+ UpdateContext* context = reinterpret_cast <UpdateContext*>(request->_tempObject );
184+ if (!context) return ;
185+
186+ // DEBUG_PRINTF_P(PSTR("HandleOTAData: %d %d %d\n"), index, len, isFinal);
187+
188+ if (context->replySent || (context->errorMessage .length ())) return ;
189+
190+ if (index == 0 ) {
191+ if (!beginOTA (request, context)) return ;
192+ }
193+
194+ // Perform validation if we haven't done it yet and we have reached the metadata offset
195+ if (!context->releaseCheckPassed && (index+len) > METADATA_OFFSET) {
196+ // Current chunk contains the metadata offset
197+ size_t availableDataAfterOffset = (index + len) - METADATA_OFFSET;
198+
199+ DEBUG_PRINTF_P (PSTR (" OTA metadata check: %d in buffer, %d received, %d available\n " releaseMetadataBuffer .size (), len, availableDataAfterOffset);
200+
201+ if (availableDataAfterOffset >= METADATA_SEARCH_RANGE) {
202+ // We have enough data to validate, one way or another
203+ const uint8_t * search_data = data;
204+ size_t search_len = len;
205+
206+ // If we have saved data, use that instead
207+ if (context->releaseMetadataBuffer .size ()) {
208+ // Add this data
209+ context->releaseMetadataBuffer .insert (context->releaseMetadataBuffer .end (), data, data+len);
210+ search_data = context->releaseMetadataBuffer .data ();
211+ search_len = context->releaseMetadataBuffer .size ();
212+ }
213+
214+ // Do the checking
215+ char errorMessage[128 ];
216+ bool OTA_ok = validateOTA (search_data, search_len, errorMessage, sizeof (errorMessage));
217+
218+ // Release buffer if there was one
219+ context->releaseMetadataBuffer = decltype (context->releaseMetadataBuffer ){};
220+
221+ if (!OTA_ok) {
222+ DEBUG_PRINTF_P (PSTR (" OTA declined: %s\n "
223+ context->errorMessage = errorMessage;
224+ context->errorMessage += F (" Enable 'Ignore firmware validation' to proceed anyway."
225+ return ;
226+ } else {
227+ DEBUG_PRINTLN (F (" OTA allowed: Release compatibility check passed"
228+ context->releaseCheckPassed = true ;
229+ }
230+ } else {
231+ // Store the data we just got for next pass
232+ context->releaseMetadataBuffer .insert (context->releaseMetadataBuffer .end (), data, data+len);
233+ }
234+ }
235+
236+ // Check if validation was still pending (shouldn't happen normally)
237+ // This is done before writing the last chunk, so endOTA can abort
238+ if (isFinal && !context->releaseCheckPassed ) {
239+ DEBUG_PRINTLN (F (" OTA failed: Validation never completed"
240+ // Don't write the last chunk to the updater: this will trip an error later
241+ context->errorMessage = F (" Release check data never arrived?"
242+ return ;
243+ }
244+
245+ // Write chunk data to OTA update (only if release check passed or still pending)
246+ if (!Update.hasError ()) {
247+ if (Update.write (data, len) != len) {
248+ DEBUG_PRINTF_P (PSTR (" OTA write failed on chunk %zu: %s\n " UPDATE_ERROR ());
249+ }
250+ }
251+
252+ if (isFinal) {
253+ DEBUG_PRINTLN (F (" OTA Update End"
254+ // Upload complete
255+ context->uploadComplete = true ;
256+ }
257+ }
0 commit comments