79 lines
2.6 KiB
Python
79 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
import socket
|
|
import unittest
|
|
|
|
from straywild_discovery.relay import (
|
|
RELAY_AUTH_PREFIX,
|
|
RelayConfig,
|
|
RelayService,
|
|
)
|
|
|
|
|
|
class RelayServiceTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.host = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
self.host.bind(("127.0.0.1", 0))
|
|
self.host.settimeout(0.35)
|
|
self.client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
self.client.bind(("127.0.0.1", 0))
|
|
self.client.settimeout(1.0)
|
|
self.relay = RelayService(RelayConfig(
|
|
bind_host="127.0.0.1",
|
|
public_host="127.0.0.1",
|
|
port_start=0,
|
|
port_end=0,
|
|
max_allocations=4,
|
|
))
|
|
self.relay.start()
|
|
|
|
def tearDown(self) -> None:
|
|
self.relay.close()
|
|
self.client.close()
|
|
self.host.close()
|
|
|
|
def test_authenticated_allocation_masks_both_peer_endpoints(self) -> None:
|
|
token = "join-capability"
|
|
route = self.relay.allocate(
|
|
token,
|
|
"127.0.0.1",
|
|
int(self.host.getsockname()[1]),
|
|
)
|
|
relay_endpoint = (route.host, route.port)
|
|
|
|
self.client.sendto(b"not authenticated", relay_endpoint)
|
|
with self.assertRaises(TimeoutError):
|
|
self.host.recvfrom(4096)
|
|
|
|
self.client.sendto(RELAY_AUTH_PREFIX + token.encode(), relay_endpoint)
|
|
self.client.sendto(b"client gameplay", relay_endpoint)
|
|
packet, relay_as_seen_by_host = self.host.recvfrom(4096)
|
|
self.assertEqual(packet, b"client gameplay")
|
|
self.assertEqual(relay_as_seen_by_host[1], route.port)
|
|
self.assertNotEqual(
|
|
relay_as_seen_by_host[1], int(self.client.getsockname()[1])
|
|
)
|
|
|
|
self.host.sendto(b"host gameplay", relay_as_seen_by_host)
|
|
packet, relay_as_seen_by_client = self.client.recvfrom(4096)
|
|
self.assertEqual(packet, b"host gameplay")
|
|
self.assertEqual(relay_as_seen_by_client, relay_endpoint)
|
|
self.assertNotEqual(
|
|
relay_as_seen_by_client[1], int(self.host.getsockname()[1])
|
|
)
|
|
|
|
def test_wrong_capability_cannot_claim_allocation(self) -> None:
|
|
route = self.relay.allocate(
|
|
"right-token",
|
|
"127.0.0.1",
|
|
int(self.host.getsockname()[1]),
|
|
)
|
|
relay_endpoint = (route.host, route.port)
|
|
self.client.sendto(RELAY_AUTH_PREFIX + b"wrong-token", relay_endpoint)
|
|
self.client.sendto(b"blocked gameplay", relay_endpoint)
|
|
with self.assertRaises(TimeoutError):
|
|
self.host.recvfrom(4096)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|