From f2a5eab726a1cf4aba4368914aa7a3eebd5ac9b6 Mon Sep 17 00:00:00 2001 From: agessaman Date: Sun, 8 Mar 2026 19:39:49 -0700 Subject: [PATCH] Refactor packet handling in PacketRouter and RepeaterHandler - Removed redundant original_path assignment in `RepeaterHandler` to streamline packet processing. - Introduced `_is_direct_final_hop` helper method in `PacketRouter` to determine if a packet is the final destination for direct routes with an empty path. - Updated comments in `PacketRouter` to clarify the handling of packets during routing, especially for direct forwarding scenarios. - Adjusted logic to ensure packets are correctly processed or delivered based on their routing status, enhancing overall packet management. --- repeater/engine.py | 1 - repeater/packet_router.py | 51 ++++++++++++++++++++++++++++++--------- tests/test_engine.py | 5 ++-- 3 files changed, 42 insertions(+), 15 deletions(-) diff --git a/repeater/engine.py b/repeater/engine.py index d06f9c8..e192df6 100644 --- a/repeater/engine.py +++ b/repeater/engine.py @@ -789,7 +789,6 @@ class RepeaterHandler(BaseHandler): self.mark_seen(packet) - original_path = list(packet.path) # Remove first hash entry (hash_size bytes) packet.path = bytearray(packet.path[hash_size:]) packet.path_len = PathUtils.encode_path_len(hash_size, hop_count - 1) diff --git a/repeater/packet_router.py b/repeater/packet_router.py index 53c5ff1..d9c02e9 100644 --- a/repeater/packet_router.py +++ b/repeater/packet_router.py @@ -13,6 +13,11 @@ from pymc_core.node.handlers.protocol_request import ProtocolRequestHandler from pymc_core.node.handlers.protocol_response import ProtocolResponseHandler from pymc_core.node.handlers.text import TextMessageHandler from pymc_core.node.handlers.trace import TraceHandler +from pymc_core.protocol.constants import ( + PH_ROUTE_MASK, + ROUTE_TYPE_DIRECT, + ROUTE_TYPE_TRANSPORT_DIRECT, +) logger = logging.getLogger("PacketRouter") @@ -29,6 +34,15 @@ def _companion_dedup_key(packet) -> str | None: return None +def _is_direct_final_hop(packet) -> bool: + """True if packet is DIRECT (or TRANSPORT_DIRECT) with empty path — we're the final destination.""" + route = getattr(packet, "header", 0) & PH_ROUTE_MASK + if route != ROUTE_TYPE_DIRECT and route != ROUTE_TYPE_TRANSPORT_DIRECT: + return False + path = getattr(packet, "path", None) + return not path or len(path) == 0 + + class PacketRouter: def __init__(self, daemon_instance): @@ -197,15 +211,14 @@ class PacketRouter: self._record_for_ui(packet, metadata) elif payload_type == AckHandler.payload_type(): - # ACK has no dest in payload (4-byte CRC only); deliver to all bridges so sender sees send_confirmed + # ACK has no dest in payload (4-byte CRC only); deliver to all bridges so sender sees send_confirmed. + # Do not set processed_by_injection so packet also reaches engine for DIRECT forwarding when we're a middle hop. companion_bridges = getattr(self.daemon, "companion_bridges", {}) for bridge in companion_bridges.values(): try: await bridge.process_received_packet(packet) except Exception as e: logger.debug(f"Companion bridge ACK error: {e}") - processed_by_injection = True - self._record_for_ui(packet, metadata) elif payload_type == TextMessageHandler.payload_type(): dest_hash = packet.payload[0] if packet.payload else None @@ -226,8 +239,7 @@ class PacketRouter: if dest_hash is not None and dest_hash in companion_bridges: if self._should_deliver_path_to_companions(packet): await companion_bridges[dest_hash].process_received_packet(packet) - processed_by_injection = True - self._record_for_ui(packet, metadata) + # Do not set processed_by_injection so packet also reaches engine for DIRECT forwarding when we're a middle hop. elif companion_bridges and self._should_deliver_path_to_companions(packet): # Dest not in bridges: path-return with ephemeral dest (e.g. multi-hop login). # Deliver to all bridges; each will try to decrypt and ignore if not relevant. @@ -241,8 +253,7 @@ class PacketRouter: dest_hash or 0, len(companion_bridges), ) - processed_by_injection = True - self._record_for_ui(packet, metadata) + # Do not set processed_by_injection so packet also reaches engine for DIRECT forwarding when we're a middle hop. elif self.daemon.path_helper: await self.daemon.path_helper.process_path_packet(packet) @@ -251,6 +262,7 @@ class PacketRouter: # Deliver to the bridge that is the destination, or to all bridges when the # response is addressed to this repeater (path-based reply: firmware sends # to first hop instead of original requester). + # Do not set processed_by_injection so packet also reaches engine for DIRECT forwarding when we're a middle hop. dest_hash = packet.payload[0] if packet.payload and len(packet.payload) >= 1 else None companion_bridges = getattr(self.daemon, "companion_bridges", {}) local_hash = getattr(self.daemon, "local_hash", None) @@ -263,8 +275,6 @@ class PacketRouter: ) except Exception as e: logger.debug(f"Companion bridge RESPONSE error: {e}") - processed_by_injection = True - self._record_for_ui(packet, metadata) elif dest_hash == local_hash and companion_bridges: # Response addressed to this repeater (e.g. path-based reply to first hop) for bridge in companion_bridges.values(): @@ -277,8 +287,6 @@ class PacketRouter: dest_hash, len(companion_bridges), ) - processed_by_injection = True - self._record_for_ui(packet, metadata) elif companion_bridges: # Dest not in bridges and not local: likely ANON_REQ response (dest = ephemeral # sender hash). Deliver to all bridges; each will try to decrypt and ignore if @@ -293,12 +301,15 @@ class PacketRouter: dest_hash or 0, len(companion_bridges), ) + if companion_bridges and _is_direct_final_hop(packet): + # DIRECT with empty path: we're the final hop; don't pass to engine (it would drop with "Direct: no path") processed_by_injection = True self._record_for_ui(packet, metadata) elif payload_type == ProtocolResponseHandler.payload_type(): # PAYLOAD_TYPE_PATH (0x08): protocol responses (telemetry, binary, etc.). # Deliver at most once per logical packet so the client is not spammed with duplicates. + # Do not set processed_by_injection so packet also reaches engine for DIRECT forwarding when we're a middle hop. companion_bridges = getattr(self.daemon, "companion_bridges", {}) if companion_bridges and self._should_deliver_path_to_companions(packet): for bridge in companion_bridges.values(): @@ -306,7 +317,14 @@ class PacketRouter: await bridge.process_received_packet(packet) except Exception as e: logger.debug(f"Companion bridge RESPONSE error: {e}") - if companion_bridges: + if companion_bridges and _is_direct_final_hop(packet): + # DIRECT with empty path: we're the final hop; ensure delivery to all bridges (anon) + if not self._should_deliver_path_to_companions(packet): + for bridge in companion_bridges.values(): + try: + await bridge.process_received_packet(packet) + except Exception as e: + logger.debug(f"Companion bridge RESPONSE (final hop) error: {e}") processed_by_injection = True self._record_for_ui(packet, metadata) @@ -322,6 +340,15 @@ class PacketRouter: if handled: processed_by_injection = True self._record_for_ui(packet, metadata) + elif companion_bridges and _is_direct_final_hop(packet): + # DIRECT with empty path: we're the final hop; deliver to all bridges for anon matching + for bridge in companion_bridges.values(): + try: + await bridge.process_received_packet(packet) + except Exception as e: + logger.debug(f"Companion bridge REQ (final hop) error: {e}") + processed_by_injection = True + self._record_for_ui(packet, metadata) elif payload_type == GroupTextHandler.payload_type(): # GRP_TXT: pass to all companions (they filter by channel); still forward diff --git a/tests/test_engine.py b/tests/test_engine.py index 7c4adf9..5669614 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -287,9 +287,10 @@ class TestDirectForward: assert result.path_len == 2 def test_single_hop_path_consumed(self, handler): - """After consuming our hash the path becomes empty — packet delivered.""" + """Single hop to us: we strip and return packet with empty path (forward so it can reach destination).""" pkt = _make_direct_packet(path=bytes([LOCAL_HASH])) result = handler.direct_forward(pkt) + assert result is not None assert list(result.path) == [] assert result.path_len == 0 @@ -1044,7 +1045,7 @@ GOOD_PACKETS = [ lambda: _make_flood_packet(payload=b"\xAB\x01\x02\x03", payload_type=4)), ("good_direct_minimal", - "Direct, 1-byte payload, single hop to us", + "Direct, 1-byte payload, single hop to us (forward with empty path)", lambda: _make_direct_packet(payload=b"\x01", path=bytes([LOCAL_HASH]))), ("good_direct_multihop",