Merge branch 'main' into chore01

This commit is contained in:
Manu Sheel Gupta
2025-09-24 22:40:38 +05:30
committed by GitHub
119 changed files with 16551 additions and 520 deletions

View File

@ -36,10 +36,48 @@ jobs:
- uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python }}
- name: Install Nim for interop testing
if: matrix.toxenv == 'interop'
run: |
echo "Installing Nim for nim-libp2p interop testing..."
curl -sSf https://nim-lang.org/choosenim/init.sh | sh -s -- -y --firstInstall
echo "$HOME/.nimble/bin" >> $GITHUB_PATH
echo "$HOME/.choosenim/toolchains/nim-stable/bin" >> $GITHUB_PATH
- name: Cache nimble packages
if: matrix.toxenv == 'interop'
uses: actions/cache@v4
with:
path: |
~/.nimble
~/.choosenim/toolchains/*/lib
key: ${{ runner.os }}-nimble-${{ hashFiles('**/nim_echo_server.nim') }}
restore-keys: |
${{ runner.os }}-nimble-
- name: Build nim interop binaries
if: matrix.toxenv == 'interop'
run: |
export PATH="$HOME/.nimble/bin:$HOME/.choosenim/toolchains/nim-stable/bin:$PATH"
cd tests/interop/nim_libp2p
./scripts/setup_nim_echo.sh
- run: |
python -m pip install --upgrade pip
python -m pip install tox
- run: |
- name: Run Tests or Generate Docs
run: |
if [[ "${{ matrix.toxenv }}" == 'docs' ]]; then
export TOXENV=docs
else
export TOXENV=py${{ matrix.python }}-${{ matrix.toxenv }}
fi
# Set PATH for nim commands during tox
if [[ "${{ matrix.toxenv }}" == 'interop' ]]; then
export PATH="$HOME/.nimble/bin:$HOME/.choosenim/toolchains/nim-stable/bin:$PATH"
fi
python -m tox run -r
windows:
@ -65,5 +103,5 @@ jobs:
if [[ "${{ matrix.toxenv }}" == "wheel" ]]; then
python -m tox run -e windows-wheel
else
python -m tox run -e py311-${{ matrix.toxenv }}
python -m tox run -e py${{ matrix.python-version }}-${{ matrix.toxenv }}
fi

4
.gitignore vendored
View File

@ -178,6 +178,10 @@ env.bak/
#lockfiles
uv.lock
poetry.lock
tests/interop/js_libp2p/js_node/node_modules/
tests/interop/js_libp2p/js_node/package-lock.json
tests/interop/js_libp2p/js_node/src/node_modules/
tests/interop/js_libp2p/js_node/src/package-lock.json
# Sphinx documentation build
_build/

View File

@ -61,12 +61,12 @@ ______________________________________________________________________
### Discovery
| **Discovery** | **Status** | **Source** |
| -------------------- | :--------: | :--------------------------------------------------------------------------------: |
| **`bootstrap`** | ✅ | [source](https://github.com/libp2p/py-libp2p/tree/main/libp2p/discovery/bootstrap) |
| **`random-walk`** | 🌱 | |
| **`mdns-discovery`** | ✅ | [source](https://github.com/libp2p/py-libp2p/tree/main/libp2p/discovery/mdns) |
| **`rendezvous`** | 🌱 | |
| **Discovery** | **Status** | **Source** |
| -------------------- | :--------: | :----------------------------------------------------------------------------------: |
| **`bootstrap`** | ✅ | [source](https://github.com/libp2p/py-libp2p/tree/main/libp2p/discovery/bootstrap) |
| **`random-walk`** | | [source](https://github.com/libp2p/py-libp2p/tree/main/libp2p/discovery/random_walk) |
| **`mdns-discovery`** | ✅ | [source](https://github.com/libp2p/py-libp2p/tree/main/libp2p/discovery/mdns) |
| **`rendezvous`** | 🌱 | |
______________________________________________________________________

View File

@ -36,12 +36,14 @@ Create a file named ``relay_node.py`` with the following content:
from libp2p.relay.circuit_v2.transport import CircuitV2Transport
from libp2p.relay.circuit_v2.config import RelayConfig
from libp2p.tools.async_service import background_trio_service
from libp2p.utils import get_wildcard_address
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger("relay_node")
async def run_relay():
listen_addr = multiaddr.Multiaddr("/ip4/0.0.0.0/tcp/9000")
# Use wildcard address to listen on all interfaces
listen_addr = get_wildcard_address(9000)
host = new_host()
config = RelayConfig(
@ -107,6 +109,7 @@ Create a file named ``destination_node.py`` with the following content:
from libp2p.relay.circuit_v2.config import RelayConfig
from libp2p.peer.peerinfo import info_from_p2p_addr
from libp2p.tools.async_service import background_trio_service
from libp2p.utils import get_wildcard_address
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger("destination_node")
@ -139,7 +142,8 @@ Create a file named ``destination_node.py`` with the following content:
Run a simple destination node that accepts connections.
This is a simplified version that doesn't use the relay functionality.
"""
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/9001")
# Create a libp2p host - use wildcard address to listen on all interfaces
listen_addr = get_wildcard_address(9001)
host = new_host()
# Configure as a relay receiver (stop)
@ -252,14 +256,15 @@ Create a file named ``source_node.py`` with the following content:
from libp2p.peer.peerinfo import info_from_p2p_addr
from libp2p.tools.async_service import background_trio_service
from libp2p.relay.circuit_v2.discovery import RelayInfo
from libp2p.utils import get_wildcard_address
# Configure logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger("source_node")
async def run_source(relay_peer_id=None, destination_peer_id=None):
# Create a libp2p host
listen_addr = multiaddr.Multiaddr("/ip4/0.0.0.0/tcp/9002")
# Create a libp2p host - use wildcard address to listen on all interfaces
listen_addr = get_wildcard_address(9002)
host = new_host()
# Configure as a relay client
@ -428,7 +433,7 @@ Running the Example
Relay node multiaddr: /ip4/127.0.0.1/tcp/9000/p2p/QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx
==================================================
Listening on: [<Multiaddr /ip4/0.0.0.0/tcp/9000/p2p/QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx>]
Listening on: [<Multiaddr /ip4/127.0.0.1/tcp/9000/p2p/QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx>]
Protocol service started
Relay service started successfully
Relay limits: RelayLimits(duration=3600, data=10485760, max_circuit_conns=8, max_reservations=4)
@ -447,7 +452,7 @@ Running the Example
Use this ID in the source node: QmPBr38KeQG2ibyL4fxq6yJWpfoVNCqJMHBdNyn1Qe4h5s
==================================================
Listening on: [<Multiaddr /ip4/0.0.0.0/tcp/9001/p2p/QmPBr38KeQG2ibyL4fxq6yJWpfoVNCqJMHBdNyn1Qe4h5s>]
Listening on: [<Multiaddr /ip4/127.0.0.1/tcp/9001/p2p/QmPBr38KeQG2ibyL4fxq6yJWpfoVNCqJMHBdNyn1Qe4h5s>]
Registered echo protocol handler
Protocol service started
Transport created
@ -469,7 +474,7 @@ Running the Example
$ python source_node.py
Source node started with ID: QmPyM56cgmFoHTgvMgGfDWRdVRQznmxCDDDg2dJ8ygVXj3
Listening on: [<Multiaddr /ip4/0.0.0.0/tcp/9002/p2p/QmPyM56cgmFoHTgvMgGfDWRdVRQznmxCDDDg2dJ8ygVXj3>]
Listening on: [<Multiaddr /ip4/127.0.0.1/tcp/9002/p2p/QmPyM56cgmFoHTgvMgGfDWRdVRQznmxCDDDg2dJ8ygVXj3>]
Protocol service started
No relay peer ID provided. Please enter the relay\'s peer ID:
Enter relay peer ID: QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx

View File

@ -0,0 +1,43 @@
QUIC Echo Demo
==============
This example demonstrates a simple ``echo`` protocol using **QUIC transport**.
QUIC provides built-in TLS security and stream multiplexing over UDP, making it an excellent transport choice for libp2p applications.
.. code-block:: console
$ python -m pip install libp2p
Collecting libp2p
...
Successfully installed libp2p-x.x.x
$ echo-quic-demo
Run this from the same folder in another console:
echo-quic-demo -d /ip4/127.0.0.1/udp/8000/quic-v1/p2p/16Uiu2HAmAsbxRR1HiGJRNVPQLNMeNsBCsXT3rDjoYBQzgzNpM5mJ
Waiting for incoming connection...
Copy the line that starts with ``echo-quic-demo -p 8001``, open a new terminal in the same
folder and paste it in:
.. code-block:: console
$ echo-quic-demo -d /ip4/127.0.0.1/udp/8000/quic-v1/p2p/16Uiu2HAmE3N7KauPTmHddYPsbMcBp2C6XAmprELX3YcFEN9iXiBu
I am 16Uiu2HAmE3N7KauPTmHddYPsbMcBp2C6XAmprELX3YcFEN9iXiBu
STARTING CLIENT CONNECTION PROCESS
CLIENT CONNECTED TO SERVER
Sent: hi, there!
Got: ECHO: hi, there!
**Key differences from TCP Echo:**
- Uses UDP instead of TCP: ``/udp/8000`` instead of ``/tcp/8000``
- Includes QUIC protocol identifier: ``/quic-v1`` in the multiaddr
- Built-in TLS security (no separate security transport needed)
- Native stream multiplexing over a single QUIC connection
.. literalinclude:: ../examples/echo/echo_quic.py
:language: python
:linenos:

View File

@ -12,7 +12,7 @@ This example demonstrates how to use the libp2p ``identify`` protocol.
$ identify-demo
First host listening. Run this from another console:
identify-demo -p 8889 -d /ip4/0.0.0.0/tcp/8888/p2p/QmUiN4R3fNrCoQugGgmmb3v35neMEjKFNrsbNGVDsRHWpM
identify-demo -p 8889 -d /ip4/127.0.0.1/tcp/8888/p2p/QmUiN4R3fNrCoQugGgmmb3v35neMEjKFNrsbNGVDsRHWpM
Waiting for incoming identify request...
@ -21,13 +21,13 @@ folder and paste it in:
.. code-block:: console
$ identify-demo -p 8889 -d /ip4/0.0.0.0/tcp/8888/p2p/QmUiN4R3fNrCoQugGgmmb3v35neMEjKFNrsbNGVDsRHWpM
dialer (host_b) listening on /ip4/0.0.0.0/tcp/8889
$ identify-demo -p 8889 -d /ip4/127.0.0.1/tcp/8888/p2p/QmUiN4R3fNrCoQugGgmmb3v35neMEjKFNrsbNGVDsRHWpM
dialer (host_b) listening on /ip4/127.0.0.1/tcp/8889
Second host connecting to peer: QmUiN4R3fNrCoQugGgmmb3v35neMEjKFNrsbNGVDsRHWpM
Starting identify protocol...
Identify response:
Public Key (Base64): CAASpgIwggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQDC6c/oNPP9X13NDQ3Xrlp3zOj+ErXIWb/A4JGwWchiDBwMhMslEX3ct8CqI0BqUYKuwdFjowqqopOJ3cS2MlqtGaiP6Dg9bvGqSDoD37BpNaRVNcebRxtB0nam9SQy3PYLbHAmz0vR4ToSiL9OLRORnGOxCtHBuR8ZZ5vS0JEni8eQMpNa7IuXwyStnuty/QjugOZudBNgYSr8+9gH722KTjput5IRL7BrpIdd4HNXGVRm4b9BjNowvHu404x3a/ifeNblpy/FbYyFJEW0looygKF7hpRHhRbRKIDZt2BqOfT1sFkbqsHE85oY859+VMzP61YELgvGwai2r7KcjkW/AgMBAAE=
Listen Addresses: ['/ip4/0.0.0.0/tcp/8888/p2p/QmUiN4R3fNrCoQugGgmmb3v35neMEjKFNrsbNGVDsRHWpM']
Listen Addresses: ['/ip4/127.0.0.1/tcp/8888/p2p/QmUiN4R3fNrCoQugGgmmb3v35neMEjKFNrsbNGVDsRHWpM']
Protocols: ['/ipfs/id/1.0.0', '/ipfs/ping/1.0.0']
Observed Address: ['/ip4/127.0.0.1/tcp/38082']
Protocol Version: ipfs/0.1.0

View File

@ -34,11 +34,11 @@ There is also a more interactive version of the example which runs as separate l
==== Starting Identify-Push Listener on port 8888 ====
Listener host ready!
Listening on: /ip4/0.0.0.0/tcp/8888/p2p/QmUiN4R3fNrCoQugGgmmb3v35neMEjKFNrsbNGVDsRHWpM
Listening on: /ip4/127.0.0.1/tcp/8888/p2p/QmUiN4R3fNrCoQugGgmmb3v35neMEjKFNrsbNGVDsRHWpM
Peer ID: QmUiN4R3fNrCoQugGgmmb3v35neMEjKFNrsbNGVDsRHWpM
Run dialer with command:
identify-push-listener-dialer-demo -d /ip4/0.0.0.0/tcp/8888/p2p/QmUiN4R3fNrCoQugGgmmb3v35neMEjKFNrsbNGVDsRHWpM
identify-push-listener-dialer-demo -d /ip4/127.0.0.1/tcp/8888/p2p/QmUiN4R3fNrCoQugGgmmb3v35neMEjKFNrsbNGVDsRHWpM
Waiting for incoming connections... (Ctrl+C to exit)
@ -47,12 +47,12 @@ folder and paste it in:
.. code-block:: console
$ identify-push-listener-dialer-demo -d /ip4/0.0.0.0/tcp/8888/p2p/QmUiN4R3fNrCoQugGgmmb3v35neMEjKFNrsbNGVDsRHWpM
$ identify-push-listener-dialer-demo -d /ip4/127.0.0.1/tcp/8888/p2p/QmUiN4R3fNrCoQugGgmmb3v35neMEjKFNrsbNGVDsRHWpM
==== Starting Identify-Push Dialer on port 8889 ====
Dialer host ready!
Listening on: /ip4/0.0.0.0/tcp/8889/p2p/QmZyXwVuTaBcDeRsSkJpOpWrSt
Listening on: /ip4/127.0.0.1/tcp/8889/p2p/QmZyXwVuTaBcDeRsSkJpOpWrSt
Connecting to peer: QmUiN4R3fNrCoQugGgmmb3v35neMEjKFNrsbNGVDsRHWpM
Successfully connected to listener!

View File

@ -15,7 +15,7 @@ This example demonstrates how to create a chat application using libp2p's PubSub
2025-04-06 23:59:17,471 - pubsub-demo - INFO - Your selected topic is: pubsub-chat
2025-04-06 23:59:17,472 - pubsub-demo - INFO - Using random available port: 33269
2025-04-06 23:59:17,490 - pubsub-demo - INFO - Node started with peer ID: QmcJnocH1d1tz3Zp4MotVDjNfNFawXHw2dpB9tMYGTXJp7
2025-04-06 23:59:17,490 - pubsub-demo - INFO - Listening on: /ip4/0.0.0.0/tcp/33269
2025-04-06 23:59:17,490 - pubsub-demo - INFO - Listening on: /ip4/127.0.0.1/tcp/33269
2025-04-06 23:59:17,490 - pubsub-demo - INFO - Initializing PubSub and GossipSub...
2025-04-06 23:59:17,491 - pubsub-demo - INFO - Pubsub and GossipSub services started.
2025-04-06 23:59:17,491 - pubsub-demo - INFO - Pubsub ready.
@ -35,7 +35,7 @@ Copy the line that starts with ``pubsub-demo -d``, open a new terminal and paste
2025-04-07 00:00:59,846 - pubsub-demo - INFO - Your selected topic is: pubsub-chat
2025-04-07 00:00:59,846 - pubsub-demo - INFO - Using random available port: 51977
2025-04-07 00:00:59,864 - pubsub-demo - INFO - Node started with peer ID: QmYQKCm95Ut1aXsjHmWVYqdaVbno1eKTYC8KbEVjqUaKaQ
2025-04-07 00:00:59,864 - pubsub-demo - INFO - Listening on: /ip4/0.0.0.0/tcp/51977
2025-04-07 00:00:59,864 - pubsub-demo - INFO - Listening on: /ip4/127.0.0.1/tcp/51977
2025-04-07 00:00:59,864 - pubsub-demo - INFO - Initializing PubSub and GossipSub...
2025-04-07 00:00:59,864 - pubsub-demo - INFO - Pubsub and GossipSub services started.
2025-04-07 00:00:59,865 - pubsub-demo - INFO - Pubsub ready.

View File

@ -23,7 +23,7 @@ The Random Walk implementation performs the following key operations:
2025-08-12 19:51:25,424 - random-walk-example - INFO - Mode: server, Port: 0 Demo interval: 30s
2025-08-12 19:51:25,426 - random-walk-example - INFO - Starting server node on port 45123
2025-08-12 19:51:25,426 - random-walk-example - INFO - Node peer ID: 16Uiu2HAm7EsNv5vvjPAehGAVfChjYjD63ZHyWogQRdzntSbAg9ef
2025-08-12 19:51:25,426 - random-walk-example - INFO - Node address: /ip4/0.0.0.0/tcp/45123/p2p/16Uiu2HAm7EsNv5vvjPAehGAVfChjYjD63ZHyWogQRdzntSbAg9ef
2025-08-12 19:51:25,426 - random-walk-example - INFO - Node address: /ip4/127.0.0.1/tcp/45123/p2p/16Uiu2HAm7EsNv5vvjPAehGAVfChjYjD63ZHyWogQRdzntSbAg9ef
2025-08-12 19:51:25,427 - random-walk-example - INFO - Initial routing table size: 0
2025-08-12 19:51:25,427 - random-walk-example - INFO - DHT service started in SERVER mode
2025-08-12 19:51:25,430 - libp2p.discovery.random_walk.rt_refresh_manager - INFO - RT Refresh Manager started

View File

@ -9,6 +9,7 @@ Examples
examples.identify_push
examples.chat
examples.echo
examples.echo_quic
examples.ping
examples.pubsub
examples.circuit_relay

View File

@ -28,6 +28,11 @@ For Python, the most common transport is TCP. Here's how to set up a basic TCP t
.. literalinclude:: ../examples/doc-examples/example_transport.py
:language: python
Also, QUIC is a modern transport protocol that provides built-in TLS security and stream multiplexing over UDP:
.. literalinclude:: ../examples/doc-examples/example_quic_transport.py
:language: python
Connection Encryption
^^^^^^^^^^^^^^^^^^^^^

View File

@ -0,0 +1,77 @@
libp2p.transport.quic package
=============================
Submodules
----------
libp2p.transport.quic.config module
-----------------------------------
.. automodule:: libp2p.transport.quic.config
:members:
:undoc-members:
:show-inheritance:
libp2p.transport.quic.connection module
---------------------------------------
.. automodule:: libp2p.transport.quic.connection
:members:
:undoc-members:
:show-inheritance:
libp2p.transport.quic.exceptions module
---------------------------------------
.. automodule:: libp2p.transport.quic.exceptions
:members:
:undoc-members:
:show-inheritance:
libp2p.transport.quic.listener module
-------------------------------------
.. automodule:: libp2p.transport.quic.listener
:members:
:undoc-members:
:show-inheritance:
libp2p.transport.quic.security module
-------------------------------------
.. automodule:: libp2p.transport.quic.security
:members:
:undoc-members:
:show-inheritance:
libp2p.transport.quic.stream module
-----------------------------------
.. automodule:: libp2p.transport.quic.stream
:members:
:undoc-members:
:show-inheritance:
libp2p.transport.quic.transport module
--------------------------------------
.. automodule:: libp2p.transport.quic.transport
:members:
:undoc-members:
:show-inheritance:
libp2p.transport.quic.utils module
----------------------------------
.. automodule:: libp2p.transport.quic.utils
:members:
:undoc-members:
:show-inheritance:
Module contents
---------------
.. automodule:: libp2p.transport.quic
:members:
:undoc-members:
:show-inheritance:

View File

@ -9,6 +9,11 @@ Subpackages
libp2p.transport.tcp
.. toctree::
:maxdepth: 4
libp2p.transport.quic
Submodules
----------

View File

@ -14,11 +14,26 @@ try:
expand_wildcard_address,
get_available_interfaces,
get_optimal_binding_address,
get_wildcard_address,
)
except ImportError:
# Fallbacks if utilities are missing
# Fallbacks if utilities are missing - use minimal network discovery
import socket
def get_available_interfaces(port: int, protocol: str = "tcp"):
return [Multiaddr(f"/ip4/0.0.0.0/{protocol}/{port}")]
# Try to get local network interfaces, fallback to loopback
addrs = []
try:
# Get hostname IP (better than hardcoded localhost)
hostname = socket.gethostname()
local_ip = socket.gethostbyname(hostname)
if local_ip != "127.0.0.1":
addrs.append(Multiaddr(f"/ip4/{local_ip}/{protocol}/{port}"))
except Exception:
pass
# Always include loopback as fallback
addrs.append(Multiaddr(f"/ip4/127.0.0.1/{protocol}/{port}"))
return addrs
def expand_wildcard_address(addr: Multiaddr, port: int | None = None):
if port is None:
@ -27,6 +42,15 @@ except ImportError:
return [Multiaddr(addr_str + f"/{port}")]
def get_optimal_binding_address(port: int, protocol: str = "tcp"):
# Try to get a non-loopback address first
interfaces = get_available_interfaces(port, protocol)
for addr in interfaces:
if "127.0.0.1" not in str(addr):
return addr
# Fallback to loopback if no other interfaces found
return Multiaddr(f"/ip4/127.0.0.1/{protocol}/{port}")
def get_wildcard_address(port: int, protocol: str = "tcp"):
return Multiaddr(f"/ip4/0.0.0.0/{protocol}/{port}")
@ -37,7 +61,10 @@ def main() -> None:
for a in interfaces:
print(f" - {a}")
wildcard_v4 = Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
# Demonstrate wildcard address as a feature
wildcard_v4 = get_wildcard_address(port)
print(f"\nWildcard address (feature): {wildcard_v4}")
expanded_v4 = expand_wildcard_address(wildcard_v4)
print("\nExpanded IPv4 wildcard:")
for a in expanded_v4:

View File

@ -2,7 +2,6 @@ import argparse
import logging
import secrets
import multiaddr
import trio
from libp2p import new_host
@ -54,18 +53,26 @@ BOOTSTRAP_PEERS = [
async def run(port: int, bootstrap_addrs: list[str]) -> None:
"""Run the bootstrap discovery example."""
from libp2p.utils.address_validation import (
find_free_port,
get_available_interfaces,
get_optimal_binding_address,
)
if port <= 0:
port = find_free_port()
# Generate key pair
secret = secrets.token_bytes(32)
key_pair = create_new_key_pair(secret)
# Create listen address
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
# Create listen addresses for all available interfaces
listen_addrs = get_available_interfaces(port)
# Register peer discovery handler
peerDiscovery.register_peer_discovered_handler(on_peer_discovery)
logger.info("🚀 Starting Bootstrap Discovery Example")
logger.info(f"📍 Listening on: {listen_addr}")
logger.info(f"🌐 Bootstrap peers: {len(bootstrap_addrs)}")
print("\n" + "=" * 60)
@ -80,7 +87,22 @@ async def run(port: int, bootstrap_addrs: list[str]) -> None:
host = new_host(key_pair=key_pair, bootstrap=bootstrap_addrs)
try:
async with host.run(listen_addrs=[listen_addr]):
async with host.run(listen_addrs=listen_addrs):
# Get all available addresses with peer ID
all_addrs = host.get_addrs()
logger.info("Listener ready, listening on:")
print("Listener ready, listening on:")
for addr in all_addrs:
logger.info(f"{addr}")
print(f"{addr}")
# Display optimal address for reference
optimal_addr = get_optimal_binding_address(port)
optimal_addr_with_peer = f"{optimal_addr}/p2p/{host.get_id().to_string()}"
logger.info(f"Optimal address: {optimal_addr_with_peer}")
print(f"Optimal address: {optimal_addr_with_peer}")
# Keep running and log peer discovery events
await trio.sleep_forever()
except KeyboardInterrupt:
@ -98,7 +120,7 @@ def main() -> None:
Usage:
python bootstrap.py -p 8000
python bootstrap.py -p 8001 --custom-bootstrap \\
"/ip4/127.0.0.1/tcp/8000/p2p/QmYourPeerID"
"/ip4/[HOST_IP]/tcp/8000/p2p/QmYourPeerID"
"""
parser = argparse.ArgumentParser(

View File

@ -1,4 +1,5 @@
import argparse
import logging
import sys
import multiaddr
@ -17,6 +18,11 @@ from libp2p.peer.peerinfo import (
info_from_p2p_addr,
)
# Configure minimal logging
logging.basicConfig(level=logging.WARNING)
logging.getLogger("multiaddr").setLevel(logging.WARNING)
logging.getLogger("libp2p").setLevel(logging.WARNING)
PROTOCOL_ID = TProtocol("/chat/1.0.0")
MAX_READ_LEN = 2**32 - 1
@ -40,9 +46,18 @@ async def write_data(stream: INetStream) -> None:
async def run(port: int, destination: str) -> None:
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
from libp2p.utils.address_validation import (
find_free_port,
get_available_interfaces,
get_optimal_binding_address,
)
if port <= 0:
port = find_free_port()
listen_addrs = get_available_interfaces(port)
host = new_host()
async with host.run(listen_addrs=[listen_addr]), trio.open_nursery() as nursery:
async with host.run(listen_addrs=listen_addrs), trio.open_nursery() as nursery:
# Start the peer-store cleanup task
nursery.start_soon(host.get_peerstore().start_cleanup_task, 60)
@ -54,10 +69,19 @@ async def run(port: int, destination: str) -> None:
host.set_stream_handler(PROTOCOL_ID, stream_handler)
# Get all available addresses with peer ID
all_addrs = host.get_addrs()
print("Listener ready, listening on:\n")
for addr in all_addrs:
print(f"{addr}")
# Use optimal address for the client command
optimal_addr = get_optimal_binding_address(port)
optimal_addr_with_peer = f"{optimal_addr}/p2p/{host.get_id().to_string()}"
print(
"Run this from the same folder in another console:\n\n"
f"chat-demo "
f"-d {host.get_addrs()[0]}\n"
f"\nRun this from the same folder in another console:\n\n"
f"chat-demo -d {optimal_addr_with_peer}\n"
)
print("Waiting for incoming connection...")
@ -86,7 +110,7 @@ def main() -> None:
where <DESTINATION> is the multiaddress of the previous listener host.
"""
example_maddr = (
"/ip4/127.0.0.1/tcp/8000/p2p/QmQn4SwGkDZKkUEpBRBvTmheQycxAHJUNmVEnjA2v1qe8Q"
"/ip4/[HOST_IP]/tcp/8000/p2p/QmQn4SwGkDZKkUEpBRBvTmheQycxAHJUNmVEnjA2v1qe8Q"
)
parser = argparse.ArgumentParser(description=description)
parser.add_argument("-p", "--port", default=0, type=int, help="source port number")

View File

@ -1,6 +1,5 @@
import secrets
import multiaddr
import trio
from libp2p import (
@ -9,9 +8,10 @@ from libp2p import (
from libp2p.crypto.secp256k1 import (
create_new_key_pair,
)
from libp2p.security.insecure.transport import (
PLAINTEXT_PROTOCOL_ID,
InsecureTransport,
from libp2p.security.insecure.transport import PLAINTEXT_PROTOCOL_ID, InsecureTransport
from libp2p.utils.address_validation import (
get_available_interfaces,
get_optimal_binding_address,
)
@ -38,17 +38,19 @@ async def main():
# Create a host with the key pair and insecure transport
host = new_host(key_pair=key_pair, sec_opt=security_options)
# Configure the listening address
# Configure the listening address using the new paradigm
port = 8000
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
listen_addrs = get_available_interfaces(port)
optimal_addr = get_optimal_binding_address(port)
# Start the host
async with host.run(listen_addrs=[listen_addr]):
async with host.run(listen_addrs=listen_addrs):
print(
"libp2p has started with insecure transport "
"(not recommended for production)"
)
print("libp2p is listening on:", host.get_addrs())
print(f"Optimal address: {optimal_addr}")
# Keep the host running
await trio.sleep_forever()

View File

@ -1,6 +1,5 @@
import secrets
import multiaddr
import trio
from libp2p import (
@ -13,6 +12,10 @@ from libp2p.security.noise.transport import (
PROTOCOL_ID as NOISE_PROTOCOL_ID,
Transport as NoiseTransport,
)
from libp2p.utils.address_validation import (
get_available_interfaces,
get_optimal_binding_address,
)
async def main():
@ -39,14 +42,16 @@ async def main():
# Create a host with the key pair and Noise security
host = new_host(key_pair=key_pair, sec_opt=security_options)
# Configure the listening address
# Configure the listening address using the new paradigm
port = 8000
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
listen_addrs = get_available_interfaces(port)
optimal_addr = get_optimal_binding_address(port)
# Start the host
async with host.run(listen_addrs=[listen_addr]):
async with host.run(listen_addrs=listen_addrs):
print("libp2p has started with Noise encryption")
print("libp2p is listening on:", host.get_addrs())
print(f"Optimal address: {optimal_addr}")
# Keep the host running
await trio.sleep_forever()

View File

@ -1,6 +1,5 @@
import secrets
import multiaddr
import trio
from libp2p import (
@ -13,6 +12,10 @@ from libp2p.security.secio.transport import (
ID as SECIO_PROTOCOL_ID,
Transport as SecioTransport,
)
from libp2p.utils.address_validation import (
get_available_interfaces,
get_optimal_binding_address,
)
async def main():
@ -32,14 +35,16 @@ async def main():
# Create a host with the key pair and SECIO security
host = new_host(key_pair=key_pair, sec_opt=security_options)
# Configure the listening address
# Configure the listening address using the new paradigm
port = 8000
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
listen_addrs = get_available_interfaces(port)
optimal_addr = get_optimal_binding_address(port)
# Start the host
async with host.run(listen_addrs=[listen_addr]):
async with host.run(listen_addrs=listen_addrs):
print("libp2p has started with SECIO encryption")
print("libp2p is listening on:", host.get_addrs())
print(f"Optimal address: {optimal_addr}")
# Keep the host running
await trio.sleep_forever()

View File

@ -1,6 +1,5 @@
import secrets
import multiaddr
import trio
from libp2p import (
@ -13,6 +12,10 @@ from libp2p.security.noise.transport import (
PROTOCOL_ID as NOISE_PROTOCOL_ID,
Transport as NoiseTransport,
)
from libp2p.utils.address_validation import (
get_available_interfaces,
get_optimal_binding_address,
)
async def main():
@ -39,14 +42,16 @@ async def main():
# Create a host with the key pair, Noise security, and mplex multiplexer
host = new_host(key_pair=key_pair, sec_opt=security_options)
# Configure the listening address
# Configure the listening address using the new paradigm
port = 8000
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
listen_addrs = get_available_interfaces(port)
optimal_addr = get_optimal_binding_address(port)
# Start the host
async with host.run(listen_addrs=[listen_addr]):
async with host.run(listen_addrs=listen_addrs):
print("libp2p has started with Noise encryption and mplex multiplexing")
print("libp2p is listening on:", host.get_addrs())
print(f"Optimal address: {optimal_addr}")
# Keep the host running
await trio.sleep_forever()

View File

@ -38,6 +38,10 @@ from libp2p.network.stream.net_stream import (
from libp2p.peer.peerinfo import (
info_from_p2p_addr,
)
from libp2p.utils.address_validation import (
get_available_interfaces,
get_optimal_binding_address,
)
PROTOCOL_ID = TProtocol("/echo/1.0.0")
@ -173,7 +177,9 @@ async def run_enhanced_demo(
"""
Run enhanced echo demo with NetStream state management.
"""
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
# Use the new address paradigm
listen_addrs = get_available_interfaces(port)
optimal_addr = get_optimal_binding_address(port)
# Generate or use provided key
if seed:
@ -185,7 +191,7 @@ async def run_enhanced_demo(
host = new_host(key_pair=create_new_key_pair(secret))
async with host.run(listen_addrs=[listen_addr]):
async with host.run(listen_addrs=listen_addrs):
print(f"Host ID: {host.get_id().to_string()}")
print("=" * 60)
@ -196,10 +202,12 @@ async def run_enhanced_demo(
# type: ignore: Stream is type of NetStream
host.set_stream_handler(PROTOCOL_ID, enhanced_echo_handler)
# Use optimal address for client command
optimal_addr_with_peer = f"{optimal_addr}/p2p/{host.get_id().to_string()}"
print(
"Run client from another console:\n"
f"python3 example_net_stream.py "
f"-d {host.get_addrs()[0]}\n"
f"-d {optimal_addr_with_peer}\n"
)
print("Waiting for connections...")
print("Press Ctrl+C to stop server")
@ -226,7 +234,7 @@ async def run_enhanced_demo(
def main() -> None:
example_maddr = (
"/ip4/127.0.0.1/tcp/8000/p2p/QmQn4SwGkDZKkUEpBRBvTmheQycxAHJUNmVEnjA2v1qe8Q"
"/ip4/[HOST_IP]/tcp/8000/p2p/QmQn4SwGkDZKkUEpBRBvTmheQycxAHJUNmVEnjA2v1qe8Q"
)
parser = argparse.ArgumentParser(

View File

@ -1,6 +1,6 @@
import secrets
import multiaddr
from multiaddr import Multiaddr
import trio
from libp2p import (
@ -16,6 +16,10 @@ from libp2p.security.noise.transport import (
PROTOCOL_ID as NOISE_PROTOCOL_ID,
Transport as NoiseTransport,
)
from libp2p.utils.address_validation import (
get_available_interfaces,
get_optimal_binding_address,
)
async def main():
@ -42,14 +46,16 @@ async def main():
# Create a host with the key pair, Noise security, and mplex multiplexer
host = new_host(key_pair=key_pair, sec_opt=security_options)
# Configure the listening address
# Configure the listening address using the new paradigm
port = 8000
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
listen_addrs = get_available_interfaces(port)
optimal_addr = get_optimal_binding_address(port)
# Start the host
async with host.run(listen_addrs=[listen_addr]):
async with host.run(listen_addrs=listen_addrs):
print("libp2p has started")
print("libp2p is listening on:", host.get_addrs())
print(f"Optimal address: {optimal_addr}")
# Connect to bootstrap peers manually
bootstrap_list = [
@ -61,7 +67,7 @@ async def main():
for addr in bootstrap_list:
try:
peer_info = info_from_p2p_addr(multiaddr.Multiaddr(addr))
peer_info = info_from_p2p_addr(Multiaddr(addr))
await host.connect(peer_info)
print(f"Connected to {peer_info.peer_id.to_string()}")
except Exception as e:

View File

@ -0,0 +1,49 @@
import secrets
import trio
from libp2p import (
new_host,
)
from libp2p.crypto.secp256k1 import (
create_new_key_pair,
)
from libp2p.utils.address_validation import (
get_available_interfaces,
get_optimal_binding_address,
)
async def main():
# Create a key pair for the host
secret = secrets.token_bytes(32)
key_pair = create_new_key_pair(secret)
# Create a host with the key pair
host = new_host(key_pair=key_pair, enable_quic=True)
# Configure the listening address using the new paradigm
port = 8000
listen_addrs = get_available_interfaces(port, protocol="udp")
# Convert TCP addresses to QUIC-v1 addresses
quic_addrs = []
for addr in listen_addrs:
addr_str = str(addr).replace("/tcp/", "/udp/") + "/quic-v1"
from multiaddr import Multiaddr
quic_addrs.append(Multiaddr(addr_str))
optimal_addr = get_optimal_binding_address(port, protocol="udp")
optimal_quic_str = str(optimal_addr).replace("/tcp/", "/udp/") + "/quic-v1"
# Start the host
async with host.run(listen_addrs=quic_addrs):
print("libp2p has started with QUIC transport")
print("libp2p is listening on:", host.get_addrs())
print(f"Optimal address: {optimal_quic_str}")
# Keep the host running
await trio.sleep_forever()
# Run the async function
trio.run(main)

View File

@ -1,6 +1,5 @@
import secrets
import multiaddr
import trio
from libp2p import (
@ -13,6 +12,10 @@ from libp2p.security.noise.transport import (
PROTOCOL_ID as NOISE_PROTOCOL_ID,
Transport as NoiseTransport,
)
from libp2p.utils.address_validation import (
get_available_interfaces,
get_optimal_binding_address,
)
async def main():
@ -39,14 +42,16 @@ async def main():
# Create a host with the key pair, Noise security, and mplex multiplexer
host = new_host(key_pair=key_pair, sec_opt=security_options)
# Configure the listening address
# Configure the listening address using the new paradigm
port = 8000
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
listen_addrs = get_available_interfaces(port)
optimal_addr = get_optimal_binding_address(port)
# Start the host
async with host.run(listen_addrs=[listen_addr]):
async with host.run(listen_addrs=listen_addrs):
print("libp2p has started")
print("libp2p is listening on:", host.get_addrs())
print(f"Optimal address: {optimal_addr}")
# Keep the host running
await trio.sleep_forever()

View File

@ -1,6 +1,5 @@
import secrets
import multiaddr
import trio
from libp2p import (
@ -9,6 +8,10 @@ from libp2p import (
from libp2p.crypto.secp256k1 import (
create_new_key_pair,
)
from libp2p.utils.address_validation import (
get_available_interfaces,
get_optimal_binding_address,
)
async def main():
@ -19,14 +22,16 @@ async def main():
# Create a host with the key pair
host = new_host(key_pair=key_pair)
# Configure the listening address
# Configure the listening address using the new paradigm
port = 8000
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
listen_addrs = get_available_interfaces(port)
optimal_addr = get_optimal_binding_address(port)
# Start the host
async with host.run(listen_addrs=[listen_addr]):
async with host.run(listen_addrs=listen_addrs):
print("libp2p has started with TCP transport")
print("libp2p is listening on:", host.get_addrs())
print(f"Optimal address: {optimal_addr}")
# Keep the host running
await trio.sleep_forever()

View File

@ -7,6 +7,7 @@ This example shows how to:
2. Use different load balancing strategies
3. Access multiple connections through the new API
4. Maintain backward compatibility
5. Use the new address paradigm for network configuration
"""
import logging
@ -15,6 +16,7 @@ import trio
from libp2p import new_swarm
from libp2p.network.swarm import ConnectionConfig, RetryConfig
from libp2p.utils import get_available_interfaces, get_optimal_binding_address
# Set up logging
logging.basicConfig(level=logging.INFO)
@ -103,10 +105,45 @@ async def example_backward_compatibility() -> None:
logger.info("Backward compatibility example completed")
async def example_network_address_paradigm() -> None:
"""Example of using the new address paradigm with multiple connections."""
logger.info("Demonstrating network address paradigm...")
# Get available interfaces using the new paradigm
port = 8000 # Example port
available_interfaces = get_available_interfaces(port)
logger.info(f"Available interfaces: {available_interfaces}")
# Get optimal binding address
optimal_address = get_optimal_binding_address(port)
logger.info(f"Optimal binding address: {optimal_address}")
# Create connection config for multiple connections with network awareness
connection_config = ConnectionConfig(
max_connections_per_peer=3, load_balancing_strategy="round_robin"
)
# Create swarm with address paradigm
swarm = new_swarm(connection_config=connection_config)
logger.info("Network address paradigm features:")
logger.info(" - get_available_interfaces() for interface discovery")
logger.info(" - get_optimal_binding_address() for smart address selection")
logger.info(" - Multiple connections with proper network binding")
await swarm.close()
logger.info("Network address paradigm example completed")
async def example_production_ready_config() -> None:
"""Example of production-ready configuration."""
logger.info("Creating swarm with production-ready configuration...")
# Get optimal network configuration using the new paradigm
port = 8001 # Example port
optimal_address = get_optimal_binding_address(port)
logger.info(f"Using optimal binding address: {optimal_address}")
# Production-ready retry configuration
retry_config = RetryConfig(
max_retries=3, # Reasonable retry limit
@ -156,6 +193,9 @@ async def main() -> None:
await example_backward_compatibility()
logger.info("-" * 30)
await example_network_address_paradigm()
logger.info("-" * 30)
await example_production_ready_config()
logger.info("-" * 30)

View File

@ -1,4 +1,5 @@
import argparse
import logging
import random
import secrets
@ -26,8 +27,14 @@ from libp2p.peer.peerinfo import (
from libp2p.utils.address_validation import (
find_free_port,
get_available_interfaces,
get_optimal_binding_address,
)
# Configure minimal logging
logging.basicConfig(level=logging.WARNING)
logging.getLogger("multiaddr").setLevel(logging.WARNING)
logging.getLogger("libp2p").setLevel(logging.WARNING)
PROTOCOL_ID = TProtocol("/echo/1.0.0")
MAX_READ_LEN = 2**32 - 1
@ -76,9 +83,13 @@ async def run(port: int, destination: str, seed: int | None = None) -> None:
for addr in listen_addr:
print(f"{addr}/p2p/{peer_id}")
# Get optimal address for display
optimal_addr = get_optimal_binding_address(port)
optimal_addr_with_peer = f"{optimal_addr}/p2p/{peer_id}"
print(
"\nRun this from the same folder in another console:\n\n"
f"echo-demo -d {host.get_addrs()[0]}\n"
f"echo-demo -d {optimal_addr_with_peer}\n"
)
print("Waiting for incoming connections...")
await trio.sleep_forever()
@ -114,7 +125,7 @@ def main() -> None:
where <DESTINATION> is the multiaddress of the previous listener host.
"""
example_maddr = (
"/ip4/127.0.0.1/tcp/8000/p2p/QmQn4SwGkDZKkUEpBRBvTmheQycxAHJUNmVEnjA2v1qe8Q"
"/ip4/[HOST_IP]/tcp/8000/p2p/QmQn4SwGkDZKkUEpBRBvTmheQycxAHJUNmVEnjA2v1qe8Q"
)
parser = argparse.ArgumentParser(description=description)
parser.add_argument("-p", "--port", default=0, type=int, help="source port number")

207
examples/echo/echo_quic.py Normal file
View File

@ -0,0 +1,207 @@
#!/usr/bin/env python3
"""
QUIC Echo Example - Fixed version with proper client/server separation
This program demonstrates a simple echo protocol using QUIC transport where a peer
listens for connections and copies back any input received on a stream.
Fixed to properly separate client and server modes - clients don't start listeners.
"""
import argparse
import logging
from multiaddr import Multiaddr
import trio
from libp2p import new_host
from libp2p.crypto.secp256k1 import create_new_key_pair
from libp2p.custom_types import TProtocol
from libp2p.network.stream.net_stream import INetStream
from libp2p.peer.peerinfo import info_from_p2p_addr
# Configure minimal logging
logging.basicConfig(level=logging.WARNING)
logging.getLogger("multiaddr").setLevel(logging.WARNING)
logging.getLogger("libp2p").setLevel(logging.WARNING)
PROTOCOL_ID = TProtocol("/echo/1.0.0")
async def _echo_stream_handler(stream: INetStream) -> None:
try:
msg = await stream.read()
await stream.write(msg)
await stream.close()
except Exception as e:
print(f"Echo handler error: {e}")
try:
await stream.close()
except: # noqa: E722
pass
async def run_server(port: int, seed: int | None = None) -> None:
"""Run echo server with QUIC transport."""
from libp2p.utils.address_validation import (
find_free_port,
get_available_interfaces,
get_optimal_binding_address,
)
if port <= 0:
port = find_free_port()
# For QUIC, we need UDP addresses - use the new address paradigm
tcp_addrs = get_available_interfaces(port)
# Convert TCP addresses to QUIC addresses
quic_addrs = []
for addr in tcp_addrs:
addr_str = str(addr).replace("/tcp/", "/udp/") + "/quic"
quic_addrs.append(Multiaddr(addr_str))
if seed:
import random
random.seed(seed)
secret_number = random.getrandbits(32 * 8)
secret = secret_number.to_bytes(length=32, byteorder="big")
else:
import secrets
secret = secrets.token_bytes(32)
# Create host with QUIC transport
host = new_host(
enable_quic=True,
key_pair=create_new_key_pair(secret),
)
# Server mode: start listener
async with host.run(listen_addrs=quic_addrs):
try:
print(f"I am {host.get_id().to_string()}")
host.set_stream_handler(PROTOCOL_ID, _echo_stream_handler)
# Get all available addresses with peer ID
all_addrs = host.get_addrs()
print("Listener ready, listening on:")
for addr in all_addrs:
print(f"{addr}")
# Use optimal address for the client command
optimal_tcp = get_optimal_binding_address(port)
optimal_quic_str = str(optimal_tcp).replace("/tcp/", "/udp/") + "/quic"
peer_id = host.get_id().to_string()
optimal_quic_with_peer = f"{optimal_quic_str}/p2p/{peer_id}"
print(
f"\nRun this from the same folder in another console:\n\n"
f"python3 ./examples/echo/echo_quic.py -d {optimal_quic_with_peer}\n"
)
print("Waiting for incoming QUIC connections...")
await trio.sleep_forever()
except KeyboardInterrupt:
print("Closing server gracefully...")
await host.close()
return
async def run_client(destination: str, seed: int | None = None) -> None:
"""Run echo client with QUIC transport."""
if seed:
import random
random.seed(seed)
secret_number = random.getrandbits(32 * 8)
secret = secret_number.to_bytes(length=32, byteorder="big")
else:
import secrets
secret = secrets.token_bytes(32)
# Create host with QUIC transport
host = new_host(
enable_quic=True,
key_pair=create_new_key_pair(secret),
)
# Client mode: NO listener, just connect
async with host.run(listen_addrs=[]): # Empty listen_addrs for client
print(f"I am {host.get_id().to_string()}")
maddr = Multiaddr(destination)
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])
msg = b"hi, there!\n"
await stream.write(msg)
response = await stream.read()
print(f"Sent: {msg.decode('utf-8')}")
print(f"Got: {response.decode('utf-8')}")
await stream.close()
await host.disconnect(info.peer_id)
async def run(port: int, destination: str, seed: int | None = None) -> None:
"""
Run echo server or client with QUIC transport.
Fixed version that properly separates client and server modes.
"""
if not destination: # Server mode
await run_server(port, seed)
else: # Client mode
await run_client(destination, seed)
def main() -> None:
"""Main function - help text updated for QUIC."""
description = """
This program demonstrates a simple echo protocol using QUIC
transport where a peer listens for connections and copies back
any input received on a stream.
QUIC provides built-in TLS security and stream multiplexing over UDP.
To use it, first run 'echo-quic-demo -p <PORT>', where <PORT> is
the UDP port number. Then, run another host with ,
'echo-quic-demo -d <DESTINATION>'
where <DESTINATION> is the QUIC multiaddress of the previous listener host.
"""
example_maddr = "/ip4/[HOST_IP]/udp/8000/quic/p2p/QmQn4SwGkDZKkUEpBRBv"
parser = argparse.ArgumentParser(description=description)
parser.add_argument("-p", "--port", default=0, type=int, help="UDP port number")
parser.add_argument(
"-d",
"--destination",
type=str,
help=f"destination multiaddr string, e.g. {example_maddr}",
)
parser.add_argument(
"-s",
"--seed",
type=int,
help="provide a seed to the random number generator",
)
args = parser.parse_args()
try:
trio.run(run, args.port, args.destination, args.seed)
except KeyboardInterrupt:
pass
if __name__ == "__main__":
main()

View File

@ -20,6 +20,11 @@ from libp2p.peer.peerinfo import (
info_from_p2p_addr,
)
# Configure minimal logging
logging.basicConfig(level=logging.WARNING)
logging.getLogger("multiaddr").setLevel(logging.WARNING)
logging.getLogger("libp2p").setLevel(logging.WARNING)
logger = logging.getLogger("libp2p.identity.identify-example")
@ -58,11 +63,19 @@ def print_identify_response(identify_response: Identify):
async def run(port: int, destination: str, use_varint_format: bool = True) -> None:
localhost_ip = "0.0.0.0"
from libp2p.utils.address_validation import (
get_available_interfaces,
get_optimal_binding_address,
)
if not destination:
# Create first host (listener)
listen_addr = multiaddr.Multiaddr(f"/ip4/{localhost_ip}/tcp/{port}")
if port <= 0:
from libp2p.utils.address_validation import find_free_port
port = find_free_port()
listen_addrs = get_available_interfaces(port)
host_a = new_host()
# Set up identify handler with specified format
@ -73,25 +86,49 @@ async def run(port: int, destination: str, use_varint_format: bool = True) -> No
host_a.set_stream_handler(IDENTIFY_PROTOCOL_ID, identify_handler)
async with (
host_a.run(listen_addrs=[listen_addr]),
host_a.run(listen_addrs=listen_addrs),
trio.open_nursery() as nursery,
):
# Start the peer-store cleanup task
nursery.start_soon(host_a.get_peerstore().start_cleanup_task, 60)
# Get the actual address and replace 0.0.0.0 with 127.0.0.1 for client
# connections
server_addr = str(host_a.get_addrs()[0])
client_addr = server_addr.replace("/ip4/0.0.0.0/", "/ip4/127.0.0.1/")
# Get all available addresses with peer ID
all_addrs = host_a.get_addrs()
format_name = "length-prefixed" if use_varint_format else "raw protobuf"
format_flag = "--raw-format" if not use_varint_format else ""
print(
f"First host listening (using {format_name} format). "
f"Run this from another console:\n\n"
f"identify-demo {format_flag} -d {client_addr}\n"
)
print("Waiting for incoming identify request...")
if use_varint_format:
format_name = "length-prefixed"
print(f"First host listening (using {format_name} format).")
print("Listener ready, listening on:\n")
for addr in all_addrs:
print(f"{addr}")
# Use optimal address for the client command
optimal_addr = get_optimal_binding_address(port)
optimal_addr_with_peer = (
f"{optimal_addr}/p2p/{host_a.get_id().to_string()}"
)
print(
f"\nRun this from the same folder in another console:\n\n"
f"identify-demo -d {optimal_addr_with_peer}\n"
)
print("Waiting for incoming identify request...")
else:
format_name = "raw protobuf"
print(f"First host listening (using {format_name} format).")
print("Listener ready, listening on:\n")
for addr in all_addrs:
print(f"{addr}")
# Use optimal address for the client command
optimal_addr = get_optimal_binding_address(port)
optimal_addr_with_peer = (
f"{optimal_addr}/p2p/{host_a.get_id().to_string()}"
)
print(
f"\nRun this from the same folder in another console:\n\n"
f"identify-demo -d {optimal_addr_with_peer}\n"
)
print("Waiting for incoming identify request...")
# Add a custom handler to show connection events
async def custom_identify_handler(stream):
@ -134,11 +171,20 @@ async def run(port: int, destination: str, use_varint_format: bool = True) -> No
else:
# Create second host (dialer)
listen_addr = multiaddr.Multiaddr(f"/ip4/{localhost_ip}/tcp/{port}")
from libp2p.utils.address_validation import (
find_free_port,
get_available_interfaces,
get_optimal_binding_address,
)
if port <= 0:
port = find_free_port()
listen_addrs = get_available_interfaces(port)
host_b = new_host()
async with (
host_b.run(listen_addrs=[listen_addr]),
host_b.run(listen_addrs=listen_addrs),
trio.open_nursery() as nursery,
):
# Start the peer-store cleanup task
@ -234,7 +280,7 @@ def main() -> None:
"""
example_maddr = (
"/ip4/127.0.0.1/tcp/8888/p2p/QmQn4SwGkDZKkUEpBRBvTmheQycxAHJUNmVEnjA2v1qe8Q"
"/ip4/[HOST_IP]/tcp/8888/p2p/QmQn4SwGkDZKkUEpBRBvTmheQycxAHJUNmVEnjA2v1qe8Q"
)
parser = argparse.ArgumentParser(description=description)
@ -258,7 +304,7 @@ def main() -> None:
# Determine format: use varint (length-prefixed) if --raw-format is specified,
# otherwise use raw protobuf format (old format)
use_varint_format = args.raw_format
use_varint_format = not args.raw_format
try:
if args.destination:

View File

@ -36,6 +36,9 @@ from libp2p.identity.identify_push import (
from libp2p.peer.peerinfo import (
info_from_p2p_addr,
)
from libp2p.utils.address_validation import (
get_available_interfaces,
)
# Configure logging
logger = logging.getLogger(__name__)
@ -207,13 +210,13 @@ async def main() -> None:
ID_PUSH, create_custom_identify_push_handler(host_2, "Host 2")
)
# Start listening on random ports using the run context manager
listen_addr_1 = multiaddr.Multiaddr("/ip4/127.0.0.1/tcp/0")
listen_addr_2 = multiaddr.Multiaddr("/ip4/127.0.0.1/tcp/0")
# Start listening on available interfaces using random ports
listen_addrs_1 = get_available_interfaces(0) # 0 for random port
listen_addrs_2 = get_available_interfaces(0) # 0 for random port
async with (
host_1.run([listen_addr_1]),
host_2.run([listen_addr_2]),
host_1.run(listen_addrs_1),
host_2.run(listen_addrs_2),
trio.open_nursery() as nursery,
):
# Start the peer-store cleanup task

View File

@ -14,7 +14,7 @@ Usage:
python identify_push_listener_dialer.py
# Then in another console, run as a dialer (default port 8889):
python identify_push_listener_dialer.py -d /ip4/127.0.0.1/tcp/8888/p2p/PEER_ID
python identify_push_listener_dialer.py -d /ip4/[HOST_IP]/tcp/8888/p2p/PEER_ID
(where PEER_ID is the peer ID displayed by the listener)
"""
@ -56,6 +56,11 @@ from libp2p.peer.peerinfo import (
info_from_p2p_addr,
)
# Configure minimal logging
logging.basicConfig(level=logging.WARNING)
logging.getLogger("multiaddr").setLevel(logging.WARNING)
logging.getLogger("libp2p").setLevel(logging.WARNING)
# Configure logging
logger = logging.getLogger("libp2p.identity.identify-push-example")
@ -194,6 +199,11 @@ async def run_listener(
port: int, use_varint_format: bool = True, raw_format_flag: bool = False
) -> None:
"""Run a host in listener mode."""
from libp2p.utils.address_validation import find_free_port, get_available_interfaces
if port <= 0:
port = find_free_port()
format_name = "length-prefixed" if use_varint_format else "raw protobuf"
print(
f"\n==== Starting Identify-Push Listener on port {port} "
@ -215,26 +225,33 @@ async def run_listener(
custom_identify_push_handler_for(host, use_varint_format=use_varint_format),
)
# Start listening
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
# Start listening on all available interfaces
listen_addrs = get_available_interfaces(port)
try:
async with host.run([listen_addr]):
addr = host.get_addrs()[0]
async with host.run(listen_addrs):
all_addrs = host.get_addrs()
logger.info("Listener host ready!")
print("Listener host ready!")
logger.info(f"Listening on: {addr}")
print(f"Listening on: {addr}")
logger.info("Listener ready, listening on:")
print("Listener ready, listening on:")
for addr in all_addrs:
logger.info(f"{addr}")
print(f"{addr}")
logger.info(f"Peer ID: {host.get_id().pretty()}")
print(f"Peer ID: {host.get_id().pretty()}")
print("\nRun dialer with command:")
# Use the first address as the default for the dialer command
default_addr = all_addrs[0]
print("\nRun this from the same folder in another console:")
if raw_format_flag:
print(f"identify-push-listener-dialer-demo -d {addr} --raw-format")
print(
f"identify-push-listener-dialer-demo -d {default_addr} --raw-format"
)
else:
print(f"identify-push-listener-dialer-demo -d {addr}")
print(f"identify-push-listener-dialer-demo -d {default_addr}")
print("\nWaiting for incoming identify/push requests... (Ctrl+C to exit)")
# Keep running until interrupted
@ -274,10 +291,12 @@ async def run_dialer(
identify_push_handler_for(host, use_varint_format=use_varint_format),
)
# Start listening on a different port
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
# Start listening on available interfaces
from libp2p.utils.address_validation import get_available_interfaces
async with host.run([listen_addr]):
listen_addrs = get_available_interfaces(port)
async with host.run(listen_addrs):
logger.info("Dialer host ready!")
print("Dialer host ready!")

View File

@ -150,26 +150,43 @@ async def run_node(
key_pair = create_new_key_pair(secrets.token_bytes(32))
host = new_host(key_pair=key_pair)
listen_addr = Multiaddr(f"/ip4/127.0.0.1/tcp/{port}")
async with host.run(listen_addrs=[listen_addr]), trio.open_nursery() as nursery:
from libp2p.utils.address_validation import (
get_available_interfaces,
get_optimal_binding_address,
)
listen_addrs = get_available_interfaces(port)
async with host.run(listen_addrs=listen_addrs), trio.open_nursery() as nursery:
# Start the peer-store cleanup task
nursery.start_soon(host.get_peerstore().start_cleanup_task, 60)
peer_id = host.get_id().pretty()
addr_str = f"/ip4/127.0.0.1/tcp/{port}/p2p/{peer_id}"
# Get all available addresses with peer ID
all_addrs = host.get_addrs()
logger.info("Listener ready, listening on:")
for addr in all_addrs:
logger.info(f"{addr}")
# Use optimal address for the bootstrap command
optimal_addr = get_optimal_binding_address(port)
optimal_addr_with_peer = f"{optimal_addr}/p2p/{host.get_id().to_string()}"
bootstrap_cmd = f"--bootstrap {optimal_addr_with_peer}"
logger.info("To connect to this node, use: %s", bootstrap_cmd)
await connect_to_bootstrap_nodes(host, bootstrap_nodes)
dht = KadDHT(host, dht_mode)
# take all peer ids from the host and add them to the dht
for peer_id in host.get_peerstore().peer_ids():
await dht.routing_table.add_peer(peer_id)
logger.info(f"Connected to bootstrap nodes: {host.get_connected_peers()}")
bootstrap_cmd = f"--bootstrap {addr_str}"
logger.info("To connect to this node, use: %s", bootstrap_cmd)
# Save server address in server mode
if dht_mode == DHTMode.SERVER:
save_server_addr(addr_str)
save_server_addr(str(optimal_addr_with_peer))
# Start the DHT service
async with background_trio_service(dht):

View File

@ -2,7 +2,6 @@ import argparse
import logging
import secrets
import multiaddr
import trio
from libp2p import (
@ -14,6 +13,11 @@ from libp2p.crypto.secp256k1 import (
)
from libp2p.discovery.events.peerDiscovery import peerDiscovery
# Configure minimal logging
logging.basicConfig(level=logging.WARNING)
logging.getLogger("multiaddr").setLevel(logging.WARNING)
logging.getLogger("libp2p").setLevel(logging.WARNING)
logger = logging.getLogger("libp2p.discovery.mdns")
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
@ -22,34 +26,43 @@ handler.setFormatter(
)
logger.addHandler(handler)
# Set root logger to DEBUG to capture all logs from dependencies
logging.getLogger().setLevel(logging.DEBUG)
def onPeerDiscovery(peerinfo: PeerInfo):
logger.info(f"Discovered: {peerinfo.peer_id}")
async def run(port: int) -> None:
from libp2p.utils.address_validation import find_free_port, get_available_interfaces
if port <= 0:
port = find_free_port()
secret = secrets.token_bytes(32)
key_pair = create_new_key_pair(secret)
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
listen_addrs = get_available_interfaces(port)
peerDiscovery.register_peer_discovered_handler(onPeerDiscovery)
print(
"Run this from the same folder in another console to "
"start another peer on a different port:\n\n"
"mdns-demo -p <ANOTHER_PORT>\n"
)
print("Waiting for mDNS peer discovery events...\n")
logger.info("Starting peer Discovery")
host = new_host(key_pair=key_pair, enable_mDNS=True)
async with host.run(listen_addrs=[listen_addr]), trio.open_nursery() as nursery:
async with host.run(listen_addrs=listen_addrs), trio.open_nursery() as nursery:
# Start the peer-store cleanup task
nursery.start_soon(host.get_peerstore().start_cleanup_task, 60)
# Get all available addresses with peer ID
all_addrs = host.get_addrs()
print("Listener ready, listening on:")
for addr in all_addrs:
print(f"{addr}")
print(
"\nRun this from the same folder in another console to "
"start another peer on a different port:\n\n"
"mdns-demo -p <ANOTHER_PORT>\n"
)
print("Waiting for mDNS peer discovery events...\n")
await trio.sleep_forever()

View File

@ -1,4 +1,5 @@
import argparse
import logging
import multiaddr
import trio
@ -16,6 +17,11 @@ from libp2p.peer.peerinfo import (
info_from_p2p_addr,
)
# Configure minimal logging
logging.basicConfig(level=logging.WARNING)
logging.getLogger("multiaddr").setLevel(logging.WARNING)
logging.getLogger("libp2p").setLevel(logging.WARNING)
PING_PROTOCOL_ID = TProtocol("/ipfs/ping/1.0.0")
PING_LENGTH = 32
RESP_TIMEOUT = 60
@ -55,20 +61,38 @@ async def send_ping(stream: INetStream) -> None:
async def run(port: int, destination: str) -> None:
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
host = new_host(listen_addrs=[listen_addr])
from libp2p.utils.address_validation import (
find_free_port,
get_available_interfaces,
get_optimal_binding_address,
)
async with host.run(listen_addrs=[listen_addr]), trio.open_nursery() as nursery:
if port <= 0:
port = find_free_port()
listen_addrs = get_available_interfaces(port)
host = new_host(listen_addrs=listen_addrs)
async with host.run(listen_addrs=listen_addrs), trio.open_nursery() as nursery:
# Start the peer-store cleanup task
nursery.start_soon(host.get_peerstore().start_cleanup_task, 60)
if not destination:
host.set_stream_handler(PING_PROTOCOL_ID, handle_ping)
# Get all available addresses with peer ID
all_addrs = host.get_addrs()
print("Listener ready, listening on:\n")
for addr in all_addrs:
print(f"{addr}")
# Use optimal address for the client command
optimal_addr = get_optimal_binding_address(port)
optimal_addr_with_peer = f"{optimal_addr}/p2p/{host.get_id().to_string()}"
print(
"Run this from the same folder in another console:\n\n"
f"ping-demo "
f"-d {host.get_addrs()[0]}\n"
f"\nRun this from the same folder in another console:\n\n"
f"ping-demo -d {optimal_addr_with_peer}\n"
)
print("Waiting for incoming connection...")
@ -94,7 +118,7 @@ def main() -> None:
"""
example_maddr = (
"/ip4/127.0.0.1/tcp/8000/p2p/QmQn4SwGkDZKkUEpBRBvTmheQycxAHJUNmVEnjA2v1qe8Q"
"/ip4/[HOST_IP]/tcp/8000/p2p/QmQn4SwGkDZKkUEpBRBvTmheQycxAHJUNmVEnjA2v1qe8Q"
)
parser = argparse.ArgumentParser(description=description)

View File

@ -102,14 +102,16 @@ async def monitor_peer_topics(pubsub, nursery, termination_event):
async def run(topic: str, destination: str | None, port: int | None) -> None:
# Initialize network settings
localhost_ip = "127.0.0.1"
from libp2p.utils.address_validation import (
get_available_interfaces,
get_optimal_binding_address,
)
if port is None or port == 0:
port = find_free_port()
logger.info(f"Using random available port: {port}")
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
listen_addrs = get_available_interfaces(port)
# Create a new libp2p host
host = new_host(
@ -138,12 +140,11 @@ async def run(topic: str, destination: str | None, port: int | None) -> None:
pubsub = Pubsub(host, gossipsub)
termination_event = trio.Event() # Event to signal termination
async with host.run(listen_addrs=[listen_addr]), trio.open_nursery() as nursery:
async with host.run(listen_addrs=listen_addrs), trio.open_nursery() as nursery:
# Start the peer-store cleanup task
nursery.start_soon(host.get_peerstore().start_cleanup_task, 60)
logger.info(f"Node started with peer ID: {host.get_id()}")
logger.info(f"Listening on: {listen_addr}")
logger.info("Initializing PubSub and GossipSub...")
async with background_trio_service(pubsub):
async with background_trio_service(gossipsub):
@ -157,10 +158,21 @@ async def run(topic: str, destination: str | None, port: int | None) -> None:
if not destination:
# Server mode
# Get all available addresses with peer ID
all_addrs = host.get_addrs()
logger.info("Listener ready, listening on:")
for addr in all_addrs:
logger.info(f"{addr}")
# Use optimal address for the client command
optimal_addr = get_optimal_binding_address(port)
optimal_addr_with_peer = (
f"{optimal_addr}/p2p/{host.get_id().to_string()}"
)
logger.info(
"Run this script in another console with:\n"
f"pubsub-demo "
f"-d /ip4/{localhost_ip}/tcp/{port}/p2p/{host.get_id()}\n"
f"\nRun this from the same folder in another console:\n\n"
f"pubsub-demo -d {optimal_addr_with_peer}\n"
)
logger.info("Waiting for peers...")
@ -182,11 +194,6 @@ async def run(topic: str, destination: str | None, port: int | None) -> None:
f"Connecting to peer: {info.peer_id} "
f"using protocols: {protocols_in_maddr}"
)
logger.info(
"Run this script in another console with:\n"
f"pubsub-demo "
f"-d /ip4/{localhost_ip}/tcp/{port}/p2p/{host.get_id()}\n"
)
try:
await host.connect(info)
logger.info(f"Connected to peer: {info.peer_id}")

View File

@ -16,7 +16,6 @@ import random
import secrets
import sys
from multiaddr import Multiaddr
import trio
from libp2p import new_host
@ -130,16 +129,24 @@ async def run_node(port: int, mode: str, demo_interval: int = 30) -> None:
# Create host and DHT
key_pair = create_new_key_pair(secrets.token_bytes(32))
host = new_host(key_pair=key_pair, bootstrap=DEFAULT_BOOTSTRAP_NODES)
listen_addr = Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
async with host.run(listen_addrs=[listen_addr]), trio.open_nursery() as nursery:
from libp2p.utils.address_validation import get_available_interfaces
listen_addrs = get_available_interfaces(port)
async with host.run(listen_addrs=listen_addrs), trio.open_nursery() as nursery:
# Start maintenance tasks
nursery.start_soon(host.get_peerstore().start_cleanup_task, 60)
nursery.start_soon(maintain_connections, host)
peer_id = host.get_id().pretty()
logger.info(f"Node peer ID: {peer_id}")
logger.info(f"Node address: /ip4/0.0.0.0/tcp/{port}/p2p/{peer_id}")
# Get all available addresses with peer ID
all_addrs = host.get_addrs()
logger.info("Listener ready, listening on:")
for addr in all_addrs:
logger.info(f"{addr}")
# Create and start DHT with Random Walk enabled
dht = KadDHT(host, dht_mode, enable_random_walk=True)

View File

@ -0,0 +1,446 @@
#!/usr/bin/env python3
"""
TCP P2P Data Transfer Test
This test proves that TCP peer-to-peer data transfer works correctly in libp2p.
This serves as a baseline to compare with WebSocket tests.
"""
import pytest
from multiaddr import Multiaddr
import trio
from libp2p import create_yamux_muxer_option, new_host
from libp2p.crypto.secp256k1 import create_new_key_pair
from libp2p.custom_types import TProtocol
from libp2p.peer.peerinfo import info_from_p2p_addr
from libp2p.security.insecure.transport import PLAINTEXT_PROTOCOL_ID, InsecureTransport
# Test protocol for data exchange
TCP_DATA_PROTOCOL = TProtocol("/test/tcp-data-exchange/1.0.0")
async def create_tcp_host_pair():
"""Create a pair of hosts configured for TCP communication."""
# Create key pairs
key_pair_a = create_new_key_pair()
key_pair_b = create_new_key_pair()
# Create security options (using plaintext for simplicity)
def security_options(kp):
return {
PLAINTEXT_PROTOCOL_ID: InsecureTransport(
local_key_pair=kp, secure_bytes_provider=None, peerstore=None
)
}
# Host A (listener) - TCP transport (default)
host_a = new_host(
key_pair=key_pair_a,
sec_opt=security_options(key_pair_a),
muxer_opt=create_yamux_muxer_option(),
listen_addrs=[Multiaddr("/ip4/127.0.0.1/tcp/0")],
)
# Host B (dialer) - TCP transport (default)
host_b = new_host(
key_pair=key_pair_b,
sec_opt=security_options(key_pair_b),
muxer_opt=create_yamux_muxer_option(),
listen_addrs=[Multiaddr("/ip4/127.0.0.1/tcp/0")],
)
return host_a, host_b
@pytest.mark.trio
async def test_tcp_basic_connection():
"""Test basic TCP connection establishment."""
host_a, host_b = await create_tcp_host_pair()
connection_established = False
async def connection_handler(stream):
nonlocal connection_established
connection_established = True
await stream.close()
host_a.set_stream_handler(TCP_DATA_PROTOCOL, connection_handler)
async with (
host_a.run(listen_addrs=[Multiaddr("/ip4/127.0.0.1/tcp/0")]),
host_b.run(listen_addrs=[]),
):
# Get host A's listen address
listen_addrs = host_a.get_addrs()
assert listen_addrs, "Host A should have listen addresses"
# Extract TCP address
tcp_addr = None
for addr in listen_addrs:
if "/tcp/" in str(addr) and "/ws" not in str(addr):
tcp_addr = addr
break
assert tcp_addr, f"No TCP address found in {listen_addrs}"
print(f"🔗 Host A listening on: {tcp_addr}")
# Create peer info for host A
peer_info = info_from_p2p_addr(tcp_addr)
# Host B connects to host A
await host_b.connect(peer_info)
print("✅ TCP connection established")
# Open a stream to test the connection
stream = await host_b.new_stream(peer_info.peer_id, [TCP_DATA_PROTOCOL])
await stream.close()
# Wait a bit for the handler to be called
await trio.sleep(0.1)
assert connection_established, "TCP connection handler should have been called"
print("✅ TCP basic connection test successful!")
@pytest.mark.trio
async def test_tcp_data_transfer():
"""Test TCP peer-to-peer data transfer."""
host_a, host_b = await create_tcp_host_pair()
# Test data
test_data = b"Hello TCP P2P Data Transfer! This is a test message."
received_data = None
transfer_complete = trio.Event()
async def data_handler(stream):
nonlocal received_data
try:
# Read the incoming data
received_data = await stream.read(len(test_data))
# Echo it back to confirm successful transfer
await stream.write(received_data)
await stream.close()
transfer_complete.set()
except Exception as e:
print(f"Handler error: {e}")
transfer_complete.set()
host_a.set_stream_handler(TCP_DATA_PROTOCOL, data_handler)
async with (
host_a.run(listen_addrs=[Multiaddr("/ip4/127.0.0.1/tcp/0")]),
host_b.run(listen_addrs=[]),
):
# Get host A's listen address
listen_addrs = host_a.get_addrs()
assert listen_addrs, "Host A should have listen addresses"
# Extract TCP address
tcp_addr = None
for addr in listen_addrs:
if "/tcp/" in str(addr) and "/ws" not in str(addr):
tcp_addr = addr
break
assert tcp_addr, f"No TCP address found in {listen_addrs}"
print(f"🔗 Host A listening on: {tcp_addr}")
# Create peer info for host A
peer_info = info_from_p2p_addr(tcp_addr)
# Host B connects to host A
await host_b.connect(peer_info)
print("✅ TCP connection established")
# Open a stream for data transfer
stream = await host_b.new_stream(peer_info.peer_id, [TCP_DATA_PROTOCOL])
print("✅ TCP stream opened")
# Send test data
await stream.write(test_data)
print(f"📤 Sent data: {test_data}")
# Read echoed data back
echoed_data = await stream.read(len(test_data))
print(f"📥 Received echo: {echoed_data}")
await stream.close()
# Wait for transfer to complete
with trio.fail_after(5.0): # 5 second timeout
await transfer_complete.wait()
# Verify data transfer
assert received_data == test_data, (
f"Data mismatch: {received_data} != {test_data}"
)
assert echoed_data == test_data, f"Echo mismatch: {echoed_data} != {test_data}"
print("✅ TCP P2P data transfer successful!")
print(f" Original: {test_data}")
print(f" Received: {received_data}")
print(f" Echoed: {echoed_data}")
@pytest.mark.trio
async def test_tcp_large_data_transfer():
"""Test TCP with larger data payloads."""
host_a, host_b = await create_tcp_host_pair()
# Large test data (10KB)
test_data = b"TCP Large Data Test! " * 500 # ~10KB
received_data = None
transfer_complete = trio.Event()
async def large_data_handler(stream):
nonlocal received_data
try:
# Read data in chunks
chunks = []
total_received = 0
expected_size = len(test_data)
while total_received < expected_size:
chunk = await stream.read(min(1024, expected_size - total_received))
if not chunk:
break
chunks.append(chunk)
total_received += len(chunk)
received_data = b"".join(chunks)
# Send back confirmation
await stream.write(b"RECEIVED_OK")
await stream.close()
transfer_complete.set()
except Exception as e:
print(f"Large data handler error: {e}")
transfer_complete.set()
host_a.set_stream_handler(TCP_DATA_PROTOCOL, large_data_handler)
async with (
host_a.run(listen_addrs=[Multiaddr("/ip4/127.0.0.1/tcp/0")]),
host_b.run(listen_addrs=[]),
):
# Get host A's listen address
listen_addrs = host_a.get_addrs()
assert listen_addrs, "Host A should have listen addresses"
# Extract TCP address
tcp_addr = None
for addr in listen_addrs:
if "/tcp/" in str(addr) and "/ws" not in str(addr):
tcp_addr = addr
break
assert tcp_addr, f"No TCP address found in {listen_addrs}"
print(f"🔗 Host A listening on: {tcp_addr}")
print(f"📊 Test data size: {len(test_data)} bytes")
# Create peer info for host A
peer_info = info_from_p2p_addr(tcp_addr)
# Host B connects to host A
await host_b.connect(peer_info)
print("✅ TCP connection established")
# Open a stream for data transfer
stream = await host_b.new_stream(peer_info.peer_id, [TCP_DATA_PROTOCOL])
print("✅ TCP stream opened")
# Send large test data in chunks
chunk_size = 1024
sent_bytes = 0
for i in range(0, len(test_data), chunk_size):
chunk = test_data[i : i + chunk_size]
await stream.write(chunk)
sent_bytes += len(chunk)
if sent_bytes % (chunk_size * 4) == 0: # Progress every 4KB
print(f"📤 Sent {sent_bytes}/{len(test_data)} bytes")
print(f"📤 Sent all {len(test_data)} bytes")
# Read confirmation
confirmation = await stream.read(1024)
print(f"📥 Received confirmation: {confirmation}")
await stream.close()
# Wait for transfer to complete
with trio.fail_after(10.0): # 10 second timeout for large data
await transfer_complete.wait()
# Verify data transfer
assert received_data is not None, "No data was received"
assert received_data == test_data, (
"Large data transfer failed:"
+ f" sizes {len(received_data)} != {len(test_data)}"
)
assert confirmation == b"RECEIVED_OK", f"Confirmation failed: {confirmation}"
print("✅ TCP large data transfer successful!")
print(f" Data size: {len(test_data)} bytes")
print(f" Received: {len(received_data)} bytes")
print(f" Match: {received_data == test_data}")
@pytest.mark.trio
async def test_tcp_bidirectional_transfer():
"""Test bidirectional data transfer over TCP."""
host_a, host_b = await create_tcp_host_pair()
# Test data
data_a_to_b = b"Message from Host A to Host B via TCP"
data_b_to_a = b"Response from Host B to Host A via TCP"
received_on_a = None
received_on_b = None
transfer_complete_a = trio.Event()
transfer_complete_b = trio.Event()
async def handler_a(stream):
nonlocal received_on_a
try:
# Read data from B
received_on_a = await stream.read(len(data_b_to_a))
print(f"🅰️ Host A received: {received_on_a}")
await stream.close()
transfer_complete_a.set()
except Exception as e:
print(f"Handler A error: {e}")
transfer_complete_a.set()
async def handler_b(stream):
nonlocal received_on_b
try:
# Read data from A
received_on_b = await stream.read(len(data_a_to_b))
print(f"🅱️ Host B received: {received_on_b}")
await stream.close()
transfer_complete_b.set()
except Exception as e:
print(f"Handler B error: {e}")
transfer_complete_b.set()
# Set up handlers on both hosts
protocol_a_to_b = TProtocol("/test/tcp-a-to-b/1.0.0")
protocol_b_to_a = TProtocol("/test/tcp-b-to-a/1.0.0")
host_a.set_stream_handler(protocol_b_to_a, handler_a)
host_b.set_stream_handler(protocol_a_to_b, handler_b)
async with (
host_a.run(listen_addrs=[Multiaddr("/ip4/127.0.0.1/tcp/0")]),
host_b.run(listen_addrs=[Multiaddr("/ip4/127.0.0.1/tcp/0")]),
):
# Get addresses
addrs_a = host_a.get_addrs()
addrs_b = host_b.get_addrs()
assert addrs_a and addrs_b, "Both hosts should have addresses"
# Extract TCP addresses
tcp_addr_a = next(
(
addr
for addr in addrs_a
if "/tcp/" in str(addr) and "/ws" not in str(addr)
),
None,
)
tcp_addr_b = next(
(
addr
for addr in addrs_b
if "/tcp/" in str(addr) and "/ws" not in str(addr)
),
None,
)
assert tcp_addr_a and tcp_addr_b, (
f"TCP addresses not found: A={addrs_a}, B={addrs_b}"
)
print(f"🔗 Host A listening on: {tcp_addr_a}")
print(f"🔗 Host B listening on: {tcp_addr_b}")
# Create peer infos
peer_info_a = info_from_p2p_addr(tcp_addr_a)
peer_info_b = info_from_p2p_addr(tcp_addr_b)
# Establish connections
await host_b.connect(peer_info_a)
await host_a.connect(peer_info_b)
print("✅ Bidirectional TCP connections established")
# Send data A -> B
stream_a_to_b = await host_a.new_stream(peer_info_b.peer_id, [protocol_a_to_b])
await stream_a_to_b.write(data_a_to_b)
print(f"📤 A->B: {data_a_to_b}")
await stream_a_to_b.close()
# Send data B -> A
stream_b_to_a = await host_b.new_stream(peer_info_a.peer_id, [protocol_b_to_a])
await stream_b_to_a.write(data_b_to_a)
print(f"📤 B->A: {data_b_to_a}")
await stream_b_to_a.close()
# Wait for both transfers to complete
with trio.fail_after(5.0):
await transfer_complete_a.wait()
await transfer_complete_b.wait()
# Verify bidirectional transfer
assert received_on_a == data_b_to_a, f"A received wrong data: {received_on_a}"
assert received_on_b == data_a_to_b, f"B received wrong data: {received_on_b}"
print("✅ TCP bidirectional data transfer successful!")
print(f" A->B: {data_a_to_b}")
print(f" B->A: {data_b_to_a}")
print(f" ✓ A got: {received_on_a}")
print(f" ✓ B got: {received_on_b}")
if __name__ == "__main__":
# Run tests directly
import logging
logging.basicConfig(level=logging.INFO)
print("🧪 Running TCP P2P Data Transfer Tests")
print("=" * 50)
async def run_all_tcp_tests():
try:
print("\n1. Testing basic TCP connection...")
await test_tcp_basic_connection()
except Exception as e:
print(f"❌ Basic TCP connection test failed: {e}")
return
try:
print("\n2. Testing TCP data transfer...")
await test_tcp_data_transfer()
except Exception as e:
print(f"❌ TCP data transfer test failed: {e}")
return
try:
print("\n3. Testing TCP large data transfer...")
await test_tcp_large_data_transfer()
except Exception as e:
print(f"❌ TCP large data transfer test failed: {e}")
return
try:
print("\n4. Testing TCP bidirectional transfer...")
await test_tcp_bidirectional_transfer()
except Exception as e:
print(f"❌ TCP bidirectional transfer test failed: {e}")
return
print("\n" + "=" * 50)
print("🏁 TCP P2P Tests Complete - All Tests PASSED!")
trio.run(run_all_tcp_tests)

View File

@ -0,0 +1,210 @@
#!/usr/bin/env python3
"""
Demo script showing the new transport integration capabilities in py-libp2p.
This script demonstrates:
1. How to use the transport registry
2. How to create transports dynamically based on multiaddrs
3. How to register custom transports
4. How the new system automatically selects the right transport
"""
import asyncio
import logging
from pathlib import Path
import sys
# Add the libp2p directory to the path so we can import it
sys.path.insert(0, str(Path(__file__).parent.parent))
import multiaddr
from libp2p.transport import (
create_transport,
create_transport_for_multiaddr,
get_supported_transport_protocols,
get_transport_registry,
register_transport,
)
from libp2p.transport.tcp.tcp import TCP
from libp2p.transport.upgrader import TransportUpgrader
# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
def demo_transport_registry():
"""Demonstrate the transport registry functionality."""
print("🔧 Transport Registry Demo")
print("=" * 50)
# Get the global registry
registry = get_transport_registry()
# Show supported protocols
supported = get_supported_transport_protocols()
print(f"Supported transport protocols: {supported}")
# Show registered transports
print("\nRegistered transports:")
for protocol in supported:
transport_class = registry.get_transport(protocol)
class_name = transport_class.__name__ if transport_class else "None"
print(f" {protocol}: {class_name}")
print()
def demo_transport_factory():
"""Demonstrate the transport factory functions."""
print("🏭 Transport Factory Demo")
print("=" * 50)
# Create a dummy upgrader for WebSocket transport
upgrader = TransportUpgrader({}, {})
# Create transports using the factory function
try:
tcp_transport = create_transport("tcp")
print(f"✅ Created TCP transport: {type(tcp_transport).__name__}")
ws_transport = create_transport("ws", upgrader)
print(f"✅ Created WebSocket transport: {type(ws_transport).__name__}")
except Exception as e:
print(f"❌ Error creating transport: {e}")
print()
def demo_multiaddr_transport_selection():
"""Demonstrate automatic transport selection based on multiaddrs."""
print("🎯 Multiaddr Transport Selection Demo")
print("=" * 50)
# Create a dummy upgrader
upgrader = TransportUpgrader({}, {})
# Test different multiaddr types
test_addrs = [
"/ip4/127.0.0.1/tcp/8080",
"/ip4/127.0.0.1/tcp/8080/ws",
"/ip6/::1/tcp/8080/ws",
"/dns4/example.com/tcp/443/ws",
]
for addr_str in test_addrs:
try:
maddr = multiaddr.Multiaddr(addr_str)
transport = create_transport_for_multiaddr(maddr, upgrader)
if transport:
print(f"{addr_str} -> {type(transport).__name__}")
else:
print(f"{addr_str} -> No transport found")
except Exception as e:
print(f"{addr_str} -> Error: {e}")
print()
def demo_custom_transport_registration():
"""Demonstrate how to register custom transports."""
print("🔧 Custom Transport Registration Demo")
print("=" * 50)
# Show current supported protocols
print(f"Before registration: {get_supported_transport_protocols()}")
# Register a custom transport (using TCP as an example)
class CustomTCPTransport(TCP):
"""Custom TCP transport for demonstration."""
def __init__(self):
super().__init__()
self.custom_flag = True
# Register the custom transport
register_transport("custom_tcp", CustomTCPTransport)
# Show updated supported protocols
print(f"After registration: {get_supported_transport_protocols()}")
# Test creating the custom transport
try:
custom_transport = create_transport("custom_tcp")
print(f"✅ Created custom transport: {type(custom_transport).__name__}")
# Check if it has the custom flag (type-safe way)
if hasattr(custom_transport, "custom_flag"):
flag_value = getattr(custom_transport, "custom_flag", "Not found")
print(f" Custom flag: {flag_value}")
else:
print(" Custom flag: Not found")
except Exception as e:
print(f"❌ Error creating custom transport: {e}")
print()
def demo_integration_with_libp2p():
"""Demonstrate how the new system integrates with libp2p."""
print("🚀 Libp2p Integration Demo")
print("=" * 50)
print("The new transport system integrates seamlessly with libp2p:")
print()
print("1. ✅ Automatic transport selection based on multiaddr")
print("2. ✅ Support for WebSocket (/ws) protocol")
print("3. ✅ Fallback to TCP for backward compatibility")
print("4. ✅ Easy registration of new transport protocols")
print("5. ✅ No changes needed to existing libp2p code")
print()
print("Example usage in libp2p:")
print(" # This will automatically use WebSocket transport")
print(" host = new_host(listen_addrs=['/ip4/127.0.0.1/tcp/8080/ws'])")
print()
print(" # This will automatically use TCP transport")
print(" host = new_host(listen_addrs=['/ip4/127.0.0.1/tcp/8080'])")
print()
print()
async def main():
"""Run all demos."""
print("🎉 Py-libp2p Transport Integration Demo")
print("=" * 60)
print()
# Run all demos
demo_transport_registry()
demo_transport_factory()
demo_multiaddr_transport_selection()
demo_custom_transport_registration()
demo_integration_with_libp2p()
print("🎯 Summary of New Features:")
print("=" * 40)
print("✅ Transport Registry: Central registry for all transport implementations")
print("✅ Dynamic Transport Selection: Automatic selection based on multiaddr")
print("✅ WebSocket Support: Full /ws protocol support")
print("✅ Extensible Architecture: Easy to add new transport protocols")
print("✅ Backward Compatibility: Existing TCP code continues to work")
print("✅ Factory Functions: Simple API for creating transports")
print()
print("🚀 The transport system is now ready for production use!")
if __name__ == "__main__":
try:
asyncio.run(main())
except KeyboardInterrupt:
print("\n👋 Demo interrupted by user")
except Exception as e:
print(f"\n❌ Demo failed with error: {e}")
import traceback
traceback.print_exc()

View File

@ -0,0 +1,220 @@
#!/usr/bin/env python3
"""
Simple TCP echo demo to verify basic libp2p functionality.
"""
import argparse
import logging
import traceback
import multiaddr
import trio
from libp2p.crypto.secp256k1 import create_new_key_pair
from libp2p.custom_types import TProtocol
from libp2p.host.basic_host import BasicHost
from libp2p.network.swarm import Swarm
from libp2p.peer.id import ID
from libp2p.peer.peerinfo import info_from_p2p_addr
from libp2p.peer.peerstore import PeerStore
from libp2p.security.insecure.transport import PLAINTEXT_PROTOCOL_ID, InsecureTransport
from libp2p.stream_muxer.yamux.yamux import Yamux
from libp2p.transport.tcp.tcp import TCP
from libp2p.transport.upgrader import TransportUpgrader
# Enable debug logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger("libp2p.tcp-example")
# Simple echo protocol
ECHO_PROTOCOL_ID = TProtocol("/echo/1.0.0")
async def echo_handler(stream):
"""Simple echo handler that echoes back any data received."""
try:
data = await stream.read(1024)
if data:
message = data.decode("utf-8", errors="replace")
print(f"📥 Received: {message}")
print(f"📤 Echoing back: {message}")
await stream.write(data)
await stream.close()
except Exception as e:
logger.error(f"Echo handler error: {e}")
await stream.close()
def create_tcp_host():
"""Create a host with TCP transport."""
# Create key pair and peer store
key_pair = create_new_key_pair()
peer_id = ID.from_pubkey(key_pair.public_key)
peer_store = PeerStore()
peer_store.add_key_pair(peer_id, key_pair)
# Create transport upgrader with plaintext security
upgrader = TransportUpgrader(
secure_transports_by_protocol={
TProtocol(PLAINTEXT_PROTOCOL_ID): InsecureTransport(key_pair)
},
muxer_transports_by_protocol={TProtocol("/yamux/1.0.0"): Yamux},
)
# Create TCP transport
transport = TCP()
# Create swarm and host
swarm = Swarm(peer_id, peer_store, upgrader, transport)
host = BasicHost(swarm)
return host
async def run(port: int, destination: str) -> None:
localhost_ip = "0.0.0.0"
if not destination:
# Create first host (listener) with TCP transport
listen_addr = multiaddr.Multiaddr(f"/ip4/{localhost_ip}/tcp/{port}")
try:
host = create_tcp_host()
logger.debug("Created TCP host")
# Set up echo handler
host.set_stream_handler(ECHO_PROTOCOL_ID, echo_handler)
async with (
host.run(listen_addrs=[listen_addr]),
trio.open_nursery() as (nursery),
):
# Start the peer-store cleanup task
nursery.start_soon(host.get_peerstore().start_cleanup_task, 60)
# Get the actual address and replace 0.0.0.0 with 127.0.0.1 for client
# connections
addrs = host.get_addrs()
logger.debug(f"Host addresses: {addrs}")
if not addrs:
print("❌ Error: No addresses found for the host")
return
server_addr = str(addrs[0])
client_addr = server_addr.replace("/ip4/0.0.0.0/", "/ip4/127.0.0.1/")
print("🌐 TCP Server Started Successfully!")
print("=" * 50)
print(f"📍 Server Address: {client_addr}")
print("🔧 Protocol: /echo/1.0.0")
print("🚀 Transport: TCP")
print()
print("📋 To test the connection, run this in another terminal:")
print(f" python test_tcp_echo.py -d {client_addr}")
print()
print("⏳ Waiting for incoming TCP connections...")
print("" * 50)
await trio.sleep_forever()
except Exception as e:
print(f"❌ Error creating TCP server: {e}")
traceback.print_exc()
return
else:
# Create second host (dialer) with TCP transport
listen_addr = multiaddr.Multiaddr(f"/ip4/{localhost_ip}/tcp/{port}")
try:
# Create a single host for client operations
host = create_tcp_host()
# Start the host for client operations
async with (
host.run(listen_addrs=[listen_addr]),
trio.open_nursery() as (nursery),
):
# Start the peer-store cleanup task
nursery.start_soon(host.get_peerstore().start_cleanup_task, 60)
maddr = multiaddr.Multiaddr(destination)
info = info_from_p2p_addr(maddr)
print("🔌 TCP Client Starting...")
print("=" * 40)
print(f"🎯 Target Peer: {info.peer_id}")
print(f"📍 Target Address: {destination}")
print()
try:
print("🔗 Connecting to TCP server...")
await host.connect(info)
print("✅ Successfully connected to TCP server!")
except Exception as e:
error_msg = str(e)
print("\n❌ Connection Failed!")
print(f" Peer ID: {info.peer_id}")
print(f" Address: {destination}")
print(f" Error: {error_msg}")
return
# Create a stream and send test data
try:
stream = await host.new_stream(info.peer_id, [ECHO_PROTOCOL_ID])
except Exception as e:
print(f"❌ Failed to create stream: {e}")
return
try:
print("🚀 Starting Echo Protocol Test...")
print("" * 40)
# Send test data
test_message = b"Hello TCP Transport!"
print(f"📤 Sending message: {test_message.decode('utf-8')}")
await stream.write(test_message)
# Read response
print("⏳ Waiting for server response...")
response = await stream.read(1024)
print(f"📥 Received response: {response.decode('utf-8')}")
await stream.close()
print("" * 40)
if response == test_message:
print("🎉 Echo test successful!")
print("✅ TCP transport is working perfectly!")
else:
print("❌ Echo test failed!")
except Exception as e:
print(f"Echo protocol error: {e}")
traceback.print_exc()
print("✅ TCP demo completed successfully!")
except Exception as e:
print(f"❌ Error creating TCP client: {e}")
traceback.print_exc()
return
def main() -> None:
description = "Simple TCP echo demo for libp2p"
parser = argparse.ArgumentParser(description=description)
parser.add_argument("-p", "--port", default=0, type=int, help="source port number")
parser.add_argument(
"-d", "--destination", type=str, help="destination multiaddr string"
)
args = parser.parse_args()
try:
trio.run(run, args.port, args.destination)
except KeyboardInterrupt:
pass
if __name__ == "__main__":
main()

View File

@ -0,0 +1,145 @@
#!/usr/bin/env python3
"""
Simple test script to verify WebSocket transport functionality.
"""
import asyncio
import logging
from pathlib import Path
import sys
# Add the libp2p directory to the path so we can import it
sys.path.insert(0, str(Path(__file__).parent))
import multiaddr
from libp2p.transport import create_transport, create_transport_for_multiaddr
from libp2p.transport.upgrader import TransportUpgrader
# Set up logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)
async def test_websocket_transport():
"""Test basic WebSocket transport functionality."""
print("🧪 Testing WebSocket Transport Functionality")
print("=" * 50)
# Create a dummy upgrader
upgrader = TransportUpgrader({}, {})
# Test creating WebSocket transport
try:
ws_transport = create_transport("ws", upgrader)
print(f"✅ WebSocket transport created: {type(ws_transport).__name__}")
# Test creating transport from multiaddr
ws_maddr = multiaddr.Multiaddr("/ip4/127.0.0.1/tcp/8080/ws")
ws_transport_from_maddr = create_transport_for_multiaddr(ws_maddr, upgrader)
print(
f"✅ WebSocket transport from multiaddr: "
f"{type(ws_transport_from_maddr).__name__}"
)
# Test creating listener
handler_called = False
async def test_handler(conn):
nonlocal handler_called
handler_called = True
print(f"✅ Connection handler called with: {type(conn).__name__}")
await conn.close()
listener = ws_transport.create_listener(test_handler)
print(f"✅ WebSocket listener created: {type(listener).__name__}")
# Test that the transport can be used
print(
f"✅ WebSocket transport supports dialing: {hasattr(ws_transport, 'dial')}"
)
print(
f"✅ WebSocket transport supports listening: "
f"{hasattr(ws_transport, 'create_listener')}"
)
print("\n🎯 WebSocket Transport Test Results:")
print("✅ Transport creation: PASS")
print("✅ Multiaddr parsing: PASS")
print("✅ Listener creation: PASS")
print("✅ Interface compliance: PASS")
except Exception as e:
print(f"❌ WebSocket transport test failed: {e}")
import traceback
traceback.print_exc()
return False
return True
async def test_transport_registry():
"""Test the transport registry functionality."""
print("\n🔧 Testing Transport Registry")
print("=" * 30)
from libp2p.transport import (
get_supported_transport_protocols,
get_transport_registry,
)
registry = get_transport_registry()
supported = get_supported_transport_protocols()
print(f"Supported protocols: {supported}")
# Test getting transports
for protocol in supported:
transport_class = registry.get_transport(protocol)
class_name = transport_class.__name__ if transport_class else "None"
print(f" {protocol}: {class_name}")
# Test creating transports through registry
upgrader = TransportUpgrader({}, {})
for protocol in supported:
try:
transport = registry.create_transport(protocol, upgrader)
if transport:
print(f"{protocol}: Created successfully")
else:
print(f"{protocol}: Failed to create")
except Exception as e:
print(f"{protocol}: Error - {e}")
async def main():
"""Run all tests."""
print("🚀 WebSocket Transport Integration Test Suite")
print("=" * 60)
print()
# Run tests
success = await test_websocket_transport()
await test_transport_registry()
print("\n" + "=" * 60)
if success:
print("🎉 All tests passed! WebSocket transport is working correctly.")
else:
print("❌ Some tests failed. Check the output above for details.")
print("\n🚀 WebSocket transport is ready for use in py-libp2p!")
if __name__ == "__main__":
try:
asyncio.run(main())
except KeyboardInterrupt:
print("\n👋 Test interrupted by user")
except Exception as e:
print(f"\n❌ Test failed with error: {e}")
import traceback
traceback.print_exc()

View File

@ -0,0 +1,448 @@
import argparse
import logging
import signal
import sys
import traceback
import multiaddr
import trio
from libp2p.abc import INotifee
from libp2p.crypto.ed25519 import create_new_key_pair as create_ed25519_key_pair
from libp2p.crypto.secp256k1 import create_new_key_pair
from libp2p.custom_types import TProtocol
from libp2p.host.basic_host import BasicHost
from libp2p.network.swarm import Swarm
from libp2p.peer.id import ID
from libp2p.peer.peerinfo import info_from_p2p_addr
from libp2p.peer.peerstore import PeerStore
from libp2p.security.insecure.transport import PLAINTEXT_PROTOCOL_ID, InsecureTransport
from libp2p.security.noise.transport import (
PROTOCOL_ID as NOISE_PROTOCOL_ID,
Transport as NoiseTransport,
)
from libp2p.stream_muxer.yamux.yamux import Yamux
from libp2p.transport.upgrader import TransportUpgrader
from libp2p.transport.websocket.transport import WebsocketTransport
# Enable debug logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger("libp2p.websocket-example")
# Suppress KeyboardInterrupt by handling SIGINT directly
def signal_handler(signum, frame):
print("✅ Clean exit completed.")
sys.exit(0)
signal.signal(signal.SIGINT, signal_handler)
# Simple echo protocol
ECHO_PROTOCOL_ID = TProtocol("/echo/1.0.0")
async def echo_handler(stream):
"""Simple echo handler that echoes back any data received."""
try:
data = await stream.read(1024)
if data:
message = data.decode("utf-8", errors="replace")
print(f"📥 Received: {message}")
print(f"📤 Echoing back: {message}")
await stream.write(data)
await stream.close()
except Exception as e:
logger.error(f"Echo handler error: {e}")
await stream.close()
def create_websocket_host(listen_addrs=None, use_plaintext=False):
"""Create a host with WebSocket transport."""
# Create key pair and peer store
key_pair = create_new_key_pair()
peer_id = ID.from_pubkey(key_pair.public_key)
peer_store = PeerStore()
peer_store.add_key_pair(peer_id, key_pair)
if use_plaintext:
# Create transport upgrader with plaintext security
upgrader = TransportUpgrader(
secure_transports_by_protocol={
TProtocol(PLAINTEXT_PROTOCOL_ID): InsecureTransport(key_pair)
},
muxer_transports_by_protocol={TProtocol("/yamux/1.0.0"): Yamux},
)
else:
# Create separate Ed25519 key for Noise protocol
noise_key_pair = create_ed25519_key_pair()
# Create Noise transport
noise_transport = NoiseTransport(
libp2p_keypair=key_pair,
noise_privkey=noise_key_pair.private_key,
early_data=None,
with_noise_pipes=False,
)
# Create transport upgrader with Noise security
upgrader = TransportUpgrader(
secure_transports_by_protocol={
TProtocol(NOISE_PROTOCOL_ID): noise_transport
},
muxer_transports_by_protocol={TProtocol("/yamux/1.0.0"): Yamux},
)
# Create WebSocket transport
transport = WebsocketTransport(upgrader)
# Create swarm and host
swarm = Swarm(peer_id, peer_store, upgrader, transport)
host = BasicHost(swarm)
return host
async def run(port: int, destination: str, use_plaintext: bool = False) -> None:
localhost_ip = "0.0.0.0"
if not destination:
# Create first host (listener) with WebSocket transport
listen_addr = multiaddr.Multiaddr(f"/ip4/{localhost_ip}/tcp/{port}/ws")
try:
host = create_websocket_host(use_plaintext=use_plaintext)
logger.debug(f"Created host with use_plaintext={use_plaintext}")
# Set up echo handler
host.set_stream_handler(ECHO_PROTOCOL_ID, echo_handler)
# Add connection event handlers for debugging
class DebugNotifee(INotifee):
async def opened_stream(self, network, stream):
pass
async def closed_stream(self, network, stream):
pass
async def connected(self, network, conn):
print(
f"🔗 New libp2p connection established: "
f"{conn.muxed_conn.peer_id}"
)
if hasattr(conn.muxed_conn, "get_security_protocol"):
security = conn.muxed_conn.get_security_protocol()
else:
security = "Unknown"
print(f" Security: {security}")
async def disconnected(self, network, conn):
print(f"🔌 libp2p connection closed: {conn.muxed_conn.peer_id}")
async def listen(self, network, multiaddr):
pass
async def listen_close(self, network, multiaddr):
pass
host.get_network().register_notifee(DebugNotifee())
# Create a cancellation token for clean shutdown
cancel_scope = trio.CancelScope()
async def signal_handler():
with trio.open_signal_receiver(signal.SIGINT, signal.SIGTERM) as (
signal_receiver
):
async for sig in signal_receiver:
print(f"\n🛑 Received signal {sig}")
print("✅ Shutting down WebSocket server...")
cancel_scope.cancel()
return
async with (
host.run(listen_addrs=[listen_addr]),
trio.open_nursery() as (nursery),
):
# Start the peer-store cleanup task
nursery.start_soon(host.get_peerstore().start_cleanup_task, 60)
# Start the signal handler
nursery.start_soon(signal_handler)
# Get the actual address and replace 0.0.0.0 with 127.0.0.1 for client
# connections
addrs = host.get_addrs()
logger.debug(f"Host addresses: {addrs}")
if not addrs:
print("❌ Error: No addresses found for the host")
print("Debug: host.get_addrs() returned empty list")
return
server_addr = str(addrs[0])
client_addr = server_addr.replace("/ip4/0.0.0.0/", "/ip4/127.0.0.1/")
print("🌐 WebSocket Server Started Successfully!")
print("=" * 50)
print(f"📍 Server Address: {client_addr}")
print("🔧 Protocol: /echo/1.0.0")
print("🚀 Transport: WebSocket (/ws)")
print()
print("📋 To test the connection, run this in another terminal:")
plaintext_flag = " --plaintext" if use_plaintext else ""
print(f" python websocket_demo.py -d {client_addr}{plaintext_flag}")
print()
print("⏳ Waiting for incoming WebSocket connections...")
print("" * 50)
# Add a custom handler to show connection events
async def custom_echo_handler(stream):
peer_id = stream.muxed_conn.peer_id
print("\n🔗 New WebSocket Connection!")
print(f" Peer ID: {peer_id}")
print(" Protocol: /echo/1.0.0")
# Show remote address in multiaddr format
try:
remote_address = stream.get_remote_address()
if remote_address:
print(f" Remote: {remote_address}")
except Exception:
print(" Remote: Unknown")
print("" * 40)
# Call the original handler
await echo_handler(stream)
print("" * 40)
print(f"✅ Echo request completed for peer: {peer_id}")
print()
# Replace the handler with our custom one
host.set_stream_handler(ECHO_PROTOCOL_ID, custom_echo_handler)
# Wait indefinitely or until cancelled
with cancel_scope:
await trio.sleep_forever()
except Exception as e:
print(f"❌ Error creating WebSocket server: {e}")
traceback.print_exc()
return
else:
# Create second host (dialer) with WebSocket transport
listen_addr = multiaddr.Multiaddr(f"/ip4/{localhost_ip}/tcp/{port}/ws")
try:
# Create a single host for client operations
host = create_websocket_host(use_plaintext=use_plaintext)
# Start the host for client operations
async with (
host.run(listen_addrs=[listen_addr]),
trio.open_nursery() as (nursery),
):
# Start the peer-store cleanup task
nursery.start_soon(host.get_peerstore().start_cleanup_task, 60)
# Add connection event handlers for debugging
class ClientDebugNotifee(INotifee):
async def opened_stream(self, network, stream):
pass
async def closed_stream(self, network, stream):
pass
async def connected(self, network, conn):
print(
f"🔗 Client: libp2p connection established: "
f"{conn.muxed_conn.peer_id}"
)
async def disconnected(self, network, conn):
print(
f"🔌 Client: libp2p connection closed: "
f"{conn.muxed_conn.peer_id}"
)
async def listen(self, network, multiaddr):
pass
async def listen_close(self, network, multiaddr):
pass
host.get_network().register_notifee(ClientDebugNotifee())
maddr = multiaddr.Multiaddr(destination)
info = info_from_p2p_addr(maddr)
print("🔌 WebSocket Client Starting...")
print("=" * 40)
print(f"🎯 Target Peer: {info.peer_id}")
print(f"📍 Target Address: {destination}")
print()
try:
print("🔗 Connecting to WebSocket server...")
print(f" Security: {'Plaintext' if use_plaintext else 'Noise'}")
await host.connect(info)
print("✅ Successfully connected to WebSocket server!")
except Exception as e:
error_msg = str(e)
print("\n❌ Connection Failed!")
print(f" Peer ID: {info.peer_id}")
print(f" Address: {destination}")
print(f" Security: {'Plaintext' if use_plaintext else 'Noise'}")
print(f" Error: {error_msg}")
print(f" Error type: {type(e).__name__}")
# Add more detailed error information for debugging
if hasattr(e, "__cause__") and e.__cause__:
print(f" Root cause: {e.__cause__}")
print(f" Root cause type: {type(e.__cause__).__name__}")
print()
print("💡 Troubleshooting:")
print(" • Make sure the WebSocket server is running")
print(" • Check that the server address is correct")
print(" • Verify the server is listening on the right port")
print(
" • Ensure both client and server use the same sec protocol"
)
if not use_plaintext:
print(" • Noise over WebSocket may have compatibility issues")
return
# Create a stream and send test data
try:
stream = await host.new_stream(info.peer_id, [ECHO_PROTOCOL_ID])
except Exception as e:
print(f"❌ Failed to create stream: {e}")
return
try:
print("🚀 Starting Echo Protocol Test...")
print("" * 40)
# Send test data
test_message = b"Hello WebSocket Transport!"
print(f"📤 Sending message: {test_message.decode('utf-8')}")
await stream.write(test_message)
# Read response
print("⏳ Waiting for server response...")
response = await stream.read(1024)
print(f"📥 Received response: {response.decode('utf-8')}")
await stream.close()
print("" * 40)
if response == test_message:
print("🎉 Echo test successful!")
print("✅ WebSocket transport is working perfectly!")
print("✅ Client completed successfully, exiting.")
else:
print("❌ Echo test failed!")
print(" Response doesn't match sent data.")
print(f" Sent: {test_message}")
print(f" Received: {response}")
except Exception as e:
error_msg = str(e)
print(f"Echo protocol error: {error_msg}")
traceback.print_exc()
finally:
# Ensure stream is closed
try:
if stream:
# Check if stream has is_closed method and use it
has_is_closed = hasattr(stream, "is_closed") and callable(
getattr(stream, "is_closed")
)
if has_is_closed:
# type: ignore[attr-defined]
if not await stream.is_closed():
await stream.close()
else:
# Fallback: just try to close the stream
await stream.close()
except Exception:
pass
# host.run() context manager handles cleanup automatically
print()
print("🎉 WebSocket Demo Completed Successfully!")
print("=" * 50)
print("✅ WebSocket transport is working perfectly!")
print("✅ Echo protocol communication successful!")
print("✅ libp2p integration verified!")
print()
print("🚀 Your WebSocket transport is ready for production use!")
# Add a small delay to ensure all cleanup is complete
await trio.sleep(0.1)
except Exception as e:
print(f"❌ Error creating WebSocket client: {e}")
traceback.print_exc()
return
def main() -> None:
description = """
This program demonstrates the libp2p WebSocket transport.
First run
'python websocket_demo.py -p <PORT> [--plaintext]' to start a WebSocket server.
Then run
'python websocket_demo.py <ANOTHER_PORT> -d <DESTINATION> [--plaintext]'
where <DESTINATION> is the multiaddress shown by the server.
By default, this example uses Noise encryption for secure communication.
Use --plaintext for testing with unencrypted communication
(not recommended for production).
"""
example_maddr = (
"/ip4/127.0.0.1/tcp/8888/ws/p2p/QmQn4SwGkDZKkUEpBRBvTmheQycxAHJUNmVEnjA2v1qe8Q"
)
parser = argparse.ArgumentParser(description=description)
parser.add_argument("-p", "--port", default=0, type=int, help="source port number")
parser.add_argument(
"-d",
"--destination",
type=str,
help=f"destination multiaddr string, e.g. {example_maddr}",
)
parser.add_argument(
"--plaintext",
action="store_true",
help=(
"use plaintext security instead of Noise encryption "
"(not recommended for production)"
),
)
args = parser.parse_args()
# Determine security mode: use Noise by default,
# plaintext if --plaintext is specified
use_plaintext = args.plaintext
try:
trio.run(run, args.port, args.destination, use_plaintext)
except KeyboardInterrupt:
# This is expected when Ctrl+C is pressed
# The signal handler already printed the shutdown message
print("✅ Clean exit completed.")
return
except Exception as e:
print(f"❌ Unexpected error: {e}")
return
if __name__ == "__main__":
main()

View File

@ -1,5 +1,12 @@
"""Libp2p Python implementation."""
import logging
import ssl
from libp2p.transport.quic.utils import is_quic_multiaddr
from typing import Any
from libp2p.transport.quic.transport import QUICTransport
from libp2p.transport.quic.config import QUICTransportConfig
from collections.abc import (
Mapping,
Sequence,
@ -18,6 +25,7 @@ from libp2p.abc import (
IPeerRouting,
IPeerStore,
ISecureTransport,
ITransport,
)
from libp2p.crypto.keys import (
KeyPair,
@ -38,10 +46,12 @@ from libp2p.host.routed_host import (
RoutedHost,
)
from libp2p.network.swarm import (
ConnectionConfig,
RetryConfig,
Swarm,
)
from libp2p.network.config import (
ConnectionConfig,
RetryConfig
)
from libp2p.peer.id import (
ID,
)
@ -72,6 +82,10 @@ from libp2p.transport.tcp.tcp import (
from libp2p.transport.upgrader import (
TransportUpgrader,
)
from libp2p.transport.transport_registry import (
create_transport_for_multiaddr,
get_supported_transport_protocols,
)
from libp2p.utils.logging import (
setup_logging,
)
@ -87,6 +101,7 @@ MUXER_YAMUX = "YAMUX"
MUXER_MPLEX = "MPLEX"
DEFAULT_NEGOTIATE_TIMEOUT = 5
logger = logging.getLogger(__name__)
def set_default_muxer(muxer_name: Literal["YAMUX", "MPLEX"]) -> None:
"""
@ -162,9 +177,13 @@ def new_swarm(
peerstore_opt: IPeerStore | None = None,
muxer_preference: Literal["YAMUX", "MPLEX"] | None = None,
listen_addrs: Sequence[multiaddr.Multiaddr] | None = None,
enable_quic: bool = False,
retry_config: Optional["RetryConfig"] = None,
connection_config: Optional["ConnectionConfig"] = None,
connection_config: ConnectionConfig | QUICTransportConfig | None = None,
tls_client_config: ssl.SSLContext | None = None,
tls_server_config: ssl.SSLContext | None = None,
) -> INetworkService:
logger.debug(f"new_swarm: enable_quic={enable_quic}, listen_addrs={listen_addrs}")
"""
Create a swarm instance based on the parameters.
@ -174,6 +193,8 @@ def new_swarm(
:param peerstore_opt: optional peerstore
:param muxer_preference: optional explicit muxer preference
:param listen_addrs: optional list of multiaddrs to listen on
:param enable_quic: enable quic for transport
:param quic_transport_opt: options for transport
:return: return a default swarm instance
Note: Yamux (/yamux/1.0.0) is the preferred stream multiplexer
@ -186,16 +207,48 @@ def new_swarm(
id_opt = generate_peer_id_from(key_pair)
transport: TCP | QUICTransport | ITransport
quic_transport_opt = connection_config if isinstance(connection_config, QUICTransportConfig) else None
if listen_addrs is None:
transport = TCP()
else:
addr = listen_addrs[0]
if addr.__contains__("tcp"):
transport = TCP()
elif addr.__contains__("quic"):
raise ValueError("QUIC not yet supported")
if enable_quic:
transport = QUICTransport(key_pair.private_key, config=quic_transport_opt)
else:
raise ValueError(f"Unknown transport in listen_addrs: {listen_addrs}")
transport = TCP()
else:
# Use transport registry to select the appropriate transport
from libp2p.transport.transport_registry import create_transport_for_multiaddr
# Create a temporary upgrader for transport selection
# We'll create the real upgrader later with the proper configuration
temp_upgrader = TransportUpgrader(
secure_transports_by_protocol={},
muxer_transports_by_protocol={}
)
addr = listen_addrs[0]
logger.debug(f"new_swarm: Creating transport for address: {addr}")
transport_maybe = create_transport_for_multiaddr(
addr,
temp_upgrader,
private_key=key_pair.private_key,
config=quic_transport_opt,
tls_client_config=tls_client_config,
tls_server_config=tls_server_config
)
if transport_maybe is None:
raise ValueError(f"Unsupported transport for listen_addrs: {listen_addrs}")
transport = transport_maybe
logger.debug(f"new_swarm: Created transport: {type(transport)}")
# If enable_quic is True but we didn't get a QUIC transport, force QUIC
if enable_quic and not isinstance(transport, QUICTransport):
logger.debug(f"new_swarm: Forcing QUIC transport (enable_quic=True but got {type(transport)})")
transport = QUICTransport(key_pair.private_key, config=quic_transport_opt)
logger.debug(f"new_swarm: Final transport type: {type(transport)}")
# Generate X25519 keypair for Noise
noise_key_pair = create_new_x25519_key_pair()
@ -236,6 +289,7 @@ def new_swarm(
muxer_transports_by_protocol=muxer_transports_by_protocol,
)
peerstore = peerstore_opt or PeerStore()
# Store our key pair in peerstore
peerstore.add_key_pair(id_opt, key_pair)
@ -261,6 +315,10 @@ def new_host(
enable_mDNS: bool = False,
bootstrap: list[str] | None = None,
negotiate_timeout: int = DEFAULT_NEGOTIATE_TIMEOUT,
enable_quic: bool = False,
quic_transport_opt: QUICTransportConfig | None = None,
tls_client_config: ssl.SSLContext | None = None,
tls_server_config: ssl.SSLContext | None = None,
) -> IHost:
"""
Create a new libp2p host based on the given parameters.
@ -274,15 +332,27 @@ def new_host(
:param listen_addrs: optional list of multiaddrs to listen on
:param enable_mDNS: whether to enable mDNS discovery
:param bootstrap: optional list of bootstrap peer addresses as strings
:param enable_quic: optinal choice to use QUIC for transport
:param quic_transport_opt: optional configuration for quic transport
:param tls_client_config: optional TLS client configuration for WebSocket transport
:param tls_server_config: optional TLS server configuration for WebSocket transport
:return: return a host instance
"""
if not enable_quic and quic_transport_opt is not None:
logger.warning(f"QUIC config provided but QUIC not enabled, ignoring QUIC config")
swarm = new_swarm(
enable_quic=enable_quic,
key_pair=key_pair,
muxer_opt=muxer_opt,
sec_opt=sec_opt,
peerstore_opt=peerstore_opt,
muxer_preference=muxer_preference,
listen_addrs=listen_addrs,
connection_config=quic_transport_opt if enable_quic else None,
tls_client_config=tls_client_config,
tls_server_config=tls_server_config
)
if disc_opt is not None:

View File

@ -5,17 +5,17 @@ from collections.abc import (
)
from typing import TYPE_CHECKING, NewType, Union, cast
from libp2p.transport.quic.stream import QUICStream
if TYPE_CHECKING:
from libp2p.abc import (
IMuxedConn,
INetStream,
ISecureTransport,
)
from libp2p.abc import IMuxedConn, IMuxedStream, INetStream, ISecureTransport
from libp2p.transport.quic.connection import QUICConnection
else:
IMuxedConn = cast(type, object)
INetStream = cast(type, object)
ISecureTransport = cast(type, object)
IMuxedStream = cast(type, object)
QUICConnection = cast(type, object)
from libp2p.io.abc import (
ReadWriteCloser,
@ -37,3 +37,6 @@ SyncValidatorFn = Callable[[ID, rpc_pb2.Message], bool]
AsyncValidatorFn = Callable[[ID, rpc_pb2.Message], Awaitable[bool]]
ValidatorFn = Union[SyncValidatorFn, AsyncValidatorFn]
UnsubscribeFn = Callable[[], Awaitable[None]]
TQUICStreamHandlerFn = Callable[[QUICStream], Awaitable[None]]
TQUICConnHandlerFn = Callable[[QUICConnection], Awaitable[None]]
MessageID = NewType("MessageID", str)

View File

@ -213,7 +213,6 @@ class BasicHost(IHost):
self,
peer_id: ID,
protocol_ids: Sequence[TProtocol],
negotitate_timeout: int = DEFAULT_NEGOTIATE_TIMEOUT,
) -> INetStream:
"""
:param peer_id: peer_id that host is connecting
@ -227,7 +226,7 @@ class BasicHost(IHost):
selected_protocol = await self.multiselect_client.select_one_of(
list(protocol_ids),
MultiselectCommunicator(net_stream),
negotitate_timeout,
self.negotiate_timeout,
)
except MultiselectClientError as error:
logger.debug("fail to open a stream to peer %s, error=%s", peer_id, error)

70
libp2p/network/config.py Normal file
View File

@ -0,0 +1,70 @@
from dataclasses import dataclass
@dataclass
class RetryConfig:
"""
Configuration for retry logic with exponential backoff.
This configuration controls how connection attempts are retried when they fail.
The retry mechanism uses exponential backoff with jitter to prevent thundering
herd problems in distributed systems.
Attributes:
max_retries: Maximum number of retry attempts before giving up.
Default: 3 attempts
initial_delay: Initial delay in seconds before the first retry.
Default: 0.1 seconds (100ms)
max_delay: Maximum delay cap in seconds to prevent excessive wait times.
Default: 30.0 seconds
backoff_multiplier: Multiplier for exponential backoff (each retry multiplies
the delay by this factor). Default: 2.0 (doubles each time)
jitter_factor: Random jitter factor (0.0-1.0) to add randomness to delays
and prevent synchronized retries. Default: 0.1 (10% jitter)
"""
max_retries: int = 3
initial_delay: float = 0.1
max_delay: float = 30.0
backoff_multiplier: float = 2.0
jitter_factor: float = 0.1
@dataclass
class ConnectionConfig:
"""
Configuration for multi-connection support.
This configuration controls how multiple connections per peer are managed,
including connection limits, timeouts, and load balancing strategies.
Attributes:
max_connections_per_peer: Maximum number of connections allowed to a single
peer. Default: 3 connections
connection_timeout: Timeout in seconds for establishing new connections.
Default: 30.0 seconds
load_balancing_strategy: Strategy for distributing streams across connections.
Options: "round_robin" (default) or "least_loaded"
"""
max_connections_per_peer: int = 3
connection_timeout: float = 30.0
load_balancing_strategy: str = "round_robin" # or "least_loaded"
def __post_init__(self) -> None:
"""Validate configuration after initialization."""
if not (
self.load_balancing_strategy == "round_robin"
or self.load_balancing_strategy == "least_loaded"
):
raise ValueError(
"Load balancing strategy can only be 'round_robin' or 'least_loaded'"
)
if self.max_connections_per_peer < 1:
raise ValueError("Max connection per peer should be atleast 1")
if self.connection_timeout < 0:
raise ValueError("Connection timeout should be positive")

View File

@ -17,6 +17,7 @@ from libp2p.stream_muxer.exceptions import (
MuxedStreamError,
MuxedStreamReset,
)
from libp2p.transport.quic.exceptions import QUICStreamClosedError, QUICStreamResetError
from .exceptions import (
StreamClosed,
@ -170,7 +171,7 @@ class NetStream(INetStream):
elif self.__stream_state == StreamState.OPEN:
self.__stream_state = StreamState.CLOSE_READ
raise StreamEOF() from error
except MuxedStreamReset as error:
except (MuxedStreamReset, QUICStreamClosedError, QUICStreamResetError) as error:
async with self._state_lock:
if self.__stream_state in [
StreamState.OPEN,
@ -199,7 +200,12 @@ class NetStream(INetStream):
try:
await self.muxed_stream.write(data)
except (MuxedStreamClosed, MuxedStreamError) as error:
except (
MuxedStreamClosed,
MuxedStreamError,
QUICStreamClosedError,
QUICStreamResetError,
) as error:
async with self._state_lock:
if self.__stream_state == StreamState.OPEN:
self.__stream_state = StreamState.CLOSE_WRITE

View File

@ -2,9 +2,9 @@ from collections.abc import (
Awaitable,
Callable,
)
from dataclasses import dataclass
import logging
import random
from typing import cast
from multiaddr import (
Multiaddr,
@ -27,6 +27,7 @@ from libp2p.custom_types import (
from libp2p.io.abc import (
ReadWriteCloser,
)
from libp2p.network.config import ConnectionConfig, RetryConfig
from libp2p.peer.id import (
ID,
)
@ -41,6 +42,9 @@ from libp2p.transport.exceptions import (
OpenConnectionError,
SecurityUpgradeFailure,
)
from libp2p.transport.quic.config import QUICTransportConfig
from libp2p.transport.quic.connection import QUICConnection
from libp2p.transport.quic.transport import QUICTransport
from libp2p.transport.upgrader import (
TransportUpgrader,
)
@ -61,59 +65,6 @@ from .exceptions import (
logger = logging.getLogger("libp2p.network.swarm")
@dataclass
class RetryConfig:
"""
Configuration for retry logic with exponential backoff.
This configuration controls how connection attempts are retried when they fail.
The retry mechanism uses exponential backoff with jitter to prevent thundering
herd problems in distributed systems.
Attributes:
max_retries: Maximum number of retry attempts before giving up.
Default: 3 attempts
initial_delay: Initial delay in seconds before the first retry.
Default: 0.1 seconds (100ms)
max_delay: Maximum delay cap in seconds to prevent excessive wait times.
Default: 30.0 seconds
backoff_multiplier: Multiplier for exponential backoff (each retry multiplies
the delay by this factor). Default: 2.0 (doubles each time)
jitter_factor: Random jitter factor (0.0-1.0) to add randomness to delays
and prevent synchronized retries. Default: 0.1 (10% jitter)
"""
max_retries: int = 3
initial_delay: float = 0.1
max_delay: float = 30.0
backoff_multiplier: float = 2.0
jitter_factor: float = 0.1
@dataclass
class ConnectionConfig:
"""
Configuration for multi-connection support.
This configuration controls how multiple connections per peer are managed,
including connection limits, timeouts, and load balancing strategies.
Attributes:
max_connections_per_peer: Maximum number of connections allowed to a single
peer. Default: 3 connections
connection_timeout: Timeout in seconds for establishing new connections.
Default: 30.0 seconds
load_balancing_strategy: Strategy for distributing streams across connections.
Options: "round_robin" (default) or "least_loaded"
"""
max_connections_per_peer: int = 3
connection_timeout: float = 30.0
load_balancing_strategy: str = "round_robin" # or "least_loaded"
def create_default_stream_handler(network: INetworkService) -> StreamHandlerFn:
async def stream_handler(stream: INetStream) -> None:
await network.get_manager().wait_finished()
@ -126,8 +77,7 @@ class Swarm(Service, INetworkService):
peerstore: IPeerStore
upgrader: TransportUpgrader
transport: ITransport
# Enhanced: Support for multiple connections per peer
connections: dict[ID, list[INetConn]] # Multiple connections per peer
connections: dict[ID, list[INetConn]]
listeners: dict[str, IListener]
common_stream_handler: StreamHandlerFn
listener_nursery: trio.Nursery | None
@ -137,7 +87,7 @@ class Swarm(Service, INetworkService):
# Enhanced: New configuration
retry_config: RetryConfig
connection_config: ConnectionConfig
connection_config: ConnectionConfig | QUICTransportConfig
_round_robin_index: dict[ID, int]
def __init__(
@ -147,7 +97,7 @@ class Swarm(Service, INetworkService):
upgrader: TransportUpgrader,
transport: ITransport,
retry_config: RetryConfig | None = None,
connection_config: ConnectionConfig | None = None,
connection_config: ConnectionConfig | QUICTransportConfig | None = None,
):
self.self_id = peer_id
self.peerstore = peerstore
@ -178,6 +128,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:
@ -370,6 +325,7 @@ class Swarm(Service, INetworkService):
# Dial peer (connection to peer does not yet exist)
# Transport dials peer (gets back a raw conn)
try:
addr = Multiaddr(f"{addr}/p2p/{peer_id}")
raw_conn = await self.transport.dial(addr)
except OpenConnectionError as error:
logger.debug("fail to dial peer %s over base transport", peer_id)
@ -377,6 +333,15 @@ class Swarm(Service, INetworkService):
f"fail to open connection to peer {peer_id}"
) from error
if isinstance(self.transport, QUICTransport) and isinstance(
raw_conn, IMuxedConn
):
logger.info(
"Skipping upgrade for QUIC, QUIC connections are already multiplexed"
)
swarm_conn = await self.add_conn(raw_conn)
return swarm_conn
logger.debug("dialed peer %s over base transport", peer_id)
# Per, https://discuss.libp2p.io/t/multistream-security/130, we first secure
@ -402,9 +367,7 @@ class Swarm(Service, INetworkService):
logger.debug("upgraded mux for peer %s", peer_id)
swarm_conn = await self.add_conn(muxed_conn)
logger.debug("successfully dialed peer %s", peer_id)
return swarm_conn
async def dial_addr(self, addr: Multiaddr, peer_id: ID) -> INetConn:
@ -427,7 +390,6 @@ class Swarm(Service, INetworkService):
:return: net stream instance
"""
logger.debug("attempting to open a stream to peer %s", peer_id)
# Get existing connections or dial new ones
connections = self.get_connections(peer_id)
if not connections:
@ -436,6 +398,10 @@ class Swarm(Service, INetworkService):
# Load balancing strategy at interface level
connection = self._select_connection(connections, peer_id)
if isinstance(self.transport, QUICTransport) and connection is not None:
conn = cast(SwarmConn, connection)
return await conn.new_stream()
try:
net_stream = await connection.new_stream()
logger.debug("successfully opened a stream to peer %s", peer_id)
@ -515,18 +481,38 @@ class Swarm(Service, INetworkService):
- Call listener listen with the multiaddr
- Map multiaddr to listener
"""
logger.debug(f"Swarm.listen called with multiaddrs: {multiaddrs}")
# We need to wait until `self.listener_nursery` is created.
logger.debug("Starting to listen")
await self.event_listener_nursery_created.wait()
success_count = 0
for maddr in multiaddrs:
logger.debug(f"Swarm.listen processing multiaddr: {maddr}")
if str(maddr) in self.listeners:
logger.debug(f"Swarm.listen: listener already exists for {maddr}")
success_count += 1
continue
async def conn_handler(
read_write_closer: ReadWriteCloser, maddr: Multiaddr = maddr
) -> None:
# No need to upgrade QUIC Connection
if isinstance(self.transport, QUICTransport):
try:
quic_conn = cast(QUICConnection, read_write_closer)
await self.add_conn(quic_conn)
peer_id = quic_conn.peer_id
logger.debug(
f"successfully opened quic connection to peer {peer_id}"
)
# NOTE: This is a intentional barrier to prevent from the
# handler exiting and closing the connection.
await self.manager.wait_finished()
except Exception:
await read_write_closer.close()
return
raw_conn = RawConnection(read_write_closer, False)
# Per, https://discuss.libp2p.io/t/multistream-security/130, we first
@ -562,13 +548,18 @@ class Swarm(Service, INetworkService):
try:
# Success
logger.debug(f"Swarm.listen: creating listener for {maddr}")
listener = self.transport.create_listener(conn_handler)
logger.debug(f"Swarm.listen: listener created for {maddr}")
self.listeners[str(maddr)] = listener
# TODO: `listener.listen` is not bounded with nursery. If we want to be
# I/O agnostic, we should change the API.
if self.listener_nursery is None:
raise SwarmException("swarm instance hasn't been run")
assert self.listener_nursery is not None # For type checker
logger.debug(f"Swarm.listen: calling listener.listen for {maddr}")
await listener.listen(maddr, self.listener_nursery)
logger.debug(f"Swarm.listen: listener.listen completed for {maddr}")
# Call notifiers since event occurred
await self.notify_listen(maddr)
@ -660,9 +651,10 @@ class Swarm(Service, INetworkService):
muxed_conn,
self,
)
logger.debug("Swarm::add_conn | starting muxed connection")
self.manager.run_task(muxed_conn.start)
await muxed_conn.event_started.wait()
logger.debug("Swarm::add_conn | starting swarm connection")
self.manager.run_task(swarm_conn.start)
await swarm_conn.event_started.wait()

View File

@ -1,3 +1,5 @@
from builtins import AssertionError
from libp2p.abc import (
IMultiselectCommunicator,
)
@ -36,7 +38,8 @@ class MultiselectCommunicator(IMultiselectCommunicator):
msg_bytes = encode_delim(msg_str.encode())
try:
await self.read_writer.write(msg_bytes)
except IOException as error:
# Handle for connection close during ongoing negotiation in QUIC
except (IOException, AssertionError, ValueError) as error:
raise MultiselectCommunicatorError(
"fail to write to multiselect communicator"
) from error

View File

@ -1,6 +1,3 @@
from ast import (
literal_eval,
)
from collections import (
defaultdict,
)
@ -22,6 +19,7 @@ from libp2p.abc import (
IPubsubRouter,
)
from libp2p.custom_types import (
MessageID,
TProtocol,
)
from libp2p.peer.id import (
@ -56,6 +54,10 @@ from .pb import (
from .pubsub import (
Pubsub,
)
from .utils import (
parse_message_id_safe,
safe_parse_message_id,
)
PROTOCOL_ID = TProtocol("/meshsub/1.0.0")
PROTOCOL_ID_V11 = TProtocol("/meshsub/1.1.0")
@ -306,7 +308,8 @@ class GossipSub(IPubsubRouter, Service):
floodsub_peers: set[ID] = {
peer_id
for peer_id in self.pubsub.peer_topics[topic]
if self.peer_protocol[peer_id] == floodsub.PROTOCOL_ID
if peer_id in self.peer_protocol
and self.peer_protocol[peer_id] == floodsub.PROTOCOL_ID
}
send_to.update(floodsub_peers)
@ -794,8 +797,8 @@ class GossipSub(IPubsubRouter, Service):
# Add all unknown message ids (ids that appear in ihave_msg but not in
# seen_seqnos) to list of messages we want to request
msg_ids_wanted: list[str] = [
msg_id
msg_ids_wanted: list[MessageID] = [
parse_message_id_safe(msg_id)
for msg_id in ihave_msg.messageIDs
if msg_id not in seen_seqnos_and_peers
]
@ -811,9 +814,9 @@ class GossipSub(IPubsubRouter, Service):
Forwards all request messages that are present in mcache to the
requesting peer.
"""
# FIXME: Update type of message ID
# FIXME: Find a better way to parse the msg ids
msg_ids: list[Any] = [literal_eval(msg) for msg in iwant_msg.messageIDs]
msg_ids: list[tuple[bytes, bytes]] = [
safe_parse_message_id(msg) for msg in iwant_msg.messageIDs
]
msgs_to_forward: list[rpc_pb2.Message] = []
for msg_id_iwant in msg_ids:
# Check if the wanted message ID is present in mcache

View File

@ -1,6 +1,10 @@
import ast
import logging
from libp2p.abc import IHost
from libp2p.custom_types import (
MessageID,
)
from libp2p.peer.envelope import consume_envelope
from libp2p.peer.id import ID
from libp2p.pubsub.pb.rpc_pb2 import RPC
@ -48,3 +52,29 @@ def maybe_consume_signed_record(msg: RPC, host: IHost, peer_id: ID) -> bool:
logger.error("Failed to update the Certified-Addr-Book: %s", e)
return False
return True
def parse_message_id_safe(msg_id_str: str) -> MessageID:
"""Safely handle message ID as string."""
return MessageID(msg_id_str)
def safe_parse_message_id(msg_id_str: str) -> tuple[bytes, bytes]:
"""
Safely parse message ID using ast.literal_eval with validation.
:param msg_id_str: String representation of message ID
:return: Tuple of (seqno, from_id) as bytes
:raises ValueError: If parsing fails
"""
try:
parsed = ast.literal_eval(msg_id_str)
if not isinstance(parsed, tuple) or len(parsed) != 2:
raise ValueError("Invalid message ID format")
seqno, from_id = parsed
if not isinstance(seqno, bytes) or not isinstance(from_id, bytes):
raise ValueError("Message ID components must be bytes")
return (seqno, from_id)
except (ValueError, SyntaxError) as e:
raise ValueError(f"Invalid message ID format: {e}")

View File

@ -9,6 +9,7 @@ from dataclasses import (
dataclass,
field,
)
from enum import Flag, auto
from libp2p.peer.peerinfo import (
PeerInfo,
@ -18,29 +19,118 @@ from .resources import (
RelayLimits,
)
DEFAULT_MIN_RELAYS = 3
DEFAULT_MAX_RELAYS = 20
DEFAULT_DISCOVERY_INTERVAL = 300 # seconds
DEFAULT_RESERVATION_TTL = 3600 # seconds
DEFAULT_MAX_CIRCUIT_DURATION = 3600 # seconds
DEFAULT_MAX_CIRCUIT_BYTES = 1024 * 1024 * 1024 # 1GB
DEFAULT_MAX_CIRCUIT_CONNS = 8
DEFAULT_MAX_RESERVATIONS = 4
MAX_RESERVATIONS_PER_IP = 8
MAX_CIRCUITS_PER_IP = 16
RESERVATION_RATE_PER_IP = 4 # per minute
CIRCUIT_RATE_PER_IP = 8 # per minute
MAX_CIRCUITS_TOTAL = 64
MAX_RESERVATIONS_TOTAL = 32
MAX_BANDWIDTH_PER_CIRCUIT = 1024 * 1024 # 1MB/s
MAX_BANDWIDTH_TOTAL = 10 * 1024 * 1024 # 10MB/s
MIN_RELAY_SCORE = 0.5
MAX_RELAY_LATENCY = 1.0 # seconds
ENABLE_AUTO_RELAY = True
AUTO_RELAY_TIMEOUT = 30 # seconds
MAX_AUTO_RELAY_ATTEMPTS = 3
RESERVATION_REFRESH_THRESHOLD = 0.8 # Refresh at 80% of TTL
MAX_CONCURRENT_RESERVATIONS = 2
# Timeout constants for different components
DEFAULT_DISCOVERY_STREAM_TIMEOUT = 10 # seconds
DEFAULT_PEER_PROTOCOL_TIMEOUT = 5 # seconds
DEFAULT_PROTOCOL_READ_TIMEOUT = 15 # seconds
DEFAULT_PROTOCOL_WRITE_TIMEOUT = 15 # seconds
DEFAULT_PROTOCOL_CLOSE_TIMEOUT = 10 # seconds
DEFAULT_DCUTR_READ_TIMEOUT = 30 # seconds
DEFAULT_DCUTR_WRITE_TIMEOUT = 30 # seconds
DEFAULT_DIAL_TIMEOUT = 10 # seconds
@dataclass
class TimeoutConfig:
"""Timeout configuration for different Circuit Relay v2 components."""
# Discovery timeouts
discovery_stream_timeout: int = DEFAULT_DISCOVERY_STREAM_TIMEOUT
peer_protocol_timeout: int = DEFAULT_PEER_PROTOCOL_TIMEOUT
# Core protocol timeouts
protocol_read_timeout: int = DEFAULT_PROTOCOL_READ_TIMEOUT
protocol_write_timeout: int = DEFAULT_PROTOCOL_WRITE_TIMEOUT
protocol_close_timeout: int = DEFAULT_PROTOCOL_CLOSE_TIMEOUT
# DCUtR timeouts
dcutr_read_timeout: int = DEFAULT_DCUTR_READ_TIMEOUT
dcutr_write_timeout: int = DEFAULT_DCUTR_WRITE_TIMEOUT
dial_timeout: int = DEFAULT_DIAL_TIMEOUT
# Relay roles enum
class RelayRole(Flag):
"""
Bit-flag enum that captures the three possible relay capabilities.
A node can combine multiple roles using bit-wise OR, for example::
RelayRole.HOP | RelayRole.STOP
"""
HOP = auto() # Act as a relay for others ("hop")
STOP = auto() # Accept relayed connections ("stop")
CLIENT = auto() # Dial through existing relays ("client")
@dataclass
class RelayConfig:
"""Configuration for Circuit Relay v2."""
# Role configuration
enable_hop: bool = False # Whether to act as a relay (hop)
enable_stop: bool = True # Whether to accept relayed connections (stop)
enable_client: bool = True # Whether to use relays for dialing
# Role configuration (bit-flags)
roles: RelayRole = RelayRole.STOP | RelayRole.CLIENT
# Resource limits
limits: RelayLimits | None = None
# Discovery configuration
bootstrap_relays: list[PeerInfo] = field(default_factory=list)
min_relays: int = 3
max_relays: int = 20
discovery_interval: int = 300 # seconds
min_relays: int = DEFAULT_MIN_RELAYS
max_relays: int = DEFAULT_MAX_RELAYS
discovery_interval: int = DEFAULT_DISCOVERY_INTERVAL
# Connection configuration
reservation_ttl: int = 3600 # seconds
max_circuit_duration: int = 3600 # seconds
max_circuit_bytes: int = 1024 * 1024 * 1024 # 1GB
reservation_ttl: int = DEFAULT_RESERVATION_TTL
max_circuit_duration: int = DEFAULT_MAX_CIRCUIT_DURATION
max_circuit_bytes: int = DEFAULT_MAX_CIRCUIT_BYTES
# Timeout configuration
timeouts: TimeoutConfig = field(default_factory=TimeoutConfig)
# ---------------------------------------------------------------------
# Backwards-compat boolean helpers. Existing code that still accesses
# ``cfg.enable_hop, cfg.enable_stop, cfg.enable_client`` will continue to work.
# ---------------------------------------------------------------------
@property
def enable_hop(self) -> bool: # pragma: no cover helper
return bool(self.roles & RelayRole.HOP)
@property
def enable_stop(self) -> bool: # pragma: no cover helper
return bool(self.roles & RelayRole.STOP)
@property
def enable_client(self) -> bool: # pragma: no cover helper
return bool(self.roles & RelayRole.CLIENT)
def __post_init__(self) -> None:
"""Initialize default values."""
@ -48,8 +138,8 @@ class RelayConfig:
self.limits = RelayLimits(
duration=self.max_circuit_duration,
data=self.max_circuit_bytes,
max_circuit_conns=8,
max_reservations=4,
max_circuit_conns=DEFAULT_MAX_CIRCUIT_CONNS,
max_reservations=DEFAULT_MAX_RESERVATIONS,
)
@ -58,20 +148,20 @@ class HopConfig:
"""Configuration specific to relay (hop) nodes."""
# Resource limits per IP
max_reservations_per_ip: int = 8
max_circuits_per_ip: int = 16
max_reservations_per_ip: int = MAX_RESERVATIONS_PER_IP
max_circuits_per_ip: int = MAX_CIRCUITS_PER_IP
# Rate limiting
reservation_rate_per_ip: int = 4 # per minute
circuit_rate_per_ip: int = 8 # per minute
reservation_rate_per_ip: int = RESERVATION_RATE_PER_IP
circuit_rate_per_ip: int = CIRCUIT_RATE_PER_IP
# Resource quotas
max_circuits_total: int = 64
max_reservations_total: int = 32
max_circuits_total: int = MAX_CIRCUITS_TOTAL
max_reservations_total: int = MAX_RESERVATIONS_TOTAL
# Bandwidth limits
max_bandwidth_per_circuit: int = 1024 * 1024 # 1MB/s
max_bandwidth_total: int = 10 * 1024 * 1024 # 10MB/s
max_bandwidth_per_circuit: int = MAX_BANDWIDTH_PER_CIRCUIT
max_bandwidth_total: int = MAX_BANDWIDTH_TOTAL
@dataclass
@ -79,14 +169,14 @@ class ClientConfig:
"""Configuration specific to relay clients."""
# Relay selection
min_relay_score: float = 0.5
max_relay_latency: float = 1.0 # seconds
min_relay_score: float = MIN_RELAY_SCORE
max_relay_latency: float = MAX_RELAY_LATENCY
# Auto-relay settings
enable_auto_relay: bool = True
auto_relay_timeout: int = 30 # seconds
max_auto_relay_attempts: int = 3
enable_auto_relay: bool = ENABLE_AUTO_RELAY
auto_relay_timeout: int = AUTO_RELAY_TIMEOUT
max_auto_relay_attempts: int = MAX_AUTO_RELAY_ATTEMPTS
# Reservation management
reservation_refresh_threshold: float = 0.8 # Refresh at 80% of TTL
max_concurrent_reservations: int = 2
reservation_refresh_threshold: float = RESERVATION_REFRESH_THRESHOLD
max_concurrent_reservations: int = MAX_CONCURRENT_RESERVATIONS

View File

@ -29,6 +29,11 @@ from libp2p.peer.id import (
from libp2p.peer.peerinfo import (
PeerInfo,
)
from libp2p.relay.circuit_v2.config import (
DEFAULT_DCUTR_READ_TIMEOUT,
DEFAULT_DCUTR_WRITE_TIMEOUT,
DEFAULT_DIAL_TIMEOUT,
)
from libp2p.relay.circuit_v2.nat import (
ReachabilityChecker,
)
@ -47,11 +52,7 @@ PROTOCOL_ID = TProtocol("/libp2p/dcutr")
# Maximum message size for DCUtR (4KiB as per spec)
MAX_MESSAGE_SIZE = 4 * 1024
# Timeouts
STREAM_READ_TIMEOUT = 30 # seconds
STREAM_WRITE_TIMEOUT = 30 # seconds
DIAL_TIMEOUT = 10 # seconds
# DCUtR protocol constants
# Maximum number of hole punch attempts per peer
MAX_HOLE_PUNCH_ATTEMPTS = 5
@ -70,7 +71,13 @@ class DCUtRProtocol(Service):
hole punching, after they have established an initial connection through a relay.
"""
def __init__(self, host: IHost):
def __init__(
self,
host: IHost,
read_timeout: int = DEFAULT_DCUTR_READ_TIMEOUT,
write_timeout: int = DEFAULT_DCUTR_WRITE_TIMEOUT,
dial_timeout: int = DEFAULT_DIAL_TIMEOUT,
):
"""
Initialize the DCUtR protocol.
@ -78,10 +85,19 @@ class DCUtRProtocol(Service):
----------
host : IHost
The libp2p host this protocol is running on
read_timeout : int
Timeout for stream read operations, in seconds
write_timeout : int
Timeout for stream write operations, in seconds
dial_timeout : int
Timeout for dial operations, in seconds
"""
super().__init__()
self.host = host
self.read_timeout = read_timeout
self.write_timeout = write_timeout
self.dial_timeout = dial_timeout
self.event_started = trio.Event()
self._hole_punch_attempts: dict[ID, int] = {}
self._direct_connections: set[ID] = set()
@ -161,7 +177,7 @@ class DCUtRProtocol(Service):
try:
# Read the CONNECT message
with trio.fail_after(STREAM_READ_TIMEOUT):
with trio.fail_after(self.read_timeout):
msg_bytes = await stream.read(MAX_MESSAGE_SIZE)
# Parse the message
@ -196,7 +212,7 @@ class DCUtRProtocol(Service):
response.type = HolePunch.CONNECT
response.ObsAddrs.extend(our_addrs)
with trio.fail_after(STREAM_WRITE_TIMEOUT):
with trio.fail_after(self.write_timeout):
await stream.write(response.SerializeToString())
logger.debug(
@ -206,7 +222,7 @@ class DCUtRProtocol(Service):
)
# Wait for SYNC message
with trio.fail_after(STREAM_READ_TIMEOUT):
with trio.fail_after(self.read_timeout):
sync_bytes = await stream.read(MAX_MESSAGE_SIZE)
# Parse the SYNC message
@ -300,7 +316,7 @@ class DCUtRProtocol(Service):
connect_msg.ObsAddrs.extend(our_addrs)
start_time = time.time()
with trio.fail_after(STREAM_WRITE_TIMEOUT):
with trio.fail_after(self.write_timeout):
await stream.write(connect_msg.SerializeToString())
logger.debug(
@ -310,7 +326,7 @@ class DCUtRProtocol(Service):
)
# Receive the peer's CONNECT message
with trio.fail_after(STREAM_READ_TIMEOUT):
with trio.fail_after(self.read_timeout):
resp_bytes = await stream.read(MAX_MESSAGE_SIZE)
# Calculate RTT
@ -349,7 +365,7 @@ class DCUtRProtocol(Service):
sync_msg = HolePunch()
sync_msg.type = HolePunch.SYNC
with trio.fail_after(STREAM_WRITE_TIMEOUT):
with trio.fail_after(self.write_timeout):
await stream.write(sync_msg.SerializeToString())
logger.debug("Sent SYNC message to %s", peer_id)
@ -468,7 +484,7 @@ class DCUtRProtocol(Service):
peer_info = PeerInfo(peer_id, [addr])
# Try to connect with timeout
with trio.fail_after(DIAL_TIMEOUT):
with trio.fail_after(self.dial_timeout):
await self.host.connect(peer_info)
logger.info("Successfully connected to %s at %s", peer_id, addr)

View File

@ -31,6 +31,11 @@ from libp2p.tools.async_service import (
Service,
)
from .config import (
DEFAULT_DISCOVERY_INTERVAL,
DEFAULT_DISCOVERY_STREAM_TIMEOUT,
DEFAULT_PEER_PROTOCOL_TIMEOUT,
)
from .pb.circuit_pb2 import (
HopMessage,
)
@ -43,10 +48,8 @@ from .protocol_buffer import (
logger = logging.getLogger("libp2p.relay.circuit_v2.discovery")
# Constants
# Discovery constants
MAX_RELAYS_TO_TRACK = 10
DEFAULT_DISCOVERY_INTERVAL = 60 # seconds
STREAM_TIMEOUT = 10 # seconds
# Extended interfaces for type checking
@ -86,6 +89,8 @@ class RelayDiscovery(Service):
auto_reserve: bool = False,
discovery_interval: int = DEFAULT_DISCOVERY_INTERVAL,
max_relays: int = MAX_RELAYS_TO_TRACK,
stream_timeout: int = DEFAULT_DISCOVERY_STREAM_TIMEOUT,
peer_protocol_timeout: int = DEFAULT_PEER_PROTOCOL_TIMEOUT,
) -> None:
"""
Initialize the discovery service.
@ -100,6 +105,10 @@ class RelayDiscovery(Service):
How often to run discovery, in seconds
max_relays : int
Maximum number of relays to track
stream_timeout : int
Timeout for stream operations during discovery, in seconds
peer_protocol_timeout : int
Timeout for checking peer protocol support, in seconds
"""
super().__init__()
@ -107,6 +116,8 @@ class RelayDiscovery(Service):
self.auto_reserve = auto_reserve
self.discovery_interval = discovery_interval
self.max_relays = max_relays
self.stream_timeout = stream_timeout
self.peer_protocol_timeout = peer_protocol_timeout
self._discovered_relays: dict[ID, RelayInfo] = {}
self._protocol_cache: dict[
ID, set[str]
@ -165,8 +176,8 @@ class RelayDiscovery(Service):
self._discovered_relays[peer_id].last_seen = time.time()
continue
# Check if peer supports the relay protocol
with trio.move_on_after(5): # Don't wait too long for protocol info
# Don't wait too long for protocol info
with trio.move_on_after(self.peer_protocol_timeout):
if await self._supports_relay_protocol(peer_id):
await self._add_relay(peer_id)
@ -264,7 +275,7 @@ class RelayDiscovery(Service):
async def _check_via_direct_connection(self, peer_id: ID) -> bool | None:
"""Check protocol support via direct connection."""
try:
with trio.fail_after(STREAM_TIMEOUT):
with trio.fail_after(self.stream_timeout):
stream = await self.host.new_stream(peer_id, [PROTOCOL_ID])
if stream:
await stream.close()
@ -370,7 +381,7 @@ class RelayDiscovery(Service):
# Open a stream to the relay with timeout
try:
with trio.fail_after(STREAM_TIMEOUT):
with trio.fail_after(self.stream_timeout):
stream = await self.host.new_stream(peer_id, [PROTOCOL_ID])
if not stream:
logger.error("Failed to open stream to relay %s", peer_id)
@ -386,7 +397,7 @@ class RelayDiscovery(Service):
peer=self.host.get_id().to_bytes(),
)
with trio.fail_after(STREAM_TIMEOUT):
with trio.fail_after(self.stream_timeout):
await stream.write(request.SerializeToString())
# Wait for response

View File

@ -5,6 +5,7 @@ This module implements the Circuit Relay v2 protocol as specified in:
https://github.com/libp2p/specs/blob/master/relay/circuit-v2.md
"""
from enum import Enum, auto
import logging
import time
from typing import (
@ -37,6 +38,15 @@ from libp2p.tools.async_service import (
Service,
)
from .config import (
DEFAULT_MAX_CIRCUIT_BYTES,
DEFAULT_MAX_CIRCUIT_CONNS,
DEFAULT_MAX_CIRCUIT_DURATION,
DEFAULT_MAX_RESERVATIONS,
DEFAULT_PROTOCOL_CLOSE_TIMEOUT,
DEFAULT_PROTOCOL_READ_TIMEOUT,
DEFAULT_PROTOCOL_WRITE_TIMEOUT,
)
from .pb.circuit_pb2 import (
HopMessage,
Limit,
@ -58,18 +68,22 @@ logger = logging.getLogger("libp2p.relay.circuit_v2")
PROTOCOL_ID = TProtocol("/libp2p/circuit/relay/2.0.0")
STOP_PROTOCOL_ID = TProtocol("/libp2p/circuit/relay/2.0.0/stop")
# Direction enum for data piping
class Pipe(Enum):
SRC_TO_DST = auto()
DST_TO_SRC = auto()
# Default limits for relay resources
DEFAULT_RELAY_LIMITS = RelayLimits(
duration=60 * 60, # 1 hour
data=1024 * 1024 * 1024, # 1GB
max_circuit_conns=8,
max_reservations=4,
duration=DEFAULT_MAX_CIRCUIT_DURATION,
data=DEFAULT_MAX_CIRCUIT_BYTES,
max_circuit_conns=DEFAULT_MAX_CIRCUIT_CONNS,
max_reservations=DEFAULT_MAX_RESERVATIONS,
)
# Stream operation timeouts
STREAM_READ_TIMEOUT = 15 # seconds
STREAM_WRITE_TIMEOUT = 15 # seconds
STREAM_CLOSE_TIMEOUT = 10 # seconds
# Stream operation constants
MAX_READ_RETRIES = 5 # Maximum number of read retries
@ -113,6 +127,9 @@ class CircuitV2Protocol(Service):
host: IHost,
limits: RelayLimits | None = None,
allow_hop: bool = False,
read_timeout: int = DEFAULT_PROTOCOL_READ_TIMEOUT,
write_timeout: int = DEFAULT_PROTOCOL_WRITE_TIMEOUT,
close_timeout: int = DEFAULT_PROTOCOL_CLOSE_TIMEOUT,
) -> None:
"""
Initialize a Circuit Relay v2 protocol instance.
@ -125,11 +142,20 @@ class CircuitV2Protocol(Service):
Resource limits for the relay
allow_hop : bool
Whether to allow this node to act as a relay
read_timeout : int
Timeout for stream read operations, in seconds
write_timeout : int
Timeout for stream write operations, in seconds
close_timeout : int
Timeout for stream close operations, in seconds
"""
self.host = host
self.limits = limits or DEFAULT_RELAY_LIMITS
self.allow_hop = allow_hop
self.read_timeout = read_timeout
self.write_timeout = write_timeout
self.close_timeout = close_timeout
self.resource_manager = RelayResourceManager(self.limits)
self._active_relays: dict[ID, tuple[INetStream, INetStream | None]] = {}
self.event_started = trio.Event()
@ -174,7 +200,7 @@ class CircuitV2Protocol(Service):
return
try:
with trio.fail_after(STREAM_CLOSE_TIMEOUT):
with trio.fail_after(self.close_timeout):
await stream.close()
except Exception:
try:
@ -216,7 +242,7 @@ class CircuitV2Protocol(Service):
while retries < max_retries:
try:
with trio.fail_after(STREAM_READ_TIMEOUT):
with trio.fail_after(self.read_timeout):
# Try reading with timeout
logger.debug(
"Attempting to read from stream (attempt %d/%d)",
@ -293,7 +319,7 @@ class CircuitV2Protocol(Service):
# First, handle the read timeout gracefully
try:
with trio.fail_after(
STREAM_READ_TIMEOUT * 2
self.read_timeout * 2
): # Double the timeout for reading
msg_bytes = await stream.read()
if not msg_bytes:
@ -414,7 +440,7 @@ class CircuitV2Protocol(Service):
"""
try:
# Read the incoming message with timeout
with trio.fail_after(STREAM_READ_TIMEOUT):
with trio.fail_after(self.read_timeout):
msg_bytes = await stream.read()
stop_msg = StopMessage()
stop_msg.ParseFromString(msg_bytes)
@ -458,8 +484,20 @@ class CircuitV2Protocol(Service):
# Start relaying data
async with trio.open_nursery() as nursery:
nursery.start_soon(self._relay_data, src_stream, stream, peer_id)
nursery.start_soon(self._relay_data, stream, src_stream, peer_id)
nursery.start_soon(
self._relay_data,
src_stream,
stream,
peer_id,
Pipe.SRC_TO_DST,
)
nursery.start_soon(
self._relay_data,
stream,
src_stream,
peer_id,
Pipe.DST_TO_SRC,
)
except trio.TooSlowError:
logger.error("Timeout reading from stop stream")
@ -509,7 +547,7 @@ class CircuitV2Protocol(Service):
ttl = self.resource_manager.reserve(peer_id)
# Send reservation success response
with trio.fail_after(STREAM_WRITE_TIMEOUT):
with trio.fail_after(self.write_timeout):
status = create_status(
code=StatusCode.OK, message="Reservation accepted"
)
@ -560,7 +598,7 @@ class CircuitV2Protocol(Service):
# Always close the stream when done with reservation
if cast(INetStreamWithExtras, stream).is_open():
try:
with trio.fail_after(STREAM_CLOSE_TIMEOUT):
with trio.fail_after(self.close_timeout):
await stream.close()
except Exception as close_err:
logger.error("Error closing stream: %s", str(close_err))
@ -596,7 +634,7 @@ class CircuitV2Protocol(Service):
self._active_relays[peer_id] = (stream, None)
# Try to connect to the destination with timeout
with trio.fail_after(STREAM_READ_TIMEOUT):
with trio.fail_after(self.read_timeout):
dst_stream = await self.host.new_stream(peer_id, [STOP_PROTOCOL_ID])
if not dst_stream:
raise ConnectionError("Could not connect to destination")
@ -648,8 +686,20 @@ class CircuitV2Protocol(Service):
# Start relaying data
async with trio.open_nursery() as nursery:
nursery.start_soon(self._relay_data, stream, dst_stream, peer_id)
nursery.start_soon(self._relay_data, dst_stream, stream, peer_id)
nursery.start_soon(
self._relay_data,
stream,
dst_stream,
peer_id,
Pipe.SRC_TO_DST,
)
nursery.start_soon(
self._relay_data,
dst_stream,
stream,
peer_id,
Pipe.DST_TO_SRC,
)
except (trio.TooSlowError, ConnectionError) as e:
logger.error("Error establishing relay connection: %s", str(e))
@ -685,6 +735,7 @@ class CircuitV2Protocol(Service):
src_stream: INetStream,
dst_stream: INetStream,
peer_id: ID,
direction: Pipe,
) -> None:
"""
Relay data between two streams.
@ -698,24 +749,27 @@ class CircuitV2Protocol(Service):
peer_id : ID
ID of the peer being relayed
direction : Pipe
Direction of data flow (``Pipe.SRC_TO_DST`` or ``Pipe.DST_TO_SRC``)
"""
try:
while True:
# Read data with retries
data = await self._read_stream_with_retry(src_stream)
if not data:
logger.info("Source stream closed/reset")
logger.info("%s closed/reset", direction.name)
break
# Write data with timeout
try:
with trio.fail_after(STREAM_WRITE_TIMEOUT):
with trio.fail_after(self.write_timeout):
await dst_stream.write(data)
except trio.TooSlowError:
logger.error("Timeout writing to destination stream")
logger.error("Timeout writing in %s", direction.name)
break
except Exception as e:
logger.error("Error writing to destination stream: %s", str(e))
logger.error("Error writing in %s: %s", direction.name, str(e))
break
# Update resource usage
@ -744,7 +798,7 @@ class CircuitV2Protocol(Service):
"""Send a status message."""
try:
logger.debug("Sending status message with code %s: %s", code, message)
with trio.fail_after(STREAM_WRITE_TIMEOUT * 2): # Double the timeout
with trio.fail_after(self.write_timeout * 2): # Double the timeout
# Create a proto Status directly
pb_status = PbStatus()
pb_status.code = cast(
@ -782,7 +836,7 @@ class CircuitV2Protocol(Service):
"""Send a status message on a STOP stream."""
try:
logger.debug("Sending stop status message with code %s: %s", code, message)
with trio.fail_after(STREAM_WRITE_TIMEOUT * 2): # Double the timeout
with trio.fail_after(self.write_timeout * 2): # Double the timeout
# Create a proto Status directly
pb_status = PbStatus()
pb_status.code = cast(

View File

@ -8,6 +8,7 @@ including reservations and connection limits.
from dataclasses import (
dataclass,
)
from enum import Enum, auto
import hashlib
import os
import time
@ -19,6 +20,18 @@ from libp2p.peer.id import (
# Import the protobuf definitions
from .pb.circuit_pb2 import Reservation as PbReservation
RANDOM_BYTES_LENGTH = 16 # 128 bits of randomness
TIMESTAMP_MULTIPLIER = 1000000 # To convert seconds to microseconds
# Reservation status enum
class ReservationStatus(Enum):
"""Lifecycle status of a relay reservation."""
ACTIVE = auto()
EXPIRED = auto()
REJECTED = auto()
@dataclass
class RelayLimits:
@ -68,8 +81,8 @@ class Reservation:
# - Peer ID to bind it to the specific peer
# - Timestamp for uniqueness
# - Hash everything for a fixed size output
random_bytes = os.urandom(16) # 128 bits of randomness
timestamp = str(int(self.created_at * 1000000)).encode()
random_bytes = os.urandom(RANDOM_BYTES_LENGTH)
timestamp = str(int(self.created_at * TIMESTAMP_MULTIPLIER)).encode()
peer_bytes = self.peer_id.to_bytes()
# Combine all elements and hash them
@ -84,6 +97,15 @@ class Reservation:
"""Check if the reservation has expired."""
return time.time() > self.expires_at
# Expose a friendly status enum
@property
def status(self) -> ReservationStatus:
"""Return the current status as a ``ReservationStatus`` enum."""
return (
ReservationStatus.EXPIRED if self.is_expired() else ReservationStatus.ACTIVE
)
def can_accept_connection(self) -> bool:
"""Check if a new connection can be accepted."""
return (

View File

@ -89,6 +89,8 @@ class CircuitV2Transport(ITransport):
auto_reserve=config.enable_client,
discovery_interval=config.discovery_interval,
max_relays=config.max_relays,
stream_timeout=config.timeouts.discovery_stream_timeout,
peer_protocol_timeout=config.timeouts.peer_protocol_timeout,
)
self.relay_counter = 0 # for round robin load balancing

View File

@ -1,3 +1,4 @@
import logging
from typing import (
cast,
)
@ -15,6 +16,8 @@ from libp2p.io.msgio import (
FixedSizeLenMsgReadWriter,
)
logger = logging.getLogger(__name__)
SIZE_NOISE_MESSAGE_LEN = 2
MAX_NOISE_MESSAGE_LEN = 2 ** (8 * SIZE_NOISE_MESSAGE_LEN) - 1
SIZE_NOISE_MESSAGE_BODY_LEN = 2
@ -50,18 +53,25 @@ class BaseNoiseMsgReadWriter(EncryptedMsgReadWriter):
self.noise_state = noise_state
async def write_msg(self, msg: bytes, prefix_encoded: bool = False) -> None:
logger.debug(f"Noise write_msg: encrypting {len(msg)} bytes")
data_encrypted = self.encrypt(msg)
if prefix_encoded:
# Manually add the prefix if needed
data_encrypted = self.prefix + data_encrypted
logger.debug(f"Noise write_msg: writing {len(data_encrypted)} encrypted bytes")
await self.read_writer.write_msg(data_encrypted)
logger.debug("Noise write_msg: write completed successfully")
async def read_msg(self, prefix_encoded: bool = False) -> bytes:
logger.debug("Noise read_msg: reading encrypted message")
noise_msg_encrypted = await self.read_writer.read_msg()
logger.debug(f"Noise read_msg: read {len(noise_msg_encrypted)} encrypted bytes")
if prefix_encoded:
return self.decrypt(noise_msg_encrypted[len(self.prefix) :])
result = self.decrypt(noise_msg_encrypted[len(self.prefix) :])
else:
return self.decrypt(noise_msg_encrypted)
result = self.decrypt(noise_msg_encrypted)
logger.debug(f"Noise read_msg: decrypted to {len(result)} bytes")
return result
async def close(self) -> None:
await self.read_writer.close()

View File

@ -1,6 +1,7 @@
from dataclasses import (
dataclass,
)
import logging
from libp2p.crypto.keys import (
PrivateKey,
@ -12,6 +13,8 @@ from libp2p.crypto.serialization import (
from .pb import noise_pb2 as noise_pb
logger = logging.getLogger(__name__)
SIGNED_DATA_PREFIX = "noise-libp2p-static-key:"
@ -48,6 +51,8 @@ def make_handshake_payload_sig(
id_privkey: PrivateKey, noise_static_pubkey: PublicKey
) -> bytes:
data = make_data_to_be_signed(noise_static_pubkey)
logger.debug(f"make_handshake_payload_sig: signing data length: {len(data)}")
logger.debug(f"make_handshake_payload_sig: signing data hex: {data.hex()}")
return id_privkey.sign(data)
@ -60,4 +65,27 @@ def verify_handshake_payload_sig(
2. signed by the private key corresponding to `id_pubkey`
"""
expected_data = make_data_to_be_signed(noise_static_pubkey)
return payload.id_pubkey.verify(expected_data, payload.id_sig)
logger.debug(
f"verify_handshake_payload_sig: payload.id_pubkey type: "
f"{type(payload.id_pubkey)}"
)
logger.debug(
f"verify_handshake_payload_sig: noise_static_pubkey type: "
f"{type(noise_static_pubkey)}"
)
logger.debug(
f"verify_handshake_payload_sig: expected_data length: {len(expected_data)}"
)
logger.debug(
f"verify_handshake_payload_sig: expected_data hex: {expected_data.hex()}"
)
logger.debug(
f"verify_handshake_payload_sig: payload.id_sig length: {len(payload.id_sig)}"
)
try:
result = payload.id_pubkey.verify(expected_data, payload.id_sig)
logger.debug(f"verify_handshake_payload_sig: verification result: {result}")
return result
except Exception as e:
logger.error(f"verify_handshake_payload_sig: verification exception: {e}")
return False

View File

@ -2,6 +2,7 @@ from abc import (
ABC,
abstractmethod,
)
import logging
from cryptography.hazmat.primitives import (
serialization,
@ -46,6 +47,8 @@ from .messages import (
verify_handshake_payload_sig,
)
logger = logging.getLogger(__name__)
class IPattern(ABC):
@abstractmethod
@ -95,6 +98,7 @@ class PatternXX(BasePattern):
self.early_data = early_data
async def handshake_inbound(self, conn: IRawConnection) -> ISecureConn:
logger.debug(f"Noise XX handshake_inbound started for peer {self.local_peer}")
noise_state = self.create_noise_state()
noise_state.set_as_responder()
noise_state.start_handshake()
@ -107,15 +111,22 @@ class PatternXX(BasePattern):
read_writer = NoiseHandshakeReadWriter(conn, noise_state)
# Consume msg#1.
logger.debug("Noise XX handshake_inbound: reading msg#1")
await read_writer.read_msg()
logger.debug("Noise XX handshake_inbound: read msg#1 successfully")
# Send msg#2, which should include our handshake payload.
logger.debug("Noise XX handshake_inbound: preparing msg#2")
our_payload = self.make_handshake_payload()
msg_2 = our_payload.serialize()
logger.debug(f"Noise XX handshake_inbound: sending msg#2 ({len(msg_2)} bytes)")
await read_writer.write_msg(msg_2)
logger.debug("Noise XX handshake_inbound: sent msg#2 successfully")
# Receive and consume msg#3.
logger.debug("Noise XX handshake_inbound: reading msg#3")
msg_3 = await read_writer.read_msg()
logger.debug(f"Noise XX handshake_inbound: read msg#3 ({len(msg_3)} bytes)")
peer_handshake_payload = NoiseHandshakePayload.deserialize(msg_3)
if handshake_state.rs is None:
@ -147,6 +158,7 @@ class PatternXX(BasePattern):
async def handshake_outbound(
self, conn: IRawConnection, remote_peer: ID
) -> ISecureConn:
logger.debug(f"Noise XX handshake_outbound started to peer {remote_peer}")
noise_state = self.create_noise_state()
read_writer = NoiseHandshakeReadWriter(conn, noise_state)
@ -159,11 +171,15 @@ class PatternXX(BasePattern):
raise NoiseStateError("Handshake state is not initialized")
# Send msg#1, which is *not* encrypted.
logger.debug("Noise XX handshake_outbound: sending msg#1")
msg_1 = b""
await read_writer.write_msg(msg_1)
logger.debug("Noise XX handshake_outbound: sent msg#1 successfully")
# Read msg#2 from the remote, which contains the public key of the peer.
logger.debug("Noise XX handshake_outbound: reading msg#2")
msg_2 = await read_writer.read_msg()
logger.debug(f"Noise XX handshake_outbound: read msg#2 ({len(msg_2)} bytes)")
peer_handshake_payload = NoiseHandshakePayload.deserialize(msg_2)
if handshake_state.rs is None:
@ -174,8 +190,27 @@ class PatternXX(BasePattern):
)
remote_pubkey = self._get_pubkey_from_noise_keypair(handshake_state.rs)
logger.debug(
f"Noise XX handshake_outbound: verifying signature for peer {remote_peer}"
)
logger.debug(
f"Noise XX handshake_outbound: remote_pubkey type: {type(remote_pubkey)}"
)
id_pubkey_repr = peer_handshake_payload.id_pubkey.to_bytes().hex()
logger.debug(
f"Noise XX handshake_outbound: peer_handshake_payload.id_pubkey: "
f"{id_pubkey_repr}"
)
if not verify_handshake_payload_sig(peer_handshake_payload, remote_pubkey):
logger.error(
f"Noise XX handshake_outbound: signature verification failed for peer "
f"{remote_peer}"
)
raise InvalidSignature
logger.debug(
f"Noise XX handshake_outbound: signature verification successful for peer "
f"{remote_peer}"
)
remote_peer_id_from_pubkey = ID.from_pubkey(peer_handshake_payload.id_pubkey)
if remote_peer_id_from_pubkey != remote_peer:
raise PeerIDMismatchesPubkey(

View File

@ -1,5 +1,3 @@
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from types import (
TracebackType,
)
@ -15,6 +13,7 @@ from libp2p.abc import (
from libp2p.stream_muxer.exceptions import (
MuxedConnUnavailable,
)
from libp2p.stream_muxer.rw_lock import ReadWriteLock
from .constants import (
HeaderTags,
@ -34,72 +33,6 @@ if TYPE_CHECKING:
)
class ReadWriteLock:
"""
A read-write lock that allows multiple concurrent readers
or one exclusive writer, implemented using Trio primitives.
"""
def __init__(self) -> None:
self._readers = 0
self._readers_lock = trio.Lock() # Protects access to _readers count
self._writer_lock = trio.Semaphore(1) # Allows only one writer at a time
async def acquire_read(self) -> None:
"""Acquire a read lock. Multiple readers can hold it simultaneously."""
try:
async with self._readers_lock:
if self._readers == 0:
await self._writer_lock.acquire()
self._readers += 1
except trio.Cancelled:
raise
async def release_read(self) -> None:
"""Release a read lock."""
async with self._readers_lock:
if self._readers == 1:
self._writer_lock.release()
self._readers -= 1
async def acquire_write(self) -> None:
"""Acquire an exclusive write lock."""
try:
await self._writer_lock.acquire()
except trio.Cancelled:
raise
def release_write(self) -> None:
"""Release the exclusive write lock."""
self._writer_lock.release()
@asynccontextmanager
async def read_lock(self) -> AsyncGenerator[None, None]:
"""Context manager for acquiring and releasing a read lock safely."""
acquire = False
try:
await self.acquire_read()
acquire = True
yield
finally:
if acquire:
with trio.CancelScope() as scope:
scope.shield = True
await self.release_read()
@asynccontextmanager
async def write_lock(self) -> AsyncGenerator[None, None]:
"""Context manager for acquiring and releasing a write lock safely."""
acquire = False
try:
await self.acquire_write()
acquire = True
yield
finally:
if acquire:
self.release_write()
class MplexStream(IMuxedStream):
"""
reference: https://github.com/libp2p/go-mplex/blob/master/stream.go

View File

@ -21,6 +21,7 @@ from libp2p.protocol_muxer.exceptions import (
MultiselectError,
)
from libp2p.protocol_muxer.multiselect import (
DEFAULT_NEGOTIATE_TIMEOUT,
Multiselect,
)
from libp2p.protocol_muxer.multiselect_client import (
@ -46,11 +47,17 @@ class MuxerMultistream:
transports: "OrderedDict[TProtocol, TMuxerClass]"
multiselect: Multiselect
multiselect_client: MultiselectClient
negotiate_timeout: int
def __init__(self, muxer_transports_by_protocol: TMuxerOptions) -> None:
def __init__(
self,
muxer_transports_by_protocol: TMuxerOptions,
negotiate_timeout: int = DEFAULT_NEGOTIATE_TIMEOUT,
) -> None:
self.transports = OrderedDict()
self.multiselect = Multiselect()
self.multistream_client = MultiselectClient()
self.negotiate_timeout = negotiate_timeout
for protocol, transport in muxer_transports_by_protocol.items():
self.add_transport(protocol, transport)
@ -80,10 +87,12 @@ class MuxerMultistream:
communicator = MultiselectCommunicator(conn)
if conn.is_initiator:
protocol = await self.multiselect_client.select_one_of(
tuple(self.transports.keys()), communicator
tuple(self.transports.keys()), communicator, self.negotiate_timeout
)
else:
protocol, _ = await self.multiselect.negotiate(communicator)
protocol, _ = await self.multiselect.negotiate(
communicator, self.negotiate_timeout
)
if protocol is None:
raise MultiselectError(
"Fail to negotiate a stream muxer protocol: no protocol selected"
@ -93,7 +102,7 @@ class MuxerMultistream:
async def new_conn(self, conn: ISecureConn, peer_id: ID) -> IMuxedConn:
communicator = MultiselectCommunicator(conn)
protocol = await self.multistream_client.select_one_of(
tuple(self.transports.keys()), communicator
tuple(self.transports.keys()), communicator, self.negotiate_timeout
)
transport_class = self.transports[protocol]
if protocol == PROTOCOL_ID:

View File

@ -0,0 +1,70 @@
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
import trio
class ReadWriteLock:
"""
A read-write lock that allows multiple concurrent readers
or one exclusive writer, implemented using Trio primitives.
"""
def __init__(self) -> None:
self._readers = 0
self._readers_lock = trio.Lock() # Protects access to _readers count
self._writer_lock = trio.Semaphore(1) # Allows only one writer at a time
async def acquire_read(self) -> None:
"""Acquire a read lock. Multiple readers can hold it simultaneously."""
try:
async with self._readers_lock:
if self._readers == 0:
await self._writer_lock.acquire()
self._readers += 1
except trio.Cancelled:
raise
async def release_read(self) -> None:
"""Release a read lock."""
async with self._readers_lock:
if self._readers == 1:
self._writer_lock.release()
self._readers -= 1
async def acquire_write(self) -> None:
"""Acquire an exclusive write lock."""
try:
await self._writer_lock.acquire()
except trio.Cancelled:
raise
def release_write(self) -> None:
"""Release the exclusive write lock."""
self._writer_lock.release()
@asynccontextmanager
async def read_lock(self) -> AsyncGenerator[None, None]:
"""Context manager for acquiring and releasing a read lock safely."""
acquire = False
try:
await self.acquire_read()
acquire = True
yield
finally:
if acquire:
with trio.CancelScope() as scope:
scope.shield = True
await self.release_read()
@asynccontextmanager
async def write_lock(self) -> AsyncGenerator[None, None]:
"""Context manager for acquiring and releasing a write lock safely."""
acquire = False
try:
await self.acquire_write()
acquire = True
yield
finally:
if acquire:
self.release_write()

View File

@ -44,6 +44,7 @@ from libp2p.stream_muxer.exceptions import (
MuxedStreamError,
MuxedStreamReset,
)
from libp2p.stream_muxer.rw_lock import ReadWriteLock
# Configure logger for this module
logger = logging.getLogger("libp2p.stream_muxer.yamux")
@ -80,6 +81,8 @@ class YamuxStream(IMuxedStream):
self.send_window = DEFAULT_WINDOW_SIZE
self.recv_window = DEFAULT_WINDOW_SIZE
self.window_lock = trio.Lock()
self.rw_lock = ReadWriteLock()
self.close_lock = trio.Lock()
async def __aenter__(self) -> "YamuxStream":
"""Enter the async context manager."""
@ -95,52 +98,54 @@ class YamuxStream(IMuxedStream):
await self.close()
async def write(self, data: bytes) -> None:
if self.send_closed:
raise MuxedStreamError("Stream is closed for sending")
async with self.rw_lock.write_lock():
if self.send_closed:
raise MuxedStreamError("Stream is closed for sending")
# Flow control: Check if we have enough send window
total_len = len(data)
sent = 0
logger.debug(f"Stream {self.stream_id}: Starts writing {total_len} bytes ")
while sent < total_len:
# Wait for available window with timeout
timeout = False
async with self.window_lock:
if self.send_window == 0:
logger.debug(
f"Stream {self.stream_id}: Window is zero, waiting for update"
# Flow control: Check if we have enough send window
total_len = len(data)
sent = 0
logger.debug(f"Stream {self.stream_id}: Starts writing {total_len} bytes ")
while sent < total_len:
# Wait for available window with timeout
timeout = False
async with self.window_lock:
if self.send_window == 0:
logger.debug(
f"Stream {self.stream_id}: "
"Window is zero, waiting for update"
)
# Release lock and wait with timeout
self.window_lock.release()
# To avoid re-acquiring the lock immediately,
with trio.move_on_after(5.0) as cancel_scope:
while self.send_window == 0 and not self.closed:
await trio.sleep(0.01)
# If we timed out, cancel the scope
timeout = cancel_scope.cancelled_caught
# Re-acquire lock
await self.window_lock.acquire()
# If we timed out waiting for window update, raise an error
if timeout:
raise MuxedStreamError(
"Timed out waiting for window update after 5 seconds."
)
if self.closed:
raise MuxedStreamError("Stream is closed")
# Calculate how much we can send now
to_send = min(self.send_window, total_len - sent)
chunk = data[sent : sent + to_send]
self.send_window -= to_send
# Send the data
header = struct.pack(
YAMUX_HEADER_FORMAT, 0, TYPE_DATA, 0, self.stream_id, len(chunk)
)
# Release lock and wait with timeout
self.window_lock.release()
# To avoid re-acquiring the lock immediately,
with trio.move_on_after(5.0) as cancel_scope:
while self.send_window == 0 and not self.closed:
await trio.sleep(0.01)
# If we timed out, cancel the scope
timeout = cancel_scope.cancelled_caught
# Re-acquire lock
await self.window_lock.acquire()
# If we timed out waiting for window update, raise an error
if timeout:
raise MuxedStreamError(
"Timed out waiting for window update after 5 seconds."
)
if self.closed:
raise MuxedStreamError("Stream is closed")
# Calculate how much we can send now
to_send = min(self.send_window, total_len - sent)
chunk = data[sent : sent + to_send]
self.send_window -= to_send
# Send the data
header = struct.pack(
YAMUX_HEADER_FORMAT, 0, TYPE_DATA, 0, self.stream_id, len(chunk)
)
await self.conn.secured_conn.write(header + chunk)
sent += to_send
await self.conn.secured_conn.write(header + chunk)
sent += to_send
async def send_window_update(self, increment: int, skip_lock: bool = False) -> None:
"""
@ -257,30 +262,32 @@ class YamuxStream(IMuxedStream):
return data
async def close(self) -> None:
if not self.send_closed:
logger.debug(f"Half-closing stream {self.stream_id} (local end)")
header = struct.pack(
YAMUX_HEADER_FORMAT, 0, TYPE_DATA, FLAG_FIN, self.stream_id, 0
)
await self.conn.secured_conn.write(header)
self.send_closed = True
async with self.close_lock:
if not self.send_closed:
logger.debug(f"Half-closing stream {self.stream_id} (local end)")
header = struct.pack(
YAMUX_HEADER_FORMAT, 0, TYPE_DATA, FLAG_FIN, self.stream_id, 0
)
await self.conn.secured_conn.write(header)
self.send_closed = True
# Only set fully closed if both directions are closed
if self.send_closed and self.recv_closed:
self.closed = True
else:
# Stream is half-closed but not fully closed
self.closed = False
# Only set fully closed if both directions are closed
if self.send_closed and self.recv_closed:
self.closed = True
else:
# Stream is half-closed but not fully closed
self.closed = False
async def reset(self) -> None:
if not self.closed:
logger.debug(f"Resetting stream {self.stream_id}")
header = struct.pack(
YAMUX_HEADER_FORMAT, 0, TYPE_DATA, FLAG_RST, self.stream_id, 0
)
await self.conn.secured_conn.write(header)
self.closed = True
self.reset_received = True # Mark as reset
async with self.close_lock:
logger.debug(f"Resetting stream {self.stream_id}")
header = struct.pack(
YAMUX_HEADER_FORMAT, 0, TYPE_DATA, FLAG_RST, self.stream_id, 0
)
await self.conn.secured_conn.write(header)
self.closed = True
self.reset_received = True # Mark as reset
def set_deadline(self, ttl: int) -> bool:
"""

View File

@ -0,0 +1,57 @@
from typing import Any
from .tcp.tcp import TCP
from .websocket.transport import WebsocketTransport
from .transport_registry import (
TransportRegistry,
create_transport_for_multiaddr,
get_transport_registry,
register_transport,
get_supported_transport_protocols,
)
from .upgrader import TransportUpgrader
from libp2p.abc import ITransport
def create_transport(protocol: str, upgrader: TransportUpgrader | None = None, **kwargs: Any) -> ITransport:
"""
Convenience function to create a transport instance.
:param protocol: The transport protocol ("tcp", "ws", "wss", or custom)
:param upgrader: Optional transport upgrader (required for WebSocket)
:param kwargs: Additional arguments for transport construction (e.g., tls_client_config, tls_server_config)
:return: Transport instance
"""
# First check if it's a built-in protocol
if protocol in ["ws", "wss"]:
if upgrader is None:
raise ValueError(f"WebSocket transport requires an upgrader")
return WebsocketTransport(
upgrader,
tls_client_config=kwargs.get("tls_client_config"),
tls_server_config=kwargs.get("tls_server_config"),
handshake_timeout=kwargs.get("handshake_timeout", 15.0)
)
elif protocol == "tcp":
return TCP()
else:
# Check if it's a custom registered transport
registry = get_transport_registry()
transport_class = registry.get_transport(protocol)
if transport_class:
transport = registry.create_transport(protocol, upgrader, **kwargs)
if transport is None:
raise ValueError(f"Failed to create transport for protocol: {protocol}")
return transport
else:
raise ValueError(f"Unsupported transport protocol: {protocol}")
__all__ = [
"TCP",
"WebsocketTransport",
"TransportRegistry",
"create_transport_for_multiaddr",
"create_transport",
"get_transport_registry",
"register_transport",
"get_supported_transport_protocols",
]

View File

View File

@ -0,0 +1,345 @@
"""
Configuration classes for QUIC transport.
"""
from dataclasses import (
dataclass,
field,
)
import ssl
from typing import Any, Literal, TypedDict
from libp2p.custom_types import TProtocol
from libp2p.network.config import ConnectionConfig
class QUICTransportKwargs(TypedDict, total=False):
"""Type definition for kwargs accepted by new_transport function."""
# Connection settings
idle_timeout: float
max_datagram_size: int
local_port: int | None
# Protocol version support
enable_draft29: bool
enable_v1: bool
# TLS settings
verify_mode: ssl.VerifyMode
alpn_protocols: list[str]
# Performance settings
max_concurrent_streams: int
connection_window: int
stream_window: int
# Logging and debugging
enable_qlog: bool
qlog_dir: str | None
# Connection management
max_connections: int
connection_timeout: float
# Protocol identifiers
PROTOCOL_QUIC_V1: TProtocol
PROTOCOL_QUIC_DRAFT29: TProtocol
@dataclass
class QUICTransportConfig(ConnectionConfig):
"""Configuration for QUIC transport."""
# Connection settings
idle_timeout: float = 30.0 # Seconds before an idle connection is closed.
max_datagram_size: int = (
1200 # Maximum size of UDP datagrams to avoid IP fragmentation.
)
local_port: int | None = (
None # Local port to bind to. If None, a random port is chosen.
)
# Protocol version support
enable_draft29: bool = True # Enable QUIC draft-29 for compatibility
enable_v1: bool = True # Enable QUIC v1 (RFC 9000)
# TLS settings
verify_mode: ssl.VerifyMode = ssl.CERT_NONE
alpn_protocols: list[str] = field(default_factory=lambda: ["libp2p"])
# Performance settings
max_concurrent_streams: int = 100 # Maximum concurrent streams per connection
connection_window: int = 1024 * 1024 # Connection flow control window
stream_window: int = 64 * 1024 # Stream flow control window
# Logging and debugging
enable_qlog: bool = False # Enable QUIC logging
qlog_dir: str | None = None # Directory for QUIC logs
# Connection management
max_connections: int = 1000 # Maximum number of connections
connection_timeout: float = 10.0 # Connection establishment timeout
MAX_CONCURRENT_STREAMS: int = 1000
"""Maximum number of concurrent streams per connection."""
MAX_INCOMING_STREAMS: int = 1000
"""Maximum number of incoming streams per connection."""
CONNECTION_HANDSHAKE_TIMEOUT: float = 60.0
"""Timeout for connection handshake (seconds)."""
MAX_OUTGOING_STREAMS: int = 1000
"""Maximum number of outgoing streams per connection."""
CONNECTION_CLOSE_TIMEOUT: int = 10
"""Timeout for opening new connection (seconds)."""
# Stream timeouts
STREAM_OPEN_TIMEOUT: float = 5.0
"""Timeout for opening new streams (seconds)."""
STREAM_ACCEPT_TIMEOUT: float = 30.0
"""Timeout for accepting incoming streams (seconds)."""
STREAM_READ_TIMEOUT: float = 30.0
"""Default timeout for stream read operations (seconds)."""
STREAM_WRITE_TIMEOUT: float = 30.0
"""Default timeout for stream write operations (seconds)."""
STREAM_CLOSE_TIMEOUT: float = 10.0
"""Timeout for graceful stream close (seconds)."""
# Flow control configuration
STREAM_FLOW_CONTROL_WINDOW: int = 1024 * 1024 # 1MB
"""Per-stream flow control window size."""
CONNECTION_FLOW_CONTROL_WINDOW: int = 1536 * 1024 # 1.5MB
"""Connection-wide flow control window size."""
# Buffer management
MAX_STREAM_RECEIVE_BUFFER: int = 2 * 1024 * 1024 # 2MB
"""Maximum receive buffer size per stream."""
STREAM_RECEIVE_BUFFER_LOW_WATERMARK: int = 64 * 1024 # 64KB
"""Low watermark for stream receive buffer."""
STREAM_RECEIVE_BUFFER_HIGH_WATERMARK: int = 512 * 1024 # 512KB
"""High watermark for stream receive buffer."""
# Stream lifecycle configuration
ENABLE_STREAM_RESET_ON_ERROR: bool = True
"""Whether to automatically reset streams on errors."""
STREAM_RESET_ERROR_CODE: int = 1
"""Default error code for stream resets."""
ENABLE_STREAM_KEEP_ALIVE: bool = False
"""Whether to enable stream keep-alive mechanisms."""
STREAM_KEEP_ALIVE_INTERVAL: float = 30.0
"""Interval for stream keep-alive pings (seconds)."""
# Resource management
ENABLE_STREAM_RESOURCE_TRACKING: bool = True
"""Whether to track stream resource usage."""
STREAM_MEMORY_LIMIT_PER_STREAM: int = 2 * 1024 * 1024 # 2MB
"""Memory limit per individual stream."""
STREAM_MEMORY_LIMIT_PER_CONNECTION: int = 100 * 1024 * 1024 # 100MB
"""Total memory limit for all streams per connection."""
# Concurrency and performance
ENABLE_STREAM_BATCHING: bool = True
"""Whether to batch multiple stream operations."""
STREAM_BATCH_SIZE: int = 10
"""Number of streams to process in a batch."""
STREAM_PROCESSING_CONCURRENCY: int = 100
"""Maximum concurrent stream processing tasks."""
# Debugging and monitoring
ENABLE_STREAM_METRICS: bool = True
"""Whether to collect stream metrics."""
ENABLE_STREAM_TIMELINE_TRACKING: bool = True
"""Whether to track stream lifecycle timelines."""
STREAM_METRICS_COLLECTION_INTERVAL: float = 60.0
"""Interval for collecting stream metrics (seconds)."""
# Error handling configuration
STREAM_ERROR_RETRY_ATTEMPTS: int = 3
"""Number of retry attempts for recoverable stream errors."""
STREAM_ERROR_RETRY_DELAY: float = 1.0
"""Initial delay between stream error retries (seconds)."""
STREAM_ERROR_RETRY_BACKOFF_FACTOR: float = 2.0
"""Backoff factor for stream error retries."""
# Protocol identifiers matching go-libp2p
PROTOCOL_QUIC_V1: TProtocol = TProtocol("quic-v1") # RFC 9000
PROTOCOL_QUIC_DRAFT29: TProtocol = TProtocol("quic") # draft-29
def __post_init__(self) -> None:
"""Validate configuration after initialization."""
if not (self.enable_draft29 or self.enable_v1):
raise ValueError("At least one QUIC version must be enabled")
if self.idle_timeout <= 0:
raise ValueError("Idle timeout must be positive")
if self.max_datagram_size < 1200:
raise ValueError("Max datagram size must be at least 1200 bytes")
# Validate timeouts
timeout_fields = [
"STREAM_OPEN_TIMEOUT",
"STREAM_ACCEPT_TIMEOUT",
"STREAM_READ_TIMEOUT",
"STREAM_WRITE_TIMEOUT",
"STREAM_CLOSE_TIMEOUT",
]
for timeout_field in timeout_fields:
if getattr(self, timeout_field) <= 0:
raise ValueError(f"{timeout_field} must be positive")
# Validate flow control windows
if self.STREAM_FLOW_CONTROL_WINDOW <= 0:
raise ValueError("STREAM_FLOW_CONTROL_WINDOW must be positive")
if self.CONNECTION_FLOW_CONTROL_WINDOW < self.STREAM_FLOW_CONTROL_WINDOW:
raise ValueError(
"CONNECTION_FLOW_CONTROL_WINDOW must be >= STREAM_FLOW_CONTROL_WINDOW"
)
# Validate buffer sizes
if self.MAX_STREAM_RECEIVE_BUFFER <= 0:
raise ValueError("MAX_STREAM_RECEIVE_BUFFER must be positive")
if self.STREAM_RECEIVE_BUFFER_HIGH_WATERMARK > self.MAX_STREAM_RECEIVE_BUFFER:
raise ValueError(
"STREAM_RECEIVE_BUFFER_HIGH_WATERMARK cannot".__add__(
"exceed MAX_STREAM_RECEIVE_BUFFER"
)
)
if (
self.STREAM_RECEIVE_BUFFER_LOW_WATERMARK
>= self.STREAM_RECEIVE_BUFFER_HIGH_WATERMARK
):
raise ValueError(
"STREAM_RECEIVE_BUFFER_LOW_WATERMARK must be < HIGH_WATERMARK"
)
# Validate memory limits
if self.STREAM_MEMORY_LIMIT_PER_STREAM <= 0:
raise ValueError("STREAM_MEMORY_LIMIT_PER_STREAM must be positive")
if self.STREAM_MEMORY_LIMIT_PER_CONNECTION <= 0:
raise ValueError("STREAM_MEMORY_LIMIT_PER_CONNECTION must be positive")
expected_stream_memory = (
self.MAX_CONCURRENT_STREAMS * self.STREAM_MEMORY_LIMIT_PER_STREAM
)
if expected_stream_memory > self.STREAM_MEMORY_LIMIT_PER_CONNECTION * 2:
# Allow some headroom, but warn if configuration seems inconsistent
import logging
logger = logging.getLogger(__name__)
logger.warning(
"Stream memory configuration may be inconsistent: "
f"{self.MAX_CONCURRENT_STREAMS} streams ×"
"{self.STREAM_MEMORY_LIMIT_PER_STREAM} bytes "
"could exceed connection limit of"
f"{self.STREAM_MEMORY_LIMIT_PER_CONNECTION} bytes"
)
def get_stream_config_dict(self) -> dict[str, Any]:
"""Get stream-specific configuration as dictionary."""
stream_config = {}
for attr_name in dir(self):
if attr_name.startswith(
("STREAM_", "MAX_", "ENABLE_STREAM", "CONNECTION_FLOW")
):
stream_config[attr_name.lower()] = getattr(self, attr_name)
return stream_config
# Additional configuration classes for specific stream features
class QUICStreamFlowControlConfig:
"""Configuration for QUIC stream flow control."""
def __init__(
self,
initial_window_size: int = 512 * 1024,
max_window_size: int = 2 * 1024 * 1024,
window_update_threshold: float = 0.5,
enable_auto_tuning: bool = True,
):
self.initial_window_size = initial_window_size
self.max_window_size = max_window_size
self.window_update_threshold = window_update_threshold
self.enable_auto_tuning = enable_auto_tuning
def create_stream_config_for_use_case(
use_case: Literal[
"high_throughput", "low_latency", "many_streams", "memory_constrained"
],
) -> QUICTransportConfig:
"""
Create optimized stream configuration for specific use cases.
Args:
use_case: One of "high_throughput", "low_latency", "many_streams","
"memory_constrained"
Returns:
Optimized QUICTransportConfig
"""
base_config = QUICTransportConfig()
if use_case == "high_throughput":
# Optimize for high throughput
base_config.STREAM_FLOW_CONTROL_WINDOW = 2 * 1024 * 1024 # 2MB
base_config.CONNECTION_FLOW_CONTROL_WINDOW = 10 * 1024 * 1024 # 10MB
base_config.MAX_STREAM_RECEIVE_BUFFER = 4 * 1024 * 1024 # 4MB
base_config.STREAM_PROCESSING_CONCURRENCY = 200
elif use_case == "low_latency":
# Optimize for low latency
base_config.STREAM_OPEN_TIMEOUT = 1.0
base_config.STREAM_READ_TIMEOUT = 5.0
base_config.STREAM_WRITE_TIMEOUT = 5.0
base_config.ENABLE_STREAM_BATCHING = False
base_config.STREAM_BATCH_SIZE = 1
elif use_case == "many_streams":
# Optimize for many concurrent streams
base_config.MAX_CONCURRENT_STREAMS = 5000
base_config.STREAM_FLOW_CONTROL_WINDOW = 128 * 1024 # 128KB
base_config.MAX_STREAM_RECEIVE_BUFFER = 256 * 1024 # 256KB
base_config.STREAM_PROCESSING_CONCURRENCY = 500
elif use_case == "memory_constrained":
# Optimize for low memory usage
base_config.MAX_CONCURRENT_STREAMS = 100
base_config.STREAM_FLOW_CONTROL_WINDOW = 64 * 1024 # 64KB
base_config.CONNECTION_FLOW_CONTROL_WINDOW = 256 * 1024 # 256KB
base_config.MAX_STREAM_RECEIVE_BUFFER = 128 * 1024 # 128KB
base_config.STREAM_MEMORY_LIMIT_PER_STREAM = 512 * 1024 # 512KB
base_config.STREAM_PROCESSING_CONCURRENCY = 50
else:
raise ValueError(f"Unknown use case: {use_case}")
return base_config

File diff suppressed because it is too large Load Diff

View File

@ -0,0 +1,391 @@
"""
QUIC Transport exceptions
"""
from typing import Any, Literal
class QUICError(Exception):
"""Base exception for all QUIC transport errors."""
def __init__(self, message: str, error_code: int | None = None):
super().__init__(message)
self.error_code = error_code
# Transport-level exceptions
class QUICTransportError(QUICError):
"""Base exception for QUIC transport operations."""
pass
class QUICDialError(QUICTransportError):
"""Error occurred during QUIC connection establishment."""
pass
class QUICListenError(QUICTransportError):
"""Error occurred during QUIC listener operations."""
pass
class QUICSecurityError(QUICTransportError):
"""Error related to QUIC security/TLS operations."""
pass
# Connection-level exceptions
class QUICConnectionError(QUICError):
"""Base exception for QUIC connection operations."""
pass
class QUICConnectionClosedError(QUICConnectionError):
"""QUIC connection has been closed."""
pass
class QUICConnectionTimeoutError(QUICConnectionError):
"""QUIC connection operation timed out."""
pass
class QUICHandshakeError(QUICConnectionError):
"""Error during QUIC handshake process."""
pass
class QUICPeerVerificationError(QUICConnectionError):
"""Error verifying peer identity during handshake."""
pass
# Stream-level exceptions
class QUICStreamError(QUICError):
"""Base exception for QUIC stream operations."""
def __init__(
self,
message: str,
stream_id: str | None = None,
error_code: int | None = None,
):
super().__init__(message, error_code)
self.stream_id = stream_id
class QUICStreamClosedError(QUICStreamError):
"""Stream is closed and cannot be used for I/O operations."""
pass
class QUICStreamResetError(QUICStreamError):
"""Stream was reset by local or remote peer."""
def __init__(
self,
message: str,
stream_id: str | None = None,
error_code: int | None = None,
reset_by_peer: bool = False,
):
super().__init__(message, stream_id, error_code)
self.reset_by_peer = reset_by_peer
class QUICStreamTimeoutError(QUICStreamError):
"""Stream operation timed out."""
pass
class QUICStreamBackpressureError(QUICStreamError):
"""Stream write blocked due to flow control."""
pass
class QUICStreamLimitError(QUICStreamError):
"""Stream limit reached (too many concurrent streams)."""
pass
class QUICStreamStateError(QUICStreamError):
"""Invalid operation for current stream state."""
def __init__(
self,
message: str,
stream_id: str | None = None,
current_state: str | None = None,
attempted_operation: str | None = None,
):
super().__init__(message, stream_id)
self.current_state = current_state
self.attempted_operation = attempted_operation
# Flow control exceptions
class QUICFlowControlError(QUICError):
"""Base exception for flow control related errors."""
pass
class QUICFlowControlViolationError(QUICFlowControlError):
"""Flow control limits were violated."""
pass
class QUICFlowControlDeadlockError(QUICFlowControlError):
"""Flow control deadlock detected."""
pass
# Resource management exceptions
class QUICResourceError(QUICError):
"""Base exception for resource management errors."""
pass
class QUICMemoryLimitError(QUICResourceError):
"""Memory limit exceeded."""
pass
class QUICConnectionLimitError(QUICResourceError):
"""Connection limit exceeded."""
pass
# Multiaddr and addressing exceptions
class QUICAddressError(QUICError):
"""Base exception for QUIC addressing errors."""
pass
class QUICInvalidMultiaddrError(QUICAddressError):
"""Invalid multiaddr format for QUIC transport."""
pass
class QUICAddressResolutionError(QUICAddressError):
"""Failed to resolve QUIC address."""
pass
class QUICProtocolError(QUICError):
"""Base exception for QUIC protocol errors."""
pass
class QUICVersionNegotiationError(QUICProtocolError):
"""QUIC version negotiation failed."""
pass
class QUICUnsupportedVersionError(QUICProtocolError):
"""Unsupported QUIC version."""
pass
# Configuration exceptions
class QUICConfigurationError(QUICError):
"""Base exception for QUIC configuration errors."""
pass
class QUICInvalidConfigError(QUICConfigurationError):
"""Invalid QUIC configuration parameters."""
pass
class QUICCertificateError(QUICConfigurationError):
"""Error with TLS certificate configuration."""
pass
def map_quic_error_code(error_code: int) -> str:
"""
Map QUIC error codes to human-readable descriptions.
Based on RFC 9000 Transport Error Codes.
"""
error_codes = {
0x00: "NO_ERROR",
0x01: "INTERNAL_ERROR",
0x02: "CONNECTION_REFUSED",
0x03: "FLOW_CONTROL_ERROR",
0x04: "STREAM_LIMIT_ERROR",
0x05: "STREAM_STATE_ERROR",
0x06: "FINAL_SIZE_ERROR",
0x07: "FRAME_ENCODING_ERROR",
0x08: "TRANSPORT_PARAMETER_ERROR",
0x09: "CONNECTION_ID_LIMIT_ERROR",
0x0A: "PROTOCOL_VIOLATION",
0x0B: "INVALID_TOKEN",
0x0C: "APPLICATION_ERROR",
0x0D: "CRYPTO_BUFFER_EXCEEDED",
0x0E: "KEY_UPDATE_ERROR",
0x0F: "AEAD_LIMIT_REACHED",
0x10: "NO_VIABLE_PATH",
}
return error_codes.get(error_code, f"UNKNOWN_ERROR_{error_code:02X}")
def create_stream_error(
error_type: str,
message: str,
stream_id: str | None = None,
error_code: int | None = None,
) -> QUICStreamError:
"""
Factory function to create appropriate stream error based on type.
Args:
error_type: Type of error ("closed", "reset", "timeout", "backpressure", etc.)
message: Error message
stream_id: Stream identifier
error_code: QUIC error code
Returns:
Appropriate QUICStreamError subclass
"""
error_type = error_type.lower()
if error_type in ("closed", "close"):
return QUICStreamClosedError(message, stream_id, error_code)
elif error_type == "reset":
return QUICStreamResetError(message, stream_id, error_code)
elif error_type == "timeout":
return QUICStreamTimeoutError(message, stream_id, error_code)
elif error_type in ("backpressure", "flow_control"):
return QUICStreamBackpressureError(message, stream_id, error_code)
elif error_type in ("limit", "stream_limit"):
return QUICStreamLimitError(message, stream_id, error_code)
elif error_type == "state":
return QUICStreamStateError(message, stream_id)
else:
return QUICStreamError(message, stream_id, error_code)
def create_connection_error(
error_type: str, message: str, error_code: int | None = None
) -> QUICConnectionError:
"""
Factory function to create appropriate connection error based on type.
Args:
error_type: Type of error ("closed", "timeout", "handshake", etc.)
message: Error message
error_code: QUIC error code
Returns:
Appropriate QUICConnectionError subclass
"""
error_type = error_type.lower()
if error_type in ("closed", "close"):
return QUICConnectionClosedError(message, error_code)
elif error_type == "timeout":
return QUICConnectionTimeoutError(message, error_code)
elif error_type == "handshake":
return QUICHandshakeError(message, error_code)
elif error_type in ("peer_verification", "verification"):
return QUICPeerVerificationError(message, error_code)
else:
return QUICConnectionError(message, error_code)
class QUICErrorContext:
"""
Context manager for handling QUIC errors with automatic error mapping.
Useful for converting low-level aioquic errors to py-libp2p QUIC errors.
"""
def __init__(self, operation: str, component: str = "quic") -> None:
self.operation = operation
self.component = component
def __enter__(self) -> "QUICErrorContext":
return self
# TODO: Fix types for exc_type
def __exit__(
self,
exc_type: type[BaseException] | None | None,
exc_val: BaseException | None,
exc_tb: Any,
) -> Literal[False]:
if exc_type is None:
return False
if exc_val is None:
return False
# Map common aioquic exceptions to our exceptions
if "ConnectionClosed" in str(exc_type):
raise QUICConnectionClosedError(
f"Connection closed during {self.operation}: {exc_val}"
) from exc_val
elif "StreamReset" in str(exc_type):
raise QUICStreamResetError(
f"Stream reset during {self.operation}: {exc_val}"
) from exc_val
elif "timeout" in str(exc_val).lower():
if "stream" in self.component.lower():
raise QUICStreamTimeoutError(
f"Timeout during {self.operation}: {exc_val}"
) from exc_val
else:
raise QUICConnectionTimeoutError(
f"Timeout during {self.operation}: {exc_val}"
) from exc_val
elif "flow control" in str(exc_val).lower():
raise QUICStreamBackpressureError(
f"Flow control error during {self.operation}: {exc_val}"
) from exc_val
# Let other exceptions propagate
return False

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

@ -0,0 +1,656 @@
"""
QUIC Stream implementation
Provides stream interface over QUIC's native multiplexing.
"""
from enum import Enum
import logging
import time
from types import TracebackType
from typing import TYPE_CHECKING, Any, cast
import trio
from .exceptions import (
QUICStreamBackpressureError,
QUICStreamClosedError,
QUICStreamResetError,
QUICStreamTimeoutError,
)
if TYPE_CHECKING:
from libp2p.abc import IMuxedStream
from libp2p.custom_types import TProtocol
from .connection import QUICConnection
else:
IMuxedStream = cast(type, object)
TProtocol = cast(type, object)
logger = logging.getLogger(__name__)
class StreamState(Enum):
"""Stream lifecycle states following libp2p patterns."""
OPEN = "open"
WRITE_CLOSED = "write_closed"
READ_CLOSED = "read_closed"
CLOSED = "closed"
RESET = "reset"
class StreamDirection(Enum):
"""Stream direction for tracking initiator."""
INBOUND = "inbound"
OUTBOUND = "outbound"
class StreamTimeline:
"""Track stream lifecycle events for debugging and monitoring."""
def __init__(self) -> None:
self.created_at = time.time()
self.opened_at: float | None = None
self.first_data_at: float | None = None
self.closed_at: float | None = None
self.reset_at: float | None = None
self.error_code: int | None = None
def record_open(self) -> None:
self.opened_at = time.time()
def record_first_data(self) -> None:
if self.first_data_at is None:
self.first_data_at = time.time()
def record_close(self) -> None:
self.closed_at = time.time()
def record_reset(self, error_code: int) -> None:
self.reset_at = time.time()
self.error_code = error_code
class QUICStream(IMuxedStream):
"""
QUIC Stream implementation following libp2p IMuxedStream interface.
Based on patterns from go-libp2p and js-libp2p, this implementation:
- Leverages QUIC's native multiplexing and flow control
- Integrates with libp2p resource management
- Provides comprehensive error handling with QUIC-specific codes
- Supports bidirectional communication with independent close semantics
- Implements proper stream lifecycle management
"""
def __init__(
self,
connection: "QUICConnection",
stream_id: int,
direction: StreamDirection,
remote_addr: tuple[str, int],
resource_scope: Any | None = None,
):
"""
Initialize QUIC stream.
Args:
connection: Parent QUIC connection
stream_id: QUIC stream identifier
direction: Stream direction (inbound/outbound)
resource_scope: Resource manager scope for memory accounting
remote_addr: Remote addr stream is connected to
"""
self._connection = connection
self._stream_id = stream_id
self._direction = direction
self._resource_scope = resource_scope
# libp2p interface compliance
self._protocol: TProtocol | None = None
self._metadata: dict[str, Any] = {}
self._remote_addr = remote_addr
# Stream state management
self._state = StreamState.OPEN
self._state_lock = trio.Lock()
# Flow control and buffering
self._receive_buffer = bytearray()
self._receive_buffer_lock = trio.Lock()
self._receive_event = trio.Event()
self._backpressure_event = trio.Event()
self._backpressure_event.set() # Initially no backpressure
# Close/reset state
self._write_closed = False
self._read_closed = False
self._close_event = trio.Event()
self._reset_error_code: int | None = None
# Lifecycle tracking
self._timeline = StreamTimeline()
self._timeline.record_open()
# Resource accounting
self._memory_reserved = 0
# Stream constant configurations
self.READ_TIMEOUT = connection._transport._config.STREAM_READ_TIMEOUT
self.WRITE_TIMEOUT = connection._transport._config.STREAM_WRITE_TIMEOUT
self.FLOW_CONTROL_WINDOW_SIZE = (
connection._transport._config.STREAM_FLOW_CONTROL_WINDOW
)
self.MAX_RECEIVE_BUFFER_SIZE = (
connection._transport._config.MAX_STREAM_RECEIVE_BUFFER
)
if self._resource_scope:
self._reserve_memory(self.FLOW_CONTROL_WINDOW_SIZE)
logger.debug(
f"Created QUIC stream {stream_id} "
f"({direction.value}, connection: {connection.remote_peer_id()})"
)
# Properties for libp2p interface compliance
@property
def protocol(self) -> TProtocol | None:
"""Get the protocol identifier for this stream."""
return self._protocol
@protocol.setter
def protocol(self, protocol_id: TProtocol) -> None:
"""Set the protocol identifier for this stream."""
self._protocol = protocol_id
self._metadata["protocol"] = protocol_id
logger.debug(f"Stream {self.stream_id} protocol set to: {protocol_id}")
@property
def stream_id(self) -> str:
"""Get stream ID as string for libp2p compatibility."""
return str(self._stream_id)
@property
def muxed_conn(self) -> "QUICConnection": # type: ignore
"""Get the parent muxed connection."""
return self._connection
@property
def state(self) -> StreamState:
"""Get current stream state."""
return self._state
@property
def direction(self) -> StreamDirection:
"""Get stream direction."""
return self._direction
@property
def is_initiator(self) -> bool:
"""Check if this stream was locally initiated."""
return self._direction == StreamDirection.OUTBOUND
# Core stream operations
async def read(self, n: int | None = None) -> bytes:
"""
Read data from the stream with QUIC flow control.
Args:
n: Maximum number of bytes to read. If None or -1, read all available.
Returns:
Data read from stream
Raises:
QUICStreamClosedError: Stream is closed
QUICStreamResetError: Stream was reset
QUICStreamTimeoutError: Read timeout exceeded
"""
if n is None:
n = -1
async with self._state_lock:
if self._state in (StreamState.CLOSED, StreamState.RESET):
raise QUICStreamClosedError(f"Stream {self.stream_id} is closed")
if self._read_closed:
# Return any remaining buffered data, then EOF
async with self._receive_buffer_lock:
if self._receive_buffer:
data = self._extract_data_from_buffer(n)
self._timeline.record_first_data()
return data
return b""
# Wait for data with timeout
timeout = self.READ_TIMEOUT
try:
with trio.move_on_after(timeout) as cancel_scope:
while True:
async with self._receive_buffer_lock:
if self._receive_buffer:
data = self._extract_data_from_buffer(n)
self._timeline.record_first_data()
return data
# Check if stream was closed while waiting
if self._read_closed:
return b""
# Wait for more data
await self._receive_event.wait()
self._receive_event = trio.Event() # Reset for next wait
if cancel_scope.cancelled_caught:
raise QUICStreamTimeoutError(f"Read timeout on stream {self.stream_id}")
return b""
except QUICStreamResetError:
# Stream was reset while reading
raise
except Exception as e:
logger.error(f"Error reading from stream {self.stream_id}: {e}")
await self._handle_stream_error(e)
raise
async def write(self, data: bytes) -> None:
"""
Write data to the stream with QUIC flow control.
Args:
data: Data to write
Raises:
QUICStreamClosedError: Stream is closed for writing
QUICStreamBackpressureError: Flow control window exhausted
QUICStreamResetError: Stream was reset
"""
if not data:
return
async with self._state_lock:
if self._state in (StreamState.CLOSED, StreamState.RESET):
raise QUICStreamClosedError(f"Stream {self.stream_id} is closed")
if self._write_closed:
raise QUICStreamClosedError(
f"Stream {self.stream_id} write side is closed"
)
try:
# Handle flow control backpressure
await self._backpressure_event.wait()
# Send data through QUIC connection
self._connection._quic.send_stream_data(self._stream_id, data)
await self._connection._transmit()
self._timeline.record_first_data()
logger.debug(f"Wrote {len(data)} bytes to stream {self.stream_id}")
except Exception as e:
logger.error(f"Error writing to stream {self.stream_id}: {e}")
# Convert QUIC-specific errors
if "flow control" in str(e).lower():
raise QUICStreamBackpressureError(f"Flow control limit reached: {e}")
await self._handle_stream_error(e)
raise
async def close(self) -> None:
"""
Close the stream gracefully (both read and write sides).
This implements proper close semantics where both sides
are closed and resources are cleaned up.
"""
async with self._state_lock:
if self._state in (StreamState.CLOSED, StreamState.RESET):
return
logger.debug(f"Closing stream {self.stream_id}")
# Close both sides
if not self._write_closed:
await self.close_write()
if not self._read_closed:
await self.close_read()
# Update state and cleanup
async with self._state_lock:
self._state = StreamState.CLOSED
await self._cleanup_resources()
self._timeline.record_close()
self._close_event.set()
logger.debug(f"Stream {self.stream_id} closed")
async def close_write(self) -> None:
"""Close the write side of the stream."""
if self._write_closed:
return
try:
# Send FIN to close write side
self._connection._quic.send_stream_data(
self._stream_id, b"", end_stream=True
)
await self._connection._transmit()
self._write_closed = True
async with self._state_lock:
if self._read_closed:
self._state = StreamState.CLOSED
else:
self._state = StreamState.WRITE_CLOSED
logger.debug(f"Stream {self.stream_id} write side closed")
except Exception as e:
logger.error(f"Error closing write side of stream {self.stream_id}: {e}")
async def close_read(self) -> None:
"""Close the read side of the stream."""
if self._read_closed:
return
try:
self._read_closed = True
async with self._state_lock:
if self._write_closed:
self._state = StreamState.CLOSED
else:
self._state = StreamState.READ_CLOSED
# Wake up any pending reads
self._receive_event.set()
logger.debug(f"Stream {self.stream_id} read side closed")
except Exception as e:
logger.error(f"Error closing read side of stream {self.stream_id}: {e}")
async def reset(self, error_code: int = 0) -> None:
"""
Reset the stream with the given error code.
Args:
error_code: QUIC error code for the reset
"""
async with self._state_lock:
if self._state == StreamState.RESET:
return
logger.debug(
f"Resetting stream {self.stream_id} with error code {error_code}"
)
self._state = StreamState.RESET
self._reset_error_code = error_code
try:
# Send QUIC reset frame
self._connection._quic.reset_stream(self._stream_id, error_code)
await self._connection._transmit()
except Exception as e:
logger.error(f"Error sending reset for stream {self.stream_id}: {e}")
finally:
# Always cleanup resources
await self._cleanup_resources()
self._timeline.record_reset(error_code)
self._close_event.set()
def is_closed(self) -> bool:
"""Check if stream is completely closed."""
return self._state in (StreamState.CLOSED, StreamState.RESET)
def is_reset(self) -> bool:
"""Check if stream was reset."""
return self._state == StreamState.RESET
def can_read(self) -> bool:
"""Check if stream can be read from."""
return not self._read_closed and self._state not in (
StreamState.CLOSED,
StreamState.RESET,
)
def can_write(self) -> bool:
"""Check if stream can be written to."""
return not self._write_closed and self._state not in (
StreamState.CLOSED,
StreamState.RESET,
)
async def handle_data_received(self, data: bytes, end_stream: bool) -> None:
"""
Handle data received from the QUIC connection.
Args:
data: Received data
end_stream: Whether this is the last data (FIN received)
"""
if self._state == StreamState.RESET:
return
if data:
async with self._receive_buffer_lock:
if len(self._receive_buffer) + len(data) > self.MAX_RECEIVE_BUFFER_SIZE:
logger.warning(
f"Stream {self.stream_id} receive buffer overflow, "
f"dropping {len(data)} bytes"
)
return
self._receive_buffer.extend(data)
self._timeline.record_first_data()
# Notify waiting readers
self._receive_event.set()
logger.debug(f"Stream {self.stream_id} received {len(data)} bytes")
if end_stream:
self._read_closed = True
async with self._state_lock:
if self._write_closed:
self._state = StreamState.CLOSED
else:
self._state = StreamState.READ_CLOSED
# Wake up readers to process remaining data and EOF
self._receive_event.set()
logger.debug(f"Stream {self.stream_id} received FIN")
async def handle_stop_sending(self, error_code: int) -> None:
"""
Handle STOP_SENDING frame from remote peer.
When a STOP_SENDING frame is received, the peer is requesting that we
stop sending data on this stream. We respond by resetting the stream.
Args:
error_code: Error code from the STOP_SENDING frame
"""
logger.debug(
f"Stream {self.stream_id} handling STOP_SENDING (error_code={error_code})"
)
self._write_closed = True
# Wake up any pending write operations
self._backpressure_event.set()
async with self._state_lock:
if self.direction == StreamDirection.OUTBOUND:
self._state = StreamState.CLOSED
elif self._read_closed:
self._state = StreamState.CLOSED
else:
# Only write side closed - add WRITE_CLOSED state if needed
self._state = StreamState.WRITE_CLOSED
# Send RESET_STREAM in response (QUIC protocol requirement)
try:
self._connection._quic.reset_stream(int(self.stream_id), error_code)
await self._connection._transmit()
logger.debug(f"Sent RESET_STREAM for stream {self.stream_id}")
except Exception as e:
logger.warning(
f"Could not send RESET_STREAM for stream {self.stream_id}: {e}"
)
async def handle_reset(self, error_code: int) -> None:
"""
Handle stream reset from remote peer.
Args:
error_code: QUIC error code from reset frame
"""
logger.debug(
f"Stream {self.stream_id} reset by peer with error code {error_code}"
)
async with self._state_lock:
self._state = StreamState.RESET
self._reset_error_code = error_code
await self._cleanup_resources()
self._timeline.record_reset(error_code)
self._close_event.set()
# Wake up any pending operations
self._receive_event.set()
self._backpressure_event.set()
async def handle_flow_control_update(self, available_window: int) -> None:
"""
Handle flow control window updates.
Args:
available_window: Available flow control window size
"""
if available_window > 0:
self._backpressure_event.set()
logger.debug(
f"Stream {self.stream_id} flow control".__add__(
f"window updated: {available_window}"
)
)
else:
self._backpressure_event = trio.Event() # Reset to blocking state
logger.debug(f"Stream {self.stream_id} flow control window exhausted")
def _extract_data_from_buffer(self, n: int) -> bytes:
"""Extract data from receive buffer with specified limit."""
if n == -1:
# Read all available data
data = bytes(self._receive_buffer)
self._receive_buffer.clear()
else:
# Read up to n bytes
data = bytes(self._receive_buffer[:n])
self._receive_buffer = self._receive_buffer[n:]
return data
async def _handle_stream_error(self, error: Exception) -> None:
"""Handle errors by resetting the stream."""
logger.error(f"Stream {self.stream_id} error: {error}")
await self.reset(error_code=1) # Generic error code
def _reserve_memory(self, size: int) -> None:
"""Reserve memory with resource manager."""
if self._resource_scope:
try:
self._resource_scope.reserve_memory(size)
self._memory_reserved += size
except Exception as e:
logger.warning(
f"Failed to reserve memory for stream {self.stream_id}: {e}"
)
def _release_memory(self, size: int) -> None:
"""Release memory with resource manager."""
if self._resource_scope and size > 0:
try:
self._resource_scope.release_memory(size)
self._memory_reserved = max(0, self._memory_reserved - size)
except Exception as e:
logger.warning(
f"Failed to release memory for stream {self.stream_id}: {e}"
)
async def _cleanup_resources(self) -> None:
"""Clean up stream resources."""
# Release all reserved memory
if self._memory_reserved > 0:
self._release_memory(self._memory_reserved)
# Clear receive buffer
async with self._receive_buffer_lock:
self._receive_buffer.clear()
# Remove from connection's stream registry
self._connection._remove_stream(self._stream_id)
logger.debug(f"Stream {self.stream_id} resources cleaned up")
# Abstact implementations
def get_remote_address(self) -> tuple[str, int]:
return self._remote_addr
async def __aenter__(self) -> "QUICStream":
"""Enter the async context manager."""
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> None:
"""Exit the async context manager and close the stream."""
logger.debug("Exiting the context and closing the stream")
await self.close()
def set_deadline(self, ttl: int) -> bool:
"""
Set a deadline for the stream. QUIC does not support deadlines natively,
so this method always returns False to indicate the operation is unsupported.
:param ttl: Time-to-live in seconds (ignored).
:return: False, as deadlines are not supported.
"""
raise NotImplementedError("QUIC does not support setting read deadlines")
# String representation for debugging
def __repr__(self) -> str:
return (
f"QUICStream(id={self.stream_id}, "
f"state={self._state.value}, "
f"direction={self._direction.value}, "
f"protocol={self._protocol})"
)
def __str__(self) -> str:
return f"QUICStream({self.stream_id})"

View File

@ -0,0 +1,491 @@
"""
QUIC Transport implementation
"""
import copy
import logging
import ssl
from typing import TYPE_CHECKING, cast
from aioquic.quic.configuration import (
QuicConfiguration,
)
from aioquic.quic.connection import (
QuicConnection as NativeQUICConnection,
)
from aioquic.quic.logger import QuicLogger
import multiaddr
import trio
from libp2p.abc import (
ITransport,
)
from libp2p.crypto.keys import (
PrivateKey,
)
from libp2p.custom_types import TProtocol, TQUICConnHandlerFn
from libp2p.peer.id import (
ID,
)
from libp2p.transport.quic.security import QUICTLSSecurityConfig
from libp2p.transport.quic.utils import (
create_client_config_from_base,
create_server_config_from_base,
get_alpn_protocols,
is_quic_multiaddr,
multiaddr_to_quic_version,
quic_multiaddr_to_endpoint,
quic_version_to_wire_format,
)
if TYPE_CHECKING:
from libp2p.network.swarm import Swarm
else:
Swarm = cast(type, object)
from .config import (
QUICTransportConfig,
)
from .connection import (
QUICConnection,
)
from .exceptions import (
QUICDialError,
QUICListenError,
QUICSecurityError,
)
from .listener import (
QUICListener,
)
from .security import (
QUICTLSConfigManager,
create_quic_security_transport,
)
QUIC_V1_PROTOCOL = QUICTransportConfig.PROTOCOL_QUIC_V1
QUIC_DRAFT29_PROTOCOL = QUICTransportConfig.PROTOCOL_QUIC_DRAFT29
logger = logging.getLogger(__name__)
class QUICTransport(ITransport):
"""
QUIC Stream implementation following libp2p IMuxedStream interface.
"""
def __init__(
self, private_key: PrivateKey, config: QUICTransportConfig | None = None
) -> None:
"""
Initialize QUIC transport with security integration.
Args:
private_key: libp2p private key for identity and TLS cert generation
config: QUIC transport configuration options
"""
self._private_key = private_key
self._peer_id = ID.from_pubkey(private_key.get_public_key())
self._config = config or QUICTransportConfig()
# Connection management
self._connections: dict[str, QUICConnection] = {}
self._listeners: list[QUICListener] = []
# Security manager for TLS integration
self._security_manager = create_quic_security_transport(
self._private_key, self._peer_id
)
# QUIC configurations for different versions
self._quic_configs: dict[TProtocol, QuicConfiguration] = {}
self._setup_quic_configurations()
# Resource management
self._closed = False
self._nursery_manager = trio.CapacityLimiter(1)
self._background_nursery: trio.Nursery | None = None
self._swarm: Swarm | None = None
logger.debug(
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
logger.debug("Transport background nursery set")
def set_swarm(self, swarm: Swarm) -> None:
"""Set the swarm for adding incoming connections."""
self._swarm = swarm
def _setup_quic_configurations(self) -> None:
"""Setup QUIC configurations."""
try:
# Get TLS configuration from security manager
server_tls_config = self._security_manager.create_server_config()
client_tls_config = self._security_manager.create_client_config()
# Base server configuration
base_server_config = QuicConfiguration(
is_client=False,
alpn_protocols=get_alpn_protocols(),
verify_mode=self._config.verify_mode,
max_datagram_frame_size=self._config.max_datagram_size,
idle_timeout=self._config.idle_timeout,
)
# Base client configuration
base_client_config = QuicConfiguration(
is_client=True,
alpn_protocols=get_alpn_protocols(),
verify_mode=self._config.verify_mode,
max_datagram_frame_size=self._config.max_datagram_size,
idle_timeout=self._config.idle_timeout,
)
# Apply TLS configuration
self._apply_tls_configuration(base_server_config, server_tls_config)
self._apply_tls_configuration(base_client_config, client_tls_config)
# QUIC v1 (RFC 9000) configurations
if self._config.enable_v1:
quic_v1_server_config = create_server_config_from_base(
base_server_config, self._security_manager, self._config
)
quic_v1_server_config.supported_versions = [
quic_version_to_wire_format(QUIC_V1_PROTOCOL)
]
quic_v1_client_config = create_client_config_from_base(
base_client_config, self._security_manager, self._config
)
quic_v1_client_config.supported_versions = [
quic_version_to_wire_format(QUIC_V1_PROTOCOL)
]
# Store both server and client configs for v1
self._quic_configs[TProtocol(f"{QUIC_V1_PROTOCOL}_server")] = (
quic_v1_server_config
)
self._quic_configs[TProtocol(f"{QUIC_V1_PROTOCOL}_client")] = (
quic_v1_client_config
)
# QUIC draft-29 configurations for compatibility
if self._config.enable_draft29:
draft29_server_config: QuicConfiguration = copy.copy(base_server_config)
draft29_server_config.supported_versions = [
quic_version_to_wire_format(QUIC_DRAFT29_PROTOCOL)
]
draft29_client_config = copy.copy(base_client_config)
draft29_client_config.supported_versions = [
quic_version_to_wire_format(QUIC_DRAFT29_PROTOCOL)
]
self._quic_configs[TProtocol(f"{QUIC_DRAFT29_PROTOCOL}_server")] = (
draft29_server_config
)
self._quic_configs[TProtocol(f"{QUIC_DRAFT29_PROTOCOL}_client")] = (
draft29_client_config
)
logger.debug("QUIC configurations initialized with libp2p TLS security")
except Exception as e:
raise QUICSecurityError(
f"Failed to setup QUIC TLS configurations: {e}"
) from e
def _apply_tls_configuration(
self, config: QuicConfiguration, tls_config: QUICTLSSecurityConfig
) -> None:
"""
Apply TLS configuration to a QUIC configuration using aioquic's actual API.
Args:
config: QuicConfiguration to update
tls_config: TLS configuration dictionary from security manager
"""
try:
config.certificate = tls_config.certificate
config.private_key = tls_config.private_key
config.certificate_chain = tls_config.certificate_chain
config.alpn_protocols = tls_config.alpn_protocols
config.verify_mode = ssl.CERT_NONE
logger.debug("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,
) -> QUICConnection:
"""
Dial a remote peer using QUIC transport with security verification.
Args:
maddr: Multiaddr of the remote peer (e.g., /ip4/1.2.3.4/udp/4001/quic-v1)
peer_id: Expected peer ID for verification
nursery: Nursery to execute the background tasks
Returns:
Raw connection interface to the remote peer
Raises:
QUICDialError: If dialing fails
QUICSecurityError: If security verification fails
"""
if self._closed:
raise QUICDialError("Transport is closed")
if not is_quic_multiaddr(maddr):
raise QUICDialError(f"Invalid QUIC multiaddr: {maddr}")
try:
# Extract connection details from multiaddr
host, port = quic_multiaddr_to_endpoint(maddr)
remote_peer_id = maddr.get_peer_id()
if remote_peer_id is not None:
remote_peer_id = ID.from_base58(remote_peer_id)
if remote_peer_id is None:
logger.error("Unable to derive peer id from multiaddr")
raise QUICDialError("Unable to derive peer id from multiaddr")
quic_version = multiaddr_to_quic_version(maddr)
# Get appropriate QUIC client configuration
config_key = TProtocol(f"{quic_version}_client")
logger.debug("config_key", config_key, self._quic_configs.keys())
config = self._quic_configs.get(config_key)
if not config:
raise QUICDialError(f"Unsupported QUIC version: {quic_version}")
config.is_client = True
config.quic_logger = QuicLogger()
# Ensure client certificate is properly set for mutual authentication
if not config.certificate or not config.private_key:
logger.warning(
"Client config missing certificate - applying TLS config"
)
client_tls_config = self._security_manager.create_client_config()
self._apply_tls_configuration(config, client_tls_config)
# Debug log to verify certificate is present
logger.info(
f"Dialing QUIC connection to {host}:{port} (version: {{quic_version}})"
)
logger.debug("Starting QUIC Connection")
# Create QUIC connection using aioquic's sans-IO core
native_quic_connection = NativeQUICConnection(configuration=config)
# Create trio-based QUIC connection wrapper with security
connection = QUICConnection(
quic_connection=native_quic_connection,
remote_addr=(host, port),
remote_peer_id=remote_peer_id,
local_peer_id=self._peer_id,
is_initiator=True,
maddr=maddr,
transport=self,
security_manager=self._security_manager,
)
logger.debug("QUIC Connection Created")
if self._background_nursery is None:
logger.error("No nursery set to execute background tasks")
raise QUICDialError("No nursery found to execute tasks")
await connection.connect(self._background_nursery)
# Store connection for management
conn_id = f"{host}:{port}"
self._connections[conn_id] = connection
return connection
except Exception as e:
logger.error(f"Failed to dial QUIC connection to {maddr}: {e}")
raise QUICDialError(f"Dial failed: {e}") from e
async def _verify_peer_identity(
self, connection: QUICConnection, expected_peer_id: ID
) -> None:
"""
Verify remote peer identity after TLS handshake.
Args:
connection: The established QUIC connection
expected_peer_id: Expected peer ID
Raises:
QUICSecurityError: If peer verification fails
"""
try:
# Get peer certificate from the connection
peer_certificate = await connection.get_peer_certificate()
if not peer_certificate:
raise QUICSecurityError("No peer certificate available")
# Verify peer identity using security manager
verified_peer_id = self._security_manager.verify_peer_identity(
peer_certificate, expected_peer_id
)
if verified_peer_id != expected_peer_id:
raise QUICSecurityError(
"Peer ID verification failed: expected "
f"{expected_peer_id}, got {verified_peer_id}"
)
logger.debug(f"Peer identity verified: {verified_peer_id}")
logger.debug(f"Peer identity verified: {verified_peer_id}")
except Exception as e:
raise QUICSecurityError(f"Peer identity verification failed: {e}") from e
def create_listener(self, handler_function: TQUICConnHandlerFn) -> QUICListener:
"""
Create a QUIC listener with integrated security.
Args:
handler_function: Function to handle new connections
Returns:
QUIC listener instance
Raises:
QUICListenError: If transport is closed
"""
if self._closed:
raise QUICListenError("Transport is closed")
# Get server configurations for the listener
server_configs = {
version: config
for version, config in self._quic_configs.items()
if version.endswith("_server")
}
listener = QUICListener(
transport=self,
handler_function=handler_function,
quic_configs=server_configs,
config=self._config,
security_manager=self._security_manager,
)
self._listeners.append(listener)
logger.debug("Created QUIC listener with security")
return listener
def can_dial(self, maddr: multiaddr.Multiaddr) -> bool:
"""
Check if this transport can dial the given multiaddr.
Args:
maddr: Multiaddr to check
Returns:
True if this transport can dial the address
"""
return is_quic_multiaddr(maddr)
def protocols(self) -> list[TProtocol]:
"""
Get supported protocol identifiers.
Returns:
List of supported protocol strings
"""
protocols = [QUIC_V1_PROTOCOL]
if self._config.enable_draft29:
protocols.append(QUIC_DRAFT29_PROTOCOL)
return protocols
def listen_order(self) -> int:
"""
Get the listen order priority for this transport.
Matches go-libp2p's ListenOrder = 1 for QUIC.
Returns:
Priority order for listening (lower = higher priority)
"""
return 1
async def close(self) -> None:
"""Close the transport and cleanup resources."""
if self._closed:
return
self._closed = True
logger.debug("Closing QUIC transport")
# Close all active connections and listeners concurrently using trio nursery
async with trio.open_nursery() as nursery:
# Close all connections
for connection in self._connections.values():
nursery.start_soon(connection.close)
# Close all listeners
for listener in self._listeners:
nursery.start_soon(listener.close)
self._connections.clear()
self._listeners.clear()
logger.debug("QUIC transport closed")
async def _cleanup_terminated_connection(self, connection: QUICConnection) -> None:
"""Clean up a terminated connection from all listeners."""
try:
for listener in self._listeners:
await listener._remove_connection_by_object(connection)
logger.debug(
"✅ TRANSPORT: Cleaned up terminated connection from all listeners"
)
except Exception as e:
logger.error(f"❌ TRANSPORT: Error cleaning up terminated connection: {e}")
def get_stats(self) -> dict[str, int | list[str] | object]:
"""Get transport statistics including security info."""
return {
"active_connections": len(self._connections),
"active_listeners": len(self._listeners),
"supported_protocols": self.protocols(),
"local_peer_id": str(self._peer_id),
"security_enabled": True,
"tls_configured": True,
}
def get_security_manager(self) -> QUICTLSConfigManager:
"""
Get the security manager for this transport.
Returns:
The QUIC TLS configuration manager
"""
return self._security_manager
def get_listener_socket(self) -> trio.socket.SocketType | None:
"""Get the socket from the first active listener."""
for listener in self._listeners:
if listener.is_listening() and listener._socket:
return listener._socket
return None

View File

@ -0,0 +1,466 @@
"""
Multiaddr utilities for QUIC transport - Module 4.
Essential utilities required for QUIC transport implementation.
Based on go-libp2p and js-libp2p QUIC implementations.
"""
import ipaddress
import logging
import ssl
from aioquic.quic.configuration import QuicConfiguration
import multiaddr
from libp2p.custom_types import TProtocol
from libp2p.transport.quic.security import QUICTLSConfigManager
from .config import QUICTransportConfig
from .exceptions import QUICInvalidMultiaddrError, QUICUnsupportedVersionError
logger = logging.getLogger(__name__)
# Protocol constants
QUIC_V1_PROTOCOL = QUICTransportConfig.PROTOCOL_QUIC_V1
QUIC_DRAFT29_PROTOCOL = QUICTransportConfig.PROTOCOL_QUIC_DRAFT29
UDP_PROTOCOL = "udp"
IP4_PROTOCOL = "ip4"
IP6_PROTOCOL = "ip6"
SERVER_CONFIG_PROTOCOL_V1 = f"{QUIC_V1_PROTOCOL}_server"
CLIENT_CONFIG_PROTCOL_V1 = f"{QUIC_V1_PROTOCOL}_client"
SERVER_CONFIG_PROTOCOL_DRAFT_29 = f"{QUIC_DRAFT29_PROTOCOL}_server"
CLIENT_CONFIG_PROTOCOL_DRAFT_29 = f"{QUIC_DRAFT29_PROTOCOL}_client"
CUSTOM_QUIC_VERSION_MAPPING: dict[str, int] = {
SERVER_CONFIG_PROTOCOL_V1: 0x00000001, # RFC 9000
CLIENT_CONFIG_PROTCOL_V1: 0x00000001, # RFC 9000
SERVER_CONFIG_PROTOCOL_DRAFT_29: 0xFF00001D, # draft-29
CLIENT_CONFIG_PROTOCOL_DRAFT_29: 0xFF00001D, # draft-29
}
# QUIC version to wire format mappings (required for aioquic)
QUIC_VERSION_MAPPINGS: dict[TProtocol, int] = {
QUIC_V1_PROTOCOL: 0x00000001, # RFC 9000
QUIC_DRAFT29_PROTOCOL: 0xFF00001D, # draft-29
}
# ALPN protocols for libp2p over QUIC
LIBP2P_ALPN_PROTOCOLS: list[str] = ["libp2p"]
def is_quic_multiaddr(maddr: multiaddr.Multiaddr) -> bool:
"""
Check if a multiaddr represents a QUIC address.
Valid QUIC multiaddrs:
- /ip4/127.0.0.1/udp/4001/quic-v1
- /ip4/127.0.0.1/udp/4001/quic
- /ip6/::1/udp/4001/quic-v1
- /ip6/::1/udp/4001/quic
Args:
maddr: Multiaddr to check
Returns:
True if the multiaddr represents a QUIC address
"""
try:
addr_str = str(maddr)
# Check for required components
has_ip = f"/{IP4_PROTOCOL}/" in addr_str or f"/{IP6_PROTOCOL}/" in addr_str
has_udp = f"/{UDP_PROTOCOL}/" in addr_str
has_quic = (
f"/{QUIC_V1_PROTOCOL}" in addr_str
or f"/{QUIC_DRAFT29_PROTOCOL}" in addr_str
or "/quic" in addr_str
)
return has_ip and has_udp and has_quic
except Exception:
return False
def quic_multiaddr_to_endpoint(maddr: multiaddr.Multiaddr) -> tuple[str, int]:
"""
Extract host and port from a QUIC multiaddr.
Args:
maddr: QUIC multiaddr
Returns:
Tuple of (host, port)
Raises:
QUICInvalidMultiaddrError: If multiaddr is not a valid QUIC address
"""
if not is_quic_multiaddr(maddr):
raise QUICInvalidMultiaddrError(f"Not a valid QUIC multiaddr: {maddr}")
try:
host = None
port = None
# Try to get IPv4 address
try:
host = maddr.value_for_protocol(multiaddr.protocols.P_IP4) # type: ignore
except Exception:
pass
# Try to get IPv6 address if IPv4 not found
if host is None:
try:
host = maddr.value_for_protocol(multiaddr.protocols.P_IP6) # type: ignore
except Exception:
pass
# Get UDP port
try:
port_str = maddr.value_for_protocol(multiaddr.protocols.P_UDP) # type: ignore
port = int(port_str)
except Exception:
pass
if host is None or port is None:
raise QUICInvalidMultiaddrError(f"Could not extract host/port from {maddr}")
return host, port
except Exception as e:
raise QUICInvalidMultiaddrError(
f"Failed to parse QUIC multiaddr {maddr}: {e}"
) from e
def multiaddr_to_quic_version(maddr: multiaddr.Multiaddr) -> TProtocol:
"""
Determine QUIC version from multiaddr.
Args:
maddr: QUIC multiaddr
Returns:
QUIC version identifier ("quic-v1" or "quic")
Raises:
QUICInvalidMultiaddrError: If multiaddr doesn't contain QUIC protocol
"""
try:
addr_str = str(maddr)
if f"/{QUIC_V1_PROTOCOL}" in addr_str:
return QUIC_V1_PROTOCOL # RFC 9000
elif f"/{QUIC_DRAFT29_PROTOCOL}" in addr_str:
return QUIC_DRAFT29_PROTOCOL # draft-29
else:
raise QUICInvalidMultiaddrError(f"No QUIC protocol found in {maddr}")
except Exception as e:
raise QUICInvalidMultiaddrError(
f"Failed to determine QUIC version from {maddr}: {e}"
) from e
def create_quic_multiaddr(
host: str, port: int, version: str = "quic-v1"
) -> multiaddr.Multiaddr:
"""
Create a QUIC multiaddr from host, port, and version.
Args:
host: IP address (IPv4 or IPv6)
port: UDP port number
version: QUIC version ("quic-v1" or "quic")
Returns:
QUIC multiaddr
Raises:
QUICInvalidMultiaddrError: If invalid parameters provided
"""
try:
# Determine IP version
try:
ip = ipaddress.ip_address(host)
if isinstance(ip, ipaddress.IPv4Address):
ip_proto = IP4_PROTOCOL
else:
ip_proto = IP6_PROTOCOL
except ValueError:
raise QUICInvalidMultiaddrError(f"Invalid IP address: {host}")
# Validate port
if not (0 <= port <= 65535):
raise QUICInvalidMultiaddrError(f"Invalid port: {port}")
# Validate and normalize QUIC version
if version == "quic-v1" or version == "/quic-v1":
quic_proto = QUIC_V1_PROTOCOL
elif version == "quic" or version == "/quic":
quic_proto = QUIC_DRAFT29_PROTOCOL
else:
raise QUICInvalidMultiaddrError(f"Invalid QUIC version: {version}")
# Construct multiaddr
addr_str = f"/{ip_proto}/{host}/{UDP_PROTOCOL}/{port}/{quic_proto}"
return multiaddr.Multiaddr(addr_str)
except Exception as e:
raise QUICInvalidMultiaddrError(f"Failed to create QUIC multiaddr: {e}") from e
def quic_version_to_wire_format(version: TProtocol) -> int:
"""
Convert QUIC version string to wire format integer for aioquic.
Args:
version: QUIC version string ("quic-v1" or "quic")
Returns:
Wire format version number
Raises:
QUICUnsupportedVersionError: If version is not supported
"""
wire_version = QUIC_VERSION_MAPPINGS.get(version)
if wire_version is None:
raise QUICUnsupportedVersionError(f"Unsupported QUIC version: {version}")
return wire_version
def custom_quic_version_to_wire_format(version: TProtocol) -> int:
"""
Convert QUIC version string to wire format integer for aioquic.
Args:
version: QUIC version string ("quic-v1" or "quic")
Returns:
Wire format version number
Raises:
QUICUnsupportedVersionError: If version is not supported
"""
wire_version = CUSTOM_QUIC_VERSION_MAPPING.get(version)
if wire_version is None:
raise QUICUnsupportedVersionError(f"Unsupported QUIC version: {version}")
return wire_version
def get_alpn_protocols() -> list[str]:
"""
Get ALPN protocols for libp2p over QUIC.
Returns:
List of ALPN protocol identifiers
"""
return LIBP2P_ALPN_PROTOCOLS.copy()
def normalize_quic_multiaddr(maddr: multiaddr.Multiaddr) -> multiaddr.Multiaddr:
"""
Normalize a QUIC multiaddr to canonical form.
Args:
maddr: Input QUIC multiaddr
Returns:
Normalized multiaddr
Raises:
QUICInvalidMultiaddrError: If not a valid QUIC multiaddr
"""
if not is_quic_multiaddr(maddr):
raise QUICInvalidMultiaddrError(f"Not a QUIC multiaddr: {maddr}")
host, port = quic_multiaddr_to_endpoint(maddr)
version = multiaddr_to_quic_version(maddr)
return create_quic_multiaddr(host, port, version)
def create_server_config_from_base(
base_config: QuicConfiguration,
security_manager: QUICTLSConfigManager | None = None,
transport_config: QUICTransportConfig | None = None,
) -> QuicConfiguration:
"""
Create a server configuration without using deepcopy.
Manually copies attributes while handling cryptography objects properly.
"""
try:
# Create new server configuration from scratch
server_config = QuicConfiguration(is_client=False)
server_config.verify_mode = ssl.CERT_NONE
# Copy basic configuration attributes (these are safe to copy)
copyable_attrs = [
"alpn_protocols",
"verify_mode",
"max_datagram_frame_size",
"idle_timeout",
"max_concurrent_streams",
"supported_versions",
"max_data",
"max_stream_data",
"stateless_retry",
"quantum_readiness_test",
]
for attr in copyable_attrs:
if hasattr(base_config, attr):
value = getattr(base_config, attr)
if value is not None:
setattr(server_config, attr, value)
# Handle cryptography objects - these need direct reference, not copying
crypto_attrs = [
"certificate",
"private_key",
"certificate_chain",
"ca_certs",
]
for attr in crypto_attrs:
if hasattr(base_config, attr):
value = getattr(base_config, attr)
if value is not None:
setattr(server_config, attr, value)
# Apply security manager configuration if available
if security_manager:
try:
server_tls_config = security_manager.create_server_config()
# Override with security manager's TLS configuration
if server_tls_config.certificate:
server_config.certificate = server_tls_config.certificate
if server_tls_config.private_key:
server_config.private_key = server_tls_config.private_key
if server_tls_config.certificate_chain:
server_config.certificate_chain = (
server_tls_config.certificate_chain
)
if server_tls_config.alpn_protocols:
server_config.alpn_protocols = server_tls_config.alpn_protocols
server_tls_config.request_client_certificate = True
if getattr(server_tls_config, "request_client_certificate", False):
server_config._libp2p_request_client_cert = True # type: ignore
else:
logger.error(
"🔧 Failed to set request_client_certificate in server config"
)
except Exception as e:
logger.warning(f"Failed to apply security manager config: {e}")
# Set transport-specific defaults if provided
if transport_config:
if server_config.idle_timeout == 0:
server_config.idle_timeout = getattr(
transport_config, "idle_timeout", 30.0
)
if server_config.max_datagram_frame_size is None:
server_config.max_datagram_frame_size = getattr(
transport_config, "max_datagram_size", 1200
)
# Ensure we have ALPN protocols
if not server_config.alpn_protocols:
server_config.alpn_protocols = ["libp2p"]
logger.debug("Successfully created server config without deepcopy")
return server_config
except Exception as e:
logger.error(f"Failed to create server config: {e}")
raise
def create_client_config_from_base(
base_config: QuicConfiguration,
security_manager: QUICTLSConfigManager | None = None,
transport_config: QUICTransportConfig | None = None,
) -> QuicConfiguration:
"""
Create a client configuration without using deepcopy.
"""
try:
# Create new client configuration from scratch
client_config = QuicConfiguration(is_client=True)
client_config.verify_mode = ssl.CERT_NONE
# Copy basic configuration attributes
copyable_attrs = [
"alpn_protocols",
"verify_mode",
"max_datagram_frame_size",
"idle_timeout",
"max_concurrent_streams",
"supported_versions",
"max_data",
"max_stream_data",
"quantum_readiness_test",
]
for attr in copyable_attrs:
if hasattr(base_config, attr):
value = getattr(base_config, attr)
if value is not None:
setattr(client_config, attr, value)
# Handle cryptography objects - these need direct reference, not copying
crypto_attrs = [
"certificate",
"private_key",
"certificate_chain",
"ca_certs",
]
for attr in crypto_attrs:
if hasattr(base_config, attr):
value = getattr(base_config, attr)
if value is not None:
setattr(client_config, attr, value)
# Apply security manager configuration if available
if security_manager:
try:
client_tls_config = security_manager.create_client_config()
# Override with security manager's TLS configuration
if client_tls_config.certificate:
client_config.certificate = client_tls_config.certificate
if client_tls_config.private_key:
client_config.private_key = client_tls_config.private_key
if client_tls_config.certificate_chain:
client_config.certificate_chain = (
client_tls_config.certificate_chain
)
if client_tls_config.alpn_protocols:
client_config.alpn_protocols = client_tls_config.alpn_protocols
except Exception as e:
logger.warning(f"Failed to apply security manager config: {e}")
# Ensure we have ALPN protocols
if not client_config.alpn_protocols:
client_config.alpn_protocols = ["libp2p"]
logger.debug("Successfully created client config without deepcopy")
return client_config
except Exception as e:
logger.error(f"Failed to create client config: {e}")
raise

View File

@ -0,0 +1,267 @@
"""
Transport registry for dynamic transport selection based on multiaddr protocols.
"""
from collections.abc import Callable
import logging
from typing import Any
from multiaddr import Multiaddr
from multiaddr.protocols import Protocol
from libp2p.abc import ITransport
from libp2p.transport.tcp.tcp import TCP
from libp2p.transport.upgrader import TransportUpgrader
from libp2p.transport.websocket.multiaddr_utils import (
is_valid_websocket_multiaddr,
)
# Import QUIC utilities here to avoid circular imports
def _get_quic_transport() -> Any:
from libp2p.transport.quic.transport import QUICTransport
return QUICTransport
def _get_quic_validation() -> Callable[[Multiaddr], bool]:
from libp2p.transport.quic.utils import is_quic_multiaddr
return is_quic_multiaddr
# Import WebsocketTransport here to avoid circular imports
def _get_websocket_transport() -> Any:
from libp2p.transport.websocket.transport import WebsocketTransport
return WebsocketTransport
logger = logging.getLogger("libp2p.transport.registry")
def _is_valid_tcp_multiaddr(maddr: Multiaddr) -> bool:
"""
Validate that a multiaddr has a valid TCP structure.
:param maddr: The multiaddr to validate
:return: True if valid TCP structure, False otherwise
"""
try:
# TCP multiaddr should have structure like /ip4/127.0.0.1/tcp/8080
# or /ip6/::1/tcp/8080
protocols: list[Protocol] = list(maddr.protocols())
# Must have at least 2 protocols: network (ip4/ip6) + tcp
if len(protocols) < 2:
return False
# First protocol should be a network protocol (ip4, ip6, dns4, dns6)
if protocols[0].name not in ["ip4", "ip6", "dns4", "dns6"]:
return False
# Second protocol should be tcp
if protocols[1].name != "tcp":
return False
# Should not have any protocols after tcp (unless it's a valid
# continuation like p2p)
# For now, we'll be strict and only allow network + tcp
if len(protocols) > 2:
# Check if the additional protocols are valid continuations
valid_continuations = ["p2p"] # Add more as needed
for i in range(2, len(protocols)):
if protocols[i].name not in valid_continuations:
return False
return True
except Exception:
return False
class TransportRegistry:
"""
Registry for mapping multiaddr protocols to transport implementations.
"""
def __init__(self) -> None:
self._transports: dict[str, type[ITransport]] = {}
self._register_default_transports()
def _register_default_transports(self) -> None:
"""Register the default transport implementations."""
# Register TCP transport for /tcp protocol
self.register_transport("tcp", TCP)
# Register WebSocket transport for /ws and /wss protocols
WebsocketTransport = _get_websocket_transport()
self.register_transport("ws", WebsocketTransport)
self.register_transport("wss", WebsocketTransport)
# Register QUIC transport for /quic and /quic-v1 protocols
QUICTransport = _get_quic_transport()
self.register_transport("quic", QUICTransport)
self.register_transport("quic-v1", QUICTransport)
def register_transport(
self, protocol: str, transport_class: type[ITransport]
) -> None:
"""
Register a transport class for a specific protocol.
:param protocol: The protocol identifier (e.g., "tcp", "ws")
:param transport_class: The transport class to register
"""
self._transports[protocol] = transport_class
logger.debug(
f"Registered transport {transport_class.__name__} for protocol {protocol}"
)
def get_transport(self, protocol: str) -> type[ITransport] | None:
"""
Get the transport class for a specific protocol.
:param protocol: The protocol identifier
:return: The transport class or None if not found
"""
return self._transports.get(protocol)
def get_supported_protocols(self) -> list[str]:
"""Get list of supported transport protocols."""
return list(self._transports.keys())
def create_transport(
self, protocol: str, upgrader: TransportUpgrader | None = None, **kwargs: Any
) -> ITransport | None:
"""
Create a transport instance for a specific protocol.
:param protocol: The protocol identifier
:param upgrader: The transport upgrader instance (required for WebSocket)
:param kwargs: Additional arguments for transport construction
:return: Transport instance or None if protocol not supported or creation fails
"""
transport_class = self.get_transport(protocol)
if transport_class is None:
return None
try:
if protocol in ["ws", "wss"]:
# WebSocket transport requires upgrader
if upgrader is None:
logger.warning(
f"WebSocket transport '{protocol}' requires upgrader"
)
return None
# Use explicit WebsocketTransport to avoid type issues
WebsocketTransport = _get_websocket_transport()
return WebsocketTransport(
upgrader,
tls_client_config=kwargs.get("tls_client_config"),
tls_server_config=kwargs.get("tls_server_config"),
handshake_timeout=kwargs.get("handshake_timeout", 15.0),
)
elif protocol in ["quic", "quic-v1"]:
# QUIC transport requires private_key
private_key = kwargs.get("private_key")
if private_key is None:
logger.warning(f"QUIC transport '{protocol}' requires private_key")
return None
# Use explicit QUICTransport to avoid type issues
QUICTransport = _get_quic_transport()
config = kwargs.get("config")
return QUICTransport(private_key, config)
else:
# TCP transport doesn't require upgrader
return transport_class()
except Exception as e:
logger.error(f"Failed to create transport for protocol {protocol}: {e}")
return None
# Global transport registry instance (lazy initialization)
_global_registry: TransportRegistry | None = None
def get_transport_registry() -> TransportRegistry:
"""Get the global transport registry instance."""
global _global_registry
if _global_registry is None:
_global_registry = TransportRegistry()
return _global_registry
def register_transport(protocol: str, transport_class: type[ITransport]) -> None:
"""Register a transport class in the global registry."""
registry = get_transport_registry()
registry.register_transport(protocol, transport_class)
def create_transport_for_multiaddr(
maddr: Multiaddr, upgrader: TransportUpgrader, **kwargs: Any
) -> ITransport | None:
"""
Create the appropriate transport for a given multiaddr.
:param maddr: The multiaddr to create transport for
:param upgrader: The transport upgrader instance
:param kwargs: Additional arguments for transport construction
(e.g., private_key for QUIC)
:return: Transport instance or None if no suitable transport found
"""
try:
# Get all protocols in the multiaddr
protocols = [proto.name for proto in maddr.protocols()]
# Check for supported transport protocols in order of preference
# We need to validate that the multiaddr structure is valid for our transports
if "quic" in protocols or "quic-v1" in protocols:
# For QUIC, we need a valid structure like:
# /ip4/127.0.0.1/udp/4001/quic
# /ip4/127.0.0.1/udp/4001/quic-v1
is_quic_multiaddr = _get_quic_validation()
if is_quic_multiaddr(maddr):
# Determine QUIC version
registry = get_transport_registry()
if "quic-v1" in protocols:
return registry.create_transport("quic-v1", upgrader, **kwargs)
else:
return registry.create_transport("quic", upgrader, **kwargs)
elif "ws" in protocols or "wss" in protocols or "tls" in protocols:
# For WebSocket, we need a valid structure like:
# /ip4/127.0.0.1/tcp/8080/ws (insecure)
# /ip4/127.0.0.1/tcp/8080/wss (secure)
# /ip4/127.0.0.1/tcp/8080/tls/ws (secure with TLS)
# /ip4/127.0.0.1/tcp/8080/tls/sni/example.com/ws (secure with SNI)
if is_valid_websocket_multiaddr(maddr):
# Determine if this is a secure WebSocket connection
registry = get_transport_registry()
if "wss" in protocols or "tls" in protocols:
return registry.create_transport("wss", upgrader, **kwargs)
else:
return registry.create_transport("ws", upgrader, **kwargs)
elif "tcp" in protocols:
# For TCP, we need a valid structure like /ip4/127.0.0.1/tcp/8080
# Check if the multiaddr has proper TCP structure
if _is_valid_tcp_multiaddr(maddr):
registry = get_transport_registry()
return registry.create_transport("tcp", upgrader)
# If no supported transport protocol found or structure is invalid, return None
logger.warning(
f"No supported transport protocol found or invalid structure in "
f"multiaddr: {maddr}"
)
return None
except Exception as e:
# Handle any errors gracefully (e.g., invalid multiaddr)
logger.warning(f"Error processing multiaddr {maddr}: {e}")
return None
def get_supported_transport_protocols() -> list[str]:
"""Get list of supported transport protocols from the global registry."""
registry = get_transport_registry()
return registry.get_supported_protocols()

View File

@ -14,6 +14,9 @@ from libp2p.protocol_muxer.exceptions import (
MultiselectClientError,
MultiselectError,
)
from libp2p.protocol_muxer.multiselect import (
DEFAULT_NEGOTIATE_TIMEOUT,
)
from libp2p.security.exceptions import (
HandshakeFailure,
)
@ -37,9 +40,12 @@ class TransportUpgrader:
self,
secure_transports_by_protocol: TSecurityOptions,
muxer_transports_by_protocol: TMuxerOptions,
negotiate_timeout: int = DEFAULT_NEGOTIATE_TIMEOUT,
):
self.security_multistream = SecurityMultistream(secure_transports_by_protocol)
self.muxer_multistream = MuxerMultistream(muxer_transports_by_protocol)
self.muxer_multistream = MuxerMultistream(
muxer_transports_by_protocol, negotiate_timeout
)
async def upgrade_security(
self,

View File

@ -0,0 +1,198 @@
import logging
import time
from typing import Any
import trio
from libp2p.io.abc import ReadWriteCloser
from libp2p.io.exceptions import IOException
logger = logging.getLogger(__name__)
class P2PWebSocketConnection(ReadWriteCloser):
"""
Wraps a WebSocketConnection to provide the raw stream interface
that libp2p protocols expect.
Implements production-ready buffer management and flow control
as recommended in the libp2p WebSocket specification.
"""
def __init__(
self,
ws_connection: Any,
ws_context: Any = None,
is_secure: bool = False,
max_buffered_amount: int = 4 * 1024 * 1024,
) -> None:
self._ws_connection = ws_connection
self._ws_context = ws_context
self._is_secure = is_secure
self._read_buffer = b""
self._read_lock = trio.Lock()
self._connection_start_time = time.time()
self._bytes_read = 0
self._bytes_written = 0
self._closed = False
self._close_lock = trio.Lock()
self._max_buffered_amount = max_buffered_amount
self._write_lock = trio.Lock()
async def write(self, data: bytes) -> None:
"""Write data with flow control and buffer management"""
if self._closed:
raise IOException("Connection is closed")
async with self._write_lock:
try:
logger.debug(f"WebSocket writing {len(data)} bytes")
# Check buffer amount for flow control
if hasattr(self._ws_connection, "bufferedAmount"):
buffered = self._ws_connection.bufferedAmount
if buffered > self._max_buffered_amount:
logger.warning(f"WebSocket buffer full: {buffered} bytes")
# In production, you might want to
# wait or implement backpressure
# For now, we'll continue but log the warning
# Send as a binary WebSocket message
await self._ws_connection.send_message(data)
self._bytes_written += len(data)
logger.debug(f"WebSocket wrote {len(data)} bytes successfully")
except Exception as e:
logger.error(f"WebSocket write failed: {e}")
self._closed = True
raise IOException from e
async def read(self, n: int | None = None) -> bytes:
"""
Read up to n bytes (if n is given), else read up to 64KiB.
This implementation provides byte-level access to WebSocket messages,
which is required for libp2p protocol compatibility.
For WebSocket compatibility with libp2p protocols, this method:
1. Buffers incoming WebSocket messages
2. Returns exactly the requested number of bytes when n is specified
3. Accumulates multiple WebSocket messages if needed to satisfy the request
4. Returns empty bytes (not raises) when connection is closed and no data
available
"""
if self._closed:
raise IOException("Connection is closed")
async with self._read_lock:
try:
# If n is None, read at least one message and return all buffered data
if n is None:
if not self._read_buffer:
try:
# Use a short timeout to avoid blocking indefinitely
with trio.fail_after(1.0): # 1 second timeout
message = await self._ws_connection.get_message()
if isinstance(message, str):
message = message.encode("utf-8")
self._read_buffer = message
except trio.TooSlowError:
# No message available within timeout
return b""
except Exception:
# Return empty bytes if no data available
# (connection closed)
return b""
result = self._read_buffer
self._read_buffer = b""
self._bytes_read += len(result)
return result
# For specific byte count requests, return UP TO n bytes (not exactly n)
# This matches TCP semantics where read(1024) returns available data
# up to 1024 bytes
# If we don't have any data buffered, try to get at least one message
if not self._read_buffer:
try:
# Use a short timeout to avoid blocking indefinitely
with trio.fail_after(1.0): # 1 second timeout
message = await self._ws_connection.get_message()
if isinstance(message, str):
message = message.encode("utf-8")
self._read_buffer = message
except trio.TooSlowError:
return b"" # No data available
except Exception:
return b""
# Now return up to n bytes from the buffer (TCP-like semantics)
if len(self._read_buffer) == 0:
return b""
# Return up to n bytes (like TCP read())
result = self._read_buffer[:n]
self._read_buffer = self._read_buffer[len(result) :]
self._bytes_read += len(result)
return result
except Exception as e:
logger.error(f"WebSocket read failed: {e}")
raise IOException from e
async def close(self) -> None:
"""Close the WebSocket connection. This method is idempotent."""
async with self._close_lock:
if self._closed:
return # Already closed
logger.debug("WebSocket connection closing")
self._closed = True
try:
# Always close the connection directly, avoid context manager issues
# The context manager may be causing cancel scope corruption
logger.debug("WebSocket closing connection directly")
await self._ws_connection.aclose()
# Exit the context manager if we have one
if self._ws_context is not None:
await self._ws_context.__aexit__(None, None, None)
except Exception as e:
logger.error(f"WebSocket close error: {e}")
# Don't raise here, as close() should be idempotent
finally:
logger.debug("WebSocket connection closed")
def is_closed(self) -> bool:
"""Check if the connection is closed"""
return self._closed
def conn_state(self) -> dict[str, Any]:
"""
Return connection state information similar to Go's ConnState() method.
:return: Dictionary containing connection state information
"""
current_time = time.time()
return {
"transport": "websocket",
"secure": self._is_secure,
"connection_duration": current_time - self._connection_start_time,
"bytes_read": self._bytes_read,
"bytes_written": self._bytes_written,
"total_bytes": self._bytes_read + self._bytes_written,
}
def get_remote_address(self) -> tuple[str, int] | None:
# Try to get remote address from the WebSocket connection
try:
remote = self._ws_connection.remote
if hasattr(remote, "address") and hasattr(remote, "port"):
return str(remote.address), int(remote.port)
elif isinstance(remote, str):
# Parse address:port format
if ":" in remote:
host, port = remote.rsplit(":", 1)
return host, int(port)
except Exception:
pass
return None

View File

@ -0,0 +1,225 @@
from collections.abc import Awaitable, Callable
import logging
import ssl
from typing import Any
from multiaddr import Multiaddr
import trio
from trio_typing import TaskStatus
from trio_websocket import serve_websocket
from libp2p.abc import IListener
from libp2p.custom_types import THandler
from libp2p.transport.upgrader import TransportUpgrader
from libp2p.transport.websocket.multiaddr_utils import parse_websocket_multiaddr
from .connection import P2PWebSocketConnection
logger = logging.getLogger("libp2p.transport.websocket.listener")
class WebsocketListener(IListener):
"""
Listen on /ip4/.../tcp/.../ws addresses, handshake WS, wrap into RawConnection.
"""
def __init__(
self,
handler: THandler,
upgrader: TransportUpgrader,
tls_config: ssl.SSLContext | None = None,
handshake_timeout: float = 15.0,
) -> None:
self._handler = handler
self._upgrader = upgrader
self._tls_config = tls_config
self._handshake_timeout = handshake_timeout
self._server = None
self._shutdown_event = trio.Event()
self._nursery: trio.Nursery | None = None
self._listeners: Any = None
self._is_wss = False # Track whether this is a WSS listener
async def listen(self, maddr: Multiaddr, nursery: trio.Nursery) -> bool:
logger.debug(f"WebsocketListener.listen called with {maddr}")
# Parse the WebSocket multiaddr to determine if it's secure
try:
parsed = parse_websocket_multiaddr(maddr)
except ValueError as e:
raise ValueError(f"Invalid WebSocket multiaddr: {e}") from e
# Check if WSS is requested but no TLS config provided
if parsed.is_wss and self._tls_config is None:
raise ValueError(
f"Cannot listen on WSS address {maddr} without TLS configuration"
)
# Store whether this is a WSS listener
self._is_wss = parsed.is_wss
# Extract host and port from the base multiaddr
host = (
parsed.rest_multiaddr.value_for_protocol("ip4")
or parsed.rest_multiaddr.value_for_protocol("ip6")
or parsed.rest_multiaddr.value_for_protocol("dns")
or parsed.rest_multiaddr.value_for_protocol("dns4")
or parsed.rest_multiaddr.value_for_protocol("dns6")
or "0.0.0.0"
)
port_str = parsed.rest_multiaddr.value_for_protocol("tcp")
if port_str is None:
raise ValueError(f"No TCP port found in multiaddr: {maddr}")
port = int(port_str)
logger.debug(
f"WebsocketListener: host={host}, port={port}, secure={parsed.is_wss}"
)
async def serve_websocket_tcp(
handler: Callable[[Any], Awaitable[None]],
port: int,
host: str,
task_status: TaskStatus[Any],
) -> None:
"""Start TCP server and handle WebSocket connections manually"""
logger.debug(
"serve_websocket_tcp %s %s (secure=%s)", host, port, parsed.is_wss
)
async def websocket_handler(request: Any) -> None:
"""Handle WebSocket requests"""
logger.debug("WebSocket request received")
try:
# Apply handshake timeout
with trio.fail_after(self._handshake_timeout):
# Accept the WebSocket connection
ws_connection = await request.accept()
logger.debug("WebSocket handshake successful")
# Create the WebSocket connection wrapper
conn = P2PWebSocketConnection(
ws_connection, is_secure=parsed.is_wss
) # type: ignore[no-untyped-call]
# Call the handler function that was passed to create_listener
# This handler will handle the security and muxing upgrades
logger.debug("Calling connection handler")
await self._handler(conn)
# Don't keep the connection alive indefinitely
# Let the handler manage the connection lifecycle
logger.debug(
"Handler completed, connection will be managed by handler"
)
except trio.TooSlowError:
logger.debug(
f"WebSocket handshake timeout after {self._handshake_timeout}s"
)
try:
await request.reject(408) # Request Timeout
except Exception:
pass
except Exception as e:
logger.debug(f"WebSocket connection error: {e}")
logger.debug(f"Error type: {type(e)}")
import traceback
logger.debug(f"Traceback: {traceback.format_exc()}")
# Reject the connection
try:
await request.reject(400)
except Exception:
pass
# Use trio_websocket.serve_websocket for proper WebSocket handling
ssl_context = self._tls_config if parsed.is_wss else None
await serve_websocket(
websocket_handler, host, port, ssl_context, task_status=task_status
)
# Store the nursery for shutdown
self._nursery = nursery
# Start the server using nursery.start() like TCP does
logger.debug("Calling nursery.start()...")
started_listeners = await nursery.start(
serve_websocket_tcp,
None, # No handler needed since it's defined inside serve_websocket_tcp
port,
host,
)
logger.debug(f"nursery.start() returned: {started_listeners}")
if started_listeners is None:
logger.error(f"Failed to start WebSocket listener for {maddr}")
return False
# Store the listeners for get_addrs() and close() - these are real
# SocketListener objects
self._listeners = started_listeners
logger.debug(
"WebsocketListener.listen returning True with WebSocketServer object"
)
return True
def get_addrs(self) -> tuple[Multiaddr, ...]:
if not hasattr(self, "_listeners") or not self._listeners:
logger.debug("No listeners available for get_addrs()")
return ()
# Handle WebSocketServer objects
if hasattr(self._listeners, "port"):
# This is a WebSocketServer object
port = self._listeners.port
# Create a multiaddr from the port with correct WSS/WS protocol
protocol = "wss" if self._is_wss else "ws"
return (Multiaddr(f"/ip4/127.0.0.1/tcp/{port}/{protocol}"),)
else:
# This is a list of listeners (like TCP)
listeners = self._listeners
# Get addresses from listeners like TCP does
return tuple(
_multiaddr_from_socket(listener.socket, self._is_wss)
for listener in listeners
)
async def close(self) -> None:
"""Close the WebSocket listener and stop accepting new connections"""
logger.debug("WebsocketListener.close called")
if hasattr(self, "_listeners") and self._listeners:
# Signal shutdown
self._shutdown_event.set()
# Close the WebSocket server
if hasattr(self._listeners, "aclose"):
# This is a WebSocketServer object
logger.debug("Closing WebSocket server")
await self._listeners.aclose()
logger.debug("WebSocket server closed")
elif isinstance(self._listeners, (list, tuple)):
# This is a list of listeners (like TCP)
logger.debug("Closing TCP listeners")
for listener in self._listeners:
await listener.aclose()
logger.debug("TCP listeners closed")
else:
# Unknown type, try to close it directly
logger.debug("Closing unknown listener type")
if hasattr(self._listeners, "close"):
self._listeners.close()
logger.debug("Unknown listener closed")
# Clear the listeners reference
self._listeners = None
logger.debug("WebsocketListener.close completed")
def _multiaddr_from_socket(
socket: trio.socket.SocketType, is_wss: bool = False
) -> Multiaddr:
"""Convert socket to multiaddr"""
ip, port = socket.getsockname()
protocol = "wss" if is_wss else "ws"
return Multiaddr(f"/ip4/{ip}/tcp/{port}/{protocol}")

View File

@ -0,0 +1,202 @@
"""
WebSocket multiaddr parsing utilities.
"""
from typing import NamedTuple
from multiaddr import Multiaddr
from multiaddr.protocols import Protocol
class ParsedWebSocketMultiaddr(NamedTuple):
"""Parsed WebSocket multiaddr information."""
is_wss: bool
sni: str | None
rest_multiaddr: Multiaddr
def parse_websocket_multiaddr(maddr: Multiaddr) -> ParsedWebSocketMultiaddr:
"""
Parse a WebSocket multiaddr and extract security information.
:param maddr: The multiaddr to parse
:return: Parsed WebSocket multiaddr information
:raises ValueError: If the multiaddr is not a valid WebSocket multiaddr
"""
# First validate that this is a valid WebSocket multiaddr
if not is_valid_websocket_multiaddr(maddr):
raise ValueError(f"Not a valid WebSocket multiaddr: {maddr}")
protocols = list(maddr.protocols())
# Find the WebSocket protocol and check for security
is_wss = False
sni = None
ws_index = -1
tls_index = -1
sni_index = -1
# Find protocol indices
for i, protocol in enumerate(protocols):
if protocol.name == "ws":
ws_index = i
elif protocol.name == "wss":
ws_index = i
is_wss = True
elif protocol.name == "tls":
tls_index = i
elif protocol.name == "sni":
sni_index = i
sni = protocol.value
if ws_index == -1:
raise ValueError("Not a WebSocket multiaddr")
# Handle /wss protocol (convert to /tls/ws internally)
if is_wss and tls_index == -1:
# Convert /wss to /tls/ws format
# Remove /wss to get the base multiaddr
without_wss = maddr.decapsulate(Multiaddr("/wss"))
return ParsedWebSocketMultiaddr(
is_wss=True, sni=None, rest_multiaddr=without_wss
)
# Handle /tls/ws and /tls/sni/.../ws formats
if tls_index != -1:
is_wss = True
# Extract the base multiaddr (everything before /tls)
# For /ip4/127.0.0.1/tcp/8080/tls/ws, we want /ip4/127.0.0.1/tcp/8080
# Use multiaddr methods to properly extract the base
rest_multiaddr = maddr
# Remove /tls/ws or /tls/sni/.../ws from the end
if sni_index != -1:
# /tls/sni/example.com/ws format
rest_multiaddr = rest_multiaddr.decapsulate(Multiaddr("/ws"))
rest_multiaddr = rest_multiaddr.decapsulate(Multiaddr(f"/sni/{sni}"))
rest_multiaddr = rest_multiaddr.decapsulate(Multiaddr("/tls"))
else:
# /tls/ws format
rest_multiaddr = rest_multiaddr.decapsulate(Multiaddr("/ws"))
rest_multiaddr = rest_multiaddr.decapsulate(Multiaddr("/tls"))
return ParsedWebSocketMultiaddr(
is_wss=is_wss, sni=sni, rest_multiaddr=rest_multiaddr
)
# Regular /ws multiaddr - remove /ws and any additional protocols
rest_multiaddr = maddr.decapsulate(Multiaddr("/ws"))
return ParsedWebSocketMultiaddr(
is_wss=False, sni=None, rest_multiaddr=rest_multiaddr
)
def is_valid_websocket_multiaddr(maddr: Multiaddr) -> bool:
"""
Validate that a multiaddr has a valid WebSocket structure.
:param maddr: The multiaddr to validate
:return: True if valid WebSocket structure, False otherwise
"""
try:
# WebSocket multiaddr should have structure like:
# /ip4/127.0.0.1/tcp/8080/ws (insecure)
# /ip4/127.0.0.1/tcp/8080/wss (secure)
# /ip4/127.0.0.1/tcp/8080/tls/ws (secure with TLS)
# /ip4/127.0.0.1/tcp/8080/tls/sni/example.com/ws (secure with SNI)
protocols: list[Protocol] = list(maddr.protocols())
# Must have at least 3 protocols: network (ip4/ip6/dns4/dns6) + tcp + ws/wss
if len(protocols) < 3:
return False
# First protocol should be a network protocol (ip4, ip6, dns, dns4, dns6)
if protocols[0].name not in ["ip4", "ip6", "dns", "dns4", "dns6"]:
return False
# Second protocol should be tcp
if protocols[1].name != "tcp":
return False
# Check for valid WebSocket protocols
ws_protocols = ["ws", "wss"]
tls_protocols = ["tls"]
sni_protocols = ["sni"]
# Find the WebSocket protocol
ws_protocol_found = False
tls_found = False
# sni_found = False # Not used currently
for i, protocol in enumerate(protocols[2:], start=2):
if protocol.name in ws_protocols:
ws_protocol_found = True
break
elif protocol.name in tls_protocols:
tls_found = True
elif protocol.name in sni_protocols:
pass # sni_found = True # Not used in current implementation
if not ws_protocol_found:
return False
# Validate protocol sequence
# For /ws: network + tcp + ws
# For /wss: network + tcp + wss
# For /tls/ws: network + tcp + tls + ws
# For /tls/sni/example.com/ws: network + tcp + tls + sni + ws
# Check if it's a simple /ws or /wss
if len(protocols) == 3:
return protocols[2].name in ["ws", "wss"]
# Check for /tls/ws or /tls/sni/.../ws patterns
if tls_found:
# Must end with /ws (not /wss when using /tls)
if protocols[-1].name != "ws":
return False
# Check for valid TLS sequence
tls_index = None
for i, protocol in enumerate(protocols[2:], start=2):
if protocol.name == "tls":
tls_index = i
break
if tls_index is None:
return False
# After tls, we can have sni, then ws
remaining_protocols = protocols[tls_index + 1 :]
if len(remaining_protocols) == 1:
# /tls/ws
return remaining_protocols[0].name == "ws"
elif len(remaining_protocols) == 2:
# /tls/sni/example.com/ws
return (
remaining_protocols[0].name == "sni"
and remaining_protocols[1].name == "ws"
)
else:
return False
# If we have more than 3 protocols but no TLS, check for valid continuations
# Allow additional protocols after the WebSocket protocol (like /p2p)
valid_continuations = ["p2p"]
# Find the WebSocket protocol index
ws_index = None
for i, protocol in enumerate(protocols):
if protocol.name in ["ws", "wss"]:
ws_index = i
break
if ws_index is not None:
# Check protocols after the WebSocket protocol
for i in range(ws_index + 1, len(protocols)):
if protocols[i].name not in valid_continuations:
return False
return True
except Exception:
return False

View File

@ -0,0 +1,229 @@
import logging
import ssl
from multiaddr import Multiaddr
from libp2p.abc import IListener, ITransport
from libp2p.custom_types import THandler
from libp2p.network.connection.raw_connection import RawConnection
from libp2p.transport.exceptions import OpenConnectionError
from libp2p.transport.upgrader import TransportUpgrader
from libp2p.transport.websocket.multiaddr_utils import parse_websocket_multiaddr
from .connection import P2PWebSocketConnection
from .listener import WebsocketListener
logger = logging.getLogger(__name__)
class WebsocketTransport(ITransport):
"""
Libp2p WebSocket transport: dial and listen on /ip4/.../tcp/.../ws and /wss
Implements production-ready WebSocket transport with:
- Flow control and buffer management
- Connection limits and rate limiting
- Proper error handling and cleanup
- Support for both WS and WSS protocols
- TLS configuration and handshake timeout
"""
def __init__(
self,
upgrader: TransportUpgrader,
tls_client_config: ssl.SSLContext | None = None,
tls_server_config: ssl.SSLContext | None = None,
handshake_timeout: float = 15.0,
max_buffered_amount: int = 4 * 1024 * 1024,
):
self._upgrader = upgrader
self._tls_client_config = tls_client_config
self._tls_server_config = tls_server_config
self._handshake_timeout = handshake_timeout
self._max_buffered_amount = max_buffered_amount
self._connection_count = 0
self._max_connections = 1000 # Production limit
async def dial(self, maddr: Multiaddr) -> RawConnection:
"""Dial a WebSocket connection to the given multiaddr."""
logger.debug(f"WebsocketTransport.dial called with {maddr}")
# Parse the WebSocket multiaddr to determine if it's secure
try:
parsed = parse_websocket_multiaddr(maddr)
except ValueError as e:
raise ValueError(f"Invalid WebSocket multiaddr: {e}") from e
# Extract host and port from the base multiaddr
host = (
parsed.rest_multiaddr.value_for_protocol("ip4")
or parsed.rest_multiaddr.value_for_protocol("ip6")
or parsed.rest_multiaddr.value_for_protocol("dns")
or parsed.rest_multiaddr.value_for_protocol("dns4")
or parsed.rest_multiaddr.value_for_protocol("dns6")
)
port_str = parsed.rest_multiaddr.value_for_protocol("tcp")
if port_str is None:
raise ValueError(f"No TCP port found in multiaddr: {maddr}")
port = int(port_str)
# Build WebSocket URL based on security
if parsed.is_wss:
ws_url = f"wss://{host}:{port}/"
else:
ws_url = f"ws://{host}:{port}/"
logger.debug(
f"WebsocketTransport.dial connecting to {ws_url} (secure={parsed.is_wss})"
)
try:
# Check connection limits
if self._connection_count >= self._max_connections:
raise OpenConnectionError(
f"Maximum connections reached: {self._max_connections}"
)
# Prepare SSL context for WSS connections
ssl_context = None
if parsed.is_wss:
if self._tls_client_config:
ssl_context = self._tls_client_config
else:
# Create default SSL context for client
ssl_context = ssl.create_default_context()
# Set SNI if available
if parsed.sni:
ssl_context.check_hostname = False
ssl_context.verify_mode = ssl.CERT_NONE
logger.debug(f"WebsocketTransport.dial opening connection to {ws_url}")
# Use a different approach: start background nursery that will persist
logger.debug("WebsocketTransport.dial establishing connection")
# Import trio-websocket functions
from trio_websocket import connect_websocket
from trio_websocket._impl import _url_to_host
# Parse the WebSocket URL to get host, port, resource
# like trio-websocket does
ws_host, ws_port, ws_resource, ws_ssl_context = _url_to_host(
ws_url, ssl_context
)
logger.debug(
f"WebsocketTransport.dial parsed URL: host={ws_host}, "
f"port={ws_port}, resource={ws_resource}"
)
# Create a background task manager for this connection
import trio
nursery_manager = trio.lowlevel.current_task().parent_nursery
if nursery_manager is None:
raise OpenConnectionError(
f"No parent nursery available for WebSocket connection to {maddr}"
)
# Apply timeout to the connection process
with trio.fail_after(self._handshake_timeout):
logger.debug("WebsocketTransport.dial connecting WebSocket")
ws = await connect_websocket(
nursery_manager, # Use the existing nursery from libp2p
ws_host,
ws_port,
ws_resource,
use_ssl=ws_ssl_context,
message_queue_size=1024, # Reasonable defaults
max_message_size=16 * 1024 * 1024, # 16MB max message
)
logger.debug("WebsocketTransport.dial WebSocket connection established")
# Create our connection wrapper with both WSS support and flow control
conn = P2PWebSocketConnection(
ws,
None,
is_secure=parsed.is_wss,
max_buffered_amount=self._max_buffered_amount,
)
logger.debug("WebsocketTransport.dial created P2PWebSocketConnection")
self._connection_count += 1
logger.debug(f"Total connections: {self._connection_count}")
return RawConnection(conn, initiator=True)
except trio.TooSlowError as e:
raise OpenConnectionError(
f"WebSocket handshake timeout after {self._handshake_timeout}s "
f"for {maddr}"
) from e
except Exception as e:
logger.error(f"Failed to dial WebSocket {maddr}: {e}")
raise OpenConnectionError(f"Failed to dial WebSocket {maddr}: {e}") from e
def create_listener(self, handler: THandler) -> IListener: # type: ignore[override]
"""
The type checker is incorrectly reporting this as an inconsistent override.
"""
logger.debug("WebsocketTransport.create_listener called")
return WebsocketListener(
handler, self._upgrader, self._tls_server_config, self._handshake_timeout
)
def resolve(self, maddr: Multiaddr) -> list[Multiaddr]:
"""
Resolve a WebSocket multiaddr, automatically adding SNI for DNS names.
Similar to Go's Resolve() method.
:param maddr: The multiaddr to resolve
:return: List of resolved multiaddrs
"""
try:
parsed = parse_websocket_multiaddr(maddr)
except ValueError as e:
logger.debug(f"Invalid WebSocket multiaddr for resolution: {e}")
return [maddr] # Return original if not a valid WebSocket multiaddr
logger.debug(
f"Parsed multiaddr {maddr}: is_wss={parsed.is_wss}, sni={parsed.sni}"
)
if not parsed.is_wss:
# No /tls/ws component, this isn't a secure websocket multiaddr
return [maddr]
if parsed.sni is not None:
# Already has SNI, return as-is
return [maddr]
# Try to extract DNS name from the base multiaddr
dns_name = None
for protocol_name in ["dns", "dns4", "dns6"]:
try:
dns_name = parsed.rest_multiaddr.value_for_protocol(protocol_name)
break
except Exception:
continue
if dns_name is None:
# No DNS name found, return original
return [maddr]
# Create new multiaddr with SNI
# For /dns/example.com/tcp/8080/wss ->
# /dns/example.com/tcp/8080/tls/sni/example.com/ws
try:
# Remove /wss and add /tls/sni/example.com/ws
without_wss = maddr.decapsulate(Multiaddr("/wss"))
sni_component = Multiaddr(f"/sni/{dns_name}")
resolved = (
without_wss.encapsulate(Multiaddr("/tls"))
.encapsulate(sni_component)
.encapsulate(Multiaddr("/ws"))
)
logger.debug(f"Resolved {maddr} to {resolved}")
return [resolved]
except Exception as e:
logger.debug(f"Failed to resolve multiaddr {maddr}: {e}")
return [maddr]

View File

@ -3,38 +3,24 @@ from __future__ import annotations
import socket
from multiaddr import Multiaddr
try:
from multiaddr.utils import ( # type: ignore
get_network_addrs,
get_thin_waist_addresses,
)
_HAS_THIN_WAIST = True
except ImportError: # pragma: no cover - only executed in older environments
_HAS_THIN_WAIST = False
get_thin_waist_addresses = None # type: ignore
get_network_addrs = None # type: ignore
from multiaddr.utils import get_network_addrs, get_thin_waist_addresses
def _safe_get_network_addrs(ip_version: int) -> list[str]:
"""
Internal safe wrapper. Returns a list of IP addresses for the requested IP version.
Falls back to minimal defaults when Thin Waist helpers are missing.
:param ip_version: 4 or 6
"""
if _HAS_THIN_WAIST and get_network_addrs:
try:
return get_network_addrs(ip_version) or []
except Exception: # pragma: no cover - defensive
return []
# Fallback behavior (very conservative)
if ip_version == 4:
return ["127.0.0.1"]
if ip_version == 6:
return ["::1"]
return []
try:
return get_network_addrs(ip_version) or []
except Exception: # pragma: no cover - defensive
# Fallback behavior (very conservative)
if ip_version == 4:
return ["127.0.0.1"]
if ip_version == 6:
return ["::1"]
return []
def find_free_port() -> int:
@ -47,16 +33,13 @@ def find_free_port() -> int:
def _safe_expand(addr: Multiaddr, port: int | None = None) -> list[Multiaddr]:
"""
Internal safe expansion wrapper. Returns a list of Multiaddr objects.
If Thin Waist isn't available, returns [addr] (identity).
"""
if _HAS_THIN_WAIST and get_thin_waist_addresses:
try:
if port is not None:
return get_thin_waist_addresses(addr, port=port) or []
return get_thin_waist_addresses(addr) or []
except Exception: # pragma: no cover - defensive
return [addr]
return [addr]
try:
if port is not None:
return get_thin_waist_addresses(addr, port=port) or []
return get_thin_waist_addresses(addr) or []
except Exception: # pragma: no cover - defensive
return [addr]
def get_available_interfaces(port: int, protocol: str = "tcp") -> list[Multiaddr]:
@ -73,8 +56,9 @@ def get_available_interfaces(port: int, protocol: str = "tcp") -> list[Multiaddr
seen_v4: set[str] = set()
for ip in _safe_get_network_addrs(4):
seen_v4.add(ip)
addrs.append(Multiaddr(f"/ip4/{ip}/{protocol}/{port}"))
if ip not in seen_v4: # Avoid duplicates
seen_v4.add(ip)
addrs.append(Multiaddr(f"/ip4/{ip}/{protocol}/{port}"))
# Ensure IPv4 loopback is always included when IPv4 interfaces are discovered
if seen_v4 and "127.0.0.1" not in seen_v4:
@ -89,8 +73,9 @@ def get_available_interfaces(port: int, protocol: str = "tcp") -> list[Multiaddr
#
# seen_v6: set[str] = set()
# for ip in _safe_get_network_addrs(6):
# seen_v6.add(ip)
# addrs.append(Multiaddr(f"/ip6/{ip}/{protocol}/{port}"))
# if ip not in seen_v6: # Avoid duplicates
# seen_v6.add(ip)
# addrs.append(Multiaddr(f"/ip6/{ip}/{protocol}/{port}"))
#
# # Always include IPv6 loopback for testing purposes when IPv6 is available
# # This ensures IPv6 functionality can be tested even without global IPv6 addresses
@ -99,7 +84,7 @@ def get_available_interfaces(port: int, protocol: str = "tcp") -> list[Multiaddr
# Fallback if nothing discovered
if not addrs:
addrs.append(Multiaddr(f"/ip4/0.0.0.0/{protocol}/{port}"))
addrs.append(Multiaddr(f"/ip4/127.0.0.1/{protocol}/{port}"))
return addrs
@ -120,6 +105,20 @@ def expand_wildcard_address(
return expanded
def get_wildcard_address(port: int, protocol: str = "tcp") -> Multiaddr:
"""
Get wildcard address (0.0.0.0) when explicitly needed.
This function provides access to wildcard binding as a feature when
explicitly required, preserving the ability to bind to all interfaces.
:param port: Port number.
:param protocol: Transport protocol.
:return: A Multiaddr with wildcard binding (0.0.0.0).
"""
return Multiaddr(f"/ip4/0.0.0.0/{protocol}/{port}")
def get_optimal_binding_address(port: int, protocol: str = "tcp") -> Multiaddr:
"""
Choose an optimal address for an example to bind to:
@ -148,13 +147,14 @@ def get_optimal_binding_address(port: int, protocol: str = "tcp") -> Multiaddr:
if "/ip4/127." in str(c) or "/ip6/::1" in str(c):
return c
# As a final fallback, produce a wildcard
return Multiaddr(f"/ip4/0.0.0.0/{protocol}/{port}")
# As a final fallback, produce a loopback address
return Multiaddr(f"/ip4/127.0.0.1/{protocol}/{port}")
__all__ = [
"get_available_interfaces",
"get_optimal_binding_address",
"get_wildcard_address",
"expand_wildcard_address",
"find_free_port",
]

View File

@ -0,0 +1,12 @@
Added experimental WebSocket transport support with basic WS and WSS functionality. This includes:
- WebSocket transport implementation with trio-websocket backend
- Support for both WS (WebSocket) and WSS (WebSocket Secure) protocols
- Basic connection management and stream handling
- TLS configuration support for WSS connections
- Multiaddr parsing for WebSocket addresses
- Integration with libp2p host and peer discovery
**Note**: This is experimental functionality. Advanced features like proxy support,
interop testing, and production examples are still in development. See
https://github.com/libp2p/py-libp2p/discussions/937 for the complete roadmap of missing features.

View File

@ -0,0 +1 @@
Add QUIC transport support for faster, more efficient peer-to-peer connections with native stream multiplexing.

View File

@ -0,0 +1 @@
Fixed message id type inconsistency in handle ihave and message id parsing improvement in handle iwant in pubsub module.

View File

@ -0,0 +1,2 @@
Updated all example scripts and core modules to use secure loopback addresses instead of wildcard addresses for network binding.
The `get_wildcard_address` function and related logic now utilize all available interfaces safely, improving security and consistency across the codebase.

View File

@ -0,0 +1 @@
Exposed timeout method in muxer multistream and updated all the usage. Added testcases to verify that timeout value is passed correctly

View File

@ -0,0 +1,6 @@
enhancement: Add write lock to `YamuxStream` to prevent concurrent write race conditions
- Implements ReadWriteLock for `YamuxStream` write operations
- Prevents data corruption from concurrent write operations
- Read operations remain lock-free due to existing `Yamux` architecture
- Resolves race conditions identified in Issue #793

View File

@ -0,0 +1,11 @@
Replace magic numbers with named constants and enums for clarity and maintainability
**Key Changes:**
- **Introduced type-safe enums** for better code clarity:
- `RelayRole(Flag)` enum with HOP, STOP, CLIENT roles supporting bitwise combinations (e.g., `RelayRole.HOP | RelayRole.STOP`)
- `ReservationStatus(Enum)` for reservation lifecycle management (ACTIVE, EXPIRED, REJECTED)
- **Replaced magic numbers with named constants** throughout the codebase, improving code maintainability and eliminating hardcoded timeout values (15s, 30s, 10s) with descriptive constant names
- **Added comprehensive timeout configuration system** with new `TimeoutConfig` dataclass supporting component-specific timeouts (discovery, protocol, DCUtR)
- **Enhanced configurability** of `RelayDiscovery`, `CircuitV2Protocol`, and `DCUtRProtocol` constructors with optional timeout parameters
- **Improved architecture consistency** with clean configuration flow across all circuit relay components
**Backward Compatibility:** All changes maintain full backward compatibility. Existing code continues to work unchanged while new timeout configuration options are available for users who need them.

View File

@ -0,0 +1 @@
Fix multiaddr dependency to use the last py-multiaddr commit hash to resolve installation issues

View File

@ -0,0 +1 @@
Updated multiaddr dependency from git repository to pip package version 0.0.11.

View File

@ -0,0 +1 @@
Fixed Windows CI/CD tests to use correct Python version instead of hardcoded Python 3.11. test 2

View File

@ -16,13 +16,14 @@ maintainers = [
{ name = "Dave Grantham", email = "dwg@linuxprogrammer.org" },
]
dependencies = [
"aioquic>=1.2.0",
"base58>=1.0.3",
"coincurve>=10.0.0",
"coincurve==21.0.0",
"exceptiongroup>=1.2.0; python_version < '3.11'",
"fastecdsa==2.3.2; sys_platform != 'win32'",
"grpcio>=1.41.0",
"lru-dict>=1.1.6",
# "multiaddr>=0.0.9",
"multiaddr @ git+https://github.com/multiformats/py-multiaddr.git@db8124e2321f316d3b7d2733c7df11d6ad9c03e6",
"multiaddr>=0.0.11",
"mypy-protobuf>=3.0.0",
"noiseprotocol>=0.3.0",
"protobuf>=4.25.0,<5.0.0",
@ -32,7 +33,7 @@ dependencies = [
"rpcudp>=3.0.0",
"trio-typing>=0.0.4",
"trio>=0.26.0",
"fastecdsa==2.3.2; sys_platform != 'win32'",
"trio-websocket>=0.11.0",
"zeroconf (>=0.147.0,<0.148.0)",
]
classifiers = [
@ -52,6 +53,7 @@ Homepage = "https://github.com/libp2p/py-libp2p"
[project.scripts]
chat-demo = "examples.chat.chat:main"
echo-demo = "examples.echo.echo:main"
echo-quic-demo="examples.echo.echo_quic:main"
ping-demo = "examples.ping.ping:main"
identify-demo = "examples.identify.identify:main"
identify-push-demo = "examples.identify_push.identify_push_demo:run_main"
@ -77,6 +79,7 @@ dev = [
"pytest>=7.0.0",
"pytest-xdist>=2.4.0",
"pytest-trio>=0.5.2",
"pytest-timeout>=2.4.0",
"factory-boy>=2.12.0,<3.0.0",
"ruff>=0.11.10",
"pyrefly (>=0.17.1,<0.18.0)",
@ -88,11 +91,12 @@ docs = [
"tomli; python_version < '3.11'",
]
test = [
"factory-boy>=2.12.0,<3.0.0",
"p2pclient==0.2.0",
"pytest>=7.0.0",
"pytest-xdist>=2.4.0",
"pytest-timeout>=2.4.0",
"pytest-trio>=0.5.2",
"factory-boy>=2.12.0,<3.0.0",
"pytest-xdist>=2.4.0",
]
[tool.setuptools]
@ -282,4 +286,5 @@ project_excludes = [
"**/*pb2.py",
"**/*.pyi",
".venv/**",
"./tests/interop/nim_libp2p",
]

View File

@ -1,6 +1,5 @@
import pytest
@pytest.fixture
def security_protocol():
return None
return None

View File

@ -250,10 +250,13 @@ def test_new_swarm_tcp_multiaddr_supported():
assert isinstance(swarm.transport, TCP)
def test_new_swarm_quic_multiaddr_raises():
def test_new_swarm_quic_multiaddr_supported():
from libp2p.transport.quic.transport import QUICTransport
addr = Multiaddr("/ip4/127.0.0.1/udp/9999/quic")
with pytest.raises(ValueError, match="QUIC not yet supported"):
new_swarm(listen_addrs=[addr])
swarm = new_swarm(listen_addrs=[addr])
assert isinstance(swarm, Swarm)
assert isinstance(swarm.transport, QUICTransport)
@pytest.mark.trio

View File

@ -1,4 +1,8 @@
import random
from unittest.mock import (
AsyncMock,
MagicMock,
)
import pytest
import trio
@ -7,6 +11,9 @@ from libp2p.pubsub.gossipsub import (
PROTOCOL_ID,
GossipSub,
)
from libp2p.pubsub.pb import (
rpc_pb2,
)
from libp2p.tools.utils import (
connect,
)
@ -754,3 +761,173 @@ async def test_single_host():
assert connected_peers == 0, (
f"Single host has {connected_peers} connections, expected 0"
)
@pytest.mark.trio
async def test_handle_ihave(monkeypatch):
async with PubsubFactory.create_batch_with_gossipsub(2) as pubsubs_gsub:
gossipsub_routers = []
for pubsub in pubsubs_gsub:
if isinstance(pubsub.router, GossipSub):
gossipsub_routers.append(pubsub.router)
gossipsubs = tuple(gossipsub_routers)
index_alice = 0
index_bob = 1
id_bob = pubsubs_gsub[index_bob].my_id
# Connect Alice and Bob
await connect(pubsubs_gsub[index_alice].host, pubsubs_gsub[index_bob].host)
await trio.sleep(0.1) # Allow connections to establish
# Mock emit_iwant to capture calls
mock_emit_iwant = AsyncMock()
monkeypatch.setattr(gossipsubs[index_alice], "emit_iwant", mock_emit_iwant)
# Create a test message ID as a string representation of a (seqno, from) tuple
test_seqno = b"1234"
test_from = id_bob.to_bytes()
test_msg_id = f"(b'{test_seqno.hex()}', b'{test_from.hex()}')"
ihave_msg = rpc_pb2.ControlIHave(messageIDs=[test_msg_id])
# Mock seen_messages.cache to avoid false positives
monkeypatch.setattr(pubsubs_gsub[index_alice].seen_messages, "cache", {})
# Simulate Bob sending IHAVE to Alice
await gossipsubs[index_alice].handle_ihave(ihave_msg, id_bob)
# Check if emit_iwant was called with the correct message ID
mock_emit_iwant.assert_called_once()
called_args = mock_emit_iwant.call_args[0]
assert called_args[0] == [test_msg_id] # Expected message IDs
assert called_args[1] == id_bob # Sender peer ID
@pytest.mark.trio
async def test_handle_iwant(monkeypatch):
async with PubsubFactory.create_batch_with_gossipsub(2) as pubsubs_gsub:
gossipsub_routers = []
for pubsub in pubsubs_gsub:
if isinstance(pubsub.router, GossipSub):
gossipsub_routers.append(pubsub.router)
gossipsubs = tuple(gossipsub_routers)
index_alice = 0
index_bob = 1
id_alice = pubsubs_gsub[index_alice].my_id
# Connect Alice and Bob
await connect(pubsubs_gsub[index_alice].host, pubsubs_gsub[index_bob].host)
await trio.sleep(0.1) # Allow connections to establish
# Mock mcache.get to return a message
test_message = rpc_pb2.Message(data=b"test_data")
test_seqno = b"1234"
test_from = id_alice.to_bytes()
# ✅ Correct: use raw tuple and str() to serialize, no hex()
test_msg_id = str((test_seqno, test_from))
mock_mcache_get = MagicMock(return_value=test_message)
monkeypatch.setattr(gossipsubs[index_bob].mcache, "get", mock_mcache_get)
# Mock write_msg to capture the sent packet
mock_write_msg = AsyncMock()
monkeypatch.setattr(gossipsubs[index_bob].pubsub, "write_msg", mock_write_msg)
# Simulate Alice sending IWANT to Bob
iwant_msg = rpc_pb2.ControlIWant(messageIDs=[test_msg_id])
await gossipsubs[index_bob].handle_iwant(iwant_msg, id_alice)
# Check if write_msg was called with the correct packet
mock_write_msg.assert_called_once()
packet = mock_write_msg.call_args[0][1]
assert isinstance(packet, rpc_pb2.RPC)
assert len(packet.publish) == 1
assert packet.publish[0] == test_message
# Verify that mcache.get was called with the correct parsed message ID
mock_mcache_get.assert_called_once()
called_msg_id = mock_mcache_get.call_args[0][0]
assert isinstance(called_msg_id, tuple)
assert called_msg_id == (test_seqno, test_from)
@pytest.mark.trio
async def test_handle_iwant_invalid_msg_id(monkeypatch):
"""
Test that handle_iwant raises ValueError for malformed message IDs.
"""
async with PubsubFactory.create_batch_with_gossipsub(2) as pubsubs_gsub:
gossipsub_routers = []
for pubsub in pubsubs_gsub:
if isinstance(pubsub.router, GossipSub):
gossipsub_routers.append(pubsub.router)
gossipsubs = tuple(gossipsub_routers)
index_alice = 0
index_bob = 1
id_alice = pubsubs_gsub[index_alice].my_id
await connect(pubsubs_gsub[index_alice].host, pubsubs_gsub[index_bob].host)
await trio.sleep(0.1)
# Malformed message ID (not a tuple string)
malformed_msg_id = "not_a_valid_msg_id"
iwant_msg = rpc_pb2.ControlIWant(messageIDs=[malformed_msg_id])
# Mock mcache.get and write_msg to ensure they are not called
mock_mcache_get = MagicMock()
monkeypatch.setattr(gossipsubs[index_bob].mcache, "get", mock_mcache_get)
mock_write_msg = AsyncMock()
monkeypatch.setattr(gossipsubs[index_bob].pubsub, "write_msg", mock_write_msg)
with pytest.raises(ValueError):
await gossipsubs[index_bob].handle_iwant(iwant_msg, id_alice)
mock_mcache_get.assert_not_called()
mock_write_msg.assert_not_called()
# Message ID that's a tuple string but not (bytes, bytes)
invalid_tuple_msg_id = "('abc', 123)"
iwant_msg = rpc_pb2.ControlIWant(messageIDs=[invalid_tuple_msg_id])
with pytest.raises(ValueError):
await gossipsubs[index_bob].handle_iwant(iwant_msg, id_alice)
mock_mcache_get.assert_not_called()
mock_write_msg.assert_not_called()
@pytest.mark.trio
async def test_handle_ihave_empty_message_ids(monkeypatch):
"""
Test that handle_ihave with an empty messageIDs list does not call emit_iwant.
"""
async with PubsubFactory.create_batch_with_gossipsub(2) as pubsubs_gsub:
gossipsub_routers = []
for pubsub in pubsubs_gsub:
if isinstance(pubsub.router, GossipSub):
gossipsub_routers.append(pubsub.router)
gossipsubs = tuple(gossipsub_routers)
index_alice = 0
index_bob = 1
id_bob = pubsubs_gsub[index_bob].my_id
# Connect Alice and Bob
await connect(pubsubs_gsub[index_alice].host, pubsubs_gsub[index_bob].host)
await trio.sleep(0.1) # Allow connections to establish
# Mock emit_iwant to capture calls
mock_emit_iwant = AsyncMock()
monkeypatch.setattr(gossipsubs[index_alice], "emit_iwant", mock_emit_iwant)
# Empty messageIDs list
ihave_msg = rpc_pb2.ControlIHave(messageIDs=[])
# Mock seen_messages.cache to avoid false positives
monkeypatch.setattr(pubsubs_gsub[index_alice].seen_messages, "cache", {})
# Simulate Bob sending IHAVE to Alice
await gossipsubs[index_alice].handle_ihave(ihave_msg, id_bob)
# emit_iwant should not be called since there are no message IDs
mock_emit_iwant.assert_not_called()

View File

@ -65,7 +65,7 @@ async def test_prune_backoff():
@pytest.mark.trio
async def test_unsubscribe_backoff():
async with PubsubFactory.create_batch_with_gossipsub(
2, heartbeat_interval=1, prune_back_off=1, unsubscribe_back_off=2
2, heartbeat_interval=0.5, prune_back_off=2, unsubscribe_back_off=4
) as pubsubs:
gsub0 = pubsubs[0].router
gsub1 = pubsubs[1].router
@ -107,7 +107,8 @@ async def test_unsubscribe_backoff():
)
# try to graft again (should succeed after backoff)
await trio.sleep(1)
# Wait longer than unsubscribe_back_off (4 seconds) + some buffer
await trio.sleep(4.5)
await gsub0.emit_graft(topic, host_1.get_id())
await trio.sleep(1)
assert host_0.get_id() in gsub1.mesh[topic], (

View File

@ -0,0 +1,108 @@
from unittest.mock import (
AsyncMock,
MagicMock,
)
import pytest
from libp2p.custom_types import (
TMuxerClass,
TProtocol,
)
from libp2p.peer.id import (
ID,
)
from libp2p.protocol_muxer.exceptions import (
MultiselectError,
)
from libp2p.stream_muxer.muxer_multistream import (
MuxerMultistream,
)
@pytest.mark.trio
async def test_muxer_timeout_configuration():
"""Test that muxer respects timeout configuration."""
muxer = MuxerMultistream({}, negotiate_timeout=1)
assert muxer.negotiate_timeout == 1
@pytest.mark.trio
async def test_select_transport_passes_timeout_to_multiselect():
"""Test that timeout is passed to multiselect client in select_transport."""
# Mock dependencies
mock_conn = MagicMock()
mock_conn.is_initiator = False
# Mock MultiselectClient
muxer = MuxerMultistream({}, negotiate_timeout=10)
muxer.multiselect.negotiate = AsyncMock(return_value=("mock_protocol", None))
muxer.transports[TProtocol("mock_protocol")] = MagicMock(return_value=MagicMock())
# Call select_transport
await muxer.select_transport(mock_conn)
# Verify that select_one_of was called with the correct timeout
args, _ = muxer.multiselect.negotiate.call_args
assert args[1] == 10
@pytest.mark.trio
async def test_new_conn_passes_timeout_to_multistream_client():
"""Test that timeout is passed to multistream client in new_conn."""
# Mock dependencies
mock_conn = MagicMock()
mock_conn.is_initiator = True
mock_peer_id = ID(b"test_peer")
mock_communicator = MagicMock()
# Mock MultistreamClient and transports
muxer = MuxerMultistream({}, negotiate_timeout=30)
muxer.multistream_client.select_one_of = AsyncMock(return_value="mock_protocol")
muxer.transports[TProtocol("mock_protocol")] = MagicMock(return_value=MagicMock())
# Call new_conn
await muxer.new_conn(mock_conn, mock_peer_id)
# Verify that select_one_of was called with the correct timeout
muxer.multistream_client.select_one_of(
tuple(muxer.transports.keys()), mock_communicator, 30
)
@pytest.mark.trio
async def test_select_transport_no_protocol_selected():
"""
Test that select_transport raises MultiselectError when no protocol is selected.
"""
# Mock dependencies
mock_conn = MagicMock()
mock_conn.is_initiator = False
# Mock Multiselect to return None
muxer = MuxerMultistream({}, negotiate_timeout=30)
muxer.multiselect.negotiate = AsyncMock(return_value=(None, None))
# Expect MultiselectError to be raised
with pytest.raises(MultiselectError, match="no protocol selected"):
await muxer.select_transport(mock_conn)
@pytest.mark.trio
async def test_add_transport_updates_precedence():
"""Test that adding a transport updates protocol precedence."""
# Mock transport classes
mock_transport1 = MagicMock(spec=TMuxerClass)
mock_transport2 = MagicMock(spec=TMuxerClass)
# Initialize muxer and add transports
muxer = MuxerMultistream({}, negotiate_timeout=30)
muxer.add_transport(TProtocol("proto1"), mock_transport1)
muxer.add_transport(TProtocol("proto2"), mock_transport2)
# Verify transport order
assert list(muxer.transports.keys()) == ["proto1", "proto2"]
# Re-add proto1 to check if it moves to the end
muxer.add_transport(TProtocol("proto1"), mock_transport1)
assert list(muxer.transports.keys()) == ["proto2", "proto1"]

View File

@ -0,0 +1,553 @@
"""
Enhanced tests for QUIC connection functionality - Module 3.
Tests all new features including advanced stream management, resource management,
error handling, and concurrent operations.
"""
from unittest.mock import AsyncMock, Mock, patch
import pytest
from multiaddr.multiaddr import Multiaddr
import trio
from libp2p.crypto.ed25519 import create_new_key_pair
from libp2p.peer.id import ID
from libp2p.transport.quic.config import QUICTransportConfig
from libp2p.transport.quic.connection import QUICConnection
from libp2p.transport.quic.exceptions import (
QUICConnectionClosedError,
QUICConnectionError,
QUICConnectionTimeoutError,
QUICPeerVerificationError,
QUICStreamLimitError,
QUICStreamTimeoutError,
)
from libp2p.transport.quic.security import QUICTLSConfigManager
from libp2p.transport.quic.stream import QUICStream, StreamDirection
class MockResourceScope:
"""Mock resource scope for testing."""
def __init__(self):
self.memory_reserved = 0
def reserve_memory(self, size):
self.memory_reserved += size
def release_memory(self, size):
self.memory_reserved = max(0, self.memory_reserved - size)
class TestQUICConnection:
"""Test suite for QUIC connection functionality."""
@pytest.fixture
def mock_quic_connection(self):
"""Create mock aioquic QuicConnection."""
mock = Mock()
mock.next_event.return_value = None
mock.datagrams_to_send.return_value = []
mock.get_timer.return_value = None
mock.connect = Mock()
mock.close = Mock()
mock.send_stream_data = Mock()
mock.reset_stream = Mock()
return mock
@pytest.fixture
def mock_quic_transport(self):
mock = Mock()
mock._config = QUICTransportConfig()
return mock
@pytest.fixture
def mock_resource_scope(self):
"""Create mock resource scope."""
return MockResourceScope()
@pytest.fixture
def quic_connection(
self,
mock_quic_connection: Mock,
mock_quic_transport: Mock,
mock_resource_scope: MockResourceScope,
):
"""Create test QUIC connection with enhanced features."""
private_key = create_new_key_pair().private_key
peer_id = ID.from_pubkey(private_key.get_public_key())
mock_security_manager = Mock()
return QUICConnection(
quic_connection=mock_quic_connection,
remote_addr=("127.0.0.1", 4001),
remote_peer_id=None,
local_peer_id=peer_id,
is_initiator=True,
maddr=Multiaddr("/ip4/127.0.0.1/udp/4001/quic"),
transport=mock_quic_transport,
resource_scope=mock_resource_scope,
security_manager=mock_security_manager,
)
@pytest.fixture
def server_connection(self, mock_quic_connection, mock_resource_scope):
"""Create server-side QUIC connection."""
private_key = create_new_key_pair().private_key
peer_id = ID.from_pubkey(private_key.get_public_key())
return QUICConnection(
quic_connection=mock_quic_connection,
remote_addr=("127.0.0.1", 4001),
remote_peer_id=peer_id,
local_peer_id=peer_id,
is_initiator=False,
maddr=Multiaddr("/ip4/127.0.0.1/udp/4001/quic"),
transport=Mock(),
resource_scope=mock_resource_scope,
)
# Basic functionality tests
def test_connection_initialization_enhanced(
self, quic_connection, mock_resource_scope
):
"""Test enhanced connection initialization."""
assert quic_connection._remote_addr == ("127.0.0.1", 4001)
assert quic_connection.is_initiator is True
assert not quic_connection.is_closed
assert not quic_connection.is_established
assert len(quic_connection._streams) == 0
assert quic_connection._resource_scope == mock_resource_scope
assert quic_connection._outbound_stream_count == 0
assert quic_connection._inbound_stream_count == 0
assert len(quic_connection._stream_accept_queue) == 0
def test_stream_id_calculation_enhanced(self):
"""Test enhanced stream ID calculation for client/server."""
# Client connection (initiator)
client_conn = QUICConnection(
quic_connection=Mock(),
remote_addr=("127.0.0.1", 4001),
remote_peer_id=None,
local_peer_id=Mock(),
is_initiator=True,
maddr=Multiaddr("/ip4/127.0.0.1/udp/4001/quic"),
transport=Mock(),
)
assert client_conn._next_stream_id == 0 # Client starts with 0
# Server connection (not initiator)
server_conn = QUICConnection(
quic_connection=Mock(),
remote_addr=("127.0.0.1", 4001),
remote_peer_id=None,
local_peer_id=Mock(),
is_initiator=False,
maddr=Multiaddr("/ip4/127.0.0.1/udp/4001/quic"),
transport=Mock(),
)
assert server_conn._next_stream_id == 1 # Server starts with 1
def test_incoming_stream_detection_enhanced(self, quic_connection):
"""Test enhanced incoming stream detection logic."""
# For client (initiator), odd stream IDs are incoming
assert quic_connection._is_incoming_stream(1) is True # Server-initiated
assert quic_connection._is_incoming_stream(0) is False # Client-initiated
assert quic_connection._is_incoming_stream(5) is True # Server-initiated
assert quic_connection._is_incoming_stream(4) is False # Client-initiated
# Stream management tests
@pytest.mark.trio
async def test_open_stream_basic(self, quic_connection):
"""Test basic stream opening."""
quic_connection._started = True
stream = await quic_connection.open_stream()
assert isinstance(stream, QUICStream)
assert stream.stream_id == "0"
assert stream.direction == StreamDirection.OUTBOUND
assert 0 in quic_connection._streams
assert quic_connection._outbound_stream_count == 1
@pytest.mark.trio
async def test_open_stream_limit_reached(self, quic_connection):
"""Test stream limit enforcement."""
quic_connection._started = True
quic_connection._outbound_stream_count = quic_connection.MAX_OUTGOING_STREAMS
with pytest.raises(QUICStreamLimitError, match="Maximum outbound streams"):
await quic_connection.open_stream()
@pytest.mark.trio
async def test_open_stream_timeout(self, quic_connection: QUICConnection):
"""Test stream opening timeout."""
quic_connection._started = True
return
# Mock the stream ID lock to simulate slow operation
async def slow_acquire():
await trio.sleep(10) # Longer than timeout
with patch.object(
quic_connection._stream_lock, "acquire", side_effect=slow_acquire
):
with pytest.raises(
QUICStreamTimeoutError, match="Stream creation timed out"
):
await quic_connection.open_stream(timeout=0.1)
@pytest.mark.trio
async def test_accept_stream_basic(self, quic_connection):
"""Test basic stream acceptance."""
# Create a mock inbound stream
mock_stream = Mock(spec=QUICStream)
mock_stream.stream_id = "1"
# Add to accept queue
quic_connection._stream_accept_queue.append(mock_stream)
quic_connection._stream_accept_event.set()
accepted_stream = await quic_connection.accept_stream(timeout=0.1)
assert accepted_stream == mock_stream
assert len(quic_connection._stream_accept_queue) == 0
@pytest.mark.trio
async def test_accept_stream_timeout(self, quic_connection):
"""Test stream acceptance timeout."""
with pytest.raises(QUICStreamTimeoutError, match="Stream accept timed out"):
await quic_connection.accept_stream(timeout=0.1)
@pytest.mark.trio
async def test_accept_stream_on_closed_connection(self, quic_connection):
"""Test stream acceptance on closed connection."""
await quic_connection.close()
with pytest.raises(QUICConnectionClosedError, match="Connection is closed"):
await quic_connection.accept_stream()
# Stream handler tests
@pytest.mark.trio
async def test_stream_handler_setting(self, quic_connection):
"""Test setting stream handler."""
async def mock_handler(stream):
pass
quic_connection.set_stream_handler(mock_handler)
assert quic_connection._stream_handler == mock_handler
# Connection lifecycle tests
@pytest.mark.trio
async def test_connection_start_client(self, quic_connection):
"""Test client connection start."""
with patch.object(
quic_connection, "_initiate_connection", new_callable=AsyncMock
) as mock_initiate:
await quic_connection.start()
assert quic_connection._started
mock_initiate.assert_called_once()
@pytest.mark.trio
async def test_connection_start_server(self, server_connection):
"""Test server connection start."""
await server_connection.start()
assert server_connection._started
assert server_connection._established
assert server_connection._connected_event.is_set()
@pytest.mark.trio
async def test_connection_start_already_started(self, quic_connection):
"""Test starting already started connection."""
quic_connection._started = True
# Should not raise error, just log warning
await quic_connection.start()
assert quic_connection._started
@pytest.mark.trio
async def test_connection_start_closed(self, quic_connection):
"""Test starting closed connection."""
quic_connection._closed = True
with pytest.raises(
QUICConnectionError, match="Cannot start a closed connection"
):
await quic_connection.start()
@pytest.mark.trio
async def test_connection_connect_with_nursery(
self, quic_connection: QUICConnection
):
"""Test connection establishment with nursery."""
quic_connection._started = True
quic_connection._established = True
quic_connection._connected_event.set()
with patch.object(
quic_connection, "_start_background_tasks", new_callable=AsyncMock
) as mock_start_tasks:
with patch.object(
quic_connection,
"_verify_peer_identity_with_security",
new_callable=AsyncMock,
) as mock_verify:
async with trio.open_nursery() as nursery:
await quic_connection.connect(nursery)
assert quic_connection._nursery == nursery
mock_start_tasks.assert_called_once()
mock_verify.assert_called_once()
@pytest.mark.trio
@pytest.mark.slow
async def test_connection_connect_timeout(
self, quic_connection: QUICConnection
) -> None:
"""Test connection establishment timeout."""
quic_connection._started = True
# Don't set connected event to simulate timeout
with patch.object(
quic_connection, "_start_background_tasks", new_callable=AsyncMock
):
async with trio.open_nursery() as nursery:
with pytest.raises(
QUICConnectionTimeoutError, match="Connection handshake timed out"
):
await quic_connection.connect(nursery)
# Resource management tests
@pytest.mark.trio
async def test_stream_removal_resource_cleanup(
self, quic_connection: QUICConnection, mock_resource_scope
):
"""Test stream removal and resource cleanup."""
quic_connection._started = True
# Create a stream
stream = await quic_connection.open_stream()
# Remove the stream
quic_connection._remove_stream(int(stream.stream_id))
assert int(stream.stream_id) not in quic_connection._streams
# Note: Count updates is async, so we can't test it directly here
# Error handling tests
@pytest.mark.trio
async def test_connection_error_handling(self, quic_connection) -> None:
"""Test connection error handling."""
error = Exception("Test error")
with patch.object(
quic_connection, "close", new_callable=AsyncMock
) as mock_close:
await quic_connection._handle_connection_error(error)
mock_close.assert_called_once()
# Statistics and monitoring tests
@pytest.mark.trio
async def test_connection_stats_enhanced(self, quic_connection) -> None:
"""Test enhanced connection statistics."""
quic_connection._started = True
# Create some streams
_stream1 = await quic_connection.open_stream()
_stream2 = await quic_connection.open_stream()
stats = quic_connection.get_stream_stats()
expected_keys = [
"total_streams",
"outbound_streams",
"inbound_streams",
"max_streams",
"stream_utilization",
"stats",
]
for key in expected_keys:
assert key in stats
assert stats["total_streams"] == 2
assert stats["outbound_streams"] == 2
assert stats["inbound_streams"] == 0
@pytest.mark.trio
async def test_get_active_streams(self, quic_connection) -> None:
"""Test getting active streams."""
quic_connection._started = True
# Create streams
stream1 = await quic_connection.open_stream()
stream2 = await quic_connection.open_stream()
active_streams = quic_connection.get_active_streams()
assert len(active_streams) == 2
assert stream1 in active_streams
assert stream2 in active_streams
@pytest.mark.trio
async def test_get_streams_by_protocol(self, quic_connection) -> None:
"""Test getting streams by protocol."""
quic_connection._started = True
# Create streams with different protocols
stream1 = await quic_connection.open_stream()
stream1.protocol = "/test/1.0.0"
stream2 = await quic_connection.open_stream()
stream2.protocol = "/other/1.0.0"
test_streams = quic_connection.get_streams_by_protocol("/test/1.0.0")
other_streams = quic_connection.get_streams_by_protocol("/other/1.0.0")
assert len(test_streams) == 1
assert len(other_streams) == 1
assert stream1 in test_streams
assert stream2 in other_streams
# Enhanced close tests
@pytest.mark.trio
async def test_connection_close_enhanced(
self, quic_connection: QUICConnection
) -> None:
"""Test enhanced connection close with stream cleanup."""
quic_connection._started = True
# Create some streams
_stream1 = await quic_connection.open_stream()
_stream2 = await quic_connection.open_stream()
await quic_connection.close()
assert quic_connection.is_closed
assert len(quic_connection._streams) == 0
# Concurrent operations tests
@pytest.mark.trio
async def test_concurrent_stream_operations(
self, quic_connection: QUICConnection
) -> None:
"""Test concurrent stream operations."""
quic_connection._started = True
async def create_stream():
return await quic_connection.open_stream()
# Create multiple streams concurrently
async with trio.open_nursery() as nursery:
for i in range(10):
nursery.start_soon(create_stream)
# Wait a bit for all to start
await trio.sleep(0.1)
# Should have created streams without conflicts
assert quic_connection._outbound_stream_count == 10
assert len(quic_connection._streams) == 10
# Connection properties tests
def test_connection_properties(self, quic_connection: QUICConnection) -> None:
"""Test connection property accessors."""
assert quic_connection.multiaddr() == quic_connection._maddr
assert quic_connection.local_peer_id() == quic_connection._local_peer_id
assert quic_connection.remote_peer_id() == quic_connection._remote_peer_id
# IRawConnection interface tests
@pytest.mark.trio
async def test_raw_connection_write(self, quic_connection: QUICConnection) -> None:
"""Test raw connection write interface."""
quic_connection._started = True
with patch.object(quic_connection, "open_stream") as mock_open:
mock_stream = AsyncMock()
mock_open.return_value = mock_stream
await quic_connection.write(b"test data")
mock_open.assert_called_once()
mock_stream.write.assert_called_once_with(b"test data")
mock_stream.close_write.assert_called_once()
@pytest.mark.trio
async def test_raw_connection_read_not_implemented(
self, quic_connection: QUICConnection
) -> None:
"""Test raw connection read raises NotImplementedError."""
with pytest.raises(NotImplementedError):
await quic_connection.read()
# Mock verification helpers
def test_mock_resource_scope_functionality(self, mock_resource_scope) -> None:
"""Test mock resource scope works correctly."""
assert mock_resource_scope.memory_reserved == 0
mock_resource_scope.reserve_memory(1000)
assert mock_resource_scope.memory_reserved == 1000
mock_resource_scope.reserve_memory(500)
assert mock_resource_scope.memory_reserved == 1500
mock_resource_scope.release_memory(600)
assert mock_resource_scope.memory_reserved == 900
mock_resource_scope.release_memory(2000) # Should not go negative
assert mock_resource_scope.memory_reserved == 0
@pytest.mark.trio
async def test_invalid_certificate_verification():
key_pair1 = create_new_key_pair()
key_pair2 = create_new_key_pair()
peer_id1 = ID.from_pubkey(key_pair1.public_key)
peer_id2 = ID.from_pubkey(key_pair2.public_key)
manager = QUICTLSConfigManager(
libp2p_private_key=key_pair1.private_key, peer_id=peer_id1
)
# Match the certificate against a different peer_id
with pytest.raises(QUICPeerVerificationError, match="Peer ID mismatch"):
manager.verify_peer_identity(manager.tls_config.certificate, peer_id2)
from cryptography.hazmat.primitives.serialization import Encoding
# --- Corrupt the certificate by tampering the DER bytes ---
cert_bytes = manager.tls_config.certificate.public_bytes(Encoding.DER)
corrupted_bytes = bytearray(cert_bytes)
# Flip some random bytes in the middle of the certificate
corrupted_bytes[len(corrupted_bytes) // 2] ^= 0xFF
from cryptography import x509
from cryptography.hazmat.backends import default_backend
# This will still parse (structurally valid), but the signature
# or fingerprint will break
corrupted_cert = x509.load_der_x509_certificate(
bytes(corrupted_bytes), backend=default_backend()
)
with pytest.raises(
QUICPeerVerificationError, match="Certificate verification failed"
):
manager.verify_peer_identity(corrupted_cert, peer_id1)

View File

@ -0,0 +1,624 @@
"""
QUIC Connection ID Management Tests
This test module covers comprehensive testing of QUIC connection ID functionality
including generation, rotation, retirement, and validation according to RFC 9000.
Tests are organized into:
1. Basic Connection ID Management
2. Connection ID Rotation and Updates
3. Connection ID Retirement
4. Error Conditions and Edge Cases
5. Integration Tests with Real Connections
"""
import secrets
import time
from typing import Any
from unittest.mock import Mock
import pytest
from aioquic.buffer import Buffer
# Import aioquic components for low-level testing
from aioquic.quic.configuration import QuicConfiguration
from aioquic.quic.connection import QuicConnection, QuicConnectionId
from multiaddr import Multiaddr
from libp2p.crypto.ed25519 import create_new_key_pair
from libp2p.peer.id import ID
from libp2p.transport.quic.config import QUICTransportConfig
from libp2p.transport.quic.connection import QUICConnection
from libp2p.transport.quic.transport import QUICTransport
class ConnectionIdTestHelper:
"""Helper class for connection ID testing utilities."""
@staticmethod
def generate_connection_id(length: int = 8) -> bytes:
"""Generate a random connection ID of specified length."""
return secrets.token_bytes(length)
@staticmethod
def create_quic_connection_id(cid: bytes, sequence: int = 0) -> QuicConnectionId:
"""Create a QuicConnectionId object."""
return QuicConnectionId(
cid=cid,
sequence_number=sequence,
stateless_reset_token=secrets.token_bytes(16),
)
@staticmethod
def extract_connection_ids_from_connection(conn: QUICConnection) -> dict[str, Any]:
"""Extract connection ID information from a QUIC connection."""
quic = conn._quic
return {
"host_cids": [cid.cid.hex() for cid in getattr(quic, "_host_cids", [])],
"peer_cid": getattr(quic, "_peer_cid", None),
"peer_cid_available": [
cid.cid.hex() for cid in getattr(quic, "_peer_cid_available", [])
],
"retire_connection_ids": getattr(quic, "_retire_connection_ids", []),
"host_cid_seq": getattr(quic, "_host_cid_seq", 0),
}
class TestBasicConnectionIdManagement:
"""Test basic connection ID management functionality."""
@pytest.fixture
def mock_quic_connection(self):
"""Create a mock QUIC connection with connection ID support."""
mock_quic = Mock(spec=QuicConnection)
mock_quic._host_cids = []
mock_quic._host_cid_seq = 0
mock_quic._peer_cid = None
mock_quic._peer_cid_available = []
mock_quic._retire_connection_ids = []
mock_quic._configuration = Mock()
mock_quic._configuration.connection_id_length = 8
mock_quic._remote_active_connection_id_limit = 8
return mock_quic
@pytest.fixture
def quic_connection(self, mock_quic_connection):
"""Create a QUICConnection instance for testing."""
private_key = create_new_key_pair().private_key
peer_id = ID.from_pubkey(private_key.get_public_key())
return QUICConnection(
quic_connection=mock_quic_connection,
remote_addr=("127.0.0.1", 4001),
remote_peer_id=peer_id,
local_peer_id=peer_id,
is_initiator=True,
maddr=Multiaddr("/ip4/127.0.0.1/udp/4001/quic"),
transport=Mock(),
)
def test_connection_id_initialization(self, quic_connection):
"""Test that connection ID tracking is properly initialized."""
# Check that connection ID tracking structures are initialized
assert hasattr(quic_connection, "_available_connection_ids")
assert hasattr(quic_connection, "_current_connection_id")
assert hasattr(quic_connection, "_retired_connection_ids")
assert hasattr(quic_connection, "_connection_id_sequence_numbers")
# Initial state should be empty
assert len(quic_connection._available_connection_ids) == 0
assert quic_connection._current_connection_id is None
assert len(quic_connection._retired_connection_ids) == 0
assert len(quic_connection._connection_id_sequence_numbers) == 0
def test_connection_id_stats_tracking(self, quic_connection):
"""Test connection ID statistics are properly tracked."""
stats = quic_connection.get_connection_id_stats()
# Check that all expected stats are present
expected_keys = [
"available_connection_ids",
"current_connection_id",
"retired_connection_ids",
"connection_ids_issued",
"connection_ids_retired",
"connection_id_changes",
"available_cid_list",
]
for key in expected_keys:
assert key in stats
# Initial values should be zero/empty
assert stats["available_connection_ids"] == 0
assert stats["current_connection_id"] is None
assert stats["retired_connection_ids"] == 0
assert stats["connection_ids_issued"] == 0
assert stats["connection_ids_retired"] == 0
assert stats["connection_id_changes"] == 0
assert stats["available_cid_list"] == []
def test_current_connection_id_getter(self, quic_connection):
"""Test getting current connection ID."""
# Initially no connection ID
assert quic_connection.get_current_connection_id() is None
# Set a connection ID
test_cid = ConnectionIdTestHelper.generate_connection_id()
quic_connection._current_connection_id = test_cid
assert quic_connection.get_current_connection_id() == test_cid
def test_connection_id_generation(self):
"""Test connection ID generation utilities."""
# Test default length
cid1 = ConnectionIdTestHelper.generate_connection_id()
assert len(cid1) == 8
assert isinstance(cid1, bytes)
# Test custom length
cid2 = ConnectionIdTestHelper.generate_connection_id(16)
assert len(cid2) == 16
# Test uniqueness
cid3 = ConnectionIdTestHelper.generate_connection_id()
assert cid1 != cid3
class TestConnectionIdRotationAndUpdates:
"""Test connection ID rotation and update mechanisms."""
@pytest.fixture
def transport_config(self):
"""Create transport configuration."""
return QUICTransportConfig(
idle_timeout=10.0,
connection_timeout=5.0,
max_concurrent_streams=100,
)
@pytest.fixture
def server_key(self):
"""Generate server private key."""
return create_new_key_pair().private_key
@pytest.fixture
def client_key(self):
"""Generate client private key."""
return create_new_key_pair().private_key
def test_connection_id_replenishment(self):
"""Test connection ID replenishment mechanism."""
# Create a real QuicConnection to test replenishment
config = QuicConfiguration(is_client=True)
config.connection_id_length = 8
quic_conn = QuicConnection(configuration=config)
# Initial state - should have some host connection IDs
initial_count = len(quic_conn._host_cids)
assert initial_count > 0
# Remove some connection IDs to trigger replenishment
while len(quic_conn._host_cids) > 2:
quic_conn._host_cids.pop()
# Trigger replenishment
quic_conn._replenish_connection_ids()
# Should have replenished up to the limit
assert len(quic_conn._host_cids) >= initial_count
# All connection IDs should have unique sequence numbers
sequences = [cid.sequence_number for cid in quic_conn._host_cids]
assert len(sequences) == len(set(sequences))
def test_connection_id_sequence_numbers(self):
"""Test connection ID sequence number management."""
config = QuicConfiguration(is_client=True)
quic_conn = QuicConnection(configuration=config)
# Get initial sequence number
initial_seq = quic_conn._host_cid_seq
# Trigger replenishment to generate new connection IDs
quic_conn._replenish_connection_ids()
# Sequence numbers should increment
assert quic_conn._host_cid_seq > initial_seq
# All host connection IDs should have sequential numbers
sequences = [cid.sequence_number for cid in quic_conn._host_cids]
sequences.sort()
# Check for proper sequence
for i in range(len(sequences) - 1):
assert sequences[i + 1] > sequences[i]
def test_connection_id_limits(self):
"""Test connection ID limit enforcement."""
config = QuicConfiguration(is_client=True)
config.connection_id_length = 8
quic_conn = QuicConnection(configuration=config)
# Set a reasonable limit
quic_conn._remote_active_connection_id_limit = 4
# Replenish connection IDs
quic_conn._replenish_connection_ids()
# Should not exceed the limit
assert len(quic_conn._host_cids) <= quic_conn._remote_active_connection_id_limit
class TestConnectionIdRetirement:
"""Test connection ID retirement functionality."""
def test_connection_id_retirement_basic(self):
"""Test basic connection ID retirement."""
config = QuicConfiguration(is_client=True)
quic_conn = QuicConnection(configuration=config)
# Create a test connection ID to retire
test_cid = ConnectionIdTestHelper.create_quic_connection_id(
ConnectionIdTestHelper.generate_connection_id(), sequence=1
)
# Add it to peer connection IDs
quic_conn._peer_cid_available.append(test_cid)
quic_conn._peer_cid_sequence_numbers.add(1)
# Retire the connection ID
quic_conn._retire_peer_cid(test_cid)
# Should be added to retirement list
assert 1 in quic_conn._retire_connection_ids
def test_connection_id_retirement_limits(self):
"""Test connection ID retirement limits."""
config = QuicConfiguration(is_client=True)
quic_conn = QuicConnection(configuration=config)
# Fill up retirement list near the limit
max_retirements = 32 # Based on aioquic's default limit
for i in range(max_retirements):
quic_conn._retire_connection_ids.append(i)
# Should be at limit
assert len(quic_conn._retire_connection_ids) == max_retirements
def test_connection_id_retirement_events(self):
"""Test that retirement generates proper events."""
config = QuicConfiguration(is_client=True)
quic_conn = QuicConnection(configuration=config)
# Create and add a host connection ID
test_cid = ConnectionIdTestHelper.create_quic_connection_id(
ConnectionIdTestHelper.generate_connection_id(), sequence=5
)
quic_conn._host_cids.append(test_cid)
# Create a retirement frame buffer
from aioquic.buffer import Buffer
buf = Buffer(capacity=16)
buf.push_uint_var(5) # sequence number to retire
buf.seek(0)
# Process retirement (this should generate an event)
try:
quic_conn._handle_retire_connection_id_frame(
Mock(), # context
0x19, # RETIRE_CONNECTION_ID frame type
buf,
)
# Check that connection ID was removed
remaining_sequences = [cid.sequence_number for cid in quic_conn._host_cids]
assert 5 not in remaining_sequences
except Exception:
# May fail due to missing context, but that's okay for this test
pass
class TestConnectionIdErrorConditions:
"""Test error conditions and edge cases in connection ID handling."""
def test_invalid_connection_id_length(self):
"""Test handling of invalid connection ID lengths."""
# Connection IDs must be 1-20 bytes according to RFC 9000
# Test too short (0 bytes) - this should be handled gracefully
empty_cid = b""
assert len(empty_cid) == 0
# Test too long (>20 bytes)
long_cid = secrets.token_bytes(21)
assert len(long_cid) == 21
# Test valid lengths
for length in range(1, 21):
valid_cid = secrets.token_bytes(length)
assert len(valid_cid) == length
def test_duplicate_sequence_numbers(self):
"""Test handling of duplicate sequence numbers."""
config = QuicConfiguration(is_client=True)
quic_conn = QuicConnection(configuration=config)
# Create two connection IDs with same sequence number
cid1 = ConnectionIdTestHelper.create_quic_connection_id(
ConnectionIdTestHelper.generate_connection_id(), sequence=10
)
cid2 = ConnectionIdTestHelper.create_quic_connection_id(
ConnectionIdTestHelper.generate_connection_id(), sequence=10
)
# Add first connection ID
quic_conn._peer_cid_available.append(cid1)
quic_conn._peer_cid_sequence_numbers.add(10)
# Adding second with same sequence should be handled appropriately
# (The implementation should prevent duplicates)
if 10 not in quic_conn._peer_cid_sequence_numbers:
quic_conn._peer_cid_available.append(cid2)
quic_conn._peer_cid_sequence_numbers.add(10)
# Should only have one entry for sequence 10
sequences = [cid.sequence_number for cid in quic_conn._peer_cid_available]
assert sequences.count(10) <= 1
def test_retire_unknown_connection_id(self):
"""Test retiring an unknown connection ID."""
config = QuicConfiguration(is_client=True)
quic_conn = QuicConnection(configuration=config)
# Try to create a buffer to retire unknown sequence number
buf = Buffer(capacity=16)
buf.push_uint_var(999) # Unknown sequence number
buf.seek(0)
# This should raise an error when processed
# (Testing the error condition, not the full processing)
unknown_sequence = 999
known_sequences = [cid.sequence_number for cid in quic_conn._host_cids]
assert unknown_sequence not in known_sequences
def test_retire_current_connection_id(self):
"""Test that retiring current connection ID is prevented."""
config = QuicConfiguration(is_client=True)
quic_conn = QuicConnection(configuration=config)
# Get current connection ID if available
if quic_conn._host_cids:
current_cid = quic_conn._host_cids[0]
current_sequence = current_cid.sequence_number
# Trying to retire current connection ID should be prevented
# This is tested by checking the sequence number logic
assert current_sequence >= 0
class TestConnectionIdIntegration:
"""Integration tests for connection ID functionality with real connections."""
@pytest.fixture
def server_config(self):
"""Server transport configuration."""
return QUICTransportConfig(
idle_timeout=10.0,
connection_timeout=5.0,
max_concurrent_streams=100,
)
@pytest.fixture
def client_config(self):
"""Client transport configuration."""
return QUICTransportConfig(
idle_timeout=10.0,
connection_timeout=5.0,
)
@pytest.fixture
def server_key(self):
"""Generate server private key."""
return create_new_key_pair().private_key
@pytest.fixture
def client_key(self):
"""Generate client private key."""
return create_new_key_pair().private_key
@pytest.mark.trio
async def test_connection_id_exchange_during_handshake(
self, server_key, client_key, server_config, client_config
):
"""Test connection ID exchange during connection handshake."""
# This test would require a full connection setup
# For now, we test the setup components
server_transport = QUICTransport(server_key, server_config)
client_transport = QUICTransport(client_key, client_config)
# Verify transports are created with proper configuration
assert server_transport._config == server_config
assert client_transport._config == client_config
# Test that connection ID tracking is available
# (Integration with actual networking would require more setup)
def test_connection_id_extraction_utilities(self):
"""Test connection ID extraction utilities."""
# Create a mock connection with some connection IDs
private_key = create_new_key_pair().private_key
peer_id = ID.from_pubkey(private_key.get_public_key())
mock_quic = Mock()
mock_quic._host_cids = [
ConnectionIdTestHelper.create_quic_connection_id(
ConnectionIdTestHelper.generate_connection_id(), i
)
for i in range(3)
]
mock_quic._peer_cid = None
mock_quic._peer_cid_available = []
mock_quic._retire_connection_ids = []
mock_quic._host_cid_seq = 3
quic_conn = QUICConnection(
quic_connection=mock_quic,
remote_addr=("127.0.0.1", 4001),
remote_peer_id=peer_id,
local_peer_id=peer_id,
is_initiator=True,
maddr=Multiaddr("/ip4/127.0.0.1/udp/4001/quic"),
transport=Mock(),
)
# Extract connection ID information
cid_info = ConnectionIdTestHelper.extract_connection_ids_from_connection(
quic_conn
)
# Verify extraction works
assert "host_cids" in cid_info
assert "peer_cid" in cid_info
assert "peer_cid_available" in cid_info
assert "retire_connection_ids" in cid_info
assert "host_cid_seq" in cid_info
# Check values
assert len(cid_info["host_cids"]) == 3
assert cid_info["host_cid_seq"] == 3
assert cid_info["peer_cid"] is None
assert len(cid_info["peer_cid_available"]) == 0
assert len(cid_info["retire_connection_ids"]) == 0
class TestConnectionIdStatistics:
"""Test connection ID statistics and monitoring."""
@pytest.fixture
def connection_with_stats(self):
"""Create a connection with connection ID statistics."""
private_key = create_new_key_pair().private_key
peer_id = ID.from_pubkey(private_key.get_public_key())
mock_quic = Mock()
mock_quic._host_cids = []
mock_quic._peer_cid = None
mock_quic._peer_cid_available = []
mock_quic._retire_connection_ids = []
return QUICConnection(
quic_connection=mock_quic,
remote_addr=("127.0.0.1", 4001),
remote_peer_id=peer_id,
local_peer_id=peer_id,
is_initiator=True,
maddr=Multiaddr("/ip4/127.0.0.1/udp/4001/quic"),
transport=Mock(),
)
def test_connection_id_stats_initialization(self, connection_with_stats):
"""Test that connection ID statistics are properly initialized."""
stats = connection_with_stats._stats
# Check that connection ID stats are present
assert "connection_ids_issued" in stats
assert "connection_ids_retired" in stats
assert "connection_id_changes" in stats
# Initial values should be zero
assert stats["connection_ids_issued"] == 0
assert stats["connection_ids_retired"] == 0
assert stats["connection_id_changes"] == 0
def test_connection_id_stats_update(self, connection_with_stats):
"""Test updating connection ID statistics."""
conn = connection_with_stats
# Add some connection IDs to tracking
test_cids = [ConnectionIdTestHelper.generate_connection_id() for _ in range(3)]
for cid in test_cids:
conn._available_connection_ids.add(cid)
# Update stats (this would normally be done by the implementation)
conn._stats["connection_ids_issued"] = len(test_cids)
# Verify stats
stats = conn.get_connection_id_stats()
assert stats["connection_ids_issued"] == 3
assert stats["available_connection_ids"] == 3
def test_connection_id_list_representation(self, connection_with_stats):
"""Test connection ID list representation in stats."""
conn = connection_with_stats
# Add some connection IDs
test_cids = [ConnectionIdTestHelper.generate_connection_id() for _ in range(2)]
for cid in test_cids:
conn._available_connection_ids.add(cid)
# Get stats
stats = conn.get_connection_id_stats()
# Check that CID list is properly formatted
assert "available_cid_list" in stats
assert len(stats["available_cid_list"]) == 2
# All entries should be hex strings
for cid_hex in stats["available_cid_list"]:
assert isinstance(cid_hex, str)
assert len(cid_hex) == 16 # 8 bytes = 16 hex chars
# Performance and stress tests
class TestConnectionIdPerformance:
"""Test connection ID performance and stress scenarios."""
def test_connection_id_generation_performance(self):
"""Test connection ID generation performance."""
start_time = time.time()
# Generate many connection IDs
cids = []
for _ in range(1000):
cid = ConnectionIdTestHelper.generate_connection_id()
cids.append(cid)
end_time = time.time()
generation_time = end_time - start_time
# Should be reasonably fast (less than 1 second for 1000 IDs)
assert generation_time < 1.0
# All should be unique
assert len(set(cids)) == len(cids)
def test_connection_id_tracking_memory(self):
"""Test memory usage of connection ID tracking."""
conn_ids = set()
# Add many connection IDs
for _ in range(1000):
cid = ConnectionIdTestHelper.generate_connection_id()
conn_ids.add(cid)
# Verify they're all stored
assert len(conn_ids) == 1000
# Clean up
conn_ids.clear()
assert len(conn_ids) == 0
if __name__ == "__main__":
# Run tests if executed directly
pytest.main([__file__, "-v"])

View File

@ -0,0 +1,418 @@
"""
Basic QUIC Echo Test
Simple test to verify the basic QUIC flow:
1. Client connects to server
2. Client sends data
3. Server receives data and echoes back
4. Client receives the echo
This test focuses on identifying where the accept_stream issue occurs.
"""
import logging
import pytest
import multiaddr
import trio
from examples.ping.ping import PING_LENGTH, PING_PROTOCOL_ID
from libp2p import new_host
from libp2p.abc import INetStream
from libp2p.crypto.secp256k1 import create_new_key_pair
from libp2p.peer.id import ID
from libp2p.peer.peerinfo import info_from_p2p_addr
from libp2p.transport.quic.config import QUICTransportConfig
from libp2p.transport.quic.connection import QUICConnection
from libp2p.transport.quic.transport import QUICTransport
from libp2p.transport.quic.utils import create_quic_multiaddr
# Set up logging to see what's happening
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)
class TestBasicQUICFlow:
"""Test basic QUIC client-server communication flow."""
@pytest.fixture
def server_key(self):
"""Generate server key pair."""
return create_new_key_pair()
@pytest.fixture
def client_key(self):
"""Generate client key pair."""
return create_new_key_pair()
@pytest.fixture
def server_config(self):
"""Simple server configuration."""
return QUICTransportConfig(
idle_timeout=10.0,
connection_timeout=5.0,
max_concurrent_streams=10,
max_connections=5,
)
@pytest.fixture
def client_config(self):
"""Simple client configuration."""
return QUICTransportConfig(
idle_timeout=10.0,
connection_timeout=5.0,
max_concurrent_streams=5,
)
@pytest.mark.trio
async def test_basic_echo_flow(
self, server_key, client_key, server_config, client_config
):
"""Test basic client-server echo flow with detailed logging."""
print("\n=== BASIC QUIC ECHO TEST ===")
# Create server components
server_transport = QUICTransport(server_key.private_key, server_config)
# Track test state
server_received_data = None
server_connection_established = False
echo_sent = False
async def echo_server_handler(connection: QUICConnection) -> None:
"""Simple echo server handler with detailed logging."""
nonlocal server_received_data, server_connection_established, echo_sent
print("🔗 SERVER: Connection handler called")
server_connection_established = True
try:
print("📡 SERVER: Waiting for incoming stream...")
# Accept stream with timeout and detailed logging
print("📡 SERVER: Calling accept_stream...")
stream = await connection.accept_stream(timeout=5.0)
if stream is None:
print("❌ SERVER: accept_stream returned None")
return
print(f"✅ SERVER: Stream accepted! Stream ID: {stream.stream_id}")
# Read data from the stream
print("📖 SERVER: Reading data from stream...")
server_data = await stream.read(1024)
if not server_data:
print("❌ SERVER: No data received from stream")
return
server_received_data = server_data.decode("utf-8", errors="ignore")
print(f"📨 SERVER: Received data: '{server_received_data}'")
# Echo the data back
echo_message = f"ECHO: {server_received_data}"
print(f"📤 SERVER: Sending echo: '{echo_message}'")
await stream.write(echo_message.encode())
echo_sent = True
print("✅ SERVER: Echo sent successfully")
# Close the stream
await stream.close()
print("🔒 SERVER: Stream closed")
except Exception as e:
print(f"❌ SERVER: Error in handler: {e}")
import traceback
traceback.print_exc()
# Create listener
listener = server_transport.create_listener(echo_server_handler)
listen_addr = create_quic_multiaddr("127.0.0.1", 0, "/quic")
# Variables to track client state
client_connected = False
client_sent_data = False
client_received_echo = None
try:
print("🚀 Starting server...")
async with trio.open_nursery() as nursery:
# Start server listener
success = await listener.listen(listen_addr, nursery)
assert success, "Failed to start server listener"
# Get server address
server_addrs = listener.get_addrs()
server_addr = multiaddr.Multiaddr(
f"{server_addrs[0]}/p2p/{ID.from_pubkey(server_key.public_key)}"
)
print(f"🔧 SERVER: Listening on {server_addr}")
# Give server a moment to be ready
await trio.sleep(0.1)
print("🚀 Starting client...")
# Create client transport
client_transport = QUICTransport(client_key.private_key, client_config)
client_transport.set_background_nursery(nursery)
try:
# Connect to server
print(f"📞 CLIENT: Connecting to {server_addr}")
connection = await client_transport.dial(server_addr)
client_connected = True
print("✅ CLIENT: Connected to server")
# Open a stream
print("📤 CLIENT: Opening stream...")
stream = await connection.open_stream()
print(f"✅ CLIENT: Stream opened with ID: {stream.stream_id}")
# Send test data
test_message = "Hello QUIC Server!"
print(f"📨 CLIENT: Sending message: '{test_message}'")
await stream.write(test_message.encode())
client_sent_data = True
print("✅ CLIENT: Message sent")
# Read echo response
print("📖 CLIENT: Waiting for echo response...")
response_data = await stream.read(1024)
if response_data:
client_received_echo = response_data.decode(
"utf-8", errors="ignore"
)
print(f"📬 CLIENT: Received echo: '{client_received_echo}'")
else:
print("❌ CLIENT: No echo response received")
print("🔒 CLIENT: Closing connection")
await connection.close()
print("🔒 CLIENT: Connection closed")
print("🔒 CLIENT: Closing transport")
await client_transport.close()
print("🔒 CLIENT: Transport closed")
except Exception as e:
print(f"❌ CLIENT: Error: {e}")
import traceback
traceback.print_exc()
finally:
await client_transport.close()
print("🔒 CLIENT: Transport closed")
# Give everything time to complete
await trio.sleep(0.5)
# Cancel nursery to stop server
nursery.cancel_scope.cancel()
finally:
# Cleanup
if not listener._closed:
await listener.close()
await server_transport.close()
# Verify the flow worked
print("\n📊 TEST RESULTS:")
print(f" Server connection established: {server_connection_established}")
print(f" Client connected: {client_connected}")
print(f" Client sent data: {client_sent_data}")
print(f" Server received data: '{server_received_data}'")
print(f" Echo sent by server: {echo_sent}")
print(f" Client received echo: '{client_received_echo}'")
# Test assertions
assert server_connection_established, "Server connection handler was not called"
assert client_connected, "Client failed to connect"
assert client_sent_data, "Client failed to send data"
assert server_received_data == "Hello QUIC Server!", (
f"Server received wrong data: '{server_received_data}'"
)
assert echo_sent, "Server failed to send echo"
assert client_received_echo == "ECHO: Hello QUIC Server!", (
f"Client received wrong echo: '{client_received_echo}'"
)
print("✅ BASIC ECHO TEST PASSED!")
@pytest.mark.trio
async def test_server_accept_stream_timeout(
self, server_key, client_key, server_config, client_config
):
"""Test what happens when server accept_stream times out."""
print("\n=== TESTING SERVER ACCEPT_STREAM TIMEOUT ===")
server_transport = QUICTransport(server_key.private_key, server_config)
accept_stream_called = False
accept_stream_timeout = False
async def timeout_test_handler(connection: QUICConnection) -> None:
"""Handler that tests accept_stream timeout."""
nonlocal accept_stream_called, accept_stream_timeout
print("🔗 SERVER: Connection established, testing accept_stream timeout")
accept_stream_called = True
try:
print("📡 SERVER: Calling accept_stream with 2 second timeout...")
stream = await connection.accept_stream(timeout=2.0)
print(f"✅ SERVER: accept_stream returned: {stream}")
except Exception as e:
print(f"⏰ SERVER: accept_stream timed out or failed: {e}")
accept_stream_timeout = True
listener = server_transport.create_listener(timeout_test_handler)
listen_addr = create_quic_multiaddr("127.0.0.1", 0, "/quic")
client_connected = False
try:
async with trio.open_nursery() as nursery:
# Start server
server_transport.set_background_nursery(nursery)
success = await listener.listen(listen_addr, nursery)
assert success
server_addr = multiaddr.Multiaddr(
f"{listener.get_addrs()[0]}/p2p/{ID.from_pubkey(server_key.public_key)}"
)
print(f"🔧 SERVER: Listening on {server_addr}")
# Create client but DON'T open a stream
async with trio.open_nursery() as client_nursery:
client_transport = QUICTransport(
client_key.private_key, client_config
)
client_transport.set_background_nursery(client_nursery)
try:
print("📞 CLIENT: Connecting (but NOT opening stream)...")
connection = await client_transport.dial(server_addr)
client_connected = True
print("✅ CLIENT: Connected (no stream opened)")
# Wait for server timeout
await trio.sleep(3.0)
await connection.close()
print("🔒 CLIENT: Connection closed")
finally:
await client_transport.close()
nursery.cancel_scope.cancel()
finally:
await listener.close()
await server_transport.close()
print("\n📊 TIMEOUT TEST RESULTS:")
print(f" Client connected: {client_connected}")
print(f" accept_stream called: {accept_stream_called}")
print(f" accept_stream timeout: {accept_stream_timeout}")
assert client_connected, "Client should have connected"
assert accept_stream_called, "accept_stream should have been called"
assert accept_stream_timeout, (
"accept_stream should have timed out when no stream was opened"
)
print("✅ TIMEOUT TEST PASSED!")
@pytest.mark.trio
async def test_yamux_stress_ping():
STREAM_COUNT = 100
listen_addr = create_quic_multiaddr("127.0.0.1", 0, "/quic")
latencies = []
failures = []
# === Server Setup ===
server_host = new_host(listen_addrs=[listen_addr])
async def handle_ping(stream: INetStream) -> None:
try:
while True:
payload = await stream.read(PING_LENGTH)
if not payload:
break
await stream.write(payload)
except Exception:
await stream.reset()
server_host.set_stream_handler(PING_PROTOCOL_ID, handle_ping)
async with server_host.run(listen_addrs=[listen_addr]):
# Give server time to start
await trio.sleep(0.1)
# === Client Setup ===
destination = str(server_host.get_addrs()[0])
maddr = multiaddr.Multiaddr(destination)
info = info_from_p2p_addr(maddr)
client_listen_addr = create_quic_multiaddr("127.0.0.1", 0, "/quic")
client_host = new_host(listen_addrs=[client_listen_addr])
async with client_host.run(listen_addrs=[client_listen_addr]):
await client_host.connect(info)
async def ping_stream(i: int):
stream = None
try:
start = trio.current_time()
stream = await client_host.new_stream(
info.peer_id, [PING_PROTOCOL_ID]
)
await stream.write(b"\x01" * PING_LENGTH)
with trio.fail_after(5):
response = await stream.read(PING_LENGTH)
if response == b"\x01" * PING_LENGTH:
latency_ms = int((trio.current_time() - start) * 1000)
latencies.append(latency_ms)
print(f"[Ping #{i}] Latency: {latency_ms} ms")
await stream.close()
except Exception as e:
print(f"[Ping #{i}] Failed: {e}")
failures.append(i)
if stream:
await stream.reset()
async with trio.open_nursery() as nursery:
for i in range(STREAM_COUNT):
nursery.start_soon(ping_stream, i)
# === Result Summary ===
print("\n📊 Ping Stress Test Summary")
print(f"Total Streams Launched: {STREAM_COUNT}")
print(f"Successful Pings: {len(latencies)}")
print(f"Failed Pings: {len(failures)}")
if failures:
print(f"❌ Failed stream indices: {failures}")
# === Assertions ===
assert len(latencies) == STREAM_COUNT, (
f"Expected {STREAM_COUNT} successful streams, got {len(latencies)}"
)
assert all(isinstance(x, int) and x >= 0 for x in latencies), (
"Invalid latencies"
)
avg_latency = sum(latencies) / len(latencies)
print(f"✅ Average Latency: {avg_latency:.2f} ms")
assert avg_latency < 1000

Some files were not shown because too many files have changed in this diff Show More