From 3188e0d36f45f8828325194d570b043bbb4815d9 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Wed, 14 Oct 2026 17:30:56 -1000 Subject: [PATCH 1/2] [ota] Fall back to a UDP data phase when the link is lossy (#20498) --- .../components/esphome/ota/ota_esphome.cpp | 111 +++- esphome/components/esphome/ota/ota_esphome.h | 108 +++- .../esphome/ota/ota_esphome_noise.cpp | 12 +- .../esphome/ota/ota_esphome_udp.cpp | 375 +++++++++++ esphome/espota2.py | 605 +++++++++++++++--- tests/integration/test_host_ota.py | 94 ++- tests/unit_tests/test_espota2.py | 494 +++++++++++++- tests/unit_tests/test_espota2_noise.py | 54 +- 8 files changed, 1710 insertions(+), 143 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..df2aa4ee37a 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,10 @@ static constexpr uint8_t SERVER_FEATURE_SUPPORTS_NOISE = 0x04; // Raw deflate, window <= OTA_INFLATE_WINDOW_SIZE. Binding once offered: the // client must then send the image size frame and a deflate stream. static constexpr uint8_t SERVER_FEATURE_SUPPORTS_DEFLATE = 0x08; +#ifdef ESPHOME_OTA_UDP +static constexpr uint8_t SERVER_FEATURE_SUPPORTS_UDP = 0x10; +static constexpr uint8_t UDP_REQUEST_V1 = 0x01; // espota2.UDP_REQUEST +#endif #ifdef USE_OTA_ENCRYPTION inline bool ESPHomeOTAComponent::noise_offered_() const { @@ -266,6 +276,12 @@ void ESPHomeOTAComponent::handle_handshake_() { this->cleanup_connection_(); return; } + // Only unauthenticated handshakes reach here; the data phase blocks inside handle_data_ + if (now - this->client_connect_time_ > OTA_HANDSHAKE_PREEMPT && this->server_->ready()) { + ESP_LOGD(TAG, "Newer connection waiting"); + this->cleanup_connection_(); + return; + } switch (this->ota_state_) { case OTAState::MAGIC_READ: { @@ -329,6 +345,9 @@ void ESPHomeOTAComponent::handle_handshake_() { static_assert(HANDSHAKE_BUF_SIZE >= 2, "handshake_buf_ must hold the 2-byte extended-protocol feature ack"); this->handshake_buf_[0] = ota::OTA_RESPONSE_FEATURE_FLAGS; this->handshake_buf_[1] = (supports_compression ? SERVER_FEATURE_SUPPORTS_COMPRESSION : 0); +#ifdef ESPHOME_OTA_UDP + this->handshake_buf_[1] |= SERVER_FEATURE_SUPPORTS_UDP; +#endif #ifdef USE_OTA_PARTITIONS this->handshake_buf_[1] |= SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS; #endif @@ -483,6 +502,10 @@ void ESPHomeOTAComponent::handle_data_() { char *sbuf = reinterpret_cast(buf); size_t image_size; ota::OTAType ota_type = ota::OTA_TYPE_UPDATE_APP; +#ifdef ESPHOME_OTA_UDP + UdpLink udp_link; + this->udp_ = nullptr; // client_connect_time_ is done with +#endif // Set socket timeouts and blocking mode (see strategy table above) struct timeval tv; @@ -498,6 +521,38 @@ void ESPHomeOTAComponent::handle_data_() { // Acknowledge auth OK - 1 byte this->data_write_byte_(ota::OTA_RESPONSE_AUTH_OK); +#ifdef ESPHOME_OTA_UDP + // The server always offers UDP on the extended protocol + if (this->extended_proto_() && (this->ota_features_ & CLIENT_FEATURE_SUPPORTS_UDP) != 0) { + // Whether the client wants UDP (0: its link looks good), answered with our token + if (!this->data_readall_(buf, 1)) { + if (this->client_left_before_start_()) + return; + this->log_read_error_(LOG_STR("UDP request")); + goto error; // NOLINT(cppcoreguidelines-avoid-goto) + } + // 1 is this UDP format; other values are free for later ones, which we decline (a zero reply keeps TCP) + if (buf[0] == UDP_REQUEST_V1) { + this->udp_open_(udp_link); + if (udp_link.reply[0] == 0) { + ESP_LOGW(TAG, "UDP unavailable, using TCP"); + } + } + if (!this->data_write_(udp_link.reply, sizeof(udp_link.reply))) + goto error; // NOLINT(cppcoreguidelines-avoid-goto) + if (udp_link.reply[0] != 0) { + this->udp_ = &udp_link; + if (!this->udp_start_()) { + // Nothing started yet: just drop the session, which the client retries as a network failure + this->log_read_error_(LOG_STR("UDP start")); + this->cleanup_connection_(); + return; + } + } + } + +#endif + if (this->extended_proto_()) { // Read ota type, 1 byte if (!this->data_readall_(buf, 1)) { @@ -598,11 +653,17 @@ void ESPHomeOTAComponent::handle_data_() { // Acknowledge Update end OK - 1 byte this->data_write_byte_(ota::OTA_RESPONSE_UPDATE_END_OK); - // Read ACK - if (!this->data_readall_(buf, 1) || buf[0] != ota::OTA_RESPONSE_OK) { - this->log_read_error_(LOG_STR("ack")); - // do not go to error, this is not fatal - } +#ifdef ESPHOME_OTA_UDP + // UDP never reports a close: wait only until the client has our responses, not for its final ack + if (this->udp_ != nullptr) { + this->udp_linger_(); + } else +#endif + // Read ACK + if (!this->data_readall_(buf, 1) || buf[0] != ota::OTA_RESPONSE_OK) { + this->log_read_error_(LOG_STR("ack")); + // do not go to error, this is not fatal + } this->cleanup_connection_(); delay(10); @@ -623,6 +684,9 @@ void ESPHomeOTAComponent::handle_data_() { error: this->data_write_byte_(static_cast(error_code)); +#ifdef ESPHOME_OTA_UDP + this->udp_linger_(); +#endif // Abort backend before cleanup - cleanup_connection_() destroys the backend. // Always call abort() unconditionally: backends register external partitions before @@ -651,7 +715,7 @@ bool ESPHomeOTAComponent::readall_(uint8_t *buf, size_t len) { return false; } - ssize_t read = this->client_->read(buf + at, len - at); + ssize_t read = this->transport_read_(buf + at, len - at); if (read == -1) { const int err = errno; if (!this->would_block_(err)) { @@ -676,6 +740,10 @@ bool ESPHomeOTAComponent::readall_(uint8_t *buf, size_t len) { return true; } bool ESPHomeOTAComponent::writeall_(const uint8_t *buf, size_t len) { +#ifdef ESPHOME_OTA_UDP + if (this->udp_ != nullptr) + return this->udp_send_(buf, len); +#endif uint32_t start = millis(); uint32_t at = 0; while (len - at > 0) { @@ -728,7 +796,20 @@ void ESPHomeOTAComponent::log_read_error_(const LogString *what) { void ESPHomeOTAComponent::log_start_(const LogString *phase) { char peername[socket::SOCKADDR_STR_LEN]; - this->client_->getpeername_to(peername); +#ifdef ESPHOME_OTA_UDP + if (this->client_ == nullptr) { + // A committed UDP session has already dropped TCP +#ifdef ESPHOME_OTA_UDP_LWIP + ipaddr_ntoa_r(&this->udp_->peer_ip, peername, sizeof(peername)); +#else + socket::format_sockaddr_to(reinterpret_cast(&this->udp_->peer), this->udp_->peer_len, + peername); +#endif + } else +#endif + { + this->client_->getpeername_to(peername); + } ESP_LOGD(TAG, "Starting %s from %s", LOG_STR_ARG(phase), peername); } @@ -840,7 +921,7 @@ ssize_t ESPHomeOTAComponent::receive_data_(uint8_t *buf, DataTransfer &xfer) { break; } #endif - read = this->client_->read(buf, requested); + read = this->transport_read_(buf, requested); if (read > 0) break; if (read == 0) { @@ -874,6 +955,11 @@ ssize_t ESPHomeOTAComponent::receive_data_(uint8_t *buf, DataTransfer &xfer) { void ESPHomeOTAComponent::send_chunk_acks_(DataTransfer &xfer) { #if USE_OTA_VERSION == 2 +#ifdef ESPHOME_OTA_UDP + // UDP acks every message itself + if (this->udp_ != nullptr) + return; +#endif while (xfer.acknowledged + OTA_BLOCK_SIZE <= xfer.total || (xfer.total == xfer.ota_size && xfer.acknowledged < xfer.ota_size)) { this->data_write_byte_(ota::OTA_RESPONSE_CHUNK_OK); @@ -883,7 +969,12 @@ void ESPHomeOTAComponent::send_chunk_acks_(DataTransfer &xfer) { } void ESPHomeOTAComponent::cleanup_connection_() { - this->client_->close(); +#ifdef ESPHOME_OTA_UDP + this->udp_ = nullptr; // the link itself lives on handle_data_'s stack +#endif + // Already gone once a UDP session committed + if (this->client_ != nullptr) + this->client_->close(); this->client_ = nullptr; this->client_connect_time_ = 0; this->handshake_buf_pos_ = 0; diff --git a/esphome/components/esphome/ota/ota_esphome.h b/esphome/components/esphome/ota/ota_esphome.h index 910a5b3d99d..fbb1ba5447b 100644 --- a/esphome/components/esphome/ota/ota_esphome.h +++ b/esphome/components/esphome/ota/ota_esphome.h @@ -15,6 +15,25 @@ #include "esphome/core/preferences.h" #include "esphome/core/hash_base.h" +// UDP data phase: raw lwIP needs a main loop lock (or a single thread); the host polls a socket instead +#ifdef USE_ESP32 +#include +#endif +#if defined(USE_ESP8266) || defined(USE_RP2) || defined(USE_LIBRETINY) || \ + (defined(USE_ESP32) && defined(CONFIG_LWIP_TCPIP_CORE_LOCKING)) +#define ESPHOME_OTA_UDP +#define ESPHOME_OTA_UDP_LWIP +#include "lwip/ip_addr.h" +#include "lwip/pbuf.h" +#include "lwip/udp.h" +#if defined(USE_LIBRETINY) && !LWIP_TCPIP_CORE_LOCKING +#undef ESPHOME_OTA_UDP +#undef ESPHOME_OTA_UDP_LWIP +#endif +#elif defined(USE_HOST) +#define ESPHOME_OTA_UDP +#endif + namespace esphome { /// ESPHomeOTAComponent provides a simple way to integrate Over-the-Air updates into your app using ArduinoOTA. @@ -65,6 +84,69 @@ class ESPHomeOTAComponent final : public ota::OTAComponent { protected: void handle_handshake_(); void handle_data_(); + +#ifdef ESPHOME_OTA_UDP + // UDP transport for the data phase; wire format in ota_esphome_udp.cpp + static constexpr size_t UDP_OUT_SIZE = 64; // our unacked responses: at most three encrypted ones + // Messages the device keeps unread: like TCP's 4 segment window, so no more memory than TCP holds + static constexpr uint8_t UDP_WINDOW = 4; +#ifdef ESPHOME_OTA_UDP_LWIP + using UdpSlot = struct pbuf *; +#else + using UdpSlot = struct OtaUdpDatagram *; +#endif + // Lives on handle_data_'s stack; the receive callback shares it under the lwIP lock + struct UdpLink { + ~UdpLink(); + void release(); // closes the socket and frees what is queued +#ifdef ESPHOME_OTA_UDP_LWIP + struct udp_pcb *pcb{nullptr}; + ip_addr_t peer_ip; +#else + std::unique_ptr sock; + struct sockaddr_storage peer; // the TCP peer; its port is replaced by peer_port + socklen_t peer_len{0}; +#endif + uint16_t peer_port{0}; // learned from the client's datagrams + // The session token, our answer to the client: top bit clear, so never an error code; all zero declines + uint8_t reply[4]{}; + UdpSlot queue[UDP_WINDOW]{}; // slot seq % UDP_WINDOW for seq in [consumed, consumed + UDP_WINDOW) + uint16_t received{0}; // every seq before it is here + uint16_t consumed{0}; // the message being read + uint16_t read_offset{0}; // payload bytes of it already read + uint16_t out_base{0}; // stream offset of out[0] + uint8_t out_len{0}; + bool committed{false}; // the client heard us over UDP + uint8_t out[UDP_OUT_SIZE]; + }; + // Fills link.reply on success; leaves it all zero (declined) otherwise + void udp_open_(UdpLink &link); + // Waits for the client to commit to UDP or fall back to TCP; false on failure + bool udp_start_(); + // Like a blocking socket read with SO_RCVTIMEO: up to len bytes, or -1 with EWOULDBLOCK + ssize_t udp_read_(uint8_t *buf, size_t len); + bool udp_send_(const uint8_t *data, size_t len); + // In a UDP session, keeps answering until the client has our responses (an error code), briefly + void udp_linger_(); + // A datagram from the client's IP; caller holds the lwIP lock (or runs in the receive callback) + static void udp_on_datagram(UdpLink &link, UdpSlot p, uint16_t port); + // Caller holds the lwIP lock (or runs in the receive callback) + static void udp_send_ack(UdpLink &link, uint16_t prompted_by); +#ifdef ESPHOME_OTA_UDP_LWIP + static void udp_recv_cb(void *arg, struct udp_pcb *pcb, struct pbuf *p, const ip_addr_t *addr, uint16_t port); + void udp_poll_() {} +#else + void udp_poll_(); // hands what the socket received to udp_on_datagram +#endif +#endif + // Data phase only (udp_ shares storage with the handshake's connect time): TCP or, once committed, the UDP stream + inline ssize_t transport_read_(uint8_t *buf, size_t len) { +#ifdef ESPHOME_OTA_UDP + if (this->udp_ != nullptr) + return this->udp_read_(buf, len); +#endif + return this->client_->read(buf, len); + } #ifdef USE_OTA_PASSWORD static constexpr size_t SHA256_HEX_SIZE = 64; // SHA256 hash as hex string (32 bytes * 2) bool handle_auth_send_(); @@ -75,7 +157,6 @@ class ESPHomeOTAComponent final : public ota::OTAComponent { #endif // USE_OTA_PASSWORD bool readall_(uint8_t *buf, size_t len); bool writeall_(const uint8_t *buf, size_t len); - inline bool write_byte_(uint8_t byte) { return this->writeall_(&byte, 1); } #ifdef USE_OTA_ENCRYPTION // Heap-allocated only while an encrypted OTA session is active. @@ -105,18 +186,21 @@ class ESPHomeOTAComponent final : public ota::OTAComponent { ssize_t noise_read_frame_blocking_(uint8_t *buf, size_t min_ciphertext, size_t max_ciphertext); bool noise_readall_(uint8_t *buf, size_t len); ssize_t noise_read_data_(uint8_t *buf, size_t capacity); - bool noise_write_byte_(uint8_t byte); + static constexpr size_t NOISE_MAX_RESPONSE = 4; // the UDP token + bool noise_write_(const uint8_t *data, size_t len); #endif // USE_OTA_ENCRYPTION // Data-phase I/O dispatch: through the noise transport when a session is // active, straight to the socket otherwise. - inline bool data_write_byte_(uint8_t byte) { + // len is at most NOISE_MAX_RESPONSE + inline bool data_write_(const uint8_t *data, size_t len) { #ifdef USE_OTA_ENCRYPTION if (this->noise_ != nullptr) - return this->noise_write_byte_(byte); + return this->noise_write_(data, len); #endif - return this->write_byte_(byte); + return this->writeall_(data, len); } + inline bool data_write_byte_(uint8_t byte) { return this->data_write_(&byte, 1); } // When encrypted, buf must have room for len + noise::MAC_SIZE bytes. inline bool data_readall_(uint8_t *buf, size_t len) { #ifdef USE_OTA_ENCRYPTION @@ -200,7 +284,21 @@ class ESPHomeOTAComponent final : public ota::OTAComponent { std::unique_ptr client_; ota::OTABackendPtr backend_; +#ifdef ESPHOME_OTA_UDP_LWIP + union { + uint32_t client_connect_time_{0}; // handshake states + UdpLink *udp_; // data phase (never both), points into handle_data_'s stack + }; + static_assert(sizeof(UdpLink *) == sizeof(uint32_t), "clearing client_connect_time_ must clear udp_"); +#else uint32_t client_connect_time_{0}; +#ifdef ESPHOME_OTA_UDP + UdpLink *udp_{nullptr}; // data phase, points into handle_data_'s stack +#endif +#endif +#if defined(ESPHOME_OTA_UDP) && defined(USE_OTA_ENCRYPTION) + static_assert(sizeof(UdpLink::reply) <= NOISE_MAX_RESPONSE, "the token goes out through noise_write_"); +#endif static constexpr size_t HANDSHAKE_BUF_SIZE = 5; // Buffer size for OTA data transfer. The upload client derives its maximum // encrypted frame plaintext from this (espota2.NOISE_MAX_PLAINTEXT is this diff --git a/esphome/components/esphome/ota/ota_esphome_noise.cpp b/esphome/components/esphome/ota/ota_esphome_noise.cpp index a271ef9d39d..02c19ed7b7d 100644 --- a/esphome/components/esphome/ota/ota_esphome_noise.cpp +++ b/esphome/components/esphome/ota/ota_esphome_noise.cpp @@ -266,13 +266,15 @@ ssize_t ESPHomeOTAComponent::noise_read_data_(uint8_t *buf, size_t capacity) { return this->noise_read_frame_blocking_(buf, noise::MAC_SIZE + 1, max_ciphertext); } -/// Blocking write of one response byte as an encrypted frame. -bool ESPHomeOTAComponent::noise_write_byte_(uint8_t byte) { - uint8_t frame[noise::FRAME_HEADER_SIZE + 1 + noise::MAC_SIZE]; - frame[noise::FRAME_HEADER_SIZE] = byte; +/// Blocking write of a short response (at most NOISE_MAX_RESPONSE bytes) as one encrypted frame. +bool ESPHomeOTAComponent::noise_write_(const uint8_t *data, size_t len) { + uint8_t frame[noise::FRAME_HEADER_SIZE + NOISE_MAX_RESPONSE + noise::MAC_SIZE]; + if (len > NOISE_MAX_RESPONSE) + return false; + memcpy(frame + noise::FRAME_HEADER_SIZE, data, len); NoiseBuffer mbuf; noise_buffer_init(mbuf); - noise_buffer_set_inout(mbuf, frame + noise::FRAME_HEADER_SIZE, 1, 1 + noise::MAC_SIZE); + noise_buffer_set_inout(mbuf, frame + noise::FRAME_HEADER_SIZE, len, len + noise::MAC_SIZE); int err = noise_cipherstate_encrypt(this->noise_->send_cipher, &mbuf); if (err != 0) { ESP_LOGW(TAG, "Encrypt: %d", err); 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..6044284755c --- /dev/null +++ b/esphome/components/esphome/ota/ota_esphome_udp.cpp @@ -0,0 +1,375 @@ +#include "ota_esphome.h" +#if defined(USE_OTA) && defined(ESPHOME_OTA_UDP) + +#include "esphome/core/application.h" +#include "esphome/core/hal.h" +#include "esphome/core/helpers.h" +#include "esphome/core/log.h" + +#ifdef ESPHOME_OTA_UDP_LWIP +#include "lwip/inet.h" +#ifdef USE_LIBRETINY +#include "lwip/tcpip.h" +#endif +#endif + +#include +#include +#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 15 s probe +static constexpr uint32_t UDP_COMMIT_TIMEOUT = 80000; +// Like the TCP socket's SO_RCVTIMEO in handle_data_, so readall_ treats both alike +static constexpr uint32_t UDP_READ_WAIT_MS = 2000; +// How often the commit wait looks at the UDP link between reads of the TCP fallback byte +static constexpr uint32_t UDP_COMMIT_POLL_MS = 100; +static constexpr uint32_t UDP_LINGER_MS = 3000; + +#ifdef ESPHOME_OTA_UDP_LWIP +#ifdef USE_LIBRETINY +// LibreTiny's LwIPLock is a no-op; take the core lock directly (the header drops UDP without one) +struct UdpLock { + UdpLock() { LOCK_TCPIP_CORE(); } + ~UdpLock() { UNLOCK_TCPIP_CORE(); } +}; +#else +using UdpLock = LwIPLock; +#endif +static uint16_t slot_len(struct pbuf *p) { return p->tot_len; } +static uint16_t slot_copy(struct pbuf *p, void *dst, uint16_t len, uint16_t offset) { + return pbuf_copy_partial(p, dst, len, offset); +} +static void slot_free(struct pbuf *p) { pbuf_free(p); } +#else +// The host has no lwIP thread: the OTA loop drains the socket itself. Empty bodies, not = default, so clang-tidy +// does not flag the guards as unused +struct UdpLock { + UdpLock() {} + ~UdpLock() {} +}; +struct OtaUdpDatagram { + uint16_t len; + uint8_t data[1500]; +}; +static uint16_t slot_len(OtaUdpDatagram *p) { return p->len; } +static uint16_t slot_copy(OtaUdpDatagram *p, void *dst, uint16_t len, uint16_t offset) { + const uint16_t n = std::min(len, p->len - offset); + memcpy(dst, p->data + offset, n); + return n; +} +static void slot_free(OtaUdpDatagram *p) { delete p; } +#endif + +static void put_u16(uint8_t *at, uint16_t value) { + at[0] = value >> 8; + at[1] = value & 0xFF; +} + +ESPHomeOTAComponent::UdpLink::~UdpLink() { this->release(); } + +void ESPHomeOTAComponent::UdpLink::release() { + UdpLock lock; +#ifdef ESPHOME_OTA_UDP_LWIP + if (this->pcb != nullptr) + udp_remove(this->pcb); + this->pcb = nullptr; +#else + this->sock = nullptr; +#endif + for (auto &p : this->queue) { + if (p != nullptr) + slot_free(p); + p = nullptr; + } +} + +void ESPHomeOTAComponent::udp_send_ack(UdpLink &link, uint16_t prompted_by) { + const uint16_t len = UDP_ACK_HEADER_SIZE + link.out_len; +#ifdef ESPHOME_OTA_UDP_LWIP + struct pbuf *p = pbuf_alloc(PBUF_TRANSPORT, len, PBUF_RAM); + if (p == nullptr) + return; // Like a lost ACK: the client asks again + auto *ack = static_cast(p->payload); +#else + uint8_t ack[UDP_ACK_HEADER_SIZE + UDP_OUT_SIZE]; +#endif + ack[0] = UDP_MSG_ACK; + memcpy(ack + 1, link.reply, 4); + put_u16(ack + 5, link.received); + put_u16(ack + 7, prompted_by); + put_u16(ack + 9, link.out_base); + uint8_t window = (UDP_WINDOW - static_cast(link.received - link.consumed)) << 4; + for (uint8_t k = 1; k < UDP_WINDOW; k++) { + if (link.queue[(link.received + k) % UDP_WINDOW] != nullptr) + window |= 1 << k; + } + ack[11] = window; + memcpy(ack + UDP_ACK_HEADER_SIZE, link.out, link.out_len); + // Best effort: a lost ACK is repeated for the client's next datagram +#ifdef ESPHOME_OTA_UDP_LWIP + udp_sendto(link.pcb, p, &link.peer_ip, link.peer_port); + pbuf_free(p); +#else + struct sockaddr_storage to = link.peer; + if (to.ss_family == AF_INET6) { + reinterpret_cast(&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 +} + +void ESPHomeOTAComponent::udp_on_datagram(UdpLink &link, UdpSlot p, uint16_t port) { + uint8_t hdr[UDP_HEADER_SIZE]; + if (slot_copy(p, hdr, UDP_HEADER_SIZE, 0) != UDP_HEADER_SIZE || memcmp(hdr + 1, link.reply, 4) != 0) { + slot_free(p); + return; + } + // The token proves the sender; answer the port it arrived from, which NAT may have rewritten + link.peer_port = port; + link.committed |= (hdr[0] & UDP_FLAG_COMMITTED) != 0; + // Drop the responses the client already has + const uint16_t acked = encode_uint16(hdr[7], hdr[8]) - link.out_base; + if (acked != 0 && acked <= link.out_len) { + link.out_len -= acked; + memmove(link.out, link.out + acked, link.out_len); + link.out_base += acked; + } + const uint16_t seq = encode_uint16(hdr[5], hdr[6]); + UdpSlot &slot = link.queue[seq % UDP_WINDOW]; + if ((hdr[0] & ~UDP_FLAG_COMMITTED) == UDP_MSG_DATA && slot_len(p) > UDP_HEADER_SIZE && + static_cast(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 { + slot_free(p); + } + udp_send_ack(link, seq); +} + +#ifdef ESPHOME_OTA_UDP_LWIP +// lwIP context: the tcpip thread (ESP32, LibreTiny), SYS (ESP8266) or an IRQ (RP2040) +void ESPHomeOTAComponent::udp_recv_cb(void *arg, struct udp_pcb * /*pcb*/, struct pbuf *p, const ip_addr_t *addr, + uint16_t port) { + auto &link = *static_cast(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); + } 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 + } + UdpLock lock; +#if LWIP_IPV6 + link.pcb = udp_new_ip_type(IPADDR_TYPE_ANY); +#else + link.pcb = udp_new(); +#endif + // On failure the reply stays all zero and ~UdpLink removes the pcb + if (link.pcb == nullptr || udp_bind(link.pcb, IP_ANY_TYPE, this->port_) != ERR_OK) + return; + udp_recv(link.pcb, udp_recv_cb, &link); +#else + link.peer = peer; + link.peer_len = peer_len; + // Same domain as the listener, so the peer address from getpeername fits + link.sock = socket::socket_ip(SOCK_DGRAM, IPPROTO_UDP); + struct sockaddr_storage local; + const socklen_t local_len = + socket::set_sockaddr_any(reinterpret_cast(&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) + return; +#endif + // Below the error codes (0x80 and up) the client checks the first byte against, and never all zero + token[0] = (token[0] & 0x7F) | 0x01; + memcpy(link.reply, token, sizeof(token)); +} + +bool ESPHomeOTAComponent::udp_start_() { + // TCP now only carries a fallback byte; short read timeouts let the UDP link be checked in between + struct timeval tv { + 0, UDP_COMMIT_POLL_MS * 1000 + }; + this->client_->setsockopt(SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)); + const uint32_t start = millis(); + while (millis() - start < UDP_COMMIT_TIMEOUT) { + this->udp_poll_(); + bool committed; + { + UdpLock lock; + committed = this->udp_->committed; + } + if (committed) { + // The client heard us over UDP, so TCP is done. Outside the lock: closing a socket waits on the tcpip thread + this->client_ = nullptr; + return true; + } + // A raw byte outside any Noise frame: the client heard nothing over UDP. Forging it only forces the ordinary + // TCP path, which dropping our datagrams forces anyway + uint8_t fallback; + const ssize_t read = this->client_->read(&fallback, 1); + if (read == 1) { + this->udp_->release(); + this->udp_ = nullptr; + tv = {UDP_READ_WAIT_MS / 1000, 0}; + this->client_->setsockopt(SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)); + return true; + } + if (read == 0) { + this->remote_closed_ = true; + return false; + } + if (!this->would_block_(errno)) + return false; + App.feed_wdt(); + } + return false; +} + +ssize_t ESPHomeOTAComponent::udp_read_(uint8_t *buf, size_t len) { + UdpLink &link = *this->udp_; + const uint32_t start = millis(); + size_t got = 0; + for (;;) { + this->udp_poll_(); + UdpSlot p; + { + UdpLock lock; + p = link.consumed != link.received ? link.queue[link.consumed % UDP_WINDOW] : nullptr; + } + if (p == nullptr) { + if (got != 0) + return got; + if (millis() - start >= UDP_READ_WAIT_MS) + break; + this->yield_and_feed_watchdog_(); + continue; + } + // The callback never touches the slot being read, so the copy runs without the lock + const uint16_t offset = UDP_HEADER_SIZE + link.read_offset; + const uint16_t n = slot_copy(p, buf + got, std::min(len - got, slot_len(p) - offset), offset); + got += n; + link.read_offset += n; + if (UDP_HEADER_SIZE + link.read_offset >= slot_len(p)) { + UdpLock lock; + slot_free(p); + link.queue[link.consumed % UDP_WINDOW] = nullptr; + // A full window just opened: tell the client, which is holding back + const bool was_full = static_cast(link.received - link.consumed) == UDP_WINDOW; + link.consumed++; + link.read_offset = 0; + if (was_full) + udp_send_ack(link, link.received - 1); + } + if (got == len) + return got; + } + errno = EWOULDBLOCK; + return -1; +} + +bool ESPHomeOTAComponent::udp_send_(const uint8_t *data, size_t len) { + UdpLink &link = *this->udp_; + UdpLock lock; + if (link.out_len + len > UDP_OUT_SIZE) + return false; + memcpy(link.out + link.out_len, data, len); + link.out_len += len; + // Prompted by nothing new, so the client takes no round trip sample from it + udp_send_ack(link, link.received - 1); + return true; +} + +void ESPHomeOTAComponent::udp_linger_() { + if (this->udp_ == nullptr) + return; + const uint32_t start = millis(); + while (millis() - start < UDP_LINGER_MS) { + this->udp_poll_(); + { + UdpLock lock; + if (this->udp_->out_len == 0) + return; + } + this->yield_and_feed_watchdog_(); + } +} + +} // namespace esphome +#endif // USE_OTA && ESPHOME_OTA_UDP diff --git a/esphome/espota2.py b/esphome/espota2.py index 14e163676d1..aaef178753a 100644 --- a/esphome/espota2.py +++ b/esphome/espota2.py @@ -70,11 +70,13 @@ CLIENT_FEATURE_SUPPORTS_SHA256_AUTH = 0x02 CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL = 0x04 CLIENT_FEATURE_SUPPORTS_NOISE = 0x08 CLIENT_FEATURE_SUPPORTS_DEFLATE = 0x10 +CLIENT_FEATURE_SUPPORTS_UDP = 0x20 SERVER_FEATURE_SUPPORTS_COMPRESSION = 0x01 SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS = 0x02 SERVER_FEATURE_SUPPORTS_NOISE = 0x04 # Binding once offered: the device then expects the image size and a deflate stream SERVER_FEATURE_SUPPORTS_DEFLATE = 0x08 +SERVER_FEATURE_SUPPORTS_UDP = 0x10 # Wire constant: the deflate bit promises a 4 KB window (OTA_INFLATE_WINDOW_SIZE) DEFLATE_WINDOW_BITS = 12 @@ -103,17 +105,64 @@ COMPRESS_LEVEL = 9 UPLOAD_BUFFER_SIZE = UPLOAD_BLOCK_SIZE * 8 # Flaky Wi-Fi links often drop the first OTA attempt, and the device may need time -# to clean up a half-open connection (its handshake watchdog runs at 20s) before it +# to clean up a half-open connection (its handshake watchdog runs at 60s) before it # accepts a new one, so wait between attempts instead of failing the upload outright. # Every resolved address is tried once, and this many extra attempts are shared # across the addresses on top of that. -EXTRA_UPLOAD_ATTEMPTS = 2 +EXTRA_UPLOAD_ATTEMPTS = 4 +# Once a connection has reached the device, a lossy link (not a dead device) is +# the likely cause, so keep trying like an API reconnect would +EXTRA_UPLOAD_ATTEMPTS_REACHED = 8 UPLOAD_RETRY_DELAY = 5.0 +# A SYN getting through is the cheap step; an unreachable address still fails fast +CONNECT_TIMEOUT = 20.0 +# Handshake waits; as patient as an API connect on a lossy link, and matching the +# device's OTA_SOCKET_TIMEOUT_HANDSHAKE +SETUP_TIMEOUT = 60.0 # Data phase timeout; must stay longer than the device's OTA_SOCKET_TIMEOUT_DATA # (105 s) so a stalled session is gone before a retry, and long enough for lwIP # to get a lost chunk ack through after the retransmit run seen in practice DATA_PHASE_TIMEOUT = 160.0 +# UDP data phase; wire format in ota_esphome_udp.cpp +UDP_MSG_DATA = 0x01 +UDP_MSG_PING = 0x02 +UDP_MSG_PROBE = 0x03 +UDP_FLAG_COMMITTED = 0x80 +UDP_MSG_ACK = 0x40 +UDP_HEADER_SIZE = 9 +UDP_ACK_HEADER_SIZE = 12 +# Stream bytes per datagram, within one Ethernet frame +UDP_MAX_PAYLOAD = 1400 +# Messages the device queues unconsumed (UDP_WINDOW in ota_esphome.h); as many +# as TCP's window holds, so UDP costs the device no more memory than TCP +UDP_WINDOW = 4 +# Probe this long before falling back to TCP, enough to ride out a radio busy +# with BLE; the device waits 80 s +UDP_PROBE_TIMEOUT = 15.0 +UDP_PROBE_INTERVAL = 0.1 +# Resend timeout bounds around srtt + 4 * rttvar; random loss never grows it +UDP_RESEND_INITIAL = 0.5 +UDP_RESEND_MIN = 0.2 +UDP_RESEND_MAX = 2.0 +# Acks still asking for the head after later messages arrived; resend it once +UDP_DUP_ACKS = 2 +# Progress acks in a row with no round trip to time: the link is slower than +# the resend timer, so back it off (random loss still yields samples) +UDP_UNTIMED_ACKS = 8 +# Refusals in a row with no ACK between: the device's UDP port is closed +UDP_REFUSALS = 3 +UDP_FALLBACK = b"\x00" +# The UDP format this client speaks; devices decline values they do not know +UDP_REQUEST = b"\x01" +# A plain handshake exchange slower than this hid a retransmit (the device's +# lwIP waits about a second before resending); UDP is used from then on +UDP_SLOW_EXCHANGE = 0.75 +# getsockopt(IPPROTO_TCP) option, size, offset and format of the retransmit +# count: tcp_info.tcpi_total_retrans (Linux), tcp_connection_info. +# tcpi_txretransmitbytes (macOS) +_TCP_RETRANS_OPTION = {"linux": (11, 104, 100, "I"), "darwin": (0x106, 80, 72, "Q")} + _LOGGER = logging.getLogger(__name__) # Authentication method lookup table: response -> (hash_func, nonce_size, name) @@ -511,6 +560,10 @@ class NoiseSocketWrapper: self._decrypt = None self._buffer = b"" + def use_transport(self, transport: UdpChannel) -> None: + """Continue the session over another byte stream.""" + self._sock = transport + # Only harmless socket controls pass through; byte-moving methods are # deliberately absent so plaintext cannot leak past the transport. def settimeout(self, timeout: float | None) -> None: @@ -613,19 +666,316 @@ class NoiseSocketWrapper: return data +class _UdpMessage: + """One datagram of the UDP stream and its send history.""" + + __slots__ = ("held", "last_send", "payload", "resend", "sends", "seq") + + def __init__(self, seq: int, payload: bytes) -> None: + self.seq = seq + self.payload = payload + self.last_send = 0.0 + self.sends = 0 + self.resend = False # duplicate acks asked for it + self.held = False # the device keeps it, waiting for an earlier one + + +class UdpChannel: + """Byte stream over numbered UDP datagrams, standing in for the socket. + + The device acks every datagram on arrival and says how many it has room + for; each is resent after TCP's srtt + 4 * rttvar, so loss slows the + transfer instead of ending it. It backs off only while acks come back + untimed, i.e. the link is slower than the timer, never on random loss. + """ + + def __init__(self, sock: socket.socket, token: bytes) -> None: + self._sock = sock + self._token = token + self._timeout = DATA_PHASE_TIMEOUT + self._next_seq = 0 + self._queue: list[_UdpMessage] = [] # not yet received, oldest first + self._rx = bytearray() + self._device_bytes = 0 + self._committed = False + self._srtt: float | None = None + self._rttvar = 0.0 + # Doubles while acks come back untimed, so a link slower than the + # timer still yields messages sent once to time (Karn, RFC 6298) + self._backoff = 1 + self._untimed = 0 + self._dup_acks = 0 + self._free = UDP_WINDOW # the device's empty slots, as of its last ACK + self._last_ping = 0.0 + self._last_progress = time.monotonic() + self._sends = 0 + # Last socket error, e.g. refused once the device's socket is gone + self._last_error: OSError | None = None + self._refusals = 0 + + def settimeout(self, timeout: float) -> None: + self._timeout = timeout + + def setsockopt(self, level: int, optname: int, value: int) -> None: + """TCP tuning does not apply.""" + + def close(self) -> None: + _LOGGER.debug( + "UDP: %d messages in %d datagrams, smoothed RTT %.0f ms", + self._next_seq, + self._sends, + (self._srtt or 0) * 1000, + ) + self._sock.close() + + def probe(self) -> bool: + """True once the device answers over UDP within UDP_PROBE_TIMEOUT.""" + deadline = time.monotonic() + UDP_PROBE_TIMEOUT + try: + while (now := time.monotonic()) < deadline: + self._send(UDP_MSG_PROBE, 0, b"") + self._receive(min(UDP_PROBE_INTERVAL, deadline - now)) + if self._committed: + return True + except OTANetworkError: + pass # rejected, e.g. by a firewall + return False + + def sendall(self, data: bytes) -> None: + for offset in range(0, len(data), UDP_MAX_PAYLOAD): + self._queue.append( + _UdpMessage(self._next_seq, data[offset : offset + UDP_MAX_PAYLOAD]) + ) + self._next_seq = (self._next_seq + 1) & 0xFFFF + # Return once everything is at least in flight, or the device spoke (an error) + while len(self._queue) > UDP_WINDOW and not self._rx: + self._pump() + + def pending(self) -> bool: + """Whether the device sent bytes nobody has read yet.""" + return bool(self._rx) + + def recv(self, amount: int) -> bytes: + while not self._rx: + self._pump() + data = bytes(self._rx[:amount]) + del self._rx[:amount] + return data + + def flush(self, timeout: float) -> None: + """Best effort: wait until the device has every message.""" + deadline = time.monotonic() + timeout + with contextlib.suppress(OTANetworkError): + while self._queue and time.monotonic() < deadline: + self._pump() + + def _rto(self) -> float: + if self._srtt is None: + rto = UDP_RESEND_INITIAL + else: + rto = max(self._srtt + 4 * self._rttvar, UDP_RESEND_MIN) + return min(rto * self._backoff, UDP_RESEND_MAX) + + def _send(self, msg_type: int, seq: int, payload: bytes) -> None: + flags = UDP_FLAG_COMMITTED if self._committed else 0 + header = ( + bytes([msg_type | flags]) + + self._token + + seq.to_bytes(2, "big") + + (self._device_bytes & 0xFFFF).to_bytes(2, "big") + ) + self._sends += 1 + try: + self._sock.send(header + payload) + except OSError as err: + # A full buffer or ICMP error is just more loss, but worth reporting + self._socket_error(err) + + def _pump(self) -> None: + now = time.monotonic() + if now - self._last_progress > self._timeout: + reason = f" (last error: {self._last_error})" if self._last_error else "" + raise OTANetworkError( + f"no progress over UDP for {self._timeout:.0f}s{reason}" + ) + rto = self._rto() + wake = now + rto + # At least the head, which doubles as the probe while the device is busy + in_flight = self._queue[: max(self._free, 1)] + for message in in_flight: + if message.held: + continue + if message.sends == 0 or message.resend or now - message.last_send >= rto: + message.last_send = now + message.sends += 1 + message.resend = False + self._send(UDP_MSG_DATA, message.seq, message.payload) + wake = min(wake, message.last_send + rto) + if all(message.held for message in in_flight): + # Nothing to send: ask for the device's responses + if now - self._last_ping >= rto: + self._last_ping = now + self._send(UDP_MSG_PING, 0, b"") + wake = self._last_ping + rto + self._receive(wake - now) + + def _receive(self, wait: float) -> None: + import select + + # Wait for the first datagram, then drain whatever else arrived + while select.select([self._sock], [], [], max(wait, 0))[0]: + wait = 0 + try: + data = self._sock.recv(2048) + except OSError as err: + self._socket_error(err) + continue + self._handle_ack(data) + + def _socket_error(self, err: OSError) -> None: + if self._last_error is None or err.errno != self._last_error.errno: + _LOGGER.debug("UDP socket error: %s", err) + self._last_error = err + # Windows reports the ICMP port unreachable as a reset + if isinstance(err, (ConnectionRefusedError, ConnectionResetError)): + self._refusals += 1 + if self._refusals >= UDP_REFUSALS: + raise OTANetworkError(f"device closed its UDP port: {err}") + + def _sample_rtt(self, sample: float) -> None: + # RFC 6298 smoothing + self._backoff = 1 + self._untimed = 0 + if self._srtt is None: + self._srtt, self._rttvar = sample, sample / 2 + else: + self._rttvar = 0.75 * self._rttvar + 0.25 * abs(self._srtt - sample) + self._srtt = 0.875 * self._srtt + 0.125 * sample + + def _handle_ack(self, data: bytes) -> None: + if ( + len(data) < UDP_ACK_HEADER_SIZE + or data[0] != UDP_MSG_ACK + or data[1:5] != self._token + ): + return + now = time.monotonic() + self._committed = True + self._refusals = 0 + received = int.from_bytes(data[5:7], "big") + prompted = int.from_bytes(data[7:9], "big") + self._free = data[11] >> 4 + if self._queue: + head = self._queue[0].seq + index = (prompted - head) & 0xFFFF + # This datagram's own round trip; only from messages sent once (Karn) + timed = index < len(self._queue) and self._queue[index].sends == 1 + if timed: + self._sample_rtt(now - self._queue[index].last_send) + done = (received - head) & 0xFFFF + if 0 < done <= len(self._queue): + del self._queue[:done] + self._last_progress = now + self._dup_acks = 0 + if not timed: + self._untimed += 1 + if self._untimed >= UDP_UNTIMED_ACKS: + self._untimed = 0 + if self._rto() < UDP_RESEND_MAX: + self._backoff *= 2 + elif done == 0 and self._queue[0].sends: + # A later message arrived while the head is missing; resend it once + self._dup_acks += 1 + if self._dup_acks == UDP_DUP_ACKS: + self._queue[0].resend = True + # Rebuilt from every ACK, so a reordered one cannot leave a stale mark + for k, message in enumerate(self._queue[:UDP_WINDOW]): + message.held = (received + k - message.seq) & 0xFFFF == 0 and bool( + data[11] & (1 << k) + ) + skip = (self._device_bytes - int.from_bytes(data[9:11], "big")) & 0xFFFF + fresh = data[UDP_ACK_HEADER_SIZE + skip :] + if skip <= len(data) - UDP_ACK_HEADER_SIZE and fresh: + self._rx += fresh + self._device_bytes += len(fresh) + self._last_progress = now + + +def _tcp_retransmits(sock: socket.socket) -> int | None: + """Segments this side retransmitted, where the OS reports it.""" + import struct + import sys + + if (option := _TCP_RETRANS_OPTION.get(sys.platform)) is None: + return None + optname, size, offset, fmt = option + try: + info = sock.getsockopt(socket.IPPROTO_TCP, optname, size) + except (OSError, TypeError): + return None + if len(info) < offset + struct.calcsize(fmt): + return None + return struct.unpack_from(fmt, info, offset)[0] + + +def _link_lossy( + tcp_sock: socket.socket, slowest_exchange: float, prefer_udp: bool +) -> bool: + """Whether the handshake showed loss, so the data phase should use UDP.""" + return ( + prefer_udp + or slowest_exchange > UDP_SLOW_EXCHANGE + or bool(_tcp_retransmits(tcp_sock)) + ) + + +def _start_udp( + sock: socket.socket | NoiseSocketWrapper, tcp_sock: socket.socket, lossy: bool +) -> UdpChannel | None: + """Move the data phase to UDP; None keeps it on TCP.""" + # A clean link asks to stay on TCP + send_check(sock, UDP_REQUEST if lossy else UDP_FALLBACK, "UDP request") + token = receive_exactly(sock, 4, "UDP token", None, decode=False) + if not lossy: + _LOGGER.debug("Link looks clean, staying on TCP") + return None + if token == bytes(4): + _LOGGER.debug("Device declined UDP") + return None + udp_sock = socket.socket(tcp_sock.family, socket.SOCK_DGRAM) + try: + udp_sock.connect(tcp_sock.getpeername()) + except OSError as err: + udp_sock.close() + _LOGGER.info("UDP unavailable (%s), using TCP", err) + else: + channel = UdpChannel(udp_sock, token) + if channel.probe(): + _LOGGER.info("Using UDP") + return channel + channel.close() + _LOGGER.info("UDP blocked, using TCP") + # Raw, outside any Noise frame: forging it only forces the TCP path + send_check(tcp_sock, UDP_FALLBACK, "UDP fallback") + return None + + def _negotiate_session( sock: socket.socket, noise_psk: str | None, plaintext_fallback: bool = False, allow_plaintext_upload: bool = False, -) -> tuple[socket.socket | NoiseSocketWrapper, int, int, bool]: +) -> tuple[socket.socket | NoiseSocketWrapper, int, int, bool, float]: """Negotiate through the Noise handshake; returns the socket to go on with - (wrapped once encrypted), version, features, extended.""" + (wrapped once encrypted), version, features, extended and the slowest + plaintext exchange in seconds.""" # Enable nodelay, we need it for phase 1 sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + exchange_start = time.monotonic() send_check(sock, MAGIC_BYTES, "magic bytes") _, version = receive_exactly(sock, 2, "version", RESPONSE_OK) + slowest_exchange = time.monotonic() - exchange_start _LOGGER.debug("Device supports OTA version %s", version) supported_versions = (OTA_VERSION_1_0, OTA_VERSION_2_0) if version not in supported_versions: @@ -639,9 +989,11 @@ def _negotiate_session( | CLIENT_FEATURE_SUPPORTS_SHA256_AUTH | CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL | CLIENT_FEATURE_SUPPORTS_DEFLATE + | CLIENT_FEATURE_SUPPORTS_UDP ) if noise_psk: features_to_send |= CLIENT_FEATURE_SUPPORTS_NOISE + exchange_start = time.monotonic() send_check(sock, features_to_send, "features") features = receive_exactly( sock, @@ -649,6 +1001,7 @@ def _negotiate_session( "features", None, # Accept any response )[0] + slowest_exchange = max(slowest_exchange, time.monotonic() - exchange_start) extended_proto = False if features == RESPONSE_FEATURE_FLAGS: @@ -705,7 +1058,7 @@ def _negotiate_session( _LOGGER.debug("Encrypted connection established") if allow_plaintext_upload: _LOGGER.warning(ALLOW_PLAINTEXT_UPLOAD_REMOVE_WARNING) - return sock, version, features, extended_proto + return sock, version, features, extended_proto, slowest_exchange def perform_ota( @@ -718,6 +1071,7 @@ def perform_ota( plaintext_fallback: bool = False, allow_plaintext_upload: bool = False, on_upload_start: Callable[[], None] | None = None, + prefer_udp: bool = False, ) -> None: # Validate up front; an out-of-range value would only surface as a # ValueError deep inside send_check, bypassing OTAError handling @@ -740,7 +1094,9 @@ def perform_ota( file_size = len(file_contents) _LOGGER.info("Uploading %s (%s bytes)", filename, file_size) - sock, version, features, extended_proto = _negotiate_session( + # sock may become a Noise wrapper or a UDP channel; UDP needs the peer + tcp_sock = sock + sock, version, features, extended_proto, slowest_exchange = _negotiate_session( sock, noise_psk, plaintext_fallback, allow_plaintext_upload ) _LOGGER.info("Connection established; device supports OTA version %s", version) @@ -832,117 +1188,153 @@ def perform_ota( if on_upload_start is not None: on_upload_start() # the device may commit anything sent from here + udp: UdpChannel | None = None + # Still under SETUP_TIMEOUT, which the device's commit wait is sized for + if extended_proto and features & SERVER_FEATURE_SUPPORTS_UDP: + udp = _start_udp( + sock, tcp_sock, _link_lossy(tcp_sock, slowest_exchange, prefer_udp) + ) + if udp is not None: + # The device drops TCP once it hears the committed channel + if isinstance(sock, NoiseSocketWrapper): + sock.use_transport(udp) + else: + sock = udp + sock.settimeout(DATA_PHASE_TIMEOUT) - - if extended_proto: - send_check(sock, ota_type, "ota type") - - upload_size = len(upload_contents) - # The device erases flash between receiving the size and acking the - # prepare, so this window shows the erase cost (near zero when the - # device erases lazily during the upload) - prepare_start = time.perf_counter() - send_check(sock, upload_size.to_bytes(SIZE_FIELD_BYTES, "big"), "binary size") - if deflate: - # Own frame: an encrypted session carries one field per frame - send_check(sock, file_size.to_bytes(SIZE_FIELD_BYTES, "big"), "image size") - receive_exactly(sock, 1, "update prepare result", RESPONSE_UPDATE_PREPARE_OK) - prepare_duration = time.perf_counter() - prepare_start - _LOGGER.info("Preparing for upload took %.2f seconds", prepare_duration) - - # The device hashes what it writes: the inflated image, else the received bytes - upload_md5 = hashlib.md5(file_contents if deflate else upload_contents).hexdigest() - _LOGGER.debug("MD5 of upload is %s", upload_md5) - - send_check(sock, upload_md5, "file checksum") - receive_exactly(sock, 1, "file checksum result", RESPONSE_BIN_MD5_OK) - - # Disable nodelay for transfer - sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 0) - # Limit send buffer (usually around 100kB) in order to have progress bar - # show the actual progress - - sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, UPLOAD_BUFFER_SIZE) - start_time = time.perf_counter() - - offset = 0 - progress = ProgressBar("Uploading") try: - while True: - chunk = upload_contents[offset : offset + UPLOAD_BLOCK_SIZE] - if not chunk: - break - offset += len(chunk) + if extended_proto: + send_check(sock, ota_type, "ota type") - try: - sock.sendall(chunk) - except OSError as err: - # A send failure can hide an error byte the device reported - # just before dropping the connection; surface that as the - # real, non-retryable cause when it is available - try: - sock.settimeout(1.0) - check_error(recv_decode(sock, 1), None) - except (OSError, OTANetworkError) as probe_err: - _LOGGER.debug( - "No device error behind the send failure: %s", probe_err - ) - raise OTANetworkError(f"sending data: {err}") from err + upload_size = len(upload_contents) + # The device erases flash between receiving the size and acking the + # prepare, so this window shows the erase cost (near zero when the + # device erases lazily during the upload) + prepare_start = time.perf_counter() + send_check(sock, upload_size.to_bytes(SIZE_FIELD_BYTES, "big"), "binary size") + if deflate: + # Own frame: an encrypted session carries one field per frame + send_check(sock, file_size.to_bytes(SIZE_FIELD_BYTES, "big"), "image size") + receive_exactly(sock, 1, "update prepare result", RESPONSE_UPDATE_PREPARE_OK) + prepare_duration = time.perf_counter() - prepare_start + _LOGGER.info("Preparing for upload took %.2f seconds", prepare_duration) + + # The device hashes what it writes: the inflated image, else the received bytes + upload_md5 = hashlib.md5( + file_contents if deflate else upload_contents + ).hexdigest() + _LOGGER.debug("MD5 of upload is %s", upload_md5) + + send_check(sock, upload_md5, "file checksum") + receive_exactly(sock, 1, "file checksum result", RESPONSE_BIN_MD5_OK) + + # Disable nodelay for transfer + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 0) + # Limit send buffer (usually around 100kB) in order to have progress bar + # show the actual progress + + sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, UPLOAD_BUFFER_SIZE) + start_time = time.perf_counter() + + offset = 0 + progress = ProgressBar("Uploading") + try: + while True: + chunk = upload_contents[offset : offset + UPLOAD_BLOCK_SIZE] + if not chunk: + break + offset += len(chunk) - if version >= OTA_VERSION_2_0: try: - receive_exactly(sock, 1, "chunk result", RESPONSE_CHUNK_OK) + sock.sendall(chunk) + except OSError as err: + # A send failure can hide an error byte the device reported + # just before dropping the connection; surface that as the + # real, non-retryable cause when it is available + try: + sock.settimeout(1.0) + check_error(recv_decode(sock, 1), None) + except (OSError, OTANetworkError) as probe_err: + _LOGGER.debug( + "No device error behind the send failure: %s", probe_err + ) + raise OTANetworkError(f"sending data: {err}") from err except OTANetworkError as err: + if udp is None: + raise + # Over UDP: surface a device error code if one arrived, and + # treat a failure on the last chunk like a lost final ack + if udp.pending(): + receive_exactly(sock, 1, "data", []) if offset < upload_size: raise - # The device already had the complete image when this ack - # was lost, so it may be committing; do not retry raise _committed_error(err) from err - progress.update(offset / upload_size) - except OTAError: - # Terminate the progress bar line before the error is logged + # The UDP channel acks every message itself, so anything the device + # sends mid upload is an error code + if udp is not None and udp.pending() and offset < upload_size: + receive_exactly(sock, 1, "data", []) + if version >= OTA_VERSION_2_0 and udp is None: + try: + receive_exactly(sock, 1, "chunk result", RESPONSE_CHUNK_OK) + except OTANetworkError as err: + if offset < upload_size: + raise + # The device already had the complete image when this ack + # was lost, so it may be committing; do not retry + raise _committed_error(err) from err + + progress.update(offset / upload_size) + except OTAError: + # Terminate the progress bar line before the error is logged + progress.done() + raise progress.done() - raise - progress.done() - # Enable nodelay for last checks - sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) - duration = time.perf_counter() - start_time + # Enable nodelay for last checks + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + duration = time.perf_counter() - start_time - _LOGGER.info("Upload took %.2f seconds, waiting for result...", duration) + _LOGGER.info("Upload took %.2f seconds, waiting for result...", duration) - # Once the device has the complete image it commits the update and - # reboots on its own; the exact commit point is not observable from - # here, so treat everything past the data phase as non-retryable. A - # re-upload could flash a device that already updated successfully. - commit_start = time.perf_counter() - try: - receive_exactly(sock, 1, "update receive result", RESPONSE_RECEIVE_OK) - receive_exactly(sock, 1, "update end result", RESPONSE_UPDATE_END_OK) - except OTANetworkError as err: - raise _committed_error(err) from err - commit_duration = time.perf_counter() - commit_start + # Once the device has the complete image it commits the update and + # reboots on its own; the exact commit point is not observable from + # here, so treat everything past the data phase as non-retryable. A + # re-upload could flash a device that already updated successfully. + commit_start = time.perf_counter() + try: + receive_exactly(sock, 1, "update receive result", RESPONSE_RECEIVE_OK) + receive_exactly(sock, 1, "update end result", RESPONSE_UPDATE_END_OK) + except OTANetworkError as err: + raise _committed_error(err) from err + commit_duration = time.perf_counter() - commit_start - # Sum of the named windows so the breakdown is self consistent; connect, - # handshake, auth, and the one MD5 round trip are not included - _LOGGER.info( - "Update took %.2f seconds (prepare %.2f, upload %.2f, commit %.2f)", - prepare_duration + duration + commit_duration, - prepare_duration, - duration, - commit_duration, - ) + # Sum of the named windows so the breakdown is self consistent; connect, + # handshake, auth, and the one MD5 round trip are not included + _LOGGER.info( + "Update took %.2f seconds (prepare %.2f, upload %.2f, commit %.2f)", + prepare_duration + duration + commit_duration, + prepare_duration, + duration, + commit_duration, + ) - try: - send_check(sock, RESPONSE_OK, "end acknowledgement") - except OTANetworkError as err: - # The device treats a missing end acknowledgement as non-fatal and is - # already rebooting into the new firmware, so the update succeeded - _LOGGER.warning("Failed sending end acknowledgement: %s", err) - _LOGGER.info("OTA successful (end acknowledgement not delivered)") - else: - _LOGGER.info("OTA successful") + try: + send_check(sock, RESPONSE_OK, "end acknowledgement") + except OTANetworkError as err: + # The device treats a missing end acknowledgement as non-fatal and is + # already rebooting into the new firmware, so the update succeeded + _LOGGER.warning("Failed sending end acknowledgement: %s", err) + _LOGGER.info("OTA successful (end acknowledgement not delivered)") + else: + _LOGGER.info("OTA successful") + if udp is not None: + # Our end acknowledgement also tells the device we have its + # responses, ending its wait before it reboots + udp.flush(2.0) + finally: + if udp is not None: + udp.close() # Do not connect logs until it is fully on time.sleep(1) @@ -1121,7 +1513,7 @@ def run_ota_impl_( # are shared across the addresses, cycling through them. Wait before an # attempt when the previous one actually reached the device, or when # revisiting an address, so a flaky link can recover and the device can - # clean up a half-open connection (its handshake watchdog runs at 20s); + # clean up a half-open connection (its handshake watchdog runs at 60s); # moving on to the next address family stays immediate. Known limitation: # a silent mid-transfer drop with no reset can wedge the device until its # 105s data timeout, which outlasts this budget; the retries target the @@ -1139,6 +1531,8 @@ def run_ota_impl_( on_upload_start() on_upload_start = None # once per upload, whatever the retries do + # A network failure means a lossy link, so later attempts send data over UDP + prefer_udp = False while attempt < total_attempts: af, socktype, _, _, sa = res[attempt % len(res)] if not reconnect and (reached_device or attempt >= len(res)): @@ -1153,7 +1547,7 @@ def run_ota_impl_( reconnect = False _LOGGER.info("Connecting to %s port %s...", sa[0], sa[1]) sock = socket.socket(af, socktype) - sock.settimeout(20.0) + sock.settimeout(CONNECT_TIMEOUT) try: sock.connect(sa) except OSError as err: @@ -1164,7 +1558,9 @@ def run_ota_impl_( continue _LOGGER.info("Connected to %s", sa[0]) + sock.settimeout(SETUP_TIMEOUT) reached_device = True + total_attempts = max(total_attempts, len(res) + EXTRA_UPLOAD_ATTEMPTS_REACHED) with contextlib.closing(sock), Path(filename).open("rb") as file_handle: try: perform_ota( @@ -1177,6 +1573,7 @@ def run_ota_impl_( encryption.plaintext_fallback, allow_plaintext_upload=allow_plaintext_upload, on_upload_start=upload_started, + prefer_udp=prefer_udp, ) except OTAKeyRejected as err: # The device sent a reject and closed on its own, so reconnect @@ -1196,6 +1593,7 @@ def run_ota_impl_( raise except OTAHandshakeNetworkError as err: last_error = str(err) + prefer_udp = True if encryption.handshake_fault_falls_back(): encryption.downgrade(last_error) continue @@ -1205,6 +1603,7 @@ def run_ota_impl_( except OTANetworkError as err: # Transient network failure; retry last_error = str(err) + prefer_udp = True _LOGGER.warning("%s", last_error) attempt += 1 continue 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 e83335a175f..04b09b4d542 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 @@ -357,6 +358,7 @@ def test_perform_ota_successful_md5_auth( | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL | espota2.CLIENT_FEATURE_SUPPORTS_DEFLATE + | espota2.CLIENT_FEATURE_SUPPORTS_UDP ] ) ) @@ -765,7 +767,7 @@ def test_run_ota_impl_successful( assert result_host == "192.168.1.100" # Verify socket was configured correctly - mock_socket.settimeout.assert_called_with(20.0) + mock_socket.settimeout.assert_called_with(espota2.SETUP_TIMEOUT) mock_socket.connect.assert_called_once_with(("192.168.1.100", 3232)) mock_socket.close.assert_called_once() @@ -853,13 +855,22 @@ def test_run_ota_impl_network_error_retry_succeeds( assert result_host == "192.168.1.100" assert mock_perform_ota.call_count == 2 mock_sleep.assert_called_once_with(espota2.UPLOAD_RETRY_DELAY) + # The failure marks the link lossy, so the retry moves the data to UDP + assert [c.kwargs["prefer_udp"] for c in mock_perform_ota.call_args_list] == [ + False, + True, + ] @pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip") def test_run_ota_impl_network_error_exhausts_attempts( mock_socket: Mock, firmware_file: Path, mock_perform_ota: Mock, mock_sleep: Mock ) -> None: - """Test run_ota_impl_ gives up after all attempts hit network errors.""" + """Test run_ota_impl_ gives up after all attempts hit network errors. + + Every attempt connected, so the budget is the longer one for a reachable + device on a lossy link. + """ mock_perform_ota.side_effect = espota2.OTANetworkError("sending data: broken pipe") result_code, result_host = espota2.run_ota_impl_( @@ -868,8 +879,8 @@ def test_run_ota_impl_network_error_exhausts_attempts( assert result_code == 1 assert result_host is None - assert mock_perform_ota.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS + 1 - assert mock_sleep.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS + assert mock_perform_ota.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS_REACHED + 1 + assert mock_sleep.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS_REACHED @pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip_dual") @@ -890,12 +901,10 @@ def test_run_ota_impl_multiple_addresses_cycle( assert mock_socket.connect.call_args_list == [ call(DUAL_STACK_SA6), call(DUAL_STACK_SA4), - call(DUAL_STACK_SA6), - call(DUAL_STACK_SA4), - ] + ] * (1 + espota2.EXTRA_UPLOAD_ATTEMPTS // 2) # No connect ever reached the device, so the delay only applies before # the revisits - assert mock_sleep.call_count == 2 + assert mock_sleep.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS @pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip_dual") @@ -1079,6 +1088,7 @@ def test_perform_ota_successful_sha256_auth( | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL | espota2.CLIENT_FEATURE_SUPPORTS_DEFLATE + | espota2.CLIENT_FEATURE_SUPPORTS_UDP ] ) ) @@ -1136,6 +1146,7 @@ def test_perform_ota_sha256_fallback_to_md5( | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL | espota2.CLIENT_FEATURE_SUPPORTS_DEFLATE + | espota2.CLIENT_FEATURE_SUPPORTS_UDP ] ) ) @@ -1246,6 +1257,7 @@ def test_perform_ota_extended_protocol_app( | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL | espota2.CLIENT_FEATURE_SUPPORTS_DEFLATE + | espota2.CLIENT_FEATURE_SUPPORTS_UDP ] ) ) @@ -1307,6 +1319,7 @@ def test_perform_ota_successful_partition_table( | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL | espota2.CLIENT_FEATURE_SUPPORTS_DEFLATE + | espota2.CLIENT_FEATURE_SUPPORTS_UDP ] ) ) @@ -1572,3 +1585,468 @@ def test_perform_ota_with_deflate(mock_socket: Mock, server_features: int) -> No assert len(payload) == sent_size < len(original_content) assert zlib.decompress(payload, -espota2.DEFLATE_WINDOW_BITS) == original_content assert sent[5] == hashlib.md5(original_content).hexdigest().encode() + + +class _LossyUdpDevice: + """The device side of the UDP channel (ota_esphome_udp.cpp) on a lossy link.""" + + def __init__(self, token: bytes, loss: float, seed: int = 1) -> None: + import random + import threading + + self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + self.sock.bind(("127.0.0.1", 0)) + self.sock.settimeout(0.05) + self.token = token + self.loss = loss + self.random = random.Random(seed) + self.next_seq = 0 + self.slots: dict[int, bytes] = {} # early messages within the window + self.received: list[bytes] = [] + self.out = bytearray() + self.out_base = 0 + self.committed = False + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, daemon=True) + + def __enter__(self) -> Self: + self._thread.start() + return self + + def __exit__(self, *exc: object) -> None: + self._stop.set() + self._thread.join() + self.sock.close() + + def _run(self) -> None: + while not self._stop.is_set(): + try: + data, peer = self.sock.recvfrom(2048) + except TimeoutError: + continue + if self.random.random() < self.loss: + continue + if data[1:5] != self.token: + continue + self.committed |= bool(data[0] & espota2.UDP_FLAG_COMMITTED) + acked = (int.from_bytes(data[7:9], "big") - self.out_base) & 0xFFFF + if acked <= len(self.out): + del self.out[:acked] + self.out_base += acked + seq = int.from_bytes(data[5:7], "big") + msg_type = data[0] & ~espota2.UDP_FLAG_COMMITTED + if ( + msg_type == espota2.UDP_MSG_DATA + and (seq - self.next_seq) & 0xFFFF < espota2.UDP_WINDOW + ): + self.slots.setdefault(seq, data[espota2.UDP_HEADER_SIZE :]) + # Consumed as soon as it is next, like a device that is never busy + while self.next_seq in self.slots: + self.received.append(self.slots.pop(self.next_seq)) + self.next_seq = (self.next_seq + 1) & 0xFFFF + if self.random.random() < self.loss: + continue + window = sum( + 1 << k + for k in range(1, espota2.UDP_WINDOW) + if (self.next_seq + k) & 0xFFFF in self.slots + ) + # Consumes at once, so every slot it does not hold is free + window |= (espota2.UDP_WINDOW - len(self.slots)) << 4 + self.sock.sendto( + bytes([espota2.UDP_MSG_ACK]) + + self.token + + self.next_seq.to_bytes(2, "big") + + seq.to_bytes(2, "big") + + (self.out_base & 0xFFFF).to_bytes(2, "big") + + bytes([window]) + + bytes(self.out), + peer, + ) + + +def _udp_client(device: _LossyUdpDevice) -> espota2.UdpChannel: + sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + sock.connect(device.sock.getsockname()) + return espota2.UdpChannel(sock, device.token) + + +@pytest.mark.parametrize("loss", [0.0, 0.5]) +def test_udp_channel_delivers_in_order(loss: float) -> None: + """The stream arrives intact and responses come back, despite loss.""" + data = bytes(range(256)) * 160 + with _LossyUdpDevice(b"\x01\x02\x03\x04", loss) as device: + channel = _udp_client(device) + assert channel.probe() + channel.sendall(data) + device.out += bytes([espota2.RESPONSE_RECEIVE_OK]) + assert channel.recv(1) == bytes([espota2.RESPONSE_RECEIVE_OK]) + channel.flush(30.0) + channel.close() + assert b"".join(device.received) == data + assert all(len(m) <= espota2.UDP_MAX_PAYLOAD for m in device.received) + assert device.committed + + +def test_start_udp_falls_back_when_blocked(mock_socket: Mock) -> None: + """No answer to the probes sends the fallback over TCP and stays on TCP.""" + with _LossyUdpDevice(b"\x05\x06\x07\x08", 1.0) as device: + mock_socket.family = socket.AF_INET + mock_socket.getpeername.return_value = device.sock.getsockname() + mock_socket.recv.side_effect = [b"\x05", b"\x06\x07\x08"] + with patch.object(espota2, "UDP_PROBE_TIMEOUT", 0.3): + assert espota2._start_udp(mock_socket, mock_socket, True) is None + assert mock_socket.sendall.call_args_list == [ + call(espota2.UDP_REQUEST), + call(espota2.UDP_FALLBACK), # raw on the TCP socket + ] + + +def test_start_udp_declined(mock_socket: Mock) -> None: + """A zero token means the device could not open UDP; nothing is probed.""" + mock_socket.recv.side_effect = [b"\x00", b"\x00\x00\x00"] + assert espota2._start_udp(mock_socket, mock_socket, True) is None + assert mock_socket.sendall.call_args_list == [call(espota2.UDP_REQUEST)] + + +def test_start_udp_clean_link_stays_on_tcp(mock_socket: Mock) -> None: + """A clean handshake asks to stay on TCP and opens no UDP socket.""" + mock_socket.recv.side_effect = [b"\x00", b"\x00\x00\x00"] + with patch.object(espota2.socket, "socket") as udp_socket: + assert espota2._start_udp(mock_socket, mock_socket, False) is None + udp_socket.assert_not_called() + assert mock_socket.sendall.call_args_list == [call(espota2.UDP_FALLBACK)] + + +@pytest.mark.parametrize( + ("slowest", "prefer_udp", "retransmits", "lossy"), + [ + (0.01, False, 0, False), + (0.01, False, None, False), + (espota2.UDP_SLOW_EXCHANGE + 0.1, False, 0, True), + (0.01, True, 0, True), + (0.01, False, 2, True), + ], +) +def test_link_lossy( + slowest: float, prefer_udp: bool, retransmits: int | None, lossy: bool +) -> None: + """A slow exchange, a retransmit or an earlier failure means UDP.""" + with patch.object(espota2, "_tcp_retransmits", return_value=retransmits): + assert espota2._link_lossy(Mock(), slowest, prefer_udp) is lossy + + +@pytest.mark.parametrize( + ("platform", "info", "expected"), + [ + ("linux", bytes(100) + (3).to_bytes(4, "little") + bytes(100), 3), + ("darwin", bytes(72) + (1460).to_bytes(8, "little") + bytes(8), 1460), + ("win32", b"", None), + ("linux", bytes(10), None), + ], +) +def test_tcp_retransmits(platform: str, info: bytes, expected: int | None) -> None: + """The retransmit counter is read where the OS reports it.""" + sock = Mock() + sock.getsockopt.return_value = info + with patch("sys.platform", platform): + assert espota2._tcp_retransmits(sock) == expected + + +def test_udp_channel_reports_socket_errors() -> None: + """A closed device port fails fast instead of waiting out the timeout.""" + gone = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + gone.bind(("127.0.0.1", 0)) + addr = gone.getsockname() + gone.close() + sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + sock.connect(addr) + channel = espota2.UdpChannel(sock, b"\x01\x01\x01\x01") + channel.settimeout(10.0) + with pytest.raises(espota2.OTANetworkError, match="closed its UDP port"): + channel.sendall(bytes(espota2.UDP_MAX_PAYLOAD * 5)) + # A probe that is refused falls back instead + assert not espota2.UdpChannel(sock, b"\x01\x01\x01\x01").probe() + channel.close() + + +def test_udp_channel_no_progress_names_last_error() -> None: + """Without refusals, the no-progress timeout names the last socket error.""" + sock = Mock() + sock.send.side_effect = BlockingIOError(35, "full") + channel = espota2.UdpChannel(sock, b"\x01\x01\x01\x01") + channel.settimeout(0.2) + channel.sendall(b"x") + with ( + patch("select.select", return_value=([], [], [])), + pytest.raises(espota2.OTANetworkError, match="no progress.*full"), + ): + while True: + channel._pump() + + +def test_udp_channel_ack_resets_refusals() -> None: + """Refusals only count in a row; an ACK in between resets them.""" + channel = espota2.UdpChannel(Mock(), b"\x01\x01\x01\x01") + ack = bytes([espota2.UDP_MSG_ACK]) + b"\x01\x01\x01\x01" + bytes(6) + b"\x40" + for _ in range(espota2.UDP_REFUSALS * 2): + channel._socket_error(ConnectionRefusedError(61, "refused")) + channel._handle_ack(ack) + # Windows' form of the same refusal counts too + with pytest.raises(espota2.OTANetworkError, match="closed its UDP port"): + for _ in range(espota2.UDP_REFUSALS): + channel._socket_error(ConnectionResetError(10054, "forcibly closed")) + + +def _udp_handshake() -> list[bytes]: + """TCP recv responses up to auth, from a device that offers UDP.""" + return [ + bytes([espota2.RESPONSE_OK]), + bytes([espota2.OTA_VERSION_2_0]), + bytes([espota2.RESPONSE_FEATURE_FLAGS]), + bytes([espota2.SERVER_FEATURE_SUPPORTS_UDP]), + bytes([espota2.RESPONSE_AUTH_OK]), + ] + + +@pytest.mark.usefixtures("mock_time") +def test_perform_ota_over_udp(mock_socket: Mock) -> None: + """The data phase moves to the channel, which acks chunks itself.""" + mock_socket.recv.side_effect = _udp_handshake() + udp = Mock() + udp.pending.return_value = False + udp.recv.side_effect = [ + bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]), + bytes([espota2.RESPONSE_BIN_MD5_OK]), + bytes([espota2.RESPONSE_RECEIVE_OK]), + bytes([espota2.RESPONSE_UPDATE_END_OK]), + ] + with patch.object(espota2, "_start_udp", return_value=udp) as start: + espota2.perform_ota( + mock_socket, None, io.BytesIO(bytes(20000)), "test.bin", prefer_udp=True + ) + assert start.call_args.args[2] is True # an earlier failure means lossy + assert udp.recv.call_count == 4 # no chunk acks + udp.flush.assert_called_once() + udp.close.assert_called_once() + + +@pytest.mark.usefixtures("mock_time") +def test_perform_ota_over_udp_device_error(mock_socket: Mock) -> None: + """A device error mid upload stops the upload and is not retried.""" + mock_socket.recv.side_effect = _udp_handshake() + mock_socket.getsockopt.return_value = b"" # no retransmit counter + udp = Mock() + udp.pending.return_value = True + udp.recv.side_effect = [ + bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]), + bytes([espota2.RESPONSE_BIN_MD5_OK]), + bytes([espota2.RESPONSE_ERROR_WRITING_FLASH]), + ] + with ( + patch.object(espota2, "_start_udp", return_value=udp), + pytest.raises(espota2.OTAError, match="Writing OTA data") as exc, + ): + espota2.perform_ota(mock_socket, None, io.BytesIO(bytes(20000)), "test.bin") + assert not isinstance(exc.value, espota2.OTANetworkError) + udp.close.assert_called() + + +@pytest.mark.usefixtures("mock_time") +def test_perform_ota_udp_declined_stays_on_tcp(mock_socket: Mock) -> None: + """Without a channel the upload continues on TCP with chunk acks.""" + mock_socket.recv.side_effect = [ + *_udp_handshake(), + bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]), + bytes([espota2.RESPONSE_BIN_MD5_OK]), + *_UPLOAD_TAIL, + ] + with patch.object(espota2, "_start_udp", return_value=None): + espota2.perform_ota( + mock_socket, None, io.BytesIO(b"x" * 100), "test.bin", prefer_udp=True + ) + assert mock_socket.sendall.call_args_list[-1] == call(bytes([espota2.RESPONSE_OK])) + + +def test_udp_channel_records_socket_errors() -> None: + """Send and receive errors are recorded like loss.""" + sock = Mock() + sock.send.side_effect = BlockingIOError(35, "full") + sock.recv.side_effect = ConnectionRefusedError(61, "refused") + channel = espota2.UdpChannel(sock, b"\x01\x01\x01\x01") + channel._send(espota2.UDP_MSG_PING, 0, b"") + assert isinstance(channel._last_error, BlockingIOError) + with patch("select.select", side_effect=[([sock], [], []), ([], [], [])]): + channel._receive(0.0) + assert isinstance(channel._last_error, ConnectionRefusedError) + + +def test_start_udp_probe_succeeds(mock_socket: Mock) -> None: + """A device that answers the probe gets the data phase over UDP.""" + with _LossyUdpDevice(b"\x05\x06\x07\x09", 0.0) as device: + mock_socket.family = socket.AF_INET + mock_socket.getpeername.return_value = device.sock.getsockname() + mock_socket.recv.side_effect = [b"\x05", b"\x06\x07\x09"] + channel = espota2._start_udp(mock_socket, mock_socket, True) + assert channel is not None + channel.close() + assert mock_socket.sendall.call_args_list == [call(espota2.UDP_REQUEST)] + + +def test_udp_channel_pings_when_all_held() -> None: + """With every message in flight already held, it asks for responses instead, + at most once per resend timeout.""" + device = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + device.bind(("127.0.0.1", 0)) + device.settimeout(1.0) + sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + sock.connect(device.getsockname()) + channel = espota2.UdpChannel(sock, b"\x01\x01\x01\x01") + channel.sendall(b"x") + channel._queue[0].held = True + with patch.object(channel, "_receive"): + channel._pump() + channel._pump() + assert device.recv(2048)[0] == espota2.UDP_MSG_PING + device.settimeout(0.1) + with pytest.raises(TimeoutError): + device.recv(2048) + channel.close() + device.close() + + +def test_udp_channel_ignores_foreign_datagrams() -> None: + """Short, non-ACK or wrong-token datagrams change nothing.""" + channel = espota2.UdpChannel(Mock(), b"\x01\x01\x01\x01") + for data in (b"\x40", bytes(espota2.UDP_ACK_HEADER_SIZE), b"\x40" + bytes(11)): + channel._handle_ack(data) + assert not channel.pending() + # An ACK before the head was ever sent is no duplicate + channel.sendall(b"x") + ack = bytes([espota2.UDP_MSG_ACK]) + b"\x01\x01\x01\x01" + bytes(6) + b"\x40" + channel._handle_ack(ack) + channel._handle_ack(ack) + assert not channel._queue[0].resend + + +def test_udp_channel_logs_each_socket_error_once( + caplog: pytest.LogCaptureFixture, +) -> None: + """A repeated error is recorded but logged only when it changes.""" + channel = espota2.UdpChannel(Mock(), b"\x01\x01\x01\x01") + with caplog.at_level(logging.DEBUG, "esphome.espota2"): + channel._socket_error(ConnectionRefusedError(61, "refused")) + channel._socket_error(ConnectionRefusedError(61, "refused again")) + assert caplog.text.count("UDP socket error") == 1 + assert "again" in str(channel._last_error) + + +def test_tcp_retransmits_unsupported_option() -> None: + """An OS that rejects the option reports nothing.""" + sock = Mock() + sock.getsockopt.side_effect = OSError("not supported") + with patch("sys.platform", "linux"): + assert espota2._tcp_retransmits(sock) is None + + +@pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip") +def test_run_ota_impl_handshake_failure_prefers_udp( + firmware_file: Path, mock_perform_ota: Mock, mock_sleep: Mock +) -> None: + """A handshake that dies after reaching the device moves the retry to UDP.""" + mock_perform_ota.side_effect = [ + espota2.OTAHandshakeNetworkError("noise handshake: timed out"), + None, + ] + result_code, _ = espota2.run_ota_impl_("test.local", 3232, None, str(firmware_file)) + assert result_code == 0 + assert [c.kwargs["prefer_udp"] for c in mock_perform_ota.call_args_list] == [ + False, + True, + ] + + +def test_start_udp_socket_error_falls_back(mock_socket: Mock) -> None: + """A UDP socket that cannot reach the peer stays on TCP.""" + mock_socket.family = socket.AF_INET + mock_socket.getpeername.side_effect = OSError("not connected") + mock_socket.recv.side_effect = [b"\x05", b"\x06\x07\x08"] + assert espota2._start_udp(mock_socket, mock_socket, True) is None + assert mock_socket.sendall.call_args_list == [ + call(espota2.UDP_REQUEST), + call(espota2.UDP_FALLBACK), + ] + + +def test_udp_channel_backs_off_when_acks_come_back_untimed() -> None: + """Progress acks that only answer resent messages mean the link is slower + than the timer: it backs off, up to its cap, until a sample resets it.""" + token = b"\x01\x01\x01\x01" + channel = espota2.UdpChannel(Mock(), token) + channel._queue = [espota2._UdpMessage(seq, b"x") for seq in range(40)] + rtos = [] + for received in range(1, 33): + for message in channel._queue[:1]: + message.sends = 2 # every message was resent before its ack + rtos.append(channel._rto()) + ack = ( + bytes([espota2.UDP_MSG_ACK]) + + token + + received.to_bytes(2, "big") + + (received - 1).to_bytes(2, "big") + + bytes(2) + + b"\x40" + ) + channel._handle_ack(ack) + n = espota2.UDP_UNTIMED_ACKS + assert rtos[0] == espota2.UDP_RESEND_INITIAL + assert rtos[n] == espota2.UDP_RESEND_INITIAL * 2 + assert rtos[-1] == espota2.UDP_RESEND_MAX + channel._sample_rtt(0.1) + assert channel._rto() == pytest.approx(0.3) # srtt 0.1 + 4 * rttvar 0.05 + + +@pytest.mark.usefixtures("mock_time") +@pytest.mark.parametrize( + ("size", "pending", "match", "retryable"), + [ + (100, False, "may have already committed", False), # one, final chunk + (20000, False, "no progress", True), # fails on an earlier chunk + (20000, True, "Writing OTA data", False), # device error waiting + ], +) +def test_perform_ota_over_udp_send_failure( + mock_socket: Mock, size: int, pending: bool, match: str, retryable: bool +) -> None: + """A UDP send failure surfaces a waiting device error, and on the last + chunk is not retried, like a lost final chunk ack over TCP.""" + mock_socket.recv.side_effect = _udp_handshake() + udp = Mock() + udp.pending.return_value = pending + udp.recv.side_effect = [ + bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]), + bytes([espota2.RESPONSE_BIN_MD5_OK]), + bytes([espota2.RESPONSE_ERROR_WRITING_FLASH]), + ] + # ota type, size and MD5 go through; the first data chunk fails + udp.sendall.side_effect = [None] * 3 + [espota2.OTANetworkError("no progress")] + with ( + patch.object(espota2, "_start_udp", return_value=udp), + pytest.raises(espota2.OTAError, match=match) as exc, + ): + espota2.perform_ota( + mock_socket, None, io.BytesIO(bytes(size)), "test.bin", prefer_udp=True + ) + assert isinstance(exc.value, espota2.OTANetworkError) is retryable + + +@pytest.mark.usefixtures("mock_time") +def test_perform_ota_tcp_send_network_error_unchanged(mock_socket: Mock) -> None: + """Over TCP an OTANetworkError from a send stays retryable, even on the + last chunk.""" + mock_socket.recv.side_effect = _no_auth_handshake(espota2.OTA_VERSION_2_0) + mock_socket.sendall.side_effect = [None] * 4 + [ + espota2.OTANetworkError("sending frame") + ] + with pytest.raises(espota2.OTANetworkError, match="sending frame"): + espota2.perform_ota(mock_socket, None, io.BytesIO(b"x" * 100), "test.bin") diff --git a/tests/unit_tests/test_espota2_noise.py b/tests/unit_tests/test_espota2_noise.py index f447e03986e..c316b55a9cf 100644 --- a/tests/unit_tests/test_espota2_noise.py +++ b/tests/unit_tests/test_espota2_noise.py @@ -76,9 +76,11 @@ class FakeEncryptedDevice(threading.Thread): prologue_features_override: int | None = None, connections: int = 1, drop_handshakes: int = 0, + udp: bool = False, reject_reason: str | None = None, ) -> None: super().__init__(daemon=True) + self.udp = udp # offer UDP; the data phase then has no chunk acks self.connections = connections self.drop_handshakes = drop_handshakes # hang up mid-handshake this many times self.reject_reason = reject_reason # refuse every handshake with this reason @@ -125,6 +127,8 @@ class FakeEncryptedDevice(threading.Thread): sock.sendall(bytes([espota2.RESPONSE_ERROR_ENCRYPTION_REQUIRED])) return server_flags = espota2.SERVER_FEATURE_SUPPORTS_NOISE if self.offer_noise else 0 + if self.udp: + server_flags |= espota2.SERVER_FEATURE_SUPPORTS_UDP sock.sendall(bytes([espota2.RESPONSE_FEATURE_FLAGS, server_flags])) if not (noise_negotiated and self.offer_noise): # A device that does not require encryption continues in @@ -217,7 +221,7 @@ class FakeEncryptedDevice(threading.Thread): acked = 0 while len(received) < size: received += recv_data(size - len(received)) - if self.version >= espota2.OTA_VERSION_2_0: + if self.version >= espota2.OTA_VERSION_2_0 and not self.udp: while acked + espota2.UPLOAD_BLOCK_SIZE <= len(received) or ( len(received) == size and acked < size ): @@ -913,6 +917,54 @@ def test_bare_block_refuses_a_device_that_cannot_encrypt( assert not any("Retrying in plaintext" in r.message for r in caplog.records) +class _TcpAsUdp: + """Stands in for a UdpChannel, carrying its byte stream over the TCP socket.""" + + def __init__(self, sock: socket.socket) -> None: + self._sock = sock + self.flushed = False + self.closed = False + + def settimeout(self, timeout: float | None) -> None: + self._sock.settimeout(timeout) + + def setsockopt(self, level: int, optname: int, value: int) -> None: + self._sock.setsockopt(level, optname, value) + + def sendall(self, data: bytes) -> None: + self._sock.sendall(data) + + def recv(self, amount: int) -> bytes: + return self._sock.recv(amount) + + def pending(self) -> bool: + return False + + def flush(self, timeout: float) -> None: + self.flushed = True + + def close(self) -> None: + self.closed = True + + +def test_encrypted_upload_over_udp() -> None: + """After the handshake, Noise frames continue on the UDP channel.""" + pytest.importorskip("aioesphomeapi.noise") + firmware = bytes(range(256)) * 80 + device = FakeEncryptedDevice(udp=True) + channels: list[_TcpAsUdp] = [] + + def start_udp(sock: Any, tcp_sock: socket.socket, lossy: bool) -> _TcpAsUdp: + channels.append(_TcpAsUdp(tcp_sock)) + return channels[0] + + with patch("time.sleep"), patch.object(espota2, "_start_udp", start_udp): + _upload(device, firmware, PSK) + device.join_and_check() + assert device.received == firmware + assert channels[0].flushed and channels[0].closed + + def test_probe_ota_key_recomputes_the_budget_after_connect() -> None: """A slow connect leaves less for the handshake than the budget had.""" clock = [0.0] From 9a5d2d97536e6204c4b9d0a8080612d3153dacdb Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Wed, 14 Oct 2026 17:38:39 -1000 Subject: [PATCH 2/2] [api] Encode action strings straight from flash (#20425) --- esphome/components/api/__init__.py | 33 ---------- esphome/components/api/api.proto | 16 ++--- esphome/components/api/api_options.proto | 2 +- esphome/components/api/api_pb2.cpp | 16 ++--- esphome/components/api/api_pb2.h | 16 ++--- esphome/components/api/api_pb2_dump.cpp | 28 ++++++--- esphome/components/api/api_server.cpp | 2 +- esphome/components/api/custom_api_device.h | 16 ++--- .../components/api/homeassistant_service.h | 32 ++-------- esphome/components/api/list_entities.cpp | 3 +- esphome/components/api/proto.h | 34 ++++++++++ esphome/components/api/user_services.cpp | 29 +++------ esphome/components/api/user_services.h | 31 +++------- .../button/homeassistant_button.cpp | 22 +++---- .../number/homeassistant_number.cpp | 29 +++++---- .../select/homeassistant_select.cpp | 28 ++++----- .../switch/homeassistant_switch.cpp | 18 +++--- .../homeassistant/text/homeassistant_text.cpp | 28 ++++----- esphome/core/defines.h | 1 - esphome/core/string_ref.h | 62 +++++++++++++------ script/api_protobuf/api_protobuf.py | 61 ++++++++++++------ .../api/test_action_metadata.py | 50 ++------------- .../core/test_progmem_string_ref.cpp | 30 +++++++++ .../api/test_api_protobuf_generator.py | 35 ++++++++++- 24 files changed, 327 insertions(+), 295 deletions(-) create mode 100644 tests/components/core/test_progmem_string_ref.cpp diff --git a/esphome/components/api/__init__.py b/esphome/components/api/__init__.py index 9611228ba1d..44d8b443676 100644 --- a/esphome/components/api/__init__.py +++ b/esphome/components/api/__init__.py @@ -251,9 +251,6 @@ def _validate_supports_response(value: Any) -> str: return cv.enum(SUPPORTS_RESPONSE_OPTIONS, lower=True)(value) -# ESP8266 copies every string of an action into a stack buffer sized by codegen; keep it small -ESP8266_ACTION_STRINGS_MAX_TOTAL = 384 - VARIABLE_SCHEMA = cv.Schema( { cv.Required(CONF_TYPE): cv.one_of(*SERVICE_ARG_NATIVE_TYPES, lower=True), @@ -484,29 +481,6 @@ def _action_strings(conf: ConfigType, has_metadata: bool) -> list[str | None]: return strings -def _action_strings_size(strings: list[str | None]) -> int: - """Bytes needed to copy every string out of flash, each with its terminator.""" - return sum( - len(string.encode("utf-8")) + 1 for string in strings if string is not None - ) - - -def _validate_esp8266_action_strings(config: ConfigType) -> ConfigType: - if not CORE.is_esp8266: - return config - actions = config.get(CONF_ACTIONS, []) - has_metadata = _has_action_metadata(actions) - for conf in actions: - size = _action_strings_size(_action_strings(conf, has_metadata)) - if size > ESP8266_ACTION_STRINGS_MAX_TOTAL: - raise cv.Invalid( - f"Action '{conf[CONF_ACTION]}' has {size} bytes of name, variable name, " - f"description and example text; ESP8266 allows at most " - f"{ESP8266_ACTION_STRINGS_MAX_TOTAL} bytes per action" - ) - return config - - def _validate_outgoing_host_ipv6(config: ConfigType) -> ConfigType: """An IPv6 host can never be parsed, so never dialed, without IPv6.""" if ( @@ -531,7 +505,6 @@ def _validate_wizard(config: ConfigType) -> ConfigType: FINAL_VALIDATE_SCHEMA = cv.All( - _validate_esp8266_action_strings, _validate_outgoing_host_ipv6, _validate_wizard, ) @@ -587,7 +560,6 @@ async def to_code(config: ConfigType) -> None: if config[CONF_HOMEASSISTANT_STATES]: cg.add_define("USE_API_HOMEASSISTANT_STATES") - scratch_size = 0 if actions: # Metadata is compiled in for every action once any action declares it, because the # string table layout is fixed by the define rather than per action @@ -639,8 +611,6 @@ async def to_code(config: ConfigType) -> None: func_args.append((native, name)) strings = _action_strings(conf, has_metadata) table = _add_action_strings(index, strings) - if CORE.is_esp8266: - scratch_size = max(scratch_size, _action_strings_size(strings)) # Template args: supports_response mode, then user service arg types templ = cg.TemplateArguments(supports_response, *service_template_args) # Key is hashed here because the name is not readable at runtime on ESP8266 @@ -667,9 +637,6 @@ async def to_code(config: ConfigType) -> None: cg.add(auto.add_action(unregister_action)) # Register all services at once - single allocation, no reallocations cg.add(var.initialize_user_services(triggers)) - if CORE.is_esp8266 and has_user_actions: - # Stack buffer that list-entities copies PROGMEM strings into, sized for the largest action - cg.add_define("API_USER_ACTION_STRINGS_SCRATCH_SIZE", max(scratch_size, 1)) if (wizard_config := config.get(wizard.CONF_WIZARD)) is not None: await wizard.to_code(wizard_config) diff --git a/esphome/components/api/api.proto b/esphome/components/api/api.proto index 41ece9275fd..679aff06b5e 100644 --- a/esphome/components/api/api.proto +++ b/esphome/components/api/api.proto @@ -1029,8 +1029,8 @@ message SubscribeHomeassistantServicesRequest { } message HomeassistantServiceMap { - string key = 1; - string value = 2; + string key = 1 [(progmem) = true]; + string value = 2 [(progmem) = true]; } message HomeassistantActionRequest { @@ -1039,7 +1039,7 @@ message HomeassistantActionRequest { option (no_delay) = true; option (ifdef) = "USE_API_HOMEASSISTANT_SERVICES"; - string service = 1; + string service = 1 [(progmem) = true]; repeated HomeassistantServiceMap data = 2 [(fixed_vector) = true]; repeated HomeassistantServiceMap data_template = 3 [(fixed_vector) = true]; repeated HomeassistantServiceMap variables = 4 [(fixed_vector) = true]; @@ -1160,21 +1160,21 @@ enum SupportsResponseType { } message ListEntitiesServicesArgument { option (ifdef) = "USE_API_USER_DEFINED_ACTIONS"; - string name = 1; + string name = 1 [(progmem) = true]; ServiceArgType type = 2; - string description = 3 [(field_ifdef) = "USE_API_USER_DEFINED_ACTION_METADATA"]; - string example = 4 [(field_ifdef) = "USE_API_USER_DEFINED_ACTION_METADATA"]; + string description = 3 [(field_ifdef) = "USE_API_USER_DEFINED_ACTION_METADATA", (progmem) = true]; + string example = 4 [(field_ifdef) = "USE_API_USER_DEFINED_ACTION_METADATA", (progmem) = true]; } message ListEntitiesServicesResponse { option (id) = 41; option (source) = SOURCE_SERVER; option (ifdef) = "USE_API_USER_DEFINED_ACTIONS"; - string name = 1; + string name = 1 [(progmem) = true]; fixed32 key = 2 [(force) = true]; repeated ListEntitiesServicesArgument args = 3 [(fixed_vector) = true]; SupportsResponseType supports_response = 4; - string description = 5 [(field_ifdef) = "USE_API_USER_DEFINED_ACTION_METADATA"]; + string description = 5 [(field_ifdef) = "USE_API_USER_DEFINED_ACTION_METADATA", (progmem) = true]; } message ExecuteServiceArgument { option (ifdef) = "USE_API_USER_DEFINED_ACTIONS"; diff --git a/esphome/components/api/api_options.proto b/esphome/components/api/api_options.proto index 8e8fa99833d..052b0c36bc7 100644 --- a/esphome/components/api/api_options.proto +++ b/esphome/components/api/api_options.proto @@ -123,7 +123,7 @@ extend google.protobuf.FieldOptions { // from an absent one (e.g. a UTC ParsedTimezone, which is all zeros). optional bool track_presence = 50020 [default=false]; - // progmem: The pointer_to_buffer bytes field points into flash (PROGMEM). + // progmem: The string field, or pointer_to_buffer bytes field, points into flash (PROGMEM). // The generated encode and dump copy it with progmem_memcpy, as ESP8266 // cannot read flash with ordinary loads. optional bool progmem = 50021 [default=false]; diff --git a/esphome/components/api/api_pb2.cpp b/esphome/components/api/api_pb2.cpp index 227af11bc08..8f7b622bc96 100644 --- a/esphome/components/api/api_pb2.cpp +++ b/esphome/components/api/api_pb2.cpp @@ -1142,8 +1142,8 @@ uint32_t NoiseEncryptionSetKeyResponse::calc_size_msg(const void *self) { #ifdef USE_API_HOMEASSISTANT_SERVICES uint8_t *HomeassistantServiceMap::encode_msg(const void *self, uint8_t *__restrict__ pos PROTO_ENCODE_DEBUG_PARAM) { const auto &msg = *static_cast(self); - pos = ProtoEncode::encode_string(pos PROTO_ENCODE_DEBUG_ARG, 1, msg.key); - pos = ProtoEncode::encode_string(pos PROTO_ENCODE_DEBUG_ARG, 2, msg.value); + pos = ProtoEncode::encode_progmem_string(pos PROTO_ENCODE_DEBUG_ARG, 1, msg.key); + pos = ProtoEncode::encode_progmem_string(pos PROTO_ENCODE_DEBUG_ARG, 2, msg.value); return pos; } uint32_t HomeassistantServiceMap::calc_size_msg(const void *self) { @@ -1155,7 +1155,7 @@ uint32_t HomeassistantServiceMap::calc_size_msg(const void *self) { } uint8_t *HomeassistantActionRequest::encode_msg(const void *self, uint8_t *__restrict__ pos PROTO_ENCODE_DEBUG_PARAM) { const auto &msg = *static_cast(self); - pos = ProtoEncode::encode_string(pos PROTO_ENCODE_DEBUG_ARG, 1, msg.service); + pos = ProtoEncode::encode_progmem_string(pos PROTO_ENCODE_DEBUG_ARG, 1, msg.service); for (auto &it : msg.data) { pos = ProtoEncode::encode_sub_message(pos PROTO_ENCODE_DEBUG_ARG, 2, it); } @@ -1326,13 +1326,13 @@ void GetTimeResponse::decode_field(void *self, uint32_t tag, const uint8_t *data uint8_t *ListEntitiesServicesArgument::encode_msg(const void *self, uint8_t *__restrict__ pos PROTO_ENCODE_DEBUG_PARAM) { const auto &msg = *static_cast(self); - pos = ProtoEncode::encode_string(pos PROTO_ENCODE_DEBUG_ARG, 1, msg.name); + pos = ProtoEncode::encode_progmem_string(pos PROTO_ENCODE_DEBUG_ARG, 1, msg.name); pos = ProtoEncode::encode_uint32(pos PROTO_ENCODE_DEBUG_ARG, 2, static_cast(msg.type)); #ifdef USE_API_USER_DEFINED_ACTION_METADATA - pos = ProtoEncode::encode_string(pos PROTO_ENCODE_DEBUG_ARG, 3, msg.description); + pos = ProtoEncode::encode_progmem_string(pos PROTO_ENCODE_DEBUG_ARG, 3, msg.description); #endif #ifdef USE_API_USER_DEFINED_ACTION_METADATA - pos = ProtoEncode::encode_string(pos PROTO_ENCODE_DEBUG_ARG, 4, msg.example); + pos = ProtoEncode::encode_progmem_string(pos PROTO_ENCODE_DEBUG_ARG, 4, msg.example); #endif return pos; } @@ -1352,14 +1352,14 @@ uint32_t ListEntitiesServicesArgument::calc_size_msg(const void *self) { uint8_t *ListEntitiesServicesResponse::encode_msg(const void *self, uint8_t *__restrict__ pos PROTO_ENCODE_DEBUG_PARAM) { const auto &msg = *static_cast(self); - pos = ProtoEncode::encode_string(pos PROTO_ENCODE_DEBUG_ARG, 1, msg.name); + pos = ProtoEncode::encode_progmem_string(pos PROTO_ENCODE_DEBUG_ARG, 1, msg.name); pos = ProtoEncode::write_tag_and_fixed32(pos PROTO_ENCODE_DEBUG_ARG, 21, msg.key); for (auto &it : msg.args) { pos = ProtoEncode::encode_sub_message(pos PROTO_ENCODE_DEBUG_ARG, 3, it); } pos = ProtoEncode::encode_uint32(pos PROTO_ENCODE_DEBUG_ARG, 4, static_cast(msg.supports_response)); #ifdef USE_API_USER_DEFINED_ACTION_METADATA - pos = ProtoEncode::encode_string(pos PROTO_ENCODE_DEBUG_ARG, 5, msg.description); + pos = ProtoEncode::encode_progmem_string(pos PROTO_ENCODE_DEBUG_ARG, 5, msg.description); #endif return pos; } diff --git a/esphome/components/api/api_pb2.h b/esphome/components/api/api_pb2.h index 7ec8955610c..fd12882dc9a 100644 --- a/esphome/components/api/api_pb2.h +++ b/esphome/components/api/api_pb2.h @@ -1389,8 +1389,8 @@ class NoiseEncryptionSetKeyResponse final : public ProtoMessage { #ifdef USE_API_HOMEASSISTANT_SERVICES class HomeassistantServiceMap final : public ProtoMessage { public: - StringRef key{nullptr, 0}; // null until set, encode only - StringRef value{nullptr, 0}; // null until set, encode only + ProgmemStringRef key{}; + ProgmemStringRef value{}; static uint8_t *encode_msg(const void *self, uint8_t *pos PROTO_ENCODE_DEBUG_PARAM); uint8_t *encode(ProtoWriteBuffer &buffer PROTO_ENCODE_DEBUG_PARAM) const { return encode_msg(this, buffer.get_pos() PROTO_ENCODE_DEBUG_ARG); @@ -1410,7 +1410,7 @@ class HomeassistantActionRequest final : public ProtoMessage { #ifdef HAS_PROTO_MESSAGE_DUMP const LogString *message_name() const override { return LOG_STR("homeassistant_action_request"); } #endif - StringRef service{nullptr, 0}; // null until set, encode only + ProgmemStringRef service{}; FixedVector data{}; FixedVector data_template{}; FixedVector variables{}; @@ -1577,13 +1577,13 @@ class GetTimeResponse final : public ProtoDecodableMessage { #ifdef USE_API_USER_DEFINED_ACTIONS class ListEntitiesServicesArgument final : public ProtoMessage { public: - StringRef name{nullptr, 0}; // null until set, encode only + ProgmemStringRef name{}; enums::ServiceArgType type{}; #ifdef USE_API_USER_DEFINED_ACTION_METADATA - StringRef description{nullptr, 0}; // null until set, encode only + ProgmemStringRef description{}; #endif #ifdef USE_API_USER_DEFINED_ACTION_METADATA - StringRef example{nullptr, 0}; // null until set, encode only + ProgmemStringRef example{}; #endif static uint8_t *encode_msg(const void *self, uint8_t *pos PROTO_ENCODE_DEBUG_PARAM); uint8_t *encode(ProtoWriteBuffer &buffer PROTO_ENCODE_DEBUG_PARAM) const { @@ -1604,12 +1604,12 @@ class ListEntitiesServicesResponse final : public ProtoMessage { #ifdef HAS_PROTO_MESSAGE_DUMP const LogString *message_name() const override { return LOG_STR("list_entities_services_response"); } #endif - StringRef name{nullptr, 0}; // null until set, encode only + ProgmemStringRef name{}; uint32_t key{0}; FixedVector args{}; enums::SupportsResponseType supports_response{}; #ifdef USE_API_USER_DEFINED_ACTION_METADATA - StringRef description{nullptr, 0}; // null until set, encode only + ProgmemStringRef description{}; #endif static uint8_t *encode_msg(const void *self, uint8_t *pos PROTO_ENCODE_DEBUG_PARAM); uint8_t *encode(ProtoWriteBuffer &buffer PROTO_ENCODE_DEBUG_PARAM) const { diff --git a/esphome/components/api/api_pb2_dump.cpp b/esphome/components/api/api_pb2_dump.cpp index 23ce6319ed1..688ec48e6e6 100644 --- a/esphome/components/api/api_pb2_dump.cpp +++ b/esphome/components/api/api_pb2_dump.cpp @@ -12,8 +12,9 @@ namespace esphome::api { #ifdef USE_ESP8266 // Out-of-line to avoid inlining strlen_P/memcpy_P at every call site -void DumpBuffer::append_p_esp8266(const char *str) { - size_t len = strlen_P(str); +void DumpBuffer::append_p_esp8266(const char *str) { this->append_p_esp8266(str, strlen_P(str)); } + +void DumpBuffer::append_p_esp8266(const char *str, size_t len) { size_t space = CAPACITY - 1 - pos_; if (len > space) len = space; @@ -134,6 +135,13 @@ static void dump_progmem_bytes_field(DumpBuffer &out, const char *field_name, co progmem_memcpy(data_buf, data, len); dump_bytes_field(out, field_name, data_buf, len, indent); } + +// Helper for string fields in flash +static void dump_progmem_string_field(DumpBuffer &out, const char *field_name, const ProgmemStringRef &value, + int indent = 2) { + append_field_prefix(out, field_name, indent); + out.append("'").append_p(value.progmem_ptr(), value.size()).append("'\n"); +} #pragma GCC diagnostic pop template<> const char *proto_enum_to_string(enums::DisconnectReason value) { @@ -1512,13 +1520,13 @@ const char *NoiseEncryptionSetKeyResponse::dump_to(DumpBuffer &out) const { #ifdef USE_API_HOMEASSISTANT_SERVICES const char *HomeassistantServiceMap::dump_to(DumpBuffer &out) const { MessageDumpHelper helper(out, ESPHOME_PSTR("HomeassistantServiceMap")); - dump_field(out, ESPHOME_PSTR("key"), this->key); - dump_field(out, ESPHOME_PSTR("value"), this->value); + dump_progmem_string_field(out, ESPHOME_PSTR("key"), this->key); + dump_progmem_string_field(out, ESPHOME_PSTR("value"), this->value); return out.c_str(); } const char *HomeassistantActionRequest::dump_to(DumpBuffer &out) const { MessageDumpHelper helper(out, ESPHOME_PSTR("HomeassistantActionRequest")); - dump_field(out, ESPHOME_PSTR("service"), this->service); + dump_progmem_string_field(out, ESPHOME_PSTR("service"), this->service); for (const auto &it : this->data) { out.append(4, ' ').append_p(ESPHOME_PSTR("data")).append(": "); it.dump_to(out); @@ -1613,19 +1621,19 @@ const char *GetTimeResponse::dump_to(DumpBuffer &out) const { #ifdef USE_API_USER_DEFINED_ACTIONS const char *ListEntitiesServicesArgument::dump_to(DumpBuffer &out) const { MessageDumpHelper helper(out, ESPHOME_PSTR("ListEntitiesServicesArgument")); - dump_field(out, ESPHOME_PSTR("name"), this->name); + dump_progmem_string_field(out, ESPHOME_PSTR("name"), this->name); dump_field(out, ESPHOME_PSTR("type"), static_cast(this->type)); #ifdef USE_API_USER_DEFINED_ACTION_METADATA - dump_field(out, ESPHOME_PSTR("description"), this->description); + dump_progmem_string_field(out, ESPHOME_PSTR("description"), this->description); #endif #ifdef USE_API_USER_DEFINED_ACTION_METADATA - dump_field(out, ESPHOME_PSTR("example"), this->example); + dump_progmem_string_field(out, ESPHOME_PSTR("example"), this->example); #endif return out.c_str(); } const char *ListEntitiesServicesResponse::dump_to(DumpBuffer &out) const { MessageDumpHelper helper(out, ESPHOME_PSTR("ListEntitiesServicesResponse")); - dump_field(out, ESPHOME_PSTR("name"), this->name); + dump_progmem_string_field(out, ESPHOME_PSTR("name"), this->name); dump_field(out, ESPHOME_PSTR("key"), this->key); for (const auto &it : this->args) { out.append(4, ' ').append_p(ESPHOME_PSTR("args")).append(": "); @@ -1634,7 +1642,7 @@ const char *ListEntitiesServicesResponse::dump_to(DumpBuffer &out) const { } dump_field(out, ESPHOME_PSTR("supports_response"), static_cast(this->supports_response)); #ifdef USE_API_USER_DEFINED_ACTION_METADATA - dump_field(out, ESPHOME_PSTR("description"), this->description); + dump_progmem_string_field(out, ESPHOME_PSTR("description"), this->description); #endif return out.c_str(); } diff --git a/esphome/components/api/api_server.cpp b/esphome/components/api/api_server.cpp index f63039a9714..a8d14256dcb 100644 --- a/esphome/components/api/api_server.cpp +++ b/esphome/components/api/api_server.cpp @@ -547,7 +547,7 @@ void APIServer::send_homeassistant_action(const HomeassistantActionRequest &call // arrive before the subscription and are lost - warn instead of failing silently. ESP_LOGW(TAG, "Home Assistant %s '%.*s' dropped; %s", call.is_event ? LOG_STR_LITERAL("event") : LOG_STR_LITERAL("action"), - static_cast(call.service.size()), call.service.empty() ? "" : call.service.c_str(), + static_cast(call.service.size()), call.service.empty() ? "" : call.service.progmem_ptr(), this->is_connected() ? LOG_STR_LITERAL("client has not subscribed to actions (yet)") : LOG_STR_LITERAL("no client connected")); } diff --git a/esphome/components/api/custom_api_device.h b/esphome/components/api/custom_api_device.h index c2b0f14bcc5..208c956eb17 100644 --- a/esphome/components/api/custom_api_device.h +++ b/esphome/components/api/custom_api_device.h @@ -247,7 +247,7 @@ class CustomAPIDevice { */ void call_homeassistant_service(const std::string &service_name) { HomeassistantActionRequest resp; - resp.service = StringRef(service_name); + resp.service = ProgmemStringRef(StringRef(service_name)); global_api_server->send_homeassistant_action(resp); } @@ -267,12 +267,12 @@ class CustomAPIDevice { */ void call_homeassistant_service(const std::string &service_name, const std::map &data) { HomeassistantActionRequest resp; - resp.service = StringRef(service_name); + resp.service = ProgmemStringRef(StringRef(service_name)); resp.data.init(data.size()); for (auto &it : data) { auto &kv = resp.data.emplace_back(); - kv.key = StringRef(it.first); - kv.value = StringRef(it.second); // data map lives until send completes + kv.key = ProgmemStringRef(StringRef(it.first)); + kv.value = ProgmemStringRef(StringRef(it.second)); // data map lives until send completes } global_api_server->send_homeassistant_action(resp); } @@ -289,7 +289,7 @@ class CustomAPIDevice { */ void fire_homeassistant_event(const std::string &event_name) { HomeassistantActionRequest resp; - resp.service = StringRef(event_name); + resp.service = ProgmemStringRef(StringRef(event_name)); resp.is_event = true; global_api_server->send_homeassistant_action(resp); } @@ -309,13 +309,13 @@ class CustomAPIDevice { */ void fire_homeassistant_event(const std::string &service_name, const std::map &data) { HomeassistantActionRequest resp; - resp.service = StringRef(service_name); + resp.service = ProgmemStringRef(StringRef(service_name)); resp.is_event = true; resp.data.init(data.size()); for (auto &it : data) { auto &kv = resp.data.emplace_back(); - kv.key = StringRef(it.first); - kv.value = StringRef(it.second); // data map lives until send completes + kv.key = ProgmemStringRef(StringRef(it.first)); + kv.value = ProgmemStringRef(StringRef(it.second)); // data map lives until send completes } global_api_server->send_homeassistant_action(resp); } diff --git a/esphome/components/api/homeassistant_service.h b/esphome/components/api/homeassistant_service.h index e7798a5d267..8964e8ca352 100644 --- a/esphome/components/api/homeassistant_service.h +++ b/esphome/components/api/homeassistant_service.h @@ -116,42 +116,18 @@ template class HomeAssistantServiceCallAction final : public Act const Field *fields = this->fields_; const size_t total = 1 + this->data_count_ + this->data_template_count_ + this->variables_count_; - // Lambda results, and on ESP8266 the RAM copies of the flash strings, must live until the send + // Lambda results must live until the send; the constant strings are encoded from flash where they are size_t lambda_count = 0; -#ifdef USE_ESP8266 - size_t flash_len = 0; -#endif - for (size_t i = 0; i < total; i++) { + for (size_t i = 0; i < total; i++) lambda_count += fields[i].fn != nullptr; -#ifdef USE_ESP8266 - if (fields[i].fn == nullptr) - flash_len += ESPHOME_strlen_P(fields[i].value); - if (fields[i].key != nullptr) - flash_len += ESPHOME_strlen_P(fields[i].key); -#endif - } FixedVector results; results.init(lambda_count); -#ifdef USE_ESP8266 - SmallBufferWithHeapFallback<128, char> flash_copy(flash_len); - char *cursor = flash_copy.get(); -#endif - auto string_ref = [&](const char *str) { -#ifdef USE_ESP8266 - size_t len = ESPHOME_strlen_P(str); - memcpy_P(cursor, str, len); - StringRef ref(cursor, len); - cursor += len; - return ref; -#else - return StringRef(str); -#endif - }; + auto string_ref = [](const char *str) { return ProgmemStringRef(str, ESPHOME_strlen_P(str)); }; auto value_ref = [&](const Field &field) { if (field.fn == nullptr) return string_ref(field.value); results.push_back(field.fn(x...)); - return StringRef(results.back()); + return ProgmemStringRef(StringRef(results.back())); }; auto fill = [&](FixedVector &dest, uint8_t count) { dest.init(count); diff --git a/esphome/components/api/list_entities.cpp b/esphome/components/api/list_entities.cpp index 507b098fb48..57ff616ca7c 100644 --- a/esphome/components/api/list_entities.cpp +++ b/esphome/components/api/list_entities.cpp @@ -99,8 +99,7 @@ ListEntitiesIterator::ListEntitiesIterator(APIConnection *client) : client_(clie static constexpr uint8_t SERVICE_YIELD_INTERVAL = 3; bool ListEntitiesIterator::on_service(UserServiceDescriptor *service) { - UserActionScratch scratch; - auto resp = service->encode_list_service_response(scratch); + auto resp = service->encode_list_service_response(); if (!this->client_->send_message(resp)) return false; // at_ is this service's index diff --git a/esphome/components/api/proto.h b/esphome/components/api/proto.h index 0fdda5067ea..e01deccbc21 100644 --- a/esphome/components/api/proto.h +++ b/esphome/components/api/proto.h @@ -418,6 +418,18 @@ class ProtoEncode { std::memcpy(pos + 2, ref.c_str(), ref.size()); return pos + 2 + ref.size(); } + /// encode_short_string_force for a string that may be in flash (PROGMEM). + [[nodiscard]] static inline uint8_t *encode_short_progmem_string_force( + uint8_t *__restrict__ pos PROTO_ENCODE_DEBUG_PARAM, uint8_t tag, const ProgmemStringRef &ref) { +#ifdef ESPHOME_DEBUG_API + assert(ref.size() < 128 && "encode_short_progmem_string_force: string exceeds max_data_length < 128"); +#endif + PROTO_ENCODE_CHECK_BOUNDS(pos, 2 + ref.size()); + pos[0] = tag; + pos[1] = static_cast(ref.size()); + progmem_memcpy(pos + 2, ref.progmem_ptr(), ref.size()); + return pos + 2 + ref.size(); + } /// Write a precomputed tag byte + 32-bit value. Outlined on embedded: one copy beats inline stores per field. [[nodiscard]] static PROTO_OUTLINE_FOR_SIZE uint8_t *write_tag_and_fixed32( uint8_t *__restrict__ pos PROTO_ENCODE_DEBUG_PARAM, uint8_t tag, uint32_t value) { @@ -481,6 +493,17 @@ class ProtoEncode { progmem_memcpy(pos, data, len); return pos + len; } + /// encode_progmem_bytes for a string field that may be in flash. + [[nodiscard]] static inline uint8_t *encode_progmem_string(uint8_t *__restrict__ pos PROTO_ENCODE_DEBUG_PARAM, + uint32_t field_id, const ProgmemStringRef &ref) { + return encode_progmem_bytes(pos PROTO_ENCODE_DEBUG_ARG, field_id, + reinterpret_cast(ref.progmem_ptr()), ref.size()); + } + [[nodiscard]] static inline uint8_t *encode_progmem_string_force(uint8_t *__restrict__ pos PROTO_ENCODE_DEBUG_PARAM, + uint32_t field_id, const ProgmemStringRef &ref) { + return encode_progmem_bytes_force(pos PROTO_ENCODE_DEBUG_ARG, field_id, + reinterpret_cast(ref.progmem_ptr()), ref.size()); + } [[nodiscard]] static inline uint8_t *encode_uint32_force(uint8_t *__restrict__ pos PROTO_ENCODE_DEBUG_PARAM, uint32_t field_id, uint32_t value) { pos = encode_field_raw(pos PROTO_ENCODE_DEBUG_ARG, field_id, 0); @@ -680,9 +703,20 @@ class DumpBuffer { return *this; } + /// Append len bytes of a PROGMEM string + DumpBuffer &append_p(const char *str, size_t len) { +#ifdef USE_ESP8266 + append_p_esp8266(str, len); +#else + append_impl_(str, len); +#endif + return *this; + } + #ifdef USE_ESP8266 /// Out-of-line ESP8266 PROGMEM append to avoid inlining strlen_P/memcpy_P at every call site void append_p_esp8266(const char *str); + void append_p_esp8266(const char *str, size_t len); #endif const char *c_str() const { return buf_; } diff --git a/esphome/components/api/user_services.cpp b/esphome/components/api/user_services.cpp index fad3cde29b9..a9836170640 100644 --- a/esphome/components/api/user_services.cpp +++ b/esphome/components/api/user_services.cpp @@ -1,47 +1,36 @@ #include "user_services.h" #include "esphome/core/hal.h" #include "esphome/core/log.h" +#include "esphome/core/progmem.h" #include "esphome/core/string_ref.h" namespace esphome::api { -StringRef UserServiceStatic::str_(size_t idx, std::span &scratch) const { +ProgmemStringRef UserServiceStatic::str_(size_t idx) const { const char *s = progmem_read_ptr(&this->strings_[idx]); if (s == nullptr) return {}; -#ifdef USE_ESP8266 - // Codegen sizes the scratch buffer for the largest service; the bound only guards other callers - if (scratch.empty()) - return {}; - size_t len = strnlen_P(s, scratch.size() - 1); - progmem_memcpy(scratch.data(), s, len); - scratch[len] = '\0'; - StringRef ref(scratch.data(), len); - scratch = scratch.subspan(len + 1); - return ref; -#else - return StringRef(s); -#endif + return ProgmemStringRef(s, ESPHOME_strlen_P(s)); } ListEntitiesServicesResponse UserServiceStatic::encode_list_service_response_( - std::span arg_types, std::span scratch) const { + std::span arg_types) const { ListEntitiesServicesResponse msg; - msg.name = this->str_(0, scratch); + msg.name = this->str_(0); msg.key = this->key_; msg.supports_response = this->supports_response_; #ifdef USE_API_USER_DEFINED_ACTION_METADATA - msg.description = this->str_(1, scratch); + msg.description = this->str_(1); #endif msg.args.init(arg_types.size()); for (size_t i = 0; i < arg_types.size(); i++) { size_t base = USER_ACTION_HEADER_STRINGS + i * USER_ACTION_ARG_STRINGS; auto &arg = msg.args.emplace_back(); arg.type = arg_types[i]; - arg.name = this->str_(base, scratch); + arg.name = this->str_(base); #ifdef USE_API_USER_DEFINED_ACTION_METADATA - arg.description = this->str_(base + 1, scratch); - arg.example = this->str_(base + 2, scratch); + arg.description = this->str_(base + 1); + arg.example = this->str_(base + 2); #endif } return msg; diff --git a/esphome/components/api/user_services.h b/esphome/components/api/user_services.h index 3b17bdb7bcd..057e04dd7a7 100644 --- a/esphome/components/api/user_services.h +++ b/esphome/components/api/user_services.h @@ -20,9 +20,8 @@ class APIServer; class UserServiceDescriptor { public: - /// Build the list-entities message. On ESP8266 the strings live in PROGMEM and are copied into - /// `scratch`, so the returned message is only valid while `scratch` is; other platforms ignore it. - virtual ListEntitiesServicesResponse encode_list_service_response(std::span scratch) = 0; + /// Build the list-entities message; its strings may point into flash (PROGMEM). + virtual ListEntitiesServicesResponse encode_list_service_response() = 0; virtual bool execute_service(const ExecuteServiceRequest &req) = 0; #ifdef USE_API_USER_DEFINED_ACTION_RESPONSES @@ -37,13 +36,6 @@ template T get_execute_arg_value(const ExecuteServiceArgument &arg); template enums::ServiceArgType to_service_arg_type(); -// Scratch buffer list-entities hands to encode_list_service_response(); only ESP8266 copies into it -#ifdef USE_ESP8266 -using UserActionScratch = std::array; -#else -using UserActionScratch = std::array; -#endif - // Non-template base for YAML-defined services so the list-entities encoder is compiled once. // All strings live in one PROGMEM pointer table emitted by codegen (see _action_strings in // __init__.py), so each service costs a single pointer of RAM. Layout: the action name, then @@ -63,12 +55,9 @@ class UserServiceStatic : public UserServiceDescriptor { : strings_(strings), key_(key), supports_response_(supports_response) {} protected: - ListEntitiesServicesResponse encode_list_service_response_(std::span arg_types, - std::span scratch) const; - /// Reference table entry `idx`; nullptr gives an empty StringRef. - /// On ESP8266 the bytes are copied out of PROGMEM into `scratch` with a terminator, and the span - /// is advanced past the copy. - StringRef str_(size_t idx, std::span &scratch) const; + ListEntitiesServicesResponse encode_list_service_response_(std::span arg_types) const; + /// Table entry `idx` as it sits in flash; nullptr gives an empty string. + ProgmemStringRef str_(size_t idx) const; const char *const *strings_; // PROGMEM pointer table, read with progmem_read_ptr() uint32_t key_; @@ -79,9 +68,9 @@ template class UserServiceBase : public UserServiceStatic { public: using UserServiceStatic::UserServiceStatic; - ListEntitiesServicesResponse encode_list_service_response(std::span scratch) override { + ListEntitiesServicesResponse encode_list_service_response() override { std::array arg_types = {to_service_arg_type()...}; - return this->encode_list_service_response_(arg_types, scratch); + return this->encode_list_service_response_(arg_types); } bool execute_service(const ExecuteServiceRequest &req) override { @@ -125,9 +114,9 @@ template class UserServiceDynamic : public UserServiceDescriptor this->key_ = fnv1_hash(this->name_.c_str()); } - ListEntitiesServicesResponse encode_list_service_response(std::span /*scratch*/) override { + ListEntitiesServicesResponse encode_list_service_response() override { ListEntitiesServicesResponse msg; - msg.name = StringRef(this->name_); + msg.name = ProgmemStringRef(StringRef(this->name_)); msg.key = this->key_; msg.supports_response = enums::SUPPORTS_RESPONSE_NONE; // Dynamic services don't support responses yet std::array arg_types = {to_service_arg_type()...}; @@ -135,7 +124,7 @@ template class UserServiceDynamic : public UserServiceDescriptor for (size_t i = 0; i < sizeof...(Ts); i++) { auto &arg = msg.args.emplace_back(); arg.type = arg_types[i]; - arg.name = StringRef(this->arg_names_[i]); + arg.name = ProgmemStringRef(StringRef(this->arg_names_[i])); } return msg; } diff --git a/esphome/components/homeassistant/button/homeassistant_button.cpp b/esphome/components/homeassistant/button/homeassistant_button.cpp index 6cf1c87eaa5..5c830c8c923 100644 --- a/esphome/components/homeassistant/button/homeassistant_button.cpp +++ b/esphome/components/homeassistant/button/homeassistant_button.cpp @@ -5,6 +5,7 @@ #include "esphome/components/api/api_pb2.h" #include "esphome/components/api/api_server.h" #include "esphome/core/log.h" +#include "esphome/core/progmem.h" #include "esphome/core/string_ref.h" namespace esphome::homeassistant { @@ -31,22 +32,21 @@ void HomeassistantButton::press_action() { } #endif - static constexpr auto SERVICE_BUTTON = StringRef::from_lit("button.press"); - static constexpr auto SERVICE_INPUT_BUTTON = StringRef::from_lit("input_button.press"); - static constexpr auto ENTITY_ID_KEY = StringRef::from_lit("entity_id"); - static constexpr char INPUT_PREFIX[] = "input_"; + // In flash on the ESP8266 and encoded from there + static const char SERVICE_BUTTON[] PROGMEM = "button.press"; + static const char SERVICE_INPUT_BUTTON[] PROGMEM = "input_button.press"; + static const char ENTITY_ID_KEY[] PROGMEM = "entity_id"; api::HomeassistantActionRequest resp; - if (strncmp(this->entity_id_, INPUT_PREFIX, sizeof(INPUT_PREFIX) - 1) == 0) { - resp.service = SERVICE_INPUT_BUTTON; - } else { - resp.service = SERVICE_BUTTON; - } + // An input helper entity id starts with the "input_" its action name starts with + const bool input_helper = ESPHOME_strncmp_P(this->entity_id_, SERVICE_INPUT_BUTTON, sizeof("input_") - 1) == 0; + resp.service = + input_helper ? ProgmemStringRef::from_lit(SERVICE_INPUT_BUTTON) : ProgmemStringRef::from_lit(SERVICE_BUTTON); resp.data.init(1); auto &entity_id_kv = resp.data.emplace_back(); - entity_id_kv.key = ENTITY_ID_KEY; - entity_id_kv.value = StringRef(this->entity_id_); + entity_id_kv.key = ProgmemStringRef::from_lit(ENTITY_ID_KEY); + entity_id_kv.value = ProgmemStringRef(StringRef(this->entity_id_)); api::global_api_server->send_homeassistant_action(resp); } diff --git a/esphome/components/homeassistant/number/homeassistant_number.cpp b/esphome/components/homeassistant/number/homeassistant_number.cpp index df92aad38eb..73efc8e5922 100644 --- a/esphome/components/homeassistant/number/homeassistant_number.cpp +++ b/esphome/components/homeassistant/number/homeassistant_number.cpp @@ -5,6 +5,7 @@ #include "esphome/components/api/api_pb2.h" #include "esphome/components/api/api_server.h" #include "esphome/core/log.h" +#include "esphome/core/progmem.h" #include "esphome/core/string_ref.h" namespace esphome::homeassistant { @@ -89,31 +90,29 @@ void HomeassistantNumber::control(float value) { this->publish_state(value); - static constexpr auto NUMBER_SERVICE_NAME = StringRef::from_lit("number.set_value"); - static constexpr auto INPUT_NUMBER_SERVICE_NAME = StringRef::from_lit("input_number.set_value"); - static constexpr char INPUT_PREFIX[] = "input_"; - static constexpr auto ENTITY_ID_KEY = StringRef::from_lit("entity_id"); - static constexpr auto VALUE_KEY = StringRef::from_lit("value"); + // In flash on the ESP8266 and encoded from there + static const char NUMBER_SERVICE_NAME[] PROGMEM = "number.set_value"; + static const char INPUT_NUMBER_SERVICE_NAME[] PROGMEM = "input_number.set_value"; + static const char ENTITY_ID_KEY[] PROGMEM = "entity_id"; + static const char VALUE_KEY[] PROGMEM = "value"; api::HomeassistantActionRequest resp; - // input_number entities have their own set_value action - if (strncmp(this->entity_id_, INPUT_PREFIX, sizeof(INPUT_PREFIX) - 1) == 0) { - resp.service = INPUT_NUMBER_SERVICE_NAME; - } else { - resp.service = NUMBER_SERVICE_NAME; - } + // An input helper entity id starts with the "input_" its action name starts with + const bool input_helper = ESPHOME_strncmp_P(this->entity_id_, INPUT_NUMBER_SERVICE_NAME, sizeof("input_") - 1) == 0; + resp.service = input_helper ? ProgmemStringRef::from_lit(INPUT_NUMBER_SERVICE_NAME) + : ProgmemStringRef::from_lit(NUMBER_SERVICE_NAME); resp.data.init(2); auto &entity_id = resp.data.emplace_back(); - entity_id.key = ENTITY_ID_KEY; - entity_id.value = StringRef(this->entity_id_); + entity_id.key = ProgmemStringRef::from_lit(ENTITY_ID_KEY); + entity_id.value = ProgmemStringRef(StringRef(this->entity_id_)); auto &entity_value = resp.data.emplace_back(); - entity_value.key = VALUE_KEY; + entity_value.key = ProgmemStringRef::from_lit(VALUE_KEY); // Stack buffer - no heap allocation; %g produces shortest representation char value_buf[16]; buf_append_printf(value_buf, sizeof(value_buf), 0, "%g", value); - entity_value.value = StringRef(value_buf); + entity_value.value = ProgmemStringRef(StringRef(value_buf)); api::global_api_server->send_homeassistant_action(resp); } diff --git a/esphome/components/homeassistant/select/homeassistant_select.cpp b/esphome/components/homeassistant/select/homeassistant_select.cpp index 189a6b1f376..4f4fc7a3295 100644 --- a/esphome/components/homeassistant/select/homeassistant_select.cpp +++ b/esphome/components/homeassistant/select/homeassistant_select.cpp @@ -6,6 +6,7 @@ #include "esphome/components/api/api_server.h" #include "esphome/core/helpers.h" #include "esphome/core/log.h" +#include "esphome/core/progmem.h" #include "esphome/core/string_ref.h" namespace esphome::homeassistant { @@ -237,27 +238,26 @@ void HomeassistantSelect::control(size_t index) { } #endif - static constexpr auto SERVICE_SELECT = StringRef::from_lit("select.select_option"); - static constexpr auto SERVICE_INPUT_SELECT = StringRef::from_lit("input_select.select_option"); - static constexpr auto ENTITY_ID_KEY = StringRef::from_lit("entity_id"); - static constexpr auto OPTION_KEY = StringRef::from_lit("option"); - static constexpr char INPUT_PREFIX[] = "input_"; + // In flash on the ESP8266 and encoded from there + static const char SERVICE_SELECT[] PROGMEM = "select.select_option"; + static const char SERVICE_INPUT_SELECT[] PROGMEM = "input_select.select_option"; + static const char ENTITY_ID_KEY[] PROGMEM = "entity_id"; + static const char OPTION_KEY[] PROGMEM = "option"; api::HomeassistantActionRequest resp; - if (strncmp(this->entity_id_, INPUT_PREFIX, sizeof(INPUT_PREFIX) - 1) == 0) { - resp.service = SERVICE_INPUT_SELECT; - } else { - resp.service = SERVICE_SELECT; - } + // An input helper entity id starts with the "input_" its action name starts with + const bool input_helper = ESPHOME_strncmp_P(this->entity_id_, SERVICE_INPUT_SELECT, sizeof("input_") - 1) == 0; + resp.service = + input_helper ? ProgmemStringRef::from_lit(SERVICE_INPUT_SELECT) : ProgmemStringRef::from_lit(SERVICE_SELECT); resp.data.init(2); auto &entity_id = resp.data.emplace_back(); - entity_id.key = ENTITY_ID_KEY; - entity_id.value = StringRef(this->entity_id_); + entity_id.key = ProgmemStringRef::from_lit(ENTITY_ID_KEY); + entity_id.value = ProgmemStringRef(StringRef(this->entity_id_)); auto &option = resp.data.emplace_back(); - option.key = OPTION_KEY; - option.value = StringRef(this->option_at(index)); + option.key = ProgmemStringRef::from_lit(OPTION_KEY); + option.value = ProgmemStringRef(StringRef(this->option_at(index))); api::global_api_server->send_homeassistant_action(resp); } diff --git a/esphome/components/homeassistant/switch/homeassistant_switch.cpp b/esphome/components/homeassistant/switch/homeassistant_switch.cpp index 841f16df200..f6fbdc1ef48 100644 --- a/esphome/components/homeassistant/switch/homeassistant_switch.cpp +++ b/esphome/components/homeassistant/switch/homeassistant_switch.cpp @@ -1,6 +1,7 @@ #include "homeassistant_switch.h" #include "esphome/components/api/api_server.h" #include "esphome/core/log.h" +#include "esphome/core/progmem.h" #include "esphome/core/string_ref.h" namespace esphome::homeassistant { @@ -47,21 +48,18 @@ void HomeassistantSwitch::write_state(bool state) { } #endif - static constexpr auto SERVICE_ON = StringRef::from_lit("homeassistant.turn_on"); - static constexpr auto SERVICE_OFF = StringRef::from_lit("homeassistant.turn_off"); - static constexpr auto ENTITY_ID_KEY = StringRef::from_lit("entity_id"); + // In flash on the ESP8266 and encoded from there + static const char SERVICE_ON[] PROGMEM = "homeassistant.turn_on"; + static const char SERVICE_OFF[] PROGMEM = "homeassistant.turn_off"; + static const char ENTITY_ID_KEY[] PROGMEM = "entity_id"; api::HomeassistantActionRequest resp; - if (state) { - resp.service = SERVICE_ON; - } else { - resp.service = SERVICE_OFF; - } + resp.service = state ? ProgmemStringRef::from_lit(SERVICE_ON) : ProgmemStringRef::from_lit(SERVICE_OFF); resp.data.init(1); auto &entity_id_kv = resp.data.emplace_back(); - entity_id_kv.key = ENTITY_ID_KEY; - entity_id_kv.value = StringRef(this->entity_id_); + entity_id_kv.key = ProgmemStringRef::from_lit(ENTITY_ID_KEY); + entity_id_kv.value = ProgmemStringRef(StringRef(this->entity_id_)); api::global_api_server->send_homeassistant_action(resp); } diff --git a/esphome/components/homeassistant/text/homeassistant_text.cpp b/esphome/components/homeassistant/text/homeassistant_text.cpp index 7155f6ee189..2a5598756a5 100644 --- a/esphome/components/homeassistant/text/homeassistant_text.cpp +++ b/esphome/components/homeassistant/text/homeassistant_text.cpp @@ -6,6 +6,7 @@ #include "esphome/components/api/api_server.h" #include "esphome/core/helpers.h" #include "esphome/core/log.h" +#include "esphome/core/progmem.h" #include "esphome/core/string_ref.h" namespace esphome::homeassistant { @@ -86,27 +87,26 @@ void HomeassistantText::control(const std::string &value) { } #endif - static constexpr auto SERVICE_TEXT = StringRef::from_lit("text.set_value"); - static constexpr auto SERVICE_INPUT_TEXT = StringRef::from_lit("input_text.set_value"); - static constexpr auto ENTITY_ID_KEY = StringRef::from_lit("entity_id"); - static constexpr auto VALUE_KEY = StringRef::from_lit("value"); - static constexpr char INPUT_PREFIX[] = "input_"; + // In flash on the ESP8266 and encoded from there + static const char SERVICE_TEXT[] PROGMEM = "text.set_value"; + static const char SERVICE_INPUT_TEXT[] PROGMEM = "input_text.set_value"; + static const char ENTITY_ID_KEY[] PROGMEM = "entity_id"; + static const char VALUE_KEY[] PROGMEM = "value"; api::HomeassistantActionRequest resp; - if (strncmp(this->entity_id_, INPUT_PREFIX, sizeof(INPUT_PREFIX) - 1) == 0) { - resp.service = SERVICE_INPUT_TEXT; - } else { - resp.service = SERVICE_TEXT; - } + // An input helper entity id starts with the "input_" its action name starts with + const bool input_helper = ESPHOME_strncmp_P(this->entity_id_, SERVICE_INPUT_TEXT, sizeof("input_") - 1) == 0; + resp.service = + input_helper ? ProgmemStringRef::from_lit(SERVICE_INPUT_TEXT) : ProgmemStringRef::from_lit(SERVICE_TEXT); resp.data.init(2); auto &entity_id_kv = resp.data.emplace_back(); - entity_id_kv.key = ENTITY_ID_KEY; - entity_id_kv.value = StringRef(this->entity_id_); + entity_id_kv.key = ProgmemStringRef::from_lit(ENTITY_ID_KEY); + entity_id_kv.value = ProgmemStringRef(StringRef(this->entity_id_)); auto &value_kv = resp.data.emplace_back(); - value_kv.key = VALUE_KEY; - value_kv.value = StringRef(value.data(), value.size()); + value_kv.key = ProgmemStringRef::from_lit(VALUE_KEY); + value_kv.value = ProgmemStringRef(value.data(), value.size()); api::global_api_server->send_homeassistant_action(resp); } diff --git a/esphome/core/defines.h b/esphome/core/defines.h index e0164c69245..3b499594944 100644 --- a/esphome/core/defines.h +++ b/esphome/core/defines.h @@ -322,7 +322,6 @@ #define USE_API_USER_DEFINED_ACTION_RESPONSES #define USE_API_USER_DEFINED_ACTION_RESPONSES_JSON #define API_MAX_SEND_QUEUE 8 -#define API_USER_ACTION_STRINGS_SCRATCH_SIZE 64 #define API_WIZARD_DATA_SIZE 1 #define API_WIZARD_INPUT_COUNT 1 #define MAX_API_CONNECTIONS 6 diff --git a/esphome/core/string_ref.h b/esphome/core/string_ref.h index e9e31411e42..7deea8ff664 100644 --- a/esphome/core/string_ref.h +++ b/esphome/core/string_ref.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -189,24 +190,6 @@ inline bool operator!=(const StringRef &lhs, const __FlashStringHelper *rhs) { r inline bool operator!=(const __FlashStringHelper *lhs, const StringRef &rhs) { return !(rhs == lhs); } -/// Codegen use only: a flash string literal with its byte length, compared without a runtime strlen. -struct ProgmemStringRef { - PGM_P str; - size_t len; -}; - -inline bool operator==(const std::string &lhs, ProgmemStringRef rhs) { - return lhs.size() == rhs.len && (rhs.len == 0 || memcmp_P(lhs.data(), rhs.str, rhs.len) == 0); -} - -inline bool operator!=(const std::string &lhs, ProgmemStringRef rhs) { return !(lhs == rhs); } - -inline bool operator==(const StringRef &lhs, ProgmemStringRef rhs) { - // A StringRef may be {nullptr, 0}, and memcmp_P with a null pointer is undefined even for 0 bytes - return lhs.size() == rhs.len && (rhs.len == 0 || memcmp_P(lhs.c_str(), rhs.str, rhs.len) == 0); -} - -inline bool operator!=(const StringRef &lhs, ProgmemStringRef rhs) { return !(lhs == rhs); } #endif // USE_ESP8266 inline bool operator<(const StringRef &lhs, const StringRef &rhs) { @@ -241,6 +224,49 @@ inline std::string operator+(const std::string &lhs, const StringRef &rhs) { str.append(rhs.c_str(), rhs.size()); return str; } + +/// A string that may live in flash (on ESP8266): pointer and length, with no byte-reading operations. +/// progmem_ptr() is safe as a "%s" log argument; compare the contents with equals(). +class ProgmemStringRef { + public: + constexpr ProgmemStringRef() = default; + constexpr ProgmemStringRef(const char *s, size_t len) : ref_(s, len) {} + constexpr explicit ProgmemStringRef(const StringRef &s) : ref_(s) {} + /// A PROGMEM char array, with its length known at compile time + template static constexpr ProgmemStringRef from_lit(const char (&s)[N]) { return {s, N - 1}; } + + constexpr const char *progmem_ptr() const { return this->ref_.c_str(); } + constexpr size_t size() const { return this->ref_.size(); } + constexpr bool empty() const { return this->ref_.empty(); } + + bool equals(const StringRef &other) const { +#ifdef USE_ESP8266 + // A StringRef may be {nullptr, 0}, and memcmp_P with a null pointer is undefined even for 0 bytes + return other.size() == this->ref_.size() && + (other.empty() || memcmp_P(other.c_str(), this->ref_.c_str(), other.size()) == 0); +#else + return other == this->ref_; +#endif + } + + // Remove before 2027.5.0 + [[deprecated( + "May point to flash; use progmem_ptr() or the owner's accessors. Removed in 2027.5.0")]] constexpr const char * + c_str() const { + return this->ref_.c_str(); + } + + private: + StringRef ref_; +}; + +#ifdef USE_ESP8266 +inline bool operator==(const StringRef &lhs, const ProgmemStringRef &rhs) { return rhs.equals(lhs); } +inline bool operator!=(const StringRef &lhs, const ProgmemStringRef &rhs) { return !rhs.equals(lhs); } +inline bool operator==(const std::string &lhs, const ProgmemStringRef &rhs) { return rhs.equals(StringRef(lhs)); } +inline bool operator!=(const std::string &lhs, const ProgmemStringRef &rhs) { return !rhs.equals(StringRef(lhs)); } +#endif + // String conversion functions for ADL compatibility (allows stoi(x) where x is StringRef) // Must be in esphome namespace for ADL to find them. Uses strtol/strtod directly to avoid heap allocation. namespace internal { diff --git a/script/api_protobuf/api_protobuf.py b/script/api_protobuf/api_protobuf.py index a143a850326..2f601f5c947 100755 --- a/script/api_protobuf/api_protobuf.py +++ b/script/api_protobuf/api_protobuf.py @@ -536,16 +536,19 @@ def create_field_type_info( f"track_presence on field '{field.name}' has no effect; it requires " "a non-repeated message field in a message that is decoded" ) - if get_field_opt(field, pb.progmem, False) and ( - field.label == FieldDescriptorProto.LABEL_REPEATED - or field.type != 12 - or not get_field_opt(field, pb.pointer_to_buffer, False) - or get_field_opt(field, pb.fixed_array_size) is not None - ): - raise ValueError( - f"progmem on field '{field.name}' requires a non-repeated bytes field " - "with pointer_to_buffer" + if get_field_opt(field, pb.progmem, False): + is_pointer_bytes = ( + field.type == FieldDescriptorProto.TYPE_BYTES + and get_field_opt(field, pb.pointer_to_buffer, False) + and get_field_opt(field, pb.fixed_array_size) is None ) + if field.label == FieldDescriptorProto.LABEL_REPEATED or not ( + field.type == FieldDescriptorProto.TYPE_STRING or is_pointer_bytes + ): + raise ValueError( + f"progmem on field '{field.name}' requires a non-repeated string field, " + "or a bytes field with pointer_to_buffer" + ) if field.label == FieldDescriptorProto.LABEL_REPEATED: # Check if this is a packed_buffer field (zero-copy packed repeated) if get_field_opt(field, pb.packed_buffer, False): @@ -1155,6 +1158,11 @@ class PointerToBufferTypeBase(TypeInfo): """Get the wire type for this field.""" return WireType.LENGTH_DELIMITED # Uses wire type 2 + @property + def progmem(self) -> bool: + """Whether the data is in flash, so encode and dump copy it with progmem_memcpy.""" + return get_field_opt(self._field, pb.progmem, False) + def get_estimated_size(self) -> int: # field ID + length varint + typical data (assume small for pointer fields) return self.calculate_field_id_size() + 2 + 16 @@ -1168,11 +1176,6 @@ class PointerToBytesBufferType(PointerToBufferTypeBase): reference_type = "const uint8_t*" const_reference_type = "const uint8_t*" - @property - def progmem(self) -> bool: - """Whether the data is in flash, so encode and dump copy it with progmem_memcpy.""" - return get_field_opt(self._field, pb.progmem, False) - @property def public_content(self) -> list[str]: # Use uint16_t for length - max packet size is well below 65535 @@ -1252,6 +1255,8 @@ class PointerToStringBufferType(PointerToBufferTypeBase): @property def public_content(self) -> list[str]: + if self.progmem: + return [f"ProgmemStringRef {self.field_name}{{}};"] if self._starts_null: return [ f"StringRef {self.field_name}{{nullptr, 0}}; // null until set, encode only" @@ -1267,9 +1272,19 @@ class PointerToStringBufferType(PointerToBufferTypeBase): ) tag = self.calculate_tag() if tag < 128: - return _encode_call( - "encode_short_string_force", str(tag), f"this->{self.field_name}" + encode_fn = ( + "encode_short_progmem_string_force" + if self.progmem + else "encode_short_string_force" ) + return _encode_call(encode_fn, str(tag), f"this->{self.field_name}") + if self.progmem: + return _encode_call( + "encode_progmem_string", + str(self.number), + f"this->{self.field_name}", + force=self.force, + ) if result := self._encode_bytes_with_precomputed_tag( f"this->{self.field_name}.c_str()", f"this->{self.field_name}.size()", @@ -1297,6 +1312,8 @@ class PointerToStringBufferType(PointerToBufferTypeBase): @property def dump_content(self) -> str: + if self.progmem: + return f'dump_progmem_string_field(out, ESPHOME_PSTR("{self.name}"), this->{self.field_name});' return f'dump_field(out, ESPHOME_PSTR("{self.name}"), this->{self.field_name});' def get_size_calculation(self, name: str, force: bool = False) -> str: @@ -3246,8 +3263,9 @@ namespace esphome::api { #ifdef USE_ESP8266 // Out-of-line to avoid inlining strlen_P/memcpy_P at every call site -void DumpBuffer::append_p_esp8266(const char *str) { - size_t len = strlen_P(str); +void DumpBuffer::append_p_esp8266(const char *str) { this->append_p_esp8266(str, strlen_P(str)); } + +void DumpBuffer::append_p_esp8266(const char *str, size_t len) { size_t space = CAPACITY - 1 - pos_; if (len > space) len = space; @@ -3368,6 +3386,13 @@ static void dump_progmem_bytes_field(DumpBuffer &out, const char *field_name, co progmem_memcpy(data_buf, data, len); dump_bytes_field(out, field_name, data_buf, len, indent); } + +// Helper for string fields in flash +static void dump_progmem_string_field(DumpBuffer &out, const char *field_name, const ProgmemStringRef &value, + int indent = 2) { + append_field_prefix(out, field_name, indent); + out.append("'").append_p(value.progmem_ptr(), value.size()).append("'\\n"); +} #pragma GCC diagnostic pop """ diff --git a/tests/component_tests/api/test_action_metadata.py b/tests/component_tests/api/test_action_metadata.py index 69c9605a675..2b7662b6655 100644 --- a/tests/component_tests/api/test_action_metadata.py +++ b/tests/component_tests/api/test_action_metadata.py @@ -8,18 +8,13 @@ import pytest from esphome.components.api import ( _action_strings, - _action_strings_size, _has_action_metadata, - _validate_esp8266_action_strings, validate_variable, ) from esphome.config_validation import Invalid -from esphome.const import PlatformFramework from esphome.core import CORE from esphome.cpp_generator import safe_exp from esphome.helpers import fnv1_hash -from tests.component_tests.helpers import get_define_value -from tests.component_tests.types import SetCoreConfigCallable CONFIG = "tests/component_tests/api/test_action_metadata.yaml" CONFIG_ESP8266 = "tests/component_tests/api/test_action_metadata_esp8266.yaml" @@ -45,7 +40,6 @@ def test_metadata_is_emitted_as_progmem_table( ) assert f"(api_action0_strings, {safe_exp(fnv1_hash('play_buzzer'))});" in main_cpp assert "USE_API_USER_DEFINED_ACTION_METADATA" in {d.name for d in CORE.defines} - assert get_define_value("API_USER_ACTION_STRINGS_SCRATCH_SIZE") is None def test_esp8266_strings_are_shared_progmem_arrays( @@ -69,15 +63,13 @@ def test_esp8266_strings_are_shared_progmem_arrays( ) -def test_esp8266_sizes_scratch_buffer_for_largest_action( +def test_esp8266_encodes_action_strings_from_flash( generate_main: Callable[[str | Path], str], ) -> None: - """ESP8266 gets a scratch buffer define equal to the byte total of the largest action.""" - generate_main(CONFIG_ESP8266) + """ESP8266 uses the same flash string table.""" + main_cpp = generate_main(CONFIG_ESP8266) - # play_buzzer: name, description, two variable names, one description, one example, - # each with a terminator - assert get_define_value("API_USER_ACTION_STRINGS_SCRATCH_SIZE") == "117" + assert "api_action0_strings[] PROGMEM" in main_cpp def test_shorthand_variables_emit_no_metadata( @@ -113,36 +105,8 @@ def test_variable_rejects_invalid(value: object) -> None: validate_variable(value) -def _oversized_action_config() -> dict: - return { - "actions": [ - { - "action": "big", - "description": "x" * 300, - "variables": {"a": {"type": "string", "example": "y" * 300}}, - } - ] - } - - -def test_esp8266_rejects_actions_over_string_budget( - set_core_config: SetCoreConfigCallable, -) -> None: - set_core_config(PlatformFramework.ESP8266_ARDUINO) - with pytest.raises(Invalid, match="ESP8266 allows at most 384 bytes"): - _validate_esp8266_action_strings(_oversized_action_config()) - - -def test_other_platforms_have_no_string_budget( - set_core_config: SetCoreConfigCallable, -) -> None: - set_core_config(PlatformFramework.ESP32_IDF) - config = _oversized_action_config() - assert _validate_esp8266_action_strings(config) is config - - -def test_empty_metadata_is_unset_and_not_counted() -> None: - """An empty description or example emits nullptr and takes no scratch space.""" +def test_empty_metadata_is_unset() -> None: + """An empty description or example emits nullptr.""" conf = { "action": "a", "description": "", @@ -150,8 +114,6 @@ def test_empty_metadata_is_unset_and_not_counted() -> None: } strings = _action_strings(conf, has_metadata=True) assert strings == ["a", None, "b", None, "ex"] - # Every emitted string counts its terminator: "a" + "b" + "ex" - assert _action_strings_size(strings) == 2 + 2 + 3 def test_empty_metadata_does_not_enable_the_define() -> None: diff --git a/tests/components/core/test_progmem_string_ref.cpp b/tests/components/core/test_progmem_string_ref.cpp new file mode 100644 index 00000000000..1d996f36db7 --- /dev/null +++ b/tests/components/core/test_progmem_string_ref.cpp @@ -0,0 +1,30 @@ +#include + +#include "esphome/core/progmem.h" +#include "esphome/core/string_ref.h" + +namespace esphome::core::testing { + +TEST(ProgmemStringRef, EqualsComparesLengthAndContent) { + ProgmemStringRef ref("Kitchen", 7); + EXPECT_TRUE(ref.equals(StringRef("Kitchen"))); + EXPECT_FALSE(ref.equals(StringRef("Kitche"))); + EXPECT_FALSE(ref.equals(StringRef("Kitchens"))); + EXPECT_FALSE(ref.equals(StringRef("Bedroom"))); +} + +TEST(ProgmemStringRef, FromLitTakesTheLengthAtCompileTime) { + static const char NAME[] PROGMEM = "entity_id"; + constexpr ProgmemStringRef ref = ProgmemStringRef::from_lit(NAME); + static_assert(ref.size() == 9); + EXPECT_TRUE(ref.equals(StringRef("entity_id"))); +} + +TEST(ProgmemStringRef, DefaultIsEmpty) { + ProgmemStringRef ref; + EXPECT_TRUE(ref.empty()); + EXPECT_EQ(ref.size(), 0u); + EXPECT_TRUE(ref.equals(StringRef(""))); +} + +} // namespace esphome::core::testing diff --git a/tests/unit_tests/components/api/test_api_protobuf_generator.py b/tests/unit_tests/components/api/test_api_protobuf_generator.py index 5e29b06ef07..9d8eab69b90 100644 --- a/tests/unit_tests/components/api/test_api_protobuf_generator.py +++ b/tests/unit_tests/components/api/test_api_protobuf_generator.py @@ -408,8 +408,39 @@ def test_progmem_rejected_in_a_decoded_message() -> None: def test_progmem_rejected_on_non_bytes_field() -> None: - """(progmem) only applies to bytes fields.""" - field = _field(9) + """(progmem) only applies to string and bytes fields.""" + field = _field(13) + field.options.Extensions[pb.progmem] = True + with pytest.raises(ValueError, match="progmem on field 'value'"): + create_field_type_info(field, needs_decode=False) + + +@pytest.mark.parametrize( + ("force", "max_len", "encode_fn"), + [ + (True, 120, "encode_short_progmem_string_force("), + (False, None, "encode_progmem_string("), + (True, None, "encode_progmem_string_force("), + ], +) +def test_progmem_string_field_copies_from_flash( + force: bool, max_len: int | None, encode_fn: str +) -> None: + """A (progmem) string field is a ProgmemStringRef that encodes and dumps through the progmem helpers.""" + field = _field(9, force=force) + field.options.Extensions[pb.progmem] = True + if max_len is not None: + field.options.Extensions[pb.max_data_length] = max_len + ti = create_field_type_info(field, needs_decode=False) + assert "ProgmemStringRef value{};" in ti.public_content + assert encode_fn in ti.encode_content + assert "encode_string" not in ti.encode_content + assert "dump_progmem_string_field(" in ti.dump_content + + +def test_progmem_rejected_on_repeated_string() -> None: + """(progmem) on a repeated string field fails instead of silently using memcpy.""" + field = _field(9, repeated=True) field.options.Extensions[pb.progmem] = True with pytest.raises(ValueError, match="progmem on field 'value'"): create_field_type_info(field, needs_decode=False)