diff --git a/lbry/torrent/tracker.py b/lbry/torrent/tracker.py index 3d548e45b..cc263b380 100644 --- a/lbry/torrent/tracker.py +++ b/lbry/torrent/tracker.py @@ -77,14 +77,15 @@ class UDPTrackerClientProtocol(asyncio.DatagramProtocol): return decode(ConnectResponse, await self.request(ConnectRequest(0x41727101980, 0, transaction_id), tracker_ip, tracker_port)) - async def announce(self, info_hash, peer_id, port, tracker_ip, tracker_port, connection_id=None): + async def announce(self, info_hash, peer_id, port, tracker_ip, tracker_port, connection_id=None, stopped=False): if not connection_id: reply = await self.connect(tracker_ip, tracker_port) connection_id = reply.connection_id # this should make the key deterministic but unique per info hash + peer id key = int.from_bytes(info_hash[:4], "big") ^ int.from_bytes(peer_id[:4], "big") ^ port transaction_id = random.getrandbits(32) - req = AnnounceRequest(connection_id, 1, transaction_id, info_hash, peer_id, 0, 0, 0, 1, 0, key, -1, port) + req = AnnounceRequest( + connection_id, 1, transaction_id, info_hash, peer_id, 0, 0, 0, 3 if stopped else 1, 0, key, -1, port) reply = await self.request(req, tracker_ip, tracker_port) return decode(AnnounceResponse, reply), connection_id @@ -112,13 +113,13 @@ class UDPTrackerClientProtocol(asyncio.DatagramProtocol): self.transport = None -async def get_peer_list(info_hash, node_id, port, tracker_ip, tracker_port): +async def get_peer_list(info_hash, node_id, port, tracker_ip, tracker_port, stopped=False): node_id = node_id or random.getrandbits(160).to_bytes(20, "big", signed=False) tracker_ip = await resolve_host(tracker_ip, tracker_port, 'udp') proto = UDPTrackerClientProtocol() transport, _ = await asyncio.get_running_loop().create_datagram_endpoint(lambda: proto, local_addr=("0.0.0.0", 0)) try: - reply, _ = await proto.announce(info_hash, node_id, port, tracker_ip, tracker_port) + reply, _ = await proto.announce(info_hash, node_id, port, tracker_ip, tracker_port, stopped=stopped) return reply.peers except asyncio.CancelledError: raise diff --git a/tests/unit/torrent/test_tracker.py b/tests/unit/torrent/test_tracker.py index 66f147f43..3c4c0fa9d 100644 --- a/tests/unit/torrent/test_tracker.py +++ b/tests/unit/torrent/test_tracker.py @@ -32,7 +32,11 @@ class UDPTrackerServerProtocol(asyncio.DatagramProtocol): # for testing. Not su 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)) + compact_address = compact_ip + req.port.to_bytes(2, "big", signed=False) + if req.event != 3: + self.peers[req.info_hash].append(compact_address) + elif compact_address in self.peers[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) @@ -57,8 +61,9 @@ class UDPTrackerClientTestCase(AsyncioTestCase): async def test_announce_using_helper_function(self): info_hash = random.getrandbits(160).to_bytes(20, "big", signed=False) peers = await get_peer_list(info_hash, None, 4444, "127.0.0.1", 59900) - self.assertEqual(len(peers), 1) self.assertEqual(peers, [CompactIPv4Peer(int.from_bytes(bytes([127, 0, 0, 1]), "big", signed=False), 4444)]) + peers = await get_peer_list(info_hash, None, 4444, "127.0.0.1", 59900, stopped=True) + self.assertEqual(peers, []) async def test_error(self): info_hash = random.getrandbits(160).to_bytes(20, "big", signed=False)