ota_esphome.cpp

Go to the documentation of this file.

#include "ota_esphome.h"
#ifdef USE_OTA_ENCRYPTION_PROVISIONED
#include "esphome/components/api/api_server.h"
#endif
#ifdef USE_OTA
#ifdef USE_OTA_PASSWORD
#include "esphome/components/sha256/sha256.h"
#endif
#include "esphome/components/network/util.h"
#include "esphome/components/socket/socket.h"
#include "esphome/components/ota/ota_backend.h"
#include "esphome/components/ota/ota_backend_esp8266.h"
#include "esphome/components/ota/ota_backend_arduino_libretiny.h"
#include "esphome/components/ota/ota_backend_arduino_rp2.h"
#include "esphome/components/ota/ota_backend_esp_idf.h"
#include "esphome/core/application.h"
#include "esphome/core/hal.h"
#include "esphome/core/helpers.h"
#include "esphome/core/log.h"
#include "esphome/core/util.h"
#ifdef USE_LWIP_FAST_SELECT
#include "esphome/core/lwip_fast_select.h"
#endif
 
#include <algorithm>
#include <cerrno>
#include <cstdio>
#include <cstring>
#include <sys/time.h>
 
namespace esphome {
 
ESPHOME_LOG_TAG(TAG, "esphome.ota");
 
#ifdef USE_OTA_ENCRYPTION
const noise::NoiseContext &ESPHomeOTAComponent::noise_context_() const {
#ifdef USE_OTA_ENCRYPTION_PROVISIONED
  // The api server holds the live key; safe mode never constructs it, and then
  // noise_ctx_ holds the saved key setup() found, if any
  if (api::global_api_server != nullptr)
    return api::global_api_server->get_noise_ctx();
#endif
  return this->noise_ctx_;
}
#endif
static constexpr uint16_t OTA_BLOCK_SIZE = 8192;
static constexpr uint32_t OTA_SOCKET_TIMEOUT_HANDSHAKE = 60000;  // milliseconds for initial handshake
// Milliseconds for data transfer. Covers the lwIP retransmit run seen in
// practice for a lost chunk ack (1.5 + 3 + 6 + 12 + 24 + 48 s); the CLI waits
// longer (espota2.DATA_PHASE_TIMEOUT) so the device is free before it retries
static constexpr uint32_t OTA_SOCKET_TIMEOUT_DATA = 105000;
// A handshake this old gives way to a waiting connection: a client retrying over a lossy link otherwise queues
// behind its own stalled attempt until both time out
static constexpr uint32_t OTA_HANDSHAKE_PREEMPT = 10000;
static constexpr uint32_t OTA_PROGRESS_INTERVAL_MS = 1000;
static constexpr size_t OTA_SIZE_FIELD_BYTES = 4;  // sizes on the wire are 4 bytes MSB first
 
// Single-instance pointer — multi-port configs are rejected in final_validate.
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
static ESPHomeOTAComponent *global_esphome_ota_component = nullptr;
 
// Called from any context (LwIP TCP/IP task, RP2040 user-IRQ).
extern "C" void esphome_wake_ota_component_any_context() {
  if (global_esphome_ota_component != nullptr) {
    global_esphome_ota_component->enable_loop_soon_any_context();
  }
}
 
void ESPHomeOTAComponent::setup() {
#ifdef USE_OTA_ENCRYPTION_PROVISIONED
  // Safe mode never constructs the api server, so read the key it saved
  noise::psk_t psk;
  if (api::global_api_server == nullptr && api::load_saved_noise_psk(psk)) {
    this->saved_psk_ = RAMAllocator<noise::psk_t>().make_unique(psk);
    if (this->saved_psk_ != nullptr) {
      this->noise_ctx_.set_psk(this->saved_psk_->data());
    }
  }
#endif
  this->server_ = socket::socket_ip_loop_monitored(SOCK_STREAM, 0).release();  // monitored for incoming connections
  if (this->server_ == nullptr) {
    this->server_failed_(LOG_STR("creation"));
    return;
  }
  int enable = 1;
  int err = this->server_->setsockopt(SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(int));
  if (err != 0) {
    this->log_socket_error_(LOG_STR("reuseaddr"));
    // we can still continue
  }
  err = this->server_->setblocking(false);
  if (err != 0) {
    this->server_failed_(LOG_STR("nonblocking"));
    return;
  }
 
  struct sockaddr_storage server;
 
  socklen_t sl = socket::set_sockaddr_any((struct sockaddr *) &server, sizeof(server), this->port_);
  if (sl == 0) {
    this->server_failed_(LOG_STR("set sockaddr"));
    return;
  }
 
  err = this->server_->bind((struct sockaddr *) &server, sl);
  if (err != 0) {
    this->server_failed_(LOG_STR("bind"));
    return;
  }
 
  err = this->server_->listen(1);  // Only one client at a time
  if (err != 0) {
    this->server_failed_(LOG_STR("listen"));
    return;
  }
 
  // loop() self-disables on its first idle tick; no explicit disable_loop() needed here.
  global_esphome_ota_component = this;
#ifdef USE_LWIP_FAST_SELECT
  // Filter fast-select wakes to this listener only. If the sock lookup returns nullptr,
  // no wakes fire and loop() falls back to the self-disable safety net.
  esphome_fast_select_set_ota_listener_sock(esphome_lwip_get_sock(this->server_->get_fd()));
#endif
 
#ifdef USE_OTA_PARTITIONS
  ota::get_running_app_position(this->running_app_offset_, this->running_app_size_);
#endif
}
 
void ESPHomeOTAComponent::dump_config() {
  char addr_buf[network::USE_ADDRESS_BUFFER_SIZE];
  ESP_LOGCONFIG(TAG,
                "Over-The-Air updates:\n"
                "  Address: %s:%u\n"
                "  Version: %d"
#ifdef USE_OTA_ENCRYPTION
                "\n  Encryption: %s"
#endif
                ,
                network::get_use_address_to(addr_buf), this->port_, USE_OTA_VERSION
#ifdef USE_OTA_ENCRYPTION_REQUIRED
                ,
                LOG_STR_LITERAL("required")
#elif defined(USE_OTA_ENCRYPTION_PROVISIONED)
                // A runtime provisioned key may not exist yet
                ,
                this->noise_context_().has_psk() ? LOG_STR_LITERAL("offered, plaintext accepted")
                                                 : LOG_STR_LITERAL("offered once the api key is provisioned")
#elif defined(USE_OTA_ENCRYPTION)
                ,
                LOG_STR_LITERAL("offered, plaintext accepted")
#endif
  );
#ifdef USE_OTA_PASSWORD
  if (!this->password_.empty()) {
    ESP_LOGCONFIG(TAG, "  Password configured");
  }
#endif
#ifdef USE_OTA_PARTITIONS
  ESP_LOGCONFIG(TAG,
                "  Partition access allowed\n"
                "  Running app:\n"
                "    Partition address: 0x%" PRIX32 "\n"
                "    Used size: %zu bytes (0x%zX)",
                this->running_app_offset_, this->running_app_size_, this->running_app_size_);
 
#ifdef USE_ESP32
  ESP_LOGCONFIG(TAG,
                "  Partition table:\n"
                "    %-12s %-4s %-8s %-10s %-10s",
                "Name", "Type", "Subtype", "Address", "Size");
  esp_partition_iterator_t it = esp_partition_find(ESP_PARTITION_TYPE_ANY, ESP_PARTITION_SUBTYPE_ANY, nullptr);
  while (it != nullptr) {
    const esp_partition_t *partition = esp_partition_get(it);
    ESP_LOGCONFIG(TAG, "    %-12s 0x%-2X 0x%-6X 0x%-8" PRIX32 " 0x%-8" PRIX32, partition->label, partition->type,
                  partition->subtype, partition->address, partition->size);
    it = esp_partition_next(it);
  }
  esp_partition_iterator_release(it);
  esp_bootloader_desc_t bootloader_desc;
  esp_err_t err = esp_ota_get_bootloader_description(nullptr, &bootloader_desc);
  ESP_LOGCONFIG(TAG, "  Bootloader: ESP-IDF %s",
                (err == ESP_OK) ? bootloader_desc.idf_ver : LOG_STR_LITERAL("version unknown"));
#endif  // USE_ESP32
#endif  // USE_OTA_PARTITIONS
}
 
void ESPHomeOTAComponent::loop() {
  // Self-disable idle loop where a wake path re-enables on listener readiness
  // (fast-select, raw-TCP accept_fn_). Host BSD select doesn't, so stay enabled.
  if (this->client_ == nullptr && !this->server_->ready()) {
#ifndef USE_HOST
    this->disable_loop();
#endif
    return;
  }
  this->handle_handshake_();
}
 
static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_COMPRESSION = 0x01;
static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_SHA256_AUTH = 0x02;
static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL = 0x04;
static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_NOISE = 0x08;
static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_DEFLATE = 0x10;
#ifdef ESPHOME_OTA_UDP
static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_UDP = 0x20;
#endif
// Noise needs the extended protocol: the prologue binds the 2-byte feature ack
static constexpr uint8_t CLIENT_NOISE_FEATURES =
    CLIENT_FEATURE_SUPPORTS_NOISE | CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL;
static constexpr uint8_t SERVER_FEATURE_SUPPORTS_COMPRESSION = 0x01;
static constexpr uint8_t SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS = 0x02;
static constexpr uint8_t SERVER_FEATURE_SUPPORTS_NOISE = 0x04;
// Raw deflate, window <= OTA_INFLATE_WINDOW_SIZE. Binding once offered: the
// client must then send the image size frame and a deflate stream.
static constexpr uint8_t SERVER_FEATURE_SUPPORTS_DEFLATE = 0x08;
#ifdef ESPHOME_OTA_UDP
static constexpr uint8_t SERVER_FEATURE_SUPPORTS_UDP = 0x10;
static constexpr uint8_t UDP_REQUEST_V1 = 0x01;  // espota2.UDP_REQUEST
#endif
 
#ifdef USE_OTA_ENCRYPTION
inline bool ESPHomeOTAComponent::noise_offered_() const {
  return (this->handshake_buf_[1] & SERVER_FEATURE_SUPPORTS_NOISE) != 0 &&
         (this->ota_features_ & CLIENT_NOISE_FEATURES) == CLIENT_NOISE_FEATURES;
}
#endif
 
inline bool ESPHomeOTAComponent::extended_proto_() const {
#ifdef USE_OTA_ENCRYPTION_REQUIRED
  // FEATURE_READ already refused every client without the extended protocol
  return true;
#else
  return (this->ota_features_ & CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL) != 0;
#endif
}
 
void ESPHomeOTAComponent::handle_handshake_() {
 
  if (this->client_ == nullptr) {
    // We already checked server_->ready() in loop(), so we can accept directly
    struct sockaddr_storage source_addr;
    socklen_t addr_len = sizeof(source_addr);
    int enable = 1;
 
    this->client_ = this->server_->accept_loop_monitored((struct sockaddr *) &source_addr, &addr_len);
    if (this->client_ == nullptr)
      return;
    int err = this->client_->setsockopt(IPPROTO_TCP, TCP_NODELAY, &enable, sizeof(int));
    if (err != 0) {
      this->log_socket_error_(LOG_STR("nodelay"));
      this->cleanup_connection_();
      return;
    }
    err = this->client_->setblocking(false);
    if (err != 0) {
      this->log_socket_error_(LOG_STR("non-blocking"));
      this->cleanup_connection_();
      return;
    }
    this->log_start_(LOG_STR("handshake"));
    this->client_connect_time_ = App.get_loop_component_start_time();
    this->handshake_buf_pos_ = 0;  // Reset handshake buffer position
    this->ota_state_ = OTAState::MAGIC_READ;
  }
 
  // Check for handshake timeout
  uint32_t now = App.get_loop_component_start_time();
  if (now - this->client_connect_time_ > OTA_SOCKET_TIMEOUT_HANDSHAKE) {
    ESP_LOGW(TAG, "Handshake timeout");
    this->cleanup_connection_();
    return;
  }
  // Only unauthenticated handshakes reach here; the data phase blocks inside handle_data_
  if (now - this->client_connect_time_ > OTA_HANDSHAKE_PREEMPT && this->server_->ready()) {
    ESP_LOGD(TAG, "Newer connection waiting");
    this->cleanup_connection_();
    return;
  }
 
  switch (this->ota_state_) {
    case OTAState::MAGIC_READ: {
      // Try to read remaining magic bytes (5 total)
      if (!this->try_read_(5, LOG_STR("read magic"))) {
        return;
      }
 
      // Validate magic bytes
      if (memcmp(this->handshake_buf_, MAGIC_BYTES, sizeof(MAGIC_BYTES)) != 0) {
        ESP_LOGW(TAG, "Magic bytes mismatch! 0x%02X-0x%02X-0x%02X-0x%02X-0x%02X", this->handshake_buf_[0],
                 this->handshake_buf_[1], this->handshake_buf_[2], this->handshake_buf_[3], this->handshake_buf_[4]);
        this->send_error_and_cleanup_(ota::OTA_RESPONSE_ERROR_MAGIC);
        return;
      }
 
      // Magic bytes valid, move to next state
      this->transition_ota_state_(OTAState::MAGIC_ACK);
      this->handshake_buf_[0] = ota::OTA_RESPONSE_OK;
      this->handshake_buf_[1] = USE_OTA_VERSION;
      [[fallthrough]];
    }
 
    case OTAState::MAGIC_ACK: {
      // Send OK and version - 2 bytes
      if (!this->try_write_(2, LOG_STR("ack magic"))) {
        return;
      }
      // All bytes sent, create backend and move to next state
      this->backend_ = ota::make_ota_backend();
      this->transition_ota_state_(OTAState::FEATURE_READ);
      [[fallthrough]];
    }
 
    case OTAState::FEATURE_READ: {
      // Read features - 1 byte
      if (!this->try_read_(1, LOG_STR("read feature"))) {
        return;
      }
      this->ota_features_ = this->handshake_buf_[0];
      ESP_LOGV(TAG, "Features: 0x%02X", this->ota_features_);
 
#ifdef USE_OTA_ENCRYPTION_REQUIRED
      // `ota: encryption:` requires the client to negotiate encryption
      if ((this->ota_features_ & CLIENT_NOISE_FEATURES) != CLIENT_NOISE_FEATURES) {
        ESP_LOGW(TAG, "Client does not support encryption");
        this->send_error_and_cleanup_(ota::OTA_RESPONSE_ERROR_ENCRYPTION_REQUIRED);
        return;
      }
#endif
 
      this->transition_ota_state_(OTAState::FEATURE_ACK);
 
      const bool supports_compression =
          (this->ota_features_ & CLIENT_FEATURE_SUPPORTS_COMPRESSION) != 0 && ota::OTABackend::supports_compression();
 
      // Compose the feature-ack response. When the client negotiates the extended protocol we emit
      // a 2-byte response (marker + server feature flags); otherwise we emit the single-byte
      // legacy response.
      if (this->extended_proto_()) {
        static_assert(HANDSHAKE_BUF_SIZE >= 2, "handshake_buf_ must hold the 2-byte extended-protocol feature ack");
        this->handshake_buf_[0] = ota::OTA_RESPONSE_FEATURE_FLAGS;
        this->handshake_buf_[1] = (supports_compression ? SERVER_FEATURE_SUPPORTS_COMPRESSION : 0);
#ifdef ESPHOME_OTA_UDP
        this->handshake_buf_[1] |= SERVER_FEATURE_SUPPORTS_UDP;
#endif
#ifdef USE_OTA_PARTITIONS
        this->handshake_buf_[1] |= SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS;
#endif
#ifdef USE_OTA_ENCRYPTION_PROVISIONED
        // A runtime provisioned key may not exist yet
        if (this->noise_context_().has_psk()) {
          this->handshake_buf_[1] |= SERVER_FEATURE_SUPPORTS_NOISE;
        }
#elif defined(USE_OTA_ENCRYPTION)
        // A yaml key always exists: validation rejects the all-zeros key
        this->handshake_buf_[1] |= SERVER_FEATURE_SUPPORTS_NOISE;
#endif
#ifdef USE_OTA_ENCRYPTION
        // Reserve the noise session before the optional inflate buffer, so the
        // required allocation is not starved by the compression window
        if (this->noise_offered_()) {
          this->noise_reserve_session_();
        }
#endif
#ifdef USE_OTA_DEFLATE
        // Offered only once the session memory is in hand; else uncompressed
        if ((this->ota_features_ & CLIENT_FEATURE_SUPPORTS_DEFLATE) != 0) {
          // Value initialized: a corrupt stream that back references the
          // window before it is filled then copies zeros, never stale memory.
          // Default placement, PSRAM first where present: the session lives for one
          // upload and keeps 4.9 KB of internal heap free while it runs
          this->inflate_ = RAMAllocator<InflateSession>().make_unique();
          if (this->inflate_ != nullptr) {
            this->handshake_buf_[1] |= SERVER_FEATURE_SUPPORTS_DEFLATE;
          } else {
            ESP_LOGW(TAG, "No memory to inflate");
          }
        }
#endif
      } else {
        this->handshake_buf_[0] =
            supports_compression ? ota::OTA_RESPONSE_SUPPORTS_COMPRESSION : ota::OTA_RESPONSE_HEADER_OK;
      }
      [[fallthrough]];
    }
 
    case OTAState::FEATURE_ACK: {
      static constexpr size_t STANDARD_PROTO_ACK_SIZE = 1;
      static constexpr size_t EXTENDED_PROTO_ACK_SIZE = 2;
      const size_t ack_size = this->extended_proto_() ? EXTENDED_PROTO_ACK_SIZE : STANDARD_PROTO_ACK_SIZE;
      if (!this->try_write_(ack_size, LOG_STR("ack feature"))) {
        return;
      }
#ifdef USE_OTA_ENCRYPTION
      // Latch the offer actually sent: a key activating between the two
      // states must not start a session the client never expects
      if (this->noise_offered_()) {
        // handshake_buf_ still holds the feature ack composed above; a
        // would-block re-entry lands here without rebuilding it
        if (!this->noise_start_session_(this->handshake_buf_[1])) {
          return;
        }
        this->transition_ota_state_(OTAState::NOISE_HANDSHAKE);
        return;
      }
#endif
#ifdef USE_OTA_PASSWORD
      // If password is set, move to auth phase
      if (!this->password_.empty()) {
        this->transition_ota_state_(OTAState::AUTH_SEND);
      } else
#endif
      {
        // No password, move directly to data phase
        this->transition_ota_state_(OTAState::DATA);
      }
      [[fallthrough]];
    }
 
#ifdef USE_OTA_PASSWORD
    case OTAState::AUTH_SEND: {
      // Non-blocking authentication send
      if (!this->handle_auth_send_()) {
        return;
      }
      this->transition_ota_state_(OTAState::AUTH_READ);
      [[fallthrough]];
    }
 
    case OTAState::AUTH_READ: {
      // Non-blocking authentication read & verify
      if (!this->handle_auth_read_()) {
        return;
      }
      this->transition_ota_state_(OTAState::DATA);
      [[fallthrough]];
    }
#endif
 
    case OTAState::DATA:
      this->handle_data_();
      return;
 
#ifdef USE_OTA_ENCRYPTION
    case OTAState::NOISE_HANDSHAKE:
      if (!this->handle_noise_handshake_()) {
        return;
      }
      this->transition_ota_state_(OTAState::DATA);
      this->handle_data_();
      return;
#endif
 
    default:
      break;
  }
}
 
void ESPHomeOTAComponent::handle_data_() {
  // Backend calls overwrite this with OK; reset to UNKNOWN before any
  // goto error that follows a successful begin()/write()
  ota::OTAResponseTypes error_code = ota::OTA_RESPONSE_ERROR_UNKNOWN;
  DataTransfer xfer;
  uint8_t buf[OTA_BUFFER_SIZE];
  char *sbuf = reinterpret_cast<char *>(buf);
  size_t image_size;
  ota::OTAType ota_type = ota::OTA_TYPE_UPDATE_APP;
#ifdef ESPHOME_OTA_UDP
  UdpLink udp_link;
  this->udp_ = nullptr;  // client_connect_time_ is done with
#endif
 
  // Set socket timeouts and blocking mode (see strategy table above)
  struct timeval tv;
  tv.tv_sec = 2;
  tv.tv_usec = 0;
  this->client_->setsockopt(SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
  this->client_->setsockopt(SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv));
  if (this->client_->setblocking(true) != 0) {
    this->log_socket_error_(LOG_STR("blocking"));
    goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
  }
 
  // Acknowledge auth OK - 1 byte
  this->data_write_byte_(ota::OTA_RESPONSE_AUTH_OK);
 
#ifdef ESPHOME_OTA_UDP
  // The server always offers UDP on the extended protocol
  if (this->extended_proto_() && (this->ota_features_ & CLIENT_FEATURE_SUPPORTS_UDP) != 0) {
    // Whether the client wants UDP (0: its link looks good), answered with our token
    if (!this->data_readall_(buf, 1)) {
      if (this->client_left_before_start_())
        return;
      this->log_read_error_(LOG_STR("UDP request"));
      goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
    }
    // 1 is this UDP format; other values are free for later ones, which we decline (a zero reply keeps TCP)
    if (buf[0] == UDP_REQUEST_V1) {
      this->udp_open_(udp_link);
      if (udp_link.reply[0] == 0) {
        ESP_LOGW(TAG, "UDP unavailable, using TCP");
      }
    }
    if (!this->data_write_(udp_link.reply, sizeof(udp_link.reply)))
      goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
    if (udp_link.reply[0] != 0) {
      this->udp_ = &udp_link;
      if (!this->udp_start_()) {
        // Nothing started yet: just drop the session, which the client retries as a network failure
        this->log_read_error_(LOG_STR("UDP start"));
        this->cleanup_connection_();
        return;
      }
    }
  }
 
#endif
 
  if (this->extended_proto_()) {
    // Read ota type, 1 byte
    if (!this->data_readall_(buf, 1)) {
      if (this->client_left_before_start_())
        return;
      this->log_read_error_(LOG_STR("OTA type"));
      goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
    }
    ota_type = static_cast<ota::OTAType>(buf[0]);
  }
  ESP_LOGV(TAG, "OTA type is 0x%02x", ota_type);
 
  // Read size, 4 bytes MSB first
  if (!this->read_size_(buf, xfer.ota_size, LOG_STR("size"))) {
    // The first request byte is the type on the extended protocol; a close after it was a cut-off request
    if (!this->extended_proto_() && this->client_left_before_start_())
      return;
    this->log_read_error_(LOG_STR("size"));
    goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
  }
  image_size = xfer.ota_size;
#ifdef USE_OTA_DEFLATE
  if (this->inflate_ != nullptr && !this->read_size_(buf, image_size, LOG_STR("image size"))) {
    this->log_read_error_(LOG_STR("image size"));
    goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
  }
#endif
 
#ifndef USE_OTA_PARTITIONS
  if (ota_type != ota::OTA_TYPE_UPDATE_APP) {
    error_code = ota::OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE;
    goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
  }
#endif
 
  // Now that we've passed authentication and are actually
  // starting the update, set the warning status and notify
  // listeners. This ensures that port scanners do not
  // accidentally trigger the update process.
  this->log_start_(LOG_STR("update"));
  this->status_set_warning();
#ifdef USE_OTA_STATE_LISTENER
  this->notify_state_(ota::OTA_STARTED, 0.0f, 0);
#endif
 
  // begin() returns quickly; flash sectors are erased incrementally during write().
  error_code = this->backend_->begin(image_size, ota_type);
  if (error_code != ota::OTA_RESPONSE_OK)
    goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
 
  // Acknowledge prepare OK - 1 byte
  this->data_write_byte_(ota::OTA_RESPONSE_UPDATE_PREPARE_OK);
 
  // Read binary MD5, 32 bytes
  if (!this->data_readall_(buf, 32)) {
    this->log_read_error_(LOG_STR("MD5 checksum"));
    error_code = ota::OTA_RESPONSE_ERROR_UNKNOWN;
    goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
  }
  sbuf[32] = '\0';
  ESP_LOGV(TAG, "Update: Binary MD5 is %s", sbuf);
  this->backend_->set_update_md5(sbuf);
 
  // Acknowledge MD5 OK - 1 byte
  this->data_write_byte_(ota::OTA_RESPONSE_BIN_MD5_OK);
 
  xfer.last_data_ms = millis();
#ifdef USE_OTA_DEFLATE
  if (this->inflate_ != nullptr) {
    error_code = this->inflate_data_(buf, image_size, xfer);
    if (error_code != ota::OTA_RESPONSE_OK)
      goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
  } else
#endif
  {
    while (xfer.total < xfer.ota_size) {
      ssize_t read = this->receive_data_(buf, xfer);
      if (read < 0) {
        error_code = ota::OTA_RESPONSE_ERROR_UNKNOWN;
        goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
      }
      error_code = this->write_flash_(buf, read);
      if (error_code != ota::OTA_RESPONSE_OK)
        goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
      this->ack_written_(xfer);
    }
  }
 
  // Acknowledge receive OK - 1 byte
  this->data_write_byte_(ota::OTA_RESPONSE_RECEIVE_OK);
 
  error_code = this->backend_->end();
  if (error_code != ota::OTA_RESPONSE_OK) {
    ESP_LOGW(TAG, "End update err %d", error_code);
    goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
  }
 
  // Acknowledge Update end OK - 1 byte
  this->data_write_byte_(ota::OTA_RESPONSE_UPDATE_END_OK);
 
#ifdef ESPHOME_OTA_UDP
  // UDP never reports a close: wait only until the client has our responses, not for its final ack
  if (this->udp_ != nullptr) {
    this->udp_linger_();
  } else
#endif
    // Read ACK
    if (!this->data_readall_(buf, 1) || buf[0] != ota::OTA_RESPONSE_OK) {
      this->log_read_error_(LOG_STR("ack"));
      // do not go to error, this is not fatal
    }
 
  this->cleanup_connection_();
  delay(10);
  ESP_LOGI(TAG, "Update complete");
  this->status_clear_warning();
#ifdef USE_OTA_STATE_LISTENER
  this->notify_state_(ota::OTA_COMPLETED, 100.0f, 0);
#endif
  delay(100);  // NOLINT
#ifdef USE_OTA_PARTITIONS
  if (ota_type == ota::OTA_TYPE_UPDATE_PARTITION_TABLE) {
    // Skip on_safe_shutdown: nvs_flash_deinit() has already invalidated open NVS handles, so
    // preferences flush would emit ESP_ERR_NVS_INVALID_HANDLE for every entry. Reboot directly.
    App.reboot();
  }
#endif
  App.safe_reboot();
 
error:
  this->data_write_byte_(static_cast<uint8_t>(error_code));
#ifdef ESPHOME_OTA_UDP
  this->udp_linger_();
#endif
 
  // Abort backend before cleanup - cleanup_connection_() destroys the backend.
  // Always call abort() unconditionally: backends register external partitions before
  // esp_ota_begin (partition table / bootloader paths), and abort() is responsible for
  // releasing those even if begin() failed before an OTA handle was opened. The IDF
  // backend's esp_ota_abort(0) is documented as harmless.
  if (this->backend_ != nullptr) {
    this->backend_->abort();
  }
 
  this->cleanup_connection_();
 
  this->status_momentary_error(5000);
#ifdef USE_OTA_STATE_LISTENER
  this->notify_state_(ota::OTA_ERROR, 0.0f, static_cast<uint8_t>(error_code));
#endif
}
 
bool ESPHomeOTAComponent::readall_(uint8_t *buf, size_t len) {
  uint32_t start = millis();
  uint32_t at = 0;
  while (len - at > 0) {
    uint32_t now = millis();
    if (now - start > OTA_SOCKET_TIMEOUT_DATA) {
      ESP_LOGW(TAG, "Timeout reading %zu bytes", len);
      return false;
    }
 
    ssize_t read = this->transport_read_(buf + at, len - at);
    if (read == -1) {
      const int err = errno;
      if (!this->would_block_(err)) {
        ESP_LOGW(TAG, "Read err %zu bytes, errno %d", len, err);
        return false;
      }
    } else if (read == 0) {
      // A partial message is a cut-off request, not a clean close; the caller reports the clean one
      this->remote_closed_ = at == 0;
      if (at > 0) {
        ESP_LOGW(TAG, "Remote closed after %u of %zu bytes", (unsigned) at, len);
      }
      return false;
    } else {
      at += read;
    }
    // read() already waited via SO_RCVTIMEO, just yield without 1ms stall
    App.feed_wdt();
    delay(0);
  }
 
  return true;
}
bool ESPHomeOTAComponent::writeall_(const uint8_t *buf, size_t len) {
#ifdef ESPHOME_OTA_UDP
  if (this->udp_ != nullptr)
    return this->udp_send_(buf, len);
#endif
  uint32_t start = millis();
  uint32_t at = 0;
  while (len - at > 0) {
    uint32_t now = millis();
    if (now - start > OTA_SOCKET_TIMEOUT_DATA) {
      ESP_LOGW(TAG, "Timeout writing %zu bytes", len);
      return false;
    }
 
    ssize_t written = this->client_->write(buf + at, len - at);
    if (written == -1) {
      const int err = errno;
      if (!this->would_block_(err)) {
        ESP_LOGW(TAG, "Write err %zu bytes, errno %d", len, err);
        return false;
      }
      // EWOULDBLOCK: on raw TCP writes never block, delay(1) prevents spinning
      this->yield_and_feed_watchdog_();
    } else {
      at += written;
      // write() may block up to SO_SNDTIMEO on BSD/lwip sockets, feed WDT
      App.feed_wdt();
    }
  }
  return true;
}
 
float ESPHomeOTAComponent::get_setup_priority() const { return setup_priority::AFTER_WIFI; }
 
void ESPHomeOTAComponent::log_socket_error_(const LogString *msg) {
  ESP_LOGW(TAG, "Socket %s: errno %d", LOG_STR_ARG(msg), errno);
}
 
bool ESPHomeOTAComponent::client_left_before_start_() {
  // Key probes and scanners hang up right after the handshake; nothing started, so no error status or callback
  if (!this->remote_closed_)
    return false;
  ESP_LOGD(TAG, "Client left after the handshake");
  this->cleanup_connection_();
  return true;
}
 
void ESPHomeOTAComponent::log_read_error_(const LogString *what) {
  if (this->remote_closed_) {
    this->log_remote_closed_(what);
    return;
  }
  ESP_LOGW(TAG, "Read %s failed", LOG_STR_ARG(what));
}
 
void ESPHomeOTAComponent::log_start_(const LogString *phase) {
  char peername[socket::SOCKADDR_STR_LEN];
#ifdef ESPHOME_OTA_UDP
  if (this->client_ == nullptr) {
    // A committed UDP session has already dropped TCP
#ifdef ESPHOME_OTA_UDP_LWIP
    ipaddr_ntoa_r(&this->udp_->peer_ip, peername, sizeof(peername));
#else
    socket::format_sockaddr_to(reinterpret_cast<const struct sockaddr *>(&this->udp_->peer), this->udp_->peer_len,
                               peername);
#endif
  } else
#endif
  {
    this->client_->getpeername_to(peername);
  }
  ESP_LOGD(TAG, "Starting %s from %s", LOG_STR_ARG(phase), peername);
}
 
void ESPHomeOTAComponent::log_remote_closed_(const LogString *during) {
  ESP_LOGW(TAG, "Remote closed at %s", LOG_STR_ARG(during));
}
 
void ESPHomeOTAComponent::server_failed_(const LogString *msg) {
  this->log_socket_error_(msg);
  // No explicit close() needed — listen sockets have no active connections on
  // failure/shutdown. Destructor handles fd cleanup (close or abort per platform).
  delete this->server_;
  this->server_ = nullptr;
  this->mark_failed();
}
 
bool ESPHomeOTAComponent::handle_read_error_(ssize_t read, const LogString *desc) {
  if (read == -1 && this->would_block_(errno)) {
    return false;  // No data yet, try again next loop
  }
 
  if (read <= 0) {
    read == 0 ? this->log_remote_closed_(desc) : this->log_socket_error_(desc);
    this->cleanup_connection_();
    return false;
  }
  return true;
}
 
bool ESPHomeOTAComponent::handle_write_error_(ssize_t written, const LogString *desc) {
  if (written == -1) {
    if (this->would_block_(errno)) {
      return false;  // Try again next loop
    }
    this->log_socket_error_(desc);
    this->cleanup_connection_();
    return false;
  }
  return true;
}
 
bool ESPHomeOTAComponent::try_read_(size_t to_read, const LogString *desc) {
  // Read bytes into handshake buffer, starting at handshake_buf_pos_
  size_t bytes_to_read = to_read - this->handshake_buf_pos_;
  ssize_t read = this->client_->read(this->handshake_buf_ + this->handshake_buf_pos_, bytes_to_read);
 
  if (!this->handle_read_error_(read, desc)) {
    return false;
  }
 
  this->handshake_buf_pos_ += read;
  // Return true only if we have all the requested bytes
  return this->handshake_buf_pos_ >= to_read;
}
 
bool ESPHomeOTAComponent::try_write_(size_t to_write, const LogString *desc) {
  // Write bytes from handshake buffer, starting at handshake_buf_pos_
  size_t bytes_to_write = to_write - this->handshake_buf_pos_;
  ssize_t written = this->client_->write(this->handshake_buf_ + this->handshake_buf_pos_, bytes_to_write);
 
  if (!this->handle_write_error_(written, desc)) {
    return false;
  }
 
  this->handshake_buf_pos_ += written;
  // Return true only if we have written all the requested bytes
  return this->handshake_buf_pos_ >= to_write;
}
 
bool ESPHomeOTAComponent::read_size_(uint8_t *buf, size_t &size, const LogString *desc) {
  if (!this->data_readall_(buf, OTA_SIZE_FIELD_BYTES))
    return false;
  size = encode_uint32(buf[0], buf[1], buf[2], buf[3]);
  ESP_LOGV(TAG, "%s is %zu bytes", LOG_STR_ARG(desc), size);
  return true;
}
 
ota::OTAResponseTypes ESPHomeOTAComponent::write_flash_(uint8_t *data, size_t len) {
  ota::OTAResponseTypes result = this->backend_->write(data, len);
  if (result != ota::OTA_RESPONSE_OK) {
    ESP_LOGW(TAG, "Flash write err %d", result);
  }
  return result;
}
 
ssize_t ESPHomeOTAComponent::receive_data_(uint8_t *buf, DataTransfer &xfer) {
  const size_t remaining = xfer.ota_size - xfer.total;
  const size_t requested = std::min(remaining, OTA_BUFFER_SIZE);
  ssize_t read;
  for (;;) {
    // A silently-vanished peer (no FIN/RST delivered, e.g. uploader killed
    // mid-transfer or NAT/router dropped state) must not wedge the device:
    // read() only fails on EOF or a real error, and lwIP TCP keepalive isn't
    // enabled here.
    if (millis() - xfer.last_data_ms > OTA_SOCKET_TIMEOUT_DATA) {
      ESP_LOGW(TAG, "No data received for %u ms", (unsigned) OTA_SOCKET_TIMEOUT_DATA);
      return -1;
    }
#ifdef USE_OTA_ENCRYPTION
    if (this->noise_ != nullptr) {
      // One frame per call; noise_read_data_ waits internally (readall_), so
      // there is no would-block retry here and failures are already logged.
      read = this->noise_read_data_(buf, requested);
      if (read <= 0) {
        if (this->remote_closed_)
          this->log_remote_closed_(LOG_STR("data"));
        return -1;
      }
      break;
    }
#endif
    read = this->transport_read_(buf, requested);
    if (read > 0)
      break;
    if (read == 0) {
      this->log_remote_closed_(LOG_STR("data"));
      return -1;
    }
    if (!this->would_block_(errno)) {
      this->log_socket_error_(LOG_STR("data"));
      return -1;
    }
    // read() already waited up to SO_RCVTIMEO for data, just feed WDT
    App.feed_wdt();
  }
 
  const uint32_t now = millis();
  xfer.last_data_ms = now;
  xfer.total += read;
  this->ack_received_(xfer);
  if (now - xfer.last_progress > OTA_PROGRESS_INTERVAL_MS) {
    xfer.last_progress = now;
    float percentage = (xfer.total * 100.0f) / xfer.ota_size;
    ESP_LOGD(TAG, "Progress: %0.1f%%", percentage);
#ifdef USE_OTA_STATE_LISTENER
    this->notify_state_(ota::OTA_IN_PROGRESS, percentage, 0);
#endif
    // feed watchdog and give other tasks a chance to run
    this->yield_and_feed_watchdog_();
  }
  return read;
}
 
void ESPHomeOTAComponent::send_chunk_acks_(DataTransfer &xfer) {
#if USE_OTA_VERSION == 2
#ifdef ESPHOME_OTA_UDP
  // UDP acks every message itself
  if (this->udp_ != nullptr)
    return;
#endif
  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
}
 
void ESPHomeOTAComponent::cleanup_connection_() {
#ifdef ESPHOME_OTA_UDP
  this->udp_ = nullptr;  // the link itself lives on handle_data_'s stack
#endif
  // Already gone once a UDP session committed
  if (this->client_ != nullptr)
    this->client_->close();
  this->client_ = nullptr;
  this->client_connect_time_ = 0;
  this->handshake_buf_pos_ = 0;
  this->ota_state_ = OTAState::IDLE;
  this->ota_features_ = 0;
  this->remote_closed_ = false;
  this->backend_ = nullptr;
#ifdef USE_OTA_PASSWORD
  this->cleanup_auth_();
#endif
#ifdef USE_OTA_ENCRYPTION
  this->noise_ = nullptr;
#endif
#ifdef USE_OTA_DEFLATE
  this->inflate_ = nullptr;
#endif
  // Intentionally no disable_loop() — letting loop() run one more iteration catches
  // any connection that queued on the listener mid-session (otherwise the wake flag,
  // set while we were in LOOP state, would be lost to enable_pending_loops_()).
}
 
void ESPHomeOTAComponent::yield_and_feed_watchdog_() {
  App.feed_wdt();
  delay(1);
}
 
#ifdef USE_OTA_PASSWORD
void ESPHomeOTAComponent::log_auth_warning_(const LogString *msg) { ESP_LOGW(TAG, "Auth: %s", LOG_STR_ARG(msg)); }
 
bool ESPHomeOTAComponent::select_auth_type_() {
  bool client_supports_sha256 = (this->ota_features_ & CLIENT_FEATURE_SUPPORTS_SHA256_AUTH) != 0;
 
  // Require SHA256
  if (!client_supports_sha256) {
    this->log_auth_warning_(LOG_STR("SHA256 required"));
    this->send_error_and_cleanup_(ota::OTA_RESPONSE_ERROR_AUTH_INVALID);
    return false;
  }
  this->auth_type_ = ota::OTA_RESPONSE_REQUEST_SHA256_AUTH;
  return true;
}
 
bool ESPHomeOTAComponent::handle_auth_send_() {
  // Initialize auth buffer if not already done
  if (!this->auth_buf_) {
    // Select auth type based on client capabilities and configuration
    if (!this->select_auth_type_()) {
      return false;
    }
 
    // Generate nonce - hasher must be created and used in same stack frame
    // CRITICAL ESP32-S3 HARDWARE SHA ACCELERATION REQUIREMENTS:
    // 1. Hash objects must NEVER be passed to another function (different stack frame)
    // 2. NO Variable Length Arrays (VLAs) - they corrupt the stack with hardware DMA
    // 3. All hash operations (init/add/calculate) must happen in the SAME function where object is created
    // Violating these causes truncated hash output (20 bytes instead of 32) or memory corruption.
    //
    // Buffer layout after AUTH_READ completes:
    //   [0]: auth_type (1 byte)
    //   [1...hex_size]: nonce (hex_size bytes) - our random nonce sent in AUTH_SEND
    //   [1+hex_size...1+2*hex_size-1]: cnonce (hex_size bytes) - client's nonce
    //   [1+2*hex_size...1+3*hex_size-1]: response (hex_size bytes) - client's hash
 
    // CRITICAL ESP32-S2/S3 HARDWARE SHA ACCELERATION: Hash object must stay in same stack frame
    // (no passing to other functions). All hash operations must happen in this function.
    sha256::SHA256 hasher;
 
    const size_t hex_size = hasher.get_size() * 2;
    const size_t nonce_len = hasher.get_size() / 4;
    const size_t auth_buf_size = 1 + 3 * hex_size;
    // Internal RAM first: 128 of these bytes go straight into the hardware SHA engine
    this->auth_buf_ =
        RAMAllocator<uint8_t>(RAMAllocator<uint8_t>::PREFER_INTERNAL).make_unique_array_for_overwrite(auth_buf_size);
    if (!this->auth_buf_) {
      this->log_auth_warning_(LOG_STR("No memory"));
      this->send_error_and_cleanup_(ota::OTA_RESPONSE_ERROR_UNKNOWN);
      return false;
    }
    this->auth_buf_pos_ = 0;
 
    char *buf = reinterpret_cast<char *>(this->auth_buf_.get() + 1);
    if (!random_bytes(reinterpret_cast<uint8_t *>(buf), nonce_len)) {
      this->log_auth_warning_(LOG_STR("Random failed"));
      this->send_error_and_cleanup_(ota::OTA_RESPONSE_ERROR_UNKNOWN);
      return false;
    }
 
    hasher.init();
    hasher.add(buf, nonce_len);
    hasher.calculate();
    this->auth_buf_[0] = this->auth_type_;
    hasher.get_hex(buf);
 
    ESP_LOGV(TAG, "Auth: Nonce is %.*s", (int) hex_size, buf);
  }
 
  // Try to write auth_type + nonce
  constexpr size_t hex_size = SHA256_HEX_SIZE;
  const size_t to_write = 1 + hex_size;
  size_t remaining = to_write - this->auth_buf_pos_;
 
  ssize_t written = this->client_->write(this->auth_buf_.get() + this->auth_buf_pos_, remaining);
  if (!this->handle_write_error_(written, LOG_STR("ack auth"))) {
    return false;
  }
 
  this->auth_buf_pos_ += written;
 
  // Check if we still have more to write
  if (this->auth_buf_pos_ < to_write) {
    return false;  // More to write, try again next loop
  }
 
  // All written, prepare for reading phase
  this->auth_buf_pos_ = 0;
  return true;
}
 
bool ESPHomeOTAComponent::handle_auth_read_() {
  constexpr size_t hex_size = SHA256_HEX_SIZE;
  const size_t to_read = hex_size * 2;  // CNonce + Response
 
  // Try to read remaining bytes (CNonce + Response)
  // We read cnonce+response starting at offset 1+hex_size (after auth_type and our nonce)
  size_t cnonce_offset = 1 + hex_size;  // Offset where cnonce should be stored in buffer
  size_t remaining = to_read - this->auth_buf_pos_;
  ssize_t read = this->client_->read(this->auth_buf_.get() + cnonce_offset + this->auth_buf_pos_, remaining);
 
  if (!this->handle_read_error_(read, LOG_STR("read auth"))) {
    return false;
  }
 
  this->auth_buf_pos_ += read;
 
  // Check if we still need more data
  if (this->auth_buf_pos_ < to_read) {
    return false;  // More to read, try again next loop
  }
 
  // We have all the data, verify it
  const char *nonce = reinterpret_cast<char *>(this->auth_buf_.get() + 1);
  const char *cnonce = nonce + hex_size;
  const char *response = cnonce + hex_size;
 
  // CRITICAL ESP32-S2/S3 HARDWARE SHA ACCELERATION: Hash object must stay in same stack frame
  // (no passing to other functions). All hash operations must happen in this function.
  sha256::SHA256 hasher;
 
  hasher.init();
  hasher.add(this->password_.c_str(), this->password_.length());
  hasher.add(nonce, hex_size * 2);  // Add both nonce and cnonce (contiguous in buffer)
  hasher.calculate();
 
  ESP_LOGV(TAG, "Auth: CNonce is %.*s", (int) hex_size, cnonce);
#if ESPHOME_LOG_LEVEL >= ESPHOME_LOG_LEVEL_VERBOSE
  char computed_hash[SHA256_HEX_SIZE + 1];  // Buffer for hex-encoded hash (max expected length + null terminator)
  hasher.get_hex(computed_hash);
  ESP_LOGV(TAG, "Auth: Result is %.*s", (int) hex_size, computed_hash);
#endif
  ESP_LOGV(TAG, "Auth: Response is %.*s", (int) hex_size, response);
 
  // Compare response
  bool matches = hasher.equals_hex(response);
 
  if (!matches) {
    this->log_auth_warning_(LOG_STR("Password mismatch"));
    this->send_error_and_cleanup_(ota::OTA_RESPONSE_ERROR_AUTH_INVALID);
    return false;
  }
 
  // Authentication successful - clean up auth state
  this->cleanup_auth_();
 
  return true;
}
 
void ESPHomeOTAComponent::cleanup_auth_() {
  this->auth_buf_ = nullptr;
  this->auth_buf_pos_ = 0;
  this->auth_type_ = 0;
}
#endif  // USE_OTA_PASSWORD
 
}  // namespace esphome
#endif