py-libp2p/libp2p/relay/circuit_v2/resources.py

"""
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