github/esphome
2
0
Fork 0
mirror of https://github.com/esphome/esphome.git synced 2026-10-20 13:06:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

[ota] Fall back to a UDP data phase when the link is lossy

Browse Source 
TCP gives up on links that lose many packets: retransmit timers back off and
the connection drops, so devices on poor Wi-Fi cannot be updated without being
moved. After the TCP handshake, a client that saw loss moves the upload to a
UDP byte stream that the device acks from lwIP's receive callback as datagrams
arrive, keeping at most four of them like TCP's window. Clean links stay on
TCP; blocked UDP falls back to TCP.
This commit is contained in:
J. Nick Koston
2026-10-10 12:15:36 -10:00
parent 721228c1d2
commit 54b34edf2d
6 changed files with 945 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
esphome/components/esphome/ota/ota_esphome.cpp
+77 -5
View File
@@ -44,11 +44,14 @@ const noise::NoiseContext &ESPHomeOTAComponent::noise_context_() const {
}
#endif
static constexpr uint16_t OTA_BLOCK_SIZE = 8192;
static constexpr uint32_t OTA_SOCKET_TIMEOUT_HANDSHAKE = 20000; // milliseconds for initial handshake
static constexpr uint32_t OTA_SOCKET_TIMEOUT_HANDSHAKE = 60000; // milliseconds for initial handshake
// Milliseconds for data transfer. Covers the lwIP retransmit run seen in
// practice for a lost chunk ack (1.5 + 3 + 6 + 12 + 24 + 48 s); the CLI waits
// longer (espota2.DATA_PHASE_TIMEOUT) so the device is free before it retries
static constexpr uint32_t OTA_SOCKET_TIMEOUT_DATA = 105000;
// A handshake this old gives way to a waiting connection: a client retrying over a lossy link otherwise queues
// behind its own stalled attempt until both time out
static constexpr uint32_t OTA_HANDSHAKE_PREEMPT = 10000;
static constexpr uint32_t OTA_PROGRESS_INTERVAL_MS = 1000;
static constexpr size_t OTA_SIZE_FIELD_BYTES = 4; // sizes on the wire are 4 bytes MSB first
@@ -199,6 +202,9 @@ static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_SHA256_AUTH = 0x02;
static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL = 0x04;
static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_NOISE = 0x08;
static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_DEFLATE = 0x10;
#ifdef ESPHOME_OTA_UDP
static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_UDP = 0x20;
#endif
// Noise needs the extended protocol: the prologue binds the 2-byte feature ack
static constexpr uint8_t CLIENT_NOISE_FEATURES =
CLIENT_FEATURE_SUPPORTS_NOISE | CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL;
@@ -208,6 +214,9 @@ static constexpr uint8_t SERVER_FEATURE_SUPPORTS_NOISE = 0x04;
// Raw deflate, window <= OTA_INFLATE_WINDOW_SIZE. Binding once offered: the
// client must then send the image size frame and a deflate stream.
static constexpr uint8_t SERVER_FEATURE_SUPPORTS_DEFLATE = 0x08;
#ifdef ESPHOME_OTA_UDP
static constexpr uint8_t SERVER_FEATURE_SUPPORTS_UDP = 0x10;
#endif
#ifdef USE_OTA_ENCRYPTION
inline bool ESPHomeOTAComponent::noise_offered_() const {
@@ -266,6 +275,12 @@ void ESPHomeOTAComponent::handle_handshake_() {
this->cleanup_connection_();
return;
}
// Only unauthenticated handshakes reach here; the data phase blocks inside handle_data_
if (now - this->client_connect_time_ > OTA_HANDSHAKE_PREEMPT && this->server_->ready()) {
ESP_LOGD(TAG, "Newer connection waiting");
this->cleanup_connection_();
return;
}
switch (this->ota_state_) {
case OTAState::MAGIC_READ: {
@@ -329,6 +344,9 @@ void ESPHomeOTAComponent::handle_handshake_() {
static_assert(HANDSHAKE_BUF_SIZE >= 2, "handshake_buf_ must hold the 2-byte extended-protocol feature ack");
this->handshake_buf_[0] = ota::OTA_RESPONSE_FEATURE_FLAGS;
this->handshake_buf_[1] = (supports_compression ? SERVER_FEATURE_SUPPORTS_COMPRESSION : 0);
#ifdef ESPHOME_OTA_UDP
this->handshake_buf_[1] |= SERVER_FEATURE_SUPPORTS_UDP;
#endif
#ifdef USE_OTA_PARTITIONS
this->handshake_buf_[1] |= SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS;
#endif
@@ -483,6 +501,10 @@ void ESPHomeOTAComponent::handle_data_() {
char *sbuf = reinterpret_cast<char *>(buf);
size_t image_size;
ota::OTAType ota_type = ota::OTA_TYPE_UPDATE_APP;
#ifdef ESPHOME_OTA_UDP
UdpLink udp_link;
this->udp_ = nullptr; // client_connect_time_ is done with
#endif
// Set socket timeouts and blocking mode (see strategy table above)
struct timeval tv;
@@ -498,6 +520,34 @@ void ESPHomeOTAComponent::handle_data_() {
// Acknowledge auth OK - 1 byte
this->data_write_byte_(ota::OTA_RESPONSE_AUTH_OK);
#ifdef ESPHOME_OTA_UDP
// The server always offers UDP on the extended protocol
if (this->extended_proto_() && (this->ota_features_ & CLIENT_FEATURE_SUPPORTS_UDP) != 0) {
// Whether the client wants UDP (0: its link looks good), answered with our token
if (!this->data_readall_(buf, 1)) {
if (this->client_left_before_start_())
return;
this->log_read_error_(LOG_STR("UDP request"));
goto error; // NOLINT(cppcoreguidelines-avoid-goto)
}
// A reply left all zero keeps the session on TCP
if (buf[0] != 0)
this->udp_open_(udp_link);
if (!this->data_write_(udp_link.reply, sizeof(udp_link.reply)))
goto error; // NOLINT(cppcoreguidelines-avoid-goto)
if (udp_link.reply[0] != 0) {
this->udp_ = &udp_link;
if (!this->udp_start_()) {
// Nothing started yet: just drop the session, which the client retries as a network failure
this->log_read_error_(LOG_STR("UDP start"));
this->cleanup_connection_();
return;
}
}
}
#endif
if (this->extended_proto_()) {
// Read ota type, 1 byte
if (!this->data_readall_(buf, 1)) {
@@ -623,6 +673,9 @@ void ESPHomeOTAComponent::handle_data_() {
error:
this->data_write_byte_(static_cast<uint8_t>(error_code));
#ifdef ESPHOME_OTA_UDP
this->udp_linger_();
#endif
// Abort backend before cleanup - cleanup_connection_() destroys the backend.
// Always call abort() unconditionally: backends register external partitions before
@@ -651,7 +704,7 @@ bool ESPHomeOTAComponent::readall_(uint8_t *buf, size_t len) {
return false;
}
ssize_t read = this->client_->read(buf + at, len - at);
ssize_t read = this->transport_read_(buf + at, len - at);
if (read == -1) {
const int err = errno;
if (!this->would_block_(err)) {
@@ -676,6 +729,10 @@ bool ESPHomeOTAComponent::readall_(uint8_t *buf, size_t len) {
return true;
}
bool ESPHomeOTAComponent::writeall_(const uint8_t *buf, size_t len) {
#ifdef ESPHOME_OTA_UDP
if (this->udp_ != nullptr)
return this->udp_send_(buf, len);
#endif
uint32_t start = millis();
uint32_t at = 0;
while (len - at > 0) {
@@ -728,7 +785,15 @@ void ESPHomeOTAComponent::log_read_error_(const LogString *what) {
void ESPHomeOTAComponent::log_start_(const LogString *phase) {
char peername[socket::SOCKADDR_STR_LEN];
this->client_->getpeername_to(peername);
#ifdef ESPHOME_OTA_UDP
if (this->client_ == nullptr) {
// A committed UDP session has already dropped TCP
ipaddr_ntoa_r(&this->udp_->peer_ip, peername, sizeof(peername));
} else
#endif
{
this->client_->getpeername_to(peername);
}
ESP_LOGD(TAG, "Starting %s from %s", LOG_STR_ARG(phase), peername);
}
@@ -840,7 +905,7 @@ ssize_t ESPHomeOTAComponent::receive_data_(uint8_t *buf, DataTransfer &xfer) {
break;
}
#endif
read = this->client_->read(buf, requested);
read = this->transport_read_(buf, requested);
if (read > 0)
break;
if (read == 0) {
@@ -874,6 +939,11 @@ ssize_t ESPHomeOTAComponent::receive_data_(uint8_t *buf, DataTransfer &xfer) {
void ESPHomeOTAComponent::send_chunk_acks_(DataTransfer &xfer) {
#if USE_OTA_VERSION == 2
#ifdef ESPHOME_OTA_UDP
// UDP acks every message itself
if (this->udp_ != nullptr)
return;
#endif
while (xfer.acknowledged + OTA_BLOCK_SIZE <= xfer.total ||
(xfer.total == xfer.ota_size && xfer.acknowledged < xfer.ota_size)) {
this->data_write_byte_(ota::OTA_RESPONSE_CHUNK_OK);
@@ -883,7 +953,9 @@ void ESPHomeOTAComponent::send_chunk_acks_(DataTransfer &xfer) {
}
void ESPHomeOTAComponent::cleanup_connection_() {
this->client_->close();
// Already gone once a UDP session committed
if (this->client_ != nullptr)
this->client_->close();
this->client_ = nullptr;
this->client_connect_time_ = 0;
this->handshake_buf_pos_ = 0;
esphome/components/esphome/ota/ota_esphome.h
+73 -5
View File
@@ -15,6 +15,20 @@
#include "esphome/core/preferences.h"
#include "esphome/core/hash_base.h"
// The UDP data phase answers datagrams from lwIP's receive callback, so it needs the raw lwIP API from the main
// loop: single threaded (ESP8266), locked (RP2040, ESP32 with core locking) or LibreTiny (declines without core
// locking)
#ifdef USE_ESP32
#include <sdkconfig.h>
#endif
#if defined(USE_ESP8266) || defined(USE_RP2) || defined(USE_LIBRETINY) || \
(defined(USE_ESP32) && defined(CONFIG_LWIP_TCPIP_CORE_LOCKING))
#define ESPHOME_OTA_UDP
#include "lwip/ip_addr.h"
struct pbuf;
struct udp_pcb;
#endif
namespace esphome {
/// ESPHomeOTAComponent provides a simple way to integrate Over-the-Air updates into your app using ArduinoOTA.
@@ -65,6 +79,51 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
protected:
void handle_handshake_();
void handle_data_();
#ifdef ESPHOME_OTA_UDP
// UDP transport for the data phase; wire format in ota_esphome_udp.cpp
static constexpr size_t UDP_OUT_SIZE = 64; // our unacked responses: at most three encrypted ones
// Messages the device keeps unread: like TCP's 4 segment window, so no more memory than TCP holds
static constexpr uint8_t UDP_WINDOW = 4;
// Lives on handle_data_'s stack; the receive callback shares it under the lwIP lock
struct UdpLink {
~UdpLink();
void release(); // closes the pcb and frees what is queued
struct udp_pcb *pcb{nullptr};
ip_addr_t peer_ip;
uint16_t peer_port{0}; // learned from the client's datagrams
// The session token, our answer to the client: top bit clear, so never an error code; all zero declines
uint8_t reply[4]{};
struct pbuf *queue[UDP_WINDOW]{}; // slot seq % UDP_WINDOW for seq in [consumed, consumed + UDP_WINDOW)
uint16_t received{0}; // every seq before it is here
uint16_t consumed{0}; // the message being read
uint16_t read_offset{0}; // payload bytes of it already read
uint16_t out_base{0}; // stream offset of out[0]
uint8_t out_len{0};
bool committed{false}; // the client heard us over UDP
uint8_t out[UDP_OUT_SIZE];
};
// Fills link.reply on success; leaves it all zero (declined) otherwise
void udp_open_(UdpLink &link);
// Waits for the client to commit to UDP or fall back to TCP; false on failure
bool udp_start_();
// Like a blocking socket read with SO_RCVTIMEO: up to len bytes, or -1 with EWOULDBLOCK
ssize_t udp_read_(uint8_t *buf, size_t len);
bool udp_send_(const uint8_t *data, size_t len);
// In a UDP session, keeps answering until the client has our responses (an error code), briefly
void udp_linger_();
static void udp_recv_cb_(void *arg, struct udp_pcb *pcb, struct pbuf *p, const ip_addr_t *addr, uint16_t port);
// Caller holds the lwIP lock (or runs in the receive callback)
static void udp_send_ack_(UdpLink &link, uint16_t prompted_by);
#endif
// The data phase reads TCP or, once committed, the UDP stream
inline ssize_t transport_read_(uint8_t *buf, size_t len) {
#ifdef ESPHOME_OTA_UDP
if (this->udp_ != nullptr)
return this->udp_read_(buf, len);
#endif
return this->client_->read(buf, len);
}
#ifdef USE_OTA_PASSWORD
static constexpr size_t SHA256_HEX_SIZE = 64; // SHA256 hash as hex string (32 bytes * 2)
bool handle_auth_send_();
@@ -75,7 +134,6 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
#endif // USE_OTA_PASSWORD
bool readall_(uint8_t *buf, size_t len);
bool writeall_(const uint8_t *buf, size_t len);
inline bool write_byte_(uint8_t byte) { return this->writeall_(&byte, 1); }
#ifdef USE_OTA_ENCRYPTION
// Heap-allocated only while an encrypted OTA session is active.
@@ -105,18 +163,21 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
ssize_t noise_read_frame_blocking_(uint8_t *buf, size_t min_ciphertext, size_t max_ciphertext);
bool noise_readall_(uint8_t *buf, size_t len);
ssize_t noise_read_data_(uint8_t *buf, size_t capacity);
bool noise_write_byte_(uint8_t byte);
static constexpr size_t NOISE_MAX_RESPONSE = 4; // the UDP token
bool noise_write_(const uint8_t *data, size_t len);
#endif // USE_OTA_ENCRYPTION
// Data-phase I/O dispatch: through the noise transport when a session is
// active, straight to the socket otherwise.
inline bool data_write_byte_(uint8_t byte) {
// len is at most NOISE_MAX_RESPONSE
inline bool data_write_(const uint8_t *data, size_t len) {
#ifdef USE_OTA_ENCRYPTION
if (this->noise_ != nullptr)
return this->noise_write_byte_(byte);
return this->noise_write_(data, len);
#endif
return this->write_byte_(byte);
return this->writeall_(data, len);
}
inline bool data_write_byte_(uint8_t byte) { return this->data_write_(&byte, 1); }
// When encrypted, buf must have room for len + noise::MAC_SIZE bytes.
inline bool data_readall_(uint8_t *buf, size_t len) {
#ifdef USE_OTA_ENCRYPTION
@@ -200,7 +261,14 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
std::unique_ptr<socket::Socket> client_;
ota::OTABackendPtr backend_;
#ifdef ESPHOME_OTA_UDP
union {
uint32_t client_connect_time_{0}; // handshake states
UdpLink *udp_; // data phase (never both), points into handle_data_'s stack
};
#else
uint32_t client_connect_time_{0};
#endif
static constexpr size_t HANDSHAKE_BUF_SIZE = 5;
// Buffer size for OTA data transfer. The upload client derives its maximum
// encrypted frame plaintext from this (espota2.NOISE_MAX_PLAINTEXT is this
esphome/components/esphome/ota/ota_esphome_noise.cpp
+5 -5
View File
@@ -266,13 +266,13 @@ ssize_t ESPHomeOTAComponent::noise_read_data_(uint8_t *buf, size_t capacity) {
return this->noise_read_frame_blocking_(buf, noise::MAC_SIZE + 1, max_ciphertext);
}
/// Blocking write of one response byte as an encrypted frame.
bool ESPHomeOTAComponent::noise_write_byte_(uint8_t byte) {
uint8_t frame[noise::FRAME_HEADER_SIZE + 1 + noise::MAC_SIZE];
frame[noise::FRAME_HEADER_SIZE] = byte;
/// Blocking write of a short response (at most NOISE_MAX_RESPONSE bytes) as one encrypted frame.
bool ESPHomeOTAComponent::noise_write_(const uint8_t *data, size_t len) {
uint8_t frame[noise::FRAME_HEADER_SIZE + NOISE_MAX_RESPONSE + noise::MAC_SIZE];
memcpy(frame + noise::FRAME_HEADER_SIZE, data, len);
NoiseBuffer mbuf;
noise_buffer_init(mbuf);
noise_buffer_set_inout(mbuf, frame + noise::FRAME_HEADER_SIZE, 1, 1 + noise::MAC_SIZE);
noise_buffer_set_inout(mbuf, frame + noise::FRAME_HEADER_SIZE, len, len + noise::MAC_SIZE);
int err = noise_cipherstate_encrypt(this->noise_->send_cipher, &mbuf);
if (err != 0) {
ESP_LOGW(TAG, "Encrypt: %d", err);
esphome/components/esphome/ota/ota_esphome_udp.cpp
+266
View File
@@ -0,0 +1,266 @@
#include "ota_esphome.h"
#if defined(USE_OTA) && defined(ESPHOME_OTA_UDP)
#include "esphome/core/application.h"
#include "esphome/core/hal.h"
#include "esphome/core/helpers.h"
#include "esphome/core/log.h"
#include "lwip/inet.h"
#include "lwip/pbuf.h"
#include "lwip/udp.h"
#ifdef USE_LIBRETINY
#include "lwip/tcpip.h"
#endif
#include <cstring>
/*
* UDP data phase: the same byte stream as TCP, cut into numbered datagrams.
* client -> device [type | COMMITTED][token u32][seq u16][device bytes received u16][payload]
* device -> client [ACK][token u32][received u16][seq that prompted it u16][stream offset u16][window u8][our bytes]
* The lwIP callback acks each datagram on arrival, so flash writes never look like loss, and keeps messages in its
* UDP_WINDOW slots until the OTA loop reads them. "window": free slots in the high nibble, and in the low nibble bit k
* set when received + k is already here (like TCP SACK).
*/
namespace esphome {
ESPHOME_LOG_TAG(TAG, "esphome.ota");
// PING (0x02) and PROBE (0x03) only prompt an ACK
static constexpr uint8_t UDP_MSG_DATA = 0x01;
static constexpr uint8_t UDP_FLAG_COMMITTED = 0x80;
static constexpr uint8_t UDP_MSG_ACK = 0x40;
static constexpr size_t UDP_HEADER_SIZE = 9;
static constexpr size_t UDP_ACK_HEADER_SIZE = 12;
// The client's 60 s wait for the token over a lossy TCP link, plus its 5 s probe
static constexpr uint32_t UDP_COMMIT_TIMEOUT = 70000;
// Like the TCP socket's SO_RCVTIMEO in handle_data_, so readall_ treats both alike
static constexpr uint32_t UDP_READ_WAIT_MS = 2000;
static constexpr uint32_t UDP_LINGER_MS = 3000;
#ifdef USE_LIBRETINY
// LibreTiny's LwIPLock is a no-op; take the core lock directly (udp_open_ declines without one)
struct UdpLock {
#if LWIP_TCPIP_CORE_LOCKING
UdpLock() { LOCK_TCPIP_CORE(); }
~UdpLock() { UNLOCK_TCPIP_CORE(); }
#endif
};
#else
using UdpLock = LwIPLock;
#endif
static void put_u16(uint8_t *at, uint16_t value) {
at[0] = value >> 8;
at[1] = value & 0xFF;
}
ESPHomeOTAComponent::UdpLink::~UdpLink() { this->release(); }
void ESPHomeOTAComponent::UdpLink::release() {
UdpLock lock;
if (this->pcb != nullptr)
udp_remove(this->pcb);
this->pcb = nullptr;
for (auto *&p : this->queue) {
if (p != nullptr)
pbuf_free(p);
p = nullptr;
}
}
void ESPHomeOTAComponent::udp_send_ack_(UdpLink &link, uint16_t prompted_by) {
const uint16_t len = UDP_ACK_HEADER_SIZE + link.out_len;
struct pbuf *p = pbuf_alloc(PBUF_TRANSPORT, len, PBUF_RAM);
if (p == nullptr)
return; // Like a lost ACK: the client asks again
auto *ack = static_cast<uint8_t *>(p->payload);
ack[0] = UDP_MSG_ACK;
memcpy(ack + 1, link.reply, 4);
put_u16(ack + 5, link.received);
put_u16(ack + 7, prompted_by);
put_u16(ack + 9, link.out_base);
uint8_t window = (UDP_WINDOW - static_cast<uint16_t>(link.received - link.consumed)) << 4;
for (uint8_t k = 1; k < UDP_WINDOW; k++) {
if (link.queue[(link.received + k) % UDP_WINDOW] != nullptr)
window |= 1 << k;
}
ack[11] = window;
memcpy(ack + UDP_ACK_HEADER_SIZE, link.out, link.out_len);
udp_sendto(link.pcb, p, &link.peer_ip, link.peer_port);
pbuf_free(p);
}
// lwIP context: the tcpip thread (ESP32, LibreTiny), SYS (ESP8266) or an IRQ (RP2040)
void ESPHomeOTAComponent::udp_recv_cb_(void *arg, struct udp_pcb * /*pcb*/, struct pbuf *p, const ip_addr_t *addr,
uint16_t port) {
auto &link = *static_cast<UdpLink *>(arg);
uint8_t hdr[UDP_HEADER_SIZE];
if (!ip_addr_cmp(addr, &link.peer_ip) || pbuf_copy_partial(p, hdr, UDP_HEADER_SIZE, 0) != UDP_HEADER_SIZE ||
memcmp(hdr + 1, link.reply, 4) != 0) {
pbuf_free(p);
return;
}
// The token proves the sender; answer the port it arrived from, which NAT may have rewritten
link.peer_port = port;
link.committed |= (hdr[0] & UDP_FLAG_COMMITTED) != 0;
// Drop the responses the client already has
const uint16_t acked = encode_uint16(hdr[7], hdr[8]) - link.out_base;
if (acked != 0 && acked <= link.out_len) {
link.out_len -= acked;
memmove(link.out, link.out + acked, link.out_len);
link.out_base += acked;
}
const uint16_t seq = encode_uint16(hdr[5], hdr[6]);
struct pbuf *&slot = link.queue[seq % UDP_WINDOW];
if ((hdr[0] & ~UDP_FLAG_COMMITTED) == UDP_MSG_DATA && p->tot_len > UDP_HEADER_SIZE &&
static_cast<uint16_t>(seq - link.consumed) < UDP_WINDOW && slot == nullptr) {
slot = p;
while (static_cast<uint16_t>(link.received - link.consumed) < UDP_WINDOW &&
link.queue[link.received % UDP_WINDOW] != nullptr)
link.received++;
} else {
pbuf_free(p);
}
udp_send_ack_(link, seq);
}
void ESPHomeOTAComponent::udp_open_(UdpLink &link) {
#if defined(USE_LIBRETINY) && !LWIP_TCPIP_CORE_LOCKING
// No safe way to call lwIP from the main loop: decline
#else
struct sockaddr_storage peer;
socklen_t peer_len = sizeof(peer);
if (this->client_->getpeername(reinterpret_cast<struct sockaddr *>(&peer), &peer_len) != 0)
return;
uint32_t ipv4;
if (socket::sockaddr_to_ipv4(reinterpret_cast<const struct sockaddr *>(&peer), &ipv4)) {
ip_addr_set_ip4_u32(&link.peer_ip, ipv4);
} else {
#if USE_NETWORK_IPV6 && LWIP_IPV6
inet6_addr_to_ip6addr(ip_2_ip6(&link.peer_ip), &reinterpret_cast<const sockaddr_in6 *>(&peer)->sin6_addr);
IP_SET_TYPE_VAL(link.peer_ip, IPADDR_TYPE_V6);
#else
return;
#endif
}
uint8_t token[4];
if (!random_bytes(token, sizeof(token)))
return;
UdpLock lock;
#if LWIP_IPV6
link.pcb = udp_new_ip_type(IPADDR_TYPE_ANY);
#else
link.pcb = udp_new();
#endif
// On failure the reply stays all zero and ~UdpLink removes the pcb
if (link.pcb == nullptr || udp_bind(link.pcb, IP_ANY_TYPE, this->port_) != ERR_OK)
return;
// Below the error codes (0x80 and up) the client checks the first byte against, and never all zero
token[0] = (token[0] & 0x7F) | 0x01;
memcpy(link.reply, token, sizeof(token));
udp_recv(link.pcb, udp_recv_cb_, &link);
#endif
}
bool ESPHomeOTAComponent::udp_start_() {
// TCP now only carries a fallback byte; its blocking reads time out every couple of seconds to look again
const uint32_t start = millis();
while (millis() - start < UDP_COMMIT_TIMEOUT) {
bool committed;
{
UdpLock lock;
committed = this->udp_->committed;
}
if (committed) {
// The client heard us over UDP, so TCP is done. Outside the lock: closing a socket waits on the tcpip thread
this->client_ = nullptr;
return true;
}
// A raw byte outside any Noise frame: the client heard nothing over UDP. Forging it only forces the ordinary
// TCP path, which dropping our datagrams forces anyway
uint8_t fallback;
const ssize_t read = this->client_->read(&fallback, 1);
if (read == 1) {
this->udp_->release();
this->udp_ = nullptr;
return true;
}
if (read == 0 || !this->would_block_(errno))
return false;
App.feed_wdt();
}
return false;
}
ssize_t ESPHomeOTAComponent::udp_read_(uint8_t *buf, size_t len) {
UdpLink &link = *this->udp_;
const uint32_t start = millis();
size_t got = 0;
for (;;) {
struct pbuf *p;
{
UdpLock lock;
p = link.consumed != link.received ? link.queue[link.consumed % UDP_WINDOW] : nullptr;
}
if (p == nullptr) {
if (got != 0)
return got;
if (millis() - start >= UDP_READ_WAIT_MS)
break;
this->yield_and_feed_watchdog_();
continue;
}
// The callback never touches the slot being read, so the copy runs without the lock
const uint16_t offset = UDP_HEADER_SIZE + link.read_offset;
const uint16_t n = pbuf_copy_partial(p, buf + got, std::min<size_t>(len - got, p->tot_len - offset), offset);
got += n;
link.read_offset += n;
if (UDP_HEADER_SIZE + link.read_offset >= p->tot_len) {
UdpLock lock;
pbuf_free(p);
link.queue[link.consumed % UDP_WINDOW] = nullptr;
// A full window just opened: tell the client, which is holding back
const bool was_full = static_cast<uint16_t>(link.received - link.consumed) == UDP_WINDOW;
link.consumed++;
link.read_offset = 0;
if (was_full)
udp_send_ack_(link, link.received - 1);
}
if (got == len)
return got;
}
errno = EWOULDBLOCK;
return -1;
}
bool ESPHomeOTAComponent::udp_send_(const uint8_t *data, size_t len) {
UdpLink &link = *this->udp_;
UdpLock lock;
if (link.out_len + len > UDP_OUT_SIZE)
return false;
memcpy(link.out + link.out_len, data, len);
link.out_len += len;
// Prompted by nothing new, so the client takes no round trip sample from it
udp_send_ack_(link, link.received - 1);
return true;
}
void ESPHomeOTAComponent::udp_linger_() {
if (this->udp_ == nullptr)
return;
const uint32_t start = millis();
while (millis() - start < UDP_LINGER_MS) {
{
UdpLock lock;
if (this->udp_->out_len == 0)
return;
}
this->yield_and_feed_watchdog_();
}
}
} // namespace esphome
#endif // USE_OTA && ESPHOME_OTA_UDP
esphome/espota2.py
+327 -5
View File
@@ -66,11 +66,13 @@ CLIENT_FEATURE_SUPPORTS_SHA256_AUTH = 0x02
CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL = 0x04
CLIENT_FEATURE_SUPPORTS_NOISE = 0x08
CLIENT_FEATURE_SUPPORTS_DEFLATE = 0x10
CLIENT_FEATURE_SUPPORTS_UDP = 0x20
SERVER_FEATURE_SUPPORTS_COMPRESSION = 0x01
SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS = 0x02
SERVER_FEATURE_SUPPORTS_NOISE = 0x04
# Binding once offered: the device then expects the image size and a deflate stream
SERVER_FEATURE_SUPPORTS_DEFLATE = 0x08
SERVER_FEATURE_SUPPORTS_UDP = 0x10
# Wire constant: the deflate bit promises a 4 KB window (OTA_INFLATE_WINDOW_SIZE)
DEFLATE_WINDOW_BITS = 12
@@ -99,17 +101,55 @@ COMPRESS_LEVEL = 9
UPLOAD_BUFFER_SIZE = UPLOAD_BLOCK_SIZE * 8
# Flaky Wi-Fi links often drop the first OTA attempt, and the device may need time
# to clean up a half-open connection (its handshake watchdog runs at 20s) before it
# to clean up a half-open connection (its handshake watchdog runs at 60s) before it
# accepts a new one, so wait between attempts instead of failing the upload outright.
# Every resolved address is tried once, and this many extra attempts are shared
# across the addresses on top of that.
EXTRA_UPLOAD_ATTEMPTS = 2
EXTRA_UPLOAD_ATTEMPTS = 4
# Once a connection has reached the device, a lossy link (not a dead device) is
# the likely cause, so keep trying like an API reconnect would
EXTRA_UPLOAD_ATTEMPTS_REACHED = 8
UPLOAD_RETRY_DELAY = 5.0
# Connect and handshake waits; as patient as an API connect on a lossy link, and
# matching the device's OTA_SOCKET_TIMEOUT_HANDSHAKE
SETUP_TIMEOUT = 60.0
# Data phase timeout; must stay longer than the device's OTA_SOCKET_TIMEOUT_DATA
# (105 s) so a stalled session is gone before a retry, and long enough for lwIP
# to get a lost chunk ack through after the retransmit run seen in practice
DATA_PHASE_TIMEOUT = 160.0
# UDP data phase; wire format in ota_esphome_udp.cpp
UDP_MSG_DATA = 0x01
UDP_MSG_PING = 0x02
UDP_MSG_PROBE = 0x03
UDP_FLAG_COMMITTED = 0x80
UDP_MSG_ACK = 0x40
UDP_HEADER_SIZE = 9
UDP_ACK_HEADER_SIZE = 12
# Stream bytes per datagram, within one Ethernet frame
UDP_MAX_PAYLOAD = 1400
# Messages the device queues unconsumed (UDP_WINDOW in ota_esphome.h); as many
# as TCP's window holds, so UDP costs the device no more memory than TCP
UDP_WINDOW = 4
# Probe this long before falling back to TCP; the device waits 70 s
UDP_PROBE_TIMEOUT = 5.0
UDP_PROBE_INTERVAL = 0.1
# Resend timeout bounds around srtt + 4 * rttvar; loss never grows it
UDP_RESEND_INITIAL = 0.5
UDP_RESEND_MIN = 0.2
UDP_RESEND_MAX = 2.0
# Acks still asking for the head after later messages arrived; resend it once
UDP_DUP_ACKS = 2
UDP_FALLBACK = b"\x00"
UDP_REQUEST = b"\x01"
# A plain handshake exchange slower than this hid a retransmit (the device's
# lwIP waits about a second before resending); UDP is used from then on
UDP_SLOW_EXCHANGE = 0.75
# getsockopt(IPPROTO_TCP) option, size, offset and format of the retransmit
# count: tcp_info.tcpi_total_retrans (Linux), tcp_connection_info.
# tcpi_txretransmitbytes (macOS)
_TCP_RETRANS_OPTION = {"linux": (11, 104, 100, "I"), "darwin": (0x106, 80, 72, "Q")}
_LOGGER = logging.getLogger(__name__)
# Authentication method lookup table: response -> (hash_func, nonce_size, name)
@@ -434,6 +474,10 @@ class NoiseSocketWrapper:
self._decrypt = None
self._buffer = b""
def use_transport(self, transport: UdpChannel) -> None:
"""Continue the session over another byte stream."""
self._sock = transport
# Only harmless socket controls pass through; byte-moving methods are
# deliberately absent so plaintext cannot leak past the transport.
def settimeout(self, timeout: float | None) -> None:
@@ -534,6 +578,252 @@ class NoiseSocketWrapper:
return data
class _UdpMessage:
"""One datagram of the UDP stream and its send history."""
__slots__ = ("held", "last_send", "payload", "resend", "sends", "seq")
def __init__(self, seq: int, payload: bytes) -> None:
self.seq = seq
self.payload = payload
self.last_send = 0.0
self.sends = 0
self.resend = False # duplicate acks asked for it
self.held = False # the device keeps it, waiting for an earlier one
class UdpChannel:
"""Byte stream over numbered UDP datagrams, standing in for the socket.
The device acks every datagram on arrival and says how many it has room
for; each is resent after TCP's srtt + 4 * rttvar, without backoff, so loss
slows the transfer instead of ending it.
"""
def __init__(self, sock: socket.socket, token: bytes) -> None:
self._sock = sock
self._token = token
self._timeout = DATA_PHASE_TIMEOUT
self._next_seq = 0
self._queue: list[_UdpMessage] = [] # not yet received, oldest first
self._rx = bytearray()
self._device_bytes = 0
self._committed = False
self._srtt: float | None = None
self._rttvar = 0.0
self._dup_acks = 0
self._free = UDP_WINDOW # the device's empty slots, as of its last ACK
self._last_ping = 0.0
self._last_progress = time.monotonic()
self._sends = 0
def settimeout(self, timeout: float) -> None:
self._timeout = timeout
def setsockopt(self, level: int, optname: int, value: int) -> None:
"""TCP tuning does not apply."""
def close(self) -> None:
_LOGGER.debug(
"UDP: %d messages in %d datagrams, smoothed RTT %.0f ms",
self._next_seq,
self._sends,
(self._srtt or 0) * 1000,
)
self._sock.close()
def probe(self) -> bool:
"""True once the device answers over UDP within UDP_PROBE_TIMEOUT."""
deadline = time.monotonic() + UDP_PROBE_TIMEOUT
while (now := time.monotonic()) < deadline:
self._send(UDP_MSG_PROBE, 0, b"")
self._receive(min(UDP_PROBE_INTERVAL, deadline - now))
if self._committed:
return True
return False
def sendall(self, data: bytes) -> None:
for offset in range(0, len(data), UDP_MAX_PAYLOAD):
self._queue.append(
_UdpMessage(self._next_seq, data[offset : offset + UDP_MAX_PAYLOAD])
)
self._next_seq = (self._next_seq + 1) & 0xFFFF
# Return once everything is at least in flight
while len(self._queue) > UDP_WINDOW:
self._pump()
def recv(self, amount: int) -> bytes:
while not self._rx:
self._pump()
data = bytes(self._rx[:amount])
del self._rx[:amount]
return data
def flush(self, timeout: float) -> None:
"""Best effort: wait until the device has every message."""
deadline = time.monotonic() + timeout
with contextlib.suppress(OTANetworkError):
while self._queue and time.monotonic() < deadline:
self._pump()
def _rto(self) -> float:
if self._srtt is None:
return UDP_RESEND_INITIAL
return min(max(self._srtt + 4 * self._rttvar, UDP_RESEND_MIN), UDP_RESEND_MAX)
def _send(self, msg_type: int, seq: int, payload: bytes) -> None:
flags = UDP_FLAG_COMMITTED if self._committed else 0
header = (
bytes([msg_type | flags])
+ self._token
+ seq.to_bytes(2, "big")
+ (self._device_bytes & 0xFFFF).to_bytes(2, "big")
)
self._sends += 1
with contextlib.suppress(OSError):
# A full buffer or ICMP error is just more loss
self._sock.send(header + payload)
def _pump(self) -> None:
now = time.monotonic()
if now - self._last_progress > self._timeout:
raise OTANetworkError(f"no progress over UDP for {self._timeout:.0f}s")
rto = self._rto()
wake = now + rto
# At least the head, which doubles as the probe while the device is busy
in_flight = self._queue[: max(self._free, 1)]
for message in in_flight:
if message.held:
continue
if message.sends == 0 or message.resend or now - message.last_send >= rto:
message.last_send = now
message.sends += 1
message.resend = False
self._send(UDP_MSG_DATA, message.seq, message.payload)
wake = min(wake, message.last_send + rto)
if all(message.held for message in in_flight):
# Nothing to send: ask for the device's responses
if now - self._last_ping >= rto:
self._last_ping = now
self._send(UDP_MSG_PING, 0, b"")
wake = self._last_ping + rto
self._receive(wake - now)
def _receive(self, wait: float) -> None:
import select
# Wait for the first datagram, then drain whatever else arrived
while select.select([self._sock], [], [], max(wait, 0))[0]:
wait = 0
try:
data = self._sock.recv(2048)
except OSError:
continue
self._handle_ack(data)
def _sample_rtt(self, sample: float) -> None:
# RFC 6298 smoothing
if self._srtt is None:
self._srtt, self._rttvar = sample, sample / 2
else:
self._rttvar = 0.75 * self._rttvar + 0.25 * abs(self._srtt - sample)
self._srtt = 0.875 * self._srtt + 0.125 * sample
def _handle_ack(self, data: bytes) -> None:
if (
len(data) < UDP_ACK_HEADER_SIZE
or data[0] != UDP_MSG_ACK
or data[1:5] != self._token
):
return
now = time.monotonic()
self._committed = True
received = int.from_bytes(data[5:7], "big")
prompted = int.from_bytes(data[7:9], "big")
self._free = data[11] >> 4
if self._queue:
head = self._queue[0].seq
index = (prompted - head) & 0xFFFF
# This datagram's own round trip; only from messages sent once (Karn)
if index < len(self._queue) and self._queue[index].sends == 1:
self._sample_rtt(now - self._queue[index].last_send)
done = (received - head) & 0xFFFF
if 0 < done <= len(self._queue):
del self._queue[:done]
self._last_progress = now
self._dup_acks = 0
elif done == 0 and self._queue[0].sends:
# A later message arrived while the head is missing; resend it once
self._dup_acks += 1
if self._dup_acks == UDP_DUP_ACKS:
self._queue[0].resend = True
# Rebuilt from every ACK, so a reordered one cannot leave a stale mark
for k, message in enumerate(self._queue[:UDP_WINDOW]):
message.held = (received + k - message.seq) & 0xFFFF == 0 and bool(
data[11] & (1 << k)
)
skip = (self._device_bytes - int.from_bytes(data[9:11], "big")) & 0xFFFF
fresh = data[UDP_ACK_HEADER_SIZE + skip :]
if skip <= len(data) - UDP_ACK_HEADER_SIZE and fresh:
self._rx += fresh
self._device_bytes += len(fresh)
self._last_progress = now
def _tcp_retransmits(sock: socket.socket) -> int | None:
"""Segments this side retransmitted, where the OS reports it."""
import struct
import sys
if (option := _TCP_RETRANS_OPTION.get(sys.platform)) is None:
return None
optname, size, offset, fmt = option
try:
info = sock.getsockopt(socket.IPPROTO_TCP, optname, size)
except (OSError, TypeError):
return None
if len(info) < offset + struct.calcsize(fmt):
return None
return struct.unpack_from(fmt, info, offset)[0]
def _link_lossy(
tcp_sock: socket.socket, slowest_exchange: float, prefer_udp: bool
) -> bool:
"""Whether the handshake showed loss, so the data phase should use UDP."""
return (
prefer_udp
or slowest_exchange > UDP_SLOW_EXCHANGE
or bool(_tcp_retransmits(tcp_sock))
)
def _start_udp(
sock: socket.socket | NoiseSocketWrapper, tcp_sock: socket.socket, lossy: bool
) -> UdpChannel | None:
"""Move the data phase to UDP; None keeps it on TCP."""
# A clean link asks to stay on TCP
send_check(sock, UDP_REQUEST if lossy else UDP_FALLBACK, "UDP request")
token = receive_exactly(sock, 4, "UDP token", None, decode=False)
if not lossy:
_LOGGER.debug("Link looks clean, staying on TCP")
return None
if token == bytes(4):
_LOGGER.debug("Device declined UDP")
return None
udp_sock = socket.socket(tcp_sock.family, socket.SOCK_DGRAM)
udp_sock.connect(tcp_sock.getpeername())
channel = UdpChannel(udp_sock, token)
if channel.probe():
_LOGGER.info("Using UDP")
return channel
channel.close()
_LOGGER.info("UDP blocked, using TCP")
# Raw, outside any Noise frame: forging it only forces the TCP path
send_check(tcp_sock, UDP_FALLBACK, "UDP fallback")
return None
def perform_ota(
sock: socket.socket,
password: str | None,
@@ -543,6 +833,7 @@ def perform_ota(
noise_psk: str | None = None,
plaintext_fallback: bool = False,
allow_plaintext_upload: bool = False,
prefer_udp: bool = False,
) -> None:
# Validate up front; an out-of-range value would only surface as a
# ValueError deep inside send_check, bypassing OTAError handling
@@ -565,11 +856,15 @@ def perform_ota(
file_size = len(file_contents)
_LOGGER.info("Uploading %s (%s bytes)", filename, file_size)
# sock may become a Noise wrapper or a UDP channel; UDP needs the peer
tcp_sock = sock
# Enable nodelay, we need it for phase 1
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
exchange_start = time.monotonic()
send_check(sock, MAGIC_BYTES, "magic bytes")
_, version = receive_exactly(sock, 2, "version", RESPONSE_OK)
slowest_exchange = time.monotonic() - exchange_start
_LOGGER.info("Connection established; device supports OTA version %s", version)
supported_versions = (OTA_VERSION_1_0, OTA_VERSION_2_0)
if version not in supported_versions:
@@ -583,9 +878,11 @@ def perform_ota(
| CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
| CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
| CLIENT_FEATURE_SUPPORTS_DEFLATE
| CLIENT_FEATURE_SUPPORTS_UDP
)
if noise_psk:
features_to_send |= CLIENT_FEATURE_SUPPORTS_NOISE
exchange_start = time.monotonic()
send_check(sock, features_to_send, "features")
features = receive_exactly(
sock,
@@ -593,6 +890,7 @@ def perform_ota(
"features",
None, # Accept any response
)[0]
slowest_exchange = max(slowest_exchange, time.monotonic() - exchange_start)
extended_proto = False
if features == RESPONSE_FEATURE_FLAGS:
@@ -751,6 +1049,18 @@ def perform_ota(
sock.settimeout(DATA_PHASE_TIMEOUT)
udp: UdpChannel | None = None
if extended_proto and features & SERVER_FEATURE_SUPPORTS_UDP:
udp = _start_udp(
sock, tcp_sock, _link_lossy(tcp_sock, slowest_exchange, prefer_udp)
)
if udp is not None:
# The device drops TCP once it hears the committed channel
if isinstance(sock, NoiseSocketWrapper):
sock.use_transport(udp)
else:
sock = udp
if extended_proto:
send_check(sock, ota_type, "ota type")
@@ -806,7 +1116,8 @@ def perform_ota(
)
raise OTANetworkError(f"sending data: {err}") from err
if version >= OTA_VERSION_2_0:
# The UDP channel acks every message itself
if version >= OTA_VERSION_2_0 and udp is None:
try:
receive_exactly(sock, 1, "chunk result", RESPONSE_CHUNK_OK)
except OTANetworkError as err:
@@ -860,6 +1171,10 @@ def perform_ota(
_LOGGER.info("OTA successful (end acknowledgement not delivered)")
else:
_LOGGER.info("OTA successful")
if udp is not None:
# Nothing waits for this; just give the end acknowledgement a chance
udp.flush(2.0)
udp.close()
# Do not connect logs until it is fully on
time.sleep(1)
@@ -904,7 +1219,7 @@ def run_ota_impl_(
# are shared across the addresses, cycling through them. Wait before an
# attempt when the previous one actually reached the device, or when
# revisiting an address, so a flaky link can recover and the device can
# clean up a half-open connection (its handshake watchdog runs at 20s);
# clean up a half-open connection (its handshake watchdog runs at 60s);
# moving on to the next address family stays immediate. Known limitation:
# a silent mid-transfer drop with no reset can wedge the device until its
# 105s data timeout, which outlasts this budget; the retries target the
@@ -914,6 +1229,8 @@ def run_ota_impl_(
reached_device = False
attempt = 0
encryption = _EncryptionAttempt(noise_psk, plaintext_fallback)
# A network failure means a lossy link, so later attempts send data over UDP
prefer_udp = False
while attempt < total_attempts:
af, socktype, _, _, sa = res[attempt % len(res)]
if reached_device or attempt >= len(res):
@@ -927,18 +1244,20 @@ def run_ota_impl_(
reached_device = False
_LOGGER.info("Connecting to %s port %s...", sa[0], sa[1])
sock = socket.socket(af, socktype)
sock.settimeout(20.0)
sock.settimeout(SETUP_TIMEOUT)
try:
sock.connect(sa)
except OSError as err:
sock.close()
_LOGGER.warning("Connecting to %s port %s failed: %s", sa[0], sa[1], err)
last_error = f"connecting to {sa[0]} failed: {err}"
prefer_udp = True
attempt += 1
continue
_LOGGER.info("Connected to %s", sa[0])
reached_device = True
total_attempts = max(total_attempts, len(res) + EXTRA_UPLOAD_ATTEMPTS_REACHED)
with contextlib.closing(sock), Path(filename).open("rb") as file_handle:
try:
perform_ota(
@@ -950,6 +1269,7 @@ def run_ota_impl_(
encryption.noise_psk,
encryption.plaintext_fallback,
allow_plaintext_upload=allow_plaintext_upload,
prefer_udp=prefer_udp,
)
except OTAEncryptionFallback as err:
# Same address and attempt budget: not a network retry
@@ -958,6 +1278,7 @@ def run_ota_impl_(
continue
except OTAHandshakeNetworkError as err:
last_error = str(err)
prefer_udp = True
if encryption.handshake_fault_falls_back():
encryption.downgrade(last_error)
continue
@@ -967,6 +1288,7 @@ def run_ota_impl_(
except OTANetworkError as err:
# Transient network failure; retry
last_error = str(err)
prefer_udp = True
_LOGGER.warning("%s", last_error)
attempt += 1
continue
tests/unit_tests/test_espota2.py
+197 -8
View File
@@ -11,6 +11,7 @@ import logging
from pathlib import Path
import socket
import struct
from typing import Self
from unittest.mock import Mock, call, patch
import zlib
@@ -356,6 +357,7 @@ def test_perform_ota_successful_md5_auth(
| espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
| espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
| espota2.CLIENT_FEATURE_SUPPORTS_DEFLATE
| espota2.CLIENT_FEATURE_SUPPORTS_UDP
]
)
)
@@ -764,7 +766,7 @@ def test_run_ota_impl_successful(
assert result_host == "192.168.1.100"
# Verify socket was configured correctly
mock_socket.settimeout.assert_called_with(20.0)
mock_socket.settimeout.assert_called_with(espota2.SETUP_TIMEOUT)
mock_socket.connect.assert_called_once_with(("192.168.1.100", 3232))
mock_socket.close.assert_called_once()
@@ -835,13 +837,22 @@ def test_run_ota_impl_network_error_retry_succeeds(
assert result_host == "192.168.1.100"
assert mock_perform_ota.call_count == 2
mock_sleep.assert_called_once_with(espota2.UPLOAD_RETRY_DELAY)
# The failure marks the link lossy, so the retry moves the data to UDP
assert [c.kwargs["prefer_udp"] for c in mock_perform_ota.call_args_list] == [
False,
True,
]
@pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip")
def test_run_ota_impl_network_error_exhausts_attempts(
mock_socket: Mock, firmware_file: Path, mock_perform_ota: Mock, mock_sleep: Mock
) -> None:
"""Test run_ota_impl_ gives up after all attempts hit network errors."""
"""Test run_ota_impl_ gives up after all attempts hit network errors.
Every attempt connected, so the budget is the longer one for a reachable
device on a lossy link.
"""
mock_perform_ota.side_effect = espota2.OTANetworkError("sending data: broken pipe")
result_code, result_host = espota2.run_ota_impl_(
@@ -850,8 +861,8 @@ def test_run_ota_impl_network_error_exhausts_attempts(
assert result_code == 1
assert result_host is None
assert mock_perform_ota.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS + 1
assert mock_sleep.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS
assert mock_perform_ota.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS_REACHED + 1
assert mock_sleep.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS_REACHED
@pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip_dual")
@@ -872,12 +883,10 @@ def test_run_ota_impl_multiple_addresses_cycle(
assert mock_socket.connect.call_args_list == [
call(DUAL_STACK_SA6),
call(DUAL_STACK_SA4),
call(DUAL_STACK_SA6),
call(DUAL_STACK_SA4),
]
] * (1 + espota2.EXTRA_UPLOAD_ATTEMPTS // 2)
# No connect ever reached the device, so the delay only applies before
# the revisits
assert mock_sleep.call_count == 2
assert mock_sleep.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS
@pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip_dual")
@@ -1061,6 +1070,7 @@ def test_perform_ota_successful_sha256_auth(
| espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
| espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
| espota2.CLIENT_FEATURE_SUPPORTS_DEFLATE
| espota2.CLIENT_FEATURE_SUPPORTS_UDP
]
)
)
@@ -1118,6 +1128,7 @@ def test_perform_ota_sha256_fallback_to_md5(
| espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
| espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
| espota2.CLIENT_FEATURE_SUPPORTS_DEFLATE
| espota2.CLIENT_FEATURE_SUPPORTS_UDP
]
)
)
@@ -1228,6 +1239,7 @@ def test_perform_ota_extended_protocol_app(
| espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
| espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
| espota2.CLIENT_FEATURE_SUPPORTS_DEFLATE
| espota2.CLIENT_FEATURE_SUPPORTS_UDP
]
)
)
@@ -1289,6 +1301,7 @@ def test_perform_ota_successful_partition_table(
| espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
| espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
| espota2.CLIENT_FEATURE_SUPPORTS_DEFLATE
| espota2.CLIENT_FEATURE_SUPPORTS_UDP
]
)
)
@@ -1554,3 +1567,179 @@ def test_perform_ota_with_deflate(mock_socket: Mock, server_features: int) -> No
assert len(payload) == sent_size < len(original_content)
assert zlib.decompress(payload, -espota2.DEFLATE_WINDOW_BITS) == original_content
assert sent[5] == hashlib.md5(original_content).hexdigest().encode()
class _LossyUdpDevice:
"""The device side of the UDP channel (ota_esphome_udp.cpp) on a lossy link."""
def __init__(self, token: bytes, loss: float, seed: int = 1) -> None:
import random
import threading
self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
self.sock.bind(("127.0.0.1", 0))
self.sock.settimeout(0.05)
self.token = token
self.loss = loss
self.random = random.Random(seed)
self.next_seq = 0
self.slots: dict[int, bytes] = {} # early messages within the window
self.received: list[bytes] = []
self.out = bytearray()
self.out_base = 0
self.committed = False
self._stop = threading.Event()
self._thread = threading.Thread(target=self._run, daemon=True)
def __enter__(self) -> Self:
self._thread.start()
return self
def __exit__(self, *exc: object) -> None:
self._stop.set()
self._thread.join()
self.sock.close()
def _run(self) -> None:
while not self._stop.is_set():
try:
data, peer = self.sock.recvfrom(2048)
except TimeoutError:
continue
if self.random.random() < self.loss:
continue
if data[1:5] != self.token:
continue
self.committed |= bool(data[0] & espota2.UDP_FLAG_COMMITTED)
acked = (int.from_bytes(data[7:9], "big") - self.out_base) & 0xFFFF
if acked <= len(self.out):
del self.out[:acked]
self.out_base += acked
seq = int.from_bytes(data[5:7], "big")
msg_type = data[0] & ~espota2.UDP_FLAG_COMMITTED
if (
msg_type == espota2.UDP_MSG_DATA
and (seq - self.next_seq) & 0xFFFF < espota2.UDP_WINDOW
):
self.slots.setdefault(seq, data[espota2.UDP_HEADER_SIZE :])
# Consumed as soon as it is next, like a device that is never busy
while self.next_seq in self.slots:
self.received.append(self.slots.pop(self.next_seq))
self.next_seq = (self.next_seq + 1) & 0xFFFF
if self.random.random() < self.loss:
continue
held = sum(
1 << k
for k in range(1, espota2.UDP_WINDOW)
if (self.next_seq + k) & 0xFFFF in self.slots
)
# Consumes at once, so every slot it does not hold is free
held |= (espota2.UDP_WINDOW - len(self.slots)) << 4
self.sock.sendto(
bytes([espota2.UDP_MSG_ACK])
+ self.token
+ self.next_seq.to_bytes(2, "big")
+ seq.to_bytes(2, "big")
+ (self.out_base & 0xFFFF).to_bytes(2, "big")
+ bytes([held])
+ bytes(self.out),
peer,
)
def _udp_client(device: _LossyUdpDevice) -> espota2.UdpChannel:
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.connect(device.sock.getsockname())
return espota2.UdpChannel(sock, device.token)
@pytest.mark.parametrize("loss", [0.0, 0.5])
def test_udp_channel_delivers_in_order(loss: float) -> None:
"""The stream arrives intact and responses come back, despite loss."""
data = bytes(range(256)) * 160
with _LossyUdpDevice(b"\x01\x02\x03\x04", loss) as device:
channel = _udp_client(device)
assert channel.probe()
channel.sendall(data)
device.out += bytes([espota2.RESPONSE_RECEIVE_OK])
assert channel.recv(1) == bytes([espota2.RESPONSE_RECEIVE_OK])
channel.flush(30.0)
channel.close()
assert b"".join(device.received) == data
assert all(len(m) <= espota2.UDP_MAX_PAYLOAD for m in device.received)
assert device.committed
def test_udp_channel_times_out_without_progress() -> None:
"""A dead link still fails, after the no-progress timeout."""
with _LossyUdpDevice(b"\x01\x01\x01\x01", 1.0) as device:
channel = _udp_client(device)
channel.settimeout(0.3)
with pytest.raises(espota2.OTANetworkError, match="no progress over UDP"):
channel.sendall(bytes(espota2.UDP_MAX_PAYLOAD * 5))
channel.close()
def test_start_udp_falls_back_when_blocked(mock_socket: Mock) -> None:
"""No answer to the probes sends the fallback over TCP and stays on TCP."""
with _LossyUdpDevice(b"\x05\x06\x07\x08", 1.0) as device:
mock_socket.family = socket.AF_INET
mock_socket.getpeername.return_value = device.sock.getsockname()
mock_socket.recv.side_effect = [b"\x05", b"\x06\x07\x08"]
with patch.object(espota2, "UDP_PROBE_TIMEOUT", 0.3):
assert espota2._start_udp(mock_socket, mock_socket, True) is None
assert mock_socket.sendall.call_args_list == [
call(espota2.UDP_REQUEST),
call(espota2.UDP_FALLBACK), # raw on the TCP socket
]
def test_start_udp_declined(mock_socket: Mock) -> None:
"""A zero token means the device could not open UDP; nothing is probed."""
mock_socket.recv.side_effect = [b"\x00", b"\x00\x00\x00"]
assert espota2._start_udp(mock_socket, mock_socket, True) is None
assert mock_socket.sendall.call_args_list == [call(espota2.UDP_REQUEST)]
def test_start_udp_clean_link_stays_on_tcp(mock_socket: Mock) -> None:
"""A clean handshake asks to stay on TCP and opens no UDP socket."""
mock_socket.recv.side_effect = [b"\x00", b"\x00\x00\x00"]
with patch.object(espota2.socket, "socket") as udp_socket:
assert espota2._start_udp(mock_socket, mock_socket, False) is None
udp_socket.assert_not_called()
assert mock_socket.sendall.call_args_list == [call(espota2.UDP_FALLBACK)]
@pytest.mark.parametrize(
("slowest", "prefer_udp", "retransmits", "lossy"),
[
(0.01, False, 0, False),
(0.01, False, None, False),
(espota2.UDP_SLOW_EXCHANGE + 0.1, False, 0, True),
(0.01, True, 0, True),
(0.01, False, 2, True),
],
)
def test_link_lossy(
slowest: float, prefer_udp: bool, retransmits: int | None, lossy: bool
) -> None:
"""A slow exchange, a retransmit or an earlier failure means UDP."""
with patch.object(espota2, "_tcp_retransmits", return_value=retransmits):
assert espota2._link_lossy(Mock(), slowest, prefer_udp) is lossy
@pytest.mark.parametrize(
("platform", "info", "expected"),
[
("linux", bytes(100) + (3).to_bytes(4, "little") + bytes(100), 3),
("darwin", bytes(72) + (1460).to_bytes(8, "little") + bytes(8), 1460),
("win32", b"", None),
("linux", bytes(10), None),
],
)
def test_tcp_retransmits(platform: str, info: bytes, expected: int | None) -> None:
"""The retransmit counter is read where the OS reports it."""
sock = Mock()
sock.getsockopt.return_value = info
with patch("sys.platform", platform):
assert espota2._tcp_retransmits(sock) == expected