diff --git a/esphome/components/esphome/ota/ota_esphome.cpp b/esphome/components/esphome/ota/ota_esphome.cpp index d226726ef0..7dfcecd67f 100644 --- a/esphome/components/esphome/ota/ota_esphome.cpp +++ b/esphome/components/esphome/ota/ota_esphome.cpp @@ -537,6 +537,7 @@ void ESPHomeOTAComponent::handle_data_() { error_code = this->write_flash_(buf, read); if (error_code != ota::OTA_RESPONSE_OK) goto error; // NOLINT(cppcoreguidelines-avoid-goto) + this->send_chunk_acks_(xfer); } } @@ -793,13 +794,6 @@ ssize_t ESPHomeOTAComponent::receive_data_(uint8_t *buf, DataTransfer &xfer) { const uint32_t now = millis(); xfer.last_data_ms = now; xfer.total += read; -#if USE_OTA_VERSION == 2 - while (xfer.acknowledged + OTA_BLOCK_SIZE <= xfer.total || - (xfer.total == xfer.ota_size && xfer.acknowledged < xfer.ota_size)) { - this->data_write_byte_(ota::OTA_RESPONSE_CHUNK_OK); - xfer.acknowledged += OTA_BLOCK_SIZE; - } -#endif if (now - xfer.last_progress > OTA_PROGRESS_INTERVAL_MS) { xfer.last_progress = now; float percentage = (xfer.total * 100.0f) / xfer.ota_size; @@ -813,19 +807,61 @@ ssize_t ESPHomeOTAComponent::receive_data_(uint8_t *buf, DataTransfer &xfer) { return read; } +void ESPHomeOTAComponent::send_chunk_acks_(DataTransfer &xfer) { +#if USE_OTA_VERSION == 2 + while (xfer.acknowledged + OTA_BLOCK_SIZE <= xfer.total || + (xfer.total == xfer.ota_size && xfer.acknowledged < xfer.ota_size)) { + this->data_write_byte_(ota::OTA_RESPONSE_CHUNK_OK); + xfer.acknowledged += OTA_BLOCK_SIZE; + } +#endif +} + #ifdef USE_OTA_DEFLATE // The window doubles as the output buffer: the decoder fills it, we flush it to // the backend, and its bytes remain available as the back-reference history for // the next windowful. +ota::OTAResponseTypes ESPHomeOTAComponent::inflate_flush_(InflateSession &session) { + const size_t produced = session.dest - session.window; + const size_t pending = produced - session.flushed; + if (pending == 0) + return ota::OTA_RESPONSE_OK; + if (pending > session.image_size - session.written) { + ESP_LOGW(TAG, "Inflate size mismatch"); + return ota::OTA_RESPONSE_ERROR_UNKNOWN; + } + ota::OTAResponseTypes result = this->write_flash_(session.window + session.flushed, pending); + if (result != ota::OTA_RESPONSE_OK) + return result; + session.flushed = produced; + session.written += pending; + // A compressible region yields many windows per socket read + App.feed_wdt(); + return ota::OTA_RESPONSE_OK; +} + ota::OTAResponseTypes ESPHomeOTAComponent::inflate_data_(uint8_t *in, size_t image_size, DataTransfer &xfer) { InflateSession &session = *this->inflate_; session.self = this; session.xfer = &xfer; session.in = in; + session.image_size = image_size; + session.written = 0; + session.error = ota::OTA_RESPONSE_OK; ota_inflate_init(&session, session.window, OTA_INFLATE_WINDOW_SIZE); - // Pulls the next compressed chunk when the decoder runs dry + // Pulls the next compressed chunk when the decoder runs dry. Everything + // received so far is decoded by then, so it is written and acked first, + // keeping a chunk ack meaning "in flash" as on the uncompressed path. session.source_read_cb = [](OtaInflateState *d) -> int { auto *s = static_cast(d); + s->error = s->self->inflate_flush_(*s); + if (s->error != ota::OTA_RESPONSE_OK) + return -1; + s->self->send_chunk_acks_(*s->xfer); + if (s->xfer->total >= s->xfer->ota_size) { + ESP_LOGW(TAG, "Inflate size mismatch"); + return -1; + } ssize_t read = s->self->receive_data_(s->in, *s->xfer); if (read <= 0) return -1; @@ -834,34 +870,32 @@ ota::OTAResponseTypes ESPHomeOTAComponent::inflate_data_(uint8_t *in, size_t ima return s->in[0]; }; - size_t written = 0; int res; do { + // The ring index wrapped to 0 exactly when the window filled, so the + // window and the output cursor stay in lockstep session.dest = session.window; session.dest_limit = session.window + OTA_INFLATE_WINDOW_SIZE; + session.flushed = 0; res = ota_inflate(&session); if (res < 0) { // eof means the read callback failed, which is already logged if (!session.eof) { ESP_LOGW(TAG, "Inflate err %d", res); } - return ota::OTA_RESPONSE_ERROR_UNKNOWN; + return session.error != ota::OTA_RESPONSE_OK ? session.error : ota::OTA_RESPONSE_ERROR_UNKNOWN; } - const size_t produced = session.dest - session.window; - // More output than announced: stop before the write and report it below - if (produced > image_size - written) - break; - ota::OTAResponseTypes write_result = this->write_flash_(session.window, produced); - if (write_result != ota::OTA_RESPONSE_OK) - return write_result; - written += produced; + ota::OTAResponseTypes flush_result = this->inflate_flush_(session); + if (flush_result != ota::OTA_RESPONSE_OK) + return flush_result; + this->send_chunk_acks_(xfer); } while (res != OTA_INFLATE_DONE); - if (written != image_size || xfer.total != xfer.ota_size) { + if (session.written != image_size || xfer.total != xfer.ota_size) { ESP_LOGW(TAG, "Inflate size mismatch"); return ota::OTA_RESPONSE_ERROR_UNKNOWN; } - ESP_LOGD(TAG, "Inflated %zu bytes from %zu", written, xfer.total); + ESP_LOGD(TAG, "Inflated %zu bytes from %zu", session.written, xfer.total); return ota::OTA_RESPONSE_OK; } #endif // USE_OTA_DEFLATE diff --git a/esphome/components/esphome/ota/ota_esphome.h b/esphome/components/esphome/ota/ota_esphome.h index 3a2b720e97..6e0c7d20a7 100644 --- a/esphome/components/esphome/ota/ota_esphome.h +++ b/esphome/components/esphome/ota/ota_esphome.h @@ -133,9 +133,10 @@ class ESPHomeOTAComponent final : public ota::OTAComponent { uint32_t last_progress{0}; }; // Receives up to OTA_BUFFER_SIZE bytes of upload data into buf, waiting up to - // the data timeout; updates xfer and sends chunk acks. Returns bytes read, -1 - // on failure (logged). + // the data timeout, and updates xfer. Returns bytes read, -1 on failure (logged). inline ssize_t receive_data_(uint8_t *buf, DataTransfer &xfer); + // Acks every received block once it is in flash, so an ack means written + inline void send_chunk_acks_(DataTransfer &xfer); // Reads a 4 byte MSB first size field into size inline bool read_size_(uint8_t *buf, size_t &size, const LogString *desc); // Writes to the backend and logs a failure @@ -204,10 +205,18 @@ class ESPHomeOTAComponent final : public ota::OTAComponent { ESPHomeOTAComponent *self; DataTransfer *xfer; uint8_t *in; // caller's buffer for the compressed input, valid during inflate_data_ + size_t image_size; + size_t written; // inflated bytes in flash + size_t flushed; // bytes of the current window already in flash + ota::OTAResponseTypes error; // first failure inside the read callback uint8_t window[OTA_INFLATE_WINDOW_SIZE]; }; +#ifndef CLANG_TIDY // static analysis sets every define at once static_assert(!ota::OTABackendPtr::element_type::supports_compression(), "USE_OTA_DEFLATE is for backends that cannot store a gzip image"); +#endif + // Writes the decoded bytes not yet in flash; dest stays put so the ring history is intact + ota::OTAResponseTypes inflate_flush_(InflateSession &session); ota::OTAResponseTypes inflate_data_(uint8_t *in, size_t image_size, DataTransfer &xfer); std::unique_ptr inflate_; #endif diff --git a/esphome/components/esphome/ota/ota_esphome_inflate.c b/esphome/components/esphome/ota/ota_esphome_inflate.c index 95603fa995..c3e1213989 100644 --- a/esphome/components/esphome/ota/ota_esphome_inflate.c +++ b/esphome/components/esphome/ota/ota_esphome_inflate.c @@ -41,6 +41,8 @@ #include "ota_esphome_inflate.h" +#include + #define TINF_OK OTA_INFLATE_OK #define TINF_DONE OTA_INFLATE_DONE #define TINF_DATA_ERROR OTA_INFLATE_DATA_ERROR @@ -353,7 +355,7 @@ static int tinf_inflate_block_data(TINF_DATA *d, TINF_TREE *lt, TINF_TREE *dt) { d->curlen = tinf_read_bits(d, LENGTH_BITS[sym], LENGTH_BASE[sym]); dist = tinf_decode_symbol(d, dt); - if (dist >= 30) { + if (dist < 0 || dist >= 30) { return TINF_DATA_ERROR; } @@ -425,8 +427,12 @@ static int tinf_inflate_uncompressed_block(TINF_DATA *d) { /* initialize decompression structure */ void ota_inflate_init(TINF_DATA *d, unsigned char *dict, unsigned int dict_len) { + d->source = NULL; + d->source_limit = NULL; + d->tag = 0; d->eof = 0; d->bitcount = 0; + d->lz_off = 0; d->bfinal = 0; d->btype = -1; d->dict_size = dict_len; diff --git a/esphome/components/ota/ota_backend_arduino_rp2.h b/esphome/components/ota/ota_backend_arduino_rp2.h index 47a23e2d0a..cae523e2b3 100644 --- a/esphome/components/ota/ota_backend_arduino_rp2.h +++ b/esphome/components/ota/ota_backend_arduino_rp2.h @@ -16,7 +16,9 @@ class ArduinoRP2OTABackend final { OTAResponseTypes end(); void abort(); // A gzip image is staged on LittleFS as is; the core's OTA stub inflates it - // into the app region at reboot, the same way the ESP8266 bootloader does + // into the app region at reboot, the same way the ESP8266 bootloader does. + // begin() then sees the gzip size, so only the staging space is checked up + // front; the inflated size is not known until the stub reads the trailer. static constexpr bool supports_compression() { return true; } private: diff --git a/tests/integration/test_host_ota.py b/tests/integration/test_host_ota.py index 75a779e6b2..da5363ef2e 100644 --- a/tests/integration/test_host_ota.py +++ b/tests/integration/test_host_ota.py @@ -186,7 +186,7 @@ async def test_host_ota_self_update( if "OTA staged at" in line: staged.set() # The host backend has no gzip support, so the upload negotiates deflate - if "Inflated " in line: + if " bytes from " in line: inflated.set() dev.on_log(line) diff --git a/tests/unit_tests/test_espota2.py b/tests/unit_tests/test_espota2.py index 956b4cc2ab..a9b9d30943 100644 --- a/tests/unit_tests/test_espota2.py +++ b/tests/unit_tests/test_espota2.py @@ -1567,7 +1567,8 @@ def test_perform_ota_gzip_wins_over_deflate(mock_socket: Mock) -> None: espota2.perform_ota(mock_socket, None, io.BytesIO(original_content), "test.bin") sent = [c[0][0] for c in mock_socket.sendall.call_args_list] - compressed = gzip.compress(original_content, compresslevel=9) - assert sent[3] == len(compressed).to_bytes(4, "big") - assert sent[4] == hashlib.md5(compressed).hexdigest().encode() - assert sent[5] == compressed + # gzip output embeds the time of compression, so compare what it holds + payload = sent[5] + assert gzip.decompress(payload) == original_content + assert sent[3] == len(payload).to_bytes(4, "big") + assert sent[4] == hashlib.md5(payload).hexdigest().encode()