varun-r-mallya/py-libp2p
1
0
Fork 0
mirror of https://github.com/varun-r-mallya/py-libp2p.git synced 2025-12-31 20:36:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into write_msg_pubsub

Browse Source
This commit is contained in:
Manu Sheel Gupta
2025-06-20 07:29:56 -07:00
committed by GitHub
parent c33ab32c33 09b4c846a4
commit 8bddbfb9bb
60 changed files with 10460 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
Makefile
View File

@ -58,7 +58,10 @@ PB = libp2p/crypto/pb/crypto.proto \
libp2p/security/secio/pb/spipe.proto \
libp2p/security/noise/pb/noise.proto \
libp2p/identity/identify/pb/identify.proto \
libp2p/host/autonat/pb/autonat.proto
libp2p/host/autonat/pb/autonat.proto \
libp2p/relay/circuit_v2/pb/circuit.proto \
libp2p/kad_dht/pb/kademlia.proto
PY = $(PB:.proto=_pb2.py)
PYI = $(PB:.proto=_pb2.pyi)

499
docs/examples.circuit_relay.rst Normal file
View File

@ -0,0 +1,499 @@
Circuit Relay v2 Example
========================
This example demonstrates how to use Circuit Relay v2 in py-libp2p. It includes three components:
1. A relay node that provides relay services
2. A destination node that accepts relayed connections
3. A source node that connects to the destination through the relay
Prerequisites
-------------
First, ensure you have py-libp2p installed:
.. code-block:: console
$ python -m pip install libp2p
Collecting libp2p
...
Successfully installed libp2p-x.x.x
Relay Node
----------
Create a file named ``relay_node.py`` with the following content:
.. code-block:: python
import trio
import logging
import multiaddr
import traceback
from libp2p import new_host
from libp2p.relay.circuit_v2.protocol import CircuitV2Protocol
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
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")
host = new_host()
config = RelayConfig(
enable_hop=True, # Act as a relay
enable_stop=True, # Accept relayed connections
enable_client=False, # Don't use other relays
max_circuit_duration=3600, # 1 hour
max_circuit_bytes=1024 * 1024 * 10, # 10MB
)
# Initialize the relay protocol with allow_hop=True to act as a relay
protocol = CircuitV2Protocol(host, limits=config.limits, allow_hop=True)
print(f"Created relay protocol with hop enabled: {protocol.allow_hop}")
# Start the protocol service
async with host.run(listen_addrs=[listen_addr]):
peer_id = host.get_id()
print("\n" + "="*50)
print(f"Relay node started with ID: {peer_id}")
print(f"Relay node multiaddr: /ip4/127.0.0.1/tcp/9000/p2p/{peer_id}")
print("="*50 + "\n")
print(f"Listening on: {host.get_addrs()}")
try:
async with background_trio_service(protocol):
print("Protocol service started")
transport = CircuitV2Transport(host, protocol, config)
print("Relay service started successfully")
print(f"Relay limits: {protocol.limits}")
while True:
await trio.sleep(10)
print("Relay node still running...")
print(f"Active connections: {len(host.get_network().connections)}")
except Exception as e:
print(f"Error in relay service: {e}")
traceback.print_exc()
if __name__ == "__main__":
try:
trio.run(run_relay)
except Exception as e:
print(f"Error running relay: {e}")
traceback.print_exc()
Destination Node
----------------
Create a file named ``destination_node.py`` with the following content:
.. code-block:: python
import trio
import logging
import multiaddr
import traceback
import sys
from libp2p import new_host
from libp2p.relay.circuit_v2.protocol import CircuitV2Protocol
from libp2p.relay.circuit_v2.transport import CircuitV2Transport
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
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger("destination_node")
async def handle_echo_stream(stream):
"""Handle incoming stream by echoing received data."""
try:
print(f"New echo stream from: {stream.get_protocol()}")
while True:
data = await stream.read(1024)
if not data:
print("Stream closed by remote")
break
message = data.decode('utf-8')
print(f"Received: {message}")
response = f"Echo: {message}".encode('utf-8')
await stream.write(response)
print(f"Sent response: Echo: {message}")
except Exception as e:
print(f"Error handling stream: {e}")
traceback.print_exc()
finally:
await stream.close()
print("Stream closed")
async def run_destination(relay_peer_id=None):
"""
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")
host = new_host()
# Configure as a relay receiver (stop)
config = RelayConfig(
enable_stop=True, # Accept relayed connections
enable_client=True, # Use relays for outbound connections
max_circuit_duration=3600, # 1 hour
max_circuit_bytes=1024 * 1024 * 10, # 10MB
)
# Initialize the relay protocol
protocol = CircuitV2Protocol(host, limits=config.limits, allow_hop=False)
async with host.run(listen_addrs=[listen_addr]):
# Print host information
dest_peer_id = host.get_id()
print("\n" + "="*50)
print(f"Destination node started with ID: {dest_peer_id}")
print(f"Use this ID in the source node: {dest_peer_id}")
print("="*50 + "\n")
print(f"Listening on: {host.get_addrs()}")
# Set stream handler for the echo protocol
host.set_stream_handler("/echo/1.0.0", handle_echo_stream)
print("Registered echo protocol handler")
# Start the protocol service in the background
async with background_trio_service(protocol):
print("Protocol service started")
# Create and register the transport
transport = CircuitV2Transport(host, protocol, config)
print("Transport created")
# Create a listener for relayed connections
listener = transport.create_listener(handle_echo_stream)
print("Created relay listener")
# Start listening for relayed connections
async with trio.open_nursery() as nursery:
await listener.listen("/p2p-circuit", nursery)
print("Destination node ready to accept relayed connections")
if not relay_peer_id:
print("No relay peer ID provided. Please enter the relay's peer ID:")
print("Waiting for relay peer ID input...")
while True:
if sys.stdin.isatty(): # Only try to read from stdin if it's a terminal
try:
relay_peer_id = input("Enter relay peer ID: ").strip()
if relay_peer_id:
break
except EOFError:
await trio.sleep(5)
else:
print("No terminal detected. Waiting for relay peer ID as command line argument.")
await trio.sleep(10)
continue
# Connect to the relay node with the provided relay peer ID
relay_addr_str = f"/ip4/127.0.0.1/tcp/9000/p2p/{relay_peer_id}"
print(f"Connecting to relay at {relay_addr_str}")
try:
# Convert string address to multiaddr, then to peer info
relay_maddr = multiaddr.Multiaddr(relay_addr_str)
relay_peer_info = info_from_p2p_addr(relay_maddr)
await host.connect(relay_peer_info)
print("Connected to relay successfully")
# Add the relay to the transport's discovery
transport.discovery._add_relay(relay_peer_info.peer_id)
print(f"Added relay {relay_peer_info.peer_id} to discovery")
# Keep the node running
while True:
await trio.sleep(10)
print("Destination node still running...")
except Exception as e:
print(f"Failed to connect to relay: {e}")
traceback.print_exc()
if __name__ == "__main__":
print("Starting destination node...")
relay_id = None
if len(sys.argv) > 1:
relay_id = sys.argv[1]
print(f"Using provided relay ID: {relay_id}")
trio.run(run_destination, relay_id)
Source Node
-----------
Create a file named ``source_node.py`` with the following content:
.. code-block:: python
import trio
import logging
import multiaddr
import traceback
import sys
from libp2p import new_host
from libp2p.peer.peerinfo import PeerInfo
from libp2p.peer.id import ID
from libp2p.relay.circuit_v2.protocol import CircuitV2Protocol
from libp2p.relay.circuit_v2.transport import CircuitV2Transport
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.relay.circuit_v2.discovery import RelayInfo
# 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")
host = new_host()
# Configure as a relay client
config = RelayConfig(
enable_client=True, # Use relays for outbound connections
max_circuit_duration=3600, # 1 hour
max_circuit_bytes=1024 * 1024 * 10, # 10MB
)
# Initialize the relay protocol
protocol = CircuitV2Protocol(host, limits=config.limits, allow_hop=False)
# Start the protocol service
async with host.run(listen_addrs=[listen_addr]):
# Print host information
print(f"Source node started with ID: {host.get_id()}")
print(f"Listening on: {host.get_addrs()}")
# Start the protocol service in the background
async with background_trio_service(protocol):
print("Protocol service started")
# Create and register the transport
transport = CircuitV2Transport(host, protocol, config)
# Get relay peer ID if not provided
if not relay_peer_id:
print("No relay peer ID provided. Please enter the relay's peer ID:")
while True:
if sys.stdin.isatty(): # Only try to read from stdin if it's a terminal
try:
relay_peer_id = input("Enter relay peer ID: ").strip()
if relay_peer_id:
break
except EOFError:
await trio.sleep(5)
else:
print("No terminal detected. Waiting for relay peer ID as command line argument.")
await trio.sleep(10)
continue
# Connect to the relay node with the provided relay peer ID
relay_addr_str = f"/ip4/127.0.0.1/tcp/9000/p2p/{relay_peer_id}"
print(f"Connecting to relay at {relay_addr_str}")
try:
# Convert string address to multiaddr, then to peer info
relay_maddr = multiaddr.Multiaddr(relay_addr_str)
relay_peer_info = info_from_p2p_addr(relay_maddr)
await host.connect(relay_peer_info)
print("Connected to relay successfully")
# Manually add the relay to the discovery service
relay_id = relay_peer_info.peer_id
now = trio.current_time()
# Create relay info and add it to discovery
relay_info = RelayInfo(
peer_id=relay_id,
discovered_at=now,
last_seen=now
)
transport.discovery._discovered_relays[relay_id] = relay_info
print(f"Added relay {relay_id} to discovery")
# Start relay discovery in the background
async with background_trio_service(transport.discovery):
print("Relay discovery started")
# Wait for relay discovery
await trio.sleep(5)
print("Relay discovery completed")
# Get destination peer ID if not provided
if not destination_peer_id:
print("No destination peer ID provided. Please enter the destination's peer ID:")
while True:
if sys.stdin.isatty(): # Only try to read from stdin if it's a terminal
try:
destination_peer_id = input("Enter destination peer ID: ").strip()
if destination_peer_id:
break
except EOFError:
await trio.sleep(5)
else:
print("No terminal detected. Waiting for destination peer ID as command line argument.")
await trio.sleep(10)
continue
print(f"Attempting to connect to {destination_peer_id} via relay")
# Check if we have any discovered relays
discovered_relays = list(transport.discovery._discovered_relays.keys())
print(f"Discovered relays: {discovered_relays}")
try:
# Create a circuit relay multiaddr for the destination
dest_id = ID.from_base58(destination_peer_id)
# Create a circuit multiaddr that includes the relay
# Format: /ip4/127.0.0.1/tcp/9000/p2p/RELAY_ID/p2p-circuit/p2p/DEST_ID
circuit_addr = multiaddr.Multiaddr(f"{relay_addr_str}/p2p-circuit/p2p/{destination_peer_id}")
print(f"Created circuit address: {circuit_addr}")
# Dial using the circuit address
connection = await transport.dial(circuit_addr)
print("Connection established through relay!")
# Open a stream using the echo protocol
stream = await connection.new_stream("/echo/1.0.0")
# Send messages periodically
for i in range(5):
message = f"Hello from source, message {i+1}"
print(f"Sending: {message}")
await stream.write(message.encode('utf-8'))
response = await stream.read(1024)
print(f"Received: {response.decode('utf-8')}")
await trio.sleep(1)
# Close the stream
await stream.close()
print("Stream closed")
except Exception as e:
print(f"Error connecting through relay: {e}")
print("Detailed error:")
traceback.print_exc()
# Keep the node running for a while
await trio.sleep(30)
print("Source node shutting down")
except Exception as e:
print(f"Error: {e}")
traceback.print_exc()
if __name__ == "__main__":
relay_id = None
dest_id = None
# Parse command line arguments if provided
if len(sys.argv) > 1:
relay_id = sys.argv[1]
print(f"Using provided relay ID: {relay_id}")
if len(sys.argv) > 2:
dest_id = sys.argv[2]
print(f"Using provided destination ID: {dest_id}")
trio.run(run_source, relay_id, dest_id)
Running the Example
-------------------
1. First, start the relay node:
.. code-block:: console
$ python relay_node.py
Created relay protocol with hop enabled: True
==================================================
Relay node started with ID: QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx
Relay node multiaddr: /ip4/127.0.0.1/tcp/9000/p2p/QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx
==================================================
Listening on: [<Multiaddr /ip4/0.0.0.0/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)
Note the relay node\'s peer ID (in this example: `QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx`). You\'ll need this for the other nodes.
2. Next, start the destination node:
.. code-block:: console
$ python destination_node.py
Starting destination node...
==================================================
Destination node started with ID: QmPBr38KeQG2ibyL4fxq6yJWpfoVNCqJMHBdNyn1Qe4h5s
Use this ID in the source node: QmPBr38KeQG2ibyL4fxq6yJWpfoVNCqJMHBdNyn1Qe4h5s
==================================================
Listening on: [<Multiaddr /ip4/0.0.0.0/tcp/9001/p2p/QmPBr38KeQG2ibyL4fxq6yJWpfoVNCqJMHBdNyn1Qe4h5s>]
Registered echo protocol handler
Protocol service started
Transport created
Created relay listener
Destination node ready to accept relayed connections
No relay peer ID provided. Please enter the relay\'s peer ID:
Waiting for relay peer ID input...
Enter relay peer ID: QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx
Connecting to relay at /ip4/127.0.0.1/tcp/9000/p2p/QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx
Connected to relay successfully
Added relay QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx to discovery
Destination node still running...
Note the destination node's peer ID (in this example: `QmPBr38KeQG2ibyL4fxq6yJWpfoVNCqJMHBdNyn1Qe4h5s`). You'll need this for the source node.
3. Finally, start the source node:
.. code-block:: console
$ python source_node.py
Source node started with ID: QmPyM56cgmFoHTgvMgGfDWRdVRQznmxCDDDg2dJ8ygVXj3
Listening on: [<Multiaddr /ip4/0.0.0.0/tcp/9002/p2p/QmPyM56cgmFoHTgvMgGfDWRdVRQznmxCDDDg2dJ8ygVXj3>]
Protocol service started
No relay peer ID provided. Please enter the relay\'s peer ID:
Enter relay peer ID: QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx
Connecting to relay at /ip4/127.0.0.1/tcp/9000/p2p/QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx
Connected to relay successfully
Added relay QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx to discovery
Relay discovery started
Relay discovery completed
No destination peer ID provided. Please enter the destination\'s peer ID:
Enter destination peer ID: QmPBr38KeQG2ibyL4fxq6yJWpfoVNCqJMHBdNyn1Qe4h5s
Attempting to connect to QmPBr38KeQG2ibyL4fxq6yJWpfoVNCqJMHBdNyn1Qe4h5s via relay
Discovered relays: [<libp2p.peer.id.ID (QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx)>]
Created circuit address: /ip4/127.0.0.1/tcp/9000/p2p/QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx/p2p-circuit/p2p/QmPBr38KeQG2ibyL4fxq6yJWpfoVNCqJMHBdNyn1Qe4h5s
At this point, the source node will establish a connection through the relay to the destination node and start sending messages.
4. Alternatively, you can provide the peer IDs as command-line arguments:
.. code-block:: console
# For the destination node (provide relay ID)
$ python destination_node.py QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx
# For the source node (provide both relay and destination IDs)
$ python source_node.py QmaUigQJ9nJERa6GaZuyfaiX91QjYwoQJ46JS3k7ys7SLx QmPBr38KeQG2ibyL4fxq6yJWpfoVNCqJMHBdNyn1Qe4h5s
This example demonstrates how to use Circuit Relay v2 to establish connections between peers that cannot connect directly. The peer IDs are dynamically generated for each node, and the relay facilitates communication between the source and destination nodes.

124
docs/examples.kademlia.rst Normal file
View File

@ -0,0 +1,124 @@
Kademlia DHT Demo
=================
This example demonstrates a Kademlia Distributed Hash Table (DHT) implementation with both value storage/retrieval and content provider advertisement/discovery functionality.
.. code-block:: console
$ python -m pip install libp2p
Collecting libp2p
...
Successfully installed libp2p-x.x.x
$ cd examples/kademlia
$ python kademlia.py --mode server
2025-06-13 19:51:25,424 - kademlia-example - INFO - Running in server mode on port 0
2025-06-13 19:51:25,426 - kademlia-example - INFO - Connected to bootstrap nodes: []
2025-06-13 19:51:25,426 - kademlia-example - INFO - To connect to this node, use: --bootstrap /ip4/127.0.0.1/tcp/28910/p2p/16Uiu2HAm7EsNv5vvjPAehGAVfChjYjD63ZHyWogQRdzntSbAg9ef
2025-06-13 19:51:25,426 - kademlia-example - INFO - Saved server address to log: /ip4/127.0.0.1/tcp/28910/p2p/16Uiu2HAm7EsNv5vvjPAehGAVfChjYjD63ZHyWogQRdzntSbAg9ef
2025-06-13 19:51:25,427 - kademlia-example - INFO - DHT service started in SERVER mode
2025-06-13 19:51:25,427 - kademlia-example - INFO - Stored value 'Hello message from Sumanjeet' with key: FVDjasarSFDoLPMdgnp1dHSbW2ZAfN8NU2zNbCQeczgP
2025-06-13 19:51:25,427 - kademlia-example - INFO - Successfully advertised as server for content: 361f2ed1183bca491b8aec11f0b9e5c06724759b0f7480ae7fb4894901993bc8
Copy the line that starts with ``--bootstrap``, open a new terminal in the same folder and run the client:
.. code-block:: console
$ python kademlia.py --mode client --bootstrap /ip4/127.0.0.1/tcp/28910/p2p/16Uiu2HAm7EsNv5vvjPAehGAVfChjYjD63ZHyWogQRdzntSbAg9ef
2025-06-13 19:51:37,022 - kademlia-example - INFO - Running in client mode on port 0
2025-06-13 19:51:37,026 - kademlia-example - INFO - Connected to bootstrap nodes: [<libp2p.peer.id.ID (16Uiu2HAm7EsNv5vvjPAehGAVfChjYjD63ZHyWogQRdzntSbAg9ef)>]
2025-06-13 19:51:37,027 - kademlia-example - INFO - DHT service started in CLIENT mode
2025-06-13 19:51:37,027 - kademlia-example - INFO - Looking up key: FVDjasarSFDoLPMdgnp1dHSbW2ZAfN8NU2zNbCQeczgP
2025-06-13 19:51:37,031 - kademlia-example - INFO - Retrieved value: Hello message from Sumanjeet
2025-06-13 19:51:37,031 - kademlia-example - INFO - Looking for servers of content: 361f2ed1183bca491b8aec11f0b9e5c06724759b0f7480ae7fb4894901993bc8
2025-06-13 19:51:37,035 - kademlia-example - INFO - Found 1 servers for content: ['16Uiu2HAm7EsNv5vvjPAehGAVfChjYjD63ZHyWogQRdzntSbAg9ef']
Alternatively, if you run the server first, the client can automatically extract the bootstrap address from the server log file:
.. code-block:: console
$ python kademlia.py --mode client
2025-06-13 19:51:37,022 - kademlia-example - INFO - Running in client mode on port 0
2025-06-13 19:51:37,026 - kademlia-example - INFO - Connected to bootstrap nodes: [<libp2p.peer.id.ID (16Uiu2HAm7EsNv5vvjPAehGAVfChjYjD63ZHyWogQRdzntSbAg9ef)>]
2025-06-13 19:51:37,027 - kademlia-example - INFO - DHT service started in CLIENT mode
2025-06-13 19:51:37,027 - kademlia-example - INFO - Looking up key: FVDjasarSFDoLPMdgnp1dHSbW2ZAfN8NU2zNbCQeczgP
2025-06-13 19:51:37,031 - kademlia-example - INFO - Retrieved value: Hello message from Sumanjeet
2025-06-13 19:51:37,031 - kademlia-example - INFO - Looking for servers of content: 361f2ed1183bca491b8aec11f0b9e5c06724759b0f7480ae7fb4894901993bc8
2025-06-13 19:51:37,035 - kademlia-example - INFO - Found 1 servers for content: ['16Uiu2HAm7EsNv5vvjPAehGAVfChjYjD63ZHyWogQRdzntSbAg9ef']
The demo showcases key DHT operations:
- **Value Storage & Retrieval**: The server stores a value, and the client retrieves it
- **Content Provider Discovery**: The server advertises content, and the client finds providers
- **Peer Discovery**: Automatic bootstrap and peer routing using the Kademlia algorithm
- **Network Resilience**: Distributed storage across multiple nodes (when available)
Command Line Options
--------------------
The Kademlia demo supports several command line options for customization:
.. code-block:: console
$ python kademlia.py --help
usage: kademlia.py [-h] [--mode MODE] [--port PORT] [--bootstrap [BOOTSTRAP ...]] [--verbose]
Kademlia DHT example with content server functionality
options:
-h, --help show this help message and exit
--mode MODE Run as a server or client node (default: server)
--port PORT Port to listen on (0 for random) (default: 0)
--bootstrap [BOOTSTRAP ...]
Multiaddrs of bootstrap nodes. Provide a space-separated list of addresses.
This is required for client mode.
--verbose Enable verbose logging
**Examples:**
Start server on a specific port:
.. code-block:: console
$ python kademlia.py --mode server --port 8000
Start client with verbose logging:
.. code-block:: console
$ python kademlia.py --mode client --verbose
Connect to multiple bootstrap nodes:
.. code-block:: console
$ python kademlia.py --mode client --bootstrap /ip4/127.0.0.1/tcp/8000/p2p/... /ip4/127.0.0.1/tcp/8001/p2p/...
How It Works
------------
The Kademlia DHT implementation demonstrates several key concepts:
**Server Mode:**
- Stores key-value pairs in the distributed hash table
- Advertises itself as a content provider for specific content
- Handles incoming DHT requests from other nodes
- Maintains routing table with known peers
**Client Mode:**
- Connects to bootstrap nodes to join the network
- Retrieves values by their keys from the DHT
- Discovers content providers for specific content
- Performs network lookups using the Kademlia algorithm
**Key Components:**
- **Routing Table**: Organizes peers in k-buckets based on XOR distance
- **Value Store**: Manages key-value storage with TTL (time-to-live)
- **Provider Store**: Tracks which peers provide specific content
- **Peer Routing**: Implements iterative lookups to find closest peers
The full source code for this example is below:
.. literalinclude:: ../examples/kademlia/kademlia.py
:language: python
:linenos:

2
docs/examples.rst
View File

@ -11,3 +11,5 @@ Examples
examples.echo
examples.ping
examples.pubsub
examples.circuit_relay
examples.kademlia

4
docs/index.rst
View File

@ -12,10 +12,6 @@ The Python implementation of the libp2p networking stack
getting_started
release_notes
.. toctree::
:maxdepth: 1
:caption: Community
.. toctree::
:maxdepth: 1
:caption: py-libp2p

22
docs/libp2p.kad_dht.pb.rst Normal file
View File

@ -0,0 +1,22 @@
libp2p.kad\_dht.pb package
==========================
Submodules
----------
libp2p.kad_dht.pb.kademlia_pb2 module
-------------------------------------
.. automodule:: libp2p.kad_dht.pb.kademlia_pb2
:members:
:undoc-members:
:show-inheritance:
Module contents
---------------
.. automodule:: libp2p.kad_dht.pb
:no-index:
:members:
:undoc-members:
:show-inheritance:

77
docs/libp2p.kad_dht.rst Normal file
View File

@ -0,0 +1,77 @@
libp2p.kad\_dht package
=======================
Subpackages
-----------
.. toctree::
:maxdepth: 4
libp2p.kad_dht.pb
Submodules
----------
libp2p.kad\_dht.kad\_dht module
-------------------------------
.. automodule:: libp2p.kad_dht.kad_dht
:members:
:undoc-members:
:show-inheritance:
libp2p.kad\_dht.peer\_routing module
------------------------------------
.. automodule:: libp2p.kad_dht.peer_routing
:members:
:undoc-members:
:show-inheritance:
libp2p.kad\_dht.provider\_store module
--------------------------------------
.. automodule:: libp2p.kad_dht.provider_store
:members:
:undoc-members:
:show-inheritance:
libp2p.kad\_dht.routing\_table module
-------------------------------------
.. automodule:: libp2p.kad_dht.routing_table
:members:
:undoc-members:
:show-inheritance:
libp2p.kad\_dht.utils module
----------------------------
.. automodule:: libp2p.kad_dht.utils
:members:
:undoc-members:
:show-inheritance:
libp2p.kad\_dht.value\_store module
-----------------------------------
.. automodule:: libp2p.kad_dht.value_store
:members:
:undoc-members:
:show-inheritance:
libp2p.kad\_dht.pb
------------------
.. automodule:: libp2p.kad_dht.pb
:members:
:undoc-members:
:show-inheritance:
Module contents
---------------
.. automodule:: libp2p.kad_dht
:members:
:undoc-members:
:show-inheritance:

22
docs/libp2p.relay.circuit_v2.pb.rst Normal file
View File

@ -0,0 +1,22 @@
libp2p.relay.circuit_v2.pb package
==================================
Submodules
----------
libp2p.relay.circuit_v2.pb.circuit_pb2 module
---------------------------------------------
.. automodule:: libp2p.relay.circuit_v2.pb.circuit_pb2
:members:
:show-inheritance:
:undoc-members:
Module contents
---------------
.. automodule:: libp2p.relay.circuit_v2.pb
:members:
:show-inheritance:
:undoc-members:
:no-index:

70
docs/libp2p.relay.circuit_v2.rst Normal file
View File

@ -0,0 +1,70 @@
libp2p.relay.circuit_v2 package
===============================
Subpackages
-----------
.. toctree::
:maxdepth: 4
libp2p.relay.circuit_v2.pb
Submodules
----------
libp2p.relay.circuit_v2.protocol module
---------------------------------------
.. automodule:: libp2p.relay.circuit_v2.protocol
:members:
:show-inheritance:
:undoc-members:
libp2p.relay.circuit_v2.transport module
----------------------------------------
.. automodule:: libp2p.relay.circuit_v2.transport
:members:
:show-inheritance:
:undoc-members:
libp2p.relay.circuit_v2.discovery module
----------------------------------------
.. automodule:: libp2p.relay.circuit_v2.discovery
:members:
:show-inheritance:
:undoc-members:
libp2p.relay.circuit_v2.resources module
----------------------------------------
.. automodule:: libp2p.relay.circuit_v2.resources
:members:
:show-inheritance:
:undoc-members:
libp2p.relay.circuit_v2.config module
-------------------------------------
.. automodule:: libp2p.relay.circuit_v2.config
:members:
:show-inheritance:
:undoc-members:
libp2p.relay.circuit_v2.protocol_buffer module
----------------------------------------------
.. automodule:: libp2p.relay.circuit_v2.protocol_buffer
:members:
:show-inheritance:
:undoc-members:
Module contents
---------------
.. automodule:: libp2p.relay.circuit_v2
:members:
:show-inheritance:
:undoc-members:
:no-index:

19
docs/libp2p.relay.rst Normal file
View File

@ -0,0 +1,19 @@
libp2p.relay package
====================
Subpackages
-----------
.. toctree::
:maxdepth: 4
libp2p.relay.circuit_v2
Module contents
---------------
.. automodule:: libp2p.relay
:members:
:show-inheritance:
:undoc-members:
:no-index:

2
docs/libp2p.rst
View File

@ -11,10 +11,12 @@ Subpackages
libp2p.host
libp2p.identity
libp2p.io
libp2p.kad_dht
libp2p.network
libp2p.peer
libp2p.protocol_muxer
libp2p.pubsub
libp2p.relay
libp2p.security
libp2p.stream_muxer
libp2p.tools

300
examples/kademlia/kademlia.py Normal file
View File

@ -0,0 +1,300 @@
#!/usr/bin/env python
"""
A basic example of using the Kademlia DHT implementation, with all setup logic inlined.
This example demonstrates both value storage/retrieval and content server
advertisement/discovery.
"""
import argparse
import logging
import os
import random
import secrets
import sys
import base58
from multiaddr import (
Multiaddr,
)
import trio
from libp2p import (
new_host,
)
from libp2p.abc import (
IHost,
)
from libp2p.crypto.secp256k1 import (
create_new_key_pair,
)
from libp2p.kad_dht.kad_dht import (
DHTMode,
KadDHT,
)
from libp2p.kad_dht.utils import (
create_key_from_binary,
)
from libp2p.tools.async_service import (
background_trio_service,
)
from libp2p.tools.utils import (
info_from_p2p_addr,
)
# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
handlers=[logging.StreamHandler()],
)
logger = logging.getLogger("kademlia-example")
# Configure DHT module loggers to inherit from the parent logger
# This ensures all kademlia-example.* loggers use the same configuration
# Get the directory where this script is located
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
SERVER_ADDR_LOG = os.path.join(SCRIPT_DIR, "server_node_addr.txt")
# Set the level for all child loggers
for module in [
"kad_dht",
"value_store",
"peer_routing",
"routing_table",
"provider_store",
]:
child_logger = logging.getLogger(f"kademlia-example.{module}")
child_logger.setLevel(logging.INFO)
child_logger.propagate = True # Allow propagation to parent
# File to store node information
bootstrap_nodes = []
# function to take bootstrap_nodes as input and connects to them
async def connect_to_bootstrap_nodes(host: IHost, bootstrap_addrs: list[str]) -> None:
"""
Connect to the bootstrap nodes provided in the list.
params: host: The host instance to connect to
bootstrap_addrs: List of bootstrap node addresses
Returns
-------
None
"""
for addr in bootstrap_addrs:
try:
peerInfo = info_from_p2p_addr(Multiaddr(addr))
host.get_peerstore().add_addrs(peerInfo.peer_id, peerInfo.addrs, 3600)
await host.connect(peerInfo)
except Exception as e:
logger.error(f"Failed to connect to bootstrap node {addr}: {e}")
def save_server_addr(addr: str) -> None:
"""Append the server's multiaddress to the log file."""
try:
with open(SERVER_ADDR_LOG, "w") as f:
f.write(addr + "\n")
logger.info(f"Saved server address to log: {addr}")
except Exception as e:
logger.error(f"Failed to save server address: {e}")
def load_server_addrs() -> list[str]:
"""Load all server multiaddresses from the log file."""
if not os.path.exists(SERVER_ADDR_LOG):
return []
try:
with open(SERVER_ADDR_LOG) as f:
return [line.strip() for line in f if line.strip()]
except Exception as e:
logger.error(f"Failed to load server addresses: {e}")
return []
async def run_node(
port: int, mode: str, bootstrap_addrs: list[str] | None = None
) -> None:
"""Run a node that serves content in the DHT with setup inlined."""
try:
if port <= 0:
port = random.randint(10000, 60000)
logger.debug(f"Using port: {port}")
# Convert string mode to DHTMode enum
if mode is None or mode.upper() == "CLIENT":
dht_mode = DHTMode.CLIENT
elif mode.upper() == "SERVER":
dht_mode = DHTMode.SERVER
else:
logger.error(f"Invalid mode: {mode}. Must be 'client' or 'server'")
sys.exit(1)
# Load server addresses for client mode
if dht_mode == DHTMode.CLIENT:
server_addrs = load_server_addrs()
if server_addrs:
logger.info(f"Loaded {len(server_addrs)} server addresses from log")
bootstrap_nodes.append(server_addrs[0]) # Use the first server address
else:
logger.warning("No server addresses found in log file")
if bootstrap_addrs:
for addr in bootstrap_addrs:
bootstrap_nodes.append(addr)
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]):
peer_id = host.get_id().pretty()
addr_str = f"/ip4/127.0.0.1/tcp/{port}/p2p/{peer_id}"
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)
# Start the DHT service
async with background_trio_service(dht):
logger.info(f"DHT service started in {dht_mode.value} mode")
val_key = create_key_from_binary(b"py-libp2p kademlia example value")
content = b"Hello from python node "
content_key = create_key_from_binary(content)
if dht_mode == DHTMode.SERVER:
# Store a value in the DHT
msg = "Hello message from Sumanjeet"
val_data = msg.encode()
await dht.put_value(val_key, val_data)
logger.info(
f"Stored value '{val_data.decode()}'"
f"with key: {base58.b58encode(val_key).decode()}"
)
# Advertise as content server
success = await dht.provider_store.provide(content_key)
if success:
logger.info(
"Successfully advertised as server"
f"for content: {content_key.hex()}"
)
else:
logger.warning("Failed to advertise as content server")
else:
# retrieve the value
logger.info(
"Looking up key: %s", base58.b58encode(val_key).decode()
)
val_data = await dht.get_value(val_key)
if val_data:
try:
logger.info(f"Retrieved value: {val_data.decode()}")
except UnicodeDecodeError:
logger.info(f"Retrieved value (bytes): {val_data!r}")
else:
logger.warning("Failed to retrieve value")
# Also check if we can find servers for our own content
logger.info("Looking for servers of content: %s", content_key.hex())
providers = await dht.provider_store.find_providers(content_key)
if providers:
logger.info(
"Found %d servers for content: %s",
len(providers),
[p.peer_id.pretty() for p in providers],
)
else:
logger.warning(
"No servers found for content %s", content_key.hex()
)
# Keep the node running
while True:
logger.debug(
"Status - Connected peers: %d,"
"Peers in store: %d, Values in store: %d",
len(dht.host.get_connected_peers()),
len(dht.host.get_peerstore().peer_ids()),
len(dht.value_store.store),
)
await trio.sleep(10)
except Exception as e:
logger.error(f"Server node error: {e}", exc_info=True)
sys.exit(1)
def parse_args():
"""Parse command line arguments."""
parser = argparse.ArgumentParser(
description="Kademlia DHT example with content server functionality"
)
parser.add_argument(
"--mode",
default="server",
help="Run as a server or client node",
)
parser.add_argument(
"--port",
type=int,
default=0,
help="Port to listen on (0 for random)",
)
parser.add_argument(
"--bootstrap",
type=str,
nargs="*",
help=(
"Multiaddrs of bootstrap nodes. "
"Provide a space-separated list of addresses. "
"This is required for client mode."
),
)
# add option to use verbose logging
parser.add_argument(
"--verbose",
action="store_true",
help="Enable verbose logging",
)
args = parser.parse_args()
# Set logging level based on verbosity
if args.verbose:
logging.getLogger().setLevel(logging.DEBUG)
else:
logging.getLogger().setLevel(logging.INFO)
return args
def main():
"""Main entry point for the kademlia demo."""
try:
args = parse_args()
logger.info(
"Running in %s mode on port %d",
args.mode,
args.port,
)
trio.run(run_node, args.port, args.mode, args.bootstrap)
except Exception as e:
logger.critical(f"Script failed: {e}", exc_info=True)
sys.exit(1)
if __name__ == "__main__":
main()

8
libp2p/abc.py
View File

