From a35cd28f9d88f9f3229c9abb0b0ba90e95dc175a Mon Sep 17 00:00:00 2001
From: Thomas Coratger <60488569+tcoratger@users.noreply.github.com>
Date: Fri, 12 Jun 2026 16:13:35 +0200
Subject: [PATCH] refactor(quic): dedup QuicConnection construction sites

Both the outbound connect path and the inbound handshake callback built a
QuicConnection, wired it onto the protocol, replayed buffered events, and
registered it in the connections map with identical steps.

Factor that shared sequence into a single private method on the manager so
the two sites share one construction path. Behavior is unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 .../networking/transport/quic/connection.py   | 43 +++++++++++--------
 1 file changed, 26 insertions(+), 17 deletions(-)

diff --git a/src/lean_spec/node/networking/transport/quic/connection.py b/src/lean_spec/node/networking/transport/quic/connection.py
index 3824a099e..341efc71c 100644
--- a/src/lean_spec/node/networking/transport/quic/connection.py
+++ b/src/lean_spec/node/networking/transport/quic/connection.py
@@ -382,6 +382,29 @@ def peer_id(self) -> PeerId:
         """Our local PeerId."""
         return self._peer_id
 
+    def _register_connection(
+        self,
+        protocol: LibP2PQuicProtocol,
+        peer_id: PeerId,
+        remote_address: str,
+    ) -> QuicConnection:
+        """
+        Wrap a handshaken protocol in a connection and register it.
+
+        Wires the connection back onto the protocol so future events reach it.
+        Replays events that arrived between handshake completion and connection
+        assignment, then tracks the connection under its peer ID.
+        """
+        connection = QuicConnection(
+            _protocol=protocol,
+            _peer_id=peer_id,
+            _remote_address=remote_address,
+        )
+        protocol.connection = connection
+        protocol._replay_buffered_events()
+        self._connections[peer_id] = connection
+        return connection
+
     async def connect(self, multiaddr: str) -> QuicConnection:
         """
         Connect to a peer at the given multiaddr.
@@ -440,16 +463,7 @@ async def connect(self, multiaddr: str) -> QuicConnection:
                 temp_key = IdentityKeypair.generate()
                 peer_id = temp_key.to_peer_id()
 
-            connection = QuicConnection(
-                _protocol=protocol,
-                _peer_id=peer_id,
-                _remote_address=multiaddr,
-            )
-            protocol.connection = connection
-            protocol._replay_buffered_events()
-
-            self._connections[peer_id] = connection
-            return connection
+            return self._register_connection(protocol, peer_id, multiaddr)
 
         except Exception as exception:
             raise QuicTransportError(f"Failed to connect: {exception}") from exception
@@ -503,14 +517,9 @@ def handle_handshake(protocol_instance: LibP2PQuicProtocol) -> None:
             remote_peer_id = temp_key.to_peer_id()
 
             remote_address = f"/ip4/{host}/udp/{port}/quic-v1/p2p/{remote_peer_id}"
-            connection = QuicConnection(
-                _protocol=protocol_instance,
-                _peer_id=remote_peer_id,
-                _remote_address=remote_address,
+            connection = self._register_connection(
+                protocol_instance, remote_peer_id, remote_address
             )
-            protocol_instance.connection = connection
-            protocol_instance._replay_buffered_events()
-            self._connections[remote_peer_id] = connection
 
             # Invoke callback asynchronously so it doesn't block event processing.
             asyncio.ensure_future(on_connection(connection))