8860c6719491e67c6267a1e483267587bb5dc3d2
[quix0rs-apt-p2p.git] / apt_p2p_Khashmir / khashmir.py
1 ## Copyright 2002-2004 Andrew Loewenstern, All Rights Reserved
2 # see LICENSE.txt for license information
3
4 """The main Khashmir program."""
5
6 import warnings
7 warnings.simplefilter("ignore", DeprecationWarning)
8
9 from datetime import datetime, timedelta
10 from random import randrange, shuffle
11 from sha import sha
12 from copy import copy
13 import os
14
15 from twisted.internet.defer import Deferred
16 from twisted.internet import protocol, reactor
17 from twisted.python import log
18 from twisted.trial import unittest
19
20 from db import DB
21 from ktable import KTable
22 from knode import KNodeBase, KNodeRead, KNodeWrite, NULL_ID
23 from khash import newID, newIDInRange
24 from actions import FindNode, FindValue, GetValue, StoreValue
25 from stats import StatsLogger
26 import krpc
27
28 class KhashmirBase(protocol.Factory):
29     """The base Khashmir class, with base functionality and find node, no key-value mappings.
30     
31     @type _Node: L{node.Node}
32     @ivar _Node: the knode implementation to use for this class of DHT
33     @type config: C{dictionary}
34     @ivar config: the configuration parameters for the DHT
35     @type port: C{int}
36     @ivar port: the port to listen on
37     @type store: L{db.DB}
38     @ivar store: the database to store nodes and key/value pairs in
39     @type node: L{node.Node}
40     @ivar node: this node
41     @type table: L{ktable.KTable}
42     @ivar table: the routing table
43     @type token_secrets: C{list} of C{string}
44     @ivar token_secrets: the current secrets to use to create tokens
45     @type stats: L{stats.StatsLogger}
46     @ivar stats: the statistics gatherer
47     @type udp: L{krpc.hostbroker}
48     @ivar udp: the factory for the KRPC protocol
49     @type listenport: L{twisted.internet.interfaces.IListeningPort}
50     @ivar listenport: the UDP listening port
51     @type next_checkpoint: L{twisted.internet.interfaces.IDelayedCall}
52     @ivar next_checkpoint: the delayed call for the next checkpoint
53     """
54     
55     _Node = KNodeBase
56     
57     def __init__(self, config, cache_dir='/tmp'):
58         """Initialize the Khashmir class and call the L{setup} method.
59         
60         @type config: C{dictionary}
61         @param config: the configuration parameters for the DHT
62         @type cache_dir: C{string}
63         @param cache_dir: the directory to store all files in
64             (optional, defaults to the /tmp directory)
65         """
66         self.config = None
67         self.setup(config, cache_dir)
68         
69     def setup(self, config, cache_dir):
70         """Setup all the Khashmir sub-modules.
71         
72         @type config: C{dictionary}
73         @param config: the configuration parameters for the DHT
74         @type cache_dir: C{string}
75         @param cache_dir: the directory to store all files in
76         """
77         self.config = config
78         self.port = config['PORT']
79         self.store = DB(os.path.join(cache_dir, 'khashmir.' + str(self.port) + '.db'))
80         self.node = self._loadSelfNode('', self.port)
81         self.table = KTable(self.node, config)
82         self.token_secrets = [newID()]
83         self.stats = StatsLogger(self.table, self.store, self.config)
84         
85         # Start listening
86         self.udp = krpc.hostbroker(self, self.stats, config)
87         self.udp.protocol = krpc.KRPC
88         self.listenport = reactor.listenUDP(self.port, self.udp)
89         
90         # Load the routing table and begin checkpointing
91         self._loadRoutingTable()
92         self.refreshTable(force = True)
93         self.next_checkpoint = reactor.callLater(60, self.checkpoint)
94
95     def Node(self, id, host = None, port = None):
96         """Create a new node.
97         
98         @see: L{node.Node.__init__}
99         """
100         n = self._Node(id, host, port)
101         n.table = self.table
102         n.conn = self.udp.connectionForAddr((n.host, n.port))
103         return n
104     
105     def __del__(self):
106         """Stop listening for packets."""
107         self.listenport.stopListening()
108         
109     def _loadSelfNode(self, host, port):
110         """Create this node, loading any previously saved one."""
111         id = self.store.getSelfNode()
112         if not id:
113             id = newID()
114         return self._Node(id, host, port)
115         
116     def checkpoint(self):
117         """Perform some periodic maintenance operations."""
118         # Create a new token secret
119         self.token_secrets.insert(0, newID())
120         if len(self.token_secrets) > 3:
121             self.token_secrets.pop()
122             
123         # Save some parameters for reloading
124         self.store.saveSelfNode(self.node.id)
125         self.store.dumpRoutingTable(self.table.buckets)
126         
127         # DHT maintenance
128         self.store.expireValues(self.config['KEY_EXPIRE'])
129         self.refreshTable()
130         
131         self.next_checkpoint = reactor.callLater(randrange(int(self.config['CHECKPOINT_INTERVAL'] * .9), 
132                                                            int(self.config['CHECKPOINT_INTERVAL'] * 1.1)), 
133                                                  self.checkpoint)
134         
135     def _loadRoutingTable(self):
136         """Load the previous routing table nodes from the database.
137         
138         It's usually a good idea to call refreshTable(force = True) after
139         loading the table.
140         """
141         nodes = self.store.getRoutingTable()
142         for rec in nodes:
143             n = self.Node(rec[0], rec[1], int(rec[2]))
144             self.table.insertNode(n, contacted = False)
145             
146     #{ Local interface
147     def addContact(self, host, port, callback=None, errback=None):
148         """Ping this node and add the contact info to the table on pong.
149         
150         @type host: C{string}
151         @param host: the IP address of the node to contact
152         @type port: C{int}
153         @param port:the port of the node to contact
154         @type callback: C{method}
155         @param callback: the method to call with the results, it must take 1
156             parameter, the contact info returned by the node
157             (optional, defaults to doing nothing with the results)
158         @type errback: C{method}
159         @param errback: the method to call if an error occurs
160             (optional, defaults to calling the callback with None)
161         """
162         n = self.Node(NULL_ID, host, port)
163         self.sendJoin(n, callback=callback, errback=errback)
164
165     def findNode(self, id, callback):
166         """Find the contact info for the K closest nodes in the global table.
167         
168         @type id: C{string}
169         @param id: the target ID to find the K closest nodes of
170         @type callback: C{method}
171         @param callback: the method to call with the results, it must take 1
172             parameter, the list of K closest nodes
173         """
174         # Start with our node
175         nodes = [copy(self.node)]
176
177         # Start the finding nodes action
178         state = FindNode(self, id, callback, self.config, self.stats)
179         reactor.callLater(0, state.goWithNodes, nodes)
180     
181     def insertNode(self, node, contacted = True):
182         """Try to insert a node in our local table, pinging oldest contact if necessary.
183         
184         If all you have is a host/port, then use L{addContact}, which calls this
185         method after receiving the PONG from the remote node. The reason for
186         the seperation is we can't insert a node into the table without its
187         node ID. That means of course the node passed into this method needs
188         to be a properly formed Node object with a valid ID.
189
190         @type node: L{node.Node}
191         @param node: the new node to try and insert
192         @type contacted: C{boolean}
193         @param contacted: whether the new node is known to be good, i.e.
194             responded to a request (optional, defaults to True)
195         """
196         old = self.table.insertNode(node, contacted=contacted)
197         if (old and old.id != self.node.id and
198             (datetime.now() - old.lastSeen) > 
199              timedelta(seconds=self.config['MIN_PING_INTERVAL'])):
200             
201             def _staleNodeHandler(err, oldnode = old, newnode = node, self = self):
202                 """The pinged node never responded, so replace it."""
203                 log.msg("ping failed (%s) %s/%s" % (self.config['PORT'], oldnode.host, oldnode.port))
204                 log.err(err)
205                 self.table.replaceStaleNode(oldnode, newnode)
206             
207             def _notStaleNodeHandler(dict, old=old, self=self):
208                 """Got a pong from the old node, so update it."""
209                 if dict['id'] == old.id:
210                     self.table.justSeenNode(old.id)
211             
212             # Bucket is full, check to see if old node is still available
213             self.stats.startedAction('ping')
214             df = old.ping(self.node.id)
215             df.addCallbacks(_notStaleNodeHandler, _staleNodeHandler)
216
217     def sendJoin(self, node, callback=None, errback=None):
218         """Join the DHT by pinging a bootstrap node.
219         
220         @type node: L{node.Node}
221         @param node: the node to send the join to
222         @type callback: C{method}
223         @param callback: the method to call with the results, it must take 1
224             parameter, the contact info returned by the node
225             (optional, defaults to doing nothing with the results)
226         @type errback: C{method}
227         @param errback: the method to call if an error occurs
228             (optional, defaults to calling the callback with None)
229         """
230
231         def _pongHandler(dict, node=node, self=self, callback=callback):
232             """Node responded properly, callback with response."""
233             n = self.Node(dict['id'], dict['_krpc_sender'][0], dict['_krpc_sender'][1])
234             self.insertNode(n)
235             if callback:
236                 callback((dict['ip_addr'], dict['port']))
237
238         def _defaultPong(err, node=node, self=self, callback=callback, errback=errback):
239             """Error occurred, fail node and errback or callback with error."""
240             log.msg("join failed (%s) %s/%s" % (self.config['PORT'], node.host, node.port))
241             log.err(err)
242             self.table.nodeFailed(node)
243             if errback:
244                 errback()
245             elif callback:
246                 callback(None)
247         
248         self.stats.startedAction('join')
249         df = node.join(self.node.id)
250         df.addCallbacks(_pongHandler, _defaultPong)
251
252     def findCloseNodes(self, callback=lambda a: None):
253         """Perform a findNode on the ID one away from our own.
254
255         This will allow us to populate our table with nodes on our network
256         closest to our own. This is called as soon as we start up with an
257         empty table.
258
259         @type callback: C{method}
260         @param callback: the method to call with the results, it must take 1
261             parameter, the list of K closest nodes
262             (optional, defaults to doing nothing with the results)
263         """
264         id = self.node.id[:-1] + chr((ord(self.node.id[-1]) + 1) % 256)
265         self.findNode(id, callback)
266
267     def refreshTable(self, force = False):
268         """Check all the buckets for those that need refreshing.
269         
270         @param force: refresh all buckets regardless of last bucket access time
271             (optional, defaults to False)
272         """
273         def callback(nodes):
274             pass
275     
276         for bucket in self.table.buckets:
277             if force or (datetime.now() - bucket.lastAccessed > 
278                          timedelta(seconds=self.config['BUCKET_STALENESS'])):
279                 # Choose a random ID in the bucket and try and find it
280                 id = newIDInRange(bucket.min, bucket.max)
281                 self.findNode(id, callback)
282
283     def shutdown(self):
284         """Closes the port and cancels pending later calls."""
285         self.listenport.stopListening()
286         try:
287             self.next_checkpoint.cancel()
288         except:
289             pass
290         self.store.close()
291     
292     def getStats(self):
293         """Gather the statistics for the DHT."""
294         return self.stats.formatHTML()
295
296     #{ Remote interface
297     def krpc_ping(self, id, _krpc_sender = None):
298         """Pong with our ID.
299         
300         @type id: C{string}
301         @param id: the node ID of the sender node
302         @type _krpc_sender: (C{string}, C{int})
303         @param _krpc_sender: the sender node's IP address and port
304         """
305         if _krpc_sender is not None:
306             n = self.Node(id, _krpc_sender[0], _krpc_sender[1])
307             self.insertNode(n, contacted = False)
308
309         return {"id" : self.node.id}
310         
311     def krpc_join(self, id, _krpc_sender = None):
312         """Add the node by responding with its address and port.
313         
314         @type id: C{string}
315         @param id: the node ID of the sender node
316         @type _krpc_sender: (C{string}, C{int})
317         @param _krpc_sender: the sender node's IP address and port
318         """
319         if _krpc_sender is not None:
320             n = self.Node(id, _krpc_sender[0], _krpc_sender[1])
321             self.insertNode(n, contacted = False)
322         else:
323             _krpc_sender = ('127.0.0.1', self.port)
324
325         return {"ip_addr" : _krpc_sender[0], "port" : _krpc_sender[1], "id" : self.node.id}
326         
327     def krpc_find_node(self, id, target, _krpc_sender = None):
328         """Find the K closest nodes to the target in the local routing table.
329         
330         @type target: C{string}
331         @param target: the target ID to find nodes for
332         @type id: C{string}
333         @param id: the node ID of the sender node
334         @type _krpc_sender: (C{string}, C{int})
335         @param _krpc_sender: the sender node's IP address and port
336         """
337         if _krpc_sender is not None:
338             n = self.Node(id, _krpc_sender[0], _krpc_sender[1])
339             self.insertNode(n, contacted = False)
340         else:
341             _krpc_sender = ('127.0.0.1', self.port)
342
343         nodes = self.table.findNodes(target)
344         nodes = map(lambda node: node.contactInfo(), nodes)
345         token = sha(self.token_secrets[0] + _krpc_sender[0]).digest()
346         return {"nodes" : nodes, "token" : token, "id" : self.node.id}
347
348
349 class KhashmirRead(KhashmirBase):
350     """The read-only Khashmir class, which can only retrieve (not store) key/value mappings."""
351
352     _Node = KNodeRead
353
354     #{ Local interface
355     def findValue(self, key, callback):
356         """Get the nodes that have values for the key from the global table.
357         
358         @type key: C{string}
359         @param key: the target key to find the values for
360         @type callback: C{method}
361         @param callback: the method to call with the results, it must take 1
362             parameter, the list of nodes with values
363         """
364         # Start with ourself
365         nodes = [copy(self.node)]
366         
367         # Search for others starting with the locally found ones
368         state = FindValue(self, key, callback, self.config, self.stats)
369         reactor.callLater(0, state.goWithNodes, nodes)
370
371     def valueForKey(self, key, callback, searchlocal = True):
372         """Get the values found for key in global table.
373         
374         Callback will be called with a list of values for each peer that
375         returns unique values. The final callback will be an empty list.
376
377         @type key: C{string}
378         @param key: the target key to get the values for
379         @type callback: C{method}
380         @param callback: the method to call with the results, it must take 2
381             parameters: the key, and the values found
382         @type searchlocal: C{boolean}
383         @param searchlocal: whether to also look for any local values
384         """
385
386         def _getValueForKey(nodes, key=key, response=callback, self=self, searchlocal=searchlocal):
387             """Use the found nodes to send requests for values to."""
388             # Get any local values
389             if searchlocal:
390                 l = self.store.retrieveValues(key)
391                 if len(l) > 0:
392                     node = copy(self.node)
393                     node.updateNumValues(len(l))
394                     nodes = nodes + [node]
395
396             state = GetValue(self, key, self.config['RETRIEVE_VALUES'], response, self.config, self.stats)
397             reactor.callLater(0, state.goWithNodes, nodes)
398             
399         # First lookup nodes that have values for the key
400         self.findValue(key, _getValueForKey)
401
402     #{ Remote interface
403     def krpc_find_value(self, id, key, _krpc_sender = None):
404         """Find the number of values stored locally for the key, and the K closest nodes.
405         
406         @type key: C{string}
407         @param key: the target key to find the values and nodes for
408         @type id: C{string}
409         @param id: the node ID of the sender node
410         @type _krpc_sender: (C{string}, C{int})
411         @param _krpc_sender: the sender node's IP address and port
412         """
413         if _krpc_sender is not None:
414             n = self.Node(id, _krpc_sender[0], _krpc_sender[1])
415             self.insertNode(n, contacted = False)
416     
417         nodes = self.table.findNodes(key)
418         nodes = map(lambda node: node.contactInfo(), nodes)
419         num_values = self.store.countValues(key)
420         return {'nodes' : nodes, 'num' : num_values, "id": self.node.id}
421
422     def krpc_get_value(self, id, key, num, _krpc_sender = None):
423         """Retrieve the values stored locally for the key.
424         
425         @type key: C{string}
426         @param key: the target key to retrieve the values for
427         @type num: C{int}
428         @param num: the maximum number of values to retrieve, or 0 to
429             retrieve all of them
430         @type id: C{string}
431         @param id: the node ID of the sender node
432         @type _krpc_sender: (C{string}, C{int})
433         @param _krpc_sender: the sender node's IP address and port
434         """
435         if _krpc_sender is not None:
436             n = self.Node(id, _krpc_sender[0], _krpc_sender[1])
437             self.insertNode(n, contacted = False)
438     
439         l = self.store.retrieveValues(key)
440         if num == 0 or num >= len(l):
441             return {'values' : l, "id": self.node.id}
442         else:
443             shuffle(l)
444             return {'values' : l[:num], "id": self.node.id}
445
446
447 class KhashmirWrite(KhashmirRead):
448     """The read-write Khashmir class, which can store and retrieve key/value mappings."""
449
450     _Node = KNodeWrite
451
452     #{ Local interface
453     def storeValueForKey(self, key, value, callback=None):
454         """Stores the value for the key in the global table.
455         
456         No status in this implementation, peers respond but don't indicate
457         status of storing values.
458
459         @type key: C{string}
460         @param key: the target key to store the value for
461         @type value: C{string}
462         @param value: the value to store with the key
463         @type callback: C{method}
464         @param callback: the method to call with the results, it must take 3
465             parameters: the key, the value stored, and the result of the store
466             (optional, defaults to doing nothing with the results)
467         """
468         def _storeValueForKey(nodes, key=key, value=value, response=callback, self=self):
469             """Use the returned K closest nodes to store the key at."""
470             if not response:
471                 def _storedValueHandler(key, value, sender):
472                     """Default callback that does nothing."""
473                     pass
474                 response = _storedValueHandler
475             action = StoreValue(self, key, value, self.config['STORE_REDUNDANCY'], response, self.config, self.stats)
476             reactor.callLater(0, action.goWithNodes, nodes)
477             
478         # First find the K closest nodes to operate on.
479         self.findNode(key, _storeValueForKey)
480                     
481     #{ Remote interface
482     def krpc_store_value(self, id, key, value, token, _krpc_sender = None):
483         """Store the value locally with the key.
484         
485         @type key: C{string}
486         @param key: the target key to store the value for
487         @type value: C{string}
488         @param value: the value to store with the key
489         @param token: the token to confirm that this peer contacted us previously
490         @type id: C{string}
491         @param id: the node ID of the sender node
492         @type _krpc_sender: (C{string}, C{int})
493         @param _krpc_sender: the sender node's IP address and port
494         """
495         if _krpc_sender is not None:
496             n = self.Node(id, _krpc_sender[0], _krpc_sender[1])
497             self.insertNode(n, contacted = False)
498         else:
499             _krpc_sender = ('127.0.0.1', self.port)
500
501         for secret in self.token_secrets:
502             this_token = sha(secret + _krpc_sender[0]).digest()
503             if token == this_token:
504                 self.store.storeValue(key, value)
505                 return {"id" : self.node.id}
506         raise krpc.KrpcError, (krpc.KRPC_ERROR_INVALID_TOKEN, 'token is invalid, do a find_nodes to get a fresh one')
507
508
509 class Khashmir(KhashmirWrite):
510     """The default Khashmir class (currently the read-write L{KhashmirWrite})."""
511     _Node = KNodeWrite
512
513
514 class SimpleTests(unittest.TestCase):
515     
516     timeout = 10
517     DHT_DEFAULTS = {'PORT': 9977, 'K': 8, 'HASH_LENGTH': 160,
518                     'CHECKPOINT_INTERVAL': 300, 'CONCURRENT_REQS': 4,
519                     'STORE_REDUNDANCY': 3, 'RETRIEVE_VALUES': -10000,
520                     'MAX_FAILURES': 3,
521                     'MIN_PING_INTERVAL': 900,'BUCKET_STALENESS': 3600,
522                     'KEY_EXPIRE': 3600, 'SPEW': False, }
523
524     def setUp(self):
525         d = self.DHT_DEFAULTS.copy()
526         d['PORT'] = 4044
527         self.a = Khashmir(d)
528         d = self.DHT_DEFAULTS.copy()
529         d['PORT'] = 4045
530         self.b = Khashmir(d)
531         
532     def tearDown(self):
533         self.a.shutdown()
534         self.b.shutdown()
535         os.unlink(self.a.store.db)
536         os.unlink(self.b.store.db)
537
538     def testAddContact(self):
539         self.failUnlessEqual(len(self.a.table.buckets), 1)
540         self.failUnlessEqual(len(self.a.table.buckets[0].l), 0)
541
542         self.failUnlessEqual(len(self.b.table.buckets), 1)
543         self.failUnlessEqual(len(self.b.table.buckets[0].l), 0)
544
545         self.a.addContact('127.0.0.1', 4045)
546         reactor.iterate()
547         reactor.iterate()
548         reactor.iterate()
549         reactor.iterate()
550
551         self.failUnlessEqual(len(self.a.table.buckets), 1)
552         self.failUnlessEqual(len(self.a.table.buckets[0].l), 1)
553         self.failUnlessEqual(len(self.b.table.buckets), 1)
554         self.failUnlessEqual(len(self.b.table.buckets[0].l), 1)
555
556     def testStoreRetrieve(self):
557         self.a.addContact('127.0.0.1', 4045)
558         reactor.iterate()
559         reactor.iterate()
560         reactor.iterate()
561         reactor.iterate()
562         self.got = 0
563         self.a.storeValueForKey(sha('foo').digest(), 'foobar')
564         reactor.iterate()
565         reactor.iterate()
566         reactor.iterate()
567         reactor.iterate()
568         reactor.iterate()
569         reactor.iterate()
570         self.a.valueForKey(sha('foo').digest(), self._cb)
571         reactor.iterate()
572         reactor.iterate()
573         reactor.iterate()
574         reactor.iterate()
575         reactor.iterate()
576         reactor.iterate()
577         reactor.iterate()
578
579     def _cb(self, key, val):
580         if not val:
581             self.failUnlessEqual(self.got, 1)
582         elif 'foobar' in val:
583             self.got = 1
584
585
586 class MultiTest(unittest.TestCase):
587     
588     timeout = 30
589     num = 20
590     DHT_DEFAULTS = {'PORT': 9977, 'K': 8, 'HASH_LENGTH': 160,
591                     'CHECKPOINT_INTERVAL': 300, 'CONCURRENT_REQS': 4,
592                     'STORE_REDUNDANCY': 3, 'RETRIEVE_VALUES': -10000,
593                     'MAX_FAILURES': 3,
594                     'MIN_PING_INTERVAL': 900,'BUCKET_STALENESS': 3600,
595                     'KEY_EXPIRE': 3600, 'SPEW': False, }
596
597     def _done(self, val):
598         self.done = 1
599         
600     def setUp(self):
601         self.l = []
602         self.startport = 4088
603         for i in range(self.num):
604             d = self.DHT_DEFAULTS.copy()
605             d['PORT'] = self.startport + i
606             self.l.append(Khashmir(d))
607         reactor.iterate()
608         reactor.iterate()
609         
610         for i in self.l:
611             i.addContact('127.0.0.1', self.l[randrange(0,self.num)].port)
612             i.addContact('127.0.0.1', self.l[randrange(0,self.num)].port)
613             i.addContact('127.0.0.1', self.l[randrange(0,self.num)].port)
614             reactor.iterate()
615             reactor.iterate()
616             reactor.iterate() 
617             
618         for i in self.l:
619             self.done = 0
620             i.findCloseNodes(self._done)
621             while not self.done:
622                 reactor.iterate()
623         for i in self.l:
624             self.done = 0
625             i.findCloseNodes(self._done)
626             while not self.done:
627                 reactor.iterate()
628
629     def tearDown(self):
630         for i in self.l:
631             i.shutdown()
632             os.unlink(i.store.db)
633             
634         reactor.iterate()
635         
636     def testStoreRetrieve(self):
637         for i in range(10):
638             K = newID()
639             V = newID()
640             
641             for a in range(3):
642                 self.done = 0
643                 def _scb(key, value, result):
644                     self.done = 1
645                 self.l[randrange(0, self.num)].storeValueForKey(K, V, _scb)
646                 while not self.done:
647                     reactor.iterate()
648
649
650                 def _rcb(key, val):
651                     if not val:
652                         self.done = 1
653                         self.failUnlessEqual(self.got, 1)
654                     elif V in val:
655                         self.got = 1
656                 for x in range(3):
657                     self.got = 0
658                     self.done = 0
659                     self.l[randrange(0, self.num)].valueForKey(K, _rcb)
660                     while not self.done:
661                         reactor.iterate()