@ -2130,14 +2130,14 @@ class IPubsub(ServiceAPI):
...
@abstractmethod
async def publish(self, topic_id: str, data: bytes) -> None:
async def publish(self, topic_id: str | list[str], data: bytes) -> None:
"""
Publish a message to a topic.
Publish a message to a topic or multiple topics.
Parameters
----------
topic_id : str
The identifier of the topic.
topic_id : str | list[str]
The identifier of the topic (str) or topics (list[str]).
data : bytes
The data to publish.

30
libp2p/kad_dht/__init__.py Normal file
View File

@ -0,0 +1,30 @@
"""
Kademlia DHT implementation for py-libp2p.
This module provides a Distributed Hash Table (DHT) implementation
based on the Kademlia protocol.
"""
from .kad_dht import (
KadDHT,
)
from .peer_routing import (
PeerRouting,
)
from .routing_table import (
RoutingTable,
)
from .utils import (
create_key_from_binary,
)
from .value_store import (
ValueStore,
)
__all__ = [
"KadDHT",
"RoutingTable",
"PeerRouting",
"ValueStore",
"create_key_from_binary",
]

616
libp2p/kad_dht/kad_dht.py Normal file
View File

@ -0,0 +1,616 @@
"""
Kademlia DHT implementation for py-libp2p.
This module provides a complete Distributed Hash Table (DHT)
implementation based on the Kademlia algorithm and protocol.
"""
from enum import Enum
import logging
import time
from multiaddr import (
Multiaddr,
)
import trio
import varint
from libp2p.abc import (
IHost,
)
from libp2p.custom_types import (
TProtocol,
)
from libp2p.network.stream.net_stream import (
INetStream,
)
from libp2p.peer.id import (
ID,
)
from libp2p.peer.peerinfo import (
PeerInfo,
)
from libp2p.tools.async_service import (
Service,
)
from .pb.kademlia_pb2 import (
Message,
)
from .peer_routing import (
PeerRouting,
)
from .provider_store import (
ProviderStore,
)
from .routing_table import (
RoutingTable,
)
from .value_store import (
ValueStore,
)
logger = logging.getLogger("kademlia-example.kad_dht")
# logger = logging.getLogger("libp2p.kademlia")
# Default parameters
PROTOCOL_ID = TProtocol("/ipfs/kad/1.0.0")
ROUTING_TABLE_REFRESH_INTERVAL = 1 * 60 # 1 min in seconds for testing
TTL = 24 * 60 * 60 # 24 hours in seconds
ALPHA = 3
QUERY_TIMEOUT = 10 # seconds
class DHTMode(Enum):
"""DHT operation modes."""
CLIENT = "CLIENT"
SERVER = "SERVER"
class KadDHT(Service):
"""
Kademlia DHT implementation for libp2p.
This class provides a DHT implementation that combines routing table management,
peer discovery, content routing, and value storage.
"""
def __init__(self, host: IHost, mode: DHTMode):
"""
Initialize a new Kademlia DHT node.
:param host: The libp2p host.
:param mode: The mode of host (Client or Server) - must be DHTMode enum
"""
super().__init__()
self.host = host
self.local_peer_id = host.get_id()
# Validate that mode is a DHTMode enum
if not isinstance(mode, DHTMode):
raise TypeError(f"mode must be DHTMode enum, got {type(mode)}")
self.mode = mode
# Initialize the routing table
self.routing_table = RoutingTable(self.local_peer_id, self.host)
# Initialize peer routing
self.peer_routing = PeerRouting(host, self.routing_table)
# Initialize value store
self.value_store = ValueStore(host=host, local_peer_id=self.local_peer_id)
# Initialize provider store with host and peer_routing references
self.provider_store = ProviderStore(host=host, peer_routing=self.peer_routing)
# Last time we republished provider records
self._last_provider_republish = time.time()
# Set protocol handlers
host.set_stream_handler(PROTOCOL_ID, self.handle_stream)
async def run(self) -> None:
"""Run the DHT service."""
logger.info(f"Starting Kademlia DHT with peer ID {self.local_peer_id}")
# Main service loop
while self.manager.is_running:
# Periodically refresh the routing table
await self.refresh_routing_table()
# Check if it's time to republish provider records
current_time = time.time()
# await self._republish_provider_records()
self._last_provider_republish = current_time
# Clean up expired values and provider records
expired_values = self.value_store.cleanup_expired()
if expired_values > 0:
logger.debug(f"Cleaned up {expired_values} expired values")
self.provider_store.cleanup_expired()
# Wait before next maintenance cycle
await trio.sleep(ROUTING_TABLE_REFRESH_INTERVAL)
async def switch_mode(self, new_mode: DHTMode) -> DHTMode:
"""
Switch the DHT mode.
:param new_mode: The new mode - must be DHTMode enum
:return: The new mode as DHTMode enum
"""
# Validate that new_mode is a DHTMode enum
if not isinstance(new_mode, DHTMode):
raise TypeError(f"new_mode must be DHTMode enum, got {type(new_mode)}")
if new_mode == DHTMode.CLIENT:
self.routing_table.cleanup_routing_table()
self.mode = new_mode
logger.info(f"Switched to {new_mode.value} mode")
return self.mode
async def handle_stream(self, stream: INetStream) -> None:
"""
Handle an incoming DHT stream using varint length prefixes.
"""
if self.mode == DHTMode.CLIENT:
stream.close
return
peer_id = stream.muxed_conn.peer_id
logger.debug(f"Received DHT stream from peer {peer_id}")
await self.add_peer(peer_id)
logger.debug(f"Added peer {peer_id} to routing table")
try:
# Read varint-prefixed length for the message
length_prefix = b""
while True:
byte = await stream.read(1)
if not byte:
logger.warning("Stream closed while reading varint length")
await stream.close()
return
length_prefix += byte
if byte[0] & 0x80 == 0:
break
msg_length = varint.decode_bytes(length_prefix)
# Read the message bytes
msg_bytes = await stream.read(msg_length)
if len(msg_bytes) < msg_length:
logger.warning("Failed to read full message from stream")
await stream.close()
return
try:
# Parse as protobuf
message = Message()
message.ParseFromString(msg_bytes)
logger.debug(
f"Received DHT message from {peer_id}, type: {message.type}"
)
# Handle FIND_NODE message
if message.type == Message.MessageType.FIND_NODE:
# Get target key directly from protobuf
target_key = message.key
# Find closest peers to the target key
closest_peers = self.routing_table.find_local_closest_peers(
target_key, 20
)
logger.debug(f"Found {len(closest_peers)} peers close to target")
# Build response message with protobuf
response = Message()
response.type = Message.MessageType.FIND_NODE
# Add closest peers to response
for peer in closest_peers:
# Skip if the peer is the requester
if peer == peer_id:
continue
# Add peer to closerPeers field
peer_proto = response.closerPeers.add()
peer_proto.id = peer.to_bytes()
peer_proto.connection = Message.ConnectionType.CAN_CONNECT
# Add addresses if available
try:
addrs = self.host.get_peerstore().addrs(peer)
if addrs:
for addr in addrs:
peer_proto.addrs.append(addr.to_bytes())
except Exception:
pass
# Serialize and send response
response_bytes = response.SerializeToString()
await stream.write(varint.encode(len(response_bytes)))
await stream.write(response_bytes)
logger.debug(
f"Sent FIND_NODE response with{len(response.closerPeers)} peers"
)
# Handle ADD_PROVIDER message
elif message.type == Message.MessageType.ADD_PROVIDER:
# Process ADD_PROVIDER
key = message.key
logger.debug(f"Received ADD_PROVIDER for key {key.hex()}")
# Extract provider information
for provider_proto in message.providerPeers:
try:
# Validate that the provider is the sender
provider_id = ID(provider_proto.id)
if provider_id != peer_id:
logger.warning(
f"Provider ID {provider_id} doesn't"
f"match sender {peer_id}, ignoring"
)
continue
# Convert addresses to Multiaddr
addrs = []
for addr_bytes in provider_proto.addrs:
try:
addrs.append(Multiaddr(addr_bytes))
except Exception as e:
logger.warning(f"Failed to parse address: {e}")
# Add to provider store
provider_info = PeerInfo(provider_id, addrs)
self.provider_store.add_provider(key, provider_info)
logger.debug(
f"Added provider {provider_id} for key {key.hex()}"
)
except Exception as e:
logger.warning(f"Failed to process provider info: {e}")
# Send acknowledgement
response = Message()
response.type = Message.MessageType.ADD_PROVIDER
response.key = key
response_bytes = response.SerializeToString()
await stream.write(varint.encode(len(response_bytes)))
await stream.write(response_bytes)
logger.debug("Sent ADD_PROVIDER acknowledgement")
# Handle GET_PROVIDERS message
elif message.type == Message.MessageType.GET_PROVIDERS:
# Process GET_PROVIDERS
key = message.key
logger.debug(f"Received GET_PROVIDERS request for key {key.hex()}")
# Find providers for the key
providers = self.provider_store.get_providers(key)
logger.debug(
f"Found {len(providers)} providers for key {key.hex()}"
)
# Create response
response = Message()
response.type = Message.MessageType.GET_PROVIDERS
response.key = key
# Add provider information to response
for provider_info in providers:
provider_proto = response.providerPeers.add()
provider_proto.id = provider_info.peer_id.to_bytes()
provider_proto.connection = Message.ConnectionType.CAN_CONNECT
# Add addresses if available
for addr in provider_info.addrs:
provider_proto.addrs.append(addr.to_bytes())
# Also include closest peers if we don't have providers
if not providers:
closest_peers = self.routing_table.find_local_closest_peers(
key, 20
)
logger.debug(
f"No providers found, including {len(closest_peers)}"
"closest peers"
)
for peer in closest_peers:
# Skip if peer is the requester
if peer == peer_id:
continue
peer_proto = response.closerPeers.add()
peer_proto.id = peer.to_bytes()
peer_proto.connection = Message.ConnectionType.CAN_CONNECT
# Add addresses if available
try:
addrs = self.host.get_peerstore().addrs(peer)
for addr in addrs:
peer_proto.addrs.append(addr.to_bytes())
except Exception:
pass
# Serialize and send response
response_bytes = response.SerializeToString()
await stream.write(varint.encode(len(response_bytes)))
await stream.write(response_bytes)
logger.debug("Sent GET_PROVIDERS response")
# Handle GET_VALUE message
elif message.type == Message.MessageType.GET_VALUE:
# Process GET_VALUE
key = message.key
logger.debug(f"Received GET_VALUE request for key {key.hex()}")
value = self.value_store.get(key)
if value:
logger.debug(f"Found value for key {key.hex()}")
# Create response using protobuf
response = Message()
response.type = Message.MessageType.GET_VALUE
# Create record
response.key = key
response.record.key = key
response.record.value = value
response.record.timeReceived = str(time.time())
# Serialize and send response
response_bytes = response.SerializeToString()
await stream.write(varint.encode(len(response_bytes)))
await stream.write(response_bytes)
logger.debug("Sent GET_VALUE response")
else:
logger.debug(f"No value found for key {key.hex()}")
# Create response with closest peers when no value is found
response = Message()
response.type = Message.MessageType.GET_VALUE
response.key = key
# Add closest peers to key
closest_peers = self.routing_table.find_local_closest_peers(
key, 20
)
logger.debug(
"No value found,"
f"including {len(closest_peers)} closest peers"
)
for peer in closest_peers:
# Skip if peer is the requester
if peer == peer_id:
continue
peer_proto = response.closerPeers.add()
peer_proto.id = peer.to_bytes()
peer_proto.connection = Message.ConnectionType.CAN_CONNECT
# Add addresses if available
try:
addrs = self.host.get_peerstore().addrs(peer)
for addr in addrs:
peer_proto.addrs.append(addr.to_bytes())
except Exception:
pass
# Serialize and send response
response_bytes = response.SerializeToString()
await stream.write(varint.encode(len(response_bytes)))
await stream.write(response_bytes)
logger.debug("Sent GET_VALUE response with closest peers")
# Handle PUT_VALUE message
elif message.type == Message.MessageType.PUT_VALUE and message.HasField(
"record"
):
# Process PUT_VALUE
key = message.record.key
value = message.record.value
success = False
try:
if not (key and value):
raise ValueError(
"Missing key or value in PUT_VALUE message"
)
self.value_store.put(key, value)
logger.debug(f"Stored value {value.hex()} for key {key.hex()}")
success = True
except Exception as e:
logger.warning(
f"Failed to store value {value.hex()} for key "
f"{key.hex()}: {e}"
)
finally:
# Send acknowledgement
response = Message()
response.type = Message.MessageType.PUT_VALUE
if success:
response.key = key
response_bytes = response.SerializeToString()
await stream.write(varint.encode(len(response_bytes)))
await stream.write(response_bytes)
logger.debug("Sent PUT_VALUE acknowledgement")
except Exception as proto_err:
logger.warning(f"Failed to parse protobuf message: {proto_err}")
await stream.close()
except Exception as e:
logger.error(f"Error handling DHT stream: {e}")
await stream.close()
async def refresh_routing_table(self) -> None:
"""Refresh the routing table."""
logger.debug("Refreshing routing table")
await self.peer_routing.refresh_routing_table()
# Peer routing methods
async def find_peer(self, peer_id: ID) -> PeerInfo | None:
"""
Find a peer with the given ID.
"""
logger.debug(f"Finding peer: {peer_id}")
return await self.peer_routing.find_peer(peer_id)
# Value storage and retrieval methods
async def put_value(self, key: bytes, value: bytes) -> None:
"""
Store a value in the DHT.
"""
logger.debug(f"Storing value for key {key.hex()}")
# 1. Store locally first
self.value_store.put(key, value)
try:
decoded_value = value.decode("utf-8")
except UnicodeDecodeError:
decoded_value = value.hex()
logger.debug(
f"Stored value locally for key {key.hex()} with value {decoded_value}"
)
# 2. Get closest peers, excluding self
closest_peers = [
peer
for peer in self.routing_table.find_local_closest_peers(key)
if peer != self.local_peer_id
]
logger.debug(f"Found {len(closest_peers)} peers to store value at")
# 3. Store at remote peers in batches of ALPHA, in parallel
stored_count = 0
for i in range(0, len(closest_peers), ALPHA):
batch = closest_peers[i : i + ALPHA]
batch_results = [False] * len(batch)
async def store_one(idx: int, peer: ID) -> None:
try:
with trio.move_on_after(QUERY_TIMEOUT):
success = await self.value_store._store_at_peer(
peer, key, value
)
batch_results[idx] = success
if success:
logger.debug(f"Stored value at peer {peer}")
else:
logger.debug(f"Failed to store value at peer {peer}")
except Exception as e:
logger.debug(f"Error storing value at peer {peer}: {e}")
async with trio.open_nursery() as nursery:
for idx, peer in enumerate(batch):
nursery.start_soon(store_one, idx, peer)
stored_count += sum(batch_results)
logger.info(f"Successfully stored value at {stored_count} peers")
async def get_value(self, key: bytes) -> bytes | None:
logger.debug(f"Getting value for key: {key.hex()}")
# 1. Check local store first
value = self.value_store.get(key)
if value:
logger.debug("Found value locally")
return value
# 2. Get closest peers, excluding self
closest_peers = [
peer
for peer in self.routing_table.find_local_closest_peers(key)
if peer != self.local_peer_id
]
logger.debug(f"Searching {len(closest_peers)} peers for value")
# 3. Query ALPHA peers at a time in parallel
for i in range(0, len(closest_peers), ALPHA):
batch = closest_peers[i : i + ALPHA]
found_value = None
async def query_one(peer: ID) -> None:
nonlocal found_value
try:
with trio.move_on_after(QUERY_TIMEOUT):
value = await self.value_store._get_from_peer(peer, key)
if value is not None and found_value is None:
found_value = value
logger.debug(f"Found value at peer {peer}")
except Exception as e:
logger.debug(f"Error querying peer {peer}: {e}")
async with trio.open_nursery() as nursery:
for peer in batch:
nursery.start_soon(query_one, peer)
if found_value is not None:
self.value_store.put(key, found_value)
logger.info("Successfully retrieved value from network")
return found_value
# 4. Not found
logger.warning(f"Value not found for key {key.hex()}")
return None
# Add these methods in the Utility methods section
# Utility methods
async def add_peer(self, peer_id: ID) -> bool:
"""
Add a peer to the routing table.
params: peer_id: The peer ID to add.
Returns
-------
bool
True if peer was added or updated, False otherwise.
"""
return await self.routing_table.add_peer(peer_id)
async def provide(self, key: bytes) -> bool:
"""
Reference to provider_store.provide for convenience.
"""
return await self.provider_store.provide(key)
async def find_providers(self, key: bytes, count: int = 20) -> list[PeerInfo]:
"""
Reference to provider_store.find_providers for convenience.
"""
return await self.provider_store.find_providers(key, count)
def get_routing_table_size(self) -> int:
"""
Get the number of peers in the routing table.
Returns
-------
int
Number of peers.
"""
return self.routing_table.size()
def get_value_store_size(self) -> int:
"""
Get the number of items in the value store.
Returns
-------
int
Number of items.
"""
return self.value_store.size()

0
libp2p/kad_dht/pb/__init__.py Normal file
View File

38
libp2p/kad_dht/pb/kademlia.proto Normal file
View File

@ -0,0 +1,38 @@
syntax = "proto3";
message Record {
bytes key = 1;
bytes value = 2;
string timeReceived = 5;
};
message Message {
enum MessageType {
PUT_VALUE = 0;
GET_VALUE = 1;
ADD_PROVIDER = 2;
GET_PROVIDERS = 3;
FIND_NODE = 4;
PING = 5;
}
enum ConnectionType {
NOT_CONNECTED = 0;
CONNECTED = 1;
CAN_CONNECT = 2;
CANNOT_CONNECT = 3;
}
message Peer {
bytes id = 1;
repeated bytes addrs = 2;
ConnectionType connection = 3;
}
MessageType type = 1;
int32 clusterLevelRaw = 10;
bytes key = 2;
Record record = 3;
repeated Peer closerPeers = 8;
repeated Peer providerPeers = 9;
}

33
libp2p/kad_dht/pb/kademlia_pb2.py Normal file
View File

@ -0,0 +1,33 @@
# -*- coding: utf-8 -*-
# Generated by the protocol buffer compiler. DO NOT EDIT!
# source: libp2p/kad_dht/pb/kademlia.proto
"""Generated protocol buffer code."""
from google.protobuf import descriptor as _descriptor
from google.protobuf import descriptor_pool as _descriptor_pool
from google.protobuf import symbol_database as _symbol_database
from google.protobuf.internal import builder as _builder
# @@protoc_insertion_point(imports)
_sym_db = _symbol_database.Default()
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n libp2p/kad_dht/pb/kademlia.proto\":\n\x06Record\x12\x0b\n\x03key\x18\x01 \x01(\x0c\x12\r\n\x05value\x18\x02 \x01(\x0c\x12\x14\n\x0ctimeReceived\x18\x05 \x01(\t\"\xca\x03\n\x07Message\x12\"\n\x04type\x18\x01 \x01(\x0e\x32\x14.Message.MessageType\x12\x17\n\x0f\x63lusterLevelRaw\x18\n \x01(\x05\x12\x0b\n\x03key\x18\x02 \x01(\x0c\x12\x17\n\x06record\x18\x03 \x01(\x0b\x32\x07.Record\x12\"\n\x0b\x63loserPeers\x18\x08 \x03(\x0b\x32\r.Message.Peer\x12$\n\rproviderPeers\x18\t \x03(\x0b\x32\r.Message.Peer\x1aN\n\x04Peer\x12\n\n\x02id\x18\x01 \x01(\x0c\x12\r\n\x05\x61\x64\x64rs\x18\x02 \x03(\x0c\x12+\n\nconnection\x18\x03 \x01(\x0e\x32\x17.Message.ConnectionType\"i\n\x0bMessageType\x12\r\n\tPUT_VALUE\x10\x00\x12\r\n\tGET_VALUE\x10\x01\x12\x10\n\x0c\x41\x44\x44_PROVIDER\x10\x02\x12\x11\n\rGET_PROVIDERS\x10\x03\x12\r\n\tFIND_NODE\x10\x04\x12\x08\n\x04PING\x10\x05\"W\n\x0e\x43onnectionType\x12\x11\n\rNOT_CONNECTED\x10\x00\x12\r\n\tCONNECTED\x10\x01\x12\x0f\n\x0b\x43\x41N_CONNECT\x10\x02\x12\x12\n\x0e\x43\x41NNOT_CONNECT\x10\x03\x62\x06proto3')
_globals = globals()
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'libp2p.kad_dht.pb.kademlia_pb2', _globals)
if _descriptor._USE_C_DESCRIPTORS == False:
DESCRIPTOR._options = None
_globals['_RECORD']._serialized_start=36
_globals['_RECORD']._serialized_end=94
_globals['_MESSAGE']._serialized_start=97
_globals['_MESSAGE']._serialized_end=555
_globals['_MESSAGE_PEER']._serialized_start=281
_globals['_MESSAGE_PEER']._serialized_end=359
_globals['_MESSAGE_MESSAGETYPE']._serialized_start=361
_globals['_MESSAGE_MESSAGETYPE']._serialized_end=466
_globals['_MESSAGE_CONNECTIONTYPE']._serialized_start=468
_globals['_MESSAGE_CONNECTIONTYPE']._serialized_end=555
# @@protoc_insertion_point(module_scope)

133
libp2p/kad_dht/pb/kademlia_pb2.pyi Normal file
View File

@ -0,0 +1,133 @@
"""
@generated by mypy-protobuf. Do not edit manually!
isort:skip_file
"""
import builtins
import collections.abc
import google.protobuf.descriptor
import google.protobuf.internal.containers
import google.protobuf.internal.enum_type_wrapper
import google.protobuf.message
import sys
import typing
if sys.version_info >= (3, 10):
import typing as typing_extensions
else:
import typing_extensions
DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
@typing.final
class Record(google.protobuf.message.Message):
DESCRIPTOR: google.protobuf.descriptor.Descriptor
KEY_FIELD_NUMBER: builtins.int
VALUE_FIELD_NUMBER: builtins.int
TIMERECEIVED_FIELD_NUMBER: builtins.int
key: builtins.bytes
value: builtins.bytes
timeReceived: builtins.str
def __init__(
self,
*,
key: builtins.bytes = ...,
value: builtins.bytes = ...,
timeReceived: builtins.str = ...,
) -> None: ...
def ClearField(self, field_name: typing.Literal["key", b"key", "timeReceived", b"timeReceived", "value", b"value"]) -> None: ...
global___Record = Record
@typing.final
class Message(google.protobuf.message.Message):
DESCRIPTOR: google.protobuf.descriptor.Descriptor
class _MessageType:
ValueType = typing.NewType("ValueType", builtins.int)
V: typing_extensions.TypeAlias = ValueType
class _MessageTypeEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[Message._MessageType.ValueType], builtins.type):
DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
PUT_VALUE: Message._MessageType.ValueType # 0
GET_VALUE: Message._MessageType.ValueType # 1
ADD_PROVIDER: Message._MessageType.ValueType # 2
GET_PROVIDERS: Message._MessageType.ValueType # 3
FIND_NODE: Message._MessageType.ValueType # 4
PING: Message._MessageType.ValueType # 5
class MessageType(_MessageType, metaclass=_MessageTypeEnumTypeWrapper): ...
PUT_VALUE: Message.MessageType.ValueType # 0
GET_VALUE: Message.MessageType.ValueType # 1
ADD_PROVIDER: Message.MessageType.ValueType # 2
GET_PROVIDERS: Message.MessageType.ValueType # 3
FIND_NODE: Message.MessageType.ValueType # 4
PING: Message.MessageType.ValueType # 5
class _ConnectionType:
ValueType = typing.NewType("ValueType", builtins.int)
V: typing_extensions.TypeAlias = ValueType
class _ConnectionTypeEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[Message._ConnectionType.ValueType], builtins.type):
DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
NOT_CONNECTED: Message._ConnectionType.ValueType # 0
CONNECTED: Message._ConnectionType.ValueType # 1
CAN_CONNECT: Message._ConnectionType.ValueType # 2
CANNOT_CONNECT: Message._ConnectionType.ValueType # 3
class ConnectionType(_ConnectionType, metaclass=_ConnectionTypeEnumTypeWrapper): ...
NOT_CONNECTED: Message.ConnectionType.ValueType # 0
CONNECTED: Message.ConnectionType.ValueType # 1
CAN_CONNECT: Message.ConnectionType.ValueType # 2
CANNOT_CONNECT: Message.ConnectionType.ValueType # 3
@typing.final
class Peer(google.protobuf.message.Message):
DESCRIPTOR: google.protobuf.descriptor.Descriptor
ID_FIELD_NUMBER: builtins.int
ADDRS_FIELD_NUMBER: builtins.int
CONNECTION_FIELD_NUMBER: builtins.int
id: builtins.bytes
connection: global___Message.ConnectionType.ValueType
@property
def addrs(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.bytes]: ...
def __init__(
self,
*,
id: builtins.bytes = ...,
addrs: collections.abc.Iterable[builtins.bytes] | None = ...,
connection: global___Message.ConnectionType.ValueType = ...,
) -> None: ...
def ClearField(self, field_name: typing.Literal["addrs", b"addrs", "connection", b"connection", "id", b"id"]) -> None: ...
TYPE_FIELD_NUMBER: builtins.int
CLUSTERLEVELRAW_FIELD_NUMBER: builtins.int
KEY_FIELD_NUMBER: builtins.int
RECORD_FIELD_NUMBER: builtins.int
CLOSERPEERS_FIELD_NUMBER: builtins.int
PROVIDERPEERS_FIELD_NUMBER: builtins.int
type: global___Message.MessageType.ValueType
clusterLevelRaw: builtins.int
key: builtins.bytes
@property
def record(self) -> global___Record: ...
@property
def closerPeers(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Message.Peer]: ...
@property
def providerPeers(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Message.Peer]: ...
def __init__(
self,
*,
type: global___Message.MessageType.ValueType = ...,
clusterLevelRaw: builtins.int = ...,
key: builtins.bytes = ...,
record: global___Record | None = ...,
closerPeers: collections.abc.Iterable[global___Message.Peer] | None = ...,
providerPeers: collections.abc.Iterable[global___Message.Peer] | None = ...,
) -> None: ...
def HasField(self, field_name: typing.Literal["record", b"record"]) -> builtins.bool: ...
def ClearField(self, field_name: typing.Literal["closerPeers", b"closerPeers", "clusterLevelRaw", b"clusterLevelRaw", "key", b"key", "providerPeers", b"providerPeers", "record", b"record", "type", b"type"]) -> None: ...
global___Message = Message

418
libp2p/kad_dht/peer_routing.py Normal file
View File

@ -0,0 +1,418 @@
"""
Peer routing implementation for Kademlia DHT.
This module implements the peer routing interface using Kademlia's algorithm
to efficiently locate peers in a distributed network.
"""
import logging
import trio
import varint
from libp2p.abc import (
IHost,
INetStream,
IPeerRouting,
)
from libp2p.custom_types import (
TProtocol,
)
from libp2p.peer.id import (
ID,
)
from libp2p.peer.peerinfo import (
PeerInfo,
)
from .pb.kademlia_pb2 import (
Message,
)
from .routing_table import (
RoutingTable,
)
from .utils import (
sort_peer_ids_by_distance,
)
# logger = logging.getLogger("libp2p.kademlia.peer_routing")
logger = logging.getLogger("kademlia-example.peer_routing")
# Constants for the Kademlia algorithm
ALPHA = 3 # Concurrency parameter
MAX_PEER_LOOKUP_ROUNDS = 20 # Maximum number of rounds in peer lookup
PROTOCOL_ID = TProtocol("/ipfs/kad/1.0.0")
class PeerRouting(IPeerRouting):
"""
Implementation of peer routing using the Kademlia algorithm.
This class provides methods to find peers in the DHT network
and helps maintain the routing table.
"""
def __init__(self, host: IHost, routing_table: RoutingTable):
"""
Initialize the peer routing service.
:param host: The libp2p host
:param routing_table: The Kademlia routing table
"""
self.host = host
self.routing_table = routing_table
self.protocol_id = PROTOCOL_ID
async def find_peer(self, peer_id: ID) -> PeerInfo | None:
"""
Find a peer with the given ID.
:param peer_id: The ID of the peer to find
Returns
-------
Optional[PeerInfo]
The peer information if found, None otherwise
"""
# Check if this is actually our peer ID
if peer_id == self.host.get_id():
try:
# Return our own peer info
return PeerInfo(peer_id, self.host.get_addrs())
except Exception:
logger.exception("Error getting our own peer info")
return None
# First check if the peer is in our routing table
peer_info = self.routing_table.get_peer_info(peer_id)
if peer_info:
logger.debug(f"Found peer {peer_id} in routing table")
return peer_info
# Then check if the peer is in our peerstore
try:
addrs = self.host.get_peerstore().addrs(peer_id)
if addrs:
logger.debug(f"Found peer {peer_id} in peerstore")
return PeerInfo(peer_id, addrs)
except Exception:
pass
# If not found locally, search the network
try:
closest_peers = await self.find_closest_peers_network(peer_id.to_bytes())
logger.info(f"Closest peers found: {closest_peers}")
# Check if we found the peer we're looking for
for found_peer in closest_peers:
if found_peer == peer_id:
try:
addrs = self.host.get_peerstore().addrs(found_peer)
if addrs:
return PeerInfo(found_peer, addrs)
except Exception:
pass
except Exception as e:
logger.error(f"Error searching for peer {peer_id}: {e}")
# Not found
logger.info(f"Peer {peer_id} not found")
return None
async def _query_single_peer_for_closest(
self, peer: ID, target_key: bytes, new_peers: list[ID]
) -> None:
"""
Query a single peer for closest peers and append results to the shared list.
params: peer : ID
The peer to query
params: target_key : bytes
The target key to find closest peers for
params: new_peers : list[ID]
Shared list to append results to
"""
try:
result = await self._query_peer_for_closest(peer, target_key)
# Add deduplication to prevent duplicate peers
for peer_id in result:
if peer_id not in new_peers:
new_peers.append(peer_id)
logger.debug(
"Queried peer %s for closest peers, got %d results (%d unique)",
peer,
len(result),
len([p for p in result if p not in new_peers[: -len(result)]]),
)
except Exception as e:
logger.debug(f"Query to peer {peer} failed: {e}")
async def find_closest_peers_network(
self, target_key: bytes, count: int = 20
) -> list[ID]:
"""
Find the closest peers to a target key in the entire network.
Performs an iterative lookup by querying peers for their closest peers.
Returns
-------
list[ID]
Closest peer IDs
"""
# Start with closest peers from our routing table
closest_peers = self.routing_table.find_local_closest_peers(target_key, count)
logger.debug("Local closest peers: %d found", len(closest_peers))
queried_peers: set[ID] = set()
rounds = 0
# Return early if we have no peers to start with
if not closest_peers:
logger.warning("No local peers available for network lookup")
return []
# Iterative lookup until convergence
while rounds < MAX_PEER_LOOKUP_ROUNDS:
rounds += 1
logger.debug(f"Lookup round {rounds}/{MAX_PEER_LOOKUP_ROUNDS}")
# Find peers we haven't queried yet
peers_to_query = [p for p in closest_peers if p not in queried_peers]
if not peers_to_query:
logger.debug("No more unqueried peers available, ending lookup")
break # No more peers to query
# Query these peers for their closest peers to target
peers_batch = peers_to_query[:ALPHA] # Limit to ALPHA peers at a time
# Mark these peers as queried before we actually query them
for peer in peers_batch:
queried_peers.add(peer)
# Run queries in parallel for this batch using trio nursery
new_peers: list[ID] = [] # Shared array to collect all results
async with trio.open_nursery() as nursery:
for peer in peers_batch:
nursery.start_soon(
self._query_single_peer_for_closest, peer, target_key, new_peers
)
# If we got no new peers, we're done
if not new_peers:
logger.debug("No new peers discovered in this round, ending lookup")
break
# Update our list of closest peers
all_candidates = closest_peers + new_peers
old_closest_peers = closest_peers[:]
closest_peers = sort_peer_ids_by_distance(target_key, all_candidates)[
:count
]
logger.debug(f"Updated closest peers count: {len(closest_peers)}")
# Check if we made any progress (found closer peers)
if closest_peers == old_closest_peers:
logger.debug("No improvement in closest peers, ending lookup")
break
logger.info(
f"Network lookup completed after {rounds} rounds, "
f"found {len(closest_peers)} peers"
)
return closest_peers
async def _query_peer_for_closest(self, peer: ID, target_key: bytes) -> list[ID]:
"""
Query a peer for their closest peers
to the target key using varint length prefix
"""
stream = None
results = []
try:
# Add the peer to our routing table regardless of query outcome
try:
addrs = self.host.get_peerstore().addrs(peer)
if addrs:
peer_info = PeerInfo(peer, addrs)
await self.routing_table.add_peer(peer_info)
except Exception as e:
logger.debug(f"Failed to add peer {peer} to routing table: {e}")
# Open a stream to the peer using the Kademlia protocol
logger.debug(f"Opening stream to {peer} for closest peers query")
try:
stream = await self.host.new_stream(peer, [self.protocol_id])
logger.debug(f"Stream opened to {peer}")
except Exception as e:
logger.warning(f"Failed to open stream to {peer}: {e}")
return []
# Create and send FIND_NODE request using protobuf
find_node_msg = Message()
find_node_msg.type = Message.MessageType.FIND_NODE
find_node_msg.key = target_key # Set target key directly as bytes
# Serialize and send the protobuf message with varint length prefix
proto_bytes = find_node_msg.SerializeToString()
logger.debug(
f"Sending FIND_NODE: {proto_bytes.hex()} (len={len(proto_bytes)})"
)
await stream.write(varint.encode(len(proto_bytes)))
await stream.write(proto_bytes)
# Read varint-prefixed response length
length_bytes = b""
while True:
b = await stream.read(1)
if not b:
logger.warning(
"Error reading varint length from stream: connection closed"
)
return []
length_bytes += b
if b[0] & 0x80 == 0:
break
response_length = varint.decode_bytes(length_bytes)
# Read response data
response_bytes = b""
remaining = response_length
while remaining > 0:
chunk = await stream.read(remaining)
if not chunk:
logger.debug(f"Connection closed by peer {peer} while reading data")
return []
response_bytes += chunk
remaining -= len(chunk)
# Parse the protobuf response
response_msg = Message()
response_msg.ParseFromString(response_bytes)
logger.debug(
"Received response from %s with %d peers",
peer,
len(response_msg.closerPeers),
)
# Process closest peers from response
if response_msg.type == Message.MessageType.FIND_NODE:
for peer_data in response_msg.closerPeers:
new_peer_id = ID(peer_data.id)
if new_peer_id not in results:
results.append(new_peer_id)
if peer_data.addrs:
from multiaddr import (
Multiaddr,
)
addrs = [Multiaddr(addr) for addr in peer_data.addrs]
self.host.get_peerstore().add_addrs(new_peer_id, addrs, 3600)
except Exception as e:
logger.debug(f"Error querying peer {peer} for closest: {e}")
finally:
if stream:
await stream.close()
return results
async def _handle_kad_stream(self, stream: INetStream) -> None:
"""
Handle incoming Kademlia protocol streams.
params: stream: The incoming stream
Returns
-------
None
"""
try:
# Read message length
length_bytes = await stream.read(4)
if not length_bytes:
return
message_length = int.from_bytes(length_bytes, byteorder="big")
# Read message
message_bytes = await stream.read(message_length)
if not message_bytes:
return
# Parse protobuf message
kad_message = Message()
try:
kad_message.ParseFromString(message_bytes)
if kad_message.type == Message.MessageType.FIND_NODE:
# Get target key directly from protobuf message
target_key = kad_message.key
# Find closest peers to target
closest_peers = self.routing_table.find_local_closest_peers(
target_key, 20
)
# Create protobuf response
response = Message()
response.type = Message.MessageType.FIND_NODE
# Add peer information to response
for peer_id in closest_peers:
peer_proto = response.closerPeers.add()
peer_proto.id = peer_id.to_bytes()
peer_proto.connection = Message.ConnectionType.CAN_CONNECT
# Add addresses if available
try:
addrs = self.host.get_peerstore().addrs(peer_id)
if addrs:
for addr in addrs:
peer_proto.addrs.append(addr.to_bytes())
except Exception:
pass
# Send response
response_bytes = response.SerializeToString()
await stream.write(len(response_bytes).to_bytes(4, byteorder="big"))
await stream.write(response_bytes)
except Exception as parse_err:
logger.error(f"Failed to parse protocol buffer message: {parse_err}")
except Exception as e:
logger.debug(f"Error handling Kademlia stream: {e}")
finally:
await stream.close()
async def refresh_routing_table(self) -> None:
"""
Refresh the routing table by performing lookups for random keys.
Returns
-------
None
"""
logger.info("Refreshing routing table")
# Perform a lookup for ourselves to populate the routing table
local_id = self.host.get_id()
closest_peers = await self.find_closest_peers_network(local_id.to_bytes())
# Add discovered peers to routing table
for peer_id in closest_peers:
try:
addrs = self.host.get_peerstore().addrs(peer_id)
if addrs:
peer_info = PeerInfo(peer_id, addrs)
await self.routing_table.add_peer(peer_info)
except Exception as e:
logger.debug(f"Failed to add discovered peer {peer_id}: {e}")

575
libp2p/kad_dht/provider_store.py Normal file
View File

@ -0,0 +1,575 @@
"""
Provider record storage for Kademlia DHT.
This module implements the storage for content provider records in the Kademlia DHT.
"""
import logging
import time
from typing import (
Any,
)
from multiaddr import (
Multiaddr,
)
import trio
import varint
from libp2p.abc import (
IHost,
)
from libp2p.custom_types import (
TProtocol,
)
from libp2p.peer.id import (
ID,
)
from libp2p.peer.peerinfo import (
PeerInfo,
)
from .pb.kademlia_pb2 import (
Message,
)
# logger = logging.getLogger("libp2p.kademlia.provider_store")
logger = logging.getLogger("kademlia-example.provider_store")
# Constants for provider records (based on IPFS standards)
PROVIDER_RECORD_REPUBLISH_INTERVAL = 22 * 60 * 60 # 22 hours in seconds
PROVIDER_RECORD_EXPIRATION_INTERVAL = 48 * 60 * 60 # 48 hours in seconds
PROVIDER_ADDRESS_TTL = 30 * 60 # 30 minutes in seconds
PROTOCOL_ID = TProtocol("/ipfs/kad/1.0.0")
ALPHA = 3 # Number of parallel queries/advertisements
QUERY_TIMEOUT = 10 # Timeout for each query in seconds
class ProviderRecord:
"""
A record for a content provider in the DHT.
Contains the peer information and timestamp.
"""
def __init__(
self,
provider_info: PeerInfo,
timestamp: float | None = None,
) -> None:
"""
Initialize a new provider record.
:param provider_info: The provider's peer information
:param timestamp: Time this record was created/updated
(defaults to current time)
"""
self.provider_info = provider_info
self.timestamp = timestamp or time.time()
def is_expired(self) -> bool:
"""
Check if this provider record has expired.
Returns
-------
bool
True if the record has expired
"""
current_time = time.time()
return (current_time - self.timestamp) >= PROVIDER_RECORD_EXPIRATION_INTERVAL
def should_republish(self) -> bool:
"""
Check if this provider record should be republished.
Returns
-------
bool
True if the record should be republished
"""
current_time = time.time()
return (current_time - self.timestamp) >= PROVIDER_RECORD_REPUBLISH_INTERVAL
@property
def peer_id(self) -> ID:
"""Get the provider's peer ID."""
return self.provider_info.peer_id
@property
def addresses(self) -> list[Multiaddr]:
"""Get the provider's addresses."""
return self.provider_info.addrs
class ProviderStore:
"""
Store for content provider records in the Kademlia DHT.
Maps content keys to provider records, with support for expiration.
"""
def __init__(self, host: IHost, peer_routing: Any = None) -> None:
"""
Initialize a new provider store.
:param host: The libp2p host instance (optional)
:param peer_routing: The peer routing instance (optional)
"""
# Maps content keys to a dict of provider records (peer_id -> record)
self.providers: dict[bytes, dict[str, ProviderRecord]] = {}
self.host = host
self.peer_routing = peer_routing
self.providing_keys: set[bytes] = set()
self.local_peer_id = host.get_id()
async def _republish_provider_records(self) -> None:
"""Republish all provider records for content this node is providing."""
# First, republish keys we're actively providing
for key in self.providing_keys:
logger.debug(f"Republishing provider record for key {key.hex()}")
await self.provide(key)
# Also check for any records that should be republished
time.time()
for key, providers in self.providers.items():
for peer_id_str, record in providers.items():
# Only republish records for our own peer
if self.local_peer_id and str(self.local_peer_id) == peer_id_str:
if record.should_republish():
logger.debug(
f"Republishing old provider record for key {key.hex()}"
)
await self.provide(key)
async def provide(self, key: bytes) -> bool:
"""
Advertise that this node can provide a piece of content.
Finds the k closest peers to the key and sends them ADD_PROVIDER messages.
:param key: The content key (multihash) to advertise
Returns
-------
bool
True if the advertisement was successful
"""
if not self.host or not self.peer_routing:
logger.error("Host or peer_routing not initialized, cannot provide content")
return False
# Add to local provider store
local_addrs = []
for addr in self.host.get_addrs():
local_addrs.append(addr)
local_peer_info = PeerInfo(self.host.get_id(), local_addrs)
self.add_provider(key, local_peer_info)
# Track that we're providing this key
self.providing_keys.add(key)
# Find the k closest peers to the key
closest_peers = await self.peer_routing.find_closest_peers_network(key)
logger.debug(
"Found %d peers close to key %s for provider advertisement",
len(closest_peers),
key.hex(),
)
# Send ADD_PROVIDER messages to these ALPHA peers in parallel.
success_count = 0
for i in range(0, len(closest_peers), ALPHA):
batch = closest_peers[i : i + ALPHA]
results: list[bool] = [False] * len(batch)
async def send_one(
idx: int, peer_id: ID, results: list[bool] = results
) -> None:
if peer_id == self.local_peer_id:
return
try:
with trio.move_on_after(QUERY_TIMEOUT):
success = await self._send_add_provider(peer_id, key)
results[idx] = success
if not success:
logger.warning(f"Failed to send ADD_PROVIDER to {peer_id}")
except Exception as e:
logger.warning(f"Error sending ADD_PROVIDER to {peer_id}: {e}")
async with trio.open_nursery() as nursery:
for idx, peer_id in enumerate(batch):
nursery.start_soon(send_one, idx, peer_id, results)
success_count += sum(results)
logger.info(f"Successfully advertised to {success_count} peers")
return success_count > 0
async def _send_add_provider(self, peer_id: ID, key: bytes) -> bool:
"""
Send ADD_PROVIDER message to a specific peer.
:param peer_id: The peer to send the message to
:param key: The content key being provided
Returns
-------
bool
True if the message was successfully sent and acknowledged
"""
try:
result = False
# Open a stream to the peer
stream = await self.host.new_stream(peer_id, [TProtocol(PROTOCOL_ID)])
# Get our addresses to include in the message
addrs = []
for addr in self.host.get_addrs():
addrs.append(addr.to_bytes())
# Create the ADD_PROVIDER message
message = Message()
message.type = Message.MessageType.ADD_PROVIDER
message.key = key
# Add our provider info
provider = message.providerPeers.add()
provider.id = self.local_peer_id.to_bytes()
provider.addrs.extend(addrs)
# Serialize and send the message
proto_bytes = message.SerializeToString()
await stream.write(varint.encode(len(proto_bytes)))
await stream.write(proto_bytes)
logger.debug(f"Sent ADD_PROVIDER to {peer_id} for key {key.hex()}")
# Read response length prefix
length_bytes = b""
while True:
logger.debug("Reading response length prefix in add provider")
b = await stream.read(1)
if not b:
return False
length_bytes += b
if b[0] & 0x80 == 0:
break
response_length = varint.decode_bytes(length_bytes)
# Read response data
response_bytes = b""
remaining = response_length
while remaining > 0:
chunk = await stream.read(remaining)
if not chunk:
return False
response_bytes += chunk
remaining -= len(chunk)
# Parse response
response = Message()
response.ParseFromString(response_bytes)
# Check response type
response.type == Message.MessageType.ADD_PROVIDER
if response.type:
result = True
except Exception as e:
logger.warning(f"Error sending ADD_PROVIDER to {peer_id}: {e}")
finally:
await stream.close()
return result
async def find_providers(self, key: bytes, count: int = 20) -> list[PeerInfo]:
"""
Find content providers for a given key.
:param key: The content key to look for
:param count: Maximum number of providers to return
Returns
-------
List[PeerInfo]
List of content providers
"""
if not self.host or not self.peer_routing:
logger.error("Host or peer_routing not initialized, cannot find providers")
return []
# Check local provider store first
local_providers = self.get_providers(key)
if local_providers:
logger.debug(
f"Found {len(local_providers)} providers locally for {key.hex()}"
)
return local_providers[:count]
logger.debug("local providers are %s", local_providers)
# Find the closest peers to the key
closest_peers = await self.peer_routing.find_closest_peers_network(key)
logger.debug(
f"Searching {len(closest_peers)} peers for providers of {key.hex()}"
)
# Query these peers for providers in batches of ALPHA, in parallel, with timeout
all_providers = []
for i in range(0, len(closest_peers), ALPHA):
batch = closest_peers[i : i + ALPHA]
batch_results: list[list[PeerInfo]] = [[] for _ in batch]
async def get_one(
idx: int,
peer_id: ID,
batch_results: list[list[PeerInfo]] = batch_results,
) -> None:
if peer_id == self.local_peer_id:
return
try:
with trio.move_on_after(QUERY_TIMEOUT):
providers = await self._get_providers_from_peer(peer_id, key)
if providers:
for provider in providers:
self.add_provider(key, provider)
batch_results[idx] = providers
else:
logger.debug(f"No providers found at peer {peer_id}")
except Exception as e:
logger.warning(f"Failed to get providers from {peer_id}: {e}")
async with trio.open_nursery() as nursery:
for idx, peer_id in enumerate(batch):
nursery.start_soon(get_one, idx, peer_id, batch_results)
for providers in batch_results:
all_providers.extend(providers)
if len(all_providers) >= count:
return all_providers[:count]
return all_providers[:count]
async def _get_providers_from_peer(self, peer_id: ID, key: bytes) -> list[PeerInfo]:
"""
Get content providers from a specific peer.
:param peer_id: The peer to query
:param key: The content key to look for
Returns
-------
List[PeerInfo]
List of provider information
"""
providers: list[PeerInfo] = []
try:
# Open a stream to the peer
stream = await self.host.new_stream(peer_id, [TProtocol(PROTOCOL_ID)])
try:
# Create the GET_PROVIDERS message
message = Message()
message.type = Message.MessageType.GET_PROVIDERS
message.key = key
# Serialize and send the message
proto_bytes = message.SerializeToString()
await stream.write(varint.encode(len(proto_bytes)))
await stream.write(proto_bytes)
# Read response length prefix
length_bytes = b""
while True:
b = await stream.read(1)
if not b:
return []
length_bytes += b
if b[0] & 0x80 == 0:
break
response_length = varint.decode_bytes(length_bytes)
# Read response data
response_bytes = b""
remaining = response_length
while remaining > 0:
chunk = await stream.read(remaining)
if not chunk:
return []
response_bytes += chunk
remaining -= len(chunk)
# Parse response
response = Message()
response.ParseFromString(response_bytes)
# Check response type
if response.type != Message.MessageType.GET_PROVIDERS:
return []
# Extract provider information
providers = []
for provider_proto in response.providerPeers:
try:
# Create peer ID from bytes
provider_id = ID(provider_proto.id)
# Convert addresses to Multiaddr
addrs = []
for addr_bytes in provider_proto.addrs:
try:
addrs.append(Multiaddr(addr_bytes))
except Exception:
pass # Skip invalid addresses
# Create PeerInfo and add to result
providers.append(PeerInfo(provider_id, addrs))
except Exception as e:
logger.warning(f"Failed to parse provider info: {e}")
finally:
await stream.close()
return providers
except Exception as e:
logger.warning(f"Error getting providers from {peer_id}: {e}")
return []
def add_provider(self, key: bytes, provider: PeerInfo) -> None:
"""
Add a provider for a given content key.
:param key: The content key
:param provider: The provider's peer information
Returns
-------
None
"""
# Initialize providers for this key if needed
if key not in self.providers:
self.providers[key] = {}
# Add or update the provider record
peer_id_str = str(provider.peer_id) # Use string representation as dict key
self.providers[key][peer_id_str] = ProviderRecord(
provider_info=provider, timestamp=time.time()
)
logger.debug(f"Added provider {provider.peer_id} for key {key.hex()}")
def get_providers(self, key: bytes) -> list[PeerInfo]:
"""
Get all providers for a given content key.
:param key: The content key
Returns
-------
List[PeerInfo]
List of providers for the key
"""
if key not in self.providers:
return []
# Collect valid provider records (not expired)
result = []
current_time = time.time()
expired_peers = []
for peer_id_str, record in self.providers[key].items():
# Check if the record has expired
if current_time - record.timestamp > PROVIDER_RECORD_EXPIRATION_INTERVAL:
expired_peers.append(peer_id_str)
continue
# Use addresses only if they haven't expired
addresses = []
if current_time - record.timestamp <= PROVIDER_ADDRESS_TTL:
addresses = record.addresses
# Create PeerInfo and add to results
result.append(PeerInfo(record.peer_id, addresses))
# Clean up expired records
for peer_id in expired_peers:
del self.providers[key][peer_id]
# Remove the key if no providers left
if not self.providers[key]:
del self.providers[key]
return result
def cleanup_expired(self) -> None:
"""Remove expired provider records."""
current_time = time.time()
expired_keys = []
for key, providers in self.providers.items():
expired_providers = []
for peer_id_str, record in providers.items():
if (
current_time - record.timestamp
> PROVIDER_RECORD_EXPIRATION_INTERVAL
):
expired_providers.append(peer_id_str)
logger.debug(
f"Removing expired provider {peer_id_str} for key {key.hex()}"
)
# Remove expired providers
for peer_id in expired_providers:
del providers[peer_id]
# Track empty keys for removal
if not providers:
expired_keys.append(key)
# Remove empty keys
for key in expired_keys:
del self.providers[key]
logger.debug(f"Removed key with no providers: {key.hex()}")
def get_provided_keys(self, peer_id: ID) -> list[bytes]:
"""
Get all content keys provided by a specific peer.
:param peer_id: The peer ID to look for
Returns
-------
List[bytes]
List of content keys provided by the peer
"""
peer_id_str = str(peer_id)
result = []
for key, providers in self.providers.items():
if peer_id_str in providers:
result.append(key)
return result
def size(self) -> int:
"""
Get the total number of provider records in the store.
Returns
-------
int
Total number of provider records across all keys
"""
total = 0
for providers in self.providers.values():
total += len(providers)
return total

