Fail fast when the device's UDP port refuses, include <algorithm>

esphome/components/esphome/ota/ota_esphome_udp.cpp

@@ -13,6 +13,7 @@
 #endif
 #endif
 
+#include <algorithm>
 #include <cstring>
 #include <memory>
 

esphome/espota2.py

@@ -142,6 +142,8 @@ 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
+# 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"
@@ -621,6 +623,7 @@ class UdpChannel:
         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
@@ -640,11 +643,14 @@ class UdpChannel:
     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
+        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:
@@ -740,6 +746,10 @@ class UdpChannel:
         if self._last_error is None or err.errno != self._last_error.errno:
             _LOGGER.debug("UDP socket error: %s", err)
         self._last_error = err
+        if isinstance(err, ConnectionRefusedError):
+            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
@@ -758,6 +768,7 @@ class UdpChannel:
             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

tests/unit_tests/test_espota2.py

@@ -1736,8 +1736,7 @@ def test_tcp_retransmits(platform: str, info: bytes, expected: int | None) -> No
 
 
 def test_udp_channel_reports_socket_errors() -> None:
-    """A dead link still fails after the no-progress timeout, naming the socket
-    error behind it."""
+    """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()
@@ -1745,14 +1744,38 @@ def test_udp_channel_reports_socket_errors() -> None:
     sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
     sock.connect(addr)
     channel = espota2.UdpChannel(sock, b"\x01\x01\x01\x01")
-    channel.settimeout(0.5)
-    with pytest.raises(
-        espota2.OTANetworkError, match="no progress over UDP.*last error"
-    ):
+    channel.settimeout(60.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)
+
+
 def _udp_handshake() -> list[bytes]:
     """TCP recv responses up to auth, from a device that offers UDP."""
     return [