fix: initial connection succesfull

This commit is contained in:
Akash Mondal
2025-06-30 11:16:08 +00:00
committed by lla-dane
parent 2689040d48
commit bbe632bd85
6 changed files with 120 additions and 79 deletions

View File

@ -115,7 +115,9 @@ async def run_client(destination: str, seed: int | None = None) -> None:
info = info_from_p2p_addr(maddr)
# Connect to server
print("STARTING CLIENT CONNECTION PROCESS")
await host.connect(info)
print("CLIENT CONNECTED TO SERVER")
# Start a stream with the destination
stream = await host.new_stream(info.peer_id, [PROTOCOL_ID])

View File

@ -40,6 +40,7 @@ from libp2p.transport.exceptions import (
OpenConnectionError,
SecurityUpgradeFailure,
)
from libp2p.transport.quic.transport import QUICTransport
from libp2p.transport.upgrader import (
TransportUpgrader,
)
@ -114,6 +115,11 @@ class Swarm(Service, INetworkService):
# Create a nursery for listener tasks.
self.listener_nursery = nursery
self.event_listener_nursery_created.set()
if isinstance(self.transport, QUICTransport):
self.transport.set_background_nursery(nursery)
self.transport.set_swarm(self)
try:
await self.manager.wait_finished()
finally:
@ -177,6 +183,14 @@ class Swarm(Service, INetworkService):
"""
Try to create a connection to peer_id with addr.
"""
# QUIC Transport
if isinstance(self.transport, QUICTransport):
raw_conn = await self.transport.dial(addr, peer_id)
print("detected QUIC connection, skipping upgrade steps")
swarm_conn = await self.add_conn(raw_conn)
print("successfully dialed peer %s via QUIC", peer_id)
return swarm_conn
try:
raw_conn = await self.transport.dial(addr)
except OpenConnectionError as error:
@ -187,14 +201,6 @@ class Swarm(Service, INetworkService):
logger.debug("dialed peer %s over base transport", peer_id)
# NEW: Check if this is a QUIC connection (already secure and muxed)
if isinstance(raw_conn, IMuxedConn):
# QUIC connections are already secure and muxed, skip upgrade steps
logger.debug("detected QUIC connection, skipping upgrade steps")
swarm_conn = await self.add_conn(raw_conn)
logger.debug("successfully dialed peer %s via QUIC", peer_id)
return swarm_conn
# Standard TCP flow - security then mux upgrade
try:
secured_conn = await self.upgrader.upgrade_security(raw_conn, True, peer_id)

View File

@ -147,7 +147,8 @@ class MultiselectClient(IMultiselectClient):
except MultiselectCommunicatorError as error:
raise MultiselectClientError() from error
if response == protocol_str:
print("Response: ", response)
if response == protocol:
return protocol
if response == PROTOCOL_NOT_FOUND_MSG:
raise MultiselectClientError("protocol not supported")

View File