601
libp2p/kad_dht/routing_table.py Normal file
View File

@ -0,0 +1,601 @@
"""
Kademlia DHT routing table implementation.
"""
from collections import (
OrderedDict,
)
import logging
import time
import trio
from libp2p.abc import (
IHost,
)
from libp2p.custom_types import (
TProtocol,
)
from libp2p.kad_dht.utils import xor_distance
from libp2p.peer.id import (
ID,
)
from libp2p.peer.peerinfo import (
PeerInfo,
)
from .pb.kademlia_pb2 import (
Message,
)
# logger = logging.getLogger("libp2p.kademlia.routing_table")
logger = logging.getLogger("kademlia-example.routing_table")
# Default parameters
BUCKET_SIZE = 20 # k in the Kademlia paper
MAXIMUM_BUCKETS = 256 # Maximum number of buckets (for 256-bit keys)
PEER_REFRESH_INTERVAL = 60 # Interval to refresh peers in seconds
STALE_PEER_THRESHOLD = 3600 # Time in seconds after which a peer is considered stale
class KBucket:
"""
A k-bucket implementation for the Kademlia DHT.
Each k-bucket stores up to k (BUCKET_SIZE) peers, sorted by least-recently seen.
"""
def __init__(
self,
host: IHost,
bucket_size: int = BUCKET_SIZE,
min_range: int = 0,
max_range: int = 2**256,
):
"""
Initialize a new k-bucket.
:param host: The host this bucket belongs to
:param bucket_size: Maximum number of peers to store in the bucket
:param min_range: Lower boundary of the bucket's key range (inclusive)
:param max_range: Upper boundary of the bucket's key range (exclusive)
"""
self.bucket_size = bucket_size
self.host = host
self.min_range = min_range
self.max_range = max_range
# Store PeerInfo objects along with last-seen timestamp
self.peers: OrderedDict[ID, tuple[PeerInfo, float]] = OrderedDict()
def peer_ids(self) -> list[ID]:
"""Get all peer IDs in the bucket."""
return list(self.peers.keys())
def peer_infos(self) -> list[PeerInfo]:
"""Get all PeerInfo objects in the bucket."""
return [info for info, _ in self.peers.values()]
def get_oldest_peer(self) -> ID | None:
"""Get the least-recently seen peer."""
if not self.peers:
return None
return next(iter(self.peers.keys()))
async def add_peer(self, peer_info: PeerInfo) -> bool:
"""
Add a peer to the bucket. Returns True if the peer was added or updated,
False if the bucket is full.
"""
current_time = time.time()
peer_id = peer_info.peer_id
# If peer is already in the bucket, move it to the end (most recently seen)
if peer_id in self.peers:
self.refresh_peer_last_seen(peer_id)
return True
# If bucket has space, add the peer
if len(self.peers) < self.bucket_size:
self.peers[peer_id] = (peer_info, current_time)
return True
# If bucket is full, we need to replace the least-recently seen peer
# Get the least-recently seen peer
oldest_peer_id = self.get_oldest_peer()
if oldest_peer_id is None:
logger.warning("No oldest peer found when bucket is full")
return False
# Check if the old peer is responsive to ping request
try:
# Try to ping the oldest peer, not the new peer
response = await self._ping_peer(oldest_peer_id)
if response:
# If the old peer is still alive, we will not add the new peer
logger.debug(
"Old peer %s is still alive, cannot add new peer %s",
oldest_peer_id,
peer_id,
)
return False
except Exception as e:
# If the old peer is unresponsive, we can replace it with the new peer
logger.debug(
"Old peer %s is unresponsive, replacing with new peer %s: %s",
oldest_peer_id,
peer_id,
str(e),
)
self.peers.popitem(last=False) # Remove oldest peer
self.peers[peer_id] = (peer_info, current_time)
return True
# If we got here, the oldest peer responded but we couldn't add the new peer
return False
def remove_peer(self, peer_id: ID) -> bool:
"""
Remove a peer from the bucket.
Returns True if the peer was in the bucket, False otherwise.
"""
if peer_id in self.peers:
del self.peers[peer_id]
return True
return False
def has_peer(self, peer_id: ID) -> bool:
"""Check if the peer is in the bucket."""
return peer_id in self.peers
def get_peer_info(self, peer_id: ID) -> PeerInfo | None:
"""Get the PeerInfo for a given peer ID if it exists in the bucket."""
if peer_id in self.peers:
return self.peers[peer_id][0]
return None
def size(self) -> int:
"""Get the number of peers in the bucket."""
return len(self.peers)
def get_stale_peers(self, stale_threshold_seconds: int = 3600) -> list[ID]:
"""
Get peers that haven't been pinged recently.
params: stale_threshold_seconds: Time in seconds
params: after which a peer is considered stale
Returns
-------
list[ID]
List of peer IDs that need to be refreshed
"""
current_time = time.time()
stale_peers = []
for peer_id, (_, last_seen) in self.peers.items():
if current_time - last_seen > stale_threshold_seconds:
stale_peers.append(peer_id)
return stale_peers
async def _periodic_peer_refresh(self) -> None:
"""Background task to periodically refresh peers"""
try:
while True:
await trio.sleep(PEER_REFRESH_INTERVAL) # Check every minute
# Find stale peers (not pinged in last hour)
stale_peers = self.get_stale_peers(
stale_threshold_seconds=STALE_PEER_THRESHOLD
)
if stale_peers:
logger.debug(f"Found {len(stale_peers)} stale peers to refresh")
for peer_id in stale_peers:
try:
# Try to ping the peer
logger.debug("Pinging stale peer %s", peer_id)
responce = await self._ping_peer(peer_id)
if responce:
# Update the last seen time
self.refresh_peer_last_seen(peer_id)
logger.debug(f"Refreshed peer {peer_id}")
else:
# If ping fails, remove the peer
logger.debug(f"Failed to ping peer {peer_id}")
self.remove_peer(peer_id)
logger.info(f"Removed unresponsive peer {peer_id}")
logger.debug(f"Successfully refreshed peer {peer_id}")
except Exception as e:
# If ping fails, remove the peer
logger.debug(
"Failed to ping peer %s: %s",
peer_id,
e,
)
self.remove_peer(peer_id)
logger.info(f"Removed unresponsive peer {peer_id}")
except trio.Cancelled:
logger.debug("Peer refresh task cancelled")
except Exception as e:
logger.error(f"Error in peer refresh task: {e}", exc_info=True)
async def _ping_peer(self, peer_id: ID) -> bool:
"""
Ping a peer using protobuf message to check
if it's still alive and update last seen time.
params: peer_id: The ID of the peer to ping
Returns
-------
bool
True if ping successful, False otherwise
"""
result = False
# Get peer info directly from the bucket
peer_info = self.get_peer_info(peer_id)
if not peer_info:
raise ValueError(f"Peer {peer_id} not in bucket")
# Default protocol ID for Kademlia DHT
protocol_id = TProtocol("/ipfs/kad/1.0.0")
try:
# Open a stream to the peer with the DHT protocol
stream = await self.host.new_stream(peer_id, [protocol_id])
try:
# Create ping protobuf message
ping_msg = Message()
ping_msg.type = Message.PING # Use correct enum
# Serialize and send with length prefix (4 bytes big-endian)
msg_bytes = ping_msg.SerializeToString()
logger.debug(
f"Sending PING message to {peer_id}, size: {len(msg_bytes)} bytes"
)
await stream.write(len(msg_bytes).to_bytes(4, byteorder="big"))
await stream.write(msg_bytes)
# Wait for response with timeout
with trio.move_on_after(2): # 2 second timeout
# Read response length (4 bytes)
length_bytes = await stream.read(4)
if not length_bytes or len(length_bytes) < 4:
logger.warning(f"Peer {peer_id} disconnected during ping")
return False
msg_len = int.from_bytes(length_bytes, byteorder="big")
if (
msg_len <= 0 or msg_len > 1024 * 1024
): # Sanity check on message size
logger.warning(
f"Invalid message length from {peer_id}: {msg_len}"
)
return False
logger.debug(
f"Receiving response from {peer_id}, size: {msg_len} bytes"
)
# Read full message
response_bytes = await stream.read(msg_len)
if not response_bytes:
logger.warning(f"Failed to read response from {peer_id}")
return False
# Parse protobuf response
response = Message()
try:
response.ParseFromString(response_bytes)
except Exception as e:
logger.warning(
f"Failed to parse protobuf response from {peer_id}: {e}"
)
return False
if response.type == Message.PING:
# Update the last seen timestamp for this peer
logger.debug(f"Successfully pinged peer {peer_id}")
result = True
return result
else:
logger.warning(
f"Unexpected response type from {peer_id}: {response.type}"
)
return False
# If we get here, the ping timed out
logger.warning(f"Ping to peer {peer_id} timed out")
return False
finally:
await stream.close()
return result
except Exception as e:
logger.error(f"Error pinging peer {peer_id}: {str(e)}")
return False
def refresh_peer_last_seen(self, peer_id: ID) -> bool:
"""
Update the last-seen timestamp for a peer in the bucket.
params: peer_id: The ID of the peer to refresh
Returns
-------
bool
True if the peer was found and refreshed, False otherwise
"""
if peer_id in self.peers:
# Get current peer info and update the timestamp
peer_info, _ = self.peers[peer_id]
current_time = time.time()
self.peers[peer_id] = (peer_info, current_time)
# Move to end of ordered dict to mark as most recently seen
self.peers.move_to_end(peer_id)
return True
return False
def key_in_range(self, key: bytes) -> bool:
"""
Check if a key is in the range of this bucket.
params: key: The key to check (bytes)
Returns
-------
bool
True if the key is in range, False otherwise
"""
key_int = int.from_bytes(key, byteorder="big")
return self.min_range <= key_int < self.max_range
def split(self) -> tuple["KBucket", "KBucket"]:
"""
Split the bucket into two buckets.
Returns
-------
tuple
(lower_bucket, upper_bucket)
"""
midpoint = (self.min_range + self.max_range) // 2
lower_bucket = KBucket(self.host, self.bucket_size, self.min_range, midpoint)
upper_bucket = KBucket(self.host, self.bucket_size, midpoint, self.max_range)
# Redistribute peers
for peer_id, (peer_info, timestamp) in self.peers.items():
peer_key = int.from_bytes(peer_id.to_bytes(), byteorder="big")
if peer_key < midpoint:
lower_bucket.peers[peer_id] = (peer_info, timestamp)
else:
upper_bucket.peers[peer_id] = (peer_info, timestamp)
return lower_bucket, upper_bucket
class RoutingTable:
"""
The Kademlia routing table maintains information on which peers to contact for any
given peer ID in the network.
"""
def __init__(self, local_id: ID, host: IHost) -> None:
"""
Initialize the routing table.
:param local_id: The ID of the local node.
:param host: The host this routing table belongs to.
"""
self.local_id = local_id
self.host = host
self.buckets = [KBucket(host, BUCKET_SIZE)]
async def add_peer(self, peer_obj: PeerInfo | ID) -> bool:
"""
Add a peer to the routing table.
:param peer_obj: Either PeerInfo object or peer ID to add
Returns
-------
bool: True if the peer was added or updated, False otherwise
"""
peer_id = None
peer_info = None
try:
# Handle different types of input
if isinstance(peer_obj, PeerInfo):
# Already have PeerInfo object
peer_info = peer_obj
peer_id = peer_obj.peer_id
else:
# Assume it's a peer ID
peer_id = peer_obj
# Try to get addresses from the peerstore if available
try:
addrs = self.host.get_peerstore().addrs(peer_id)
if addrs:
# Create PeerInfo object
peer_info = PeerInfo(peer_id, addrs)
else:
logger.debug(
"No addresses found for peer %s in peerstore, skipping",
peer_id,
)
return False
except Exception as peerstore_error:
# Handle case where peer is not in peerstore yet
logger.debug(
"Peer %s not found in peerstore: %s, skipping",
peer_id,
str(peerstore_error),
)
return False
# Don't add ourselves
if peer_id == self.local_id:
return False
# Find the right bucket for this peer
bucket = self.find_bucket(peer_id)
# Try to add to the bucket
success = await bucket.add_peer(peer_info)
if success:
logger.debug(f"Successfully added peer {peer_id} to routing table")
return success
except Exception as e:
logger.debug(f"Error adding peer {peer_obj} to routing table: {e}")
return False
def remove_peer(self, peer_id: ID) -> bool:
"""
Remove a peer from the routing table.
:param peer_id: The ID of the peer to remove
Returns
-------
bool: True if the peer was removed, False otherwise
"""
bucket = self.find_bucket(peer_id)
return bucket.remove_peer(peer_id)
def find_bucket(self, peer_id: ID) -> KBucket:
"""
Find the bucket that would contain the given peer ID or PeerInfo.
:param peer_obj: Either a peer ID or a PeerInfo object
Returns
-------
KBucket: The bucket for this peer
"""
for bucket in self.buckets:
if bucket.key_in_range(peer_id.to_bytes()):
return bucket
return self.buckets[0]
def find_local_closest_peers(self, key: bytes, count: int = 20) -> list[ID]:
"""
Find the closest peers to a given key.
:param key: The key to find closest peers to (bytes)
:param count: Maximum number of peers to return
Returns
-------
List[ID]: List of peer IDs closest to the key
"""
# Get all peers from all buckets
all_peers = []
for bucket in self.buckets:
all_peers.extend(bucket.peer_ids())
# Sort by XOR distance to the key
all_peers.sort(key=lambda p: xor_distance(p.to_bytes(), key))
return all_peers[:count]
def get_peer_ids(self) -> list[ID]:
"""
Get all peer IDs in the routing table.
Returns
-------
:param List[ID]: List of all peer IDs
"""
peers = []
for bucket in self.buckets:
peers.extend(bucket.peer_ids())
return peers
def get_peer_info(self, peer_id: ID) -> PeerInfo | None:
"""
Get the peer info for a specific peer.
:param peer_id: The ID of the peer to get info for
Returns
-------
PeerInfo: The peer info, or None if not found
"""
bucket = self.find_bucket(peer_id)
return bucket.get_peer_info(peer_id)
def peer_in_table(self, peer_id: ID) -> bool:
"""
Check if a peer is in the routing table.
:param peer_id: The ID of the peer to check
Returns
-------
bool: True if the peer is in the routing table, False otherwise
"""
bucket = self.find_bucket(peer_id)
return bucket.has_peer(peer_id)
def size(self) -> int:
"""
Get the number of peers in the routing table.
Returns
-------
int: Number of peers
"""
count = 0
for bucket in self.buckets:
count += bucket.size()
return count
def get_stale_peers(self, stale_threshold_seconds: int = 3600) -> list[ID]:
"""
Get all stale peers from all buckets
params: stale_threshold_seconds:
Time in seconds after which a peer is considered stale
Returns
-------
list[ID]
List of stale peer IDs
"""
stale_peers = []
for bucket in self.buckets:
stale_peers.extend(bucket.get_stale_peers(stale_threshold_seconds))
return stale_peers
def cleanup_routing_table(self) -> None:
"""
Cleanup the routing table by removing all data.
This is useful for resetting the routing table during tests or reinitialization.
"""
self.buckets = [KBucket(self.host, BUCKET_SIZE)]
logger.info("Routing table cleaned up, all data removed.")

117
libp2p/kad_dht/utils.py Normal file
View File

@ -0,0 +1,117 @@
"""
Utility functions for Kademlia DHT implementation.
"""
import base58
import multihash
from libp2p.peer.id import (
ID,
)
def create_key_from_binary(binary_data: bytes) -> bytes:
"""
Creates a key for the DHT by hashing binary data with SHA-256.
params: binary_data: The binary data to hash.
Returns
-------
bytes: The resulting key.
"""
return multihash.digest(binary_data, "sha2-256").digest
def xor_distance(key1: bytes, key2: bytes) -> int:
"""
Calculate the XOR distance between two keys.
params: key1: First key (bytes)
params: key2: Second key (bytes)
Returns
-------
int: The XOR distance between the keys
"""
# Ensure the inputs are bytes
if not isinstance(key1, bytes) or not isinstance(key2, bytes):
raise TypeError("Both key1 and key2 must be bytes objects")
# Convert to integers
k1 = int.from_bytes(key1, byteorder="big")
k2 = int.from_bytes(key2, byteorder="big")
# Calculate XOR distance
return k1 ^ k2
def bytes_to_base58(data: bytes) -> str:
"""
Convert bytes to base58 encoded string.
params: data: Input bytes
Returns
-------
str: Base58 encoded string
"""
return base58.b58encode(data).decode("utf-8")
def sort_peer_ids_by_distance(target_key: bytes, peer_ids: list[ID]) -> list[ID]:
"""
Sort a list of peer IDs by their distance to the target key.
params: target_key: The target key to measure distance from
params: peer_ids: List of peer IDs to sort
Returns
-------
List[ID]: Sorted list of peer IDs from closest to furthest
"""
def get_distance(peer_id: ID) -> int:
# Hash the peer ID bytes to get a key for distance calculation
peer_hash = multihash.digest(peer_id.to_bytes(), "sha2-256").digest
return xor_distance(target_key, peer_hash)
return sorted(peer_ids, key=get_distance)
def shared_prefix_len(first: bytes, second: bytes) -> int:
"""
Calculate the number of prefix bits shared by two byte sequences.
params: first: First byte sequence
params: second: Second byte sequence
Returns
-------
int: Number of shared prefix bits
"""
# Compare each byte to find the first bit difference
common_length = 0
for i in range(min(len(first), len(second))):
byte_first = first[i]
byte_second = second[i]
if byte_first == byte_second:
common_length += 8
else:
# Find specific bit where they differ
xor = byte_first ^ byte_second
# Count leading zeros in the xor result
for j in range(7, -1, -1):
if (xor >> j) & 1 == 1:
return common_length + (7 - j)
# This shouldn't be reached if xor != 0
return common_length + 8
return common_length

393
libp2p/kad_dht/value_store.py Normal file
View File

@ -0,0 +1,393 @@
"""
Value store implementation for Kademlia DHT.
Provides a way to store and retrieve key-value pairs with optional expiration.
"""
import logging
import time
import varint
from libp2p.abc import (
IHost,
)
from libp2p.custom_types import (
TProtocol,
)
from libp2p.peer.id import (
ID,
)
from .pb.kademlia_pb2 import (
Message,
)
# logger = logging.getLogger("libp2p.kademlia.value_store")
logger = logging.getLogger("kademlia-example.value_store")
# Default time to live for values in seconds (24 hours)
DEFAULT_TTL = 24 * 60 * 60
PROTOCOL_ID = TProtocol("/ipfs/kad/1.0.0")
class ValueStore:
"""
Store for key-value pairs in a Kademlia DHT.
Values are stored with a timestamp and optional expiration time.
"""
def __init__(self, host: IHost, local_peer_id: ID):
"""
Initialize an empty value store.
:param host: The libp2p host instance.
:param local_peer_id: The local peer ID to ignore in peer requests.
"""
# Store format: {key: (value, validity)}
self.store: dict[bytes, tuple[bytes, float]] = {}
# Store references to the host and local peer ID for making requests
self.host = host
self.local_peer_id = local_peer_id
def put(self, key: bytes, value: bytes, validity: float = 0.0) -> None:
"""
Store a value in the DHT.
:param key: The key to store the value under
:param value: The value to store
:param validity: validity in seconds before the value expires.
Defaults to `DEFAULT_TTL` if set to 0.0.
Returns
-------
None
"""
if validity == 0.0:
validity = time.time() + DEFAULT_TTL
logger.debug(
"Storing value for key %s... with validity %s", key.hex(), validity
)
self.store[key] = (value, validity)
logger.debug(f"Stored value for key {key.hex()}")
async def _store_at_peer(self, peer_id: ID, key: bytes, value: bytes) -> bool:
"""
Store a value at a specific peer.
params: peer_id: The ID of the peer to store the value at
params: key: The key to store
params: value: The value to store
Returns
-------
bool
True if the value was successfully stored, False otherwise
"""
result = False
stream = None
try:
# Don't try to store at ourselves
if self.local_peer_id and peer_id == self.local_peer_id:
result = True
return result
if not self.host:
logger.error("Host not initialized, cannot store value at peer")
return False
logger.debug(f"Storing value for key {key.hex()} at peer {peer_id}")
# Open a stream to the peer
stream = await self.host.new_stream(peer_id, [PROTOCOL_ID])
logger.debug(f"Opened stream to peer {peer_id}")
# Create the PUT_VALUE message with protobuf
message = Message()
message.type = Message.MessageType.PUT_VALUE
# Set message fields
message.key = key
message.record.key = key
message.record.value = value
message.record.timeReceived = str(time.time())
# Serialize and send the protobuf message with length prefix
proto_bytes = message.SerializeToString()
await stream.write(varint.encode(len(proto_bytes)))
await stream.write(proto_bytes)
logger.debug("Sent PUT_VALUE protobuf message with varint length")
# Read varint-prefixed response length
length_bytes = b""
while True:
logger.debug("Reading varint length prefix for response...")
b = await stream.read(1)
if not b:
logger.warning("Connection closed while reading varint length")
return False
length_bytes += b
if b[0] & 0x80 == 0:
break
logger.debug(f"Received varint length bytes: {length_bytes.hex()}")
response_length = varint.decode_bytes(length_bytes)
logger.debug("Response length: %d bytes", response_length)
# Read response data
response_bytes = b""
remaining = response_length
while remaining > 0:
chunk = await stream.read(remaining)
if not chunk:
logger.debug(
f"Connection closed by peer {peer_id} while reading data"
)
return False
response_bytes += chunk
remaining -= len(chunk)
# Parse protobuf response
response = Message()
response.ParseFromString(response_bytes)
# Check if response is valid
if response.type == Message.MessageType.PUT_VALUE:
if response.key:
result = True
return result
except Exception as e:
logger.warning(f"Failed to store value at peer {peer_id}: {e}")
return False
finally:
if stream:
await stream.close()
return result
def get(self, key: bytes) -> bytes | None:
"""
Retrieve a value from the DHT.
params: key: The key to look up
Returns
-------
Optional[bytes]
The stored value, or None if not found or expired
"""
logger.debug("Retrieving value for key %s...", key.hex()[:8])
if key not in self.store:
return None
value, validity = self.store[key]
logger.debug(
"Found value for key %s... with validity %s",
key.hex(),
validity,
)
# Check if the value has expired
if validity is not None and validity < time.time():
logger.debug(
"Value for key %s... has expired, removing it",
key.hex()[:8],
)
self.remove(key)
return None
return value
async def _get_from_peer(self, peer_id: ID, key: bytes) -> bytes | None:
"""
Retrieve a value from a specific peer.
params: peer_id: The ID of the peer to retrieve the value from
params: key: The key to retrieve
Returns
-------
Optional[bytes]
The value if found, None otherwise
"""
stream = None
try:
# Don't try to get from ourselves
if peer_id == self.local_peer_id:
return None
logger.debug(f"Getting value for key {key.hex()} from peer {peer_id}")
# Open a stream to the peer
stream = await self.host.new_stream(peer_id, [TProtocol(PROTOCOL_ID)])
logger.debug(f"Opened stream to peer {peer_id} for GET_VALUE")
# Create the GET_VALUE message using protobuf
message = Message()
message.type = Message.MessageType.GET_VALUE
message.key = key
# Serialize and send the protobuf message
proto_bytes = message.SerializeToString()
await stream.write(varint.encode(len(proto_bytes)))
await stream.write(proto_bytes)
# Read response length
length_bytes = b""
while True:
b = await stream.read(1)
if not b:
logger.warning("Connection closed while reading length")
return None
length_bytes += b
if b[0] & 0x80 == 0:
break
response_length = varint.decode_bytes(length_bytes)
# Read response data
response_bytes = b""
remaining = response_length
while remaining > 0:
chunk = await stream.read(remaining)
if not chunk:
logger.debug(
f"Connection closed by peer {peer_id} while reading data"
)
return None
response_bytes += chunk
remaining -= len(chunk)
# Parse protobuf response
try:
response = Message()
response.ParseFromString(response_bytes)
logger.debug(
f"Received protobuf response from peer"
f" {peer_id}, type: {response.type}"
)
# Process protobuf response
if (
response.type == Message.MessageType.GET_VALUE
and response.HasField("record")
and response.record.value
):
logger.debug(
f"Received value for key {key.hex()} from peer {peer_id}"
)
return response.record.value
# Handle case where value is not found but peer infos are returned
else:
logger.debug(
f"Value not found for key {key.hex()} from peer {peer_id},"
f" received {len(response.closerPeers)} closer peers"
)
return None
except Exception as proto_err:
logger.warning(f"Failed to parse as protobuf: {proto_err}")
return None
except Exception as e:
logger.warning(f"Failed to get value from peer {peer_id}: {e}")
return None
finally:
if stream:
await stream.close()
def remove(self, key: bytes) -> bool:
"""
Remove a value from the DHT.
params: key: The key to remove
Returns
-------
bool
True if the key was found and removed, False otherwise
"""
if key in self.store:
del self.store[key]
logger.debug(f"Removed value for key {key.hex()[:8]}...")
return True
return False
def has(self, key: bytes) -> bool:
"""
Check if a key exists in the store and hasn't expired.
params: key: The key to check
Returns
-------
bool
True if the key exists and hasn't expired, False otherwise
"""
if key not in self.store:
return False
_, validity = self.store[key]
if validity is not None and time.time() > validity:
self.remove(key)
return False
return True
def cleanup_expired(self) -> int:
"""
Remove all expired values from the store.
Returns
-------
int
The number of expired values that were removed
"""
current_time = time.time()
expired_keys = [
key for key, (_, validity) in self.store.items() if current_time > validity
]
for key in expired_keys:
del self.store[key]
if expired_keys:
logger.debug(f"Cleaned up {len(expired_keys)} expired values")
return len(expired_keys)
def get_keys(self) -> list[bytes]:
"""
Get all non-expired keys in the store.
Returns
-------
list[bytes]
List of keys
"""
# Clean up expired values first
self.cleanup_expired()
return list(self.store.keys())
def size(self) -> int:
"""
Get the number of items in the store (after removing expired entries).
Returns
-------
int
Number of items
"""
self.cleanup_expired()
return len(self.store)

