From 54b34edf2d7f5844cbf3f97686907609a7c41b59 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sat, 10 Oct 2026 12:15:36 -1000 Subject: [PATCH] [ota] Fall back to a UDP data phase when the link is lossy TCP gives up on links that lose many packets: retransmit timers back off and the connection drops, so devices on poor Wi-Fi cannot be updated without being moved. After the TCP handshake, a client that saw loss moves the upload to a UDP byte stream that the device acks from lwIP's receive callback as datagrams arrive, keeping at most four of them like TCP's window. Clean links stay on TCP; blocked UDP falls back to TCP. --- .../components/esphome/ota/ota_esphome.cpp | 82 ++++- esphome/components/esphome/ota/ota_esphome.h | 78 +++- .../esphome/ota/ota_esphome_noise.cpp | 10 +- .../esphome/ota/ota_esphome_udp.cpp | 266 ++++++++++++++ esphome/espota2.py | 332 +++++++++++++++++- tests/unit_tests/test_espota2.py | 205 ++++++++++- 6 files changed, 945 insertions(+), 28 deletions(-) create mode 100644 esphome/components/esphome/ota/ota_esphome_udp.cpp diff --git a/esphome/components/esphome/ota/ota_esphome.cpp b/esphome/components/esphome/ota/ota_esphome.cpp index e6b19c42f0c..e59d3fdb9e6 100644 --- a/esphome/components/esphome/ota/ota_esphome.cpp +++ b/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,9 @@ 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; +#endif #ifdef USE_OTA_ENCRYPTION inline bool ESPHomeOTAComponent::noise_offered_() const { @@ -266,6 +275,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 +344,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 +501,10 @@ void ESPHomeOTAComponent::handle_data_() { char *sbuf = reinterpret_cast(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 +520,34 @@ 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) + } + // A reply left all zero keeps the session on TCP + if (buf[0] != 0) + this->udp_open_(udp_link); + 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)) { @@ -623,6 +673,9 @@ void ESPHomeOTAComponent::handle_data_() { error: this->data_write_byte_(static_cast(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 +704,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 +729,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 +785,15 @@ 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 + ipaddr_ntoa_r(&this->udp_->peer_ip, peername, sizeof(peername)); + } else +#endif + { + this->client_->getpeername_to(peername); + } ESP_LOGD(TAG, "Starting %s from %s", LOG_STR_ARG(phase), peername); } @@ -840,7 +905,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 +939,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 +953,9 @@ void ESPHomeOTAComponent::send_chunk_acks_(DataTransfer &xfer) { } void ESPHomeOTAComponent::cleanup_connection_() { - this->client_->close(); + // 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; diff --git a/esphome/components/esphome/ota/ota_esphome.h b/esphome/components/esphome/ota/ota_esphome.h index 910a5b3d99d..63dc7579557 100644 --- a/esphome/components/esphome/ota/ota_esphome.h +++ b/esphome/components/esphome/ota/ota_esphome.h @@ -15,6 +15,20 @@ #include "esphome/core/preferences.h" #include "esphome/core/hash_base.h" +// The UDP data phase answers datagrams from lwIP's receive callback, so it needs the raw lwIP API from the main +// loop: single threaded (ESP8266), locked (RP2040, ESP32 with core locking) or LibreTiny (declines without core +// locking) +#ifdef USE_ESP32 +#include +#endif +#if defined(USE_ESP8266) || defined(USE_RP2) || defined(USE_LIBRETINY) || \ + (defined(USE_ESP32) && defined(CONFIG_LWIP_TCPIP_CORE_LOCKING)) +#define ESPHOME_OTA_UDP +#include "lwip/ip_addr.h" +struct pbuf; +struct udp_pcb; +#endif + namespace esphome { /// ESPHomeOTAComponent provides a simple way to integrate Over-the-Air updates into your app using ArduinoOTA. @@ -65,6 +79,51 @@ 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; + // Lives on handle_data_'s stack; the receive callback shares it under the lwIP lock + struct UdpLink { + ~UdpLink(); + void release(); // closes the pcb and frees what is queued + struct udp_pcb *pcb{nullptr}; + ip_addr_t peer_ip; + 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]{}; + struct pbuf *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_(); + static void udp_recv_cb_(void *arg, struct udp_pcb *pcb, struct pbuf *p, const ip_addr_t *addr, 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); +#endif + // The data phase reads 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 +134,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 +163,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 +261,14 @@ class ESPHomeOTAComponent final : public ota::OTAComponent { std::unique_ptr client_; ota::OTABackendPtr backend_; +#ifdef ESPHOME_OTA_UDP + union { + uint32_t client_connect_time_{0}; // handshake states + UdpLink *udp_; // data phase (never both), points into handle_data_'s stack + }; +#else uint32_t client_connect_time_{0}; +#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 diff --git a/esphome/components/esphome/ota/ota_esphome_noise.cpp b/esphome/components/esphome/ota/ota_esphome_noise.cpp index a271ef9d39d..7ae5e43df3f 100644 --- a/esphome/components/esphome/ota/ota_esphome_noise.cpp +++ b/esphome/components/esphome/ota/ota_esphome_noise.cpp @@ -266,13 +266,13 @@ 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]; + 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); diff --git a/esphome/components/esphome/ota/ota_esphome_udp.cpp b/esphome/components/esphome/ota/ota_esphome_udp.cpp new file mode 100644 index 00000000000..b87243e2cbb --- /dev/null +++ b/esphome/components/esphome/ota/ota_esphome_udp.cpp @@ -0,0 +1,266 @@ +#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" + +#include "lwip/inet.h" +#include "lwip/pbuf.h" +#include "lwip/udp.h" +#ifdef USE_LIBRETINY +#include "lwip/tcpip.h" +#endif + +#include + +/* + * 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 5 s probe +static constexpr uint32_t UDP_COMMIT_TIMEOUT = 70000; +// Like the TCP socket's SO_RCVTIMEO in handle_data_, so readall_ treats both alike +static constexpr uint32_t UDP_READ_WAIT_MS = 2000; +static constexpr uint32_t UDP_LINGER_MS = 3000; + +#ifdef USE_LIBRETINY +// LibreTiny's LwIPLock is a no-op; take the core lock directly (udp_open_ declines without one) +struct UdpLock { +#if LWIP_TCPIP_CORE_LOCKING + UdpLock() { LOCK_TCPIP_CORE(); } + ~UdpLock() { UNLOCK_TCPIP_CORE(); } +#endif +}; +#else +using UdpLock = LwIPLock; +#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; + if (this->pcb != nullptr) + udp_remove(this->pcb); + this->pcb = nullptr; + for (auto *&p : this->queue) { + if (p != nullptr) + pbuf_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; + 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(p->payload); + 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(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); + udp_sendto(link.pcb, p, &link.peer_ip, link.peer_port); + pbuf_free(p); +} + +// 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(arg); + uint8_t hdr[UDP_HEADER_SIZE]; + if (!ip_addr_cmp(addr, &link.peer_ip) || pbuf_copy_partial(p, hdr, UDP_HEADER_SIZE, 0) != UDP_HEADER_SIZE || + memcmp(hdr + 1, link.reply, 4) != 0) { + pbuf_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]); + struct pbuf *&slot = link.queue[seq % UDP_WINDOW]; + if ((hdr[0] & ~UDP_FLAG_COMMITTED) == UDP_MSG_DATA && p->tot_len > UDP_HEADER_SIZE && + static_cast(seq - link.consumed) < UDP_WINDOW && slot == nullptr) { + slot = p; + while (static_cast(link.received - link.consumed) < UDP_WINDOW && + link.queue[link.received % UDP_WINDOW] != nullptr) + link.received++; + } else { + pbuf_free(p); + } + udp_send_ack_(link, seq); +} + +void ESPHomeOTAComponent::udp_open_(UdpLink &link) { +#if defined(USE_LIBRETINY) && !LWIP_TCPIP_CORE_LOCKING + // No safe way to call lwIP from the main loop: decline +#else + struct sockaddr_storage peer; + socklen_t peer_len = sizeof(peer); + if (this->client_->getpeername(reinterpret_cast(&peer), &peer_len) != 0) + return; + uint32_t ipv4; + if (socket::sockaddr_to_ipv4(reinterpret_cast(&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(&peer)->sin6_addr); + IP_SET_TYPE_VAL(link.peer_ip, IPADDR_TYPE_V6); +#else + return; +#endif + } + uint8_t token[4]; + if (!random_bytes(token, sizeof(token))) + return; + 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; + // 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)); + udp_recv(link.pcb, udp_recv_cb_, &link); +#endif +} + +bool ESPHomeOTAComponent::udp_start_() { + // TCP now only carries a fallback byte; its blocking reads time out every couple of seconds to look again + const uint32_t start = millis(); + while (millis() - start < UDP_COMMIT_TIMEOUT) { + 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; + return true; + } + if (read == 0 || !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 (;;) { + struct pbuf *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 = pbuf_copy_partial(p, buf + got, std::min(len - got, p->tot_len - offset), offset); + got += n; + link.read_offset += n; + if (UDP_HEADER_SIZE + link.read_offset >= p->tot_len) { + UdpLock lock; + pbuf_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(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) { + { + UdpLock lock; + if (this->udp_->out_len == 0) + return; + } + this->yield_and_feed_watchdog_(); + } +} + +} // namespace esphome +#endif // USE_OTA && ESPHOME_OTA_UDP diff --git a/esphome/espota2.py b/esphome/espota2.py index dbdf3233660..d8d4734a7c9 100644 --- a/esphome/espota2.py +++ b/esphome/espota2.py @@ -66,11 +66,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 @@ -99,17 +101,55 @@ 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 +# Connect and 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; the device waits 70 s +UDP_PROBE_TIMEOUT = 5.0 +UDP_PROBE_INTERVAL = 0.1 +# Resend timeout bounds around srtt + 4 * rttvar; 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 +UDP_FALLBACK = b"\x00" +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) @@ -434,6 +474,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: @@ -534,6 +578,252 @@ 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, without backoff, so loss + slows the transfer instead of ending it. + """ + + 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 + 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 + + 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 + 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 + 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 + while len(self._queue) > UDP_WINDOW: + self._pump() + + 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: + return UDP_RESEND_INITIAL + return min(max(self._srtt + 4 * self._rttvar, UDP_RESEND_MIN), 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 + with contextlib.suppress(OSError): + # A full buffer or ICMP error is just more loss + self._sock.send(header + payload) + + def _pump(self) -> None: + now = time.monotonic() + if now - self._last_progress > self._timeout: + raise OTANetworkError(f"no progress over UDP for {self._timeout:.0f}s") + 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: + continue + self._handle_ack(data) + + def _sample_rtt(self, sample: float) -> None: + # RFC 6298 smoothing + 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 + 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) + if index < len(self._queue) and self._queue[index].sends == 1: + 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 + 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) + udp_sock.connect(tcp_sock.getpeername()) + 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 perform_ota( sock: socket.socket, password: str | None, @@ -543,6 +833,7 @@ def perform_ota( noise_psk: str | None = None, plaintext_fallback: bool = False, allow_plaintext_upload: bool = False, + 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 @@ -565,11 +856,15 @@ def perform_ota( file_size = len(file_contents) _LOGGER.info("Uploading %s (%s bytes)", filename, file_size) + # sock may become a Noise wrapper or a UDP channel; UDP needs the peer + tcp_sock = sock # 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.info("Connection established; device supports OTA version %s", version) supported_versions = (OTA_VERSION_1_0, OTA_VERSION_2_0) if version not in supported_versions: @@ -583,9 +878,11 @@ def perform_ota( | 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, @@ -593,6 +890,7 @@ def perform_ota( "features", None, # Accept any response )[0] + slowest_exchange = max(slowest_exchange, time.monotonic() - exchange_start) extended_proto = False if features == RESPONSE_FEATURE_FLAGS: @@ -751,6 +1049,18 @@ def perform_ota( sock.settimeout(DATA_PHASE_TIMEOUT) + udp: UdpChannel | None = None + 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 + if extended_proto: send_check(sock, ota_type, "ota type") @@ -806,7 +1116,8 @@ def perform_ota( ) raise OTANetworkError(f"sending data: {err}") from err - if version >= OTA_VERSION_2_0: + # The UDP channel acks every message itself + if version >= OTA_VERSION_2_0 and udp is None: try: receive_exactly(sock, 1, "chunk result", RESPONSE_CHUNK_OK) except OTANetworkError as err: @@ -860,6 +1171,10 @@ def perform_ota( _LOGGER.info("OTA successful (end acknowledgement not delivered)") else: _LOGGER.info("OTA successful") + if udp is not None: + # Nothing waits for this; just give the end acknowledgement a chance + udp.flush(2.0) + udp.close() # Do not connect logs until it is fully on time.sleep(1) @@ -904,7 +1219,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 @@ -914,6 +1229,8 @@ def run_ota_impl_( reached_device = False attempt = 0 encryption = _EncryptionAttempt(noise_psk, plaintext_fallback) + # 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 reached_device or attempt >= len(res): @@ -927,18 +1244,20 @@ def run_ota_impl_( reached_device = False _LOGGER.info("Connecting to %s port %s...", sa[0], sa[1]) sock = socket.socket(af, socktype) - sock.settimeout(20.0) + sock.settimeout(SETUP_TIMEOUT) try: sock.connect(sa) except OSError as err: sock.close() _LOGGER.warning("Connecting to %s port %s failed: %s", sa[0], sa[1], err) last_error = f"connecting to {sa[0]} failed: {err}" + prefer_udp = True attempt += 1 continue _LOGGER.info("Connected to %s", sa[0]) 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( @@ -950,6 +1269,7 @@ def run_ota_impl_( encryption.noise_psk, encryption.plaintext_fallback, allow_plaintext_upload=allow_plaintext_upload, + prefer_udp=prefer_udp, ) except OTAEncryptionFallback as err: # Same address and attempt budget: not a network retry @@ -958,6 +1278,7 @@ def run_ota_impl_( continue except OTAHandshakeNetworkError as err: last_error = str(err) + prefer_udp = True if encryption.handshake_fault_falls_back(): encryption.downgrade(last_error) continue @@ -967,6 +1288,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 diff --git a/tests/unit_tests/test_espota2.py b/tests/unit_tests/test_espota2.py index f4c5935059c..f024a32b95b 100644 --- a/tests/unit_tests/test_espota2.py +++ b/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 @@ -356,6 +357,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 ] ) ) @@ -764,7 +766,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() @@ -835,13 +837,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_( @@ -850,8 +861,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") @@ -872,12 +883,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") @@ -1061,6 +1070,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 ] ) ) @@ -1118,6 +1128,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 ] ) ) @@ -1228,6 +1239,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 ] ) ) @@ -1289,6 +1301,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 ] ) ) @@ -1554,3 +1567,179 @@ 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 + held = 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 + held |= (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([held]) + + 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_udp_channel_times_out_without_progress() -> None: + """A dead link still fails, after the no-progress timeout.""" + with _LossyUdpDevice(b"\x01\x01\x01\x01", 1.0) as device: + channel = _udp_client(device) + channel.settimeout(0.3) + with pytest.raises(espota2.OTANetworkError, match="no progress over UDP"): + channel.sendall(bytes(espota2.UDP_MAX_PAYLOAD * 5)) + channel.close() + + +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