isolate test_keygen data

[tahoe-lafs/tahoe-lafs.git] / src / allmydata / key_generator.py

diff --git a/src/allmydata/key_generator.py b/src/allmydata/key_generator.py
index aa735c4c5492ac75c2ea42a6775595a073e5e76d..63bf51b833f31be7ca567cfa8654398164a5286a 100644 (file)
--- a/src/allmydata/key_generator.py
+++ b/src/allmydata/key_generator.py
@@ -2,28 +2,36 @@
 import os
 import time
 
-import foolscap
-from foolscap.eventual import eventually
+from foolscap.api import Referenceable, Tub
 from zope.interface import implements
 from twisted.internet import reactor
 from twisted.application import service
-from twisted.python import log
+from allmydata.util import log, fileutil
 
 from pycryptopp.publickey import rsa
 from allmydata.interfaces import RIKeyGenerator
 
-class KeyGenerator(foolscap.Referenceable):
+class KeyGenerator(service.MultiService, Referenceable):
     implements(RIKeyGenerator)
 
-    DEFAULT_KEY_SIZE = 2048
     pool_size = 16 # no. keys to keep on hand in the pool
     pool_refresh_delay = 6 # no. sec to wait after a fetch before generating new keys
     verbose = False
 
-    def __init__(self):
+    def __init__(self, default_key_size=2048):
+        service.MultiService.__init__(self)
         self.keypool = []
         self.last_fetch = 0
-        eventually(self.maybe_refill_pool)
+        self.default_key_size = default_key_size
+
+    def startService(self):
+        self.timer = reactor.callLater(0, self.maybe_refill_pool)
+        return service.MultiService.startService(self)
+
+    def stopService(self):
+        if self.timer.active():
+            self.timer.cancel()
+        return service.MultiService.stopService(self)
 
     def __repr__(self):
         return '<KeyGenerator[%s]>' % (len(self.keypool),)
@@ -34,16 +42,20 @@ class KeyGenerator(foolscap.Referenceable):
 
     def reset_timer(self):
         self.last_fetch = time.time()
-        reactor.callLater(self.pool_refresh_delay, self.maybe_refill_pool)
+        if self.timer.active():
+            self.timer.reset(self.pool_refresh_delay)
+        else:
+            self.timer = reactor.callLater(self.pool_refresh_delay, self.maybe_refill_pool)
 
     def maybe_refill_pool(self):
         now = time.time()
         if self.last_fetch + self.pool_refresh_delay < now:
             self.vlog('%s refilling pool' % (self,))
             while len(self.keypool) < self.pool_size:
-                self.keypool.append(self.gen_key(self.DEFAULT_KEY_SIZE))
+                self.keypool.append(self.gen_key(self.default_key_size))
         else:
             self.vlog('%s not refilling pool' % (self,))
+            reactor.callLater(1, self.maybe_refill_pool)
 
     def gen_key(self, key_size):
         self.vlog('%s generating key size %s' % (self, key_size, ))
@@ -53,8 +65,10 @@ class KeyGenerator(foolscap.Referenceable):
 
     def remote_get_rsa_key_pair(self, key_size):
         self.vlog('%s remote_get_key' % (self,))
-        if key_size != self.DEFAULT_KEY_SIZE or not self.keypool:
-            return self.gen_key(key_size)
+        if key_size != self.default_key_size or not self.keypool:
+            key = self.gen_key(key_size)
+            self.reset_timer()
+            return key
         else:
             self.reset_timer()
             return self.keypool.pop()
@@ -62,28 +76,36 @@ class KeyGenerator(foolscap.Referenceable):
 class KeyGeneratorService(service.MultiService):
     furl_file = 'key_generator.furl'
 
-    def __init__(self):
+    def __init__(self, basedir='.', display_furl=True, default_key_size=2048):
         service.MultiService.__init__(self)
-        self.tub = foolscap.Tub(certFile='key_generator.pem')
+        self.basedir = basedir
+        fileutil.make_dirs(self.basedir)
+        self.tub = Tub(certFile=os.path.join(self.basedir, 'key_generator.pem'))
+        self.tub.setOption("expose-remote-exception-types", False)
         self.tub.setServiceParent(self)
-        self.key_generator = KeyGenerator()
+        self.key_generator = KeyGenerator(default_key_size=default_key_size)
+        self.key_generator.setServiceParent(self)
 
         portnum = self.get_portnum()
         self.listener = self.tub.listenOn(portnum or 'tcp:0')
         d = self.tub.setLocationAutomatically()
         if portnum is None:
             d.addCallback(self.save_portnum)
-        d.addCallback(self.tub_ready)
+        d.addCallback(self.tub_ready, display_furl)
         d.addErrback(log.err)
 
     def get_portnum(self):
-        if os.path.exists('portnum'):
-            return file('portnum', 'rb').read().strip()
+        portnumfile = os.path.join(self.basedir, 'portnum')
+        if os.path.exists(portnumfile):
+            return file(portnumfile, 'rb').read().strip()
 
     def save_portnum(self, junk):
         portnum = self.listener.getPortnum()
-        file('portnum', 'wb').write('%d\n' % (portnum,))
-
-    def tub_ready(self, junk):
-        self.keygen_furl = self.tub.registerReference(self.key_generator, furlFile=self.furl_file)
-        print 'key generator at:', self.keygen_furl 
+        portnumfile = os.path.join(self.basedir, 'portnum')
+        file(portnumfile, 'wb').write('%d\n' % (portnum,))
+
+    def tub_ready(self, junk, display_furl):
+        kgf = os.path.join(self.basedir, self.furl_file)
+        self.keygen_furl = self.tub.registerReference(self.key_generator, furlFile=kgf)
+        if display_furl:
+            print 'key generator at:', self.keygen_furl