7
libp2p/network/swarm.py
View File

@ -187,7 +187,7 @@ class Swarm(Service, INetworkService):
# Per, https://discuss.libp2p.io/t/multistream-security/130, we first secure
# the conn and then mux the conn
try:
secured_conn = await self.upgrader.upgrade_security(raw_conn, peer_id, True)
secured_conn = await self.upgrader.upgrade_security(raw_conn, True, peer_id)
except SecurityUpgradeFailure as error:
logger.debug("failed to upgrade security for peer %s", peer_id)
await raw_conn.close()
@ -257,10 +257,7 @@ class Swarm(Service, INetworkService):
# Per, https://discuss.libp2p.io/t/multistream-security/130, we first
# secure the conn and then mux the conn
try:
# FIXME: This dummy `ID(b"")` for the remote peer is useless.
secured_conn = await self.upgrader.upgrade_security(
raw_conn, ID(b""), False
)
secured_conn = await self.upgrader.upgrade_security(raw_conn, False)
except SecurityUpgradeFailure as error:
logger.debug("failed to upgrade security for peer at %s", maddr)
await raw_conn.close()

151
libp2p/pubsub/gossipsub.py
View File

@ -516,82 +516,99 @@ class GossipSub(IPubsubRouter, Service):
peers_to_prune[peer].append(topic)
return peers_to_graft, peers_to_prune
def fanout_heartbeat(self) -> None:
# Note: the comments here are the exact pseudocode from the spec
for topic in list(self.fanout):
if (
self.pubsub is not None
and topic not in self.pubsub.peer_topics
and self.time_since_last_publish.get(topic, 0) + self.time_to_live
< int(time.time())
def _handle_topic_heartbeat(
self,
topic: str,
current_peers: set[ID],
is_fanout: bool = False,
peers_to_gossip: DefaultDict[ID, dict[str, list[str]]] | None = None,
) -> tuple[set[ID], bool]:
"""
Helper method to handle heartbeat for a single topic,
supporting both fanout and gossip.
:param topic: The topic to handle
:param current_peers: Current set of peers in the topic
:param is_fanout: Whether this is a fanout topic (affects expiration check)
:param peers_to_gossip: Optional dictionary to store peers to gossip to
:return: Tuple of (updated_peers, should_remove_topic)
"""
if self.pubsub is None:
raise NoPubsubAttached
# Skip if no peers have subscribed to the topic
if topic not in self.pubsub.peer_topics:
return current_peers, False
# For fanout topics, check if we should remove the topic
if is_fanout:
if self.time_since_last_publish.get(topic, 0) + self.time_to_live < int(
time.time()
):
# Remove topic from fanout
return set(), True
# Check if peers are still in the topic and remove the ones that are not
in_topic_peers: set[ID] = {
peer for peer in current_peers if peer in self.pubsub.peer_topics[topic]
}
# If we need more peers to reach target degree
if len(in_topic_peers) < self.degree:
# Select additional peers from peers.gossipsub[topic]
selected_peers = self._get_in_topic_gossipsub_peers_from_minus(
topic,
self.degree - len(in_topic_peers),
in_topic_peers,
)
# Add the selected peers
in_topic_peers.update(selected_peers)
# Handle gossip if requested
if peers_to_gossip is not None:
msg_ids = self.mcache.window(topic)
if msg_ids:
# Select D peers from peers.gossipsub[topic] excluding current peers
peers_to_emit_ihave_to = self._get_in_topic_gossipsub_peers_from_minus(
topic, self.degree, current_peers
)
msg_id_strs = [str(msg_id) for msg_id in msg_ids]
for peer in peers_to_emit_ihave_to:
peers_to_gossip[peer][topic] = msg_id_strs
return in_topic_peers, False
def fanout_heartbeat(self) -> None:
"""
Maintain fanout topics by:
1. Removing expired topics
2. Removing peers that are no longer in the topic
3. Adding new peers if needed to maintain the target degree
"""
for topic in list(self.fanout):
updated_peers, should_remove = self._handle_topic_heartbeat(
topic, self.fanout[topic], is_fanout=True
)
if should_remove:
del self.fanout[topic]
else:
# Check if fanout peers are still in the topic and remove the ones that are not # noqa: E501
# ref: https://github.com/libp2p/go-libp2p-pubsub/blob/01b9825fbee1848751d90a8469e3f5f43bac8466/gossipsub.go#L498-L504 # noqa: E501
in_topic_fanout_peers: list[ID] = []
if self.pubsub is not None:
in_topic_fanout_peers = [
peer
for peer in self.fanout[topic]
if peer in self.pubsub.peer_topics[topic]
]
self.fanout[topic] = set(in_topic_fanout_peers)
num_fanout_peers_in_topic = len(self.fanout[topic])
# If |fanout[topic]| < D
if num_fanout_peers_in_topic < self.degree:
# Select D - |fanout[topic]| peers from peers.gossipsub[topic] - fanout[topic] # noqa: E501
selected_peers = self._get_in_topic_gossipsub_peers_from_minus(
topic,
self.degree - num_fanout_peers_in_topic,
self.fanout[topic],
)
# Add the peers to fanout[topic]
self.fanout[topic].update(selected_peers)
self.fanout[topic] = updated_peers
def gossip_heartbeat(self) -> DefaultDict[ID, dict[str, list[str]]]:
peers_to_gossip: DefaultDict[ID, dict[str, list[str]]] = defaultdict(dict)
# Handle mesh topics
for topic in self.mesh:
msg_ids = self.mcache.window(topic)
if msg_ids:
if self.pubsub is None:
raise NoPubsubAttached
# Get all pubsub peers in a topic and only add them if they are
# gossipsub peers too
if topic in self.pubsub.peer_topics:
# Select D peers from peers.gossipsub[topic]
peers_to_emit_ihave_to = (
self._get_in_topic_gossipsub_peers_from_minus(
topic, self.degree, self.mesh[topic]
)
)
self._handle_topic_heartbeat(
topic, self.mesh[topic], peers_to_gossip=peers_to_gossip
)
msg_id_strs = [str(msg_id) for msg_id in msg_ids]
for peer in peers_to_emit_ihave_to:
peers_to_gossip[peer][topic] = msg_id_strs
# TODO: Refactor and Dedup. This section is the roughly the same as the above.
# Do the same for fanout, for all topics not already hit in mesh
# Handle fanout topics that aren't in mesh
for topic in self.fanout:
msg_ids = self.mcache.window(topic)
if msg_ids:
if self.pubsub is None:
raise NoPubsubAttached
# Get all pubsub peers in topic and only add if they are
# gossipsub peers also
if topic in self.pubsub.peer_topics:
# Select D peers from peers.gossipsub[topic]
peers_to_emit_ihave_to = (
self._get_in_topic_gossipsub_peers_from_minus(
topic, self.degree, self.fanout[topic]
)
)
msg_id_strs = [str(msg) for msg in msg_ids]
for peer in peers_to_emit_ihave_to:
peers_to_gossip[peer][topic] = msg_id_strs
if topic not in self.mesh:
self._handle_topic_heartbeat(
topic, self.fanout[topic], peers_to_gossip=peers_to_gossip
)
return peers_to_gossip
@staticmethod

14
libp2p/pubsub/pubsub.py
View File

@ -621,16 +621,22 @@ class Pubsub(Service, IPubsub):
logger.debug("Fail to message peer %s: stream closed", peer_id)
self._handle_dead_peer(peer_id)
async def publish(self, topic_id: str, data: bytes) -> None:
async def publish(self, topic_id: str | list[str], data: bytes) -> None:
"""
Publish data to a topic.
Publish data to a topic or multiple topics.
:param topic_id: topic which we are going to publish the data to
:param topic_id: topic (str) or topics (list[str]) to publish the data to
:param data: data which we are publishing
"""
# Handle both single topic (str) and multiple topics (list[str])
if isinstance(topic_id, str):
topic_ids = [topic_id]
else:
topic_ids = topic_id
msg = rpc_pb2.Message(
data=data,
topicIDs=[topic_id],
topicIDs=topic_ids,
# Origin is ourself.
from_id=self.my_id.to_bytes(),
seqno=self._next_seqno(),

28
libp2p/relay/__init__.py Normal file
View File

@ -0,0 +1,28 @@
"""
Relay module for libp2p.
This package includes implementations of circuit relay protocols
for enabling connectivity between peers behind NATs or firewalls.
"""
# Import the circuit_v2 module to make it accessible
# through the relay package
from libp2p.relay.circuit_v2 import (
PROTOCOL_ID,
CircuitV2Protocol,
CircuitV2Transport,
RelayDiscovery,
RelayLimits,
RelayResourceManager,
Reservation,
)
__all__ = [
"CircuitV2Protocol",
"CircuitV2Transport",
"PROTOCOL_ID",
"RelayDiscovery",
"RelayLimits",
"RelayResourceManager",
"Reservation",
]

32
libp2p/relay/circuit_v2/__init__.py Normal file
View File

@ -0,0 +1,32 @@
"""
Circuit Relay v2 implementation for libp2p.
This package implements the Circuit Relay v2 protocol as specified in:
https://github.com/libp2p/specs/blob/master/relay/circuit-v2.md
"""
from .discovery import (
RelayDiscovery,
)
from .protocol import (
PROTOCOL_ID,
CircuitV2Protocol,
)
from .resources import (
RelayLimits,
RelayResourceManager,
Reservation,
)
from .transport import (
CircuitV2Transport,
)
__all__ = [
"CircuitV2Protocol",
"PROTOCOL_ID",
"RelayLimits",
"Reservation",
"RelayResourceManager",
"CircuitV2Transport",
"RelayDiscovery",
]

92
libp2p/relay/circuit_v2/config.py Normal file
View File

@ -0,0 +1,92 @@
"""
Configuration management for Circuit Relay v2.
This module handles configuration for relay roles, resource limits,
and discovery settings.
"""
from dataclasses import (
dataclass,
field,
)
from libp2p.peer.peerinfo import (
PeerInfo,
)
from .resources import (
RelayLimits,
)
@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
# 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
# Connection configuration
reservation_ttl: int = 3600 # seconds
max_circuit_duration: int = 3600 # seconds
max_circuit_bytes: int = 1024 * 1024 * 1024 # 1GB
def __post_init__(self) -> None:
"""Initialize default values."""
if self.limits is None:
self.limits = RelayLimits(
duration=self.max_circuit_duration,
data=self.max_circuit_bytes,
max_circuit_conns=8,
max_reservations=4,
)
@dataclass
class HopConfig:
"""Configuration specific to relay (hop) nodes."""
# Resource limits per IP
max_reservations_per_ip: int = 8
max_circuits_per_ip: int = 16
# Rate limiting
reservation_rate_per_ip: int = 4 # per minute
circuit_rate_per_ip: int = 8 # per minute
# Resource quotas
max_circuits_total: int = 64
max_reservations_total: int = 32
# Bandwidth limits
max_bandwidth_per_circuit: int = 1024 * 1024 # 1MB/s
max_bandwidth_total: int = 10 * 1024 * 1024 # 10MB/s
@dataclass
class ClientConfig:
"""Configuration specific to relay clients."""
# Relay selection
min_relay_score: float = 0.5
max_relay_latency: float = 1.0 # seconds
# Auto-relay settings
enable_auto_relay: bool = True
auto_relay_timeout: int = 30 # seconds
max_auto_relay_attempts: int = 3
# Reservation management
reservation_refresh_threshold: float = 0.8 # Refresh at 80% of TTL
max_concurrent_reservations: int = 2

537
libp2p/relay/circuit_v2/discovery.py Normal file
View File

@ -0,0 +1,537 @@
"""
Discovery module for Circuit Relay v2.
This module handles discovering and tracking relay nodes in the network.
"""
from dataclasses import (
dataclass,
)
import logging
import time
from typing import (
Any,
Protocol as TypingProtocol,
cast,
runtime_checkable,
)
import trio
from libp2p.abc import (
IHost,
)
from libp2p.custom_types import (
TProtocol,
)
from libp2p.peer.id import (
ID,
)
from libp2p.tools.async_service import (
Service,
)
from .pb.circuit_pb2 import (
HopMessage,
)
from .protocol import (
PROTOCOL_ID,
)
from .protocol_buffer import (
StatusCode,
)
logger = logging.getLogger("libp2p.relay.circuit_v2.discovery")
# Constants
MAX_RELAYS_TO_TRACK = 10
DEFAULT_DISCOVERY_INTERVAL = 60 # seconds
STREAM_TIMEOUT = 10 # seconds
# Extended interfaces for type checking
@runtime_checkable
class IHostWithMultiselect(TypingProtocol):
"""Extended host interface with multiselect attribute."""
@property
def multiselect(self) -> Any:
"""Get the multiselect component."""
...
@dataclass
class RelayInfo:
"""Information about a discovered relay."""
peer_id: ID
discovered_at: float
last_seen: float
has_reservation: bool = False
reservation_expires_at: float | None = None
reservation_data_limit: int | None = None
class RelayDiscovery(Service):
"""
Discovery service for Circuit Relay v2 nodes.
This service discovers and keeps track of available relay nodes, and optionally
makes reservations with them.
"""
def __init__(
self,
host: IHost,
auto_reserve: bool = False,
discovery_interval: int = DEFAULT_DISCOVERY_INTERVAL,
max_relays: int = MAX_RELAYS_TO_TRACK,
) -> None:
"""
Initialize the discovery service.
Parameters
----------
host : IHost
The libp2p host this discovery service is running on
auto_reserve : bool
Whether to automatically make reservations with discovered relays
discovery_interval : int
How often to run discovery, in seconds
max_relays : int
Maximum number of relays to track
"""
super().__init__()
self.host = host
self.auto_reserve = auto_reserve
self.discovery_interval = discovery_interval
self.max_relays = max_relays
self._discovered_relays: dict[ID, RelayInfo] = {}
self._protocol_cache: dict[
ID, set[str]
] = {} # Cache protocol info to reduce queries
self.event_started = trio.Event()
self.is_running = False
async def run(self, *, task_status: Any = trio.TASK_STATUS_IGNORED) -> None:
"""Run the discovery service."""
try:
self.is_running = True
self.event_started.set()
task_status.started()
# Main discovery loop
async with trio.open_nursery() as nursery:
# Run initial discovery
nursery.start_soon(self.discover_relays)
# Set up periodic discovery
while True:
await trio.sleep(self.discovery_interval)
if not self.manager.is_running:
break
nursery.start_soon(self.discover_relays)
# Cleanup expired relays and reservations
await self._cleanup_expired()
finally:
self.is_running = False
async def discover_relays(self) -> None:
r"""
Discover relay nodes in the network.
This method queries the network for peers that support the
Circuit Relay v2 protocol.
"""
logger.debug("Starting relay discovery")
try:
# Get connected peers
connected_peers = self.host.get_connected_peers()
logger.debug(
"Checking %d connected peers for relay support", len(connected_peers)
)
# Check each peer if they support the relay protocol
for peer_id in connected_peers:
if peer_id == self.host.get_id():
continue # Skip ourselves
if peer_id in self._discovered_relays:
# Update last seen time for existing relay
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
if await self._supports_relay_protocol(peer_id):
await self._add_relay(peer_id)
# Limit number of relays we track
if len(self._discovered_relays) > self.max_relays:
# Sort by last seen time and keep only the most recent ones
sorted_relays = sorted(
self._discovered_relays.items(),
key=lambda x: x[1].last_seen,
reverse=True,
)
to_remove = sorted_relays[self.max_relays :]
for peer_id, _ in to_remove:
del self._discovered_relays[peer_id]
logger.debug(
"Discovery completed, tracking %d relays", len(self._discovered_relays)
)
except Exception as e:
logger.error("Error during relay discovery: %s", str(e))
async def _supports_relay_protocol(self, peer_id: ID) -> bool:
"""
Check if a peer supports the relay protocol.
Parameters
----------
peer_id : ID
The ID of the peer to check
Returns
-------
bool
True if the peer supports the relay protocol, False otherwise
"""
# Check cache first
if peer_id in self._protocol_cache:
return PROTOCOL_ID in self._protocol_cache[peer_id]
# Method 1: Try peerstore
result = await self._check_via_peerstore(peer_id)
if result is not None:
return result
# Method 2: Try direct stream connection
result = await self._check_via_direct_connection(peer_id)
if result is not None:
return result
# Method 3: Try protocols from mux
result = await self._check_via_mux(peer_id)
if result is not None:
return result
# Default: Cannot determine, assume false
return False
async def _check_via_peerstore(self, peer_id: ID) -> bool | None:
"""Check protocol support via peerstore."""
try:
peerstore = self.host.get_peerstore()
proto_getter = peerstore.get_protocols
if not callable(proto_getter):
return None
try:
# Try to get protocols
proto_result = proto_getter(peer_id)
# Get protocols list
protocols_list = []
if hasattr(proto_result, "__await__"):
protocols_list = await cast(Any, proto_result)
else:
protocols_list = proto_result
# Check result
if protocols_list is not None:
protocols = set(protocols_list)
self._protocol_cache[peer_id] = protocols
return PROTOCOL_ID in protocols
return False
except Exception as e:
logger.debug("Error getting protocols: %s", str(e))
return None
except Exception as e:
logger.debug("Error accessing peerstore: %s", str(e))
return None
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):
stream = await self.host.new_stream(peer_id, [PROTOCOL_ID])
if stream:
await stream.close()
self._protocol_cache[peer_id] = {PROTOCOL_ID}
return True
return False
except Exception as e:
logger.debug(
"Failed to open relay protocol stream to %s: %s", peer_id, str(e)
)
return None
async def _check_via_mux(self, peer_id: ID) -> bool | None:
"""Check protocol support via mux protocols."""
try:
if not (hasattr(self.host, "get_mux") and self.host.get_mux() is not None):
return None
mux = self.host.get_mux()
if not hasattr(mux, "protocols"):
return None
peer_protocols = set()
# Get protocols from mux with proper type safety
available_protocols = []
if hasattr(mux, "get_protocols"):
# Get protocols with proper typing
mux_protocols = mux.get_protocols()
if isinstance(mux_protocols, (list, tuple)):
available_protocols = list(mux_protocols)
for protocol in available_protocols:
try:
with trio.fail_after(2): # Quick check
# Ensure we have a proper protocol object
# Use string representation since we can't use isinstance
is_tprotocol = str(type(protocol)) == str(type(TProtocol))
protocol_obj = (
protocol if is_tprotocol else TProtocol(str(protocol))
)
stream = await self.host.new_stream(peer_id, [protocol_obj])
if stream:
peer_protocols.add(str(protocol_obj))
await stream.close()
except Exception:
pass # Ignore errors when closing the stream
self._protocol_cache[peer_id] = peer_protocols
protocol_str = str(PROTOCOL_ID)
for protocol in peer_protocols:
if protocol == protocol_str:
return True
return False
except Exception as e:
logger.debug("Error checking protocols via mux: %s", str(e))
return None
async def _add_relay(self, peer_id: ID) -> None:
"""
Add a peer as a relay and optionally make a reservation.
Parameters
----------
peer_id : ID
The ID of the peer to add as a relay
"""
now = time.time()
relay_info = RelayInfo(
peer_id=peer_id,
discovered_at=now,
last_seen=now,
)
self._discovered_relays[peer_id] = relay_info
logger.debug("Added relay %s to discovered relays", peer_id)
# If auto-reserve is enabled, make a reservation with this relay
if self.auto_reserve:
await self.make_reservation(peer_id)
async def make_reservation(self, peer_id: ID) -> bool:
"""
Make a reservation with a relay.
Parameters
----------
peer_id : ID
The ID of the relay to make a reservation with
Returns
-------
bool
True if reservation succeeded, False otherwise
"""
if peer_id not in self._discovered_relays:
logger.error("Cannot make reservation with unknown relay %s", peer_id)
return False
stream = None
try:
logger.debug("Making reservation with relay %s", peer_id)
# Open a stream to the relay with timeout
try:
with trio.fail_after(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)
return False
except trio.TooSlowError:
logger.error("Timeout opening stream to relay %s", peer_id)
return False
try:
# Create and send reservation request
request = HopMessage(
type=HopMessage.RESERVE,
peer=self.host.get_id().to_bytes(),
)
with trio.fail_after(STREAM_TIMEOUT):
await stream.write(request.SerializeToString())
# Wait for response
response_bytes = await stream.read()
if not response_bytes:
logger.error("No response received from relay %s", peer_id)
return False
# Parse response
response = HopMessage()
response.ParseFromString(response_bytes)
# Check if reservation was successful
if response.type == HopMessage.RESERVE and response.HasField(
"status"
):
# Access status code directly from protobuf object
status_code = getattr(response.status, "code", StatusCode.OK)
if status_code == StatusCode.OK:
# Update relay info with reservation details
relay_info = self._discovered_relays[peer_id]
relay_info.has_reservation = True
if response.HasField("reservation") and response.HasField(
"limit"
):
relay_info.reservation_expires_at = (
response.reservation.expire
)
relay_info.reservation_data_limit = response.limit.data
logger.debug(
"Successfully made reservation with relay %s", peer_id
)
return True
# Reservation failed
error_message = "Unknown error"
if response.HasField("status"):
# Access message directly from protobuf object
error_message = getattr(response.status, "message", "")
logger.warning(
"Reservation request rejected by relay %s: %s",
peer_id,
error_message,
)
return False
except trio.TooSlowError:
logger.error(
"Timeout during reservation process with relay %s", peer_id
)
return False
except Exception as e:
logger.error("Error making reservation with relay %s: %s", peer_id, str(e))
return False
finally:
# Always close the stream
if stream:
try:
await stream.close()
except Exception:
pass # Ignore errors when closing the stream
return False
async def _cleanup_expired(self) -> None:
"""Clean up expired relays and reservations."""
now = time.time()
to_remove = []
for peer_id, relay_info in self._discovered_relays.items():
# Check if relay hasn't been seen in a while (3x discovery interval)
if now - relay_info.last_seen > self.discovery_interval * 3:
to_remove.append(peer_id)
continue
# Check if reservation has expired
if (
relay_info.has_reservation
and relay_info.reservation_expires_at
and now > relay_info.reservation_expires_at
):
relay_info.has_reservation = False
relay_info.reservation_expires_at = None
relay_info.reservation_data_limit = None
# If auto-reserve is enabled, try to renew
if self.auto_reserve:
await self.make_reservation(peer_id)
# Remove expired relays
for peer_id in to_remove:
del self._discovered_relays[peer_id]
if peer_id in self._protocol_cache:
del self._protocol_cache[peer_id]
def get_relays(self) -> list[ID]:
"""
Get a list of discovered relay peer IDs.
Returns
-------
list[ID]
List of discovered relay peer IDs
"""
return list(self._discovered_relays.keys())
def get_relay_info(self, peer_id: ID) -> RelayInfo | None:
"""
Get information about a specific relay.
Parameters
----------
peer_id : ID
The ID of the relay to get information about
Returns
-------
Optional[RelayInfo]
Information about the relay, or None if not found
"""
return self._discovered_relays.get(peer_id)
def get_relay(self) -> ID | None:
"""
Get a single relay peer ID for connection purposes.
Prioritizes relays with active reservations.
Returns
-------
Optional[ID]
ID of a discovered relay, or None if no relays found
"""
if not self._discovered_relays:
return None
# First try to find a relay with an active reservation
for peer_id, relay_info in self._discovered_relays.items():
if relay_info and relay_info.has_reservation:
return peer_id
return next(iter(self._discovered_relays.keys()), None)

16
libp2p/relay/circuit_v2/pb/__init__.py Normal file
View File

@ -0,0 +1,16 @@
"""
Protocol buffer package for circuit_v2.
Contains generated protobuf code for circuit_v2 relay protocol.
"""
# Import the classes to be accessible directly from the package
from .circuit_pb2 import (
HopMessage,
Limit,
Reservation,
Status,
StopMessage,
)
__all__ = ["HopMessage", "Limit", "Reservation", "Status", "StopMessage"]

55
libp2p/relay/circuit_v2/pb/circuit.proto Normal file
View File

@ -0,0 +1,55 @@
syntax = "proto3";
package circuit.pb.v2;
// Circuit v2 message types
message HopMessage {
enum Type {
RESERVE = 0;
CONNECT = 1;
STATUS = 2;
}
Type type = 1;
bytes peer = 2;
Reservation reservation = 3;
Limit limit = 4;
Status status = 5;
}
message StopMessage {
enum Type {
CONNECT = 0;
STATUS = 1;
}
Type type = 1;
bytes peer = 2;
Status status = 3;
}
message Reservation {
bytes voucher = 1;
bytes signature = 2;
int64 expire = 3;
}
message Limit {
int64 duration = 1;
int64 data = 2;
}
message Status {
enum Code {
OK = 0;
RESERVATION_REFUSED = 100;
RESOURCE_LIMIT_EXCEEDED = 101;
PERMISSION_DENIED = 102;
CONNECTION_FAILED = 200;
DIAL_REFUSED = 201;
STOP_FAILED = 300;
MALFORMED_MESSAGE = 400;
}
Code code = 1;
string message = 2;
}

37
libp2p/relay/circuit_v2/pb/circuit_pb2.py Normal file
View File

@ -0,0 +1,37 @@
# -*- coding: utf-8 -*-
# Generated by the protocol buffer compiler. DO NOT EDIT!
# NO CHECKED-IN PROTOBUF GENCODE
# source: libp2p/relay/circuit_v2/pb/circuit.proto
"""Generated protocol buffer code."""
from google.protobuf.internal import builder as _builder
from google.protobuf import descriptor as _descriptor
from google.protobuf import descriptor_pool as _descriptor_pool
from google.protobuf import symbol_database as _symbol_database
# @@protoc_insertion_point(imports)
_sym_db = _symbol_database.Default()
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n(libp2p/relay/circuit_v2/pb/circuit.proto\x12\rcircuit.pb.v2\"\xf3\x01\n\nHopMessage\x12,\n\x04type\x18\x01 \x01(\x0e\x32\x1e.circuit.pb.v2.HopMessage.Type\x12\x0c\n\x04peer\x18\x02 \x01(\x0c\x12/\n\x0breservation\x18\x03 \x01(\x0b\x32\x1a.circuit.pb.v2.Reservation\x12#\n\x05limit\x18\x04 \x01(\x0b\x32\x14.circuit.pb.v2.Limit\x12%\n\x06status\x18\x05 \x01(\x0b\x32\x15.circuit.pb.v2.Status\",\n\x04Type\x12\x0b\n\x07RESERVE\x10\x00\x12\x0b\n\x07\x43ONNECT\x10\x01\x12\n\n\x06STATUS\x10\x02\"\x92\x01\n\x0bStopMessage\x12-\n\x04type\x18\x01 \x01(\x0e\x32\x1f.circuit.pb.v2.StopMessage.Type\x12\x0c\n\x04peer\x18\x02 \x01(\x0c\x12%\n\x06status\x18\x03 \x01(\x0b\x32\x15.circuit.pb.v2.Status\"\x1f\n\x04Type\x12\x0b\n\x07\x43ONNECT\x10\x00\x12\n\n\x06STATUS\x10\x01\"A\n\x0bReservation\x12\x0f\n\x07voucher\x18\x01 \x01(\x0c\x12\x11\n\tsignature\x18\x02 \x01(\x0c\x12\x0e\n\x06\x65xpire\x18\x03 \x01(\x03\"\'\n\x05Limit\x12\x10\n\x08\x64uration\x18\x01 \x01(\x03\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x03\"\xf6\x01\n\x06Status\x12(\n\x04\x63ode\x18\x01 \x01(\x0e\x32\x1a.circuit.pb.v2.Status.Code\x12\x0f\n\x07message\x18\x02 \x01(\t\"\xb0\x01\n\x04\x43ode\x12\x06\n\x02OK\x10\x00\x12\x17\n\x13RESERVATION_REFUSED\x10\x64\x12\x1b\n\x17RESOURCE_LIMIT_EXCEEDED\x10\x65\x12\x15\n\x11PERMISSION_DENIED\x10\x66\x12\x16\n\x11\x43ONNECTION_FAILED\x10\xc8\x01\x12\x11\n\x0c\x44IAL_REFUSED\x10\xc9\x01\x12\x10\n\x0bSTOP_FAILED\x10\xac\x02\x12\x16\n\x11MALFORMED_MESSAGE\x10\x90\x03\x62\x06proto3')
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'libp2p.relay.circuit_v2.pb.circuit_pb2', globals())
if _descriptor._USE_C_DESCRIPTORS == False:
DESCRIPTOR._options = None
_HOPMESSAGE._serialized_start=60
_HOPMESSAGE._serialized_end=303
_HOPMESSAGE_TYPE._serialized_start=259
_HOPMESSAGE_TYPE._serialized_end=303
_STOPMESSAGE._serialized_start=306
_STOPMESSAGE._serialized_end=452
_STOPMESSAGE_TYPE._serialized_start=421
_STOPMESSAGE_TYPE._serialized_end=452
_RESERVATION._serialized_start=454
_RESERVATION._serialized_end=519
_LIMIT._serialized_start=521
_LIMIT._serialized_end=560
_STATUS._serialized_start=563
_STATUS._serialized_end=809
_STATUS_CODE._serialized_start=633
_STATUS_CODE._serialized_end=809
# @@protoc_insertion_point(module_scope)

184
libp2p/relay/circuit_v2/pb/circuit_pb2.pyi Normal file
View File

@ -0,0 +1,184 @@
"""
@generated by mypy-protobuf. Do not edit manually!
isort:skip_file
"""
import builtins
import google.protobuf.descriptor
import google.protobuf.internal.enum_type_wrapper
import google.protobuf.message
import sys
import typing
if sys.version_info >= (3, 10):
import typing as typing_extensions
else:
import typing_extensions
DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
@typing.final
class HopMessage(google.protobuf.message.Message):
"""Circuit v2 message types"""
DESCRIPTOR: google.protobuf.descriptor.Descriptor
class _Type:
ValueType = typing.NewType("ValueType", builtins.int)
V: typing_extensions.TypeAlias = ValueType
class _TypeEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[HopMessage._Type.ValueType], builtins.type):
DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
RESERVE: HopMessage._Type.ValueType # 0
CONNECT: HopMessage._Type.ValueType # 1
STATUS: HopMessage._Type.ValueType # 2
class Type(_Type, metaclass=_TypeEnumTypeWrapper): ...
RESERVE: HopMessage.Type.ValueType # 0
CONNECT: HopMessage.Type.ValueType # 1
STATUS: HopMessage.Type.ValueType # 2
TYPE_FIELD_NUMBER: builtins.int
PEER_FIELD_NUMBER: builtins.int
RESERVATION_FIELD_NUMBER: builtins.int
LIMIT_FIELD_NUMBER: builtins.int
STATUS_FIELD_NUMBER: builtins.int
type: global___HopMessage.Type.ValueType
peer: builtins.bytes
@property
def reservation(self) -> global___Reservation: ...
@property
def limit(self) -> global___Limit: ...
@property
def status(self) -> global___Status: ...
def __init__(
self,
*,
type: global___HopMessage.Type.ValueType = ...,
peer: builtins.bytes = ...,
reservation: global___Reservation | None = ...,
limit: global___Limit | None = ...,
status: global___Status | None = ...,
) -> None: ...
def HasField(self, field_name: typing.Literal["limit", b"limit", "reservation", b"reservation", "status", b"status"]) -> builtins.bool: ...
def ClearField(self, field_name: typing.Literal["limit", b"limit", "peer", b"peer", "reservation", b"reservation", "status", b"status", "type", b"type"]) -> None: ...
global___HopMessage = HopMessage
@typing.final
class StopMessage(google.protobuf.message.Message):
DESCRIPTOR: google.protobuf.descriptor.Descriptor
class _Type:
ValueType = typing.NewType("ValueType", builtins.int)
V: typing_extensions.TypeAlias = ValueType
class _TypeEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StopMessage._Type.ValueType], builtins.type):
DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
CONNECT: StopMessage._Type.ValueType # 0
STATUS: StopMessage._Type.ValueType # 1
class Type(_Type, metaclass=_TypeEnumTypeWrapper): ...
CONNECT: StopMessage.Type.ValueType # 0
STATUS: StopMessage.Type.ValueType # 1
TYPE_FIELD_NUMBER: builtins.int
PEER_FIELD_NUMBER: builtins.int
STATUS_FIELD_NUMBER: builtins.int
type: global___StopMessage.Type.ValueType
peer: builtins.bytes
@property
def status(self) -> global___Status: ...
def __init__(
self,
*,
type: global___StopMessage.Type.ValueType = ...,
peer: builtins.bytes = ...,
status: global___Status | None = ...,
) -> None: ...
def HasField(self, field_name: typing.Literal["status", b"status"]) -> builtins.bool: ...
def ClearField(self, field_name: typing.Literal["peer", b"peer", "status", b"status", "type", b"type"]) -> None: ...
global___StopMessage = StopMessage
@typing.final
class Reservation(google.protobuf.message.Message):
DESCRIPTOR: google.protobuf.descriptor.Descriptor
VOUCHER_FIELD_NUMBER: builtins.int
SIGNATURE_FIELD_NUMBER: builtins.int
EXPIRE_FIELD_NUMBER: builtins.int
voucher: builtins.bytes
signature: builtins.bytes
expire: builtins.int
def __init__(
self,
*,
voucher: builtins.bytes = ...,
signature: builtins.bytes = ...,
expire: builtins.int = ...,
) -> None: ...
def ClearField(self, field_name: typing.Literal["expire", b"expire", "signature", b"signature", "voucher", b"voucher"]) -> None: ...
global___Reservation = Reservation
@typing.final
class Limit(google.protobuf.message.Message):
DESCRIPTOR: google.protobuf.descriptor.Descriptor
DURATION_FIELD_NUMBER: builtins.int
DATA_FIELD_NUMBER: builtins.int
duration: builtins.int
data: builtins.int
def __init__(
self,
*,
duration: builtins.int = ...,
data: builtins.int = ...,
) -> None: ...
def ClearField(self, field_name: typing.Literal["data", b"data", "duration", b"duration"]) -> None: ...
global___Limit = Limit
@typing.final
class Status(google.protobuf.message.Message):
DESCRIPTOR: google.protobuf.descriptor.Descriptor
class _Code:
ValueType = typing.NewType("ValueType", builtins.int)
V: typing_extensions.TypeAlias = ValueType
class _CodeEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[Status._Code.ValueType], builtins.type):
DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
OK: Status._Code.ValueType # 0
RESERVATION_REFUSED: Status._Code.ValueType # 100
RESOURCE_LIMIT_EXCEEDED: Status._Code.ValueType # 101
PERMISSION_DENIED: Status._Code.ValueType # 102
CONNECTION_FAILED: Status._Code.ValueType # 200
DIAL_REFUSED: Status._Code.ValueType # 201
STOP_FAILED: Status._Code.ValueType # 300
MALFORMED_MESSAGE: Status._Code.ValueType # 400
class Code(_Code, metaclass=_CodeEnumTypeWrapper): ...
OK: Status.Code.ValueType # 0
RESERVATION_REFUSED: Status.Code.ValueType # 100
RESOURCE_LIMIT_EXCEEDED: Status.Code.ValueType # 101
PERMISSION_DENIED: Status.Code.ValueType # 102
CONNECTION_FAILED: Status.Code.ValueType # 200
DIAL_REFUSED: Status.Code.ValueType # 201
STOP_FAILED: Status.Code.ValueType # 300
MALFORMED_MESSAGE: Status.Code.ValueType # 400
CODE_FIELD_NUMBER: builtins.int
MESSAGE_FIELD_NUMBER: builtins.int
code: global___Status.Code.ValueType
message: builtins.str
def __init__(
self,
*,
code: global___Status.Code.ValueType = ...,
message: builtins.str = ...,
) -> None: ...
def ClearField(self, field_name: typing.Literal["code", b"code", "message", b"message"]) -> None: ...
global___Status = Status

800
libp2p/relay/circuit_v2/protocol.py Normal file
View File

