Merge pull request #522 from agessaman/fix/companion-early-and-multipart-acks

fix(router): confirm companion sends from early direct and multipart ACKs

repeater/packet_router.py

@@ -13,10 +13,13 @@ from openhop_core.node.handlers.protocol_request import ProtocolRequestHandler
 from openhop_core.node.handlers.protocol_response import ProtocolResponseHandler
 from openhop_core.node.handlers.text import TextMessageHandler
 from openhop_core.node.handlers.trace import TraceHandler
+from openhop_core.protocol import Packet
 from openhop_core.protocol.constants import (
     PAYLOAD_TYPE_GRP_DATA,
     PAYLOAD_TYPE_GRP_TXT,
     PH_ROUTE_MASK,
+    PH_TYPE_MASK,
+    PH_TYPE_SHIFT,
     ROUTE_TYPE_DIRECT,
     ROUTE_TYPE_TRANSPORT_DIRECT,
 )
@@ -62,6 +65,35 @@ def _is_direct_intermediate_hop(packet) -> bool:
     )
 
 
+def _unwrap_multipart_ack(packet):
+    """Return the ACK carried by a MULTIPART ACK wrapper as its own packet, or None.
+
+    Mirrors MeshCore Mesh.cpp's PAYLOAD_TYPE_MULTIPART case: the inner packet keeps
+    the wrapper's route and path and drops the 1-byte remaining/type prefix. The
+    header's payload type is rewritten to ACK so companion bridges dispatch it to
+    their ACK handler; firmware skips that step only because it calls onAckRecv
+    directly.
+    """
+    payload = bytes(getattr(packet, "payload", b"") or b"")
+    if len(payload) < 5 or (payload[0] & 0x0F) != AckHandler.payload_type():
+        return None
+    inner = Packet()
+    inner.header = (packet.header & ~(PH_TYPE_MASK << PH_TYPE_SHIFT)) | (
+        AckHandler.payload_type() << PH_TYPE_SHIFT
+    )
+    inner.path_len = packet.path_len
+    inner.path = bytearray(packet.path)
+    inner.payload = bytearray(payload[1:])
+    inner.payload_len = len(inner.payload)
+    inner.transport_codes = list(packet.transport_codes)
+    inner._snr = getattr(packet, "_snr", 0)
+    inner._rssi = getattr(packet, "_rssi", 0)
+    inner._rx_radio_id = getattr(packet, "_rx_radio_id", None)
+    inner._injected_for_tx = getattr(packet, "_injected_for_tx", False)
+    inner._injected_origin_hash = getattr(packet, "_injected_origin_hash", None)
+    return inner
+
+
 def _is_expected_drop_reason(reason) -> bool:
     # Membership is the fast path for a DropReason member or its exact text; the
     # startswith scan covers detailed variants that embed a suffix.
@@ -607,6 +639,12 @@ class PacketRouter:
                 if len(getattr(packet, "payload", b"")) >= 4:
                     ack_crc = int.from_bytes(packet.payload[:4], "little")
                     await self._register_ack_with_dispatcher(ack_crc, "ACK")
+                    # Firmware calls onAckRecv here before checking whether this
+                    # node is the next hop, so a co-hosted companion that sent the
+                    # message is confirmed by the ACK it overhears on the way back.
+                    # Engine forwarding is unaffected.
+                    companion_bridges = self._companion_bridges_for_packet(packet, metadata)
+                    await self._fan_out_to_bridges(packet, companion_bridges, context="ACK")
 
         # Route to specific handlers for parsing only
         elif payload_type == TraceHandler.payload_type():
@@ -693,12 +731,14 @@ class PacketRouter:
 
         elif payload_type == MultipartAckHandler.payload_type():
             # MULTIPART ACK wrapper: low nibble of first byte is embedded payload type.
-            if (
-                len(getattr(packet, "payload", b"")) >= 5
-                and (packet.payload[0] & 0x0F) == AckHandler.payload_type()
-            ):
-                ack_crc = int.from_bytes(packet.payload[1:5], "little")
+            # Bridges have no MULTIPART handler, so hand them the unwrapped ACK.
+            # Policy and origin exclusion are decided on the received wrapper.
+            inner_ack = _unwrap_multipart_ack(packet)
+            if inner_ack is not None:
+                ack_crc = int.from_bytes(inner_ack.payload[:4], "little")
                 await self._register_ack_with_dispatcher(ack_crc, "multi-ACK")
