mirror of
https://github.com/varun-r-mallya/py-libp2p.git
synced 2026-02-12 16:10:57 +00:00
Feat/587-circuit-relay (#611)
* feat: implemented setup of circuit relay and test cases * chore: remove test files to be rewritten * added 1 test suite for protocol * added 1 test suite for discovery * fixed protocol timeouts and message types to handle reservations and stream operations. * Resolved merge conflict in libp2p/tools/utils.py by combining timeout approach with retry mechanism * fix: linting issues * docs: updated documentation with circuit-relay * chore: added enums, improved typing, security and examples * fix: created proper __init__ file to ensure importability * fix: replace transport_opt with listen_addrs in examples, fixed typing and improved code * fix type checking issues across relay module and test suite * regenerated circuit_pb2 file protobuf version 3 * fixed circuit relay example and moved imports to top in test_security_multistream * chore: moved imports to the top * chore: fixed linting of test_circuit_v2_transport.py --------- Co-authored-by: Manu Sheel Gupta <manusheel.edu@gmail.com>
This commit is contained in:
254
libp2p/relay/circuit_v2/resources.py
Normal file
254
libp2p/relay/circuit_v2/resources.py
Normal 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
|
||||
Reference in New Issue
Block a user