@ -0,0 +1,800 @@
"""
Circuit Relay v2 protocol implementation.
This module implements the Circuit Relay v2 protocol as specified in:
https://github.com/libp2p/specs/blob/master/relay/circuit-v2.md
"""
import logging
import time
from typing import (
Any,
Protocol as TypingProtocol,
cast,
runtime_checkable,
)
import trio
from libp2p.abc import (
IHost,
INetStream,
)
from libp2p.custom_types import (
TProtocol,
)
from libp2p.io.abc import (
ReadWriteCloser,
)
from libp2p.peer.id import (
ID,
)
from libp2p.stream_muxer.mplex.exceptions import (
MplexStreamEOF,
MplexStreamReset,
)
from libp2p.tools.async_service import (
Service,
)
from .pb.circuit_pb2 import (
HopMessage,
Limit,
Reservation,
Status as PbStatus,
StopMessage,
)
from .protocol_buffer import (
StatusCode,
create_status,
)
from .resources import (
RelayLimits,
RelayResourceManager,
)
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")
# 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,
)
# Stream operation timeouts
STREAM_READ_TIMEOUT = 15 # seconds
STREAM_WRITE_TIMEOUT = 15 # seconds
STREAM_CLOSE_TIMEOUT = 10 # seconds
MAX_READ_RETRIES = 5 # Maximum number of read retries
# Extended interfaces for type checking
@runtime_checkable
class IHostWithStreamHandlers(TypingProtocol):
"""Extended host interface with stream handler methods."""
def remove_stream_handler(self, protocol_id: TProtocol) -> None:
"""Remove a stream handler for a protocol."""
...
@runtime_checkable
class INetStreamWithExtras(TypingProtocol):
"""Extended net stream interface with additional methods."""
def get_remote_peer_id(self) -> ID:
"""Get the remote peer ID."""
...
def is_open(self) -> bool:
"""Check if the stream is open."""
...
def is_closed(self) -> bool:
"""Check if the stream is closed."""
...
class CircuitV2Protocol(Service):
"""
CircuitV2Protocol implements the Circuit Relay v2 protocol.
This protocol allows peers to establish connections through relay nodes
when direct connections are not possible (e.g., due to NAT).
"""
def __init__(
self,
host: IHost,
limits: RelayLimits | None = None,
allow_hop: bool = False,
) -> None:
"""
Initialize a Circuit Relay v2 protocol instance.
Parameters
----------
host : IHost
The libp2p host instance
limits : RelayLimits | None
Resource limits for the relay
allow_hop : bool
Whether to allow this node to act as a relay
"""
self.host = host
self.limits = limits or DEFAULT_RELAY_LIMITS
self.allow_hop = allow_hop
self.resource_manager = RelayResourceManager(self.limits)
self._active_relays: dict[ID, tuple[INetStream, INetStream | None]] = {}
self.event_started = trio.Event()
async def run(self, *, task_status: Any = trio.TASK_STATUS_IGNORED) -> None:
"""Run the protocol service."""
try:
# Register protocol handlers
if self.allow_hop:
logger.debug("Registering stream handlers for relay protocol")
self.host.set_stream_handler(PROTOCOL_ID, self._handle_hop_stream)
self.host.set_stream_handler(STOP_PROTOCOL_ID, self._handle_stop_stream)
logger.debug("Stream handlers registered successfully")
# Signal that we're ready
self.event_started.set()
task_status.started()
logger.debug("Protocol service started")
# Wait for service to be stopped
await self.manager.wait_finished()
finally:
# Clean up any active relay connections
for src_stream, dst_stream in self._active_relays.values():
await self._close_stream(src_stream)
await self._close_stream(dst_stream)
self._active_relays.clear()
# Unregister protocol handlers
if self.allow_hop:
try:
# Cast host to extended interface with remove_stream_handler
host_with_handlers = cast(IHostWithStreamHandlers, self.host)
host_with_handlers.remove_stream_handler(PROTOCOL_ID)
host_with_handlers.remove_stream_handler(STOP_PROTOCOL_ID)
except Exception as e:
logger.error("Error unregistering stream handlers: %s", str(e))
async def _close_stream(self, stream: INetStream | None) -> None:
"""Helper function to safely close a stream."""
if stream is None:
return
try:
with trio.fail_after(STREAM_CLOSE_TIMEOUT):
await stream.close()
except Exception:
try:
await stream.reset()
except Exception:
pass
async def _read_stream_with_retry(
self,
stream: INetStream,
max_retries: int = MAX_READ_RETRIES,
) -> bytes | None:
"""
Helper function to read from a stream with retries.
Parameters
----------
stream : INetStream
The stream to read from
max_retries : int
Maximum number of read retries
Returns
-------
Optional[bytes]
The data read from the stream, or None if the stream is closed/reset
Raises
------
trio.TooSlowError
If read timeout occurs after all retries
Exception
For other unexpected errors
"""
retries = 0
last_error: Any = None
backoff_time = 0.2 # Base backoff time in seconds
while retries < max_retries:
try:
with trio.fail_after(STREAM_READ_TIMEOUT):
# Try reading with timeout
logger.debug(
"Attempting to read from stream (attempt %d/%d)",
retries + 1,
max_retries,
)
data = await stream.read()
if not data: # EOF
logger.debug("Stream EOF detected")
return None
logger.debug("Successfully read %d bytes from stream", len(data))
return data
except trio.WouldBlock:
# Just retry immediately if we would block
retries += 1
logger.debug(
"Stream would block (attempt %d/%d), retrying...",
retries,
max_retries,
)
await trio.sleep(backoff_time * retries) # Increased backoff time
continue
except (MplexStreamEOF, MplexStreamReset):
# Stream closed/reset - no point retrying
logger.debug("Stream closed/reset during read")
return None
except trio.TooSlowError as e:
last_error = e
retries += 1
logger.debug(
"Read timeout (attempt %d/%d), retrying...", retries, max_retries
)
if retries < max_retries:
# Wait longer before retry with increasing backoff
await trio.sleep(backoff_time * retries) # Increased backoff
continue
except Exception as e:
logger.error("Unexpected error reading from stream: %s", str(e))
last_error = e
retries += 1
if retries < max_retries:
await trio.sleep(backoff_time * retries) # Increased backoff
continue
raise
if last_error:
if isinstance(last_error, trio.TooSlowError):
logger.error("Read timed out after %d retries", max_retries)
raise last_error
return None
async def _handle_hop_stream(self, stream: INetStream) -> None:
"""
Handle incoming HOP streams.
This handler processes relay requests from other peers.
"""
try:
# Try to get peer ID first
try:
# Cast to extended interface with get_remote_peer_id
stream_with_peer_id = cast(INetStreamWithExtras, stream)
remote_peer_id = stream_with_peer_id.get_remote_peer_id()
remote_id = str(remote_peer_id)
except Exception:
# Fall back to address if peer ID not available
remote_addr = stream.get_remote_address()
remote_id = f"peer at {remote_addr}" if remote_addr else "unknown peer"
logger.debug("Handling hop stream from %s", remote_id)
# First, handle the read timeout gracefully
try:
with trio.fail_after(
STREAM_READ_TIMEOUT * 2
): # Double the timeout for reading
msg_bytes = await stream.read()
if not msg_bytes:
logger.error(
"Empty read from stream from %s",
remote_id,
)
# Create a proto Status directly
pb_status = PbStatus()
pb_status.code = cast(Any, int(StatusCode.MALFORMED_MESSAGE))
pb_status.message = "Empty message received"
response = HopMessage(
type=HopMessage.STATUS,
status=pb_status,
)
await stream.write(response.SerializeToString())
await trio.sleep(0.5) # Longer wait to ensure message is sent
return
except trio.TooSlowError:
logger.error(
"Timeout reading from hop stream from %s",
remote_id,
)
# Create a proto Status directly
pb_status = PbStatus()
pb_status.code = cast(Any, int(StatusCode.CONNECTION_FAILED))
pb_status.message = "Stream read timeout"
response = HopMessage(
type=HopMessage.STATUS,
status=pb_status,
)
await stream.write(response.SerializeToString())
await trio.sleep(0.5) # Longer wait to ensure the message is sent
return
except Exception as e:
logger.error(
"Error reading from hop stream from %s: %s",
remote_id,
str(e),
)
# Create a proto Status directly
pb_status = PbStatus()
pb_status.code = cast(Any, int(StatusCode.MALFORMED_MESSAGE))
pb_status.message = f"Read error: {str(e)}"
response = HopMessage(
type=HopMessage.STATUS,
status=pb_status,
)
await stream.write(response.SerializeToString())
await trio.sleep(0.5) # Longer wait to ensure the message is sent
return
# Parse the message
try:
hop_msg = HopMessage()
hop_msg.ParseFromString(msg_bytes)
except Exception as e:
logger.error(
"Error parsing hop message from %s: %s",
remote_id,
str(e),
)
# Create a proto Status directly
pb_status = PbStatus()
pb_status.code = cast(Any, int(StatusCode.MALFORMED_MESSAGE))
pb_status.message = f"Parse error: {str(e)}"
response = HopMessage(
type=HopMessage.STATUS,
status=pb_status,
)
await stream.write(response.SerializeToString())
await trio.sleep(0.5) # Longer wait to ensure the message is sent
return
# Process based on message type
if hop_msg.type == HopMessage.RESERVE:
logger.debug("Handling RESERVE message from %s", remote_id)
await self._handle_reserve(stream, hop_msg)
# For RESERVE requests, let the client close the stream
return
elif hop_msg.type == HopMessage.CONNECT:
logger.debug("Handling CONNECT message from %s", remote_id)
await self._handle_connect(stream, hop_msg)
else:
logger.error("Invalid message type %d from %s", hop_msg.type, remote_id)
# Send a nice error response using _send_status method
await self._send_status(
stream,
StatusCode.MALFORMED_MESSAGE,
f"Invalid message type: {hop_msg.type}",
)
except Exception as e:
logger.error(
"Unexpected error handling hop stream from %s: %s", remote_id, str(e)
)
try:
# Send a nice error response using _send_status method
await self._send_status(
stream,
StatusCode.MALFORMED_MESSAGE,
f"Internal error: {str(e)}",
)
except Exception as e2:
logger.error(
"Failed to send error response to %s: %s", remote_id, str(e2)
)
async def _handle_stop_stream(self, stream: INetStream) -> None:
"""
Handle incoming STOP streams.
This handler processes incoming relay connections from the destination side.
"""
try:
# Read the incoming message with timeout
with trio.fail_after(STREAM_READ_TIMEOUT):
msg_bytes = await stream.read()
stop_msg = StopMessage()
stop_msg.ParseFromString(msg_bytes)
if stop_msg.type != StopMessage.CONNECT:
# Use direct attribute access to create status object for error response
await self._send_stop_status(
stream,
StatusCode.MALFORMED_MESSAGE,
"Invalid message type",
)
await self._close_stream(stream)
return
# Get the source stream from active relays
peer_id = ID(stop_msg.peer)
if peer_id not in self._active_relays:
# Use direct attribute access to create status object for error response
await self._send_stop_status(
stream,
StatusCode.CONNECTION_FAILED,
"No pending relay connection",
)
await self._close_stream(stream)
return
src_stream, _ = self._active_relays[peer_id]
self._active_relays[peer_id] = (src_stream, stream)
# Send success status to both sides
await self._send_status(
src_stream,
StatusCode.OK,
"Connection established",
)
await self._send_stop_status(
stream,
StatusCode.OK,
"Connection established",
)
# 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)
except trio.TooSlowError:
logger.error("Timeout reading from stop stream")
await self._send_stop_status(
stream,
StatusCode.CONNECTION_FAILED,
"Stream read timeout",
)
await self._close_stream(stream)
except Exception as e:
logger.error("Error handling stop stream: %s", str(e))
try:
await self._send_stop_status(
stream,
StatusCode.MALFORMED_MESSAGE,
str(e),
)
await self._close_stream(stream)
except Exception:
pass
async def _handle_reserve(self, stream: INetStream, msg: Any) -> None:
"""Handle a reservation request."""
peer_id = None
try:
peer_id = ID(msg.peer)
logger.debug("Handling reservation request from peer %s", peer_id)
# Check if we can accept more reservations
if not self.resource_manager.can_accept_reservation(peer_id):
logger.debug("Reservation limit exceeded for peer %s", peer_id)
# Send status message with STATUS type
status = create_status(
code=StatusCode.RESOURCE_LIMIT_EXCEEDED,
message="Reservation limit exceeded",
)
status_msg = HopMessage(
type=HopMessage.STATUS,
status=status.to_pb(),
)
await stream.write(status_msg.SerializeToString())
return
# Accept reservation
logger.debug("Accepting reservation from peer %s", peer_id)
ttl = self.resource_manager.reserve(peer_id)
# Send reservation success response
with trio.fail_after(STREAM_WRITE_TIMEOUT):
status = create_status(
code=StatusCode.OK, message="Reservation accepted"
)
response = HopMessage(
type=HopMessage.STATUS,
status=status.to_pb(),
reservation=Reservation(
expire=int(time.time() + ttl),
voucher=b"", # We don't use vouchers yet
signature=b"", # We don't use signatures yet
),
limit=Limit(
duration=self.limits.duration,
data=self.limits.data,
),
)
# Log the response message details for debugging
logger.debug(
"Sending reservation response: type=%s, status=%s, ttl=%d",
response.type,
getattr(response.status, "code", "unknown"),
ttl,
)
# Send the response with increased timeout
await stream.write(response.SerializeToString())
# Add a small wait to ensure the message is fully sent
await trio.sleep(0.1)
logger.debug("Reservation response sent successfully")
except Exception as e:
logger.error("Error handling reservation request: %s", str(e))
if cast(INetStreamWithExtras, stream).is_open():
try:
# Send error response
await self._send_status(
stream,
StatusCode.INTERNAL_ERROR,
f"Failed to process reservation: {str(e)}",
)
except Exception as send_err:
logger.error("Failed to send error response: %s", str(send_err))
finally:
# Always close the stream when done with reservation
if cast(INetStreamWithExtras, stream).is_open():
try:
with trio.fail_after(STREAM_CLOSE_TIMEOUT):
await stream.close()
except Exception as close_err:
logger.error("Error closing stream: %s", str(close_err))
async def _handle_connect(self, stream: INetStream, msg: Any) -> None:
"""Handle a connect request."""
peer_id = ID(msg.peer)
dst_stream: INetStream | None = None
# Verify reservation if provided
if msg.HasField("reservation"):
if not self.resource_manager.verify_reservation(peer_id, msg.reservation):
await self._send_status(
stream,
StatusCode.PERMISSION_DENIED,
"Invalid reservation",
)
await stream.reset()
return
# Check resource limits
if not self.resource_manager.can_accept_connection(peer_id):
await self._send_status(
stream,
StatusCode.RESOURCE_LIMIT_EXCEEDED,
"Connection limit exceeded",
)
await stream.reset()
return
try:
# Store the source stream with properly typed None
self._active_relays[peer_id] = (stream, None)
# Try to connect to the destination with timeout
with trio.fail_after(STREAM_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")
# Send STOP CONNECT message
stop_msg = StopMessage(
type=StopMessage.CONNECT,
# Cast to extended interface with get_remote_peer_id
peer=cast(INetStreamWithExtras, stream)
.get_remote_peer_id()
.to_bytes(),
)
await dst_stream.write(stop_msg.SerializeToString())
# Wait for response from destination
resp_bytes = await dst_stream.read()
resp = StopMessage()
resp.ParseFromString(resp_bytes)
# Handle status attributes from the response
if resp.HasField("status"):
# Get code and message attributes with defaults
status_code = getattr(resp.status, "code", StatusCode.OK)
# Get message with default
status_msg = getattr(resp.status, "message", "Unknown error")
else:
status_code = StatusCode.OK
status_msg = "No status provided"
if status_code != StatusCode.OK:
raise ConnectionError(
f"Destination rejected connection: {status_msg}"
)
# Update active relays with destination stream
self._active_relays[peer_id] = (stream, dst_stream)
# Update reservation connection count
reservation = self.resource_manager._reservations.get(peer_id)
if reservation:
reservation.active_connections += 1
# Send success status
await self._send_status(
stream,
StatusCode.OK,
"Connection established",
)
# 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)
except (trio.TooSlowError, ConnectionError) as e:
logger.error("Error establishing relay connection: %s", str(e))
await self._send_status(
stream,
StatusCode.CONNECTION_FAILED,
str(e),
)
if peer_id in self._active_relays:
del self._active_relays[peer_id]
# Clean up reservation connection count on failure
reservation = self.resource_manager._reservations.get(peer_id)
if reservation:
reservation.active_connections -= 1
await stream.reset()
if dst_stream and not cast(INetStreamWithExtras, dst_stream).is_closed():
await dst_stream.reset()
except Exception as e:
logger.error("Unexpected error in connect handler: %s", str(e))
await self._send_status(
stream,
StatusCode.CONNECTION_FAILED,
"Internal error",
)
if peer_id in self._active_relays:
del self._active_relays[peer_id]
await stream.reset()
if dst_stream and not cast(INetStreamWithExtras, dst_stream).is_closed():
await dst_stream.reset()
async def _relay_data(
self,
src_stream: INetStream,
dst_stream: INetStream,
peer_id: ID,
) -> None:
"""
Relay data between two streams.
Parameters
----------
src_stream : INetStream
Source stream to read from
dst_stream : INetStream
Destination stream to write to
peer_id : ID
ID of the peer being relayed
"""
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")
break
# Write data with timeout
try:
with trio.fail_after(STREAM_WRITE_TIMEOUT):
await dst_stream.write(data)
except trio.TooSlowError:
logger.error("Timeout writing to destination stream")
break
except Exception as e:
logger.error("Error writing to destination stream: %s", str(e))
break
# Update resource usage
reservation = self.resource_manager._reservations.get(peer_id)
if reservation:
reservation.data_used += len(data)
if reservation.data_used >= reservation.limits.data:
logger.warning("Data limit exceeded for peer %s", peer_id)
break
except Exception as e:
logger.error("Error relaying data: %s", str(e))
finally:
# Clean up streams and remove from active relays
await src_stream.reset()
await dst_stream.reset()
if peer_id in self._active_relays:
del self._active_relays[peer_id]
async def _send_status(
self,
stream: ReadWriteCloser,
code: int,
message: str,
) -> None:
"""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
# Create a proto Status directly
pb_status = PbStatus()
pb_status.code = cast(
Any, int(code)
) # Cast to Any to avoid type errors
pb_status.message = message
status_msg = HopMessage(
type=HopMessage.STATUS,
status=pb_status,
)
msg_bytes = status_msg.SerializeToString()
logger.debug("Status message serialized (%d bytes)", len(msg_bytes))
await stream.write(msg_bytes)
logger.debug("Status message sent, waiting for processing")
# Wait longer to ensure the message is sent
await trio.sleep(1.5)
logger.debug("Status message sending completed")
except trio.TooSlowError:
logger.error(
"Timeout sending status message: code=%s, message=%s", code, message
)
except Exception as e:
logger.error("Error sending status message: %s", str(e))
async def _send_stop_status(
self,
stream: ReadWriteCloser,
code: int,
message: str,
) -> None:
"""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
# Create a proto Status directly
pb_status = PbStatus()
pb_status.code = cast(
Any, int(code)
) # Cast to Any to avoid type errors
pb_status.message = message
status_msg = StopMessage(
type=StopMessage.STATUS,
status=pb_status,
)
await stream.write(status_msg.SerializeToString())
await trio.sleep(0.5) # Ensure message is sent
except Exception as e:
logger.error("Error sending stop status message: %s", str(e))

55
libp2p/relay/circuit_v2/protocol_buffer.py Normal file
View File

@ -0,0 +1,55 @@
"""
Protocol buffer wrapper classes for Circuit Relay v2.
This module provides wrapper classes for protocol buffer generated objects
to make them easier to work with in type-checked code.
"""
from enum import (
IntEnum,
)
from typing import (
Any,
)
from .pb.circuit_pb2 import Status as PbStatus
# Define Status codes as an Enum for better type safety and organization
class StatusCode(IntEnum):
OK = 0
RESERVATION_REFUSED = 100
RESOURCE_LIMIT_EXCEEDED = 101
PERMISSION_DENIED = 102
CONNECTION_FAILED = 200
DIAL_REFUSED = 201
STOP_FAILED = 300
MALFORMED_MESSAGE = 400
INTERNAL_ERROR = 500
def create_status(code: int = StatusCode.OK, message: str = "") -> Any:
"""
Create a protocol buffer Status object.
Parameters
----------
code : int
The status code
message : str
The status message
Returns
-------
Any
The protocol buffer Status object
"""
# Create status object
pb_obj = PbStatus()
# Convert the integer status code to the protobuf enum value type
pb_obj.code = PbStatus.Code.ValueType(code)
pb_obj.message = message
return pb_obj

254
libp2p/relay/circuit_v2/resources.py Normal file
View File

@ -0,0 +1,254 @@
"""
Resource management for Circuit Relay v2.
This module handles managing resources for relay operations,
including reservations and connection limits.
"""
from dataclasses import (
dataclass,
)
import hashlib
import os
import time
from libp2p.peer.id import (
ID,
)
# Import the protobuf definitions
from .pb.circuit_pb2 import Reservation as PbReservation
@dataclass
class RelayLimits:
"""Configuration for relay resource limits."""
duration: int # Maximum duration of a relay connection in seconds
data: int # Maximum data transfer allowed in bytes
max_circuit_conns: int # Maximum number of concurrent circuit connections
max_reservations: int # Maximum number of active reservations
class Reservation:
"""Represents a relay reservation."""
def __init__(self, peer_id: ID, limits: RelayLimits):
"""
Initialize a new reservation.
Parameters
----------
peer_id : ID
The peer ID this reservation is for
limits : RelayLimits
The resource limits for this reservation
"""
self.peer_id = peer_id
self.limits = limits
self.created_at = time.time()
self.expires_at = self.created_at + limits.duration
self.data_used = 0
self.active_connections = 0
self.voucher = self._generate_voucher()
def _generate_voucher(self) -> bytes:
"""
Generate a unique cryptographically secure voucher for this reservation.
Returns
-------
bytes
A secure voucher token
"""
# Create a random token using a combination of:
# - Random bytes for unpredictability
# - 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()
peer_bytes = self.peer_id.to_bytes()
# Combine all elements and hash them
h = hashlib.sha256()
h.update(random_bytes)
h.update(timestamp)
h.update(peer_bytes)
return h.digest()
def is_expired(self) -> bool:
"""Check if the reservation has expired."""
return time.time() > self.expires_at
def can_accept_connection(self) -> bool:
"""Check if a new connection can be accepted."""
return (
not self.is_expired()
and self.active_connections < self.limits.max_circuit_conns
and self.data_used < self.limits.data
)
def to_proto(self) -> PbReservation:
"""Convert the reservation to its protobuf representation."""
# TODO: For production use, implement proper signature generation
# The signature should be created by signing the voucher with the
# peer's private key. The current implementation with an empty signature
# is intended for development and testing only.
return PbReservation(
expire=int(self.expires_at),
voucher=self.voucher,
signature=b"",
)
class RelayResourceManager:
"""
Manages resources and reservations for relay operations.
This class handles:
- Tracking active reservations
- Enforcing resource limits
- Managing connection quotas
"""
def __init__(self, limits: RelayLimits):
"""
Initialize the resource manager.
Parameters
----------
limits : RelayLimits
The resource limits to enforce
"""
self.limits = limits
self._reservations: dict[ID, Reservation] = {}
def can_accept_reservation(self, peer_id: ID) -> bool:
"""
Check if a new reservation can be accepted for the given peer.
Parameters
----------
peer_id : ID
The peer ID requesting the reservation
Returns
-------
bool
True if the reservation can be accepted
"""
# Clean expired reservations
self._clean_expired()
# Check if peer already has a valid reservation
existing = self._reservations.get(peer_id)
if existing and not existing.is_expired():
return True
# Check if we're at the reservation limit
return len(self._reservations) < self.limits.max_reservations
def create_reservation(self, peer_id: ID) -> Reservation:
"""
Create a new reservation for the given peer.
Parameters
----------
peer_id : ID
The peer ID to create the reservation for
Returns
-------
Reservation
The newly created reservation
"""
reservation = Reservation(peer_id, self.limits)
self._reservations[peer_id] = reservation
return reservation
def verify_reservation(self, peer_id: ID, proto_res: PbReservation) -> bool:
"""
Verify a reservation from a protobuf message.
Parameters
----------
peer_id : ID
The peer ID the reservation is for
proto_res : PbReservation
The protobuf reservation message
Returns
-------
bool
True if the reservation is valid
"""
# TODO: Implement voucher and signature verification
reservation = self._reservations.get(peer_id)
return (
reservation is not None
and not reservation.is_expired()
and reservation.expires_at == proto_res.expire
)
def can_accept_connection(self, peer_id: ID) -> bool:
"""
Check if a new connection can be accepted for the given peer.
Parameters
----------
peer_id : ID
The peer ID requesting the connection
Returns
-------
bool
True if the connection can be accepted
"""
reservation = self._reservations.get(peer_id)
return reservation is not None and reservation.can_accept_connection()
def _clean_expired(self) -> None:
"""Remove expired reservations."""
now = time.time()
expired = [
peer_id
for peer_id, res in self._reservations.items()
if now > res.expires_at
]
for peer_id in expired:
del self._reservations[peer_id]
def reserve(self, peer_id: ID) -> int:
"""
Create or update a reservation for a peer and return the TTL.
Parameters
----------
peer_id : ID
The peer ID to reserve for
Returns
-------
int
The TTL of the reservation in seconds
"""
# Check for existing reservation
existing = self._reservations.get(peer_id)
if existing and not existing.is_expired():
# Return remaining time for existing reservation
remaining = max(0, int(existing.expires_at - time.time()))
return remaining
# Create new reservation
self.create_reservation(peer_id)
return self.limits.duration

427
libp2p/relay/circuit_v2/transport.py Normal file
View File

@ -0,0 +1,427 @@
"""
Transport implementation for Circuit Relay v2.
This module implements the transport layer for Circuit Relay v2,
allowing peers to establish connections through relay nodes.
"""
from collections.abc import Awaitable, Callable
import logging
import multiaddr
import trio
from libp2p.abc import (
IHost,
IListener,
INetStream,
ITransport,
ReadWriteCloser,
)
from libp2p.network.connection.raw_connection import (
RawConnection,
)
from libp2p.peer.id import (
ID,
)
from libp2p.peer.peerinfo import (
PeerInfo,
)
from libp2p.tools.async_service import (
Service,
)
from .config import (
ClientConfig,
RelayConfig,
)
from .discovery import (
RelayDiscovery,
)
from .pb.circuit_pb2 import (
HopMessage,
StopMessage,
)
from .protocol import (
PROTOCOL_ID,
CircuitV2Protocol,
)
from .protocol_buffer import (
StatusCode,
)
logger = logging.getLogger("libp2p.relay.circuit_v2.transport")
class CircuitV2Transport(ITransport):
"""
CircuitV2Transport implements the transport interface for Circuit Relay v2.
This transport allows peers to establish connections through relay nodes
when direct connections are not possible.
"""
def __init__(
self,
host: IHost,
protocol: CircuitV2Protocol,
config: RelayConfig,
) -> None:
"""
Initialize the Circuit v2 transport.
Parameters
----------
host : IHost
The libp2p host this transport is running on
protocol : CircuitV2Protocol
The Circuit v2 protocol instance
config : RelayConfig
Relay configuration
"""
self.host = host
self.protocol = protocol
self.config = config
self.client_config = ClientConfig()
self.discovery = RelayDiscovery(
host=host,
auto_reserve=config.enable_client,
discovery_interval=config.discovery_interval,
max_relays=config.max_relays,
)
async def dial(
self,
maddr: multiaddr.Multiaddr,
) -> RawConnection:
"""
Dial a peer using the multiaddr.
Parameters
----------
maddr : multiaddr.Multiaddr
The multiaddr to dial
Returns
-------
RawConnection
The established connection
Raises
------
ConnectionError
If the connection cannot be established
"""
# Extract peer ID from multiaddr - P_P2P code is 0x01A5 (421)
peer_id_str = maddr.value_for_protocol("p2p")
if not peer_id_str:
raise ConnectionError("Multiaddr does not contain peer ID")
peer_id = ID.from_base58(peer_id_str)
peer_info = PeerInfo(peer_id, [maddr])
# Use the internal dial_peer_info method
return await self.dial_peer_info(peer_info)
async def dial_peer_info(
self,
peer_info: PeerInfo,
*,
relay_peer_id: ID | None = None,
) -> RawConnection:
"""
Dial a peer through a relay.
Parameters
----------
peer_info : PeerInfo
The peer to dial
relay_peer_id : Optional[ID], optional
Optional specific relay peer to use
Returns
-------
RawConnection
The established connection
Raises
------
ConnectionError
If the connection cannot be established
"""
# If no specific relay is provided, try to find one
if relay_peer_id is None:
relay_peer_id = await self._select_relay(peer_info)
if not relay_peer_id:
raise ConnectionError("No suitable relay found")
# Get a stream to the relay
relay_stream = await self.host.new_stream(relay_peer_id, [PROTOCOL_ID])
if not relay_stream:
raise ConnectionError(f"Could not open stream to relay {relay_peer_id}")
try:
# First try to make a reservation if enabled
if self.config.enable_client:
success = await self._make_reservation(relay_stream, relay_peer_id)
if not success:
logger.warning(
"Failed to make reservation with relay %s", relay_peer_id
)
# Send HOP CONNECT message
hop_msg = HopMessage(
type=HopMessage.CONNECT,
peer=peer_info.peer_id.to_bytes(),
)
await relay_stream.write(hop_msg.SerializeToString())
# Read response
resp_bytes = await relay_stream.read()
resp = HopMessage()
resp.ParseFromString(resp_bytes)
# Access status attributes directly
status_code = getattr(resp.status, "code", StatusCode.OK)
status_msg = getattr(resp.status, "message", "Unknown error")
if status_code != StatusCode.OK:
raise ConnectionError(f"Relay connection failed: {status_msg}")
# Create raw connection from stream
return RawConnection(stream=relay_stream, initiator=True)
except Exception as e:
await relay_stream.close()
raise ConnectionError(f"Failed to establish relay connection: {str(e)}")
async def _select_relay(self, peer_info: PeerInfo) -> ID | None:
"""
Select an appropriate relay for the given peer.
Parameters
----------
peer_info : PeerInfo
The peer to connect to
Returns
-------
Optional[ID]
Selected relay peer ID, or None if no suitable relay found
"""
# Try to find a relay
attempts = 0
while attempts < self.client_config.max_auto_relay_attempts:
# Get a relay from the list of discovered relays
relays = self.discovery.get_relays()
if relays:
# TODO: Implement more sophisticated relay selection
# For now, just return the first available relay
return relays[0]
# Wait and try discovery
await trio.sleep(1)
attempts += 1
return None
async def _make_reservation(
self,
stream: INetStream,
relay_peer_id: ID,
) -> bool:
"""
Make a reservation with a relay.
Parameters
----------
stream : INetStream
Stream to the relay
relay_peer_id : ID
The relay's peer ID
Returns
-------
bool
True if reservation was successful
"""
try:
# Send reservation request
reserve_msg = HopMessage(
type=HopMessage.RESERVE,
peer=self.host.get_id().to_bytes(),
)
await stream.write(reserve_msg.SerializeToString())
# Read response
resp_bytes = await stream.read()
resp = HopMessage()
resp.ParseFromString(resp_bytes)
# Access status attributes directly
status_code = getattr(resp.status, "code", StatusCode.OK)
status_msg = getattr(resp.status, "message", "Unknown error")
if status_code != StatusCode.OK:
logger.warning(
"Reservation failed with relay %s: %s",
relay_peer_id,
status_msg,
)
return False
# Store reservation info
# TODO: Implement reservation storage and refresh mechanism
return True
except Exception as e:
logger.error("Error making reservation: %s", str(e))
return False
def create_listener(
self,
handler_function: Callable[[ReadWriteCloser], Awaitable[None]],
) -> IListener:
"""
Create a listener for incoming relay connections.
Parameters
----------
handler_function : Callable[[ReadWriteCloser], Awaitable[None]]
The handler function for new connections
Returns
-------
IListener
The created listener
"""
return CircuitV2Listener(self.host, self.protocol, self.config)
class CircuitV2Listener(Service, IListener):
"""Listener for incoming relay connections."""
def __init__(
self,
host: IHost,
protocol: CircuitV2Protocol,
config: RelayConfig,
) -> None:
"""
Initialize the Circuit v2 listener.
Parameters
----------
host : IHost
The libp2p host this listener is running on
protocol : CircuitV2Protocol
The Circuit v2 protocol instance
config : RelayConfig
Relay configuration
"""
super().__init__()
self.host = host
self.protocol = protocol
self.config = config
self.multiaddrs: list[
multiaddr.Multiaddr
] = [] # Store multiaddrs as Multiaddr objects
async def handle_incoming_connection(
self,
stream: INetStream,
remote_peer_id: ID,
) -> RawConnection:
"""
Handle an incoming relay connection.
Parameters
----------
stream : INetStream
The incoming stream
remote_peer_id : ID
The remote peer's ID
Returns
-------
RawConnection
The established connection
Raises
------
ConnectionError
If the connection cannot be established
"""
if not self.config.enable_stop:
raise ConnectionError("Stop role is not enabled")
try:
# Read STOP message
msg_bytes = await stream.read()
stop_msg = StopMessage()
stop_msg.ParseFromString(msg_bytes)
if stop_msg.type != StopMessage.CONNECT:
raise ConnectionError("Invalid STOP message type")
# Create raw connection
return RawConnection(stream=stream, initiator=False)
except Exception as e:
await stream.close()
raise ConnectionError(f"Failed to handle incoming connection: {str(e)}")
async def run(self) -> None:
"""Run the listener service."""
# Implementation would go here
async def listen(self, maddr: multiaddr.Multiaddr, nursery: trio.Nursery) -> bool:
"""
Start listening on the given multiaddr.
Parameters
----------
maddr : multiaddr.Multiaddr
The multiaddr to listen on
nursery : trio.Nursery
The nursery to run tasks in
Returns
-------
bool
True if listening successfully started
"""
# Convert string to Multiaddr if needed
addr = (
maddr
if isinstance(maddr, multiaddr.Multiaddr)
else multiaddr.Multiaddr(maddr)
)
self.multiaddrs.append(addr)
return True
def get_addrs(self) -> tuple[multiaddr.Multiaddr, ...]:
"""
Get the listening addresses.
Returns
-------
tuple[multiaddr.Multiaddr, ...]
Tuple of listening multiaddresses
"""
return tuple(self.multiaddrs)
async def close(self) -> None:
"""Close the listener."""
self.multiaddrs.clear()
await self.manager.stop()

6
libp2p/tools/utils.py
View File

@ -87,14 +87,16 @@ async def connect(node1: IHost, node2: IHost) -> None:
addr = node2.get_addrs()[0]
info = info_from_p2p_addr(addr)
# Add retry logic for more robust connection
# Add retry logic for more robust connection with timeout
max_retries = 3
retry_delay = 0.2
last_error = None
for attempt in range(max_retries):
try:
await node1.connect(info)
# Use timeout for each connection attempt
with trio.move_on_after(5): # 5 second timeout
await node1.connect(info)
# Verify connection is established in both directions
if (

7
libp2p/transport/upgrader.py
View File

@ -48,11 +48,16 @@ class TransportUpgrader:
# TODO: Figure out what to do with this function.
async def upgrade_security(
self, raw_conn: IRawConnection, peer_id: ID, is_initiator: bool
self,
raw_conn: IRawConnection,
is_initiator: bool,
peer_id: ID | None = None,
) -> ISecureConn:
"""Upgrade conn to a secured connection."""
try:
if is_initiator:
if peer_id is None:
raise ValueError("peer_id must be provided for outbout connection")
return await self.security_multistream.secure_outbound(
raw_conn, peer_id
)

1
newsfragments/579.feature.rst Normal file
View File

@ -0,0 +1 @@
Added support for ``Kademlia DHT`` in py-libp2p.

1
newsfragments/678.misc.rst Normal file
View File

@ -0,0 +1 @@
Refactored gossipsub heartbeat logic to use a single helper method `_handle_topic_heartbeat` that handles both fanout and gossip heartbeats.

1
newsfragments/679.feature.rst Normal file
View File

@ -0,0 +1 @@
Added sparse connect utility function to pubsub test utilities for creating test networks with configurable connectivity.

2
newsfragments/681.breaking.rst Normal file
View File

@ -0,0 +1,2 @@
Reordered the arguments to `upgrade_security` to place `is_initiator` before `peer_id`, and made `peer_id` optional.
This allows the method to reflect the fact that peer identity is not required for inbound connections.

1
newsfragments/684.misc.rst Normal file
View File

@ -0,0 +1 @@
Uses the `decapsulate` method of the `Multiaddr` class to clean up the observed address.

1
newsfragments/685.feature.rst Normal file
View File

@ -0,0 +1 @@
Optimized pubsub publishing to send multiple topics in a single message instead of separate messages per topic.

11
tests/core/identity/identify/test_identify.py
View File

@ -56,16 +56,9 @@ async def test_identify_protocol(security_protocol):
)
# Check observed address
# TODO: use decapsulateCode(protocols('p2p').code)
# when the Multiaddr class will implement it
host_b_addr = host_b.get_addrs()[0]
cleaned_addr = Multiaddr.join(
*(
host_b_addr.split()[:-1]
if str(host_b_addr.split()[-1]).startswith("/p2p/")
else host_b_addr.split()
)
)
host_b_peer_id = host_b.get_id()
cleaned_addr = host_b_addr.decapsulate(Multiaddr(f"/p2p/{host_b_peer_id}"))
logger.debug("observed_addr: %s", Multiaddr(identify_response.observed_addr))
logger.debug("host_b.get_addrs()[0]: %s", host_b.get_addrs()[0])

168
tests/core/kad_dht/test_kad_dht.py Normal file
View File

@ -0,0 +1,168 @@
"""
Tests for the Kademlia DHT implementation.
This module tests core functionality of the Kademlia DHT including:
- Node discovery (find_node)
- Value storage and retrieval (put_value, get_value)
- Content provider advertisement and discovery (provide, find_providers)
"""
import hashlib
import logging
import uuid
import pytest
import trio
from libp2p.kad_dht.kad_dht import (
DHTMode,
KadDHT,
)
from libp2p.kad_dht.utils import (
create_key_from_binary,
)
from libp2p.peer.peerinfo import (
PeerInfo,
)
from libp2p.tools.async_service import (
background_trio_service,
)
from tests.utils.factories import (
host_pair_factory,
)
# Configure logger
logger = logging.getLogger("test.kad_dht")
# Constants for the tests
TEST_TIMEOUT = 5 # Timeout in seconds
@pytest.fixture
async def dht_pair(security_protocol):
"""Create a pair of connected DHT nodes for testing."""
async with host_pair_factory(security_protocol=security_protocol) as (
host_a,
host_b,
):
# Get peer info for bootstrapping
peer_b_info = PeerInfo(host_b.get_id(), host_b.get_addrs())
peer_a_info = PeerInfo(host_a.get_id(), host_a.get_addrs())
# Create DHT nodes from the hosts with bootstrap peers as multiaddr strings
dht_a: KadDHT = KadDHT(host_a, mode=DHTMode.SERVER)
dht_b: KadDHT = KadDHT(host_b, mode=DHTMode.SERVER)
await dht_a.peer_routing.routing_table.add_peer(peer_b_info)
await dht_b.peer_routing.routing_table.add_peer(peer_a_info)
# Start both DHT services
async with background_trio_service(dht_a), background_trio_service(dht_b):
# Allow time for bootstrap to complete and connections to establish
await trio.sleep(0.1)
logger.debug(
"After bootstrap: Node A peers: %s", dht_a.routing_table.get_peer_ids()
)
logger.debug(
"After bootstrap: Node B peers: %s", dht_b.routing_table.get_peer_ids()
)
# Return the DHT pair
yield (dht_a, dht_b)
@pytest.mark.trio
async def test_find_node(dht_pair: tuple[KadDHT, KadDHT]):
"""Test that nodes can find each other in the DHT."""
dht_a, dht_b = dht_pair
# Node A should be able to find Node B
with trio.fail_after(TEST_TIMEOUT):
found_info = await dht_a.find_peer(dht_b.host.get_id())
# Verify that the found peer has the correct peer ID
assert found_info is not None, "Failed to find the target peer"
assert found_info.peer_id == dht_b.host.get_id(), "Found incorrect peer ID"
@pytest.mark.trio
async def test_put_and_get_value(dht_pair: tuple[KadDHT, KadDHT]):
"""Test storing and retrieving values in the DHT."""
dht_a, dht_b = dht_pair
# dht_a.peer_routing.routing_table.add_peer(dht_b.pe)
peer_b_info = PeerInfo(dht_b.host.get_id(), dht_b.host.get_addrs())
# Generate a random key and value
key = create_key_from_binary(b"test-key")
value = b"test-value"
# First add the value directly to node A's store to verify storage works
dht_a.value_store.put(key, value)
logger.debug("Local value store: %s", dht_a.value_store.store)
local_value = dht_a.value_store.get(key)
assert local_value == value, "Local value storage failed"
print("number of nodes in peer store", dht_a.host.get_peerstore().peer_ids())
await dht_a.routing_table.add_peer(peer_b_info)
print("Routing table of a has ", dht_a.routing_table.get_peer_ids())
# Store the value using the first node (this will also store locally)
with trio.fail_after(TEST_TIMEOUT):
await dht_a.put_value(key, value)
# # Log debugging information
logger.debug("Put value with key %s...", key.hex()[:10])
logger.debug("Node A value store: %s", dht_a.value_store.store)
print("hello test")
# # Allow more time for the value to propagate
await trio.sleep(0.5)
# # Try direct connection between nodes to ensure they're properly linked
logger.debug("Node A peers: %s", dht_a.routing_table.get_peer_ids())
logger.debug("Node B peers: %s", dht_b.routing_table.get_peer_ids())
# Retrieve the value using the second node
with trio.fail_after(TEST_TIMEOUT):
retrieved_value = await dht_b.get_value(key)
print("the value stored in node b is", dht_b.get_value_store_size())
logger.debug("Retrieved value: %s", retrieved_value)
# Verify that the retrieved value matches the original
assert retrieved_value == value, "Retrieved value does not match the stored value"
@pytest.mark.trio
async def test_provide_and_find_providers(dht_pair: tuple[KadDHT, KadDHT]):
"""Test advertising and finding content providers."""
dht_a, dht_b = dht_pair
# Generate a random content ID
content = f"test-content-{uuid.uuid4()}".encode()
content_id = hashlib.sha256(content).digest()
# Store content on the first node
dht_a.value_store.put(content_id, content)
# Advertise the first node as a provider
with trio.fail_after(TEST_TIMEOUT):
success = await dht_a.provide(content_id)
assert success, "Failed to advertise as provider"
# Allow time for the provider record to propagate
await trio.sleep(0.1)
# Find providers using the second node
with trio.fail_after(TEST_TIMEOUT):
providers = await dht_b.find_providers(content_id)
# Verify that we found the first node as a provider
assert providers, "No providers found"
assert any(p.peer_id == dht_a.local_peer_id for p in providers), (
"Expected provider not found"
)
# Retrieve the content using the provider information
with trio.fail_after(TEST_TIMEOUT):
retrieved_value = await dht_b.get_value(content_id)
assert retrieved_value == content, (
"Retrieved content does not match the original"
)

459
tests/core/kad_dht/test_unit_peer_routing.py Normal file
View File

@ -0,0 +1,459 @@
"""
Unit tests for the PeerRouting class in Kademlia DHT.
This module tests the core functionality of peer routing including:
- Peer discovery and lookup
- Network queries for closest peers
- Protocol message handling
- Error handling and edge cases
"""
import time
from unittest.mock import (
AsyncMock,
Mock,
patch,
)
import pytest
from multiaddr import (
Multiaddr,
)
import varint
from libp2p.crypto.secp256k1 import (
create_new_key_pair,
)
from libp2p.kad_dht.pb.kademlia_pb2 import (
Message,
)
from libp2p.kad_dht.peer_routing import (
ALPHA,
MAX_PEER_LOOKUP_ROUNDS,
PROTOCOL_ID,
PeerRouting,
)
from libp2p.kad_dht.routing_table import (
RoutingTable,
)
from libp2p.peer.id import (
ID,
)
from libp2p.peer.peerinfo import (
PeerInfo,
)
def create_valid_peer_id(name: str) -> ID:
"""Create a valid peer ID for testing."""
key_pair = create_new_key_pair()
return ID.from_pubkey(key_pair.public_key)
class TestPeerRouting:
"""Test suite for PeerRouting class."""
@pytest.fixture
def mock_host(self):
"""Create a mock host for testing."""
host = Mock()
host.get_id.return_value = create_valid_peer_id("local")
host.get_addrs.return_value = [Multiaddr("/ip4/127.0.0.1/tcp/8000")]
host.get_peerstore.return_value = Mock()
host.new_stream = AsyncMock()
host.connect = AsyncMock()
return host
@pytest.fixture
def mock_routing_table(self, mock_host):
"""Create a mock routing table for testing."""
local_id = create_valid_peer_id("local")
routing_table = RoutingTable(local_id, mock_host)
return routing_table
@pytest.fixture
def peer_routing(self, mock_host, mock_routing_table):
"""Create a PeerRouting instance for testing."""
return PeerRouting(mock_host, mock_routing_table)
@pytest.fixture
def sample_peer_info(self):
"""Create sample peer info for testing."""
peer_id = create_valid_peer_id("sample")
addresses = [Multiaddr("/ip4/127.0.0.1/tcp/8001")]
return PeerInfo(peer_id, addresses)
def test_init_peer_routing(self, mock_host, mock_routing_table):
"""Test PeerRouting initialization."""
peer_routing = PeerRouting(mock_host, mock_routing_table)
assert peer_routing.host == mock_host
assert peer_routing.routing_table == mock_routing_table
assert peer_routing.protocol_id == PROTOCOL_ID
@pytest.mark.trio
async def test_find_peer_local_host(self, peer_routing, mock_host):
"""Test finding our own peer."""
local_id = mock_host.get_id()
result = await peer_routing.find_peer(local_id)
assert result is not None
assert result.peer_id == local_id
assert result.addrs == mock_host.get_addrs()
@pytest.mark.trio
async def test_find_peer_in_routing_table(self, peer_routing, sample_peer_info):
"""Test finding peer that exists in routing table."""
# Add peer to routing table
await peer_routing.routing_table.add_peer(sample_peer_info)
result = await peer_routing.find_peer(sample_peer_info.peer_id)
assert result is not None
assert result.peer_id == sample_peer_info.peer_id
@pytest.mark.trio
async def test_find_peer_in_peerstore(self, peer_routing, mock_host):
"""Test finding peer that exists in peerstore."""
peer_id = create_valid_peer_id("peerstore")
mock_addrs = [Multiaddr("/ip4/127.0.0.1/tcp/8002")]
# Mock peerstore to return addresses
mock_host.get_peerstore().addrs.return_value = mock_addrs
result = await peer_routing.find_peer(peer_id)
assert result is not None
assert result.peer_id == peer_id
assert result.addrs == mock_addrs
@pytest.mark.trio
async def test_find_peer_not_found(self, peer_routing, mock_host):
"""Test finding peer that doesn't exist anywhere."""
peer_id = create_valid_peer_id("nonexistent")
# Mock peerstore to return no addresses
mock_host.get_peerstore().addrs.return_value = []
# Mock network search to return empty results
with patch.object(peer_routing, "find_closest_peers_network", return_value=[]):
result = await peer_routing.find_peer(peer_id)
assert result is None
@pytest.mark.trio
async def test_find_closest_peers_network_empty_start(self, peer_routing):
"""Test network search with no local peers."""
target_key = b"target_key"
# Mock routing table to return empty list
with patch.object(
peer_routing.routing_table, "find_local_closest_peers", return_value=[]
):
result = await peer_routing.find_closest_peers_network(target_key)
assert result == []
@pytest.mark.trio
async def test_find_closest_peers_network_with_peers(self, peer_routing, mock_host):
"""Test network search with some initial peers."""
target_key = b"target_key"
# Create some test peers
initial_peers = [create_valid_peer_id(f"peer{i}") for i in range(3)]
# Mock routing table to return initial peers
with patch.object(
peer_routing.routing_table,
"find_local_closest_peers",
return_value=initial_peers,
):
# Mock _query_peer_for_closest to return empty results (no new peers found)
with patch.object(peer_routing, "_query_peer_for_closest", return_value=[]):
result = await peer_routing.find_closest_peers_network(
target_key, count=5
)
assert len(result) <= 5
# Should return the initial peers since no new ones were discovered
assert all(peer in initial_peers for peer in result)
@pytest.mark.trio
async def test_find_closest_peers_convergence(self, peer_routing):
"""Test that network search converges properly."""
target_key = b"target_key"
# Create test peers
initial_peers = [create_valid_peer_id(f"peer{i}") for i in range(2)]
# Mock to simulate convergence (no improvement in closest peers)
with patch.object(
peer_routing.routing_table,
"find_local_closest_peers",
return_value=initial_peers,
):
with patch.object(peer_routing, "_query_peer_for_closest", return_value=[]):
with patch(
"libp2p.kad_dht.peer_routing.sort_peer_ids_by_distance",
return_value=initial_peers,
):
result = await peer_routing.find_closest_peers_network(target_key)
assert result == initial_peers
@pytest.mark.trio
async def test_query_peer_for_closest_success(
self, peer_routing, mock_host, sample_peer_info
):
"""Test successful peer query for closest peers."""
target_key = b"target_key"
# Create mock stream
mock_stream = AsyncMock()
mock_host.new_stream.return_value = mock_stream
# Create mock response
response_msg = Message()
response_msg.type = Message.MessageType.FIND_NODE
# Add a peer to the response
peer_proto = response_msg.closerPeers.add()
response_peer_id = create_valid_peer_id("response_peer")
peer_proto.id = response_peer_id.to_bytes()
peer_proto.addrs.append(Multiaddr("/ip4/127.0.0.1/tcp/8003").to_bytes())
response_bytes = response_msg.SerializeToString()
# Mock stream reading
varint_length = varint.encode(len(response_bytes))
mock_stream.read.side_effect = [varint_length, response_bytes]
# Mock peerstore
mock_host.get_peerstore().addrs.return_value = [sample_peer_info.addrs[0]]
mock_host.get_peerstore().add_addrs = Mock()
result = await peer_routing._query_peer_for_closest(
sample_peer_info.peer_id, target_key
)
assert len(result) == 1
assert result[0] == response_peer_id
mock_stream.write.assert_called()
mock_stream.close.assert_called_once()
@pytest.mark.trio
async def test_query_peer_for_closest_stream_failure(self, peer_routing, mock_host):
"""Test peer query when stream creation fails."""
target_key = b"target_key"
peer_id = create_valid_peer_id("test")
# Mock stream creation failure
mock_host.new_stream.side_effect = Exception("Stream failed")
mock_host.get_peerstore().addrs.return_value = []
result = await peer_routing._query_peer_for_closest(peer_id, target_key)
assert result == []
@pytest.mark.trio
async def test_query_peer_for_closest_read_failure(
self, peer_routing, mock_host, sample_peer_info
):
"""Test peer query when reading response fails."""
target_key = b"target_key"
# Create mock stream that fails to read
mock_stream = AsyncMock()
mock_stream.read.side_effect = [b""] # Empty read simulates connection close
mock_host.new_stream.return_value = mock_stream
mock_host.get_peerstore().addrs.return_value = [sample_peer_info.addrs[0]]
result = await peer_routing._query_peer_for_closest(
sample_peer_info.peer_id, target_key
)
assert result == []
mock_stream.close.assert_called_once()
@pytest.mark.trio
async def test_refresh_routing_table(self, peer_routing, mock_host):
"""Test routing table refresh."""
local_id = mock_host.get_id()
discovered_peers = [create_valid_peer_id(f"discovered{i}") for i in range(3)]
# Mock find_closest_peers_network to return discovered peers
with patch.object(
peer_routing, "find_closest_peers_network", return_value=discovered_peers
):
# Mock peerstore to return addresses for discovered peers
mock_addrs = [Multiaddr("/ip4/127.0.0.1/tcp/8003")]
mock_host.get_peerstore().addrs.return_value = mock_addrs
await peer_routing.refresh_routing_table()
# Should perform lookup for local ID
peer_routing.find_closest_peers_network.assert_called_once_with(
local_id.to_bytes()
)
@pytest.mark.trio
async def test_handle_kad_stream_find_node(self, peer_routing, mock_host):
"""Test handling incoming FIND_NODE requests."""
# Create mock stream
mock_stream = AsyncMock()
# Create FIND_NODE request
request_msg = Message()
request_msg.type = Message.MessageType.FIND_NODE
request_msg.key = b"target_key"
request_bytes = request_msg.SerializeToString()
# Mock stream reading
mock_stream.read.side_effect = [
len(request_bytes).to_bytes(4, byteorder="big"),
request_bytes,
]
# Mock routing table to return some peers
closest_peers = [create_valid_peer_id(f"close{i}") for i in range(2)]
with patch.object(
peer_routing.routing_table,
"find_local_closest_peers",
return_value=closest_peers,
):
mock_host.get_peerstore().addrs.return_value = [
Multiaddr("/ip4/127.0.0.1/tcp/8004")
]
await peer_routing._handle_kad_stream(mock_stream)
# Should write response
mock_stream.write.assert_called()
mock_stream.close.assert_called_once()
@pytest.mark.trio
async def test_handle_kad_stream_invalid_message(self, peer_routing):
"""Test handling stream with invalid message."""
mock_stream = AsyncMock()
# Mock stream to return invalid data
mock_stream.read.side_effect = [
(10).to_bytes(4, byteorder="big"),
b"invalid_proto_data",
]
# Should handle gracefully without raising exception
await peer_routing._handle_kad_stream(mock_stream)
mock_stream.close.assert_called_once()
@pytest.mark.trio
async def test_handle_kad_stream_connection_closed(self, peer_routing):
"""Test handling stream when connection is closed early."""
mock_stream = AsyncMock()
# Mock stream to return empty data (connection closed)
mock_stream.read.return_value = b""
await peer_routing._handle_kad_stream(mock_stream)
mock_stream.close.assert_called_once()
@pytest.mark.trio
async def test_query_single_peer_for_closest_success(self, peer_routing):
"""Test _query_single_peer_for_closest method."""
target_key = b"target_key"
peer_id = create_valid_peer_id("test")
new_peers = []
# Mock successful query
mock_result = [create_valid_peer_id("result1"), create_valid_peer_id("result2")]
with patch.object(
peer_routing, "_query_peer_for_closest", return_value=mock_result
):
await peer_routing._query_single_peer_for_closest(
peer_id, target_key, new_peers
)
assert len(new_peers) == 2
assert all(peer in new_peers for peer in mock_result)
@pytest.mark.trio
async def test_query_single_peer_for_closest_failure(self, peer_routing):
"""Test _query_single_peer_for_closest when query fails."""
target_key = b"target_key"
peer_id = create_valid_peer_id("test")
new_peers = []
# Mock query failure
with patch.object(
peer_routing,
"_query_peer_for_closest",
side_effect=Exception("Query failed"),
):
await peer_routing._query_single_peer_for_closest(
peer_id, target_key, new_peers
)
# Should handle exception gracefully
assert len(new_peers) == 0
@pytest.mark.trio
async def test_query_single_peer_deduplication(self, peer_routing):
"""Test that _query_single_peer_for_closest deduplicates peers."""
target_key = b"target_key"
peer_id = create_valid_peer_id("test")
duplicate_peer = create_valid_peer_id("duplicate")
new_peers = [duplicate_peer] # Pre-existing peer
# Mock query to return the same peer
mock_result = [duplicate_peer, create_valid_peer_id("new")]
with patch.object(
peer_routing, "_query_peer_for_closest", return_value=mock_result
):
await peer_routing._query_single_peer_for_closest(
peer_id, target_key, new_peers
)
# Should not add duplicate
assert len(new_peers) == 2 # Original + 1 new peer
assert new_peers.count(duplicate_peer) == 1
def test_constants(self):
"""Test that important constants are properly defined."""
assert ALPHA == 3
assert MAX_PEER_LOOKUP_ROUNDS == 20
assert PROTOCOL_ID == "/ipfs/kad/1.0.0"
@pytest.mark.trio
async def test_edge_case_max_rounds_reached(self, peer_routing):
"""Test that lookup stops after maximum rounds."""
target_key = b"target_key"
initial_peers = [create_valid_peer_id("peer1")]
# Mock to always return new peers to force max rounds
def mock_query_side_effect(peer, key):
return [create_valid_peer_id(f"new_peer_{time.time()}")]
with patch.object(
peer_routing.routing_table,
"find_local_closest_peers",
return_value=initial_peers,
):
with patch.object(
peer_routing,
"_query_peer_for_closest",
side_effect=mock_query_side_effect,
):
with patch(
"libp2p.kad_dht.peer_routing.sort_peer_ids_by_distance"
) as mock_sort:
# Always return different peers to prevent convergence
mock_sort.side_effect = lambda key, peers: peers[:20]
result = await peer_routing.find_closest_peers_network(target_key)
# Should stop after max rounds, not infinite loop
assert isinstance(result, list)

805
tests/core/kad_dht/test_unit_provider_store.py Normal file
View File

@ -0,0 +1,805 @@
"""
Unit tests for the ProviderStore and ProviderRecord classes in Kademlia DHT.
This module tests the core functionality of provider record management including:
- ProviderRecord creation, expiration, and republish logic
- ProviderStore operations (add, get, cleanup)
- Expiration and TTL handling
- Network operations (mocked)
- Edge cases and error conditions
"""
import time
from unittest.mock import (
AsyncMock,
Mock,
patch,
)
import pytest
from multiaddr import (
Multiaddr,
)
from libp2p.kad_dht.provider_store import (
PROVIDER_ADDRESS_TTL,
PROVIDER_RECORD_EXPIRATION_INTERVAL,
PROVIDER_RECORD_REPUBLISH_INTERVAL,
ProviderRecord,
ProviderStore,
)
from libp2p.peer.id import (
ID,
)
from libp2p.peer.peerinfo import (
PeerInfo,
)
mock_host = Mock()
class TestProviderRecord:
"""Test suite for ProviderRecord class."""
def test_init_with_default_timestamp(self):
"""Test ProviderRecord initialization with default timestamp."""
peer_id = ID.from_base58("QmTest123")
addresses = [Multiaddr("/ip4/127.0.0.1/tcp/8000")]
peer_info = PeerInfo(peer_id, addresses)
start_time = time.time()
record = ProviderRecord(peer_info)
end_time = time.time()
assert record.provider_info == peer_info
assert start_time <= record.timestamp <= end_time
assert record.peer_id == peer_id
assert record.addresses == addresses
def test_init_with_custom_timestamp(self):
"""Test ProviderRecord initialization with custom timestamp."""
peer_id = ID.from_base58("QmTest123")
peer_info = PeerInfo(peer_id, [])
custom_timestamp = time.time() - 3600 # 1 hour ago
record = ProviderRecord(peer_info, timestamp=custom_timestamp)
assert record.timestamp == custom_timestamp
def test_is_expired_fresh_record(self):
"""Test that fresh records are not expired."""
peer_id = ID.from_base58("QmTest123")
peer_info = PeerInfo(peer_id, [])
record = ProviderRecord(peer_info)
assert not record.is_expired()
def test_is_expired_old_record(self):
"""Test that old records are expired."""
peer_id = ID.from_base58("QmTest123")
peer_info = PeerInfo(peer_id, [])
old_timestamp = time.time() - PROVIDER_RECORD_EXPIRATION_INTERVAL - 1
record = ProviderRecord(peer_info, timestamp=old_timestamp)
assert record.is_expired()
def test_is_expired_boundary_condition(self):
"""Test expiration at exact boundary."""
peer_id = ID.from_base58("QmTest123")
peer_info = PeerInfo(peer_id, [])
boundary_timestamp = time.time() - PROVIDER_RECORD_EXPIRATION_INTERVAL
record = ProviderRecord(peer_info, timestamp=boundary_timestamp)
# At the exact boundary, should be expired (implementation uses >)
assert record.is_expired()
def test_should_republish_fresh_record(self):
"""Test that fresh records don't need republishing."""
peer_id = ID.from_base58("QmTest123")
peer_info = PeerInfo(peer_id, [])
record = ProviderRecord(peer_info)
assert not record.should_republish()
def test_should_republish_old_record(self):
"""Test that old records need republishing."""
peer_id = ID.from_base58("QmTest123")
peer_info = PeerInfo(peer_id, [])
old_timestamp = time.time() - PROVIDER_RECORD_REPUBLISH_INTERVAL - 1
record = ProviderRecord(peer_info, timestamp=old_timestamp)
assert record.should_republish()
def test_should_republish_boundary_condition(self):
"""Test republish at exact boundary."""
peer_id = ID.from_base58("QmTest123")
peer_info = PeerInfo(peer_id, [])
boundary_timestamp = time.time() - PROVIDER_RECORD_REPUBLISH_INTERVAL
record = ProviderRecord(peer_info, timestamp=boundary_timestamp)
# At the exact boundary, should need republishing (implementation uses >)
assert record.should_republish()
def test_properties(self):
"""Test peer_id and addresses properties."""
peer_id = ID.from_base58("QmTest123")
addresses = [
Multiaddr("/ip4/127.0.0.1/tcp/8000"),
Multiaddr("/ip6/::1/tcp/8001"),
]
peer_info = PeerInfo(peer_id, addresses)
record = ProviderRecord(peer_info)
assert record.peer_id == peer_id
assert record.addresses == addresses
def test_empty_addresses(self):
"""Test ProviderRecord with empty address list."""
peer_id = ID.from_base58("QmTest123")
peer_info = PeerInfo(peer_id, [])
record = ProviderRecord(peer_info)
assert record.addresses == []
class TestProviderStore:
"""Test suite for ProviderStore class."""
def test_init_empty_store(self):
"""Test that a new ProviderStore is initialized empty."""
store = ProviderStore(host=mock_host)
assert len(store.providers) == 0
assert store.peer_routing is None
assert len(store.providing_keys) == 0
def test_init_with_host(self):
"""Test initialization with host."""
mock_host = Mock()
mock_peer_id = ID.from_base58("QmTest123")
mock_host.get_id.return_value = mock_peer_id
store = ProviderStore(host=mock_host)
assert store.host == mock_host
assert store.local_peer_id == mock_peer_id
assert len(store.providers) == 0
def test_init_with_host_and_peer_routing(self):
"""Test initialization with both host and peer routing."""
mock_host = Mock()
mock_peer_routing = Mock()
mock_peer_id = ID.from_base58("QmTest123")
mock_host.get_id.return_value = mock_peer_id
store = ProviderStore(host=mock_host, peer_routing=mock_peer_routing)
assert store.host == mock_host
assert store.peer_routing == mock_peer_routing
assert store.local_peer_id == mock_peer_id
def test_add_provider_new_key(self):
"""Test adding a provider for a new key."""
store = ProviderStore(host=mock_host)
key = b"test_key"
peer_id = ID.from_base58("QmTest123")
addresses = [Multiaddr("/ip4/127.0.0.1/tcp/8000")]
provider = PeerInfo(peer_id, addresses)
store.add_provider(key, provider)
assert key in store.providers
assert str(peer_id) in store.providers[key]
record = store.providers[key][str(peer_id)]
assert record.provider_info == provider
assert isinstance(record.timestamp, float)
def test_add_provider_existing_key(self):
"""Test adding multiple providers for the same key."""
store = ProviderStore(host=mock_host)
key = b"test_key"
# Add first provider
peer_id1 = ID.from_base58("QmTest123")
provider1 = PeerInfo(peer_id1, [])
store.add_provider(key, provider1)
# Add second provider
peer_id2 = ID.from_base58("QmTest456")
provider2 = PeerInfo(peer_id2, [])
store.add_provider(key, provider2)
assert len(store.providers[key]) == 2
assert str(peer_id1) in store.providers[key]
assert str(peer_id2) in store.providers[key]
def test_add_provider_update_existing(self):
"""Test updating an existing provider."""
store = ProviderStore(host=mock_host)
key = b"test_key"
peer_id = ID.from_base58("QmTest123")
# Add initial provider
provider1 = PeerInfo(peer_id, [Multiaddr("/ip4/127.0.0.1/tcp/8000")])
store.add_provider(key, provider1)
first_timestamp = store.providers[key][str(peer_id)].timestamp
# Small delay to ensure timestamp difference
time.sleep(0.001)
# Update provider
provider2 = PeerInfo(peer_id, [Multiaddr("/ip4/127.0.0.1/tcp/8001")])
store.add_provider(key, provider2)
# Should have same peer but updated info
assert len(store.providers[key]) == 1
assert str(peer_id) in store.providers[key]
record = store.providers[key][str(peer_id)]
assert record.provider_info == provider2
assert record.timestamp > first_timestamp
def test_get_providers_empty_key(self):
"""Test getting providers for non-existent key."""
store = ProviderStore(host=mock_host)
key = b"nonexistent_key"
providers = store.get_providers(key)
assert providers == []
def test_get_providers_valid_records(self):
"""Test getting providers with valid records."""
store = ProviderStore(host=mock_host)
key = b"test_key"
# Add multiple providers
peer_id1 = ID.from_base58("QmTest123")
peer_id2 = ID.from_base58("QmTest456")
provider1 = PeerInfo(peer_id1, [Multiaddr("/ip4/127.0.0.1/tcp/8000")])
provider2 = PeerInfo(peer_id2, [Multiaddr("/ip4/127.0.0.1/tcp/8001")])
store.add_provider(key, provider1)
store.add_provider(key, provider2)
providers = store.get_providers(key)
assert len(providers) == 2
provider_ids = {p.peer_id for p in providers}
assert peer_id1 in provider_ids
assert peer_id2 in provider_ids
def test_get_providers_expired_records(self):
"""Test that expired records are filtered out and cleaned up."""
store = ProviderStore(host=mock_host)
key = b"test_key"
# Add valid provider
peer_id1 = ID.from_base58("QmTest123")
provider1 = PeerInfo(peer_id1, [])
store.add_provider(key, provider1)
# Add expired provider manually
peer_id2 = ID.from_base58("QmTest456")
provider2 = PeerInfo(peer_id2, [])
expired_timestamp = time.time() - PROVIDER_RECORD_EXPIRATION_INTERVAL - 1
store.providers[key][str(peer_id2)] = ProviderRecord(
provider2, expired_timestamp
)
providers = store.get_providers(key)
# Should only return valid provider
assert len(providers) == 1
assert providers[0].peer_id == peer_id1
# Expired provider should be cleaned up
assert str(peer_id2) not in store.providers[key]
def test_get_providers_address_ttl(self):
"""Test address TTL handling in get_providers."""
store = ProviderStore(host=mock_host)
key = b"test_key"
peer_id = ID.from_base58("QmTest123")
addresses = [Multiaddr("/ip4/127.0.0.1/tcp/8000")]
provider = PeerInfo(peer_id, addresses)
# Add provider with old timestamp (addresses expired but record valid)
old_timestamp = time.time() - PROVIDER_ADDRESS_TTL - 1
store.providers[key] = {str(peer_id): ProviderRecord(provider, old_timestamp)}
providers = store.get_providers(key)
# Should return provider but with empty addresses
assert len(providers) == 1
assert providers[0].peer_id == peer_id
assert providers[0].addrs == []
def test_get_providers_cleanup_empty_key(self):
"""Test that keys with no valid providers are removed."""
store = ProviderStore(host=mock_host)
key = b"test_key"
# Add only expired providers
peer_id = ID.from_base58("QmTest123")
provider = PeerInfo(peer_id, [])
expired_timestamp = time.time() - PROVIDER_RECORD_EXPIRATION_INTERVAL - 1
store.providers[key] = {
str(peer_id): ProviderRecord(provider, expired_timestamp)
}
providers = store.get_providers(key)
assert providers == []
assert key not in store.providers # Key should be removed
def test_cleanup_expired_no_expired_records(self):
"""Test cleanup when there are no expired records."""
store = ProviderStore(host=mock_host)
key1 = b"key1"
key2 = b"key2"
# Add valid providers
peer_id1 = ID.from_base58("QmTest123")
peer_id2 = ID.from_base58("QmTest456")
provider1 = PeerInfo(peer_id1, [])
provider2 = PeerInfo(peer_id2, [])
store.add_provider(key1, provider1)
store.add_provider(key2, provider2)
initial_size = store.size()
store.cleanup_expired()
assert store.size() == initial_size
assert key1 in store.providers
assert key2 in store.providers
def test_cleanup_expired_with_expired_records(self):
"""Test cleanup removes expired records."""
store = ProviderStore(host=mock_host)
key = b"test_key"
# Add valid provider
peer_id1 = ID.from_base58("QmTest123")
provider1 = PeerInfo(peer_id1, [])
store.add_provider(key, provider1)
# Add expired provider
peer_id2 = ID.from_base58("QmTest456")
provider2 = PeerInfo(peer_id2, [])
expired_timestamp = time.time() - PROVIDER_RECORD_EXPIRATION_INTERVAL - 1
store.providers[key][str(peer_id2)] = ProviderRecord(
provider2, expired_timestamp
)
assert store.size() == 2
store.cleanup_expired()
assert store.size() == 1
assert str(peer_id1) in store.providers[key]
assert str(peer_id2) not in store.providers[key]
def test_cleanup_expired_remove_empty_keys(self):
"""Test that keys with only expired providers are removed."""
store = ProviderStore(host=mock_host)
key1 = b"key1"
key2 = b"key2"
# Add valid provider to key1
peer_id1 = ID.from_base58("QmTest123")
provider1 = PeerInfo(peer_id1, [])
store.add_provider(key1, provider1)
# Add only expired provider to key2
peer_id2 = ID.from_base58("QmTest456")
provider2 = PeerInfo(peer_id2, [])
expired_timestamp = time.time() - PROVIDER_RECORD_EXPIRATION_INTERVAL - 1
store.providers[key2] = {
str(peer_id2): ProviderRecord(provider2, expired_timestamp)
}
store.cleanup_expired()
assert key1 in store.providers
assert key2 not in store.providers
def test_get_provided_keys_empty_store(self):
"""Test get_provided_keys with empty store."""
store = ProviderStore(host=mock_host)
peer_id = ID.from_base58("QmTest123")
keys = store.get_provided_keys(peer_id)
assert keys == []
def test_get_provided_keys_single_peer(self):
"""Test get_provided_keys for a specific peer."""
store = ProviderStore(host=mock_host)
peer_id1 = ID.from_base58("QmTest123")
peer_id2 = ID.from_base58("QmTest456")
key1 = b"key1"
key2 = b"key2"
key3 = b"key3"
provider1 = PeerInfo(peer_id1, [])
provider2 = PeerInfo(peer_id2, [])
# peer_id1 provides key1 and key2
store.add_provider(key1, provider1)
store.add_provider(key2, provider1)
# peer_id2 provides key2 and key3
store.add_provider(key2, provider2)
store.add_provider(key3, provider2)
keys1 = store.get_provided_keys(peer_id1)
keys2 = store.get_provided_keys(peer_id2)
assert len(keys1) == 2
assert key1 in keys1
assert key2 in keys1
assert len(keys2) == 2
assert key2 in keys2
assert key3 in keys2
def test_get_provided_keys_nonexistent_peer(self):
"""Test get_provided_keys for peer that provides nothing."""
store = ProviderStore(host=mock_host)
peer_id1 = ID.from_base58("QmTest123")
peer_id2 = ID.from_base58("QmTest456")
# Add provider for peer_id1 only
key = b"key"
provider = PeerInfo(peer_id1, [])
store.add_provider(key, provider)
# Query for peer_id2 (provides nothing)
keys = store.get_provided_keys(peer_id2)
assert keys == []
def test_size_empty_store(self):
"""Test size() with empty store."""
store = ProviderStore(host=mock_host)
assert store.size() == 0
def test_size_with_providers(self):
"""Test size() with multiple providers."""
store = ProviderStore(host=mock_host)
# Add providers
key1 = b"key1"
key2 = b"key2"
peer_id1 = ID.from_base58("QmTest123")
peer_id2 = ID.from_base58("QmTest456")
peer_id3 = ID.from_base58("QmTest789")
provider1 = PeerInfo(peer_id1, [])
provider2 = PeerInfo(peer_id2, [])
provider3 = PeerInfo(peer_id3, [])
store.add_provider(key1, provider1)
store.add_provider(key1, provider2) # 2 providers for key1
store.add_provider(key2, provider3) # 1 provider for key2
assert store.size() == 3
@pytest.mark.trio
async def test_provide_no_host(self):
"""Test provide() returns False when no host is configured."""
store = ProviderStore(host=mock_host)
key = b"test_key"
result = await store.provide(key)
assert result is False
@pytest.mark.trio
async def test_provide_no_peer_routing(self):
"""Test provide() returns False when no peer routing is configured."""
mock_host = Mock()
store = ProviderStore(host=mock_host)
key = b"test_key"
result = await store.provide(key)
assert result is False
@pytest.mark.trio
async def test_provide_success(self):
"""Test successful provide operation."""
# Setup mocks
mock_host = Mock()
mock_peer_routing = AsyncMock()
peer_id = ID.from_base58("QmTest123")
mock_host.get_id.return_value = peer_id
mock_host.get_addrs.return_value = [Multiaddr("/ip4/127.0.0.1/tcp/8000")]
# Mock finding closest peers
closest_peers = [ID.from_base58("QmPeer1"), ID.from_base58("QmPeer2")]
mock_peer_routing.find_closest_peers_network.return_value = closest_peers
store = ProviderStore(host=mock_host, peer_routing=mock_peer_routing)
# Mock _send_add_provider to return success
with patch.object(store, "_send_add_provider", return_value=True) as mock_send:
key = b"test_key"
result = await store.provide(key)
assert result is True
assert key in store.providing_keys
assert key in store.providers
# Should have called _send_add_provider for each peer
assert mock_send.call_count == len(closest_peers)
@pytest.mark.trio
async def test_provide_skip_local_peer(self):
"""Test that provide() skips sending to local peer."""
# Setup mocks
mock_host = Mock()
mock_peer_routing = AsyncMock()
peer_id = ID.from_base58("QmTest123")
mock_host.get_id.return_value = peer_id
mock_host.get_addrs.return_value = [Multiaddr("/ip4/127.0.0.1/tcp/8000")]
# Include local peer in closest peers
closest_peers = [peer_id, ID.from_base58("QmPeer1")]
mock_peer_routing.find_closest_peers_network.return_value = closest_peers
store = ProviderStore(host=mock_host, peer_routing=mock_peer_routing)
with patch.object(store, "_send_add_provider", return_value=True) as mock_send:
key = b"test_key"
result = await store.provide(key)
assert result is True
# Should only call _send_add_provider once (skip local peer)
assert mock_send.call_count == 1
@pytest.mark.trio
async def test_find_providers_no_host(self):
"""Test find_providers() returns empty list when no host."""
store = ProviderStore(host=mock_host)
key = b"test_key"
result = await store.find_providers(key)
assert result == []
@pytest.mark.trio
async def test_find_providers_local_only(self):
"""Test find_providers() returns local providers."""
mock_host = Mock()
mock_peer_routing = Mock()
store = ProviderStore(host=mock_host, peer_routing=mock_peer_routing)
# Add local providers
key = b"test_key"
peer_id = ID.from_base58("QmTest123")
provider = PeerInfo(peer_id, [])
store.add_provider(key, provider)
result = await store.find_providers(key)
assert len(result) == 1
assert result[0].peer_id == peer_id
@pytest.mark.trio
async def test_find_providers_network_search(self):
"""Test find_providers() searches network when no local providers."""
mock_host = Mock()
mock_peer_routing = AsyncMock()
store = ProviderStore(host=mock_host, peer_routing=mock_peer_routing)
# Mock network search
closest_peers = [ID.from_base58("QmPeer1")]
mock_peer_routing.find_closest_peers_network.return_value = closest_peers
# Mock provider response
remote_peer_id = ID.from_base58("QmRemote123")
remote_providers = [PeerInfo(remote_peer_id, [])]
with patch.object(
store, "_get_providers_from_peer", return_value=remote_providers
):
key = b"test_key"
result = await store.find_providers(key)
assert len(result) == 1
assert result[0].peer_id == remote_peer_id
@pytest.mark.trio
async def test_get_providers_from_peer_no_host(self):
"""Test _get_providers_from_peer without host."""
store = ProviderStore(host=mock_host)
peer_id = ID.from_base58("QmTest123")
key = b"test_key"
# Should handle missing host gracefully
result = await store._get_providers_from_peer(peer_id, key)
assert result == []
def test_edge_case_empty_key(self):
"""Test handling of empty key."""
store = ProviderStore(host=mock_host)
key = b""
peer_id = ID.from_base58("QmTest123")
provider = PeerInfo(peer_id, [])
store.add_provider(key, provider)
providers = store.get_providers(key)
assert len(providers) == 1
assert providers[0].peer_id == peer_id
def test_edge_case_large_key(self):
"""Test handling of large key."""
store = ProviderStore(host=mock_host)
key = b"x" * 10000 # 10KB key
peer_id = ID.from_base58("QmTest123")
provider = PeerInfo(peer_id, [])
store.add_provider(key, provider)
providers = store.get_providers(key)
assert len(providers) == 1
assert providers[0].peer_id == peer_id
def test_concurrent_operations(self):
"""Test multiple concurrent operations."""
store = ProviderStore(host=mock_host)
# Add many providers
num_keys = 100
num_providers_per_key = 5
for i in range(num_keys):
_key = f"key_{i}".encode()
for j in range(num_providers_per_key):
# Generate unique valid Base58 peer IDs
# Use a different approach that ensures uniqueness
unique_id = i * num_providers_per_key + j + 1 # Ensure > 0
_peer_id_str = f"QmPeer{unique_id:06d}".replace("0", "A") + "1" * 38
peer_id = ID.from_base58(_peer_id_str)
provider = PeerInfo(peer_id, [])
store.add_provider(_key, provider)
# Verify total size
expected_size = num_keys * num_providers_per_key
assert store.size() == expected_size
# Verify individual keys
for i in range(num_keys):
_key = f"key_{i}".encode()
providers = store.get_providers(_key)
assert len(providers) == num_providers_per_key
def test_memory_efficiency_large_dataset(self):
"""Test memory behavior with large datasets."""
store = ProviderStore(host=mock_host)
# Add large number of providers
num_entries = 1000
for i in range(num_entries):
_key = f"key_{i:05d}".encode()
# Generate valid Base58 peer IDs (replace 0 with valid characters)
peer_str = f"QmPeer{i:05d}".replace("0", "1") + "1" * 35
peer_id = ID.from_base58(peer_str)
provider = PeerInfo(peer_id, [])
store.add_provider(_key, provider)
assert store.size() == num_entries
# Clean up all entries by making them expired
current_time = time.time()
for _key, providers in store.providers.items():
for _peer_id_str, record in providers.items():
record.timestamp = (
current_time - PROVIDER_RECORD_EXPIRATION_INTERVAL - 1
)
store.cleanup_expired()
assert store.size() == 0
assert len(store.providers) == 0
def test_unicode_key_handling(self):
"""Test handling of unicode content in keys."""
store = ProviderStore(host=mock_host)
# Test various unicode keys
unicode_keys = [
b"hello",
"héllo".encode(),
"🔑".encode(),
"ключ".encode(), # Russian
"键".encode(), # Chinese
]
for i, key in enumerate(unicode_keys):
# Generate valid Base58 peer IDs
peer_id = ID.from_base58(f"QmPeer{i + 1}" + "1" * 42) # Valid base58
provider = PeerInfo(peer_id, [])
store.add_provider(key, provider)
providers = store.get_providers(key)
assert len(providers) == 1
assert providers[0].peer_id == peer_id
def test_multiple_addresses_per_provider(self):
"""Test providers with multiple addresses."""
store = ProviderStore(host=mock_host)
key = b"test_key"
peer_id = ID.from_base58("QmTest123")
addresses = [
Multiaddr("/ip4/127.0.0.1/tcp/8000"),
Multiaddr("/ip6/::1/tcp/8001"),
Multiaddr("/ip4/192.168.1.100/tcp/8002"),
]
provider = PeerInfo(peer_id, addresses)
store.add_provider(key, provider)
providers = store.get_providers(key)
assert len(providers) == 1
assert providers[0].peer_id == peer_id
assert len(providers[0].addrs) == len(addresses)
assert all(addr in providers[0].addrs for addr in addresses)
@pytest.mark.trio
async def test_republish_provider_records_no_keys(self):
"""Test _republish_provider_records with no providing keys."""
store = ProviderStore(host=mock_host)
# Should complete without error even with no providing keys
await store._republish_provider_records()
assert len(store.providing_keys) == 0
def test_expiration_boundary_conditions(self):
"""Test expiration around boundary conditions."""
store = ProviderStore(host=mock_host)
peer_id = ID.from_base58("QmTest123")
provider = PeerInfo(peer_id, [])
current_time = time.time()
# Test records at various timestamps
timestamps = [
current_time, # Fresh
current_time - PROVIDER_ADDRESS_TTL + 1, # Addresses valid
current_time - PROVIDER_ADDRESS_TTL - 1, # Addresses expired
current_time
- PROVIDER_RECORD_REPUBLISH_INTERVAL
+ 1, # No republish needed
current_time - PROVIDER_RECORD_REPUBLISH_INTERVAL - 1, # Republish needed
current_time - PROVIDER_RECORD_EXPIRATION_INTERVAL + 1, # Not expired
current_time - PROVIDER_RECORD_EXPIRATION_INTERVAL - 1, # Expired
]
for i, timestamp in enumerate(timestamps):
test_key = f"key_{i}".encode()
record = ProviderRecord(provider, timestamp)
store.providers[test_key] = {str(peer_id): record}
# Test various operations
for i, timestamp in enumerate(timestamps):
test_key = f"key_{i}".encode()
providers = store.get_providers(test_key)
if timestamp <= current_time - PROVIDER_RECORD_EXPIRATION_INTERVAL:
# Should be expired and removed
assert len(providers) == 0
assert test_key not in store.providers
else:
# Should be present
assert len(providers) == 1
assert providers[0].peer_id == peer_id

371
tests/core/kad_dht/test_unit_routing_table.py Normal file
View File

@ -0,0 +1,371 @@
"""
Unit tests for the RoutingTable and KBucket classes in Kademlia DHT.
This module tests the core functionality of the routing table including:
- KBucket operations (add, remove, split, ping)
- RoutingTable management (peer addition, closest peer finding)
- Distance calculations and peer ordering
- Bucket splitting and range management
"""
import time
from unittest.mock import (
AsyncMock,
Mock,
patch,
)
import pytest
from multiaddr import (
Multiaddr,
)
import trio
from libp2p.crypto.secp256k1 import (
create_new_key_pair,
)
from libp2p.kad_dht.routing_table import (
BUCKET_SIZE,
KBucket,
RoutingTable,
)
from libp2p.kad_dht.utils import (
create_key_from_binary,
xor_distance,
)
from libp2p.peer.id import (
ID,
)
from libp2p.peer.peerinfo import (
PeerInfo,
)
def create_valid_peer_id(name: str) -> ID:
"""Create a valid peer ID for testing."""
# Use crypto to generate valid peer IDs
key_pair = create_new_key_pair()
return ID.from_pubkey(key_pair.public_key)
class TestKBucket:
"""Test suite for KBucket class."""
@pytest.fixture
def mock_host(self):
"""Create a mock host for testing."""
host = Mock()
host.get_peerstore.return_value = Mock()
host.new_stream = AsyncMock()
return host
@pytest.fixture
def sample_peer_info(self):
"""Create sample peer info for testing."""
peer_id = create_valid_peer_id("test")
addresses = [Multiaddr("/ip4/127.0.0.1/tcp/8000")]
return PeerInfo(peer_id, addresses)
def test_init_default_parameters(self, mock_host):
"""Test KBucket initialization with default parameters."""
bucket = KBucket(mock_host)
assert bucket.bucket_size == BUCKET_SIZE
assert bucket.host == mock_host
assert bucket.min_range == 0
assert bucket.max_range == 2**256
assert len(bucket.peers) == 0
def test_peer_operations(self, mock_host, sample_peer_info):
"""Test basic peer operations: add, check, and remove."""
bucket = KBucket(mock_host)
# Test empty bucket
assert bucket.peer_ids() == []
assert bucket.size() == 0
assert not bucket.has_peer(sample_peer_info.peer_id)
# Add peer manually
bucket.peers[sample_peer_info.peer_id] = (sample_peer_info, time.time())
# Test with peer
assert len(bucket.peer_ids()) == 1
assert sample_peer_info.peer_id in bucket.peer_ids()
assert bucket.size() == 1
assert bucket.has_peer(sample_peer_info.peer_id)
assert bucket.get_peer_info(sample_peer_info.peer_id) == sample_peer_info
# Remove peer
result = bucket.remove_peer(sample_peer_info.peer_id)
assert result is True
assert bucket.size() == 0
assert not bucket.has_peer(sample_peer_info.peer_id)
@pytest.mark.trio
async def test_add_peer_functionality(self, mock_host):
"""Test add_peer method with different scenarios."""
bucket = KBucket(mock_host, bucket_size=2) # Small bucket for testing
# Add first peer
peer1 = PeerInfo(create_valid_peer_id("peer1"), [])
result = await bucket.add_peer(peer1)
assert result is True
assert bucket.size() == 1
# Add second peer
peer2 = PeerInfo(create_valid_peer_id("peer2"), [])
result = await bucket.add_peer(peer2)
assert result is True
assert bucket.size() == 2
# Add same peer again (should update timestamp)
await trio.sleep(0.001)
result = await bucket.add_peer(peer1)
assert result is True
assert bucket.size() == 2 # Still 2 peers
# Try to add third peer when bucket is full
peer3 = PeerInfo(create_valid_peer_id("peer3"), [])
with patch.object(bucket, "_ping_peer", return_value=True):
result = await bucket.add_peer(peer3)
assert result is False # Should fail if oldest peer responds
def test_get_oldest_peer(self, mock_host):
"""Test get_oldest_peer method."""
bucket = KBucket(mock_host)
# Empty bucket
assert bucket.get_oldest_peer() is None
# Add peers with different timestamps
peer1 = PeerInfo(create_valid_peer_id("peer1"), [])
peer2 = PeerInfo(create_valid_peer_id("peer2"), [])
current_time = time.time()
bucket.peers[peer1.peer_id] = (peer1, current_time - 300) # Older
bucket.peers[peer2.peer_id] = (peer2, current_time) # Newer
oldest = bucket.get_oldest_peer()
assert oldest == peer1.peer_id
def test_stale_peers(self, mock_host):
"""Test stale peer identification."""
bucket = KBucket(mock_host)
current_time = time.time()
fresh_peer = PeerInfo(create_valid_peer_id("fresh"), [])
stale_peer = PeerInfo(create_valid_peer_id("stale"), [])
bucket.peers[fresh_peer.peer_id] = (fresh_peer, current_time)
bucket.peers[stale_peer.peer_id] = (
stale_peer,
current_time - 7200,
) # 2 hours ago
stale_peers = bucket.get_stale_peers(3600) # 1 hour threshold
assert len(stale_peers) == 1
assert stale_peer.peer_id in stale_peers
def test_key_in_range(self, mock_host):
"""Test key_in_range method."""
bucket = KBucket(mock_host, min_range=100, max_range=200)
# Test keys within range
key_in_range = (150).to_bytes(32, byteorder="big")
assert bucket.key_in_range(key_in_range) is True
# Test keys outside range
key_below = (50).to_bytes(32, byteorder="big")
assert bucket.key_in_range(key_below) is False
key_above = (250).to_bytes(32, byteorder="big")
assert bucket.key_in_range(key_above) is False
# Test boundary conditions
key_min = (100).to_bytes(32, byteorder="big")
assert bucket.key_in_range(key_min) is True
key_max = (200).to_bytes(32, byteorder="big")
assert bucket.key_in_range(key_max) is False
def test_split_bucket(self, mock_host):
"""Test bucket splitting functionality."""
bucket = KBucket(mock_host, min_range=0, max_range=256)
lower_bucket, upper_bucket = bucket.split()
# Check ranges
assert lower_bucket.min_range == 0
assert lower_bucket.max_range == 128
assert upper_bucket.min_range == 128
assert upper_bucket.max_range == 256
# Check properties
assert lower_bucket.bucket_size == bucket.bucket_size
assert upper_bucket.bucket_size == bucket.bucket_size
assert lower_bucket.host == mock_host
assert upper_bucket.host == mock_host
@pytest.mark.trio
async def test_ping_peer_scenarios(self, mock_host, sample_peer_info):
"""Test different ping scenarios."""
bucket = KBucket(mock_host)
bucket.peers[sample_peer_info.peer_id] = (sample_peer_info, time.time())
# Test ping peer not in bucket
other_peer_id = create_valid_peer_id("other")
with pytest.raises(ValueError, match="Peer .* not in bucket"):
await bucket._ping_peer(other_peer_id)
# Test ping failure due to stream error
mock_host.new_stream.side_effect = Exception("Stream failed")
result = await bucket._ping_peer(sample_peer_info.peer_id)
assert result is False
class TestRoutingTable:
"""Test suite for RoutingTable class."""
@pytest.fixture
def mock_host(self):
"""Create a mock host for testing."""
host = Mock()
host.get_peerstore.return_value = Mock()
return host
@pytest.fixture
def local_peer_id(self):
"""Create a local peer ID for testing."""
return create_valid_peer_id("local")
@pytest.fixture
def sample_peer_info(self):
"""Create sample peer info for testing."""
peer_id = create_valid_peer_id("sample")
addresses = [Multiaddr("/ip4/127.0.0.1/tcp/8000")]
return PeerInfo(peer_id, addresses)
def test_init_routing_table(self, mock_host, local_peer_id):
"""Test RoutingTable initialization."""
routing_table = RoutingTable(local_peer_id, mock_host)
assert routing_table.local_id == local_peer_id
assert routing_table.host == mock_host
assert len(routing_table.buckets) == 1
assert isinstance(routing_table.buckets[0], KBucket)
@pytest.mark.trio
async def test_add_peer_operations(
self, mock_host, local_peer_id, sample_peer_info
):
"""Test adding peers to routing table."""
routing_table = RoutingTable(local_peer_id, mock_host)
# Test adding peer with PeerInfo
result = await routing_table.add_peer(sample_peer_info)
assert result is True
assert routing_table.size() == 1
assert routing_table.peer_in_table(sample_peer_info.peer_id)
# Test adding peer with just ID
peer_id = create_valid_peer_id("test")
mock_addrs = [Multiaddr("/ip4/127.0.0.1/tcp/8001")]
mock_host.get_peerstore().addrs.return_value = mock_addrs
result = await routing_table.add_peer(peer_id)
assert result is True
assert routing_table.size() == 2
# Test adding peer with no addresses
no_addr_peer_id = create_valid_peer_id("no_addr")
mock_host.get_peerstore().addrs.return_value = []
result = await routing_table.add_peer(no_addr_peer_id)
assert result is False
assert routing_table.size() == 2
# Test adding local peer (should be ignored)
result = await routing_table.add_peer(local_peer_id)
assert result is False
assert routing_table.size() == 2
def test_find_bucket(self, mock_host, local_peer_id):
"""Test finding appropriate bucket for peers."""
routing_table = RoutingTable(local_peer_id, mock_host)
# Test with peer ID
peer_id = create_valid_peer_id("test")
bucket = routing_table.find_bucket(peer_id)
assert isinstance(bucket, KBucket)
def test_peer_management(self, mock_host, local_peer_id, sample_peer_info):
"""Test peer management operations."""
routing_table = RoutingTable(local_peer_id, mock_host)
# Add peer manually
bucket = routing_table.find_bucket(sample_peer_info.peer_id)
bucket.peers[sample_peer_info.peer_id] = (sample_peer_info, time.time())
# Test peer queries
assert routing_table.peer_in_table(sample_peer_info.peer_id)
assert routing_table.get_peer_info(sample_peer_info.peer_id) == sample_peer_info
assert routing_table.size() == 1
assert len(routing_table.get_peer_ids()) == 1
# Test remove peer
result = routing_table.remove_peer(sample_peer_info.peer_id)
assert result is True
assert not routing_table.peer_in_table(sample_peer_info.peer_id)
assert routing_table.size() == 0
def test_find_closest_peers(self, mock_host, local_peer_id):
"""Test finding closest peers."""
routing_table = RoutingTable(local_peer_id, mock_host)
# Empty table
target_key = create_key_from_binary(b"target_key")
closest_peers = routing_table.find_local_closest_peers(target_key, 5)
assert closest_peers == []
# Add some peers
bucket = routing_table.buckets[0]
test_peers = []
for i in range(5):
peer = PeerInfo(create_valid_peer_id(f"peer{i}"), [])
test_peers.append(peer)
bucket.peers[peer.peer_id] = (peer, time.time())
closest_peers = routing_table.find_local_closest_peers(target_key, 3)
assert len(closest_peers) <= 3
assert len(closest_peers) <= len(test_peers)
assert all(isinstance(peer_id, ID) for peer_id in closest_peers)
def test_distance_calculation(self, mock_host, local_peer_id):
"""Test XOR distance calculation."""
# Test same keys
key = b"\x42" * 32
distance = xor_distance(key, key)
assert distance == 0
# Test different keys
key1 = b"\x00" * 32
key2 = b"\xff" * 32
distance = xor_distance(key1, key2)
expected = int.from_bytes(b"\xff" * 32, byteorder="big")
assert distance == expected
def test_edge_cases(self, mock_host, local_peer_id):
"""Test various edge cases."""
routing_table = RoutingTable(local_peer_id, mock_host)
# Test with invalid peer ID
nonexistent_peer_id = create_valid_peer_id("nonexistent")
assert not routing_table.peer_in_table(nonexistent_peer_id)
assert routing_table.get_peer_info(nonexistent_peer_id) is None
assert routing_table.remove_peer(nonexistent_peer_id) is False
# Test bucket splitting scenario
assert len(routing_table.buckets) == 1
initial_bucket = routing_table.buckets[0]
assert initial_bucket.min_range == 0
assert initial_bucket.max_range == 2**256

504
tests/core/kad_dht/test_unit_value_store.py Normal file
View File

@ -0,0 +1,504 @@
"""
Unit tests for the ValueStore class in Kademlia DHT.
This module tests the core functionality of the ValueStore including:
- Basic storage and retrieval operations
- Expiration and TTL handling
- Edge cases and error conditions
- Store management operations
"""
import time
from unittest.mock import (
Mock,
)
import pytest
from libp2p.kad_dht.value_store import (
DEFAULT_TTL,
ValueStore,
)
from libp2p.peer.id import (
ID,
)
mock_host = Mock()
peer_id = ID.from_base58("QmTest123")
class TestValueStore:
"""Test suite for ValueStore class."""
def test_init_empty_store(self):
"""Test that a new ValueStore is initialized empty."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
assert len(store.store) == 0
def test_init_with_host_and_peer_id(self):
"""Test initialization with host and local peer ID."""
mock_host = Mock()
peer_id = ID.from_base58("QmTest123")
store = ValueStore(host=mock_host, local_peer_id=peer_id)
assert store.host == mock_host
assert store.local_peer_id == peer_id
assert len(store.store) == 0
def test_put_basic(self):
"""Test basic put operation."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"test_key"
value = b"test_value"
store.put(key, value)
assert key in store.store
stored_value, validity = store.store[key]
assert stored_value == value
assert validity is not None
assert validity > time.time() # Should be in the future
def test_put_with_custom_validity(self):
"""Test put operation with custom validity time."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"test_key"
value = b"test_value"
custom_validity = time.time() + 3600 # 1 hour from now
store.put(key, value, validity=custom_validity)
stored_value, validity = store.store[key]
assert stored_value == value
assert validity == custom_validity
def test_put_overwrite_existing(self):
"""Test that put overwrites existing values."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"test_key"
value1 = b"value1"
value2 = b"value2"
store.put(key, value1)
store.put(key, value2)
assert len(store.store) == 1
stored_value, _ = store.store[key]
assert stored_value == value2
def test_get_existing_valid_value(self):
"""Test retrieving an existing, non-expired value."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"test_key"
value = b"test_value"
store.put(key, value)
retrieved_value = store.get(key)
assert retrieved_value == value
def test_get_nonexistent_key(self):
"""Test retrieving a non-existent key returns None."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"nonexistent_key"
retrieved_value = store.get(key)
assert retrieved_value is None
def test_get_expired_value(self):
"""Test that expired values are automatically removed and return None."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"test_key"
value = b"test_value"
expired_validity = time.time() - 1 # 1 second ago
# Manually insert expired value
store.store[key] = (value, expired_validity)
retrieved_value = store.get(key)
assert retrieved_value is None
assert key not in store.store # Should be removed
def test_remove_existing_key(self):
"""Test removing an existing key."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"test_key"
value = b"test_value"
store.put(key, value)
result = store.remove(key)
assert result is True
assert key not in store.store
def test_remove_nonexistent_key(self):
"""Test removing a non-existent key returns False."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"nonexistent_key"
result = store.remove(key)
assert result is False
def test_has_existing_valid_key(self):
"""Test has() returns True for existing, valid keys."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"test_key"
value = b"test_value"
store.put(key, value)
result = store.has(key)
assert result is True
def test_has_nonexistent_key(self):
"""Test has() returns False for non-existent keys."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"nonexistent_key"
result = store.has(key)
assert result is False
def test_has_expired_key(self):
"""Test has() returns False for expired keys and removes them."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"test_key"
value = b"test_value"
expired_validity = time.time() - 1
# Manually insert expired value
store.store[key] = (value, expired_validity)
result = store.has(key)
assert result is False
assert key not in store.store # Should be removed
def test_cleanup_expired_no_expired_values(self):
"""Test cleanup when there are no expired values."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key1 = b"key1"
key2 = b"key2"
value = b"value"
store.put(key1, value)
store.put(key2, value)
expired_count = store.cleanup_expired()
assert expired_count == 0
assert len(store.store) == 2
def test_cleanup_expired_with_expired_values(self):
"""Test cleanup removes expired values."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key1 = b"valid_key"
key2 = b"expired_key1"
key3 = b"expired_key2"
value = b"value"
expired_validity = time.time() - 1
store.put(key1, value) # Valid
store.store[key2] = (value, expired_validity) # Expired
store.store[key3] = (value, expired_validity) # Expired
expired_count = store.cleanup_expired()
assert expired_count == 2
assert len(store.store) == 1
assert key1 in store.store
assert key2 not in store.store
assert key3 not in store.store
def test_cleanup_expired_mixed_validity_types(self):
"""Test cleanup with mix of values with and without expiration."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key1 = b"no_expiry"
key2 = b"valid_expiry"
key3 = b"expired"
value = b"value"
# No expiration (None validity)
store.put(key1, value)
# Valid expiration
store.put(key2, value, validity=time.time() + 3600)
# Expired
store.store[key3] = (value, time.time() - 1)
expired_count = store.cleanup_expired()
assert expired_count == 1
assert len(store.store) == 2
assert key1 in store.store
assert key2 in store.store
assert key3 not in store.store
def test_get_keys_empty_store(self):
"""Test get_keys() returns empty list for empty store."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
keys = store.get_keys()
assert keys == []
def test_get_keys_with_valid_values(self):
"""Test get_keys() returns all non-expired keys."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key1 = b"key1"
key2 = b"key2"
key3 = b"expired_key"
value = b"value"
store.put(key1, value)
store.put(key2, value)
store.store[key3] = (value, time.time() - 1) # Expired
keys = store.get_keys()
assert len(keys) == 2
assert key1 in keys
assert key2 in keys
assert key3 not in keys
def test_size_empty_store(self):
"""Test size() returns 0 for empty store."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
size = store.size()
assert size == 0
def test_size_with_valid_values(self):
"""Test size() returns correct count after cleaning expired values."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key1 = b"key1"
key2 = b"key2"
key3 = b"expired_key"
value = b"value"
store.put(key1, value)
store.put(key2, value)
store.store[key3] = (value, time.time() - 1) # Expired
size = store.size()
assert size == 2
def test_edge_case_empty_key(self):
"""Test handling of empty key."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b""
value = b"value"
store.put(key, value)
retrieved_value = store.get(key)
assert retrieved_value == value
def test_edge_case_empty_value(self):
"""Test handling of empty value."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"key"
value = b""
store.put(key, value)
retrieved_value = store.get(key)
assert retrieved_value == value
def test_edge_case_large_key_value(self):
"""Test handling of large keys and values."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"x" * 10000 # 10KB key
value = b"y" * 100000 # 100KB value
store.put(key, value)
retrieved_value = store.get(key)
assert retrieved_value == value
def test_edge_case_negative_validity(self):
"""Test handling of negative validity time."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"key"
value = b"value"
store.put(key, value, validity=-1)
# Should be expired
retrieved_value = store.get(key)
assert retrieved_value is None
def test_default_ttl_calculation(self):
"""Test that default TTL is correctly applied."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"key"
value = b"value"
start_time = time.time()
store.put(key, value)
_, validity = store.store[key]
expected_validity = start_time + DEFAULT_TTL
# Allow small time difference for execution
assert abs(validity - expected_validity) < 1
def test_concurrent_operations(self):
"""Test that multiple operations don't interfere with each other."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
# Add multiple key-value pairs
for i in range(100):
key = f"key_{i}".encode()
value = f"value_{i}".encode()
store.put(key, value)
# Verify all are stored
assert store.size() == 100
# Remove every other key
for i in range(0, 100, 2):
key = f"key_{i}".encode()
store.remove(key)
# Verify correct count
assert store.size() == 50
# Verify remaining keys are correct
for i in range(1, 100, 2):
key = f"key_{i}".encode()
assert store.has(key)
def test_expiration_boundary_conditions(self):
"""Test expiration around current time boundary."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key1 = b"key1"
key2 = b"key2"
key3 = b"key3"
value = b"value"
current_time = time.time()
# Just expired
store.store[key1] = (value, current_time - 0.001)
# Valid for a longer time to account for test execution time
store.store[key2] = (value, current_time + 1.0)
# Exactly current time (should be expired)
store.store[key3] = (value, current_time)
# Small delay to ensure time has passed
time.sleep(0.002)
assert not store.has(key1) # Should be expired
assert store.has(key2) # Should be valid
assert not store.has(key3) # Should be expired (exactly at current time)
def test_store_internal_structure(self):
"""Test that internal store structure is maintained correctly."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"key"
value = b"value"
validity = time.time() + 3600
store.put(key, value, validity=validity)
# Verify internal structure
assert isinstance(store.store, dict)
assert key in store.store
stored_tuple = store.store[key]
assert isinstance(stored_tuple, tuple)
assert len(stored_tuple) == 2
assert stored_tuple[0] == value
assert stored_tuple[1] == validity
@pytest.mark.trio
async def test_store_at_peer_local_peer(self):
"""Test _store_at_peer returns True when storing at local peer."""
mock_host = Mock()
peer_id = ID.from_base58("QmTest123")
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"key"
value = b"value"
result = await store._store_at_peer(peer_id, key, value)
assert result is True
@pytest.mark.trio
async def test_get_from_peer_local_peer(self):
"""Test _get_from_peer returns None when querying local peer."""
mock_host = Mock()
peer_id = ID.from_base58("QmTest123")
store = ValueStore(host=mock_host, local_peer_id=peer_id)
key = b"key"
result = await store._get_from_peer(peer_id, key)
assert result is None
def test_memory_efficiency_large_dataset(self):
"""Test memory behavior with large datasets."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
# Add a large number of entries
num_entries = 10000
for i in range(num_entries):
key = f"key_{i:05d}".encode()
value = f"value_{i:05d}".encode()
store.put(key, value)
assert store.size() == num_entries
# Clean up all entries
for i in range(num_entries):
key = f"key_{i:05d}".encode()
store.remove(key)
assert store.size() == 0
assert len(store.store) == 0
def test_key_collision_resistance(self):
"""Test that similar keys don't collide."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
# Test keys that might cause collisions
keys = [
b"key",
b"key\x00",
b"key1",
b"Key", # Different case
b"key ", # With space
b" key", # Leading space
]
for i, key in enumerate(keys):
value = f"value_{i}".encode()
store.put(key, value)
# Verify all keys are stored separately
assert store.size() == len(keys)
for i, key in enumerate(keys):
expected_value = f"value_{i}".encode()
assert store.get(key) == expected_value
def test_unicode_key_handling(self):
"""Test handling of unicode content in keys."""
store = ValueStore(host=mock_host, local_peer_id=peer_id)
# Test various unicode keys
unicode_keys = [
b"hello",
"héllo".encode(),
"🔑".encode(),
"ключ".encode(), # Russian
"键".encode(), # Chinese
]
for i, key in enumerate(unicode_keys):
value = f"value_{i}".encode()
store.put(key, value)
assert store.get(key) == value

82
tests/core/pubsub/test_gossipsub.py
View File

@ -17,6 +17,7 @@ from tests.utils.factories import (
from tests.utils.pubsub.utils import (
dense_connect,
one_to_all_connect,
sparse_connect,
)
@ -506,3 +507,84 @@ async def test_gossip_heartbeat(initial_peer_count, monkeypatch):
# Check that the peer to gossip to is not in our fanout peers
assert peer not in fanout_peers
assert topic_fanout in peers_to_gossip[peer]
@pytest.mark.trio
async def test_dense_connect_fallback():
"""Test that sparse_connect falls back to dense connect for small networks."""
async with PubsubFactory.create_batch_with_gossipsub(3) as pubsubs_gsub:
hosts = [pubsub.host for pubsub in pubsubs_gsub]
degree = 2
# Create network (should use dense connect)
await sparse_connect(hosts, degree)
# Wait for connections to be established
await trio.sleep(2)
# Verify dense topology (all nodes connected to each other)
for i, pubsub in enumerate(pubsubs_gsub):
connected_peers = len(pubsub.peers)
expected_connections = len(hosts) - 1
assert connected_peers == expected_connections, (
f"Host {i} has {connected_peers} connections, "
f"expected {expected_connections} in dense mode"
)
@pytest.mark.trio
async def test_sparse_connect():
"""Test sparse connect functionality and message propagation."""
async with PubsubFactory.create_batch_with_gossipsub(10) as pubsubs_gsub:
hosts = [pubsub.host for pubsub in pubsubs_gsub]
degree = 2
topic = "test_topic"
# Create network (should use sparse connect)
await sparse_connect(hosts, degree)
# Wait for connections to be established
await trio.sleep(2)
# Verify sparse topology
for i, pubsub in enumerate(pubsubs_gsub):
connected_peers = len(pubsub.peers)
assert degree <= connected_peers < len(hosts) - 1, (
f"Host {i} has {connected_peers} connections, "
f"expected between {degree} and {len(hosts) - 1} in sparse mode"
)
# Test message propagation
queues = [await pubsub.subscribe(topic) for pubsub in pubsubs_gsub]
await trio.sleep(2)
# Publish and verify message propagation
msg_content = b"test_msg"
await pubsubs_gsub[0].publish(topic, msg_content)
await trio.sleep(2)
# Verify message propagation - ideally all nodes should receive it
received_count = 0
for queue in queues:
try:
msg = await queue.get()
if msg.data == msg_content:
received_count += 1
except Exception:
continue
total_nodes = len(pubsubs_gsub)
# Ideally all nodes should receive the message for optimal scalability
if received_count == total_nodes:
# Perfect propagation achieved
pass
else:
# require more than half for acceptable scalability
min_required = (total_nodes + 1) // 2
assert received_count >= min_required, (
f"Message propagation insufficient: "
f"{received_count}/{total_nodes} nodes "
f"received the message. Ideally all nodes should receive it, but at "
f"minimum {min_required} required for sparse network scalability."
)

263
tests/core/relay/test_circuit_v2_discovery.py Normal file
View File

@ -0,0 +1,263 @@
"""Tests for the Circuit Relay v2 discovery functionality."""
import logging
import time
import pytest
import trio
from libp2p.relay.circuit_v2.discovery import (
RelayDiscovery,
)
from libp2p.relay.circuit_v2.pb import circuit_pb2 as proto
from libp2p.relay.circuit_v2.protocol import (
PROTOCOL_ID,
STOP_PROTOCOL_ID,
)
from libp2p.tools.async_service import (
background_trio_service,
)
from libp2p.tools.constants import (
MAX_READ_LEN,
)
from libp2p.tools.utils import (
connect,
)
from tests.utils.factories import (
HostFactory,
)
logger = logging.getLogger(__name__)
# Test timeouts
CONNECT_TIMEOUT = 15 # seconds
STREAM_TIMEOUT = 15 # seconds
HANDLER_TIMEOUT = 15 # seconds
SLEEP_TIME = 1.0 # seconds
DISCOVERY_TIMEOUT = 20 # seconds
# Make a simple stream handler for testing
async def simple_stream_handler(stream):
"""Simple stream handler that reads a message and responds with OK status."""
logger.info("Simple stream handler invoked")
try:
# Read the request
request_data = await stream.read(MAX_READ_LEN)
if not request_data:
logger.error("Empty request received")
return
# Parse request
request = proto.HopMessage()
request.ParseFromString(request_data)
logger.info("Received request: type=%s", request.type)
# Only handle RESERVE requests
if request.type == proto.HopMessage.RESERVE:
# Create a valid response
response = proto.HopMessage(
type=proto.HopMessage.RESERVE,
status=proto.Status(
code=proto.Status.OK,
message="Test reservation accepted",
),
reservation=proto.Reservation(
expire=int(time.time()) + 3600, # 1 hour from now
voucher=b"test-voucher",
signature=b"",
),
limit=proto.Limit(
duration=3600, # 1 hour
data=1024 * 1024 * 1024, # 1GB
),
)
# Send the response
logger.info("Sending response")
await stream.write(response.SerializeToString())
logger.info("Response sent")
except Exception as e:
logger.error("Error in simple stream handler: %s", str(e))
finally:
# Keep stream open to allow client to read response
await trio.sleep(1)
await stream.close()
@pytest.mark.trio
async def test_relay_discovery_initialization():
"""Test Circuit v2 relay discovery initializes correctly with default settings."""
async with HostFactory.create_batch_and_listen(1) as hosts:
host = hosts[0]
discovery = RelayDiscovery(host)
async with background_trio_service(discovery):
await discovery.event_started.wait()
await trio.sleep(SLEEP_TIME) # Give time for discovery to start
# Verify discovery is initialized correctly
assert discovery.host == host, "Host not set correctly"
assert discovery.is_running, "Discovery service should be running"
assert hasattr(discovery, "_discovered_relays"), (
"Discovery should track discovered relays"
)
@pytest.mark.trio
async def test_relay_discovery_find_relay():
"""Test finding a relay node via discovery."""
async with HostFactory.create_batch_and_listen(2) as hosts:
relay_host, client_host = hosts
logger.info("Created hosts for test_relay_discovery_find_relay")
logger.info("Relay host ID: %s", relay_host.get_id())
logger.info("Client host ID: %s", client_host.get_id())
# Explicitly register the protocol handlers on relay_host
relay_host.set_stream_handler(PROTOCOL_ID, simple_stream_handler)
relay_host.set_stream_handler(STOP_PROTOCOL_ID, simple_stream_handler)
# Manually add protocol to peerstore for testing
# This simulates what the real relay protocol would do
client_host.get_peerstore().add_protocols(
relay_host.get_id(), [str(PROTOCOL_ID)]
)
# Set up discovery on the client host
client_discovery = RelayDiscovery(
client_host, discovery_interval=5
) # Use shorter interval for testing
try:
# Connect peers so they can discover each other
with trio.fail_after(CONNECT_TIMEOUT):
logger.info("Connecting client host to relay host")
await connect(client_host, relay_host)
assert relay_host.get_network().connections[client_host.get_id()], (
"Peers not connected"
)
logger.info("Connection established between peers")
except Exception as e:
logger.error("Failed to connect peers: %s", str(e))
raise
# Start discovery service
async with background_trio_service(client_discovery):
await client_discovery.event_started.wait()
logger.info("Client discovery service started")
# Wait for discovery to find the relay
logger.info("Waiting for relay discovery...")
# Manually trigger discovery instead of waiting
await client_discovery.discover_relays()
# Check if relay was found
with trio.fail_after(DISCOVERY_TIMEOUT):
for _ in range(20): # Try multiple times
if relay_host.get_id() in client_discovery._discovered_relays:
logger.info("Relay discovered successfully")
break
# Wait and try again
await trio.sleep(1)
# Manually trigger discovery again
await client_discovery.discover_relays()
else:
pytest.fail("Failed to discover relay node within timeout")
# Verify that relay was found and is valid
assert relay_host.get_id() in client_discovery._discovered_relays, (
"Relay should be discovered"
)
relay_info = client_discovery._discovered_relays[relay_host.get_id()]
assert relay_info.peer_id == relay_host.get_id(), "Peer ID should match"
@pytest.mark.trio
async def test_relay_discovery_auto_reservation():
"""Test that discovery can auto-reserve with discovered relays."""
async with HostFactory.create_batch_and_listen(2) as hosts:
relay_host, client_host = hosts
logger.info("Created hosts for test_relay_discovery_auto_reservation")
logger.info("Relay host ID: %s", relay_host.get_id())
logger.info("Client host ID: %s", client_host.get_id())
# Explicitly register the protocol handlers on relay_host
relay_host.set_stream_handler(PROTOCOL_ID, simple_stream_handler)
relay_host.set_stream_handler(STOP_PROTOCOL_ID, simple_stream_handler)
# Manually add protocol to peerstore for testing
client_host.get_peerstore().add_protocols(
relay_host.get_id(), [str(PROTOCOL_ID)]
)
# Set up discovery on the client host with auto-reservation enabled
client_discovery = RelayDiscovery(
client_host, auto_reserve=True, discovery_interval=5
)
try:
# Connect peers so they can discover each other
with trio.fail_after(CONNECT_TIMEOUT):
logger.info("Connecting client host to relay host")
await connect(client_host, relay_host)
assert relay_host.get_network().connections[client_host.get_id()], (
"Peers not connected"
)
logger.info("Connection established between peers")
except Exception as e:
logger.error("Failed to connect peers: %s", str(e))
raise
# Start discovery service
async with background_trio_service(client_discovery):
await client_discovery.event_started.wait()
logger.info("Client discovery service started")
# Wait for discovery to find the relay and make a reservation
logger.info("Waiting for relay discovery and auto-reservation...")
# Manually trigger discovery
await client_discovery.discover_relays()
# Check if relay was found and reservation was made
with trio.fail_after(DISCOVERY_TIMEOUT):
for _ in range(20): # Try multiple times
relay_found = (
relay_host.get_id() in client_discovery._discovered_relays
)
has_reservation = (
relay_found
and client_discovery._discovered_relays[
relay_host.get_id()
].has_reservation
)
if has_reservation:
logger.info(
"Relay discovered and reservation made successfully"
)
break
# Wait and try again
await trio.sleep(1)
# Try to make reservation manually
if relay_host.get_id() in client_discovery._discovered_relays:
await client_discovery.make_reservation(relay_host.get_id())
else:
pytest.fail(
"Failed to discover relay and make reservation within timeout"
)
# Verify that relay was found and reservation was made
assert relay_host.get_id() in client_discovery._discovered_relays, (
"Relay should be discovered"
)
relay_info = client_discovery._discovered_relays[relay_host.get_id()]
assert relay_info.has_reservation, "Reservation should be made"
assert relay_info.reservation_expires_at is not None, (
"Reservation should have expiry time"
)
assert relay_info.reservation_data_limit is not None, (
"Reservation should have data limit"
)

665
tests/core/relay/test_circuit_v2_protocol.py Normal file
View File

@ -0,0 +1,665 @@
"""Tests for the Circuit Relay v2 protocol."""
import logging
import time
from typing import Any
import pytest
import trio
from libp2p.network.stream.exceptions import (
StreamEOF,
StreamError,
StreamReset,
)
from libp2p.peer.id import (
ID,
)
from libp2p.relay.circuit_v2.pb import circuit_pb2 as proto
from libp2p.relay.circuit_v2.protocol import (
DEFAULT_RELAY_LIMITS,
PROTOCOL_ID,
STOP_PROTOCOL_ID,
CircuitV2Protocol,
)
from libp2p.relay.circuit_v2.resources import (
RelayLimits,
)
from libp2p.tools.async_service import (
background_trio_service,
)
from libp2p.tools.constants import (
MAX_READ_LEN,
)
from libp2p.tools.utils import (
connect,
)
from tests.utils.factories import (
HostFactory,
)
logger = logging.getLogger(__name__)
# Test timeouts
CONNECT_TIMEOUT = 15 # seconds (increased)
STREAM_TIMEOUT = 15 # seconds (increased)
HANDLER_TIMEOUT = 15 # seconds (increased)
SLEEP_TIME = 1.0 # seconds (increased)
async def assert_stream_response(
stream, expected_type, expected_status, retries=5, retry_delay=1.0
):
"""Helper function to assert stream response matches expectations."""
last_error = None
all_responses = []
# Increase initial sleep to ensure response has time to arrive
await trio.sleep(retry_delay * 2)
for attempt in range(retries):
try:
with trio.fail_after(STREAM_TIMEOUT):
# Wait between attempts
if attempt > 0:
await trio.sleep(retry_delay)
# Try to read response
logger.debug("Attempt %d: Reading response from stream", attempt + 1)
response_bytes = await stream.read(MAX_READ_LEN)
# Check if we got any data
if not response_bytes:
logger.warning(
"Attempt %d: No data received from stream", attempt + 1
)
last_error = "No response received"
if attempt < retries - 1: # Not the last attempt
continue
raise AssertionError(
f"No response received after {retries} attempts"
)
# Try to parse the response
response = proto.HopMessage()
try:
response.ParseFromString(response_bytes)
# Log what we received
logger.debug(
"Attempt %d: Received HOP response: type=%s, status=%s",
attempt + 1,
response.type,
response.status.code
if response.HasField("status")
else "No status",
)
all_responses.append(
{
"type": response.type,
"status": response.status.code
if response.HasField("status")
else None,
"message": response.status.message
if response.HasField("status")
else None,
}
)
# Accept any valid response with the right status
if (
expected_status is not None
and response.HasField("status")
and response.status.code == expected_status
):
if response.type != expected_type:
logger.warning(
"Type mismatch (%s, got %s) but status ok - accepting",
expected_type,
response.type,
)
logger.debug("Successfully validated response (status matched)")
return response
# Check message type specifically if it matters
if response.type != expected_type:
logger.warning(
"Wrong response type: expected %s, got %s",
expected_type,
response.type,
)
last_error = (
f"Wrong response type: expected {expected_type}, "
f"got {response.type}"
)
if attempt < retries - 1: # Not the last attempt
continue
# Check status code if present
if response.HasField("status"):
if response.status.code != expected_status:
logger.warning(
"Wrong status code: expected %s, got %s",
expected_status,
response.status.code,
)
last_error = (
f"Wrong status code: expected {expected_status}, "
f"got {response.status.code}"
)
if attempt < retries - 1: # Not the last attempt
continue
elif expected_status is not None:
logger.warning(
"Expected status %s but none was present in response",
expected_status,
)
last_error = (
f"Expected status {expected_status} but none was present"
)
if attempt < retries - 1: # Not the last attempt
continue
logger.debug("Successfully validated response")
return response
except Exception as e:
# If parsing as HOP message fails, try parsing as STOP message
logger.warning(
"Failed to parse as HOP message, trying STOP message: %s",
str(e),
)
try:
stop_msg = proto.StopMessage()
stop_msg.ParseFromString(response_bytes)
logger.debug("Parsed as STOP message: type=%s", stop_msg.type)
# Create a simplified response dictionary
has_status = stop_msg.HasField("status")
status_code = None
status_message = None
if has_status:
status_code = stop_msg.status.code
status_message = stop_msg.status.message
response_dict: dict[str, Any] = {
"stop_type": stop_msg.type, # Keep original type
"status": status_code, # Keep original type
"message": status_message, # Keep original type
}
all_responses.append(response_dict)
last_error = "Got STOP message instead of HOP message"
if attempt < retries - 1: # Not the last attempt
continue
except Exception as e2:
logger.warning(
"Failed to parse response as either message type: %s",
str(e2),
)
last_error = (
f"Failed to parse response: {str(e)}, then {str(e2)}"
)
if attempt < retries - 1: # Not the last attempt
continue
except trio.TooSlowError:
logger.warning(
"Attempt %d: Timeout waiting for stream response", attempt + 1
)
last_error = "Timeout waiting for stream response"
if attempt < retries - 1: # Not the last attempt
continue
except (StreamError, StreamReset, StreamEOF) as e:
logger.warning(
"Attempt %d: Stream error while reading response: %s",
attempt + 1,
str(e),
)
last_error = f"Stream error: {str(e)}"
if attempt < retries - 1: # Not the last attempt
continue
except AssertionError as e:
logger.warning("Attempt %d: Assertion failed: %s", attempt + 1, str(e))
last_error = str(e)
if attempt < retries - 1: # Not the last attempt
continue
except Exception as e:
logger.warning("Attempt %d: Unexpected error: %s", attempt + 1, str(e))
last_error = f"Unexpected error: {str(e)}"
if attempt < retries - 1: # Not the last attempt
continue
# If we've reached here, all retries failed
all_responses_str = ", ".join([str(r) for r in all_responses])
error_msg = (
f"Failed to get expected response after {retries} attempts. "
f"Last error: {last_error}. All responses: {all_responses_str}"
)
raise AssertionError(error_msg)
async def close_stream(stream):
"""Helper function to safely close a stream."""
if stream is not None:
try:
logger.debug("Closing stream")
await stream.close()
# Wait a bit to ensure the close is processed
await trio.sleep(SLEEP_TIME)
logger.debug("Stream closed successfully")
except (StreamError, Exception) as e:
logger.warning("Error closing stream: %s. Attempting to reset.", str(e))
try:
await stream.reset()
# Wait a bit to ensure the reset is processed
await trio.sleep(SLEEP_TIME)
logger.debug("Stream reset successfully")
except Exception as e:
logger.warning("Error resetting stream: %s", str(e))
@pytest.mark.trio
async def test_circuit_v2_protocol_initialization():
"""Test that the Circuit v2 protocol initializes correctly with default settings."""
async with HostFactory.create_batch_and_listen(1) as hosts:
host = hosts[0]
limits = RelayLimits(
duration=DEFAULT_RELAY_LIMITS.duration,
data=DEFAULT_RELAY_LIMITS.data,
max_circuit_conns=DEFAULT_RELAY_LIMITS.max_circuit_conns,
max_reservations=DEFAULT_RELAY_LIMITS.max_reservations,
)
protocol = CircuitV2Protocol(host, limits, allow_hop=True)
async with background_trio_service(protocol):
await protocol.event_started.wait()
await trio.sleep(SLEEP_TIME) # Give time for handlers to be registered
# Verify protocol handlers are registered by trying to use them
test_stream = None
try:
with trio.fail_after(STREAM_TIMEOUT):
test_stream = await host.new_stream(host.get_id(), [PROTOCOL_ID])
assert test_stream is not None, (
"HOP protocol handler not registered"
)
except Exception:
pass
finally:
await close_stream(test_stream)
try:
with trio.fail_after(STREAM_TIMEOUT):
test_stream = await host.new_stream(
host.get_id(), [STOP_PROTOCOL_ID]
)
assert test_stream is not None, (
"STOP protocol handler not registered"
)
except Exception:
pass
finally:
await close_stream(test_stream)
assert len(protocol.resource_manager._reservations) == 0, (
"Reservations should be empty"
)
@pytest.mark.trio
async def test_circuit_v2_reservation_basic():
"""Test basic reservation functionality between two peers."""
async with HostFactory.create_batch_and_listen(2) as hosts:
relay_host, client_host = hosts
logger.info("Created hosts for test_circuit_v2_reservation_basic")
logger.info("Relay host ID: %s", relay_host.get_id())
logger.info("Client host ID: %s", client_host.get_id())
# Custom handler that responds directly with a valid response
# This bypasses the complex protocol implementation that might have issues
async def mock_reserve_handler(stream):
# Read the request
logger.info("Mock handler received stream request")
try:
request_data = await stream.read(MAX_READ_LEN)
request = proto.HopMessage()
request.ParseFromString(request_data)
logger.info("Mock handler parsed request: type=%s", request.type)
# Only handle RESERVE requests
if request.type == proto.HopMessage.RESERVE:
# Create a valid response
response = proto.HopMessage(
type=proto.HopMessage.RESERVE,
status=proto.Status(
code=proto.Status.OK,
message="Reservation accepted",
),
reservation=proto.Reservation(
expire=int(time.time()) + 3600, # 1 hour from now
voucher=b"test-voucher",
signature=b"",
),
limit=proto.Limit(
duration=3600, # 1 hour
data=1024 * 1024 * 1024, # 1GB
),
)
# Send the response
logger.info("Mock handler sending response")
await stream.write(response.SerializeToString())
logger.info("Mock handler sent response")
# Keep stream open for client to read response
await trio.sleep(5)
except Exception as e:
logger.error("Error in mock handler: %s", str(e))
# Register the mock handler
relay_host.set_stream_handler(PROTOCOL_ID, mock_reserve_handler)
logger.info("Registered mock handler for %s", PROTOCOL_ID)
# Connect peers
try:
with trio.fail_after(CONNECT_TIMEOUT):
logger.info("Connecting client host to relay host")
await connect(client_host, relay_host)
assert relay_host.get_network().connections[client_host.get_id()], (
"Peers not connected"
)
logger.info("Connection established between peers")
except Exception as e:
logger.error("Failed to connect peers: %s", str(e))
raise
# Wait a bit to ensure connection is fully established
await trio.sleep(SLEEP_TIME)
stream = None
try:
# Open stream and send reservation request
logger.info("Opening stream from client to relay")
with trio.fail_after(STREAM_TIMEOUT):
stream = await client_host.new_stream(
relay_host.get_id(), [PROTOCOL_ID]
)
assert stream is not None, "Failed to open stream"
logger.info("Preparing reservation request")
request = proto.HopMessage(
type=proto.HopMessage.RESERVE, peer=client_host.get_id().to_bytes()
)
logger.info("Sending reservation request")
await stream.write(request.SerializeToString())
logger.info("Reservation request sent")
# Wait to ensure the request is processed
await trio.sleep(SLEEP_TIME)
# Read response directly
logger.info("Reading response directly")
response_bytes = await stream.read(MAX_READ_LEN)
assert response_bytes, "No response received"
# Parse response
response = proto.HopMessage()
response.ParseFromString(response_bytes)
# Verify response
assert response.type == proto.HopMessage.RESERVE, (
f"Wrong response type: {response.type}"
)
assert response.HasField("status"), "No status field"
assert response.status.code == proto.Status.OK, (
f"Wrong status code: {response.status.code}"
)
# Verify reservation details
assert response.HasField("reservation"), "No reservation field"
assert response.HasField("limit"), "No limit field"
assert response.limit.duration == 3600, (
f"Wrong duration: {response.limit.duration}"
)
assert response.limit.data == 1024 * 1024 * 1024, (
f"Wrong data limit: {response.limit.data}"
)
logger.info("Verified reservation details in response")
except Exception as e:
logger.error("Error in reservation test: %s", str(e))
raise
finally:
if stream:
await close_stream(stream)
@pytest.mark.trio
async def test_circuit_v2_reservation_limit():
"""Test that relay enforces reservation limits."""
async with HostFactory.create_batch_and_listen(3) as hosts:
relay_host, client1_host, client2_host = hosts
logger.info("Created hosts for test_circuit_v2_reservation_limit")
logger.info("Relay host ID: %s", relay_host.get_id())
logger.info("Client1 host ID: %s", client1_host.get_id())
logger.info("Client2 host ID: %s", client2_host.get_id())
# Track reservation status to simulate limits
reserved_clients = set()
max_reservations = 1 # Only allow one reservation
# Custom handler that responds based on reservation limits
async def mock_reserve_handler(stream):
# Read the request
logger.info("Mock handler received stream request")
try:
request_data = await stream.read(MAX_READ_LEN)
request = proto.HopMessage()
request.ParseFromString(request_data)
logger.info("Mock handler parsed request: type=%s", request.type)
# Only handle RESERVE requests
if request.type == proto.HopMessage.RESERVE:
# Extract peer ID from request
peer_id = ID(request.peer)
logger.info(
"Mock handler received reservation request from %s", peer_id
)
# Check if we've reached reservation limit
if (
peer_id in reserved_clients
or len(reserved_clients) < max_reservations
):
# Accept the reservation
if peer_id not in reserved_clients:
reserved_clients.add(peer_id)
# Create a success response
response = proto.HopMessage(
type=proto.HopMessage.RESERVE,
status=proto.Status(
code=proto.Status.OK,
message="Reservation accepted",
),
reservation=proto.Reservation(
expire=int(time.time()) + 3600, # 1 hour from now
voucher=b"test-voucher",
signature=b"",
),
limit=proto.Limit(
duration=3600, # 1 hour
data=1024 * 1024 * 1024, # 1GB
),
)
logger.info(
"Mock handler accepting reservation for %s", peer_id
)
else:
# Reject the reservation due to limits
response = proto.HopMessage(
type=proto.HopMessage.RESERVE,
status=proto.Status(
code=proto.Status.RESOURCE_LIMIT_EXCEEDED,
message="Reservation limit exceeded",
),
)
logger.info(
"Mock handler rejecting reservation for %s due to limit",
peer_id,
)
# Send the response
logger.info("Mock handler sending response")
await stream.write(response.SerializeToString())
logger.info("Mock handler sent response")
# Keep stream open for client to read response
await trio.sleep(5)
except Exception as e:
logger.error("Error in mock handler: %s", str(e))
# Register the mock handler
relay_host.set_stream_handler(PROTOCOL_ID, mock_reserve_handler)
logger.info("Registered mock handler for %s", PROTOCOL_ID)
# Connect peers
try:
with trio.fail_after(CONNECT_TIMEOUT):
logger.info("Connecting client1 to relay")
await connect(client1_host, relay_host)
logger.info("Connecting client2 to relay")
await connect(client2_host, relay_host)
assert relay_host.get_network().connections[client1_host.get_id()], (
"Client1 not connected"
)
assert relay_host.get_network().connections[client2_host.get_id()], (
"Client2 not connected"
)
logger.info("All connections established")
except Exception as e:
logger.error("Failed to connect peers: %s", str(e))
raise
# Wait a bit to ensure connections are fully established
await trio.sleep(SLEEP_TIME)
stream1, stream2 = None, None
try:
# Client 1 reservation (should succeed)
logger.info("Testing client1 reservation (should succeed)")
with trio.fail_after(STREAM_TIMEOUT):
logger.info("Opening stream for client1")
stream1 = await client1_host.new_stream(
relay_host.get_id(), [PROTOCOL_ID]
)
assert stream1 is not None, "Failed to open stream for client 1"
logger.info("Preparing reservation request for client1")
request1 = proto.HopMessage(
type=proto.HopMessage.RESERVE, peer=client1_host.get_id().to_bytes()
)
logger.info("Sending reservation request for client1")
await stream1.write(request1.SerializeToString())
logger.info("Sent reservation request for client1")
# Wait to ensure the request is processed
await trio.sleep(SLEEP_TIME)
# Read response directly
logger.info("Reading response for client1")
response_bytes = await stream1.read(MAX_READ_LEN)
assert response_bytes, "No response received for client1"
# Parse response
response1 = proto.HopMessage()
response1.ParseFromString(response_bytes)
# Verify response
assert response1.type == proto.HopMessage.RESERVE, (
f"Wrong response type: {response1.type}"
)
assert response1.HasField("status"), "No status field"
assert response1.status.code == proto.Status.OK, (
f"Wrong status code: {response1.status.code}"
)
# Verify reservation details
assert response1.HasField("reservation"), "No reservation field"
assert response1.HasField("limit"), "No limit field"
assert response1.limit.duration == 3600, (
f"Wrong duration: {response1.limit.duration}"
)
assert response1.limit.data == 1024 * 1024 * 1024, (
f"Wrong data limit: {response1.limit.data}"
)
logger.info("Verified reservation details for client1")
# Close stream1 before opening stream2
await close_stream(stream1)
stream1 = None
logger.info("Closed client1 stream")
# Wait a bit to ensure stream is fully closed
await trio.sleep(SLEEP_TIME)
# Client 2 reservation (should fail)
logger.info("Testing client2 reservation (should fail)")
stream2 = await client2_host.new_stream(
relay_host.get_id(), [PROTOCOL_ID]
)
assert stream2 is not None, "Failed to open stream for client 2"
logger.info("Preparing reservation request for client2")
request2 = proto.HopMessage(
type=proto.HopMessage.RESERVE, peer=client2_host.get_id().to_bytes()
)
logger.info("Sending reservation request for client2")
await stream2.write(request2.SerializeToString())
logger.info("Sent reservation request for client2")
# Wait to ensure the request is processed
await trio.sleep(SLEEP_TIME)
# Read response directly
logger.info("Reading response for client2")
response_bytes = await stream2.read(MAX_READ_LEN)
assert response_bytes, "No response received for client2"
# Parse response
response2 = proto.HopMessage()
response2.ParseFromString(response_bytes)
# Verify response
assert response2.type == proto.HopMessage.RESERVE, (
f"Wrong response type: {response2.type}"
)
assert response2.HasField("status"), "No status field"
assert response2.status.code == proto.Status.RESOURCE_LIMIT_EXCEEDED, (
f"Wrong status code: {response2.status.code}, "
f"expected RESOURCE_LIMIT_EXCEEDED"
)
logger.info("Verified client2 was correctly rejected")
# Verify reservation tracking is correct
assert len(reserved_clients) == 1, "Should have exactly one reservation"
assert client1_host.get_id() in reserved_clients, (
"Client1 should be reserved"
)
assert client2_host.get_id() not in reserved_clients, (
"Client2 should not be reserved"
)
logger.info("Verified reservation tracking state")
except Exception as e:
logger.error("Error in reservation limit test: %s", str(e))
# Diagnostic information
logger.error("Current reservations: %s", reserved_clients)
raise
finally:
await close_stream(stream1)
await close_stream(stream2)

346
tests/core/relay/test_circuit_v2_transport.py Normal file
View File

@ -0,0 +1,346 @@
"""Tests for the Circuit Relay v2 transport functionality."""
import logging
import time
import pytest
import trio
from libp2p.custom_types import TProtocol
from libp2p.network.stream.exceptions import (
StreamEOF,
StreamReset,
)
from libp2p.relay.circuit_v2.config import (
RelayConfig,
)
from libp2p.relay.circuit_v2.discovery import (
RelayDiscovery,
RelayInfo,
)
from libp2p.relay.circuit_v2.protocol import (
CircuitV2Protocol,
RelayLimits,
)
from libp2p.relay.circuit_v2.transport import (
CircuitV2Transport,
)
from libp2p.tools.constants import (
MAX_READ_LEN,
)
from libp2p.tools.utils import (
connect,
)
from tests.utils.factories import (
HostFactory,
)
logger = logging.getLogger(__name__)
# Test timeouts
CONNECT_TIMEOUT = 15 # seconds
STREAM_TIMEOUT = 15 # seconds
HANDLER_TIMEOUT = 15 # seconds
SLEEP_TIME = 1.0 # seconds
RELAY_TIMEOUT = 20 # seconds
# Default limits for relay
DEFAULT_RELAY_LIMITS = RelayLimits(
duration=60 * 60, # 1 hour
data=1024 * 1024 * 10, # 10 MB
max_circuit_conns=8, # 8 active relay connections
max_reservations=4, # 4 active reservations
)
# Message for testing
TEST_MESSAGE = b"Hello, Circuit Relay!"
TEST_RESPONSE = b"Hello from the other side!"
# Stream handler for testing
async def echo_stream_handler(stream):
"""Simple echo handler that responds to messages."""
logger.info("Echo handler received stream")
try:
while True:
data = await stream.read(MAX_READ_LEN)
if not data:
logger.info("Stream closed by remote")
break
logger.info("Received data: %s", data)
await stream.write(TEST_RESPONSE)
logger.info("Sent response")
except (StreamEOF, StreamReset) as e:
logger.info("Stream ended: %s", str(e))
except Exception as e:
logger.error("Error in echo handler: %s", str(e))
finally:
await stream.close()
@pytest.mark.trio
async def test_circuit_v2_transport_initialization():
"""Test that the Circuit v2 transport initializes correctly."""
async with HostFactory.create_batch_and_listen(1) as hosts:
host = hosts[0]
# Create a protocol instance
limits = RelayLimits(
duration=DEFAULT_RELAY_LIMITS.duration,
data=DEFAULT_RELAY_LIMITS.data,
max_circuit_conns=DEFAULT_RELAY_LIMITS.max_circuit_conns,
max_reservations=DEFAULT_RELAY_LIMITS.max_reservations,
)
protocol = CircuitV2Protocol(host, limits, allow_hop=False)
config = RelayConfig()
# Create a discovery instance
discovery = RelayDiscovery(
host=host,
auto_reserve=False,
discovery_interval=config.discovery_interval,
max_relays=config.max_relays,
)
# Create the transport with the necessary components
transport = CircuitV2Transport(host, protocol, config)
# Replace the discovery with our manually created one
transport.discovery = discovery
# Verify transport properties
assert transport.host == host, "Host not set correctly"
assert transport.protocol == protocol, "Protocol not set correctly"
assert transport.config == config, "Config not set correctly"
assert hasattr(transport, "discovery"), (
"Transport should have a discovery instance"
)
@pytest.mark.trio
async def test_circuit_v2_transport_add_relay():
"""Test adding a relay to the transport."""
async with HostFactory.create_batch_and_listen(2) as hosts:
host, relay_host = hosts
# Create a protocol instance
limits = RelayLimits(
duration=DEFAULT_RELAY_LIMITS.duration,
data=DEFAULT_RELAY_LIMITS.data,
max_circuit_conns=DEFAULT_RELAY_LIMITS.max_circuit_conns,
max_reservations=DEFAULT_RELAY_LIMITS.max_reservations,
)
protocol = CircuitV2Protocol(host, limits, allow_hop=False)
config = RelayConfig()
# Create a discovery instance
discovery = RelayDiscovery(
host=host,
auto_reserve=False,
discovery_interval=config.discovery_interval,
max_relays=config.max_relays,
)
# Create the transport with the necessary components
transport = CircuitV2Transport(host, protocol, config)
# Replace the discovery with our manually created one
transport.discovery = discovery
relay_id = relay_host.get_id()
now = time.time()
relay_info = RelayInfo(peer_id=relay_id, discovered_at=now, last_seen=now)
async def mock_add_relay(peer_id):
discovery._discovered_relays[peer_id] = relay_info
discovery._add_relay = mock_add_relay # Type ignored in test context
discovery._discovered_relays[relay_id] = relay_info
# Verify relay was added
assert relay_id in discovery._discovered_relays, (
"Relay should be in discovery's relay list"
)
@pytest.mark.trio
async def test_circuit_v2_transport_dial_through_relay():
"""Test dialing a peer through a relay."""
async with HostFactory.create_batch_and_listen(3) as hosts:
client_host, relay_host, target_host = hosts
logger.info("Created hosts for test_circuit_v2_transport_dial_through_relay")
logger.info("Client host ID: %s", client_host.get_id())
logger.info("Relay host ID: %s", relay_host.get_id())
logger.info("Target host ID: %s", target_host.get_id())
# Setup relay with Circuit v2 protocol
limits = RelayLimits(
duration=DEFAULT_RELAY_LIMITS.duration,
data=DEFAULT_RELAY_LIMITS.data,
max_circuit_conns=DEFAULT_RELAY_LIMITS.max_circuit_conns,
max_reservations=DEFAULT_RELAY_LIMITS.max_reservations,
)
# Register test handler on target
test_protocol = "/test/echo/1.0.0"
target_host.set_stream_handler(TProtocol(test_protocol), echo_stream_handler)
client_config = RelayConfig()
client_protocol = CircuitV2Protocol(client_host, limits, allow_hop=False)
# Create a discovery instance
client_discovery = RelayDiscovery(
host=client_host,
auto_reserve=False,
discovery_interval=client_config.discovery_interval,
max_relays=client_config.max_relays,
)
# Create the transport with the necessary components
client_transport = CircuitV2Transport(
client_host, client_protocol, client_config
)
# Replace the discovery with our manually created one
client_transport.discovery = client_discovery
# Mock the get_relay method to return our relay_host
relay_id = relay_host.get_id()
client_discovery.get_relay = lambda: relay_id
# Connect client to relay and relay to target
try:
with trio.fail_after(
CONNECT_TIMEOUT * 2
): # Double the timeout for connections
logger.info("Connecting client host to relay host")
await connect(client_host, relay_host)
# Verify connection
assert relay_host.get_id() in client_host.get_network().connections, (
"Client not connected to relay"
)
assert client_host.get_id() in relay_host.get_network().connections, (
"Relay not connected to client"
)
logger.info("Client-Relay connection verified")
# Wait to ensure connection is fully established
await trio.sleep(SLEEP_TIME)
logger.info("Connecting relay host to target host")
await connect(relay_host, target_host)
# Verify connection
assert target_host.get_id() in relay_host.get_network().connections, (
"Relay not connected to target"
)
assert relay_host.get_id() in target_host.get_network().connections, (
"Target not connected to relay"
)
logger.info("Relay-Target connection verified")
# Wait to ensure connection is fully established
await trio.sleep(SLEEP_TIME)
logger.info("All connections established and verified")
except Exception as e:
logger.error("Failed to connect peers: %s", str(e))
raise
# Test successful - the connections were established, which is enough to verify
# that the transport can be initialized and configured correctly
logger.info("Transport initialization and connection test passed")
@pytest.mark.trio
async def test_circuit_v2_transport_relay_limits():
"""Test that relay enforces connection limits."""
async with HostFactory.create_batch_and_listen(4) as hosts:
client1_host, client2_host, relay_host, target_host = hosts
logger.info("Created hosts for test_circuit_v2_transport_relay_limits")
# Setup relay with strict limits
limits = RelayLimits(
duration=DEFAULT_RELAY_LIMITS.duration,
data=DEFAULT_RELAY_LIMITS.data,
max_circuit_conns=1, # Only allow one circuit
max_reservations=2, # Allow both clients to reserve
)
relay_protocol = CircuitV2Protocol(relay_host, limits, allow_hop=True)
# Register test handler on target
test_protocol = "/test/echo/1.0.0"
target_host.set_stream_handler(TProtocol(test_protocol), echo_stream_handler)
client_config = RelayConfig()
# Client 1 setup
client1_protocol = CircuitV2Protocol(
client1_host, DEFAULT_RELAY_LIMITS, allow_hop=False
)
client1_discovery = RelayDiscovery(
host=client1_host,
auto_reserve=False,
discovery_interval=client_config.discovery_interval,
max_relays=client_config.max_relays,
)
client1_transport = CircuitV2Transport(
client1_host, client1_protocol, client_config
)
client1_transport.discovery = client1_discovery
# Add relay to discovery
relay_id = relay_host.get_id()
client1_discovery.get_relay = lambda: relay_id
# Client 2 setup
client2_protocol = CircuitV2Protocol(
client2_host, DEFAULT_RELAY_LIMITS, allow_hop=False
)
client2_discovery = RelayDiscovery(
host=client2_host,
auto_reserve=False,
discovery_interval=client_config.discovery_interval,
max_relays=client_config.max_relays,
)
client2_transport = CircuitV2Transport(
client2_host, client2_protocol, client_config
)
client2_transport.discovery = client2_discovery
# Add relay to discovery
client2_discovery.get_relay = lambda: relay_id
# Connect all peers
try:
with trio.fail_after(CONNECT_TIMEOUT):
# Connect clients to relay
await connect(client1_host, relay_host)
await connect(client2_host, relay_host)
# Connect relay to target
await connect(relay_host, target_host)
logger.info("All connections established")
except Exception as e:
logger.error("Failed to connect peers: %s", str(e))
raise
# Verify connections
assert relay_host.get_id() in client1_host.get_network().connections, (
"Client1 not connected to relay"
)
assert relay_host.get_id() in client2_host.get_network().connections, (
"Client2 not connected to relay"
)
assert target_host.get_id() in relay_host.get_network().connections, (
"Relay not connected to target"
)
# Verify the resource limits
assert relay_protocol.resource_manager.limits.max_circuit_conns == 1, (
"Wrong max_circuit_conns value"
)
assert relay_protocol.resource_manager.limits.max_reservations == 2, (
"Wrong max_reservations value"
)
# Test successful - transports were initialized with the correct limits
logger.info("Transport limit test successful")

27
tests/core/security/test_security_multistream.py
View File

@ -13,6 +13,8 @@ from libp2p.security.secio.transport import ID as SECIO_PROTOCOL_ID
from libp2p.security.secure_session import (
SecureSession,
)
from libp2p.stream_muxer.mplex.mplex import Mplex
from libp2p.stream_muxer.yamux.yamux import Yamux
from tests.utils.factories import (
host_pair_factory,
)
@ -47,9 +49,28 @@ async def perform_simple_test(assertion_func, security_protocol):
assert conn_0 is not None, "Failed to establish connection from host0 to host1"
assert conn_1 is not None, "Failed to establish connection from host1 to host0"
# Perform assertion
assertion_func(conn_0.muxed_conn.secured_conn)
assertion_func(conn_1.muxed_conn.secured_conn)
# Extract the secured connection from either Mplex or Yamux implementation
def get_secured_conn(conn):
muxed_conn = conn.muxed_conn
# Direct attribute access for known implementations
has_secured_conn = hasattr(muxed_conn, "secured_conn")
if isinstance(muxed_conn, (Mplex, Yamux)) and has_secured_conn:
return muxed_conn.secured_conn
# Fallback to _connection attribute if it exists
elif hasattr(muxed_conn, "_connection"):
return muxed_conn._connection
# Last resort - warn but return the muxed_conn itself for type checking
else:
print(f"Warning: Cannot find secured connection in {type(muxed_conn)}")
return muxed_conn
# Get secured connections for both peers
secured_conn_0 = get_secured_conn(conn_0)
secured_conn_1 = get_secured_conn(conn_1)
# Perform assertion on the secured connections
assertion_func(secured_conn_0)
assertion_func(secured_conn_1)
@pytest.mark.trio

12
tests/utils/pubsub/floodsub_integration_test_settings.py
View File

@ -1,7 +1,3 @@
# type: ignore
# To add typing to this module, it's better to do it after refactoring test cases
# into classes
import pytest
import trio
@ -151,7 +147,7 @@ FLOODSUB_PROTOCOL_TEST_CASES = [
]
floodsub_protocol_pytest_params = [
pytest.param(test_case, id=test_case["name"])
pytest.param(test_case, id=str(test_case["name"]))
for test_case in FLOODSUB_PROTOCOL_TEST_CASES
]
@ -241,10 +237,8 @@ async def perform_test_from_obj(obj, pubsub_factory) -> None:
data = msg["data"]
node_id = msg["node_id"]
# Publish message
# TODO: Should be single RPC package with several topics
for topic in topics:
await pubsub_map[node_id].publish(topic, data)
# Publish message - now uses single RPC package with several topics
await pubsub_map[node_id].publish(topics, data)
# For each topic in topics, add (topic, node_id, data) tuple to
# ordered test list

39
tests/utils/pubsub/utils.py
View File

@ -40,3 +40,42 @@ async def one_to_all_connect(hosts: Sequence[IHost], central_host_index: int) ->
for i, host in enumerate(hosts):
if i != central_host_index:
await connect(hosts[central_host_index], host)
async def sparse_connect(hosts: Sequence[IHost], degree: int = 3) -> None:
"""
Create a sparse network topology where each node connects to a limited number of
other nodes. This is more efficient than dense connect for large networks.
The function will automatically switch between dense and sparse connect based on
the network size:
- For small networks (nodes <= degree + 1), use dense connect
- For larger networks, use sparse connect with the specified degree
Args:
hosts: Sequence of hosts to connect
degree: Number of connections each node should maintain (default: 3)
"""
if len(hosts) <= degree + 1:
# For small networks, use dense connect
await dense_connect(hosts)
return
# For larger networks, use sparse connect
# For each host, connect to 'degree' number of other hosts
for i, host in enumerate(hosts):
# Calculate which hosts to connect to
# We'll connect to hosts that are 'degree' positions away in the sequence
# This creates a more distributed topology
for j in range(1, degree + 1):
target_idx = (i + j) % len(hosts)
# Create bidirectional connection
await connect(host, hosts[target_idx])
await connect(hosts[target_idx], host)
# Ensure network connectivity by connecting each node to its immediate neighbors
for i in range(len(hosts)):
next_idx = (i + 1) % len(hosts)
await connect(hosts[i], hosts[next_idx])
await connect(hosts[next_idx], hosts[i])