Require ESPHome ASH auto-ACKing when using a serial proxy

bellows/ash.py

@@ -353,7 +353,7 @@ def parse_frame(
 
 
 class AshProtocol(asyncio.Protocol):
-    def __init__(self, ezsp_protocol) -> None:
+    def __init__(self, ezsp_protocol, *, suppress_acks: bool = False) -> None:
         self._ezsp_protocol = ezsp_protocol
         self._transport = None
         self._buffer = bytearray()
@@ -363,11 +363,17 @@ def __init__(self, ezsp_protocol) -> None:
         self._tx_seq: int = 0
         self._rx_seq: int = 0
         self._t_rx_ack = T_RX_ACK_INIT
+        self._suppress_acks = suppress_acks
 
         self._ncp_reset_code: t.NcpResetCode | None = None
         self._ncp_state: NcpState = NcpState.CONNECTED
         self._in_reject_condition: bool = False
 
+    @property
+    def suppress_acks(self) -> bool:
+        """Whether something upstream acknowledges the NCP's frames on our behalf."""
+        return self._suppress_acks
+
     def connection_made(self, transport):
         self._transport = transport
         self._ezsp_protocol.connection_made(self)
@@ -555,12 +561,20 @@ def frame_received(self, frame: AshFrame) -> None:
         else:
             raise TypeError(f"Unknown frame received: {frame}")  # pragma: no cover
 
+    def _send_ack(self) -> None:
+        # Something upstream may be ACKing on our behalf, in which case our own ACKs are
+        # pure overhead. NAKs are still sent: see `_reject_frame`.
+        if self._suppress_acks:
+            return
+
+        self._write_frame(AckFrame(res=0, ncp_ready=0, ack_num=self._rx_seq))
+
     def data_frame_received(self, frame: DataFrame) -> None:
         # The Host may not piggyback acknowledgments and should promptly send an ACK
         # frame when it receives a DATA frame.
         if frame.frm_num == self._rx_seq:
             self._rx_seq = (frame.frm_num + 1) % 8
-            self._write_frame(AckFrame(res=0, ncp_ready=0, ack_num=self._rx_seq))
+            self._send_ack()
 
             # Clear reject condition on valid in-sequence DATA frame
             self._clear_reject_condition()
@@ -569,7 +583,7 @@ def data_frame_received(self, frame: DataFrame) -> None:
         elif frame.re_tx:
             # Retransmitted frames must be immediately ACKed even if they are out of
             # sequence
-            self._write_frame(AckFrame(res=0, ncp_ready=0, ack_num=self._rx_seq))
+            self._send_ack()
         else:
             # Out-of-sequence non-retransmitted frame: reject it
             _LOGGER.debug("Received out-of-sequence frame: %r", frame)

bellows/uart.py

@@ -1,6 +1,7 @@
 import asyncio
 from asyncio import timeout as asyncio_timeout
 import logging
+import urllib.parse
 
 import zigpy.config
 import zigpy.serial
@@ -12,6 +13,19 @@
 LOGGER = logging.getLogger(__name__)
 RESET_TIMEOUT = 2.5
 
+# An ESPHome port in this mode engages its tap, which ACKs frames on our behalf. The
+# mode does not name a protocol: which one the tap speaks is the device's configuration.
+ACKLESS_MODE = "protocol"
+
+# URL schemes served by an ESPHome serial proxy, which is the only transport that can take
+# the acknowledging off our hands
+ESPHOME_SCHEMES = frozenset({"esphome", "esphome-hass"})
+
+
+def is_esphome_url(path: str) -> bool:
+    """Whether the path is served by an ESPHome serial proxy."""
+    return urllib.parse.urlparse(path).scheme in ESPHOME_SCHEMES
+
 
 class Gateway(zigpy.serial.SerialProtocol):
     def __init__(self, api, connection_done_future=None):
@@ -110,20 +124,35 @@ async def _connect(config, api):
 
     connection_done_future = loop.create_future()
 
+    path = config[zigpy.config.CONF_DEVICE_PATH]
+
     gateway = Gateway(api, connection_done_future)
-    protocol = AshProtocol(gateway)
+    # serialx opens an ESPHome port in the mode we ask for or raises, so whether the tap
+    # acknowledges for us is known before the first frame arrives
+    protocol = AshProtocol(gateway, suppress_acks=is_esphome_url(path))
 
     flow_control = config[zigpy.config.CONF_DEVICE_FLOW_CONTROL]
 
+    extra_kwargs = {}
+
+    if is_esphome_url(path):
+        extra_kwargs["mode"] = ACKLESS_MODE
+
     transport, _ = await zigpy.serial.create_serial_connection(
         loop,
         lambda: protocol,
-        url=config[zigpy.config.CONF_DEVICE_PATH],
+        url=path,
         baudrate=config[zigpy.config.CONF_DEVICE_BAUDRATE],
         xonxoff=flow_control in (None, "software"),
         rtscts=flow_control == "hardware",
+        **extra_kwargs,
     )
 
+    if is_esphome_url(path):
+        LOGGER.debug(
+            "Port acknowledges NCP frames on our behalf, suppressing our own ACKs"
+        )
+
     await gateway.wait_until_connected()
 
     thread_safe_protocol = ThreadsafeProxy(gateway, loop)

pyproject.toml

@@ -18,6 +18,7 @@ dependencies = [
     "click-log>=0.2.1",
     "voluptuous",
     "zigpy>=2.1.0",
+    "serialx>=1.13.0"
 ]
 
 [tool.setuptools.packages.find]

tests/test_ash.py

@@ -834,6 +834,54 @@ async def test_reject_condition_firmware_amplification_bug() -> None:
     assert len(transport.write.mock_calls) == 0
 
 
+@pytest.mark.parametrize("suppress_acks", [False, True])
+async def test_suppress_acks_still_naks(suppress_acks: bool) -> None:
+    """ACKs can be suppressed but NAKs must always be sent."""
+    ezsp = MagicMock()
+    protocol = ash.AshProtocol(ezsp, suppress_acks=suppress_acks)
+    transport = MagicMock()
+    transport.is_closing.return_value = False
+    protocol.connection_made(transport)
+    protocol._write_frame = MagicMock(wraps=protocol._write_frame)
+
+    # An in-sequence DATA frame is passed up but only ACKed if ACKs are not suppressed
+    good_frame = ash.DataFrame(frm_num=0, re_tx=0, ack_num=0, ezsp_frame=b"good")
+    protocol.data_received(good_frame.to_bytes() + bytes([ash.Reserved.FLAG]))
+
+    assert ezsp.data_received.mock_calls == [call(b"good")]
+
+    if suppress_acks:
+        assert protocol._write_frame.mock_calls == []
+    else:
+        assert protocol._write_frame.mock_calls == [
+            call(ash.AckFrame(res=0, ncp_ready=0, ack_num=1))
+        ]
+
+    protocol._write_frame.reset_mock()
+
+    # A retransmitted out-of-sequence frame is ACKed under the same rules
+    retx_frame = ash.DataFrame(frm_num=4, re_tx=1, ack_num=0, ezsp_frame=b"retx")
+    protocol.data_received(retx_frame.to_bytes() + bytes([ash.Reserved.FLAG]))
+
+    if suppress_acks:
+        assert protocol._write_frame.mock_calls == []
+    else:
+        assert protocol._write_frame.mock_calls == [
+            call(ash.AckFrame(res=0, ncp_ready=0, ack_num=1))
+        ]
+
+    protocol._write_frame.reset_mock()
+
+    # But a NAK is sent for an out-of-sequence frame, even with ACKs suppressed
+    bad_frame = ash.DataFrame(frm_num=4, re_tx=0, ack_num=0, ezsp_frame=b"bad")
+    protocol.data_received(bad_frame.to_bytes() + bytes([ash.Reserved.FLAG]))
+
+    assert protocol._in_reject_condition is True
+    assert protocol._write_frame.mock_calls == [
+        call(ash.NakFrame(res=0, ncp_ready=0, ack_num=1))
+    ]
+
+
 def test_ncp_failure_comparison() -> None:
     exc1 = ash.NcpFailure(code=t.NcpResetCode.ERROR_EXCEEDED_MAXIMUM_ACK_TIMEOUT_COUNT)
     exc2 = ash.NcpFailure(code=t.NcpResetCode.RESET_POWER_ON)

tests/test_uart.py

@@ -150,6 +150,79 @@ async def mock_connect(loop, protocol_factory, *args, **kwargs):
     assert len(threads) == 0
 
 
+@pytest.mark.parametrize(
+    ("path", "suppressed"),
+    [
+        ("esphome://example.com:6053/?port_name=Zigbee", True),
+        ("/dev/serial", False),
+        # A non-proxy transport is never asked to acknowledge for us
+        ("socket://example.com:1234", False),
+    ],
+)
+async def test_connect_suppress_acks(path, suppressed, monkeypatch):
+    """Suppression is decided before connecting: a frame can arrive mid-handshake."""
+    appmock = MagicMock()
+    transport = MagicMock()
+    protocols = []
+
+    async def mockconnect(loop, protocol_factory, **kwargs):
+        protocol = protocol_factory()
+        protocols.append(protocol)
+        assert protocol.suppress_acks is suppressed
+        loop.call_soon(protocol.connection_made, transport)
+        return transport, protocol
+
+    monkeypatch.setattr(zigpy.serial, "create_serial_connection", mockconnect)
+
+    gw = await uart.connect(
+        conf.SCHEMA_DEVICE(
+            {conf.CONF_DEVICE_PATH: path, conf.CONF_DEVICE_BAUDRATE: 115200}
+        ),
+        appmock,
+        use_thread=False,
+    )
+
+    assert len(protocols) == 1
+    assert protocols[0].suppress_acks is suppressed
+
+    gw.close()
+
+
+@pytest.mark.parametrize(
+    ("path", "requests_mode"),
+    [
+        ("esphome://example.com:6053/?port_name=Zigbee", True),
+        ("esphome-hass://example.com:6053/?port_name=Zigbee", True),
+        ("socket://example.com:1234", False),
+        ("/dev/serial", False),
+    ],
+)
+async def test_connect_requests_ackless_mode(path, requests_mode, monkeypatch):
+    """The mode is asked for by us, not carried in the stored path."""
+    kwargs_seen = {}
+    transport = MagicMock()
+
+    async def mockconnect(loop, protocol_factory, **kwargs):
+        kwargs_seen.update(kwargs)
+        protocol = protocol_factory()
+        loop.call_soon(protocol.connection_made, transport)
+        return transport, protocol
+
+    monkeypatch.setattr(zigpy.serial, "create_serial_connection", mockconnect)
+
+    gw = await uart.connect(
+        conf.SCHEMA_DEVICE(
+            {conf.CONF_DEVICE_PATH: path, conf.CONF_DEVICE_BAUDRATE: 115200}
+        ),
+        MagicMock(),
+        use_thread=False,
+    )
+
+    assert kwargs_seen.get("mode") == (uart.ACKLESS_MODE if requests_mode else None)
+
+    gw.close()
+
+
 @pytest.fixture
 async def gw():
     gw = uart.Gateway(MagicMock())