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()