312 lines
11 KiB
Python
312 lines
11 KiB
Python
"""Bounded allocation-based UDP relay for ENet gameplay traffic."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import selectors
|
|
import secrets
|
|
import socket
|
|
import threading
|
|
import time
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
|
|
RELAY_AUTH_PREFIX = b"straywild_RELAY_AUTH_V1 "
|
|
|
|
|
|
class RelayError(Exception):
|
|
"""A relay allocation could not be created safely."""
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class RelayConfig:
|
|
bind_host: str = "0.0.0.0"
|
|
public_host: str = "127.0.0.1"
|
|
port_start: int = 20_000
|
|
port_end: int = 22_047
|
|
max_allocations: int = 2_048
|
|
authentication_ttl_seconds: float = 12.0
|
|
idle_ttl_seconds: float = 120.0
|
|
max_packet_bytes: int = 4_096
|
|
max_packets_per_second: float = 512.0
|
|
packet_burst: float = 1_024.0
|
|
max_bytes_per_second: float = 1_500_000.0
|
|
byte_burst: float = 3_000_000.0
|
|
|
|
def validate(self) -> None:
|
|
if not self.bind_host or not self.public_host:
|
|
raise ValueError("relay bind and public hosts are required")
|
|
ephemeral_ports = self.port_start == 0 and self.port_end == 0
|
|
if not ephemeral_ports and (
|
|
self.port_start < 1 or self.port_end > 65_535
|
|
):
|
|
raise ValueError("relay ports must be between 1 and 65535")
|
|
if self.port_end < self.port_start:
|
|
raise ValueError("relay port range is invalid")
|
|
if self.max_allocations < 1:
|
|
raise ValueError("relay allocation limit must be positive")
|
|
if (
|
|
not ephemeral_ports
|
|
and self.max_allocations > self.port_end - self.port_start + 1
|
|
):
|
|
raise ValueError("relay allocation limit exceeds the port range")
|
|
if self.authentication_ttl_seconds <= 0 or self.idle_ttl_seconds <= 0:
|
|
raise ValueError("relay timeouts must be positive")
|
|
if self.max_packet_bytes < 512:
|
|
raise ValueError("relay packet limit is too small")
|
|
if min(
|
|
self.max_packets_per_second,
|
|
self.packet_burst,
|
|
self.max_bytes_per_second,
|
|
self.byte_burst,
|
|
) <= 0:
|
|
raise ValueError("relay traffic limits must be positive")
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class RelayRoute:
|
|
host: str
|
|
port: int
|
|
|
|
def as_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"transport": "relay",
|
|
"address": self.host,
|
|
"port": self.port,
|
|
"ip_privacy": "relayed",
|
|
}
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _Allocation:
|
|
udp_socket: socket.socket
|
|
port: int
|
|
authorization_packet: bytes
|
|
host_endpoint: tuple[str, int]
|
|
created_at: float
|
|
last_activity_at: float
|
|
client_endpoint: tuple[str, int] | None = None
|
|
packet_tokens: float = 0.0
|
|
byte_tokens: float = 0.0
|
|
token_updated_at: float = 0.0
|
|
|
|
|
|
class RelayService:
|
|
"""Runs many isolated relay sockets on one bounded selector thread."""
|
|
|
|
def __init__(
|
|
self,
|
|
config: RelayConfig,
|
|
clock: Callable[[], float] = time.monotonic,
|
|
) -> None:
|
|
config.validate()
|
|
self.config = config
|
|
self._clock = clock
|
|
self._selector = selectors.DefaultSelector()
|
|
self._allocations: dict[int, _Allocation] = {}
|
|
self._lock = threading.RLock()
|
|
self._stop = threading.Event()
|
|
self._thread: threading.Thread | None = None
|
|
|
|
def start(self) -> None:
|
|
with self._lock:
|
|
if self._thread is not None:
|
|
return
|
|
self._stop.clear()
|
|
self._thread = threading.Thread(
|
|
target=self._run,
|
|
name="discovery-relay",
|
|
daemon=True,
|
|
)
|
|
self._thread.start()
|
|
|
|
def close(self) -> None:
|
|
self._stop.set()
|
|
thread = self._thread
|
|
if thread is not None:
|
|
thread.join(timeout=2.0)
|
|
with self._lock:
|
|
for allocation in list(self._allocations.values()):
|
|
self._close_allocation_locked(allocation)
|
|
self._thread = None
|
|
self._selector.close()
|
|
|
|
def allocation_count(self) -> int:
|
|
with self._lock:
|
|
return len(self._allocations)
|
|
|
|
def allocate(
|
|
self,
|
|
authorization_token: str,
|
|
host_address: str,
|
|
host_port: int,
|
|
) -> RelayRoute:
|
|
if (
|
|
not authorization_token
|
|
or not host_address
|
|
or host_port < 1
|
|
or host_port > 65_535
|
|
):
|
|
raise RelayError("invalid relay allocation request")
|
|
authorization_packet = RELAY_AUTH_PREFIX + authorization_token.encode("ascii")
|
|
now = self._clock()
|
|
with self._lock:
|
|
self._purge_expired_locked(now)
|
|
if len(self._allocations) >= self.config.max_allocations:
|
|
raise RelayError("relay capacity reached")
|
|
try:
|
|
udp_socket = self._bind_available_socket_locked(host_address)
|
|
except OSError as error:
|
|
raise RelayError("could not bind a relay socket") from error
|
|
port = int(udp_socket.getsockname()[1])
|
|
allocation = _Allocation(
|
|
udp_socket=udp_socket,
|
|
port=port,
|
|
authorization_packet=authorization_packet,
|
|
host_endpoint=(host_address, host_port),
|
|
created_at=now,
|
|
last_activity_at=now,
|
|
packet_tokens=self.config.packet_burst,
|
|
byte_tokens=self.config.byte_burst,
|
|
token_updated_at=now,
|
|
)
|
|
self._allocations[port] = allocation
|
|
try:
|
|
self._selector.register(
|
|
udp_socket, selectors.EVENT_READ, allocation
|
|
)
|
|
except (OSError, ValueError) as error:
|
|
self._allocations.pop(port, None)
|
|
udp_socket.close()
|
|
raise RelayError("could not activate a relay socket") from error
|
|
return RelayRoute(self.config.public_host, port)
|
|
|
|
def _bind_available_socket_locked(self, host_address: str) -> socket.socket:
|
|
family = socket.AF_INET6 if ":" in host_address else socket.AF_INET
|
|
bind_host = self.config.bind_host
|
|
if family == socket.AF_INET6 and bind_host == "0.0.0.0":
|
|
bind_host = "::"
|
|
if family == socket.AF_INET and bind_host == "::":
|
|
bind_host = "0.0.0.0"
|
|
if self.config.port_start == 0 and self.config.port_end == 0:
|
|
udp_socket = socket.socket(family, socket.SOCK_DGRAM)
|
|
udp_socket.setblocking(False)
|
|
udp_socket.bind((bind_host, 0))
|
|
return udp_socket
|
|
range_size = self.config.port_end - self.config.port_start + 1
|
|
start_offset = secrets.randbelow(range_size)
|
|
for offset in range(range_size):
|
|
port = self.config.port_start + ((start_offset + offset) % range_size)
|
|
if port in self._allocations:
|
|
continue
|
|
udp_socket = socket.socket(family, socket.SOCK_DGRAM)
|
|
udp_socket.setblocking(False)
|
|
try:
|
|
udp_socket.bind((bind_host, port))
|
|
return udp_socket
|
|
except OSError:
|
|
udp_socket.close()
|
|
raise RelayError("no relay port is available")
|
|
|
|
def _run(self) -> None:
|
|
while not self._stop.is_set():
|
|
with self._lock:
|
|
events = self._selector.select(timeout=0.05)
|
|
for key, _mask in events:
|
|
allocation = key.data
|
|
if isinstance(allocation, _Allocation):
|
|
self._receive_locked(allocation)
|
|
self._purge_expired_locked(self._clock())
|
|
|
|
def _receive_locked(self, allocation: _Allocation) -> None:
|
|
try:
|
|
packet, source = allocation.udp_socket.recvfrom(
|
|
self.config.max_packet_bytes + 1
|
|
)
|
|
except (BlockingIOError, OSError):
|
|
return
|
|
if not packet or len(packet) > self.config.max_packet_bytes:
|
|
return
|
|
source_endpoint = (str(source[0]), int(source[1]))
|
|
now = self._clock()
|
|
if source_endpoint == allocation.host_endpoint:
|
|
if (
|
|
allocation.client_endpoint is not None
|
|
and self._consume_budget_locked(allocation, len(packet), now)
|
|
):
|
|
self._send_locked(allocation, packet, allocation.client_endpoint)
|
|
allocation.last_activity_at = now
|
|
return
|
|
if packet.startswith(RELAY_AUTH_PREFIX):
|
|
if (
|
|
packet == allocation.authorization_packet
|
|
and now - allocation.created_at
|
|
<= self.config.authentication_ttl_seconds
|
|
and (
|
|
allocation.client_endpoint is None
|
|
or allocation.client_endpoint == source_endpoint
|
|
)
|
|
):
|
|
allocation.client_endpoint = source_endpoint
|
|
allocation.authorization_packet = b""
|
|
allocation.last_activity_at = now
|
|
return
|
|
if allocation.client_endpoint != source_endpoint:
|
|
return
|
|
if not self._consume_budget_locked(allocation, len(packet), now):
|
|
return
|
|
self._send_locked(allocation, packet, allocation.host_endpoint)
|
|
allocation.last_activity_at = now
|
|
|
|
@staticmethod
|
|
def _send_locked(
|
|
allocation: _Allocation,
|
|
packet: bytes,
|
|
destination: tuple[str, int],
|
|
) -> None:
|
|
try:
|
|
allocation.udp_socket.sendto(packet, destination)
|
|
except OSError:
|
|
return
|
|
|
|
def _consume_budget_locked(
|
|
self, allocation: _Allocation, packet_bytes: int, now: float
|
|
) -> bool:
|
|
elapsed = max(0.0, now - allocation.token_updated_at)
|
|
allocation.packet_tokens = min(
|
|
self.config.packet_burst,
|
|
allocation.packet_tokens + elapsed * self.config.max_packets_per_second,
|
|
)
|
|
allocation.byte_tokens = min(
|
|
self.config.byte_burst,
|
|
allocation.byte_tokens + elapsed * self.config.max_bytes_per_second,
|
|
)
|
|
allocation.token_updated_at = now
|
|
if allocation.packet_tokens < 1.0 or allocation.byte_tokens < packet_bytes:
|
|
return False
|
|
allocation.packet_tokens -= 1.0
|
|
allocation.byte_tokens -= packet_bytes
|
|
return True
|
|
|
|
def _purge_expired_locked(self, now: float) -> None:
|
|
expired = [
|
|
allocation
|
|
for allocation in self._allocations.values()
|
|
if (
|
|
allocation.client_endpoint is None
|
|
and now - allocation.created_at
|
|
> self.config.authentication_ttl_seconds
|
|
)
|
|
or now - allocation.last_activity_at > self.config.idle_ttl_seconds
|
|
]
|
|
for allocation in expired:
|
|
self._close_allocation_locked(allocation)
|
|
|
|
def _close_allocation_locked(self, allocation: _Allocation) -> None:
|
|
self._allocations.pop(allocation.port, None)
|
|
try:
|
|
self._selector.unregister(allocation.udp_socket)
|
|
except (KeyError, ValueError):
|
|
pass
|
|
allocation.udp_socket.close()
|