py-libp2p/libp2p/kademlia/kad_peerinfo.py
2019-10-24 20:10:45 +02:00

154 lines
4.2 KiB
Python

import heapq
from operator import itemgetter
import random
from typing import List
from multiaddr import Multiaddr
from libp2p.peer.id import ID
from libp2p.peer.peerinfo import PeerInfo
from .utils import digest
P_IP = "ip4"
P_UDP = "udp"
class KadPeerInfo(PeerInfo):
def __init__(self, peer_id, addrs):
super(KadPeerInfo, self).__init__(peer_id, addrs)
self.peer_id_bytes = peer_id.to_bytes()
self.xor_id = peer_id.xor_id
self.addrs = addrs
self.ip = self.addrs[0].value_for_protocol(P_IP) if addrs else None
self.port = int(self.addrs[0].value_for_protocol(P_UDP)) if addrs else None
def same_home_as(self, node):
return sorted(self.addrs) == sorted(node.addrs)
def distance_to(self, node):
"""Get the distance between this node and another."""
return self.xor_id ^ node.xor_id
def __iter__(self):
"""
Enables use of Node as a tuple - i.e., tuple(node) works.
"""
return iter([self.peer_id_bytes, self.ip, self.port])
def __repr__(self):
return repr([self.xor_id, self.ip, self.port, self.peer_id_bytes])
def __str__(self):
return "%s:%s" % (self.ip, str(self.port))
def encode(self):
return (
str(self.peer_id_bytes)
+ "\n"
+ str("/ip4/" + str(self.ip) + "/udp/" + str(self.port))
)
class KadPeerHeap:
"""A heap of peers ordered by distance to a given node."""
def __init__(self, node, maxsize):
"""
Constructor.
@param node: The node to measure all distnaces from.
@param maxsize: The maximum size that this heap can grow to.
"""
self.node = node
self.heap = []
self.contacted = set()
self.maxsize = maxsize
def remove(self, peers):
"""
Remove a list of peer ids from this heap.
Note that while this heap retains a constant visible size (based
on the iterator), it's actual size may be quite a bit larger
than what's exposed. Therefore, removal of nodes may not change
the visible size as previously added nodes suddenly become
visible.
"""
peers = set(peers)
if not peers:
return
nheap = []
for distance, node in self.heap:
if node.peer_id_bytes not in peers:
heapq.heappush(nheap, (distance, node))
self.heap = nheap
def get_node(self, node_id):
for _, node in self.heap:
if node.peer_id_bytes == node_id:
return node
return None
def have_contacted_all(self):
return len(self.get_uncontacted()) == 0
def get_ids(self):
return [n.peer_id_bytes for n in self]
def mark_contacted(self, node):
self.contacted.add(node.peer_id_bytes)
def popleft(self):
return heapq.heappop(self.heap)[1] if self else None
def push(self, nodes):
"""
Push nodes onto heap.
@param nodes: This can be a single item or a C{list}.
"""
if not isinstance(nodes, list):
nodes = [nodes]
for node in nodes:
if node not in self:
distance = self.node.distance_to(node)
heapq.heappush(self.heap, (distance, node))
def __len__(self):
return min(len(self.heap), self.maxsize)
def __iter__(self):
nodes = heapq.nsmallest(self.maxsize, self.heap)
return iter(map(itemgetter(1), nodes))
def __contains__(self, node):
for _, other in self.heap:
if node.peer_id_bytes == other.peer_id_bytes:
return True
return False
def get_uncontacted(self):
return [n for n in self if n.peer_id_bytes not in self.contacted]
def create_kad_peerinfo(node_id_bytes=None, sender_ip=None, sender_port=None):
node_id = (
ID(node_id_bytes) if node_id_bytes else ID(digest(random.getrandbits(255)))
)
addrs: List[Multiaddr]
if sender_ip and sender_port:
addrs = [
Multiaddr(
"/" + P_IP + "/" + str(sender_ip) + "/" + P_UDP + "/" + str(sender_port)
)
]
else:
addrs = []
return KadPeerInfo(node_id, addrs)