Merge remote-tracking branch 'origin/dev' into ota-udp-transport

# Conflicts:
#	esphome/espota2.py
#	tests/unit_tests/test_espota2_noise.py

esphome/components/esphome/ota/ota_esphome.cpp

@@ -44,11 +44,14 @@ const noise::NoiseContext &ESPHomeOTAComponent::noise_context_() const {
 }
 #endif
 static constexpr uint16_t OTA_BLOCK_SIZE = 8192;
-static constexpr uint32_t OTA_SOCKET_TIMEOUT_HANDSHAKE = 20000;  // milliseconds for initial handshake
+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
 
@@ -199,6 +202,9 @@ 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;
@@ -208,6 +214,10 @@ 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 {
@@ -266,6 +276,12 @@ void ESPHomeOTAComponent::handle_handshake_() {
     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: {
@@ -329,6 +345,9 @@ void ESPHomeOTAComponent::handle_handshake_() {
         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
@@ -483,6 +502,10 @@ void ESPHomeOTAComponent::handle_data_() {
   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;
@@ -498,6 +521,38 @@ void ESPHomeOTAComponent::handle_data_() {
   // 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)) {
@@ -598,11 +653,17 @@ void ESPHomeOTAComponent::handle_data_() {
   // Acknowledge Update end OK - 1 byte
   this->data_write_byte_(ota::OTA_RESPONSE_UPDATE_END_OK);
 
-  // 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
-  }
+#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);
@@ -623,6 +684,9 @@ void ESPHomeOTAComponent::handle_data_() {
 
 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
@@ -651,7 +715,7 @@ bool ESPHomeOTAComponent::readall_(uint8_t *buf, size_t len) {
       return false;
     }
 
-    ssize_t read = this->client_->read(buf + at, len - at);
+    ssize_t read = this->transport_read_(buf + at, len - at);
     if (read == -1) {
       const int err = errno;
       if (!this->would_block_(err)) {
@@ -676,6 +740,10 @@ bool ESPHomeOTAComponent::readall_(uint8_t *buf, size_t len) {
   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) {
@@ -728,7 +796,20 @@ void ESPHomeOTAComponent::log_read_error_(const LogString *what) {
 
 void ESPHomeOTAComponent::log_start_(const LogString *phase) {
   char peername[socket::SOCKADDR_STR_LEN];
-  this->client_->getpeername_to(peername);
+#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);
 }
 
@@ -840,7 +921,7 @@ ssize_t ESPHomeOTAComponent::receive_data_(uint8_t *buf, DataTransfer &xfer) {
       break;
     }
 #endif
-    read = this->client_->read(buf, requested);
+    read = this->transport_read_(buf, requested);
     if (read > 0)
       break;
     if (read == 0) {
@@ -874,6 +955,11 @@ ssize_t ESPHomeOTAComponent::receive_data_(uint8_t *buf, DataTransfer &xfer) {
 
 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);
@@ -883,7 +969,12 @@ void ESPHomeOTAComponent::send_chunk_acks_(DataTransfer &xfer) {
 }
 
 void ESPHomeOTAComponent::cleanup_connection_() {
-  this->client_->close();
+#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;

esphome/components/esphome/ota/ota_esphome.h

@@ -15,6 +15,25 @@
 #include "esphome/core/preferences.h"
 #include "esphome/core/hash_base.h"
 
+// UDP data phase: raw lwIP needs a main loop lock (or a single thread); the host polls a socket instead
+#ifdef USE_ESP32
+#include <sdkconfig.h>
+#endif
+#if defined(USE_ESP8266) || defined(USE_RP2) || defined(USE_LIBRETINY) || \
+    (defined(USE_ESP32) && defined(CONFIG_LWIP_TCPIP_CORE_LOCKING))
+#define ESPHOME_OTA_UDP
+#define ESPHOME_OTA_UDP_LWIP
+#include "lwip/ip_addr.h"
+#include "lwip/pbuf.h"
+#include "lwip/udp.h"
+#if defined(USE_LIBRETINY) && !LWIP_TCPIP_CORE_LOCKING
+#undef ESPHOME_OTA_UDP
+#undef ESPHOME_OTA_UDP_LWIP
+#endif
+#elif defined(USE_HOST)
+#define ESPHOME_OTA_UDP
+#endif
+
 namespace esphome {
 
 /// ESPHomeOTAComponent provides a simple way to integrate Over-the-Air updates into your app using ArduinoOTA.
@@ -65,6 +84,69 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
  protected:
   void handle_handshake_();
   void handle_data_();
+
+#ifdef ESPHOME_OTA_UDP
+  // UDP transport for the data phase; wire format in ota_esphome_udp.cpp
+  static constexpr size_t UDP_OUT_SIZE = 64;  // our unacked responses: at most three encrypted ones
+  // Messages the device keeps unread: like TCP's 4 segment window, so no more memory than TCP holds
+  static constexpr uint8_t UDP_WINDOW = 4;
+#ifdef ESPHOME_OTA_UDP_LWIP
+  using UdpSlot = struct pbuf *;
+#else
+  using UdpSlot = struct OtaUdpDatagram *;
+#endif
+  // Lives on handle_data_'s stack; the receive callback shares it under the lwIP lock
+  struct UdpLink {
+    ~UdpLink();
+    void release();  // closes the socket and frees what is queued
+#ifdef ESPHOME_OTA_UDP_LWIP
+    struct udp_pcb *pcb{nullptr};
+    ip_addr_t peer_ip;
+#else
+    std::unique_ptr<socket::Socket> sock;
+    struct sockaddr_storage peer;  // the TCP peer; its port is replaced by peer_port
+    socklen_t peer_len{0};
+#endif
+    uint16_t peer_port{0};  // learned from the client's datagrams
+    // The session token, our answer to the client: top bit clear, so never an error code; all zero declines
+    uint8_t reply[4]{};
+    UdpSlot queue[UDP_WINDOW]{};  // slot seq % UDP_WINDOW for seq in [consumed, consumed + UDP_WINDOW)
+    uint16_t received{0};         // every seq before it is here
+    uint16_t consumed{0};         // the message being read
+    uint16_t read_offset{0};      // payload bytes of it already read
+    uint16_t out_base{0};         // stream offset of out[0]
+    uint8_t out_len{0};
+    bool committed{false};  // the client heard us over UDP
+    uint8_t out[UDP_OUT_SIZE];
+  };
+  // Fills link.reply on success; leaves it all zero (declined) otherwise
+  void udp_open_(UdpLink &link);
+  // Waits for the client to commit to UDP or fall back to TCP; false on failure
+  bool udp_start_();
+  // Like a blocking socket read with SO_RCVTIMEO: up to len bytes, or -1 with EWOULDBLOCK
+  ssize_t udp_read_(uint8_t *buf, size_t len);
+  bool udp_send_(const uint8_t *data, size_t len);
+  // In a UDP session, keeps answering until the client has our responses (an error code), briefly
+  void udp_linger_();
+  // A datagram from the client's IP; caller holds the lwIP lock (or runs in the receive callback)
+  static void udp_on_datagram(UdpLink &link, UdpSlot p, uint16_t port);
+  // Caller holds the lwIP lock (or runs in the receive callback)
+  static void udp_send_ack(UdpLink &link, uint16_t prompted_by);
+#ifdef ESPHOME_OTA_UDP_LWIP
+  static void udp_recv_cb(void *arg, struct udp_pcb *pcb, struct pbuf *p, const ip_addr_t *addr, uint16_t port);
+  void udp_poll_() {}
+#else
+  void udp_poll_();  // hands what the socket received to udp_on_datagram
+#endif
+#endif
+  // Data phase only (udp_ shares storage with the handshake's connect time): TCP or, once committed, the UDP stream
+  inline ssize_t transport_read_(uint8_t *buf, size_t len) {
+#ifdef ESPHOME_OTA_UDP
+    if (this->udp_ != nullptr)
+      return this->udp_read_(buf, len);
+#endif
+    return this->client_->read(buf, len);
+  }
 #ifdef USE_OTA_PASSWORD
   static constexpr size_t SHA256_HEX_SIZE = 64;  // SHA256 hash as hex string (32 bytes * 2)
   bool handle_auth_send_();
@@ -75,7 +157,6 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
 #endif  // USE_OTA_PASSWORD
   bool readall_(uint8_t *buf, size_t len);
   bool writeall_(const uint8_t *buf, size_t len);
-  inline bool write_byte_(uint8_t byte) { return this->writeall_(&byte, 1); }
 
 #ifdef USE_OTA_ENCRYPTION
   // Heap-allocated only while an encrypted OTA session is active.
@@ -105,18 +186,21 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
   ssize_t noise_read_frame_blocking_(uint8_t *buf, size_t min_ciphertext, size_t max_ciphertext);
   bool noise_readall_(uint8_t *buf, size_t len);
   ssize_t noise_read_data_(uint8_t *buf, size_t capacity);
-  bool noise_write_byte_(uint8_t byte);
+  static constexpr size_t NOISE_MAX_RESPONSE = 4;  // the UDP token
+  bool noise_write_(const uint8_t *data, size_t len);
 #endif  // USE_OTA_ENCRYPTION
 
   // Data-phase I/O dispatch: through the noise transport when a session is
   // active, straight to the socket otherwise.
-  inline bool data_write_byte_(uint8_t byte) {
+  // len is at most NOISE_MAX_RESPONSE
+  inline bool data_write_(const uint8_t *data, size_t len) {
 #ifdef USE_OTA_ENCRYPTION
     if (this->noise_ != nullptr)
-      return this->noise_write_byte_(byte);
+      return this->noise_write_(data, len);
 #endif
-    return this->write_byte_(byte);
+    return this->writeall_(data, len);
   }
+  inline bool data_write_byte_(uint8_t byte) { return this->data_write_(&byte, 1); }
   // When encrypted, buf must have room for len + noise::MAC_SIZE bytes.
   inline bool data_readall_(uint8_t *buf, size_t len) {
 #ifdef USE_OTA_ENCRYPTION
@@ -200,7 +284,21 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
   std::unique_ptr<socket::Socket> client_;
   ota::OTABackendPtr backend_;
 
+#ifdef ESPHOME_OTA_UDP_LWIP
+  union {
+    uint32_t client_connect_time_{0};  // handshake states
+    UdpLink *udp_;                     // data phase (never both), points into handle_data_'s stack
+  };
+  static_assert(sizeof(UdpLink *) == sizeof(uint32_t), "clearing client_connect_time_ must clear udp_");
+#else
   uint32_t client_connect_time_{0};
+#ifdef ESPHOME_OTA_UDP
+  UdpLink *udp_{nullptr};  // data phase, points into handle_data_'s stack
+#endif
+#endif
+#if defined(ESPHOME_OTA_UDP) && defined(USE_OTA_ENCRYPTION)
+  static_assert(sizeof(UdpLink::reply) <= NOISE_MAX_RESPONSE, "the token goes out through noise_write_");
+#endif
   static constexpr size_t HANDSHAKE_BUF_SIZE = 5;
   // Buffer size for OTA data transfer. The upload client derives its maximum
   // encrypted frame plaintext from this (espota2.NOISE_MAX_PLAINTEXT is this

esphome/components/esphome/ota/ota_esphome_noise.cpp

@@ -266,13 +266,15 @@ ssize_t ESPHomeOTAComponent::noise_read_data_(uint8_t *buf, size_t capacity) {
   return this->noise_read_frame_blocking_(buf, noise::MAC_SIZE + 1, max_ciphertext);
 }
 
-/// Blocking write of one response byte as an encrypted frame.
-bool ESPHomeOTAComponent::noise_write_byte_(uint8_t byte) {
-  uint8_t frame[noise::FRAME_HEADER_SIZE + 1 + noise::MAC_SIZE];
-  frame[noise::FRAME_HEADER_SIZE] = byte;
+/// Blocking write of a short response (at most NOISE_MAX_RESPONSE bytes) as one encrypted frame.
+bool ESPHomeOTAComponent::noise_write_(const uint8_t *data, size_t len) {
+  uint8_t frame[noise::FRAME_HEADER_SIZE + NOISE_MAX_RESPONSE + noise::MAC_SIZE];
+  if (len > NOISE_MAX_RESPONSE)
+    return false;
+  memcpy(frame + noise::FRAME_HEADER_SIZE, data, len);
   NoiseBuffer mbuf;
   noise_buffer_init(mbuf);
-  noise_buffer_set_inout(mbuf, frame + noise::FRAME_HEADER_SIZE, 1, 1 + noise::MAC_SIZE);
+  noise_buffer_set_inout(mbuf, frame + noise::FRAME_HEADER_SIZE, len, len + noise::MAC_SIZE);
   int err = noise_cipherstate_encrypt(this->noise_->send_cipher, &mbuf);
   if (err != 0) {
     ESP_LOGW(TAG, "Encrypt: %d", err);

esphome/components/esphome/ota/ota_esphome_udp.cpp

@@ -0,0 +1,375 @@
+#include "ota_esphome.h"
+#if defined(USE_OTA) && defined(ESPHOME_OTA_UDP)
+
+#include "esphome/core/application.h"
+#include "esphome/core/hal.h"
+#include "esphome/core/helpers.h"
+#include "esphome/core/log.h"
+
+#ifdef ESPHOME_OTA_UDP_LWIP
+#include "lwip/inet.h"
+#ifdef USE_LIBRETINY
+#include "lwip/tcpip.h"
+#endif
+#endif
+
+#include <algorithm>
+#include <cstring>
+#include <memory>
+
+/*
+ * UDP data phase: the same byte stream as TCP, cut into numbered datagrams.
+ *   client -> device  [type | COMMITTED][token u32][seq u16][device bytes received u16][payload]
+ *   device -> client  [ACK][token u32][received u16][seq that prompted it u16][stream offset u16][window u8][our bytes]
+ * The lwIP callback acks each datagram on arrival, so flash writes never look like loss, and keeps messages in its
+ * UDP_WINDOW slots until the OTA loop reads them. "window": free slots in the high nibble, and in the low nibble bit k
+ * set when received + k is already here (like TCP SACK).
+ */
+
+namespace esphome {
+
+ESPHOME_LOG_TAG(TAG, "esphome.ota");
+
+// PING (0x02) and PROBE (0x03) only prompt an ACK
+static constexpr uint8_t UDP_MSG_DATA = 0x01;
+static constexpr uint8_t UDP_FLAG_COMMITTED = 0x80;
+static constexpr uint8_t UDP_MSG_ACK = 0x40;
+static constexpr size_t UDP_HEADER_SIZE = 9;
+static constexpr size_t UDP_ACK_HEADER_SIZE = 12;
+// The client's 60 s wait for the token over a lossy TCP link, plus its 15 s probe
+static constexpr uint32_t UDP_COMMIT_TIMEOUT = 80000;
+// Like the TCP socket's SO_RCVTIMEO in handle_data_, so readall_ treats both alike
+static constexpr uint32_t UDP_READ_WAIT_MS = 2000;
+// How often the commit wait looks at the UDP link between reads of the TCP fallback byte
+static constexpr uint32_t UDP_COMMIT_POLL_MS = 100;
+static constexpr uint32_t UDP_LINGER_MS = 3000;
+
+#ifdef ESPHOME_OTA_UDP_LWIP
+#ifdef USE_LIBRETINY
+// LibreTiny's LwIPLock is a no-op; take the core lock directly (the header drops UDP without one)
+struct UdpLock {
+  UdpLock() { LOCK_TCPIP_CORE(); }
+  ~UdpLock() { UNLOCK_TCPIP_CORE(); }
+};
+#else
+using UdpLock = LwIPLock;
+#endif
+static uint16_t slot_len(struct pbuf *p) { return p->tot_len; }
+static uint16_t slot_copy(struct pbuf *p, void *dst, uint16_t len, uint16_t offset) {
+  return pbuf_copy_partial(p, dst, len, offset);
+}
+static void slot_free(struct pbuf *p) { pbuf_free(p); }
+#else
+// The host has no lwIP thread: the OTA loop drains the socket itself. Empty bodies, not = default, so clang-tidy
+// does not flag the guards as unused
+struct UdpLock {
+  UdpLock() {}
+  ~UdpLock() {}
+};
+struct OtaUdpDatagram {
+  uint16_t len;
+  uint8_t data[1500];
+};
+static uint16_t slot_len(OtaUdpDatagram *p) { return p->len; }
+static uint16_t slot_copy(OtaUdpDatagram *p, void *dst, uint16_t len, uint16_t offset) {
+  const uint16_t n = std::min<uint16_t>(len, p->len - offset);
+  memcpy(dst, p->data + offset, n);
+  return n;
+}
+static void slot_free(OtaUdpDatagram *p) { delete p; }
+#endif
+
+static void put_u16(uint8_t *at, uint16_t value) {
+  at[0] = value >> 8;
+  at[1] = value & 0xFF;
+}
+
+ESPHomeOTAComponent::UdpLink::~UdpLink() { this->release(); }
+
+void ESPHomeOTAComponent::UdpLink::release() {
+  UdpLock lock;
+#ifdef ESPHOME_OTA_UDP_LWIP
+  if (this->pcb != nullptr)
+    udp_remove(this->pcb);
+  this->pcb = nullptr;
+#else
+  this->sock = nullptr;
+#endif
+  for (auto &p : this->queue) {
+    if (p != nullptr)
+      slot_free(p);
+    p = nullptr;
+  }
+}
+
+void ESPHomeOTAComponent::udp_send_ack(UdpLink &link, uint16_t prompted_by) {
+  const uint16_t len = UDP_ACK_HEADER_SIZE + link.out_len;
+#ifdef ESPHOME_OTA_UDP_LWIP
+  struct pbuf *p = pbuf_alloc(PBUF_TRANSPORT, len, PBUF_RAM);
+  if (p == nullptr)
+    return;  // Like a lost ACK: the client asks again
+  auto *ack = static_cast<uint8_t *>(p->payload);
+#else
+  uint8_t ack[UDP_ACK_HEADER_SIZE + UDP_OUT_SIZE];
+#endif
+  ack[0] = UDP_MSG_ACK;
+  memcpy(ack + 1, link.reply, 4);
+  put_u16(ack + 5, link.received);
+  put_u16(ack + 7, prompted_by);
+  put_u16(ack + 9, link.out_base);
+  uint8_t window = (UDP_WINDOW - static_cast<uint16_t>(link.received - link.consumed)) << 4;
+  for (uint8_t k = 1; k < UDP_WINDOW; k++) {
+    if (link.queue[(link.received + k) % UDP_WINDOW] != nullptr)
+      window |= 1 << k;
+  }
+  ack[11] = window;
+  memcpy(ack + UDP_ACK_HEADER_SIZE, link.out, link.out_len);
+  // Best effort: a lost ACK is repeated for the client's next datagram
+#ifdef ESPHOME_OTA_UDP_LWIP
+  udp_sendto(link.pcb, p, &link.peer_ip, link.peer_port);
+  pbuf_free(p);
+#else
+  struct sockaddr_storage to = link.peer;
+  if (to.ss_family == AF_INET6) {
+    reinterpret_cast<sockaddr_in6 *>(&to)->sin6_port = htons(link.peer_port);
+  } else {
+    reinterpret_cast<sockaddr_in *>(&to)->sin_port = htons(link.peer_port);
+  }
+  link.sock->sendto(ack, len, 0, reinterpret_cast<const struct sockaddr *>(&to), link.peer_len);
+#endif
+}
+
+void ESPHomeOTAComponent::udp_on_datagram(UdpLink &link, UdpSlot p, uint16_t port) {
+  uint8_t hdr[UDP_HEADER_SIZE];
+  if (slot_copy(p, hdr, UDP_HEADER_SIZE, 0) != UDP_HEADER_SIZE || memcmp(hdr + 1, link.reply, 4) != 0) {
+    slot_free(p);
+    return;
+  }
+  // The token proves the sender; answer the port it arrived from, which NAT may have rewritten
+  link.peer_port = port;
+  link.committed |= (hdr[0] & UDP_FLAG_COMMITTED) != 0;
+  // Drop the responses the client already has
+  const uint16_t acked = encode_uint16(hdr[7], hdr[8]) - link.out_base;
+  if (acked != 0 && acked <= link.out_len) {
+    link.out_len -= acked;
+    memmove(link.out, link.out + acked, link.out_len);
+    link.out_base += acked;
+  }
+  const uint16_t seq = encode_uint16(hdr[5], hdr[6]);
+  UdpSlot &slot = link.queue[seq % UDP_WINDOW];
+  if ((hdr[0] & ~UDP_FLAG_COMMITTED) == UDP_MSG_DATA && slot_len(p) > UDP_HEADER_SIZE &&
+      static_cast<uint16_t>(seq - link.consumed) < UDP_WINDOW && slot == nullptr) {
+    slot = p;
+    while (static_cast<uint16_t>(link.received - link.consumed) < UDP_WINDOW &&
+           link.queue[link.received % UDP_WINDOW] != nullptr)
+      link.received++;
+  } else {
+    slot_free(p);
+  }
+  udp_send_ack(link, seq);
+}
+
+#ifdef ESPHOME_OTA_UDP_LWIP
+// lwIP context: the tcpip thread (ESP32, LibreTiny), SYS (ESP8266) or an IRQ (RP2040)
+void ESPHomeOTAComponent::udp_recv_cb(void *arg, struct udp_pcb * /*pcb*/, struct pbuf *p, const ip_addr_t *addr,
+                                      uint16_t port) {
+  auto &link = *static_cast<UdpLink *>(arg);
+  if (!ip_addr_cmp(addr, &link.peer_ip)) {
+    pbuf_free(p);
+    return;
+  }
+  udp_on_datagram(link, p, port);
+}
+#else
+static bool same_ip(const struct sockaddr_storage &a, const struct sockaddr_storage &b) {
+  uint32_t ipv4_a, ipv4_b;
+  const bool v4_a = socket::sockaddr_to_ipv4(reinterpret_cast<const struct sockaddr *>(&a), &ipv4_a);
+  const bool v4_b = socket::sockaddr_to_ipv4(reinterpret_cast<const struct sockaddr *>(&b), &ipv4_b);
+  if (v4_a || v4_b)
+    return v4_a && v4_b && ipv4_a == ipv4_b;
+  return a.ss_family == AF_INET6 && b.ss_family == AF_INET6 &&
+         memcmp(&reinterpret_cast<const sockaddr_in6 *>(&a)->sin6_addr,
+                &reinterpret_cast<const sockaddr_in6 *>(&b)->sin6_addr, sizeof(in6_addr)) == 0;
+}
+
+void ESPHomeOTAComponent::udp_poll_() {
+  UdpLink &link = *this->udp_;
+  for (;;) {
+    uint8_t buf[sizeof(OtaUdpDatagram::data)];
+    struct sockaddr_storage from;
+    socklen_t from_len = sizeof(from);
+    const ssize_t len = link.sock->recvfrom(buf, sizeof(buf), reinterpret_cast<struct sockaddr *>(&from), &from_len);
+    if (len < 0)
+      return;
+    if (!same_ip(from, link.peer))
+      continue;
+    auto *datagram = new OtaUdpDatagram;  // NOLINT(cppcoreguidelines-owning-memory)
+    datagram->len = len;
+    memcpy(datagram->data, buf, len);
+    const uint16_t port = from.ss_family == AF_INET6 ? ntohs(reinterpret_cast<sockaddr_in6 *>(&from)->sin6_port)
+                                                     : ntohs(reinterpret_cast<sockaddr_in *>(&from)->sin_port);
+    udp_on_datagram(link, datagram, port);
+  }
+}
+#endif
+
+void ESPHomeOTAComponent::udp_open_(UdpLink &link) {
+  struct sockaddr_storage peer;
+  socklen_t peer_len = sizeof(peer);
+  if (this->client_->getpeername(reinterpret_cast<struct sockaddr *>(&peer), &peer_len) != 0)
+    return;
+  uint8_t token[4];
+  if (!random_bytes(token, sizeof(token)))
+    return;
+#ifdef ESPHOME_OTA_UDP_LWIP
+  uint32_t ipv4;
+  if (socket::sockaddr_to_ipv4(reinterpret_cast<const struct sockaddr *>(&peer), &ipv4)) {
+    ip_addr_set_ip4_u32(&link.peer_ip, ipv4);
+  } else {
+#if USE_NETWORK_IPV6 && LWIP_IPV6
+    inet6_addr_to_ip6addr(ip_2_ip6(&link.peer_ip), &reinterpret_cast<const sockaddr_in6 *>(&peer)->sin6_addr);
+    IP_SET_TYPE_VAL(link.peer_ip, IPADDR_TYPE_V6);
+#else
+    return;
+#endif
+  }
+  UdpLock lock;
+#if LWIP_IPV6
+  link.pcb = udp_new_ip_type(IPADDR_TYPE_ANY);
+#else
+  link.pcb = udp_new();
+#endif
+  // On failure the reply stays all zero and ~UdpLink removes the pcb
+  if (link.pcb == nullptr || udp_bind(link.pcb, IP_ANY_TYPE, this->port_) != ERR_OK)
+    return;
+  udp_recv(link.pcb, udp_recv_cb, &link);
+#else
+  link.peer = peer;
+  link.peer_len = peer_len;
+  // Same domain as the listener, so the peer address from getpeername fits
+  link.sock = socket::socket_ip(SOCK_DGRAM, IPPROTO_UDP);
+  struct sockaddr_storage local;
+  const socklen_t local_len =
+      socket::set_sockaddr_any(reinterpret_cast<struct sockaddr *>(&local), sizeof(local), this->port_);
+  if (link.sock == nullptr || local_len == 0 ||
+      link.sock->bind(reinterpret_cast<struct sockaddr *>(&local), local_len) != 0 ||
+      link.sock->setblocking(false) != 0)
+    return;
+#endif
+  // Below the error codes (0x80 and up) the client checks the first byte against, and never all zero
+  token[0] = (token[0] & 0x7F) | 0x01;
+  memcpy(link.reply, token, sizeof(token));
+}
+
+bool ESPHomeOTAComponent::udp_start_() {
+  // TCP now only carries a fallback byte; short read timeouts let the UDP link be checked in between
+  struct timeval tv {
+    0, UDP_COMMIT_POLL_MS * 1000
+  };
+  this->client_->setsockopt(SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
+  const uint32_t start = millis();
+  while (millis() - start < UDP_COMMIT_TIMEOUT) {
+    this->udp_poll_();
+    bool committed;
+    {
+      UdpLock lock;
+      committed = this->udp_->committed;
+    }
+    if (committed) {
+      // The client heard us over UDP, so TCP is done. Outside the lock: closing a socket waits on the tcpip thread
+      this->client_ = nullptr;
+      return true;
+    }
+    // A raw byte outside any Noise frame: the client heard nothing over UDP. Forging it only forces the ordinary
+    // TCP path, which dropping our datagrams forces anyway
+    uint8_t fallback;
+    const ssize_t read = this->client_->read(&fallback, 1);
+    if (read == 1) {
+      this->udp_->release();
+      this->udp_ = nullptr;
+      tv = {UDP_READ_WAIT_MS / 1000, 0};
+      this->client_->setsockopt(SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
+      return true;
+    }
+    if (read == 0) {
+      this->remote_closed_ = true;
+      return false;
+    }
+    if (!this->would_block_(errno))
+      return false;
+    App.feed_wdt();
+  }
+  return false;
+}
+
+ssize_t ESPHomeOTAComponent::udp_read_(uint8_t *buf, size_t len) {
+  UdpLink &link = *this->udp_;
+  const uint32_t start = millis();
+  size_t got = 0;
+  for (;;) {
+    this->udp_poll_();
+    UdpSlot p;
+    {
+      UdpLock lock;
+      p = link.consumed != link.received ? link.queue[link.consumed % UDP_WINDOW] : nullptr;
+    }
+    if (p == nullptr) {
+      if (got != 0)
+        return got;
+      if (millis() - start >= UDP_READ_WAIT_MS)
+        break;
+      this->yield_and_feed_watchdog_();
+      continue;
+    }
+    // The callback never touches the slot being read, so the copy runs without the lock
+    const uint16_t offset = UDP_HEADER_SIZE + link.read_offset;
+    const uint16_t n = slot_copy(p, buf + got, std::min<size_t>(len - got, slot_len(p) - offset), offset);
+    got += n;
+    link.read_offset += n;
+    if (UDP_HEADER_SIZE + link.read_offset >= slot_len(p)) {
+      UdpLock lock;
+      slot_free(p);
+      link.queue[link.consumed % UDP_WINDOW] = nullptr;
+      // A full window just opened: tell the client, which is holding back
+      const bool was_full = static_cast<uint16_t>(link.received - link.consumed) == UDP_WINDOW;
+      link.consumed++;
+      link.read_offset = 0;
+      if (was_full)
+        udp_send_ack(link, link.received - 1);
+    }
+    if (got == len)
+      return got;
+  }
+  errno = EWOULDBLOCK;
+  return -1;
+}
+
+bool ESPHomeOTAComponent::udp_send_(const uint8_t *data, size_t len) {
+  UdpLink &link = *this->udp_;
+  UdpLock lock;
+  if (link.out_len + len > UDP_OUT_SIZE)
+    return false;
+  memcpy(link.out + link.out_len, data, len);
+  link.out_len += len;
+  // Prompted by nothing new, so the client takes no round trip sample from it
+  udp_send_ack(link, link.received - 1);
+  return true;
+}
+
+void ESPHomeOTAComponent::udp_linger_() {
+  if (this->udp_ == nullptr)
+    return;
+  const uint32_t start = millis();
+  while (millis() - start < UDP_LINGER_MS) {
+    this->udp_poll_();
+    {
+      UdpLock lock;
+      if (this->udp_->out_len == 0)
+        return;
+    }
+    this->yield_and_feed_watchdog_();
+  }
+}
+
+}  // namespace esphome
+#endif  // USE_OTA && ESPHOME_OTA_UDP

esphome/espota2.py

@@ -70,11 +70,13 @@ CLIENT_FEATURE_SUPPORTS_SHA256_AUTH = 0x02
 CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL = 0x04
 CLIENT_FEATURE_SUPPORTS_NOISE = 0x08
 CLIENT_FEATURE_SUPPORTS_DEFLATE = 0x10
+CLIENT_FEATURE_SUPPORTS_UDP = 0x20
 SERVER_FEATURE_SUPPORTS_COMPRESSION = 0x01
 SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS = 0x02
 SERVER_FEATURE_SUPPORTS_NOISE = 0x04
 # Binding once offered: the device then expects the image size and a deflate stream
 SERVER_FEATURE_SUPPORTS_DEFLATE = 0x08
+SERVER_FEATURE_SUPPORTS_UDP = 0x10
 
 # Wire constant: the deflate bit promises a 4 KB window (OTA_INFLATE_WINDOW_SIZE)
 DEFLATE_WINDOW_BITS = 12
@@ -103,17 +105,64 @@ COMPRESS_LEVEL = 9
 UPLOAD_BUFFER_SIZE = UPLOAD_BLOCK_SIZE * 8
 
 # Flaky Wi-Fi links often drop the first OTA attempt, and the device may need time
-# to clean up a half-open connection (its handshake watchdog runs at 20s) before it
+# to clean up a half-open connection (its handshake watchdog runs at 60s) before it
 # accepts a new one, so wait between attempts instead of failing the upload outright.
 # Every resolved address is tried once, and this many extra attempts are shared
 # across the addresses on top of that.
-EXTRA_UPLOAD_ATTEMPTS = 2
+EXTRA_UPLOAD_ATTEMPTS = 4
+# Once a connection has reached the device, a lossy link (not a dead device) is
+# the likely cause, so keep trying like an API reconnect would
+EXTRA_UPLOAD_ATTEMPTS_REACHED = 8
 UPLOAD_RETRY_DELAY = 5.0
+# A SYN getting through is the cheap step; an unreachable address still fails fast
+CONNECT_TIMEOUT = 20.0
+# Handshake waits; as patient as an API connect on a lossy link, and matching the
+# device's OTA_SOCKET_TIMEOUT_HANDSHAKE
+SETUP_TIMEOUT = 60.0
 # Data phase timeout; must stay longer than the device's OTA_SOCKET_TIMEOUT_DATA
 # (105 s) so a stalled session is gone before a retry, and long enough for lwIP
 # to get a lost chunk ack through after the retransmit run seen in practice
 DATA_PHASE_TIMEOUT = 160.0
 
+# UDP data phase; wire format in ota_esphome_udp.cpp
+UDP_MSG_DATA = 0x01
+UDP_MSG_PING = 0x02
+UDP_MSG_PROBE = 0x03
+UDP_FLAG_COMMITTED = 0x80
+UDP_MSG_ACK = 0x40
+UDP_HEADER_SIZE = 9
+UDP_ACK_HEADER_SIZE = 12
+# Stream bytes per datagram, within one Ethernet frame
+UDP_MAX_PAYLOAD = 1400
+# Messages the device queues unconsumed (UDP_WINDOW in ota_esphome.h); as many
+# as TCP's window holds, so UDP costs the device no more memory than TCP
+UDP_WINDOW = 4
+# Probe this long before falling back to TCP, enough to ride out a radio busy
+# with BLE; the device waits 80 s
+UDP_PROBE_TIMEOUT = 15.0
+UDP_PROBE_INTERVAL = 0.1
+# Resend timeout bounds around srtt + 4 * rttvar; random loss never grows it
+UDP_RESEND_INITIAL = 0.5
+UDP_RESEND_MIN = 0.2
+UDP_RESEND_MAX = 2.0
+# Acks still asking for the head after later messages arrived; resend it once
+UDP_DUP_ACKS = 2
+# Progress acks in a row with no round trip to time: the link is slower than
+# the resend timer, so back it off (random loss still yields samples)
+UDP_UNTIMED_ACKS = 8
+# Refusals in a row with no ACK between: the device's UDP port is closed
+UDP_REFUSALS = 3
+UDP_FALLBACK = b"\x00"
+# The UDP format this client speaks; devices decline values they do not know
+UDP_REQUEST = b"\x01"
+# A plain handshake exchange slower than this hid a retransmit (the device's
+# lwIP waits about a second before resending); UDP is used from then on
+UDP_SLOW_EXCHANGE = 0.75
+# getsockopt(IPPROTO_TCP) option, size, offset and format of the retransmit
+# count: tcp_info.tcpi_total_retrans (Linux), tcp_connection_info.
+# tcpi_txretransmitbytes (macOS)
+_TCP_RETRANS_OPTION = {"linux": (11, 104, 100, "I"), "darwin": (0x106, 80, 72, "Q")}
+
 _LOGGER = logging.getLogger(__name__)
 
 # Authentication method lookup table: response -> (hash_func, nonce_size, name)
@@ -511,6 +560,10 @@ class NoiseSocketWrapper:
         self._decrypt = None
         self._buffer = b""
 
+    def use_transport(self, transport: UdpChannel) -> None:
+        """Continue the session over another byte stream."""
+        self._sock = transport
+
     # Only harmless socket controls pass through; byte-moving methods are
     # deliberately absent so plaintext cannot leak past the transport.
     def settimeout(self, timeout: float | None) -> None:
@@ -613,19 +666,316 @@ class NoiseSocketWrapper:
         return data
 
 
+class _UdpMessage:
+    """One datagram of the UDP stream and its send history."""
+
+    __slots__ = ("held", "last_send", "payload", "resend", "sends", "seq")
+
+    def __init__(self, seq: int, payload: bytes) -> None:
+        self.seq = seq
+        self.payload = payload
+        self.last_send = 0.0
+        self.sends = 0
+        self.resend = False  # duplicate acks asked for it
+        self.held = False  # the device keeps it, waiting for an earlier one
+
+
+class UdpChannel:
+    """Byte stream over numbered UDP datagrams, standing in for the socket.
+
+    The device acks every datagram on arrival and says how many it has room
+    for; each is resent after TCP's srtt + 4 * rttvar, so loss slows the
+    transfer instead of ending it. It backs off only while acks come back
+    untimed, i.e. the link is slower than the timer, never on random loss.
+    """
+
+    def __init__(self, sock: socket.socket, token: bytes) -> None:
+        self._sock = sock
+        self._token = token
+        self._timeout = DATA_PHASE_TIMEOUT
+        self._next_seq = 0
+        self._queue: list[_UdpMessage] = []  # not yet received, oldest first
+        self._rx = bytearray()
+        self._device_bytes = 0
+        self._committed = False
+        self._srtt: float | None = None
+        self._rttvar = 0.0
+        # Doubles while acks come back untimed, so a link slower than the
+        # timer still yields messages sent once to time (Karn, RFC 6298)
+        self._backoff = 1
+        self._untimed = 0
+        self._dup_acks = 0
+        self._free = UDP_WINDOW  # the device's empty slots, as of its last ACK
+        self._last_ping = 0.0
+        self._last_progress = time.monotonic()
+        self._sends = 0
+        # Last socket error, e.g. refused once the device's socket is gone
+        self._last_error: OSError | None = None
+        self._refusals = 0
+
+    def settimeout(self, timeout: float) -> None:
+        self._timeout = timeout
+
+    def setsockopt(self, level: int, optname: int, value: int) -> None:
+        """TCP tuning does not apply."""
+
+    def close(self) -> None:
+        _LOGGER.debug(
+            "UDP: %d messages in %d datagrams, smoothed RTT %.0f ms",
+            self._next_seq,
+            self._sends,
+            (self._srtt or 0) * 1000,
+        )
+        self._sock.close()
+
+    def probe(self) -> bool:
+        """True once the device answers over UDP within UDP_PROBE_TIMEOUT."""
+        deadline = time.monotonic() + UDP_PROBE_TIMEOUT
+        try:
+            while (now := time.monotonic()) < deadline:
+                self._send(UDP_MSG_PROBE, 0, b"")
+                self._receive(min(UDP_PROBE_INTERVAL, deadline - now))
+                if self._committed:
+                    return True
+        except OTANetworkError:
+            pass  # rejected, e.g. by a firewall
+        return False
+
+    def sendall(self, data: bytes) -> None:
+        for offset in range(0, len(data), UDP_MAX_PAYLOAD):
+            self._queue.append(
+                _UdpMessage(self._next_seq, data[offset : offset + UDP_MAX_PAYLOAD])
+            )
+            self._next_seq = (self._next_seq + 1) & 0xFFFF
+        # Return once everything is at least in flight, or the device spoke (an error)
+        while len(self._queue) > UDP_WINDOW and not self._rx:
+            self._pump()
+
+    def pending(self) -> bool:
+        """Whether the device sent bytes nobody has read yet."""
+        return bool(self._rx)
+
+    def recv(self, amount: int) -> bytes:
+        while not self._rx:
+            self._pump()
+        data = bytes(self._rx[:amount])
+        del self._rx[:amount]
+        return data
+
+    def flush(self, timeout: float) -> None:
+        """Best effort: wait until the device has every message."""
+        deadline = time.monotonic() + timeout
+        with contextlib.suppress(OTANetworkError):
+            while self._queue and time.monotonic() < deadline:
+                self._pump()
+
+    def _rto(self) -> float:
+        if self._srtt is None:
+            rto = UDP_RESEND_INITIAL
+        else:
+            rto = max(self._srtt + 4 * self._rttvar, UDP_RESEND_MIN)
+        return min(rto * self._backoff, UDP_RESEND_MAX)
+
+    def _send(self, msg_type: int, seq: int, payload: bytes) -> None:
+        flags = UDP_FLAG_COMMITTED if self._committed else 0
+        header = (
+            bytes([msg_type | flags])
+            + self._token
+            + seq.to_bytes(2, "big")
+            + (self._device_bytes & 0xFFFF).to_bytes(2, "big")
+        )
+        self._sends += 1
+        try:
+            self._sock.send(header + payload)
+        except OSError as err:
+            # A full buffer or ICMP error is just more loss, but worth reporting
+            self._socket_error(err)
+
+    def _pump(self) -> None:
+        now = time.monotonic()
+        if now - self._last_progress > self._timeout:
+            reason = f" (last error: {self._last_error})" if self._last_error else ""
+            raise OTANetworkError(
+                f"no progress over UDP for {self._timeout:.0f}s{reason}"
+            )
+        rto = self._rto()
+        wake = now + rto
+        # At least the head, which doubles as the probe while the device is busy
+        in_flight = self._queue[: max(self._free, 1)]
+        for message in in_flight:
+            if message.held:
+                continue
+            if message.sends == 0 or message.resend or now - message.last_send >= rto:
+                message.last_send = now
+                message.sends += 1
+                message.resend = False
+                self._send(UDP_MSG_DATA, message.seq, message.payload)
+            wake = min(wake, message.last_send + rto)
+        if all(message.held for message in in_flight):
+            # Nothing to send: ask for the device's responses
+            if now - self._last_ping >= rto:
+                self._last_ping = now
+                self._send(UDP_MSG_PING, 0, b"")
+            wake = self._last_ping + rto
+        self._receive(wake - now)
+
+    def _receive(self, wait: float) -> None:
+        import select
+
+        # Wait for the first datagram, then drain whatever else arrived
+        while select.select([self._sock], [], [], max(wait, 0))[0]:
+            wait = 0
+            try:
+                data = self._sock.recv(2048)
+            except OSError as err:
+                self._socket_error(err)
+                continue
+            self._handle_ack(data)
+
+    def _socket_error(self, err: OSError) -> None:
+        if self._last_error is None or err.errno != self._last_error.errno:
+            _LOGGER.debug("UDP socket error: %s", err)
+        self._last_error = err
+        # Windows reports the ICMP port unreachable as a reset
+        if isinstance(err, (ConnectionRefusedError, ConnectionResetError)):
+            self._refusals += 1
+            if self._refusals >= UDP_REFUSALS:
+                raise OTANetworkError(f"device closed its UDP port: {err}")
+
+    def _sample_rtt(self, sample: float) -> None:
+        # RFC 6298 smoothing
+        self._backoff = 1
+        self._untimed = 0
+        if self._srtt is None:
+            self._srtt, self._rttvar = sample, sample / 2
+        else:
+            self._rttvar = 0.75 * self._rttvar + 0.25 * abs(self._srtt - sample)
+            self._srtt = 0.875 * self._srtt + 0.125 * sample
+
+    def _handle_ack(self, data: bytes) -> None:
+        if (
+            len(data) < UDP_ACK_HEADER_SIZE
+            or data[0] != UDP_MSG_ACK
+            or data[1:5] != self._token
+        ):
+            return
+        now = time.monotonic()
+        self._committed = True
+        self._refusals = 0
+        received = int.from_bytes(data[5:7], "big")
+        prompted = int.from_bytes(data[7:9], "big")
+        self._free = data[11] >> 4
+        if self._queue:
+            head = self._queue[0].seq
+            index = (prompted - head) & 0xFFFF
+            # This datagram's own round trip; only from messages sent once (Karn)
+            timed = index < len(self._queue) and self._queue[index].sends == 1
+            if timed:
+                self._sample_rtt(now - self._queue[index].last_send)
+            done = (received - head) & 0xFFFF
+            if 0 < done <= len(self._queue):
+                del self._queue[:done]
+                self._last_progress = now
+                self._dup_acks = 0
+                if not timed:
+                    self._untimed += 1
+                    if self._untimed >= UDP_UNTIMED_ACKS:
+                        self._untimed = 0
+                        if self._rto() < UDP_RESEND_MAX:
+                            self._backoff *= 2
+            elif done == 0 and self._queue[0].sends:
+                # A later message arrived while the head is missing; resend it once
+                self._dup_acks += 1
+                if self._dup_acks == UDP_DUP_ACKS:
+                    self._queue[0].resend = True
+            # Rebuilt from every ACK, so a reordered one cannot leave a stale mark
+            for k, message in enumerate(self._queue[:UDP_WINDOW]):
+                message.held = (received + k - message.seq) & 0xFFFF == 0 and bool(
+                    data[11] & (1 << k)
+                )
+        skip = (self._device_bytes - int.from_bytes(data[9:11], "big")) & 0xFFFF
+        fresh = data[UDP_ACK_HEADER_SIZE + skip :]
+        if skip <= len(data) - UDP_ACK_HEADER_SIZE and fresh:
+            self._rx += fresh
+            self._device_bytes += len(fresh)
+            self._last_progress = now
+
+
+def _tcp_retransmits(sock: socket.socket) -> int | None:
+    """Segments this side retransmitted, where the OS reports it."""
+    import struct
+    import sys
+
+    if (option := _TCP_RETRANS_OPTION.get(sys.platform)) is None:
+        return None
+    optname, size, offset, fmt = option
+    try:
+        info = sock.getsockopt(socket.IPPROTO_TCP, optname, size)
+    except (OSError, TypeError):
+        return None
+    if len(info) < offset + struct.calcsize(fmt):
+        return None
+    return struct.unpack_from(fmt, info, offset)[0]
+
+
+def _link_lossy(
+    tcp_sock: socket.socket, slowest_exchange: float, prefer_udp: bool
+) -> bool:
+    """Whether the handshake showed loss, so the data phase should use UDP."""
+    return (
+        prefer_udp
+        or slowest_exchange > UDP_SLOW_EXCHANGE
+        or bool(_tcp_retransmits(tcp_sock))
+    )
+
+
+def _start_udp(
+    sock: socket.socket | NoiseSocketWrapper, tcp_sock: socket.socket, lossy: bool
+) -> UdpChannel | None:
+    """Move the data phase to UDP; None keeps it on TCP."""
+    # A clean link asks to stay on TCP
+    send_check(sock, UDP_REQUEST if lossy else UDP_FALLBACK, "UDP request")
+    token = receive_exactly(sock, 4, "UDP token", None, decode=False)
+    if not lossy:
+        _LOGGER.debug("Link looks clean, staying on TCP")
+        return None
+    if token == bytes(4):
+        _LOGGER.debug("Device declined UDP")
+        return None
+    udp_sock = socket.socket(tcp_sock.family, socket.SOCK_DGRAM)
+    try:
+        udp_sock.connect(tcp_sock.getpeername())
+    except OSError as err:
+        udp_sock.close()
+        _LOGGER.info("UDP unavailable (%s), using TCP", err)
+    else:
+        channel = UdpChannel(udp_sock, token)
+        if channel.probe():
+            _LOGGER.info("Using UDP")
+            return channel
+        channel.close()
+        _LOGGER.info("UDP blocked, using TCP")
+    # Raw, outside any Noise frame: forging it only forces the TCP path
+    send_check(tcp_sock, UDP_FALLBACK, "UDP fallback")
+    return None
+
+
 def _negotiate_session(
     sock: socket.socket,
     noise_psk: str | None,
     plaintext_fallback: bool = False,
     allow_plaintext_upload: bool = False,
-) -> tuple[socket.socket | NoiseSocketWrapper, int, int, bool]:
+) -> tuple[socket.socket | NoiseSocketWrapper, int, int, bool, float]:
     """Negotiate through the Noise handshake; returns the socket to go on with
-    (wrapped once encrypted), version, features, extended."""
+    (wrapped once encrypted), version, features, extended and the slowest
+    plaintext exchange in seconds."""
     # Enable nodelay, we need it for phase 1
     sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+    exchange_start = time.monotonic()
     send_check(sock, MAGIC_BYTES, "magic bytes")
 
     _, version = receive_exactly(sock, 2, "version", RESPONSE_OK)
+    slowest_exchange = time.monotonic() - exchange_start
     _LOGGER.debug("Device supports OTA version %s", version)
     supported_versions = (OTA_VERSION_1_0, OTA_VERSION_2_0)
     if version not in supported_versions:
@@ -639,9 +989,11 @@ def _negotiate_session(
         | CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
         | CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
         | CLIENT_FEATURE_SUPPORTS_DEFLATE
+        | CLIENT_FEATURE_SUPPORTS_UDP
     )
     if noise_psk:
         features_to_send |= CLIENT_FEATURE_SUPPORTS_NOISE
+    exchange_start = time.monotonic()
     send_check(sock, features_to_send, "features")
     features = receive_exactly(
         sock,
@@ -649,6 +1001,7 @@ def _negotiate_session(
         "features",
         None,  # Accept any response
     )[0]
+    slowest_exchange = max(slowest_exchange, time.monotonic() - exchange_start)
 
     extended_proto = False
     if features == RESPONSE_FEATURE_FLAGS:
@@ -705,7 +1058,7 @@ def _negotiate_session(
         _LOGGER.debug("Encrypted connection established")
         if allow_plaintext_upload:
             _LOGGER.warning(ALLOW_PLAINTEXT_UPLOAD_REMOVE_WARNING)
-    return sock, version, features, extended_proto
+    return sock, version, features, extended_proto, slowest_exchange
 
 
 def perform_ota(
@@ -718,6 +1071,7 @@ def perform_ota(
     plaintext_fallback: bool = False,
     allow_plaintext_upload: bool = False,
     on_upload_start: Callable[[], None] | None = None,
+    prefer_udp: bool = False,
 ) -> None:
     # Validate up front; an out-of-range value would only surface as a
     # ValueError deep inside send_check, bypassing OTAError handling
@@ -740,7 +1094,9 @@ def perform_ota(
     file_size = len(file_contents)
     _LOGGER.info("Uploading %s (%s bytes)", filename, file_size)
 
-    sock, version, features, extended_proto = _negotiate_session(
+    # sock may become a Noise wrapper or a UDP channel; UDP needs the peer
+    tcp_sock = sock
+    sock, version, features, extended_proto, slowest_exchange = _negotiate_session(
         sock, noise_psk, plaintext_fallback, allow_plaintext_upload
     )
     _LOGGER.info("Connection established; device supports OTA version %s", version)
@@ -832,117 +1188,153 @@ def perform_ota(
     if on_upload_start is not None:
         on_upload_start()  # the device may commit anything sent from here
 
+    udp: UdpChannel | None = None
+    # Still under SETUP_TIMEOUT, which the device's commit wait is sized for
+    if extended_proto and features & SERVER_FEATURE_SUPPORTS_UDP:
+        udp = _start_udp(
+            sock, tcp_sock, _link_lossy(tcp_sock, slowest_exchange, prefer_udp)
+        )
+        if udp is not None:
+            # The device drops TCP once it hears the committed channel
+            if isinstance(sock, NoiseSocketWrapper):
+                sock.use_transport(udp)
+            else:
+                sock = udp
+
     sock.settimeout(DATA_PHASE_TIMEOUT)
-
-    if extended_proto:
-        send_check(sock, ota_type, "ota type")
-
-    upload_size = len(upload_contents)
-    # The device erases flash between receiving the size and acking the
-    # prepare, so this window shows the erase cost (near zero when the
-    # device erases lazily during the upload)
-    prepare_start = time.perf_counter()
-    send_check(sock, upload_size.to_bytes(SIZE_FIELD_BYTES, "big"), "binary size")
-    if deflate:
-        # Own frame: an encrypted session carries one field per frame
-        send_check(sock, file_size.to_bytes(SIZE_FIELD_BYTES, "big"), "image size")
-    receive_exactly(sock, 1, "update prepare result", RESPONSE_UPDATE_PREPARE_OK)
-    prepare_duration = time.perf_counter() - prepare_start
-    _LOGGER.info("Preparing for upload took %.2f seconds", prepare_duration)
-
-    # The device hashes what it writes: the inflated image, else the received bytes
-    upload_md5 = hashlib.md5(file_contents if deflate else upload_contents).hexdigest()
-    _LOGGER.debug("MD5 of upload is %s", upload_md5)
-
-    send_check(sock, upload_md5, "file checksum")
-    receive_exactly(sock, 1, "file checksum result", RESPONSE_BIN_MD5_OK)
-
-    # Disable nodelay for transfer
-    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 0)
-    # Limit send buffer (usually around 100kB) in order to have progress bar
-    # show the actual progress
-
-    sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, UPLOAD_BUFFER_SIZE)
-    start_time = time.perf_counter()
-
-    offset = 0
-    progress = ProgressBar("Uploading")
     try:
-        while True:
-            chunk = upload_contents[offset : offset + UPLOAD_BLOCK_SIZE]
-            if not chunk:
-                break
-            offset += len(chunk)
+        if extended_proto:
+            send_check(sock, ota_type, "ota type")
 
-            try:
-                sock.sendall(chunk)
-            except OSError as err:
-                # A send failure can hide an error byte the device reported
-                # just before dropping the connection; surface that as the
-                # real, non-retryable cause when it is available
-                try:
-                    sock.settimeout(1.0)
-                    check_error(recv_decode(sock, 1), None)
-                except (OSError, OTANetworkError) as probe_err:
-                    _LOGGER.debug(
-                        "No device error behind the send failure: %s", probe_err
-                    )
-                raise OTANetworkError(f"sending data: {err}") from err
+        upload_size = len(upload_contents)
+        # The device erases flash between receiving the size and acking the
+        # prepare, so this window shows the erase cost (near zero when the
+        # device erases lazily during the upload)
+        prepare_start = time.perf_counter()
+        send_check(sock, upload_size.to_bytes(SIZE_FIELD_BYTES, "big"), "binary size")
+        if deflate:
+            # Own frame: an encrypted session carries one field per frame
+            send_check(sock, file_size.to_bytes(SIZE_FIELD_BYTES, "big"), "image size")
+        receive_exactly(sock, 1, "update prepare result", RESPONSE_UPDATE_PREPARE_OK)
+        prepare_duration = time.perf_counter() - prepare_start
+        _LOGGER.info("Preparing for upload took %.2f seconds", prepare_duration)
+
+        # The device hashes what it writes: the inflated image, else the received bytes
+        upload_md5 = hashlib.md5(
+            file_contents if deflate else upload_contents
+        ).hexdigest()
+        _LOGGER.debug("MD5 of upload is %s", upload_md5)
+
+        send_check(sock, upload_md5, "file checksum")
+        receive_exactly(sock, 1, "file checksum result", RESPONSE_BIN_MD5_OK)
+
+        # Disable nodelay for transfer
+        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 0)
+        # Limit send buffer (usually around 100kB) in order to have progress bar
+        # show the actual progress
+
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, UPLOAD_BUFFER_SIZE)
+        start_time = time.perf_counter()
+
+        offset = 0
+        progress = ProgressBar("Uploading")
+        try:
+            while True:
+                chunk = upload_contents[offset : offset + UPLOAD_BLOCK_SIZE]
+                if not chunk:
+                    break
+                offset += len(chunk)
 
-            if version >= OTA_VERSION_2_0:
                 try:
-                    receive_exactly(sock, 1, "chunk result", RESPONSE_CHUNK_OK)
+                    sock.sendall(chunk)
+                except OSError as err:
+                    # A send failure can hide an error byte the device reported
+                    # just before dropping the connection; surface that as the
+                    # real, non-retryable cause when it is available
+                    try:
+                        sock.settimeout(1.0)
+                        check_error(recv_decode(sock, 1), None)
+                    except (OSError, OTANetworkError) as probe_err:
+                        _LOGGER.debug(
+                            "No device error behind the send failure: %s", probe_err
+                        )
+                    raise OTANetworkError(f"sending data: {err}") from err
                 except OTANetworkError as err:
+                    if udp is None:
+                        raise
+                    # Over UDP: surface a device error code if one arrived, and
+                    # treat a failure on the last chunk like a lost final ack
+                    if udp.pending():
+                        receive_exactly(sock, 1, "data", [])
                     if offset < upload_size:
                         raise
-                    # The device already had the complete image when this ack
-                    # was lost, so it may be committing; do not retry
                     raise _committed_error(err) from err
 
-            progress.update(offset / upload_size)
-    except OTAError:
-        # Terminate the progress bar line before the error is logged
+                # The UDP channel acks every message itself, so anything the device
+                # sends mid upload is an error code
+                if udp is not None and udp.pending() and offset < upload_size:
+                    receive_exactly(sock, 1, "data", [])
+                if version >= OTA_VERSION_2_0 and udp is None:
+                    try:
+                        receive_exactly(sock, 1, "chunk result", RESPONSE_CHUNK_OK)
+                    except OTANetworkError as err:
+                        if offset < upload_size:
+                            raise
+                        # The device already had the complete image when this ack
+                        # was lost, so it may be committing; do not retry
+                        raise _committed_error(err) from err
+
+                progress.update(offset / upload_size)
+        except OTAError:
+            # Terminate the progress bar line before the error is logged
+            progress.done()
+            raise
         progress.done()
-        raise
-    progress.done()
 
-    # Enable nodelay for last checks
-    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
-    duration = time.perf_counter() - start_time
+        # Enable nodelay for last checks
+        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+        duration = time.perf_counter() - start_time
 
-    _LOGGER.info("Upload took %.2f seconds, waiting for result...", duration)
+        _LOGGER.info("Upload took %.2f seconds, waiting for result...", duration)
 
-    # Once the device has the complete image it commits the update and
-    # reboots on its own; the exact commit point is not observable from
-    # here, so treat everything past the data phase as non-retryable. A
-    # re-upload could flash a device that already updated successfully.
-    commit_start = time.perf_counter()
-    try:
-        receive_exactly(sock, 1, "update receive result", RESPONSE_RECEIVE_OK)
-        receive_exactly(sock, 1, "update end result", RESPONSE_UPDATE_END_OK)
-    except OTANetworkError as err:
-        raise _committed_error(err) from err
-    commit_duration = time.perf_counter() - commit_start
+        # Once the device has the complete image it commits the update and
+        # reboots on its own; the exact commit point is not observable from
+        # here, so treat everything past the data phase as non-retryable. A
+        # re-upload could flash a device that already updated successfully.
+        commit_start = time.perf_counter()
+        try:
+            receive_exactly(sock, 1, "update receive result", RESPONSE_RECEIVE_OK)
+            receive_exactly(sock, 1, "update end result", RESPONSE_UPDATE_END_OK)
+        except OTANetworkError as err:
+            raise _committed_error(err) from err
+        commit_duration = time.perf_counter() - commit_start
 
-    # Sum of the named windows so the breakdown is self consistent; connect,
-    # handshake, auth, and the one MD5 round trip are not included
-    _LOGGER.info(
-        "Update took %.2f seconds (prepare %.2f, upload %.2f, commit %.2f)",
-        prepare_duration + duration + commit_duration,
-        prepare_duration,
-        duration,
-        commit_duration,
-    )
+        # Sum of the named windows so the breakdown is self consistent; connect,
+        # handshake, auth, and the one MD5 round trip are not included
+        _LOGGER.info(
+            "Update took %.2f seconds (prepare %.2f, upload %.2f, commit %.2f)",
+            prepare_duration + duration + commit_duration,
+            prepare_duration,
+            duration,
+            commit_duration,
+        )
 
-    try:
-        send_check(sock, RESPONSE_OK, "end acknowledgement")
-    except OTANetworkError as err:
-        # The device treats a missing end acknowledgement as non-fatal and is
-        # already rebooting into the new firmware, so the update succeeded
-        _LOGGER.warning("Failed sending end acknowledgement: %s", err)
-        _LOGGER.info("OTA successful (end acknowledgement not delivered)")
-    else:
-        _LOGGER.info("OTA successful")
+        try:
+            send_check(sock, RESPONSE_OK, "end acknowledgement")
+        except OTANetworkError as err:
+            # The device treats a missing end acknowledgement as non-fatal and is
+            # already rebooting into the new firmware, so the update succeeded
+            _LOGGER.warning("Failed sending end acknowledgement: %s", err)
+            _LOGGER.info("OTA successful (end acknowledgement not delivered)")
+        else:
+            _LOGGER.info("OTA successful")
+        if udp is not None:
+            # Our end acknowledgement also tells the device we have its
+            # responses, ending its wait before it reboots
+            udp.flush(2.0)
+    finally:
+        if udp is not None:
+            udp.close()
 
     # Do not connect logs until it is fully on
     time.sleep(1)
@@ -1121,7 +1513,7 @@ def run_ota_impl_(
     # are shared across the addresses, cycling through them. Wait before an
     # attempt when the previous one actually reached the device, or when
     # revisiting an address, so a flaky link can recover and the device can
-    # clean up a half-open connection (its handshake watchdog runs at 20s);
+    # clean up a half-open connection (its handshake watchdog runs at 60s);
     # moving on to the next address family stays immediate. Known limitation:
     # a silent mid-transfer drop with no reset can wedge the device until its
     # 105s data timeout, which outlasts this budget; the retries target the
@@ -1139,6 +1531,8 @@ def run_ota_impl_(
             on_upload_start()
             on_upload_start = None  # once per upload, whatever the retries do
 
+    # A network failure means a lossy link, so later attempts send data over UDP
+    prefer_udp = False
     while attempt < total_attempts:
         af, socktype, _, _, sa = res[attempt % len(res)]
         if not reconnect and (reached_device or attempt >= len(res)):
@@ -1153,7 +1547,7 @@ def run_ota_impl_(
         reconnect = False
         _LOGGER.info("Connecting to %s port %s...", sa[0], sa[1])
         sock = socket.socket(af, socktype)
-        sock.settimeout(20.0)
+        sock.settimeout(CONNECT_TIMEOUT)
         try:
             sock.connect(sa)
         except OSError as err:
@@ -1164,7 +1558,9 @@ def run_ota_impl_(
             continue
 
         _LOGGER.info("Connected to %s", sa[0])
+        sock.settimeout(SETUP_TIMEOUT)
         reached_device = True
+        total_attempts = max(total_attempts, len(res) + EXTRA_UPLOAD_ATTEMPTS_REACHED)
         with contextlib.closing(sock), Path(filename).open("rb") as file_handle:
             try:
                 perform_ota(
@@ -1177,6 +1573,7 @@ def run_ota_impl_(
                     encryption.plaintext_fallback,
                     allow_plaintext_upload=allow_plaintext_upload,
                     on_upload_start=upload_started,
+                    prefer_udp=prefer_udp,
                 )
             except OTAKeyRejected as err:
                 # The device sent a reject and closed on its own, so reconnect
@@ -1196,6 +1593,7 @@ def run_ota_impl_(
                 raise
             except OTAHandshakeNetworkError as err:
                 last_error = str(err)
+                prefer_udp = True
                 if encryption.handshake_fault_falls_back():
                     encryption.downgrade(last_error)
                     continue
@@ -1205,6 +1603,7 @@ def run_ota_impl_(
             except OTANetworkError as err:
                 # Transient network failure; retry
                 last_error = str(err)
+                prefer_udp = True
                 _LOGGER.warning("%s", last_error)
                 attempt += 1
                 continue

tests/integration/test_host_ota.py

@@ -37,6 +37,59 @@ DEVICE_NAME = "host-ota-test"
 API_KEY = "AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8="
 
 
+_REAL_COMPRESS = zlib.compress
+
+
+def _corrupt_compress(data: bytes, *args: object, **kwargs: object) -> bytes:
+    """Reserved block type in the first header: rejected by the decoder on
+    every build, unlike a flipped data bit that may only fail the MD5."""
+    out = bytearray(_REAL_COMPRESS(data, *args, **kwargs))
+    out[0] |= 0x06
+    return bytes(out)
+
+
+@contextmanager
+def _udp_mode(
+    monkeypatch: pytest.MonkeyPatch,
+    force: bool = True,
+    loss: float = 0.0,
+    blocked: bool = False,
+) -> Generator[list[bool]]:
+    """Force the UDP data phase (unless force is False) and drop datagrams both
+    ways; yields one entry per upload, True when it used UDP."""
+    import random
+
+    rng = random.Random(1)
+    used: list[bool] = []
+    real_start = espota2._start_udp
+    real_send = espota2.UdpChannel._send
+    real_ack = espota2.UdpChannel._handle_ack
+
+    def start(*args: object) -> espota2.UdpChannel | None:
+        channel = real_start(*args)
+        used.append(channel is not None)
+        return channel
+
+    def send(self: espota2.UdpChannel, *args: object) -> None:
+        if not blocked and rng.random() >= loss:
+            real_send(self, *args)
+
+    def ack(self: espota2.UdpChannel, data: bytes) -> None:
+        if rng.random() >= loss:
+            real_ack(self, data)
+
+    with monkeypatch.context() as m:
+        if force:
+            m.setattr(espota2, "UDP_SLOW_EXCHANGE", -1.0)
+        if blocked:
+            # Nothing answers, so do not wait the full probe
+            m.setattr(espota2, "UDP_PROBE_TIMEOUT", 1.0)
+        m.setattr(espota2, "_start_udp", start)
+        m.setattr(espota2.UdpChannel, "_send", send)
+        m.setattr(espota2.UdpChannel, "_handle_ack", ack)
+        yield used
+
+
 @contextmanager
 def _reserve_port() -> Generator[tuple[int, socket.socket]]:
     s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
@@ -272,18 +325,9 @@ async def test_host_ota_deflate(
             errors.append(line)
         dev.on_log(line)
 
-    real_compress = zlib.compress
-
-    def corrupt_compress(data: bytes, *args: object, **kwargs: object) -> bytes:
-        # Reserved block type in the first header: rejected by the decoder on
-        # every build, unlike a flipped data bit that may only fail the MD5
-        out = bytearray(real_compress(data, *args, **kwargs))
-        out[0] |= 0x06
-        return bytes(out)
-
     def overlong_compress(data: bytes, *args: object, **kwargs: object) -> bytes:
         """A stream that inflates past the size the client announced."""
-        return real_compress(data + bytes(8192), *args, **kwargs)
+        return _REAL_COMPRESS(data + bytes(8192), *args, **kwargs)
 
     async with run_binary(dev.binary_path, line_callback=on_log) as (proc, _lines):
         dev.proc = proc
@@ -301,7 +345,7 @@ async def test_host_ota_deflate(
 
         # A corrupt stream fails the upload and leaves the device running
         with monkeypatch.context() as m:
-            m.setattr(zlib, "compress", corrupt_compress)
+            m.setattr(zlib, "compress", _corrupt_compress)
             await dev.refused_ota(None, None, "corrupt deflate stream was accepted")
         assert errors, "device did not report the corrupt stream"
 
@@ -318,6 +362,29 @@ async def test_host_ota_deflate(
         await dev.ota(None, None, "upload after a rejected stream failed")
         assert dev.inflates == 2
 
+        # Localhost is clean, so the client stays on TCP
+        with _udp_mode(monkeypatch, force=False) as used:
+            await dev.ota(None, None, "TCP upload failed")
+        assert used == [False], "a clean link asked for UDP"
+
+        # A lossy link moves the data phase to UDP and completes through loss
+        with _udp_mode(monkeypatch, loss=0.2) as used:
+            await dev.ota(None, None, "UDP upload with loss failed")
+        assert used == [True], "the upload did not use UDP"
+        assert dev.inflates == 4
+
+        # Blocked UDP falls back to TCP
+        with _udp_mode(monkeypatch, blocked=True) as used:
+            await dev.ota(None, None, "fallback to TCP failed")
+        assert used == [False], "blocked UDP did not fall back to TCP"
+
+        # A device error mid upload surfaces at once, not as a 160 s timeout
+        with _udp_mode(monkeypatch) as used, monkeypatch.context() as m:
+            m.setattr(zlib, "compress", _corrupt_compress)
+            async with asyncio.timeout(30):
+                await dev.refused_ota(None, None, "corrupt stream over UDP accepted")
+        assert used == [True]
+
 
 @pytest.mark.asyncio
 async def test_host_ota_encrypted(
@@ -325,6 +392,7 @@ async def test_host_ota_encrypted(
     write_yaml_config: ConfigWriter,
     compile_esphome: CompileFunction,
     reserved_tcp_port: tuple[int, socket.socket],
+    monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     """A client that leaves right after the handshake, as a key probe does,
     is a clean close, not an OTA error; a plaintext upload is refused; an
@@ -348,6 +416,10 @@ async def test_host_ota_encrypted(
             None, None, "plaintext upload to an encrypted device must fail"
         )
         await dev.ota(None, API_KEY, "encrypted OTA reported failure")
+        # Noise frames ride the UDP stream unchanged, through loss
+        with _udp_mode(monkeypatch, loss=0.2) as used:
+            await dev.ota(None, API_KEY, "encrypted UDP upload failed")
+        assert used == [True]
 
 
 @pytest.mark.asyncio

tests/unit_tests/test_espota2.py

@@ -11,6 +11,7 @@ import logging
 from pathlib import Path
 import socket
 import struct
+from typing import Self
 from unittest.mock import Mock, call, patch
 import zlib
 
@@ -357,6 +358,7 @@ def test_perform_ota_successful_md5_auth(
                 | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
                 | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
                 | espota2.CLIENT_FEATURE_SUPPORTS_DEFLATE
+                | espota2.CLIENT_FEATURE_SUPPORTS_UDP
             ]
         )
     )
@@ -765,7 +767,7 @@ def test_run_ota_impl_successful(
     assert result_host == "192.168.1.100"
 
     # Verify socket was configured correctly
-    mock_socket.settimeout.assert_called_with(20.0)
+    mock_socket.settimeout.assert_called_with(espota2.SETUP_TIMEOUT)
     mock_socket.connect.assert_called_once_with(("192.168.1.100", 3232))
     mock_socket.close.assert_called_once()
 
@@ -853,13 +855,22 @@ def test_run_ota_impl_network_error_retry_succeeds(
     assert result_host == "192.168.1.100"
     assert mock_perform_ota.call_count == 2
     mock_sleep.assert_called_once_with(espota2.UPLOAD_RETRY_DELAY)
+    # The failure marks the link lossy, so the retry moves the data to UDP
+    assert [c.kwargs["prefer_udp"] for c in mock_perform_ota.call_args_list] == [
+        False,
+        True,
+    ]
 
 
 @pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip")
 def test_run_ota_impl_network_error_exhausts_attempts(
     mock_socket: Mock, firmware_file: Path, mock_perform_ota: Mock, mock_sleep: Mock
 ) -> None:
-    """Test run_ota_impl_ gives up after all attempts hit network errors."""
+    """Test run_ota_impl_ gives up after all attempts hit network errors.
+
+    Every attempt connected, so the budget is the longer one for a reachable
+    device on a lossy link.
+    """
     mock_perform_ota.side_effect = espota2.OTANetworkError("sending data: broken pipe")
 
     result_code, result_host = espota2.run_ota_impl_(
@@ -868,8 +879,8 @@ def test_run_ota_impl_network_error_exhausts_attempts(
 
     assert result_code == 1
     assert result_host is None
-    assert mock_perform_ota.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS + 1
-    assert mock_sleep.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS
+    assert mock_perform_ota.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS_REACHED + 1
+    assert mock_sleep.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS_REACHED
 
 
 @pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip_dual")
@@ -890,12 +901,10 @@ def test_run_ota_impl_multiple_addresses_cycle(
     assert mock_socket.connect.call_args_list == [
         call(DUAL_STACK_SA6),
         call(DUAL_STACK_SA4),
-        call(DUAL_STACK_SA6),
-        call(DUAL_STACK_SA4),
-    ]
+    ] * (1 + espota2.EXTRA_UPLOAD_ATTEMPTS // 2)
     # No connect ever reached the device, so the delay only applies before
     # the revisits
-    assert mock_sleep.call_count == 2
+    assert mock_sleep.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS
 
 
 @pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip_dual")
@@ -1079,6 +1088,7 @@ def test_perform_ota_successful_sha256_auth(
                 | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
                 | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
                 | espota2.CLIENT_FEATURE_SUPPORTS_DEFLATE
+                | espota2.CLIENT_FEATURE_SUPPORTS_UDP
             ]
         )
     )
@@ -1136,6 +1146,7 @@ def test_perform_ota_sha256_fallback_to_md5(
                 | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
                 | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
                 | espota2.CLIENT_FEATURE_SUPPORTS_DEFLATE
+                | espota2.CLIENT_FEATURE_SUPPORTS_UDP
             ]
         )
     )
@@ -1246,6 +1257,7 @@ def test_perform_ota_extended_protocol_app(
                 | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
                 | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
                 | espota2.CLIENT_FEATURE_SUPPORTS_DEFLATE
+                | espota2.CLIENT_FEATURE_SUPPORTS_UDP
             ]
         )
     )
@@ -1307,6 +1319,7 @@ def test_perform_ota_successful_partition_table(
                 | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
                 | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
                 | espota2.CLIENT_FEATURE_SUPPORTS_DEFLATE
+                | espota2.CLIENT_FEATURE_SUPPORTS_UDP
             ]
         )
     )
@@ -1572,3 +1585,468 @@ def test_perform_ota_with_deflate(mock_socket: Mock, server_features: int) -> No
     assert len(payload) == sent_size < len(original_content)
     assert zlib.decompress(payload, -espota2.DEFLATE_WINDOW_BITS) == original_content
     assert sent[5] == hashlib.md5(original_content).hexdigest().encode()
+
+
+class _LossyUdpDevice:
+    """The device side of the UDP channel (ota_esphome_udp.cpp) on a lossy link."""
+
+    def __init__(self, token: bytes, loss: float, seed: int = 1) -> None:
+        import random
+        import threading
+
+        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self.sock.bind(("127.0.0.1", 0))
+        self.sock.settimeout(0.05)
+        self.token = token
+        self.loss = loss
+        self.random = random.Random(seed)
+        self.next_seq = 0
+        self.slots: dict[int, bytes] = {}  # early messages within the window
+        self.received: list[bytes] = []
+        self.out = bytearray()
+        self.out_base = 0
+        self.committed = False
+        self._stop = threading.Event()
+        self._thread = threading.Thread(target=self._run, daemon=True)
+
+    def __enter__(self) -> Self:
+        self._thread.start()
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self._stop.set()
+        self._thread.join()
+        self.sock.close()
+
+    def _run(self) -> None:
+        while not self._stop.is_set():
+            try:
+                data, peer = self.sock.recvfrom(2048)
+            except TimeoutError:
+                continue
+            if self.random.random() < self.loss:
+                continue
+            if data[1:5] != self.token:
+                continue
+            self.committed |= bool(data[0] & espota2.UDP_FLAG_COMMITTED)
+            acked = (int.from_bytes(data[7:9], "big") - self.out_base) & 0xFFFF
+            if acked <= len(self.out):
+                del self.out[:acked]
+                self.out_base += acked
+            seq = int.from_bytes(data[5:7], "big")
+            msg_type = data[0] & ~espota2.UDP_FLAG_COMMITTED
+            if (
+                msg_type == espota2.UDP_MSG_DATA
+                and (seq - self.next_seq) & 0xFFFF < espota2.UDP_WINDOW
+            ):
+                self.slots.setdefault(seq, data[espota2.UDP_HEADER_SIZE :])
+            # Consumed as soon as it is next, like a device that is never busy
+            while self.next_seq in self.slots:
+                self.received.append(self.slots.pop(self.next_seq))
+                self.next_seq = (self.next_seq + 1) & 0xFFFF
+            if self.random.random() < self.loss:
+                continue
+            window = sum(
+                1 << k
+                for k in range(1, espota2.UDP_WINDOW)
+                if (self.next_seq + k) & 0xFFFF in self.slots
+            )
+            # Consumes at once, so every slot it does not hold is free
+            window |= (espota2.UDP_WINDOW - len(self.slots)) << 4
+            self.sock.sendto(
+                bytes([espota2.UDP_MSG_ACK])
+                + self.token
+                + self.next_seq.to_bytes(2, "big")
+                + seq.to_bytes(2, "big")
+                + (self.out_base & 0xFFFF).to_bytes(2, "big")
+                + bytes([window])
+                + bytes(self.out),
+                peer,
+            )
+
+
+def _udp_client(device: _LossyUdpDevice) -> espota2.UdpChannel:
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    sock.connect(device.sock.getsockname())
+    return espota2.UdpChannel(sock, device.token)
+
+
+@pytest.mark.parametrize("loss", [0.0, 0.5])
+def test_udp_channel_delivers_in_order(loss: float) -> None:
+    """The stream arrives intact and responses come back, despite loss."""
+    data = bytes(range(256)) * 160
+    with _LossyUdpDevice(b"\x01\x02\x03\x04", loss) as device:
+        channel = _udp_client(device)
+        assert channel.probe()
+        channel.sendall(data)
+        device.out += bytes([espota2.RESPONSE_RECEIVE_OK])
+        assert channel.recv(1) == bytes([espota2.RESPONSE_RECEIVE_OK])
+        channel.flush(30.0)
+        channel.close()
+    assert b"".join(device.received) == data
+    assert all(len(m) <= espota2.UDP_MAX_PAYLOAD for m in device.received)
+    assert device.committed
+
+
+def test_start_udp_falls_back_when_blocked(mock_socket: Mock) -> None:
+    """No answer to the probes sends the fallback over TCP and stays on TCP."""
+    with _LossyUdpDevice(b"\x05\x06\x07\x08", 1.0) as device:
+        mock_socket.family = socket.AF_INET
+        mock_socket.getpeername.return_value = device.sock.getsockname()
+        mock_socket.recv.side_effect = [b"\x05", b"\x06\x07\x08"]
+        with patch.object(espota2, "UDP_PROBE_TIMEOUT", 0.3):
+            assert espota2._start_udp(mock_socket, mock_socket, True) is None
+    assert mock_socket.sendall.call_args_list == [
+        call(espota2.UDP_REQUEST),
+        call(espota2.UDP_FALLBACK),  # raw on the TCP socket
+    ]
+
+
+def test_start_udp_declined(mock_socket: Mock) -> None:
+    """A zero token means the device could not open UDP; nothing is probed."""
+    mock_socket.recv.side_effect = [b"\x00", b"\x00\x00\x00"]
+    assert espota2._start_udp(mock_socket, mock_socket, True) is None
+    assert mock_socket.sendall.call_args_list == [call(espota2.UDP_REQUEST)]
+
+
+def test_start_udp_clean_link_stays_on_tcp(mock_socket: Mock) -> None:
+    """A clean handshake asks to stay on TCP and opens no UDP socket."""
+    mock_socket.recv.side_effect = [b"\x00", b"\x00\x00\x00"]
+    with patch.object(espota2.socket, "socket") as udp_socket:
+        assert espota2._start_udp(mock_socket, mock_socket, False) is None
+    udp_socket.assert_not_called()
+    assert mock_socket.sendall.call_args_list == [call(espota2.UDP_FALLBACK)]
+
+
+@pytest.mark.parametrize(
+    ("slowest", "prefer_udp", "retransmits", "lossy"),
+    [
+        (0.01, False, 0, False),
+        (0.01, False, None, False),
+        (espota2.UDP_SLOW_EXCHANGE + 0.1, False, 0, True),
+        (0.01, True, 0, True),
+        (0.01, False, 2, True),
+    ],
+)
+def test_link_lossy(
+    slowest: float, prefer_udp: bool, retransmits: int | None, lossy: bool
+) -> None:
+    """A slow exchange, a retransmit or an earlier failure means UDP."""
+    with patch.object(espota2, "_tcp_retransmits", return_value=retransmits):
+        assert espota2._link_lossy(Mock(), slowest, prefer_udp) is lossy
+
+
+@pytest.mark.parametrize(
+    ("platform", "info", "expected"),
+    [
+        ("linux", bytes(100) + (3).to_bytes(4, "little") + bytes(100), 3),
+        ("darwin", bytes(72) + (1460).to_bytes(8, "little") + bytes(8), 1460),
+        ("win32", b"", None),
+        ("linux", bytes(10), None),
+    ],
+)
+def test_tcp_retransmits(platform: str, info: bytes, expected: int | None) -> None:
+    """The retransmit counter is read where the OS reports it."""
+    sock = Mock()
+    sock.getsockopt.return_value = info
+    with patch("sys.platform", platform):
+        assert espota2._tcp_retransmits(sock) == expected
+
+
+def test_udp_channel_reports_socket_errors() -> None:
+    """A closed device port fails fast instead of waiting out the timeout."""
+    gone = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    gone.bind(("127.0.0.1", 0))
+    addr = gone.getsockname()
+    gone.close()
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    sock.connect(addr)
+    channel = espota2.UdpChannel(sock, b"\x01\x01\x01\x01")
+    channel.settimeout(10.0)
+    with pytest.raises(espota2.OTANetworkError, match="closed its UDP port"):
+        channel.sendall(bytes(espota2.UDP_MAX_PAYLOAD * 5))
+    # A probe that is refused falls back instead
+    assert not espota2.UdpChannel(sock, b"\x01\x01\x01\x01").probe()
+    channel.close()
+
+
+def test_udp_channel_no_progress_names_last_error() -> None:
+    """Without refusals, the no-progress timeout names the last socket error."""
+    sock = Mock()
+    sock.send.side_effect = BlockingIOError(35, "full")
+    channel = espota2.UdpChannel(sock, b"\x01\x01\x01\x01")
+    channel.settimeout(0.2)
+    channel.sendall(b"x")
+    with (
+        patch("select.select", return_value=([], [], [])),
+        pytest.raises(espota2.OTANetworkError, match="no progress.*full"),
+    ):
+        while True:
+            channel._pump()
+
+
+def test_udp_channel_ack_resets_refusals() -> None:
+    """Refusals only count in a row; an ACK in between resets them."""
+    channel = espota2.UdpChannel(Mock(), b"\x01\x01\x01\x01")
+    ack = bytes([espota2.UDP_MSG_ACK]) + b"\x01\x01\x01\x01" + bytes(6) + b"\x40"
+    for _ in range(espota2.UDP_REFUSALS * 2):
+        channel._socket_error(ConnectionRefusedError(61, "refused"))
+        channel._handle_ack(ack)
+    # Windows' form of the same refusal counts too
+    with pytest.raises(espota2.OTANetworkError, match="closed its UDP port"):
+        for _ in range(espota2.UDP_REFUSALS):
+            channel._socket_error(ConnectionResetError(10054, "forcibly closed"))
+
+
+def _udp_handshake() -> list[bytes]:
+    """TCP recv responses up to auth, from a device that offers UDP."""
+    return [
+        bytes([espota2.RESPONSE_OK]),
+        bytes([espota2.OTA_VERSION_2_0]),
+        bytes([espota2.RESPONSE_FEATURE_FLAGS]),
+        bytes([espota2.SERVER_FEATURE_SUPPORTS_UDP]),
+        bytes([espota2.RESPONSE_AUTH_OK]),
+    ]
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_over_udp(mock_socket: Mock) -> None:
+    """The data phase moves to the channel, which acks chunks itself."""
+    mock_socket.recv.side_effect = _udp_handshake()
+    udp = Mock()
+    udp.pending.return_value = False
+    udp.recv.side_effect = [
+        bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]),
+        bytes([espota2.RESPONSE_BIN_MD5_OK]),
+        bytes([espota2.RESPONSE_RECEIVE_OK]),
+        bytes([espota2.RESPONSE_UPDATE_END_OK]),
+    ]
+    with patch.object(espota2, "_start_udp", return_value=udp) as start:
+        espota2.perform_ota(
+            mock_socket, None, io.BytesIO(bytes(20000)), "test.bin", prefer_udp=True
+        )
+    assert start.call_args.args[2] is True  # an earlier failure means lossy
+    assert udp.recv.call_count == 4  # no chunk acks
+    udp.flush.assert_called_once()
+    udp.close.assert_called_once()
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_over_udp_device_error(mock_socket: Mock) -> None:
+    """A device error mid upload stops the upload and is not retried."""
+    mock_socket.recv.side_effect = _udp_handshake()
+    mock_socket.getsockopt.return_value = b""  # no retransmit counter
+    udp = Mock()
+    udp.pending.return_value = True
+    udp.recv.side_effect = [
+        bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]),
+        bytes([espota2.RESPONSE_BIN_MD5_OK]),
+        bytes([espota2.RESPONSE_ERROR_WRITING_FLASH]),
+    ]
+    with (
+        patch.object(espota2, "_start_udp", return_value=udp),
+        pytest.raises(espota2.OTAError, match="Writing OTA data") as exc,
+    ):
+        espota2.perform_ota(mock_socket, None, io.BytesIO(bytes(20000)), "test.bin")
+    assert not isinstance(exc.value, espota2.OTANetworkError)
+    udp.close.assert_called()
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_udp_declined_stays_on_tcp(mock_socket: Mock) -> None:
+    """Without a channel the upload continues on TCP with chunk acks."""
+    mock_socket.recv.side_effect = [
+        *_udp_handshake(),
+        bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]),
+        bytes([espota2.RESPONSE_BIN_MD5_OK]),
+        *_UPLOAD_TAIL,
+    ]
+    with patch.object(espota2, "_start_udp", return_value=None):
+        espota2.perform_ota(
+            mock_socket, None, io.BytesIO(b"x" * 100), "test.bin", prefer_udp=True
+        )
+    assert mock_socket.sendall.call_args_list[-1] == call(bytes([espota2.RESPONSE_OK]))
+
+
+def test_udp_channel_records_socket_errors() -> None:
+    """Send and receive errors are recorded like loss."""
+    sock = Mock()
+    sock.send.side_effect = BlockingIOError(35, "full")
+    sock.recv.side_effect = ConnectionRefusedError(61, "refused")
+    channel = espota2.UdpChannel(sock, b"\x01\x01\x01\x01")
+    channel._send(espota2.UDP_MSG_PING, 0, b"")
+    assert isinstance(channel._last_error, BlockingIOError)
+    with patch("select.select", side_effect=[([sock], [], []), ([], [], [])]):
+        channel._receive(0.0)
+    assert isinstance(channel._last_error, ConnectionRefusedError)
+
+
+def test_start_udp_probe_succeeds(mock_socket: Mock) -> None:
+    """A device that answers the probe gets the data phase over UDP."""
+    with _LossyUdpDevice(b"\x05\x06\x07\x09", 0.0) as device:
+        mock_socket.family = socket.AF_INET
+        mock_socket.getpeername.return_value = device.sock.getsockname()
+        mock_socket.recv.side_effect = [b"\x05", b"\x06\x07\x09"]
+        channel = espota2._start_udp(mock_socket, mock_socket, True)
+        assert channel is not None
+        channel.close()
+    assert mock_socket.sendall.call_args_list == [call(espota2.UDP_REQUEST)]
+
+
+def test_udp_channel_pings_when_all_held() -> None:
+    """With every message in flight already held, it asks for responses instead,
+    at most once per resend timeout."""
+    device = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    device.bind(("127.0.0.1", 0))
+    device.settimeout(1.0)
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    sock.connect(device.getsockname())
+    channel = espota2.UdpChannel(sock, b"\x01\x01\x01\x01")
+    channel.sendall(b"x")
+    channel._queue[0].held = True
+    with patch.object(channel, "_receive"):
+        channel._pump()
+        channel._pump()
+    assert device.recv(2048)[0] == espota2.UDP_MSG_PING
+    device.settimeout(0.1)
+    with pytest.raises(TimeoutError):
+        device.recv(2048)
+    channel.close()
+    device.close()
+
+
+def test_udp_channel_ignores_foreign_datagrams() -> None:
+    """Short, non-ACK or wrong-token datagrams change nothing."""
+    channel = espota2.UdpChannel(Mock(), b"\x01\x01\x01\x01")
+    for data in (b"\x40", bytes(espota2.UDP_ACK_HEADER_SIZE), b"\x40" + bytes(11)):
+        channel._handle_ack(data)
+    assert not channel.pending()
+    # An ACK before the head was ever sent is no duplicate
+    channel.sendall(b"x")
+    ack = bytes([espota2.UDP_MSG_ACK]) + b"\x01\x01\x01\x01" + bytes(6) + b"\x40"
+    channel._handle_ack(ack)
+    channel._handle_ack(ack)
+    assert not channel._queue[0].resend
+
+
+def test_udp_channel_logs_each_socket_error_once(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """A repeated error is recorded but logged only when it changes."""
+    channel = espota2.UdpChannel(Mock(), b"\x01\x01\x01\x01")
+    with caplog.at_level(logging.DEBUG, "esphome.espota2"):
+        channel._socket_error(ConnectionRefusedError(61, "refused"))
+        channel._socket_error(ConnectionRefusedError(61, "refused again"))
+    assert caplog.text.count("UDP socket error") == 1
+    assert "again" in str(channel._last_error)
+
+
+def test_tcp_retransmits_unsupported_option() -> None:
+    """An OS that rejects the option reports nothing."""
+    sock = Mock()
+    sock.getsockopt.side_effect = OSError("not supported")
+    with patch("sys.platform", "linux"):
+        assert espota2._tcp_retransmits(sock) is None
+
+
+@pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip")
+def test_run_ota_impl_handshake_failure_prefers_udp(
+    firmware_file: Path, mock_perform_ota: Mock, mock_sleep: Mock
+) -> None:
+    """A handshake that dies after reaching the device moves the retry to UDP."""
+    mock_perform_ota.side_effect = [
+        espota2.OTAHandshakeNetworkError("noise handshake: timed out"),
+        None,
+    ]
+    result_code, _ = espota2.run_ota_impl_("test.local", 3232, None, str(firmware_file))
+    assert result_code == 0
+    assert [c.kwargs["prefer_udp"] for c in mock_perform_ota.call_args_list] == [
+        False,
+        True,
+    ]
+
+
+def test_start_udp_socket_error_falls_back(mock_socket: Mock) -> None:
+    """A UDP socket that cannot reach the peer stays on TCP."""
+    mock_socket.family = socket.AF_INET
+    mock_socket.getpeername.side_effect = OSError("not connected")
+    mock_socket.recv.side_effect = [b"\x05", b"\x06\x07\x08"]
+    assert espota2._start_udp(mock_socket, mock_socket, True) is None
+    assert mock_socket.sendall.call_args_list == [
+        call(espota2.UDP_REQUEST),
+        call(espota2.UDP_FALLBACK),
+    ]
+
+
+def test_udp_channel_backs_off_when_acks_come_back_untimed() -> None:
+    """Progress acks that only answer resent messages mean the link is slower
+    than the timer: it backs off, up to its cap, until a sample resets it."""
+    token = b"\x01\x01\x01\x01"
+    channel = espota2.UdpChannel(Mock(), token)
+    channel._queue = [espota2._UdpMessage(seq, b"x") for seq in range(40)]
+    rtos = []
+    for received in range(1, 33):
+        for message in channel._queue[:1]:
+            message.sends = 2  # every message was resent before its ack
+        rtos.append(channel._rto())
+        ack = (
+            bytes([espota2.UDP_MSG_ACK])
+            + token
+            + received.to_bytes(2, "big")
+            + (received - 1).to_bytes(2, "big")
+            + bytes(2)
+            + b"\x40"
+        )
+        channel._handle_ack(ack)
+    n = espota2.UDP_UNTIMED_ACKS
+    assert rtos[0] == espota2.UDP_RESEND_INITIAL
+    assert rtos[n] == espota2.UDP_RESEND_INITIAL * 2
+    assert rtos[-1] == espota2.UDP_RESEND_MAX
+    channel._sample_rtt(0.1)
+    assert channel._rto() == pytest.approx(0.3)  # srtt 0.1 + 4 * rttvar 0.05
+
+
+@pytest.mark.usefixtures("mock_time")
+@pytest.mark.parametrize(
+    ("size", "pending", "match", "retryable"),
+    [
+        (100, False, "may have already committed", False),  # one, final chunk
+        (20000, False, "no progress", True),  # fails on an earlier chunk
+        (20000, True, "Writing OTA data", False),  # device error waiting
+    ],
+)
+def test_perform_ota_over_udp_send_failure(
+    mock_socket: Mock, size: int, pending: bool, match: str, retryable: bool
+) -> None:
+    """A UDP send failure surfaces a waiting device error, and on the last
+    chunk is not retried, like a lost final chunk ack over TCP."""
+    mock_socket.recv.side_effect = _udp_handshake()
+    udp = Mock()
+    udp.pending.return_value = pending
+    udp.recv.side_effect = [
+        bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]),
+        bytes([espota2.RESPONSE_BIN_MD5_OK]),
+        bytes([espota2.RESPONSE_ERROR_WRITING_FLASH]),
+    ]
+    # ota type, size and MD5 go through; the first data chunk fails
+    udp.sendall.side_effect = [None] * 3 + [espota2.OTANetworkError("no progress")]
+    with (
+        patch.object(espota2, "_start_udp", return_value=udp),
+        pytest.raises(espota2.OTAError, match=match) as exc,
+    ):
+        espota2.perform_ota(
+            mock_socket, None, io.BytesIO(bytes(size)), "test.bin", prefer_udp=True
+        )
+    assert isinstance(exc.value, espota2.OTANetworkError) is retryable
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_tcp_send_network_error_unchanged(mock_socket: Mock) -> None:
+    """Over TCP an OTANetworkError from a send stays retryable, even on the
+    last chunk."""
+    mock_socket.recv.side_effect = _no_auth_handshake(espota2.OTA_VERSION_2_0)
+    mock_socket.sendall.side_effect = [None] * 4 + [
+        espota2.OTANetworkError("sending frame")
+    ]
+    with pytest.raises(espota2.OTANetworkError, match="sending frame"):
+        espota2.perform_ota(mock_socket, None, io.BytesIO(b"x" * 100), "test.bin")

tests/unit_tests/test_espota2_noise.py

@@ -76,9 +76,11 @@ class FakeEncryptedDevice(threading.Thread):
         prologue_features_override: int | None = None,
         connections: int = 1,
         drop_handshakes: int = 0,
+        udp: bool = False,
         reject_reason: str | None = None,
     ) -> None:
         super().__init__(daemon=True)
+        self.udp = udp  # offer UDP; the data phase then has no chunk acks
         self.connections = connections
         self.drop_handshakes = drop_handshakes  # hang up mid-handshake this many times
         self.reject_reason = reject_reason  # refuse every handshake with this reason
@@ -125,6 +127,8 @@ class FakeEncryptedDevice(threading.Thread):
             sock.sendall(bytes([espota2.RESPONSE_ERROR_ENCRYPTION_REQUIRED]))
             return
         server_flags = espota2.SERVER_FEATURE_SUPPORTS_NOISE if self.offer_noise else 0
+        if self.udp:
+            server_flags |= espota2.SERVER_FEATURE_SUPPORTS_UDP
         sock.sendall(bytes([espota2.RESPONSE_FEATURE_FLAGS, server_flags]))
         if not (noise_negotiated and self.offer_noise):
             # A device that does not require encryption continues in
@@ -217,7 +221,7 @@ class FakeEncryptedDevice(threading.Thread):
         acked = 0
         while len(received) < size:
             received += recv_data(size - len(received))
-            if self.version >= espota2.OTA_VERSION_2_0:
+            if self.version >= espota2.OTA_VERSION_2_0 and not self.udp:
                 while acked + espota2.UPLOAD_BLOCK_SIZE <= len(received) or (
                     len(received) == size and acked < size
                 ):
@@ -906,6 +910,54 @@ def test_bare_block_refuses_a_device_that_cannot_encrypt(
     assert not any("Retrying in plaintext" in r.message for r in caplog.records)
 
 
+class _TcpAsUdp:
+    """Stands in for a UdpChannel, carrying its byte stream over the TCP socket."""
+
+    def __init__(self, sock: socket.socket) -> None:
+        self._sock = sock
+        self.flushed = False
+        self.closed = False
+
+    def settimeout(self, timeout: float | None) -> None:
+        self._sock.settimeout(timeout)
+
+    def setsockopt(self, level: int, optname: int, value: int) -> None:
+        self._sock.setsockopt(level, optname, value)
+
+    def sendall(self, data: bytes) -> None:
+        self._sock.sendall(data)
+
+    def recv(self, amount: int) -> bytes:
+        return self._sock.recv(amount)
+
+    def pending(self) -> bool:
+        return False
+
+    def flush(self, timeout: float) -> None:
+        self.flushed = True
+
+    def close(self) -> None:
+        self.closed = True
+
+
+def test_encrypted_upload_over_udp() -> None:
+    """After the handshake, Noise frames continue on the UDP channel."""
+    pytest.importorskip("aioesphomeapi.noise")
+    firmware = bytes(range(256)) * 80
+    device = FakeEncryptedDevice(udp=True)
+    channels: list[_TcpAsUdp] = []
+
+    def start_udp(sock: Any, tcp_sock: socket.socket, lossy: bool) -> _TcpAsUdp:
+        channels.append(_TcpAsUdp(tcp_sock))
+        return channels[0]
+
+    with patch("time.sleep"), patch.object(espota2, "_start_udp", start_udp):
+        _upload(device, firmware, PSK)
+    device.join_and_check()
+    assert device.received == firmware
+    assert channels[0].flushed and channels[0].closed
+
+
 def test_probe_ota_key_recomputes_the_budget_after_connect() -> None:
     """A slow connect leaves less for the handshake than the budget had."""
     clock = [0.0]