341 lines
11 KiB
Python
341 lines
11 KiB
Python
# The contents of this file are subject to the BitTorrent Open Source License
|
|
# Version 1.1 (the License). You may not copy or use this file, in either
|
|
# source code or executable form, except in compliance with the License. You
|
|
# may obtain a copy of the License at http://www.bittorrent.com/license/.
|
|
#
|
|
# Software distributed under the License is distributed on an AS IS basis,
|
|
# WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
|
|
# for the specific language governing rights and limitations under the
|
|
# License.
|
|
|
|
from BitTorrent.platform import bttime as time
|
|
from bisect import *
|
|
from types import *
|
|
|
|
import khash as hash
|
|
import const
|
|
from const import K, HASH_LENGTH, NULL_ID, MAX_FAILURES, MIN_PING_INTERVAL
|
|
from node import Node
|
|
|
|
class KTable:
|
|
"""local routing table for a kademlia like distributed hash table"""
|
|
def __init__(self, node):
|
|
# this is the root node, a.k.a. US!
|
|
self.node = node
|
|
self.buckets = [KBucket([], 0L, 2L**HASH_LENGTH)]
|
|
self.insertNode(node)
|
|
|
|
def _bucketIndexForInt(self, num):
|
|
"""the index of the bucket that should hold int"""
|
|
return bisect_left(self.buckets, num)
|
|
|
|
def bucketForInt(self, num):
|
|
return self.buckets[self._bucketIndexForInt(num)]
|
|
|
|
def findNodes(self, id, invalid=True):
|
|
"""
|
|
return K nodes in our own local table closest to the ID.
|
|
"""
|
|
|
|
if isinstance(id, str):
|
|
num = hash.intify(id)
|
|
elif isinstance(id, Node):
|
|
num = id.num
|
|
elif isinstance(id, int) or isinstance(id, long):
|
|
num = id
|
|
else:
|
|
raise TypeError, "findNodes requires an int, string, or Node"
|
|
|
|
nodes = []
|
|
i = self._bucketIndexForInt(num)
|
|
|
|
# if this node is already in our table then return it
|
|
try:
|
|
node = self.buckets[i].getNodeWithInt(num)
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
return [node]
|
|
|
|
# don't have the node, get the K closest nodes
|
|
nodes = nodes + self.buckets[i].l
|
|
if not invalid:
|
|
nodes = [a for a in nodes if not a.invalid]
|
|
if len(nodes) < K:
|
|
# need more nodes
|
|
min = i - 1
|
|
max = i + 1
|
|
while len(nodes) < K and (min >= 0 or max < len(self.buckets)):
|
|
#ASw: note that this requires K be even
|
|
if min >= 0:
|
|
nodes = nodes + self.buckets[min].l
|
|
if max < len(self.buckets):
|
|
nodes = nodes + self.buckets[max].l
|
|
min = min - 1
|
|
max = max + 1
|
|
if not invalid:
|
|
nodes = [a for a in nodes if not a.invalid]
|
|
|
|
nodes.sort(lambda a, b, num=num: cmp(num ^ a.num, num ^ b.num))
|
|
return nodes[:K]
|
|
|
|
def _splitBucket(self, a):
|
|
diff = (a.max - a.min) / 2
|
|
b = KBucket([], a.max - diff, a.max)
|
|
self.buckets.insert(self.buckets.index(a.min) + 1, b)
|
|
a.max = a.max - diff
|
|
# transfer nodes to new bucket
|
|
for anode in a.l[:]:
|
|
if anode.num >= a.max:
|
|
a.removeNode(anode)
|
|
b.addNode(anode)
|
|
|
|
def replaceStaleNode(self, stale, new):
|
|
"""this is used by clients to replace a node returned by insertNode after
|
|
it fails to respond to a Pong message"""
|
|
i = self._bucketIndexForInt(stale.num)
|
|
|
|
if self.buckets[i].hasNode(stale):
|
|
self.buckets[i].removeNode(stale)
|
|
if new and self.buckets[i].hasNode(new):
|
|
self.buckets[i].seenNode(new)
|
|
elif new:
|
|
self.buckets[i].addNode(new)
|
|
|
|
return
|
|
|
|
def insertNode(self, node, contacted=1, nocheck=False):
|
|
"""
|
|
this insert the node, returning None if successful, returns the oldest node in the bucket if it's full
|
|
the caller responsible for pinging the returned node and calling replaceStaleNode if it is found to be stale!!
|
|
contacted means that yes, we contacted THEM and we know the node is reachable
|
|
"""
|
|
if node.id == NULL_ID or node.id == self.node.id:
|
|
return
|
|
|
|
if contacted:
|
|
node.updateLastSeen()
|
|
|
|
# get the bucket for this node
|
|
i = self._bucketIndexForInt(node.num)
|
|
# check to see if node is in the bucket already
|
|
if self.buckets[i].hasNode(node):
|
|
it = self.buckets[i].l.index(node.num)
|
|
xnode = self.buckets[i].l[it]
|
|
if contacted:
|
|
node.age = xnode.age
|
|
self.buckets[i].seenNode(node)
|
|
elif xnode.lastSeen != 0 and xnode.port == node.port and xnode.host == node.host:
|
|
xnode.updateLastSeen()
|
|
return
|
|
|
|
# we don't have this node, check to see if the bucket is full
|
|
if not self.buckets[i].bucketFull():
|
|
# no, append this node and return
|
|
self.buckets[i].addNode(node)
|
|
return
|
|
|
|
# full bucket, check to see if any nodes are invalid
|
|
t = time()
|
|
def ls(a, b):
|
|
if a.lastSeen > b.lastSeen:
|
|
return 1
|
|
elif b.lastSeen > a.lastSeen:
|
|
return -1
|
|
return 0
|
|
|
|
invalid = [x for x in self.buckets[i].invalid.values() if x.invalid]
|
|
if len(invalid) and not nocheck:
|
|
invalid.sort(ls)
|
|
while invalid and not self.buckets[i].hasNode(invalid[0]):
|
|
del(self.buckets[i].invalid[invalid[0].num])
|
|
invalid = invalid[1:]
|
|
if invalid and (invalid[0].lastSeen == 0 and invalid[0].fails < MAX_FAILURES):
|
|
return invalid[0]
|
|
elif invalid:
|
|
self.replaceStaleNode(invalid[0], node)
|
|
return
|
|
|
|
stale = [n for n in self.buckets[i].l if (t - n.lastSeen) > MIN_PING_INTERVAL]
|
|
if len(stale) and not nocheck:
|
|
stale.sort(ls)
|
|
return stale[0]
|
|
|
|
# bucket is full and all nodes are valid, check to see if self.node is in the bucket
|
|
if not (self.buckets[i].min <= self.node < self.buckets[i].max):
|
|
return
|
|
|
|
# this bucket is full and contains our node, split the bucket
|
|
if len(self.buckets) >= HASH_LENGTH:
|
|
# our table is FULL, this is really unlikely
|
|
print "Hash Table is FULL! Increase K!"
|
|
return
|
|
|
|
self._splitBucket(self.buckets[i])
|
|
|
|
# now that the bucket is split and balanced, try to insert the node again
|
|
return self.insertNode(node, contacted)
|
|
|
|
def justSeenNode(self, id):
|
|
"""call this any time you get a message from a node
|
|
it will update it in the table if it's there """
|
|
try:
|
|
n = self.findNodes(id)[0]
|
|
except IndexError:
|
|
return None
|
|
else:
|
|
tstamp = n.lastSeen
|
|
n.updateLastSeen()
|
|
bucket = self.bucketForInt(n.num)
|
|
bucket.seenNode(n)
|
|
return tstamp
|
|
|
|
def invalidateNode(self, n):
|
|
"""
|
|
forget about node n - use when you know that node is invalid
|
|
"""
|
|
n.invalid = True
|
|
self.bucket = self.bucketForInt(n.num)
|
|
self.bucket.invalidateNode(n)
|
|
|
|
def nodeFailed(self, node):
|
|
""" call this when a node fails to respond to a message, to invalidate that node """
|
|
try:
|
|
n = self.findNodes(node.num)[0]
|
|
except IndexError:
|
|
return None
|
|
else:
|
|
if n.msgFailed() >= const.MAX_FAILURES:
|
|
self.invalidateNode(n)
|
|
|
|
def numPeers(self):
|
|
""" estimated number of connectable nodes in global table """
|
|
return 8 * (2 ** (len(self.buckets) - 1))
|
|
|
|
class KBucket:
|
|
__slots__ = ('min', 'max', 'lastAccessed')
|
|
def __init__(self, contents, min, max):
|
|
self.l = contents
|
|
self.index = {}
|
|
self.invalid = {}
|
|
self.min = min
|
|
self.max = max
|
|
self.lastAccessed = time()
|
|
|
|
def touch(self):
|
|
self.lastAccessed = time()
|
|
|
|
def lacmp(self, a, b):
|
|
if a.lastSeen > b.lastSeen:
|
|
return 1
|
|
elif b.lastSeen > a.lastSeen:
|
|
return -1
|
|
return 0
|
|
|
|
def sort(self):
|
|
self.l.sort(self.lacmp)
|
|
|
|
def getNodeWithInt(self, num):
|
|
try:
|
|
node = self.index[num]
|
|
except KeyError:
|
|
raise ValueError
|
|
return node
|
|
|
|
def addNode(self, node):
|
|
if len(self.l) >= K:
|
|
return
|
|
if self.index.has_key(node.num):
|
|
return
|
|
self.l.append(node)
|
|
self.index[node.num] = node
|
|
self.touch()
|
|
|
|
def removeNode(self, node):
|
|
assert self.index.has_key(node.num)
|
|
del(self.l[self.l.index(node.num)])
|
|
del(self.index[node.num])
|
|
try:
|
|
del(self.invalid[node.num])
|
|
except KeyError:
|
|
pass
|
|
self.touch()
|
|
|
|
def invalidateNode(self, node):
|
|
self.invalid[node.num] = node
|
|
|
|
def seenNode(self, node):
|
|
try:
|
|
del(self.invalid[node.num])
|
|
except KeyError:
|
|
pass
|
|
it = self.l.index(node.num)
|
|
del(self.l[it])
|
|
self.l.append(node)
|
|
self.index[node.num] = node
|
|
|
|
def hasNode(self, node):
|
|
return self.index.has_key(node.num)
|
|
|
|
def bucketFull(self):
|
|
return len(self.l) >= K
|
|
|
|
def __repr__(self):
|
|
return "<KBucket %d items (%d to %d)>" % (len(self.l), self.min, self.max)
|
|
|
|
## Comparators
|
|
# necessary for bisecting list of buckets with a hash expressed as an integer or a distance
|
|
# compares integer or node object with the bucket's range
|
|
def __lt__(self, a):
|
|
if isinstance(a, Node): a = a.num
|
|
return self.max <= a
|
|
def __le__(self, a):
|
|
if isinstance(a, Node): a = a.num
|
|
return self.min < a
|
|
def __gt__(self, a):
|
|
if isinstance(a, Node): a = a.num
|
|
return self.min > a
|
|
def __ge__(self, a):
|
|
if isinstance(a, Node): a = a.num
|
|
return self.max >= a
|
|
def __eq__(self, a):
|
|
if isinstance(a, Node): a = a.num
|
|
return self.min <= a and self.max > a
|
|
def __ne__(self, a):
|
|
if isinstance(a, Node): a = a.num
|
|
return self.min >= a or self.max < a
|
|
|
|
|
|
### UNIT TESTS ###
|
|
import unittest
|
|
|
|
class TestKTable(unittest.TestCase):
|
|
def setUp(self):
|
|
self.a = Node().init(hash.newID(), 'localhost', 2002)
|
|
self.t = KTable(self.a)
|
|
|
|
def testAddNode(self):
|
|
self.b = Node().init(hash.newID(), 'localhost', 2003)
|
|
self.t.insertNode(self.b)
|
|
self.assertEqual(len(self.t.buckets[0].l), 1)
|
|
self.assertEqual(self.t.buckets[0].l[0], self.b)
|
|
|
|
def testRemove(self):
|
|
self.testAddNode()
|
|
self.t.invalidateNode(self.b)
|
|
self.assertEqual(len(self.t.buckets[0].l), 0)
|
|
|
|
def testFail(self):
|
|
self.testAddNode()
|
|
for i in range(const.MAX_FAILURES - 1):
|
|
self.t.nodeFailed(self.b)
|
|
self.assertEqual(len(self.t.buckets[0].l), 1)
|
|
self.assertEqual(self.t.buckets[0].l[0], self.b)
|
|
|
|
self.t.nodeFailed(self.b)
|
|
self.assertEqual(len(self.t.buckets[0].l), 0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|