src/allmydata/test/test_keygen.py

   1
   2 import os
   3 from twisted.trial import unittest
   4 from twisted.application import service
   5
   6 from foolscap import Tub, eventual
   7
   8 from allmydata import key_generator
   9 from allmydata.util import testutil
  10 from pycryptopp.publickey import rsa
  11
  12 def flush_but_dont_ignore(res):
  13     d = eventual.flushEventualQueue()
  14     def _done(ignored):
  15         return res
  16     d.addCallback(_done)
  17     return d
  18
  19 class KeyGenService(unittest.TestCase, testutil.PollMixin):
  20     def setUp(self):
  21         self.parent = service.MultiService()
  22         self.parent.startService()
  23
  24         self.tub = t = Tub()
  25         t.setServiceParent(self.parent)
  26         t.listenOn("tcp:0")
  27         t.setLocationAutomatically()
  28         return eventual.fireEventually()
  29
  30     def tearDown(self):
  31         d = self.parent.stopService()
  32         d.addCallback(eventual.fireEventually)
  33         d.addBoth(flush_but_dont_ignore)
  34         return d
  35
  36     def test_key_gen_service(self):
  37         def p(junk, msg):
  38             #import time
  39             #print time.asctime(), msg
  40             return junk
  41
  42         #print 'starting key generator service'
  43         keysize = 522
  44         kgs = key_generator.KeyGeneratorService(display_furl=False, default_key_size=keysize)
  45         kgs.key_generator.verbose = True
  46         kgs.setServiceParent(self.parent)
  47         kgs.key_generator.pool_size = 8
  48
  49         def keypool_full():
  50             return len(kgs.key_generator.keypool) == kgs.key_generator.pool_size
  51
  52         # first wait for key gen pool to fill up
  53         d = eventual.fireEventually()
  54         d.addCallback(p, 'waiting for pool to fill up')
  55         d.addCallback(lambda junk: self.poll(keypool_full))
  56
  57         d.addCallback(p, 'grabbing a few keys')
  58         # grab a few keys, check that pool size shrinks
  59         def get_key(junk=None):
  60             d = self.tub.getReference(kgs.keygen_furl)
  61             d.addCallback(lambda kg: kg.callRemote('get_rsa_key_pair', keysize))
  62             return d
  63
  64         def check_poolsize(junk, size):
  65             self.failUnlessEqual(len(kgs.key_generator.keypool), size)
  66
  67         n_keys_to_waste = 4
  68         for i in range(n_keys_to_waste):
  69             d.addCallback(get_key)
  70         d.addCallback(check_poolsize, kgs.key_generator.pool_size - n_keys_to_waste)
  71
  72         d.addCallback(p, 'checking a key works')
  73         # check that a retrieved key is actually useful
  74         d.addCallback(get_key)
  75         def check_key_works(keys):
  76             verifying_key, signing_key = keys
  77             v = rsa.create_verifying_key_from_string(verifying_key)
  78             s = rsa.create_signing_key_from_string(signing_key)
  79             junk = os.urandom(42)
  80             sig = s.sign(junk)
  81             self.failUnless(v.verify(junk, sig))
  82         d.addCallback(check_key_works)
  83
  84         d.addCallback(p, 'checking pool exhaustion')
  85         # exhaust the pool
  86         for i in range(kgs.key_generator.pool_size):
  87             d.addCallback(get_key)
  88         d.addCallback(check_poolsize, 0)
  89
  90         # and check it still works (will gen key synchronously on demand)
  91         d.addCallback(get_key)
  92         d.addCallback(check_key_works)
  93
  94         d.addCallback(p, 'checking pool replenishment')
  95         # check that the pool will refill
  96         d.addCallback(lambda junk: self.poll(keypool_full))
  97
  98         return d