From 5578559b8566b7dc67439449bb503c34c38324ff Mon Sep 17 00:00:00 2001
From: robk-tahoe <robk-tahoe@allmydata.com>
Date: Tue, 1 Apr 2008 18:45:13 -0700
Subject: [PATCH] added offloaded key generation

this adds a new service to pre-generate RSA key pairs.  This allows
the expensive (i.e. slow) key generation to be placed into a process
outside the node, so that the node's reactor will not block when it
needs a key pair, but instead can retrieve them from a pool of already
generated key pairs in the key-generator service.

it adds a tahoe create-key-generator command which initialises an
empty dir with a tahoe-key-generator.tac file which can then be run
via twistd.  it stashes its .pem and portnum for furl stability and
writes the furl of the key gen service to key_generator.furl, also
printing it to stdout.

by placing a key_generator.furl file into the nodes config directory
(e.g. ~/.tahoe) a node will attempt to connect to such a service, and
will use that when creating mutable files (i.e. directories) whenever
possible.  if the keygen service is unavailable, it will perform the
key generation locally instead, as before.
---
 src/allmydata/client.py            | 38 ++++++++++++-
 src/allmydata/dirnode.py           |  4 +-
 src/allmydata/interfaces.py        | 12 ++++
 src/allmydata/key_generator.py     | 89 ++++++++++++++++++++++++++++++
 src/allmydata/mutable.py           | 17 +++---
 src/allmydata/scripts/keygen.py    | 52 +++++++++++++++++
 src/allmydata/scripts/runner.py    | 10 +++-
 src/allmydata/test/common.py       |  2 +-
 src/allmydata/test/test_mutable.py |  2 +-
 9 files changed, 211 insertions(+), 15 deletions(-)
 create mode 100644 src/allmydata/key_generator.py
 create mode 100644 src/allmydata/scripts/keygen.py

diff --git a/src/allmydata/client.py b/src/allmydata/client.py
index 94880ec2..9a9f0960 100644
--- a/src/allmydata/client.py
+++ b/src/allmydata/client.py
@@ -8,6 +8,7 @@ from twisted.internet import reactor
 from twisted.application.internet import TimerService
 from foolscap import Referenceable
 from foolscap.logging import log
+from pycryptopp.publickey import rsa
 
 import allmydata
 from allmydata.storage import StorageServer
@@ -72,6 +73,10 @@ class Client(node.Node, testutil.PollMixin):
         if run_helper:
             self.init_helper()
         self.init_client()