+                companion_bridges = self._companion_bridges_for_packet(packet, metadata)
+                await self._fan_out_to_bridges(inner_ack, companion_bridges, context="multi-ACK")
 
         elif payload_type == TextMessageHandler.payload_type():
             # Same one-byte dest-hash collision handling as the login path above:

tests/test_packet_router.py

@@ -21,7 +21,7 @@ import time
 import unittest
 from unittest.mock import AsyncMock, MagicMock, patch
 
-from openhop_core.companion import CompanionBridge
+from openhop_core.companion import CompanionBridge, Contact
 from openhop_core.node.handlers.ack import AckHandler
 from openhop_core.node.handlers.advert import AdvertHandler
 from openhop_core.node.handlers.control import ControlHandler
@@ -38,6 +38,7 @@ from openhop_core.node.handlers.trace import TraceHandler
 from openhop_core.protocol.constants import (
     PAYLOAD_TYPE_GRP_DATA,
     PAYLOAD_TYPE_GRP_TXT,
+    PAYLOAD_TYPE_MULTIPART,
     ROUTE_TYPE_DIRECT,
     ROUTE_TYPE_FLOOD,
     ROUTE_TYPE_TRANSPORT_DIRECT,
@@ -1165,8 +1166,8 @@ class TestPacketRouterRoutingBranches(unittest.IsolatedAsyncioTestCase):
                 bridge.process_received_packet.assert_awaited_once_with(packet)
                 daemon.repeater_handler.assert_not_awaited()
 
-    async def test_direct_intermediate_ack_notifies_waiter_without_companion_delivery(self):
-        """MeshCore's early ACK notification does not make an intermediate hop a recipient."""
+    async def test_direct_intermediate_ack_notifies_waiter_and_companions(self):
+        """MeshCore calls onAckRecv on an early-received direct ACK, then forwards it."""
         daemon = _make_daemon()
         bridge = _make_bridge()
         daemon.companion_bridges = {0x01: bridge}
@@ -1179,7 +1180,7 @@ class TestPacketRouterRoutingBranches(unittest.IsolatedAsyncioTestCase):
         await router._route_packet(packet)
 
         daemon.dispatcher._register_ack_received.assert_awaited_once_with(0x95AFAB4D)
-        bridge.process_received_packet.assert_not_awaited()
+        bridge.process_received_packet.assert_awaited_once_with(packet)
         daemon.repeater_handler.assert_awaited_once()
 
     async def test_direct_intermediate_control_high_bit_is_released_without_local_delivery(self):
@@ -1541,3 +1542,133 @@ class TestCompanionDeliveryFailureHandling(unittest.IsolatedAsyncioTestCase):
         await router._route_packet(pkt)
 
         daemon.repeater_handler.assert_awaited_once()
+
+
+class TestCompanionAckConfirmation(unittest.IsolatedAsyncioTestCase):
+    """ACK shapes MeshCore passes to onAckRecv must confirm a real companion send.
+
+    Each case sends a DM from a real CompanionBridge so its ACK CRC is pending,
+    routes the ACK the firmware would accept, and checks send_confirmed fired.
+    """
+
+    async def _bridge_with_pending_send(self):
+        async def inject_packet(packet, wait_for_ack=False, expected_crc=None):
+            return True
+
+        bridge = CompanionBridge(LocalIdentity(), inject_packet, node_name="Sender")
+        peer = LocalIdentity().get_public_key()
+        bridge.add_update_contact(Contact(public_key=peer, name="Peer"))
+        result = await bridge.send_text_message(peer, "hello", wait_for_ack=False)
+        self.assertTrue(result.success)
+        self.assertNotEqual(result.expected_ack, 0)
+        confirmed = []
+
+        async def on_confirmed(crc, trip_ms):
+            confirmed.append(crc)
+
+        bridge.on_send_confirmed(on_confirmed)
+        return bridge, result.expected_ack, confirmed
+
+    @staticmethod
+    def _packet(payload_type, route_type, path, payload):
+        packet = Packet()
+        packet.header = (payload_type << 2) | route_type
+        packet.path = bytearray(path)
+        packet.path_len = len(path)
+        packet.payload = bytearray(payload)
+        packet.payload_len = len(payload)
+        return packet
+
+    async def _route(self, bridge, packet, policy_engine=None):
+        daemon = _make_daemon()
+        if policy_engine is not None:
+            daemon.repeater_handler.policy_engine = policy_engine
+        daemon.companion_bridges = {0x01: bridge}
+        router = PacketRouter(daemon)
+        await router._route_packet(packet)
+        return daemon
+
+    async def test_direct_ack_with_1_hop_left(self):
+        bridge, crc, confirmed = await self._bridge_with_pending_send()
+        # 6-byte firmware ACK: CRC (LE) + ext-attempt + random byte.
+        packet = self._packet(
+            AckHandler.payload_type(),
+            ROUTE_TYPE_DIRECT,
+            [0x42],
+            crc.to_bytes(4, "little") + b"\x00\x7f",
+        )
+
+        daemon = await self._route(bridge, packet)
+
+        self.assertEqual(confirmed, [crc])
+        daemon.dispatcher._register_ack_received.assert_awaited_once_with(crc)
+        daemon.repeater_handler.assert_awaited_once()  # still forwarded
+
+    async def test_multipart_ack_direct_final(self):
+        bridge, crc, confirmed = await self._bridge_with_pending_send()
+        # remaining=1 in the high nibble, embedded type ACK in the low nibble.
+        packet = self._packet(
+            PAYLOAD_TYPE_MULTIPART,
+            ROUTE_TYPE_DIRECT,
+            [],
+            bytes([(1 << 4) | AckHandler.payload_type()]) + crc.to_bytes(4, "little"),
+        )
+
+        daemon = await self._route(bridge, packet)
+
+        self.assertEqual(confirmed, [crc])
+        daemon.dispatcher._register_ack_received.assert_awaited_once_with(crc)
+
+    async def test_multipart_ack_flood(self):
+        bridge, crc, confirmed = await self._bridge_with_pending_send()
+        packet = self._packet(
+            PAYLOAD_TYPE_MULTIPART,
+            ROUTE_TYPE_FLOOD,
+            [0x11, 0x22],
+            bytes([AckHandler.payload_type()]) + crc.to_bytes(4, "little"),
+        )
+
+        daemon = await self._route(bridge, packet)
+
+        self.assertEqual(confirmed, [crc])
+        daemon.dispatcher._register_ack_received.assert_awaited_once_with(crc)
+
+    async def test_policy_drop_blocks_early_and_multipart_ack_confirmation(self):
+        drop_all = PolicyEngine({"enabled": True, "default_action": "drop", "rules": []})
+        cases = {
+            "direct ACK with 1 hop left": (
+                AckHandler.payload_type(),
+                ROUTE_TYPE_DIRECT,
+                [0x42],
+                lambda crc: crc.to_bytes(4, "little"),
+            ),
+            "MULTIPART ACK flood": (
+                PAYLOAD_TYPE_MULTIPART,
+                ROUTE_TYPE_FLOOD,
+                [0x11],
+                lambda crc: bytes([AckHandler.payload_type()]) + crc.to_bytes(4, "little"),
+            ),
+        }
+        for name, (ptype, route, path, build) in cases.items():
+            with self.subTest(name):
+                bridge, crc, confirmed = await self._bridge_with_pending_send()
+                packet = self._packet(ptype, route, path, build(crc))
+
+                await self._route(bridge, packet, policy_engine=drop_all)
+
+                self.assertEqual(confirmed, [])
+
+    async def test_multipart_ack_withheld_from_originating_companion(self):
+        bridge, crc, confirmed = await self._bridge_with_pending_send()
+        packet = self._packet(
+            PAYLOAD_TYPE_MULTIPART,
+            ROUTE_TYPE_DIRECT,
+            [],
+            bytes([AckHandler.payload_type()]) + crc.to_bytes(4, "little"),
+        )
+        packet._injected_for_tx = True
+        packet._injected_origin_hash = "0x01"
+
+        await self._route(bridge, packet)
+
+        self.assertEqual(confirmed, [])