From c276053301c5caa5566a0095e0a3e3540e8103b8 Mon Sep 17 00:00:00 2001 From: Victor Shyba Date: Sat, 12 Mar 2022 02:17:37 -0300 Subject: [PATCH] move server implementation to tracker module --- lbry/torrent/tracker.py | 43 ++++++++++++++++++++++++++++ tests/unit/torrent/test_tracker.py | 46 +----------------------------- 2 files changed, 44 insertions(+), 45 deletions(-) diff --git a/lbry/torrent/tracker.py b/lbry/torrent/tracker.py index 037c5bbb2..d29242790 100644 --- a/lbry/torrent/tracker.py +++ b/lbry/torrent/tracker.py @@ -6,6 +6,7 @@ import asyncio import logging import time from collections import namedtuple +from functools import reduce from typing import Optional from lbry.utils import resolve_host, async_timed_cache, cache_concurrent @@ -220,3 +221,45 @@ class TrackerClient: def subscribe_hash(info_hash: bytes, on_data): TrackerClient.EVENT_CONTROLLER.add(('search', info_hash, on_data)) + + +class UDPTrackerServerProtocol(asyncio.DatagramProtocol): # for testing. Not suitable for production + def __init__(self): + self.transport = None + self.known_conns = set() + self.peers = {} + + def connection_made(self, transport: asyncio.DatagramTransport) -> None: + self.transport = transport + + def add_peer(self, info_hash, ip_address: str, port: int): + self.peers.setdefault(info_hash, []) + self.peers[info_hash].append(encode_peer(ip_address, port)) + + def datagram_received(self, data: bytes, address: (str, int)) -> None: + if len(data) < 16: + return + action = int.from_bytes(data[8:12], "big", signed=False) + if action == 0: + req = decode(ConnectRequest, data) + connection_id = random.getrandbits(32) + self.known_conns.add(connection_id) + return self.transport.sendto(encode(ConnectResponse(0, req.transaction_id, connection_id)), address) + elif action == 1: + req = decode(AnnounceRequest, data) + if req.connection_id not in self.known_conns: + resp = encode(ErrorResponse(3, req.transaction_id, b'Connection ID missmatch.\x00')) + else: + compact_address = encode_peer(address[0], req.port) + if req.event != 3: + self.add_peer(req.info_hash, address[0], req.port) + elif compact_address in self.peers.get(req.info_hash, []): + self.peers[req.info_hash].remove(compact_address) + peers = [decode(CompactIPv4Peer, peer) for peer in self.peers[req.info_hash]] + resp = encode(AnnounceResponse(1, req.transaction_id, 1700, 0, len(peers), peers)) + return self.transport.sendto(resp, address) + + +def encode_peer(ip_address: str, port: int): + compact_ip = reduce(lambda buff, x: buff + bytearray([int(x)]), ip_address.split('.'), bytearray()) + return compact_ip + port.to_bytes(2, "big", signed=False) diff --git a/tests/unit/torrent/test_tracker.py b/tests/unit/torrent/test_tracker.py index 4bb733361..b240ac60e 100644 --- a/tests/unit/torrent/test_tracker.py +++ b/tests/unit/torrent/test_tracker.py @@ -1,52 +1,8 @@ import asyncio import random -from functools import reduce from lbry.testcase import AsyncioTestCase -from lbry.torrent.tracker import encode, decode, CompactIPv4Peer, ConnectRequest, \ - ConnectResponse, AnnounceRequest, ErrorResponse, AnnounceResponse, TrackerClient, subscribe_hash - - -class UDPTrackerServerProtocol(asyncio.DatagramProtocol): # for testing. Not suitable for production - def __init__(self): - self.transport = None - self.known_conns = set() - self.peers = {} - - def connection_made(self, transport: asyncio.DatagramTransport) -> None: - self.transport = transport - - def add_peer(self, info_hash, ip_address: str, port: int): - self.peers.setdefault(info_hash, []) - self.peers[info_hash].append(encode_peer(ip_address, port)) - - def datagram_received(self, data: bytes, address: (str, int)) -> None: - if len(data) < 16: - return - action = int.from_bytes(data[8:12], "big", signed=False) - if action == 0: - req = decode(ConnectRequest, data) - connection_id = random.getrandbits(32) - self.known_conns.add(connection_id) - return self.transport.sendto(encode(ConnectResponse(0, req.transaction_id, connection_id)), address) - elif action == 1: - req = decode(AnnounceRequest, data) - if req.connection_id not in self.known_conns: - resp = encode(ErrorResponse(3, req.transaction_id, b'Connection ID missmatch.\x00')) - else: - compact_address = encode_peer(address[0], req.port) - if req.event != 3: - self.add_peer(req.info_hash, address[0], req.port) - elif compact_address in self.peers.get(req.info_hash, []): - self.peers[req.info_hash].remove(compact_address) - peers = [decode(CompactIPv4Peer, peer) for peer in self.peers[req.info_hash]] - resp = encode(AnnounceResponse(1, req.transaction_id, 1700, 0, len(peers), peers)) - return self.transport.sendto(resp, address) - - -def encode_peer(ip_address: str, port: int): - compact_ip = reduce(lambda buff, x: buff + bytearray([int(x)]), ip_address.split('.'), bytearray()) - return compact_ip + port.to_bytes(2, "big", signed=False) +from lbry.torrent.tracker import CompactIPv4Peer, TrackerClient, subscribe_hash, UDPTrackerServerProtocol, encode_peer class UDPTrackerClientTestCase(AsyncioTestCase):