DHT normalizes key lengths by chopping or adding 0.

[quix0rs-apt-p2p.git] / apt_dht_Khashmir / DHT.py

diff --git a/apt_dht_Khashmir/DHT.py b/apt_dht_Khashmir/DHT.py
index 087201edb24d00c815f4498ae83c5ea584d609ec..42e2e9afd4489f33a1ddc385f1938b18301f5ec4 100644 (file)
--- a/apt_dht_Khashmir/DHT.py
+++ b/apt_dht_Khashmir/DHT.py
@@ -1,7 +1,8 @@
 
 import os, sha, random
 
-from twisted.internet import defer
+from twisted.internet import defer, reactor
+from twisted.internet.abstract import isIPAddress
 from twisted.trial import unittest
 from zope.interface import implements
 
@@ -30,7 +31,7 @@ class DHT:
         """See L{apt_dht.interfaces.IDHT}."""
         self.config_parser = config
         self.section = section
-        self.config = []
+        self.config = {}
         self.cache_dir = self.config_parser.get('DEFAULT', 'cache_dir')
         self.bootstrap = self.config_parser.getstringlist(section, 'BOOTSTRAP')
         self.bootstrap_node = self.config_parser.getboolean(section, 'BOOTSTRAP_NODE')
@@ -59,9 +60,16 @@ class DHT:
         for node in self.bootstrap:
             host, port = node.rsplit(':', 1)
             port = int(port)
-            self.khashmir.addContact(host, port, self._join_single)
+            if isIPAddress(host):
+                self._join_gotIP(host, port)
+            else:
+                reactor.resolve(host).addCallback(self._join_gotIP, port)
         
         return self.joining
+
+    def _join_gotIP(self, ip, port):
+        """Called after an IP address has been found for a single bootstrap node."""
+        self.khashmir.addContact(ip, port, self._join_single)
     
     def _join_single(self):
         """Called when a single bootstrap node has been added."""
@@ -92,6 +100,15 @@ class DHT:
             self.joined = False
             self.khashmir.shutdown()
         
+    def normalizeKey(self, key):
+        """Normalize a key's length suitable for insertion in the DHT."""
+        key_bytes = (self.config['HASH_LENGTH'] - 1) // 8 + 1
+        if len(key) < key_bytes:
+            key = key + '\000'*(key_bytes - len(key))
+        elif len(key) > key_bytes:
+            key = key[:key_bytes]
+        return key
+    
     def getValue(self, key):
         """See L{apt_dht.interfaces.IDHT}."""
         if self.config is None:
@@ -100,12 +117,13 @@ class DHT:
             raise DHTError, "have not joined a network yet"
 
         d = defer.Deferred()
+        key = self.normalizeKey(key)
         if key not in self.retrieving:
             self.khashmir.valueForKey(key, self._getValue)
         self.retrieving.setdefault(key, []).append(d)
         return d
         
-    def _getValue(self, key, result = -1):
+    def _getValue(self, key, result):
         if result:
             self.retrieved.setdefault(key, []).extend(result)
         else:
@@ -125,6 +143,7 @@ class DHT:
         if not self.joined:
             raise DHTError, "have not joined a network yet"
 
+        key = self.normalizeKey(key)
         if key in self.storing and value in self.storing[key]:
             raise DHTError, "already storing that key with the same value"
 
@@ -167,6 +186,13 @@ class TestSimpleDHT(unittest.TestCase):
         self.b.bootstrap = ["127.0.0.1:4044"]
         self.b.cache_dir = '/tmp'
         
+    def test_normalizeKey(self):
+        self.failUnless(self.a.normalizeKey('12345678901234567890') == '12345678901234567890')
+        self.failUnless(self.a.normalizeKey('12345678901234567') == '12345678901234567\000\000\000')
+        self.failUnless(self.a.normalizeKey('1234567890123456789012345') == '12345678901234567890')
+        self.failUnless(self.a.normalizeKey('1234567890123456789') == '1234567890123456789\000')
+        self.failUnless(self.a.normalizeKey('123456789012345678901') == '12345678901234567890')
+    
     def test_bootstrap_join(self):
         d = self.a.join()
         return d