From bfa0a5fc94a12f59af086491db714d21e63a08fb Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sat, 10 Oct 2026 12:56:41 -1000 Subject: [PATCH] Add host UDP support with end to end tests, surface device errors over UDP The host platform now offers UDP through a socket drained by the wait loops, so the host OTA integration tests cover UDP with loss, the TCP fallback, a device error mid upload and Noise. Over UDP the device no longer waits for the final ack, the client stops and reads a device error as soon as one arrives, socket errors are named in the no progress error, and the device notices the commit within 100 ms. LibreTiny without core locking drops the UDP code at build time. --- .../components/esphome/ota/ota_esphome.cpp | 21 ++- esphome/components/esphome/ota/ota_esphome.h | 54 ++++-- .../esphome/ota/ota_esphome_udp.cpp | 174 ++++++++++++++---- esphome/espota2.py | 37 +++- tests/integration/test_host_ota.py | 94 ++++++++-- tests/unit_tests/test_espota2.py | 34 ++-- 6 files changed, 328 insertions(+), 86 deletions(-) diff --git a/esphome/components/esphome/ota/ota_esphome.cpp b/esphome/components/esphome/ota/ota_esphome.cpp index 218c080e788..3f4400e57f1 100644 --- a/esphome/components/esphome/ota/ota_esphome.cpp +++ b/esphome/components/esphome/ota/ota_esphome.cpp @@ -649,11 +649,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); @@ -789,7 +795,12 @@ void ESPHomeOTAComponent::log_start_(const LogString *phase) { #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(&this->udp_->peer), this->udp_->peer_len, + peername); +#endif } else #endif { diff --git a/esphome/components/esphome/ota/ota_esphome.h b/esphome/components/esphome/ota/ota_esphome.h index 0c529d9f815..fbb1ba5447b 100644 --- a/esphome/components/esphome/ota/ota_esphome.h +++ b/esphome/components/esphome/ota/ota_esphome.h @@ -15,18 +15,23 @@ #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) +// UDP data phase: raw lwIP needs a main loop lock (or a single thread); the host polls a socket instead #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 +#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 { @@ -85,20 +90,31 @@ class ESPHomeOTAComponent final : public ota::OTAComponent { 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 pcb and frees what is queued + 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 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]{}; - 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] + 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]; @@ -112,9 +128,16 @@ class ESPHomeOTAComponent final : public ota::OTAComponent { 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); + // 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) { @@ -261,17 +284,20 @@ class ESPHomeOTAComponent final : public ota::OTAComponent { std::unique_ptr client_; ota::OTABackendPtr backend_; -#ifdef ESPHOME_OTA_UDP +#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_"); -#ifdef USE_OTA_ENCRYPTION - static_assert(sizeof(UdpLink::reply) <= NOISE_MAX_RESPONSE, "the token goes out through noise_write_"); -#endif #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 diff --git a/esphome/components/esphome/ota/ota_esphome_udp.cpp b/esphome/components/esphome/ota/ota_esphome_udp.cpp index ddc73dfa9aa..acf46e160c1 100644 --- a/esphome/components/esphome/ota/ota_esphome_udp.cpp +++ b/esphome/components/esphome/ota/ota_esphome_udp.cpp @@ -6,14 +6,15 @@ #include "esphome/core/helpers.h" #include "esphome/core/log.h" +#ifdef ESPHOME_OTA_UDP_LWIP #include "lwip/inet.h" -#include "lwip/pbuf.h" -#include "lwip/udp.h" #ifdef USE_LIBRETINY #include "lwip/tcpip.h" #endif +#endif #include +#include /* * UDP data phase: the same byte stream as TCP, cut into numbered datagrams. @@ -38,23 +39,44 @@ static constexpr size_t UDP_ACK_HEADER_SIZE = 12; 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; +// 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 (udp_open_ declines without one) +// LibreTiny's LwIPLock is a no-op; take the core lock directly (the header drops UDP without one) struct UdpLock { -#if LWIP_TCPIP_CORE_LOCKING UdpLock() { LOCK_TCPIP_CORE(); } ~UdpLock() { UNLOCK_TCPIP_CORE(); } -#else - // Empty bodies, not = default, so clang-tidy does not flag the guards as unused - UdpLock() {} - ~UdpLock() {} -#endif }; #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(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; @@ -65,22 +87,30 @@ 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; - for (auto *&p : this->queue) { +#else + this->sock = nullptr; +#endif + for (auto &p : this->queue) { if (p != nullptr) - pbuf_free(p); + 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(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); @@ -93,18 +123,25 @@ void ESPHomeOTAComponent::udp_send_ack(UdpLink &link, uint16_t prompted_by) { } 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(&to)->sin6_port = htons(link.peer_port); + } else { + reinterpret_cast(&to)->sin_port = htons(link.peer_port); + } + link.sock->sendto(ack, len, 0, reinterpret_cast(&to), link.peer_len); +#endif } -// 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); +void ESPHomeOTAComponent::udp_on_datagram(UdpLink &link, UdpSlot p, uint16_t port) { 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); + 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 @@ -118,27 +155,72 @@ void ESPHomeOTAComponent::udp_recv_cb(void *arg, struct udp_pcb * /*pcb*/, struc 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 && + UdpSlot &slot = link.queue[seq % UDP_WINDOW]; + if ((hdr[0] & ~UDP_FLAG_COMMITTED) == UDP_MSG_DATA && slot_len(p) > 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); + slot_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 +#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(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(&a), &ipv4_a); + const bool v4_b = socket::sockaddr_to_ipv4(reinterpret_cast(&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(&a)->sin6_addr, + &reinterpret_cast(&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(&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(&from)->sin6_port) + : ntohs(reinterpret_cast(&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(&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(&peer), &ipv4)) { ip_addr_set_ip4_u32(&link.peer_ip, ipv4); @@ -150,9 +232,6 @@ void ESPHomeOTAComponent::udp_open_(UdpLink &link) { 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); @@ -160,19 +239,40 @@ void ESPHomeOTAComponent::udp_open_(UdpLink &link) { 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) + if (link.pcb == nullptr || udp_bind(link.pcb, IP_ANY_TYPE, this->port_) != ERR_OK) { + ESP_LOGW(TAG, "UDP unavailable, using TCP"); 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(&local), sizeof(local), this->port_); + if (link.sock == nullptr || local_len == 0 || + link.sock->bind(reinterpret_cast(&local), local_len) != 0 || + link.sock->setblocking(false) != 0) { + ESP_LOGW(TAG, "UDP unavailable, using TCP"); + 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)); - 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 + // 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; @@ -190,6 +290,8 @@ bool ESPHomeOTAComponent::udp_start_() { 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->would_block_(errno)) @@ -204,7 +306,8 @@ ssize_t ESPHomeOTAComponent::udp_read_(uint8_t *buf, size_t len) { const uint32_t start = millis(); size_t got = 0; for (;;) { - struct pbuf *p; + this->udp_poll_(); + UdpSlot p; { UdpLock lock; p = link.consumed != link.received ? link.queue[link.consumed % UDP_WINDOW] : nullptr; @@ -219,12 +322,12 @@ ssize_t ESPHomeOTAComponent::udp_read_(uint8_t *buf, size_t len) { } // 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); + const uint16_t n = slot_copy(p, buf + got, std::min(len - got, slot_len(p) - offset), offset); got += n; link.read_offset += n; - if (UDP_HEADER_SIZE + link.read_offset >= p->tot_len) { + if (UDP_HEADER_SIZE + link.read_offset >= slot_len(p)) { UdpLock lock; - pbuf_free(p); + 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(link.received - link.consumed) == UDP_WINDOW; @@ -257,6 +360,7 @@ void ESPHomeOTAComponent::udp_linger_() { return; const uint32_t start = millis(); while (millis() - start < UDP_LINGER_MS) { + this->udp_poll_(); { UdpLock lock; if (this->udp_->out_len == 0) diff --git a/esphome/espota2.py b/esphome/espota2.py index 4c27df5e2a3..b6939d8c8d3 100644 --- a/esphome/espota2.py +++ b/esphome/espota2.py @@ -619,6 +619,8 @@ class UdpChannel: 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 def settimeout(self, timeout: float) -> None: self._timeout = timeout @@ -651,10 +653,14 @@ class UdpChannel: _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: + # 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() @@ -683,14 +689,19 @@ class UdpChannel: + (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 + 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: - raise OTANetworkError(f"no progress over UDP for {self._timeout:.0f}s") + 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 @@ -720,10 +731,16 @@ class UdpChannel: wait = 0 try: data = self._sock.recv(2048) - except OSError: + 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 + def _sample_rtt(self, sample: float) -> None: # RFC 6298 smoothing if self._srtt is None: @@ -1122,7 +1139,10 @@ def perform_ota( ) raise OTANetworkError(f"sending data: {err}") from err - # The UDP channel acks every message itself + # 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) @@ -1178,7 +1198,8 @@ def perform_ota( else: _LOGGER.info("OTA successful") if udp is not None: - # Nothing waits for this; just give the end acknowledgement a chance + # 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: diff --git a/tests/integration/test_host_ota.py b/tests/integration/test_host_ota.py index 17c4dea25f9..2102363b903 100644 --- a/tests/integration/test_host_ota.py +++ b/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 diff --git a/tests/unit_tests/test_espota2.py b/tests/unit_tests/test_espota2.py index f024a32b95b..6373c00ff8d 100644 --- a/tests/unit_tests/test_espota2.py +++ b/tests/unit_tests/test_espota2.py @@ -1628,20 +1628,20 @@ class _LossyUdpDevice: self.next_seq = (self.next_seq + 1) & 0xFFFF if self.random.random() < self.loss: continue - held = sum( + 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 - held |= (espota2.UDP_WINDOW - len(self.slots)) << 4 + 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([held]) + + bytes([window]) + bytes(self.out), peer, ) @@ -1670,16 +1670,6 @@ def test_udp_channel_delivers_in_order(loss: float) -> None: 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: @@ -1743,3 +1733,21 @@ def test_tcp_retransmits(platform: str, info: bytes, expected: int | None) -> No 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 dead link still fails after the no-progress timeout, naming the socket + error behind it.""" + 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(0.5) + with pytest.raises( + espota2.OTANetworkError, match="no progress over UDP.*last error" + ): + channel.sendall(bytes(espota2.UDP_MAX_PAYLOAD * 5)) + channel.close()