@ -3,11 +3,12 @@ QUIC Connection implementation.
Uses aioquic's sans-IO core with trio for async operations.
"""
from collections.abc import Awaitable, Callable
import logging
import socket
from sys import stdout
import time
from typing import TYPE_CHECKING, Any, Optional, Set
from typing import TYPE_CHECKING, Any, Optional
from aioquic.quic import events
from aioquic.quic.connection import QuicConnection
@ -75,7 +76,7 @@ class QUICConnection(IRawConnection, IMuxedConn):
self,
quic_connection: QuicConnection,
remote_addr: tuple[str, int],
peer_id: ID | None,
peer_id: ID,
local_peer_id: ID,
is_initiator: bool,
maddr: multiaddr.Multiaddr,
@ -102,7 +103,7 @@ class QUICConnection(IRawConnection, IMuxedConn):
"""
self._quic = quic_connection
self._remote_addr = remote_addr
self._peer_id = peer_id
self.peer_id = peer_id
self._local_peer_id = local_peer_id
self.__is_initiator = is_initiator
self._maddr = maddr
@ -147,12 +148,14 @@ class QUICConnection(IRawConnection, IMuxedConn):
self._background_tasks_started = False
self._nursery: trio.Nursery | None = None
self._event_processing_task: Any | None = None
self.on_close: Callable[[], Awaitable[None]] | None = None
self.event_started = trio.Event()
# *** NEW: Connection ID tracking - CRITICAL for fixing the original issue ***
self._available_connection_ids: Set[bytes] = set()
self._current_connection_id: Optional[bytes] = None
self._retired_connection_ids: Set[bytes] = set()
self._connection_id_sequence_numbers: Set[int] = set()
self._available_connection_ids: set[bytes] = set()
self._current_connection_id: bytes | None = None
self._retired_connection_ids: set[bytes] = set()
self._connection_id_sequence_numbers: set[int] = set()
# Event processing control
self._event_processing_active = False
@ -235,7 +238,7 @@ class QUICConnection(IRawConnection, IMuxedConn):
def remote_peer_id(self) -> ID | None:
"""Get the remote peer ID."""
return self._peer_id
return self.peer_id
# *** NEW: Connection ID management methods ***
def get_connection_id_stats(self) -> dict[str, Any]:
@ -252,7 +255,7 @@ class QUICConnection(IRawConnection, IMuxedConn):
"available_cid_list": [cid.hex() for cid in self._available_connection_ids],
}
def get_current_connection_id(self) -> Optional[bytes]:
def get_current_connection_id(self) -> bytes | None:
"""Get the current connection ID."""
return self._current_connection_id
@ -273,7 +276,8 @@ class QUICConnection(IRawConnection, IMuxedConn):
raise QUICConnectionError("Cannot start a closed connection")
self._started = True
logger.debug(f"Starting QUIC connection to {self._peer_id}")
self.event_started.set()
logger.debug(f"Starting QUIC connection to {self.peer_id}")
try:
# If this is a client connection, we need to establish the connection
@ -284,7 +288,7 @@ class QUICConnection(IRawConnection, IMuxedConn):
self._established = True
self._connected_event.set()
logger.debug(f"QUIC connection to {self._peer_id} started")
logger.debug(f"QUIC connection to {self.peer_id} started")
except Exception as e:
logger.error(f"Failed to start connection: {e}")
@ -356,7 +360,7 @@ class QUICConnection(IRawConnection, IMuxedConn):
await self._verify_peer_identity_with_security()
self._established = True
logger.info(f"QUIC connection established with {self._peer_id}")
logger.info(f"QUIC connection established with {self.peer_id}")
except Exception as e:
logger.error(f"Failed to establish connection: {e}")
@ -491,17 +495,16 @@ class QUICConnection(IRawConnection, IMuxedConn):
# Verify peer identity using security manager
verified_peer_id = self._security_manager.verify_peer_identity(
self._peer_certificate,
self._peer_id, # Expected peer ID for outbound connections
self.peer_id, # Expected peer ID for outbound connections
)
# Update peer ID if it wasn't known (inbound connections)
if not self._peer_id:
self._peer_id = verified_peer_id
if not self.peer_id:
self.peer_id = verified_peer_id
logger.info(f"Discovered peer ID from certificate: {verified_peer_id}")
elif self._peer_id != verified_peer_id:
elif self.peer_id != verified_peer_id:
raise QUICPeerVerificationError(
f"Peer ID mismatch: expected {self._peer_id}, "
f"got {verified_peer_id}"
f"Peer ID mismatch: expected {self.peer_id}, got {verified_peer_id}"
)
self._peer_verified = True
@ -605,7 +608,7 @@ class QUICConnection(IRawConnection, IMuxedConn):
info: dict[str, bool | Any | None] = {
"peer_verified": self._peer_verified,
"handshake_complete": self._handshake_completed,
"peer_id": str(self._peer_id) if self._peer_id else None,
"peer_id": str(self.peer_id) if self.peer_id else None,
"local_peer_id": str(self._local_peer_id),
"is_initiator": self.__is_initiator,
"has_certificate": self._peer_certificate is not None,
@ -1188,7 +1191,7 @@ class QUICConnection(IRawConnection, IMuxedConn):
return
self._closed = True
logger.debug(f"Closing QUIC connection to {self._peer_id}")
logger.debug(f"Closing QUIC connection to {self.peer_id}")
try:
# Close all streams gracefully
@ -1213,8 +1216,12 @@ class QUICConnection(IRawConnection, IMuxedConn):
except Exception:
pass
if self.on_close:
await self.on_close()
# Close QUIC connection
self._quic.close()
if self._socket:
await self._transmit() # Send close frames
@ -1226,7 +1233,7 @@ class QUICConnection(IRawConnection, IMuxedConn):
self._streams.clear()
self._closed_event.set()
logger.debug(f"QUIC connection to {self._peer_id} closed")
logger.debug(f"QUIC connection to {self.peer_id} closed")
except Exception as e:
logger.error(f"Error during connection close: {e}")
@ -1266,6 +1273,7 @@ class QUICConnection(IRawConnection, IMuxedConn):
QUICStreamClosedError: If stream is closed for reading.
QUICStreamResetError: If stream was reset.
QUICStreamTimeoutError: If read timeout occurs.
"""
# This method doesn't make sense for a muxed connection
# It's here for interface compatibility but should not be used
@ -1325,7 +1333,7 @@ class QUICConnection(IRawConnection, IMuxedConn):
def __repr__(self) -> str:
return (
f"QUICConnection(peer={self._peer_id}, "
f"QUICConnection(peer={self.peer_id}, "
f"addr={self._remote_addr}, "
f"initiator={self.__is_initiator}, "
f"verified={self._peer_verified}, "
@ -1335,4 +1343,4 @@ class QUICConnection(IRawConnection, IMuxedConn):
)
def __str__(self) -> str:
return f"QUICConnection({self._peer_id})"
return f"QUICConnection({self.peer_id})"

View File

@ -12,18 +12,19 @@ from typing import TYPE_CHECKING
from aioquic.quic import events
from aioquic.quic.configuration import QuicConfiguration
from aioquic.quic.connection import QuicConnection
from aioquic.quic.packet import QuicPacketType
from multiaddr import Multiaddr
import trio
from libp2p.abc import IListener
from libp2p.custom_types import THandler, TProtocol
from libp2p.custom_types import (
TProtocol,
TQUICConnHandlerFn,
)
from libp2p.transport.quic.security import (
LIBP2P_TLS_EXTENSION_OID,
QUICTLSConfigManager,
)
from libp2p.custom_types import TQUICConnHandlerFn
from libp2p.custom_types import TQUICStreamHandlerFn
from aioquic.quic.packet import QuicPacketType
from .config import QUICTransportConfig
from .connection import QUICConnection
@ -1099,12 +1100,21 @@ class QUICListener(IListener):
if not is_quic_multiaddr(maddr):
raise QUICListenError(f"Invalid QUIC multiaddr: {maddr}")
if self._transport._background_nursery:
active_nursery = self._transport._background_nursery
logger.debug("Using transport background nursery for listener")
elif nursery:
active_nursery = nursery
logger.debug("Using provided nursery for listener")
else:
raise QUICListenError("No nursery available")
try:
host, port = quic_multiaddr_to_endpoint(maddr)
# Create and configure socket
self._socket = await self._create_socket(host, port)
self._nursery = nursery
self._nursery = active_nursery
# Get the actual bound address
bound_host, bound_port = self._socket.getsockname()
@ -1115,7 +1125,7 @@ class QUICListener(IListener):
self._listening = True
# Start packet handling loop
nursery.start_soon(self._handle_incoming_packets)
active_nursery.start_soon(self._handle_incoming_packets)
logger.info(
f"QUIC listener started on {bound_maddr} with connection ID support"
@ -1217,33 +1227,22 @@ class QUICListener(IListener):
async def _handle_new_established_connection(
self, connection: QUICConnection
) -> None:
"""Handle newly established connection with proper stream management."""
"""Handle newly established connection by adding to swarm."""
try:
logger.debug(
f"Handling new established connection from {connection._remote_addr}"
f"New QUIC connection established from {connection._remote_addr}"
)
# Accept incoming streams and pass them to the handler
while not connection.is_closed:
try:
print(f"🔧 CONN_HANDLER: Waiting for stream...")
stream = await connection.accept_stream(timeout=1.0)
print(f"✅ CONN_HANDLER: Accepted stream {stream.stream_id}")
if self._nursery:
# Pass STREAM to handler, not connection
self._nursery.start_soon(self._handler, stream)
print(
f"✅ CONN_HANDLER: Started handler for stream {stream.stream_id}"
)
except trio.TooSlowError:
continue # Timeout is normal
except Exception as e:
logger.error(f"Error accepting stream: {e}")
break
if self._transport._swarm:
logger.debug("Adding QUIC connection directly to swarm")
await self._transport._swarm.add_conn(connection)
logger.debug("Successfully added QUIC connection to swarm")
else:
logger.error("No swarm available for QUIC connection")
await connection.close()
except Exception as e:
logger.error(f"Error in connection handler: {e}")
logger.error(f"Error adding QUIC connection to swarm: {e}")
await connection.close()
def get_addrs(self) -> tuple[Multiaddr]:

View File

@ -9,6 +9,7 @@ import copy
import logging
import ssl
import sys
from typing import TYPE_CHECKING, cast
from aioquic.quic.configuration import (
QuicConfiguration,
@ -21,13 +22,12 @@ import multiaddr
import trio
from libp2p.abc import (
IRawConnection,
ITransport,
)
from libp2p.crypto.keys import (
PrivateKey,
)
from libp2p.custom_types import THandler, TProtocol, TQUICConnHandlerFn
from libp2p.custom_types import TProtocol, TQUICConnHandlerFn
from libp2p.peer.id import (
ID,
)
@ -40,6 +40,11 @@ from libp2p.transport.quic.utils import (
quic_version_to_wire_format,
)
if TYPE_CHECKING:
from libp2p.network.swarm import Swarm
else:
Swarm = cast(type, object)
from .config import (
QUICTransportConfig,
)
@ -112,10 +117,20 @@ class QUICTransport(ITransport):
# Resource management
self._closed = False
self._nursery_manager = trio.CapacityLimiter(1)
self._background_nursery: trio.Nursery | None = None
logger.info(
f"Initialized QUIC transport with security for peer {self._peer_id}"
)
self._swarm = None
print(f"Initialized QUIC transport with security for peer {self._peer_id}")
def set_background_nursery(self, nursery: trio.Nursery) -> None:
"""Set the nursery to use for background tasks (called by swarm)."""
self._background_nursery = nursery
print("Transport background nursery set")
def set_swarm(self, swarm) -> None:
"""Set the swarm for adding incoming connections."""
self._swarm = swarm
def _setup_quic_configurations(self) -> None:
"""Setup QUIC configurations."""
@ -184,7 +199,7 @@ class QUICTransport(ITransport):
draft29_client_config
)
logger.info("QUIC configurations initialized with libp2p TLS security")
print("QUIC configurations initialized with libp2p TLS security")
except Exception as e:
raise QUICSecurityError(
@ -214,14 +229,13 @@ class QUICTransport(ITransport):
config.verify_mode = ssl.CERT_NONE
logger.debug("Successfully applied TLS configuration to QUIC config")
print("Successfully applied TLS configuration to QUIC config")
except Exception as e:
raise QUICSecurityError(f"Failed to apply TLS configuration: {e}") from e
async def dial(
self, maddr: multiaddr.Multiaddr, peer_id: ID | None = None
) -> QUICConnection:
# type: ignore
async def dial(self, maddr: multiaddr.Multiaddr, peer_id: ID) -> QUICConnection:
"""
Dial a remote peer using QUIC transport with security verification.
@ -243,6 +257,9 @@ class QUICTransport(ITransport):
if not is_quic_multiaddr(maddr):
raise QUICDialError(f"Invalid QUIC multiaddr: {maddr}")
if not peer_id:
raise QUICDialError("Peer id cannot be null")
try:
# Extract connection details from multiaddr
host, port = quic_multiaddr_to_endpoint(maddr)
@ -257,9 +274,7 @@ class QUICTransport(ITransport):
config.is_client = True
config.quic_logger = QuicLogger()
logger.debug(
f"Dialing QUIC connection to {host}:{port} (version: {quic_version})"
)
print(f"Dialing QUIC connection to {host}:{port} (version: {quic_version})")
print("Start QUIC Connection")
# Create QUIC connection using aioquic's sans-IO core
@ -279,8 +294,18 @@ class QUICTransport(ITransport):
)
# Establish connection using trio
async with trio.open_nursery() as nursery:
await connection.connect(nursery)
if self._background_nursery:
# Use swarm's long-lived nursery - background tasks persist!
await connection.connect(self._background_nursery)
print("Using background nursery for connection tasks")
else:
# Fallback to temporary nursery (with warning)
print(
"No background nursery available. Connection background tasks "
"may be cancelled when dial completes."
)
async with trio.open_nursery() as temp_nursery:
await connection.connect(temp_nursery)
# Verify peer identity after TLS handshake
if peer_id:
@ -290,7 +315,7 @@ class QUICTransport(ITransport):
conn_id = f"{host}:{port}:{peer_id}"
self._connections[conn_id] = connection
logger.info(f"Successfully dialed secure QUIC connection to {peer_id}")
print(f"Successfully dialed secure QUIC connection to {peer_id}")
return connection
except Exception as e:
@ -329,7 +354,7 @@ class QUICTransport(ITransport):
f"{expected_peer_id}, got {verified_peer_id}"
)
logger.info(f"Peer identity verified: {verified_peer_id}")
print(f"Peer identity verified: {verified_peer_id}")
print(f"Peer identity verified: {verified_peer_id}")
except Exception as e:
@ -368,7 +393,7 @@ class QUICTransport(ITransport):
)
self._listeners.append(listener)
logger.debug("Created QUIC listener with security")
print("Created QUIC listener with security")
return listener
def can_dial(self, maddr: multiaddr.Multiaddr) -> bool:
@ -414,7 +439,7 @@ class QUICTransport(ITransport):
return
self._closed = True
logger.info("Closing QUIC transport")
print("Closing QUIC transport")
# Close all active connections and listeners concurrently using trio nursery
async with trio.open_nursery() as nursery:
@ -429,7 +454,7 @@ class QUICTransport(ITransport):
self._connections.clear()
self._listeners.clear()
logger.info("QUIC transport closed")
print("QUIC transport closed")
def get_stats(self) -> dict[str, int | list[str] | object]:
"""Get transport statistics including security info."""