straywild-discovery-server/straywild_discovery/registry.py

481 lines
17 KiB
Python

"""Thread-safe ephemeral room leases for the straywild directory."""
from __future__ import annotations
from dataclasses import dataclass, replace
import hmac
import secrets
import threading
import time
import unicodedata
import uuid
from collections import deque
from collections.abc import Callable, Mapping
from typing import Any
ROOM_NAME_MAX_LENGTH = 48
VERSION_MAX_LENGTH = 32
MAX_ROOM_CAPACITY = 128
DEFAULT_MAX_ROOMS = 4096
DEFAULT_MAX_ROOMS_PER_ADDRESS = 32
JOIN_ATTEMPT_TTL_SECONDS = 12.0
MAX_ACTIVE_JOIN_ATTEMPTS_PER_ROOM = 16
JOIN_RATE_WINDOW_SECONDS = 60.0
MAX_JOIN_ATTEMPTS_PER_ADDRESS_PER_WINDOW = 10
MAX_JOIN_ATTEMPTS_PER_ROOM_PER_WINDOW = 60
MAX_JOIN_ATTEMPTS_GLOBAL_PER_WINDOW = 2048
class RegistryError(Exception):
"""Base class for errors safe to translate into API responses."""
class ValidationError(RegistryError):
pass
class RoomNotFoundError(RegistryError):
pass
class LeaseAuthorizationError(RegistryError):
pass
class RoomLimitError(RegistryError):
pass
@dataclass(frozen=True, slots=True)
class RoomAdvertisement:
room_id: str
room_name: str
address: str
port: int
current_players: int
max_players: int
game_version: str
protocol_version: int
host_kind: str
connection_mode: str
created_at: float
updated_at: float
expires_at: float
verified: bool = False
def public_dict(self) -> dict[str, Any]:
return {
"room_id": self.room_id,
"room_name": self.room_name,
"current_players": self.current_players,
"max_players": self.max_players,
"game_version": self.game_version,
"protocol_version": self.protocol_version,
"host_kind": self.host_kind,
"created_at": _iso_utc(self.created_at),
"updated_at": _iso_utc(self.updated_at),
"expires_at": _iso_utc(self.expires_at),
"verified": self.verified,
"connection_mode": self.connection_mode,
"ip_privacy": (
"relayed" if self.connection_mode == "relay" else "peer_visible"
),
}
def join_route_dict(self) -> dict[str, Any]:
"""Return the transient route disclosed only after an explicit join."""
return {
"transport": "direct",
"address": self.address,
"port": self.port,
"ip_privacy": "peer_visible",
}
@dataclass(slots=True)
class _RoomLease:
advertisement: RoomAdvertisement
token: str
verification_token: str
@dataclass(slots=True)
class _JoinAttempt:
room_id: str
expires_at: float
address: str = ""
port: int = 0
def _iso_utc(timestamp: float) -> str:
from datetime import datetime, timezone
return datetime.fromtimestamp(timestamp, timezone.utc).isoformat().replace("+00:00", "Z")
def _clean_text(value: Any, field_name: str, max_length: int) -> str:
if not isinstance(value, str):
raise ValidationError(f"{field_name} must be a string")
normalized = unicodedata.normalize("NFC", value).strip()
if not normalized:
raise ValidationError(f"{field_name} cannot be empty")
if len(normalized) > max_length:
raise ValidationError(f"{field_name} cannot exceed {max_length} characters")
if any(unicodedata.category(character).startswith("C") for character in normalized):
raise ValidationError(f"{field_name} contains unsupported control characters")
return normalized
def _clean_int(value: Any, field_name: str, minimum: int, maximum: int) -> int:
if isinstance(value, bool) or not isinstance(value, int):
raise ValidationError(f"{field_name} must be an integer")
if value < minimum or value > maximum:
raise ValidationError(f"{field_name} must be between {minimum} and {maximum}")
return value
class RoomRegistry:
"""Stores short-lived room advertisements; no state survives a restart."""
def __init__(
self,
ttl_seconds: float = 45.0,
max_rooms: int = DEFAULT_MAX_ROOMS,
max_rooms_per_address: int = DEFAULT_MAX_ROOMS_PER_ADDRESS,
relay_available: bool = False,
clock: Callable[[], float] = time.time,
) -> None:
if ttl_seconds <= 0.0:
raise ValueError("ttl_seconds must be positive")
if max_rooms <= 0 or max_rooms_per_address <= 0:
raise ValueError("room limits must be positive")
self._ttl_seconds = float(ttl_seconds)
self._max_rooms = int(max_rooms)
self._max_rooms_per_address = int(max_rooms_per_address)
self._relay_available = bool(relay_available)
self._clock = clock
self._rooms: dict[str, _RoomLease] = {}
self._join_attempts: dict[str, _JoinAttempt] = {}
self._join_attempts_by_address: dict[str, deque[float]] = {}
self._join_attempts_by_room: dict[str, deque[float]] = {}
self._global_join_attempts: deque[float] = deque()
self._lock = threading.RLock()
@property
def ttl_seconds(self) -> float:
return self._ttl_seconds
def create(
self, address: str, payload: Mapping[str, Any]
) -> tuple[RoomAdvertisement, str, str]:
now = self._clock()
room_id = str(uuid.uuid4())
token = secrets.token_urlsafe(32)
verification_token = secrets.token_urlsafe(32)
advertisement = self._build_advertisement(room_id, address, payload, now, now)
with self._lock:
self._purge_locked(now)
if len(self._rooms) >= self._max_rooms:
raise RoomLimitError("the public room directory is at capacity")
address_room_count = sum(
1
for lease in self._rooms.values()
if lease.advertisement.address == address
)
if address_room_count >= self._max_rooms_per_address:
raise RoomLimitError("too many active rooms from this address")
self._rooms[room_id] = _RoomLease(
advertisement, token, verification_token
)
return advertisement, token, verification_token
def update(
self,
room_id: str,
token: str,
address: str,
payload: Mapping[str, Any],
) -> RoomAdvertisement:
now = self._clock()
with self._lock:
self._purge_locked(now)
lease = self._authorized_lease_locked(room_id, token)
advertisement = self._build_advertisement(
room_id,
lease.advertisement.address if lease.advertisement.verified else address,
payload,
lease.advertisement.created_at,
now,
endpoint_port=(
lease.advertisement.port if lease.advertisement.verified else None
),
verified=lease.advertisement.verified,
)
lease.advertisement = advertisement
return advertisement
def delete(self, room_id: str, token: str) -> None:
now = self._clock()
with self._lock:
self._purge_locked(now)
self._authorized_lease_locked(room_id, token)
del self._rooms[room_id]
def list_rooms(
self,
*,
game_version: str | None = None,
protocol_version: int | None = None,
) -> list[RoomAdvertisement]:
now = self._clock()
with self._lock:
self._purge_locked(now)
rooms = [
lease.advertisement
for lease in self._rooms.values()
if lease.advertisement.verified
]
if game_version is not None:
rooms = [room for room in rooms if room.game_version == game_version]
if protocol_version is not None:
rooms = [room for room in rooms if room.protocol_version == protocol_version]
return sorted(rooms, key=lambda room: (-room.updated_at, room.room_name.casefold()))
def room_count(self) -> int:
now = self._clock()
with self._lock:
self._purge_locked(now)
return len(self._rooms)
def get_public_room(
self,
room_id: str,
*,
game_version: str | None = None,
protocol_version: int | None = None,
) -> RoomAdvertisement | None:
now = self._clock()
with self._lock:
self._purge_locked(now)
lease = self._rooms.get(room_id)
room = lease.advertisement if lease is not None else None
if room is None or not room.verified:
return None
if game_version is not None and room.game_version != game_version:
return None
if protocol_version is not None and room.protocol_version != protocol_version:
return None
return room
def verify_endpoint(
self,
room_id: str,
verification_token: str,
address: str,
port: int,
) -> RoomAdvertisement:
if not address or port < 1 or port > 65_535:
raise ValidationError("invalid observed UDP endpoint")
now = self._clock()
with self._lock:
self._purge_locked(now)
lease = self._rooms.get(room_id)
if lease is None:
raise RoomNotFoundError("room does not exist or its lease expired")
if not hmac.compare_digest(
lease.verification_token, verification_token
):
raise LeaseAuthorizationError("invalid endpoint verification token")
lease.advertisement = replace(
lease.advertisement,
address=address,
port=port,
updated_at=now,
expires_at=now + self._ttl_seconds,
verified=True,
)
return lease.advertisement
def create_join_attempt(
self, room_id: str, requester_address: str
) -> tuple[str, RoomAdvertisement]:
if not requester_address:
raise ValidationError("invalid join source")
now = self._clock()
with self._lock:
self._purge_locked(now)
lease = self._rooms.get(room_id)
if lease is None or not lease.advertisement.verified:
raise RoomNotFoundError("room is not available")
self._enforce_join_rate_limits_locked(room_id, requester_address, now)
active_count = sum(
1
for attempt in self._join_attempts.values()
if attempt.room_id == room_id
)
if active_count >= MAX_ACTIVE_JOIN_ATTEMPTS_PER_ROOM:
raise RoomLimitError("too many pending joins for this room")
token = secrets.token_urlsafe(32)
self._join_attempts[token] = _JoinAttempt(
room_id=room_id,
expires_at=now + JOIN_ATTEMPT_TTL_SECONDS,
)
self._join_attempts_by_address.setdefault(
requester_address, deque()
).append(now)
self._join_attempts_by_room.setdefault(room_id, deque()).append(now)
self._global_join_attempts.append(now)
return token, lease.advertisement
def register_join_endpoint(
self, token: str, address: str, port: int
) -> None:
if not address or port < 1 or port > 65_535:
raise ValidationError("invalid observed join endpoint")
now = self._clock()
with self._lock:
self._purge_locked(now)
attempt = self._join_attempts.get(token)
if attempt is None:
raise RoomNotFoundError("join attempt does not exist or expired")
attempt.address = address
attempt.port = port
def cancel_join_attempt(self, token: str) -> None:
with self._lock:
self._join_attempts.pop(token, None)
def consume_join_endpoints(
self, room_id: str, lease_token: str
) -> list[dict[str, Any]]:
now = self._clock()
with self._lock:
self._purge_locked(now)
self._authorized_lease_locked(room_id, lease_token)
consumed_tokens = [
token
for token, attempt in self._join_attempts.items()
if attempt.room_id == room_id and attempt.address and attempt.port > 0
]
endpoints = [
{
"address": self._join_attempts[token].address,
"port": self._join_attempts[token].port,
}
for token in consumed_tokens
]
for token in consumed_tokens:
del self._join_attempts[token]
return endpoints
def _build_advertisement(
self,
room_id: str,
address: str,
payload: Mapping[str, Any],
created_at: float,
now: float,
endpoint_port: int | None = None,
verified: bool = False,
) -> RoomAdvertisement:
room_name = _clean_text(payload.get("room_name"), "room_name", ROOM_NAME_MAX_LENGTH)
supplied_port = _clean_int(payload.get("port"), "port", 1, 65_535)
port = endpoint_port if endpoint_port is not None else supplied_port
current_players = _clean_int(
payload.get("current_players"), "current_players", 0, MAX_ROOM_CAPACITY
)
max_players = _clean_int(
payload.get("max_players"), "max_players", 1, MAX_ROOM_CAPACITY
)
if current_players > max_players:
raise ValidationError("current_players cannot exceed max_players")
game_version = _clean_text(
payload.get("game_version"), "game_version", VERSION_MAX_LENGTH
)
protocol_version = _clean_int(
payload.get("protocol_version"), "protocol_version", 1, 2_147_483_647
)
host_kind = payload.get("host_kind", "player")
if host_kind not in {"player", "dedicated"}:
raise ValidationError("host_kind must be player or dedicated")
connection_mode = payload.get("connection_mode", "direct")
if connection_mode not in {"direct", "relay"}:
raise ValidationError("connection_mode must be direct or relay")
if connection_mode == "relay" and host_kind == "dedicated":
raise ValidationError("dedicated servers cannot use the privacy relay")
if connection_mode == "relay" and not self._relay_available:
raise ValidationError("the privacy relay is not available")
return RoomAdvertisement(
room_id=room_id,
room_name=room_name,
address=address,
port=port,
current_players=current_players,
max_players=max_players,
game_version=game_version,
protocol_version=protocol_version,
host_kind=host_kind,
connection_mode=connection_mode,
created_at=created_at,
updated_at=now,
expires_at=now + self._ttl_seconds,
verified=verified,
)
def _authorized_lease_locked(self, room_id: str, token: str) -> _RoomLease:
lease = self._rooms.get(room_id)
if lease is None:
raise RoomNotFoundError("room does not exist or its lease expired")
if not token or not hmac.compare_digest(lease.token, token):
raise LeaseAuthorizationError("invalid room lease token")
return lease
def _enforce_join_rate_limits_locked(
self, room_id: str, requester_address: str, now: float
) -> None:
cutoff = now - JOIN_RATE_WINDOW_SECONDS
self._prune_history(self._global_join_attempts, cutoff)
address_history = self._join_attempts_by_address.setdefault(
requester_address, deque()
)
room_history = self._join_attempts_by_room.setdefault(room_id, deque())
self._prune_history(address_history, cutoff)
self._prune_history(room_history, cutoff)
if len(address_history) >= MAX_JOIN_ATTEMPTS_PER_ADDRESS_PER_WINDOW:
raise RoomLimitError("too many join requests; try again shortly")
if len(room_history) >= MAX_JOIN_ATTEMPTS_PER_ROOM_PER_WINDOW:
raise RoomLimitError("this room is receiving too many join requests")
if len(self._global_join_attempts) >= MAX_JOIN_ATTEMPTS_GLOBAL_PER_WINDOW:
raise RoomLimitError("the join service is busy; try again shortly")
@staticmethod
def _prune_history(history: deque[float], cutoff: float) -> None:
while history and history[0] <= cutoff:
history.popleft()
def _purge_locked(self, now: float) -> None:
expired = [
room_id
for room_id, lease in self._rooms.items()
if lease.advertisement.expires_at <= now
]
for room_id in expired:
del self._rooms[room_id]
expired_attempts = [
token
for token, attempt in self._join_attempts.items()
if attempt.expires_at <= now or attempt.room_id not in self._rooms
]
for token in expired_attempts:
del self._join_attempts[token]
cutoff = now - JOIN_RATE_WINDOW_SECONDS
self._prune_history(self._global_join_attempts, cutoff)
for histories in [
self._join_attempts_by_address,
self._join_attempts_by_room,
]:
for key in list(histories):
self._prune_history(histories[key], cutoff)
if not histories[key]:
del histories[key]