+        self._key_generator = None
+        key_gen_furl = self.get_config('key_generator.furl')
+        if key_gen_furl:
+            self.init_key_gen(key_gen_furl)
         # ControlServer and Helper are attached after Tub startup
 
         hotline_file = os.path.join(self.basedir,
@@ -196,6 +201,20 @@ class Client(node.Node, testutil.PollMixin):
         d.addCallback(_publish)
         d.addErrback(log.err, facility="tahoe.init", level=log.BAD)
 
+    def init_key_gen(self, key_gen_furl):
+        d = self.when_tub_ready()
+        def _subscribe(self):
+            self.tub.connectTo(key_gen_furl, self._got_key_generator)
+        d.addCallback(_subscribe)
+        d.addErrback(log.err, facility="tahoe.init", level=log.BAD)
+
+    def _got_key_generator(self, key_generator):
+        self._key_generator = key_generator
+        key_generator.notifyOnDisconnect(self._lost_key_generator)
+
+    def _lost_key_generator(self):
+        self._key_generator = None
+
     def init_web(self, webport):
         self.log("init_web(webport=%s)", args=(webport,))
 
@@ -281,16 +300,31 @@ class Client(node.Node, testutil.PollMixin):
 
     def create_empty_dirnode(self):
         n = NewDirectoryNode(self)
-        d = n.create()
+        d = n.create(self._generate_pubprivkeys)
         d.addCallback(lambda res: n)
         return d
 
     def create_mutable_file(self, contents=""):
         n = MutableFileNode(self)
-        d = n.create(contents)
+        d = n.create(contents, self._generate_pubprivkeys)
         d.addCallback(lambda res: n)
         return d
 
+    def _generate_pubprivkeys(self, key_size):
+        if self._key_generator:
+            d = self._key_generator.callRemote('get_rsa_key_pair', key_size)
+            def make_key_objs((verifying_key, signing_key)):
+                v = rsa.create_verifying_key_from_string(verifying_key)
+                s = rsa.create_signing_key_from_string(signing_key)
+                return v, s
+            d.addCallback(make_key_objs)
+            return d
+        else:
+            # RSA key generation for a 2048 bit key takes between 0.8 and 3.2 secs
+            signer = rsa.generate(key_size)
+            verifier = signer.get_verifying_key()
+            return verifier, signer
+
     def upload(self, uploadable):
         uploader = self.getServiceNamed("uploader")
         return uploader.upload(uploadable)
diff --git a/src/allmydata/dirnode.py b/src/allmydata/dirnode.py
index 1bc29496..35b53a84 100644
--- a/src/allmydata/dirnode.py
+++ b/src/allmydata/dirnode.py
@@ -53,7 +53,7 @@ class NewDirectoryNode:
         self._node.init_from_uri(self._uri.get_filenode_uri())
         return self
 
-    def create(self):
+    def create(self, keypair_generator=None):
         """
         Returns a deferred that eventually fires with self once the directory
         has been created (distributed across a set of storage servers).
@@ -62,7 +62,7 @@ class NewDirectoryNode:
         # URI to create our own.
         self._node = self.filenode_class(self._client)
         empty_contents = self._pack_contents({})
-        d = self._node.create(empty_contents)
+        d = self._node.create(empty_contents, keypair_generator)
         d.addCallback(self._filenode_created)
         return d
     def _filenode_created(self, res):
diff --git a/src/allmydata/interfaces.py b/src/allmydata/interfaces.py
index 190288a0..2c797e9b 100644
--- a/src/allmydata/interfaces.py
+++ b/src/allmydata/interfaces.py
@@ -1714,3 +1714,15 @@ class IStatsProducer(Interface):
         to be monitored, and numeric values.
         """
 
+class RIKeyGenerator(RemoteInterface):
+    __remote_name__ = "RIKeyGenerator.tahoe.allmydata.com"
+    """
+    Provides a service offering to make RSA key pairs.
+    """
+
+    def get_rsa_key_pair(key_size=int):
+        """
+        @param key_size: the size of the signature key.
+        @return: tuple(verifying_key, signing_key)
+        """
+        return TupleOf(str, str)
diff --git a/src/allmydata/key_generator.py b/src/allmydata/key_generator.py
new file mode 100644
index 00000000..aa735c4c
--- /dev/null
+++ b/src/allmydata/key_generator.py
@@ -0,0 +1,89 @@
+
+import os
+import time
+
+import foolscap
+from foolscap.eventual import eventually
+from zope.interface import implements
+from twisted.internet import reactor
+from twisted.application import service
+from twisted.python import log
+
+from pycryptopp.publickey import rsa
+from allmydata.interfaces import RIKeyGenerator
+
+class KeyGenerator(foolscap.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):
+        self.keypool = []
+        self.last_fetch = 0
+        eventually(self.maybe_refill_pool)
+
+    def __repr__(self):
+        return '<KeyGenerator[%s]>' % (len(self.keypool),)
+
+    def vlog(self, msg):
+        if self.verbose:
+            log.msg(msg)
+
+    def reset_timer(self):
+        self.last_fetch = time.time()
+        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))
+        else:
+            self.vlog('%s not refilling pool' % (self,))
+
+    def gen_key(self, key_size):
+        self.vlog('%s generating key size %s' % (self, key_size, ))
+        signer = rsa.generate(key_size)
+        verifier = signer.get_verifying_key()
+        return verifier.serialize(), signer.serialize()
+
+    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)
+        else:
+            self.reset_timer()
+            return self.keypool.pop()
+
+class KeyGeneratorService(service.MultiService):
+    furl_file = 'key_generator.furl'
+
+    def __init__(self):
+        service.MultiService.__init__(self)
+        self.tub = foolscap.Tub(certFile='key_generator.pem')
+        self.tub.setServiceParent(self)
+        self.key_generator = KeyGenerator()
+
+        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.addErrback(log.err)
+
+    def get_portnum(self):
+        if os.path.exists('portnum'):
+            return file('portnum', '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 
diff --git a/src/allmydata/mutable.py b/src/allmydata/mutable.py
index 981cb86f..3555eb5d 100644
--- a/src/allmydata/mutable.py
+++ b/src/allmydata/mutable.py
@@ -1647,7 +1647,7 @@ class MutableFileNode:
         self._encprivkey = None
         return self
 
-    def create(self, initial_contents):
+    def create(self, initial_contents, keypair_generator=None):
         """Call this when the filenode is first created. This will generate
         the keys, generate the initial shares, wait until at least numpeers
         are connected, allocate shares, and upload the initial
@@ -1656,7 +1656,7 @@ class MutableFileNode:
         """
         self._required_shares, self._total_shares = self.DEFAULT_ENCODING
 
-        d = defer.maybeDeferred(self._generate_pubprivkeys)
+        d = defer.maybeDeferred(self._generate_pubprivkeys, keypair_generator)
         def _generated( (pubkey, privkey) ):
             self._pubkey, self._privkey = pubkey, privkey
             pubkey_s = self._pubkey.serialize()
@@ -1675,11 +1675,14 @@ class MutableFileNode:
         d.addCallback(_generated)
         return d
 
-    def _generate_pubprivkeys(self):
-        # RSA key generation for a 2048 bit key takes between 0.8 and 3.2 secs
-        signer = rsa.generate(self.SIGNATURE_KEY_SIZE)
-        verifier = signer.get_verifying_key()
-        return verifier, signer
+    def _generate_pubprivkeys(self, keypair_generator):
+        if keypair_generator:
+            return keypair_generator(self.SIGNATURE_KEY_SIZE)
+        else:
+            # RSA key generation for a 2048 bit key takes between 0.8 and 3.2 secs
+            signer = rsa.generate(self.SIGNATURE_KEY_SIZE)
+            verifier = signer.get_verifying_key()
+            return verifier, signer
 
     def _publish(self, initial_contents):
         p = self.publish_class(self)
diff --git a/src/allmydata/scripts/keygen.py b/src/allmydata/scripts/keygen.py
new file mode 100644
index 00000000..5798bc68
--- /dev/null
+++ b/src/allmydata/scripts/keygen.py
@@ -0,0 +1,52 @@
+
+import os, sys
+from twisted.python import usage
+#from allmydata.scripts.common import BasedirMixin, NoDefaultBasedirMixin
+
+class CreateKeyGeneratorOptions(usage.Options):
+    optParameters = [
+        ["basedir", "C", None, "which directory to create the client in"],
+        ]
+
+keygen_tac = """
+# -*- python -*-
+
+from allmydata import key_generator
+from twisted.application import service
+
+k = key_generator.KeyGeneratorService(verbose=False)
+#k.key_generator.verbose = False
+#k.key_generator.DEFAULT_KEY_SIZE = 2048
+#k.key_generator.pool_size = 16
+#k.key_generator.pool_refresh_delay = 6
+
+application = service.Application("allmydata_key_generator")
+k.setServiceParent(application)
+"""
+
+def create_key_generator(config, out=sys.stdout, err=sys.stderr):
+    basedir = config['basedir']
+    if not basedir:
+        print >>err, "a basedir was not provided, please use --basedir or -C"
+        return -1
+    if os.path.exists(basedir):
+        if os.listdir(basedir):
+            print >>err, "The base directory \"%s\", which is \"%s\" is not empty." % (basedir, os.path.abspath(basedir))
+            print >>err, "To avoid clobbering anything, I am going to quit now."
+            print >>err, "Please use a different directory, or empty this one."
+            return -1
+        # we're willing to use an empty directory
+    else:
+        os.mkdir(basedir)
+    f = open(os.path.join(basedir, "tahoe-key-generator.tac"), "wb")
+    f.write(keygen_tac)
+    f.close()
+
+subCommands = [
+    ["create-key-generator", None, CreateKeyGeneratorOptions, "Create a key generator service."],
+]
+
+dispatch = {
+    "create-key-generator": create_key_generator,
+    }
+
diff --git a/src/allmydata/scripts/runner.py b/src/allmydata/scripts/runner.py
index a77810b0..eb8bbae8 100644
--- a/src/allmydata/scripts/runner.py
+++ b/src/allmydata/scripts/runner.py
@@ -4,9 +4,13 @@ from cStringIO import StringIO
 from twisted.python import usage
 
 from allmydata.scripts.common import BaseOptions
-import debug, create_node, startstop_node, cli
+import debug, create_node, startstop_node, cli, keygen
 
-_general_commands = create_node.subCommands + debug.subCommands + cli.subCommands
+_general_commands = ( create_node.subCommands
+                    + keygen.subCommands
+                    + debug.subCommands
+                    + cli.subCommands
+                    )
 
 class Options(BaseOptions, usage.Options):
     synopsis = "Usage:  tahoe <command> [command options]"
@@ -60,6 +64,8 @@ def runner(argv, run_by_human=True, stdout=sys.stdout, stderr=sys.stderr,
         rc = debug.dispatch[command](so, stdout, stderr)
     elif command in cli.dispatch:
         rc = cli.dispatch[command](so, stdout, stderr)
+    elif command in keygen.dispatch:
+        rc = keygen.dispatch[command](so, stdout, stderr)
     elif command in ac_dispatch:
         rc = ac_dispatch[command](so, stdout, stderr)
     else:
diff --git a/src/allmydata/test/common.py b/src/allmydata/test/common.py
index e272d49e..3a52dc5d 100644
--- a/src/allmydata/test/common.py
+++ b/src/allmydata/test/common.py
@@ -75,7 +75,7 @@ class FakeMutableFileNode:
         self.client = client
         self.my_uri = make_mutable_file_uri()
         self.storage_index = self.my_uri.storage_index
-    def create(self, initial_contents):
+    def create(self, initial_contents, key_generator=None):
         self.all_contents[self.storage_index] = initial_contents
         return defer.succeed(self)
     def init_from_uri(self, myuri):
diff --git a/src/allmydata/test/test_mutable.py b/src/allmydata/test/test_mutable.py
index 2c75ce32..2880b54e 100644
--- a/src/allmydata/test/test_mutable.py
+++ b/src/allmydata/test/test_mutable.py
@@ -33,7 +33,7 @@ class FakeFilenode(mutable.MutableFileNode):
     def init_from_uri(self, myuri):
         mutable.MutableFileNode.init_from_uri(self, myuri)
         return self
-    def _generate_pubprivkeys(self):
+    def _generate_pubprivkeys(self, key_size):
         count = self.counter.next()
         return FakePubKey(count), FakePrivKey(count)
     def _publish(self, initial_contents):
-- 
2.45.2