feat: add UDP NAT rendezvous

This commit is contained in:
Alexander Sellite 2026-08-10 19:49:47 -04:00
parent 6a7c7a676d
commit 804e19c51a
7 changed files with 370 additions and 27 deletions

View file

@ -2,7 +2,7 @@
from __future__ import annotations
from dataclasses import dataclass
from dataclasses import dataclass, replace
import hmac
import secrets
import threading
@ -18,6 +18,8 @@ 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_JOIN_ATTEMPTS_PER_ROOM = 64
class RegistryError(Exception):
@ -53,6 +55,7 @@ class RoomAdvertisement:
created_at: float
updated_at: float
expires_at: float
verified: bool = False
def public_dict(self) -> dict[str, Any]:
return {
@ -67,6 +70,7 @@ class RoomAdvertisement:
"created_at": _iso_utc(self.created_at),
"updated_at": _iso_utc(self.updated_at),
"expires_at": _iso_utc(self.expires_at),
"verified": self.verified,
}
@ -74,6 +78,15 @@ class RoomAdvertisement:
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:
@ -122,16 +135,20 @@ class RoomRegistry:
self._max_rooms_per_address = int(max_rooms_per_address)
self._clock = clock
self._rooms: dict[str, _RoomLease] = {}
self._join_attempts: dict[str, _JoinAttempt] = {}
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]:
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)
@ -144,8 +161,10 @@ class RoomRegistry:
)
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)
return advertisement, token
self._rooms[room_id] = _RoomLease(
advertisement, token, verification_token
)
return advertisement, token, verification_token
def update(
self,
@ -160,10 +179,14 @@ class RoomRegistry:
lease = self._authorized_lease_locked(room_id, token)
advertisement = self._build_advertisement(
room_id,
address,
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
@ -184,7 +207,11 @@ class RoomRegistry:
now = self._clock()
with self._lock:
self._purge_locked(now)
rooms = [lease.advertisement for lease in self._rooms.values()]
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:
@ -197,6 +224,93 @@ class RoomRegistry:
self._purge_locked(now)
return len(self._rooms)
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) -> str:
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")
active_count = sum(
1
for attempt in self._join_attempts.values()
if attempt.room_id == room_id
)
if active_count >= MAX_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,
)
return token
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 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,
@ -204,9 +318,12 @@ class RoomRegistry:
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)
port = _clean_int(payload.get("port"), "port", 1, 65_535)
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
)
@ -233,6 +350,7 @@ class RoomRegistry:
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:
@ -251,3 +369,10 @@ class RoomRegistry:
]
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]