move udp server to test file, add link to BEP15
This commit is contained in:
parent
4a0bf8a702
commit
006391dd26
2 changed files with 37 additions and 33 deletions
|
@ -3,9 +3,9 @@ import struct
|
||||||
import asyncio
|
import asyncio
|
||||||
import logging
|
import logging
|
||||||
from collections import namedtuple
|
from collections import namedtuple
|
||||||
from functools import reduce
|
|
||||||
|
|
||||||
log = logging.getLogger(__name__)
|
log = logging.getLogger(__name__)
|
||||||
|
# see: http://bittorrent.org/beps/bep_0015.html and http://xbtt.sourceforge.net/udp_tracker_protocol.html
|
||||||
ConnectRequest = namedtuple("ConnectRequest", ["connection_id", "action", "transaction_id"])
|
ConnectRequest = namedtuple("ConnectRequest", ["connection_id", "action", "transaction_id"])
|
||||||
ConnectResponse = namedtuple("ConnectResponse", ["action", "transaction_id", "connection_id"])
|
ConnectResponse = namedtuple("ConnectResponse", ["action", "transaction_id", "connection_id"])
|
||||||
AnnounceRequest = namedtuple("AnnounceRequest",
|
AnnounceRequest = namedtuple("AnnounceRequest",
|
||||||
|
@ -108,34 +108,3 @@ class UDPTrackerClientProtocol(asyncio.DatagramProtocol):
|
||||||
|
|
||||||
def connection_lost(self, exc: Exception = None) -> None:
|
def connection_lost(self, exc: Exception = None) -> None:
|
||||||
self.transport = None
|
self.transport = None
|
||||||
|
|
||||||
|
|
||||||
class UDPTrackerServerProtocol(asyncio.DatagramProtocol):
|
|
||||||
def __init__(self):
|
|
||||||
self.transport = None
|
|
||||||
self.known_conns = set()
|
|
||||||
self.peers = {}
|
|
||||||
|
|
||||||
def connection_made(self, transport: asyncio.DatagramTransport) -> None:
|
|
||||||
self.transport = transport
|
|
||||||
|
|
||||||
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:
|
|
||||||
self.peers.setdefault(req.info_hash, [])
|
|
||||||
compact_ip = reduce(lambda buff, x: buff + bytearray([int(x)]), address[0].split('.'), bytearray())
|
|
||||||
self.peers[req.info_hash].append(compact_ip + req.port.to_bytes(2, "big", signed=False))
|
|
||||||
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)
|
|
||||||
|
|
|
@ -1,6 +1,41 @@
|
||||||
|
import asyncio
|
||||||
import random
|
import random
|
||||||
|
from functools import reduce
|
||||||
|
|
||||||
from lbry.testcase import AsyncioTestCase
|
from lbry.testcase import AsyncioTestCase
|
||||||
from lbry.torrent.tracker import UDPTrackerClientProtocol, UDPTrackerServerProtocol, CompactIPv4Peer
|
from lbry.torrent.tracker import UDPTrackerClientProtocol, encode, decode, CompactIPv4Peer, ConnectRequest, \
|
||||||
|
ConnectResponse, AnnounceRequest, ErrorResponse, AnnounceResponse
|
||||||
|
|
||||||
|
|
||||||
|
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 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:
|
||||||
|
self.peers.setdefault(req.info_hash, [])
|
||||||
|
compact_ip = reduce(lambda buff, x: buff + bytearray([int(x)]), address[0].split('.'), bytearray())
|
||||||
|
self.peers[req.info_hash].append(compact_ip + req.port.to_bytes(2, "big", signed=False))
|
||||||
|
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)
|
||||||
|
|
||||||
|
|
||||||
class UDPTrackerClientTestCase(AsyncioTestCase):
|
class UDPTrackerClientTestCase(AsyncioTestCase):
|
||||||
|
|
Loading…
Reference in a new issue