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
  29     def tearDown(self):
  30         d = self.parent.stopService()
  31         d.addCallback(eventual.fireEventually)
  32         d.addBoth(flush_but_dont_ignore)
  33         return d
  34
  35     def test_key_gen_service(self):
  36         def p(junk, msg):
  37             #import time
  38             #print time.asctime(), msg
  39             return junk
  40
  41         #print 'starting key generator service'
  42         kgs = key_generator.KeyGeneratorService(display_furl=False)
  43         kgs.key_generator.verbose = True
  44         kgs.setServiceParent(self.parent)
  45         kgs.key_generator.pool_size = 8
  46         keysize = kgs.key_generator.DEFAULT_KEY_SIZE
  47
  48         def keypool_full():
  49             return len(kgs.key_generator.keypool) == kgs.key_generator.pool_size
  50
  51         # first wait for key gen pool to fill up
  52         d = eventual.fireEventually()
  53         d.addCallback(p, 'waiting for pool to fill up')
  54         d.addCallback(lambda junk: self.poll(keypool_full))
  55
  56         d.addCallback(p, 'grabbing a few keys')
  57         # grab a few keys, check that pool size shrinks
  58         def get_key(junk=None):
  59             d = self.tub.getReference(kgs.keygen_furl)
  60             d.addCallback(lambda kg: kg.callRemote('get_rsa_key_pair', keysize))
  61             return d
  62
  63         def check_poolsize(junk, size):
  64             self.failUnlessEqual(len(kgs.key_generator.keypool), size)
  65
  66         n_keys_to_waste = 4
  67         for i in range(n_keys_to_waste):
  68             d.addCallback(get_key)
  69         d.addCallback(check_poolsize, kgs.key_generator.pool_size - n_keys_to_waste)
  70
  71         d.addCallback(p, 'checking a key works')
  72         # check that a retrieved key is actually useful
  73         d.addCallback(get_key)
  74         def check_key_works(keys):
  75             verifying_key, signing_key = keys
  76             v = rsa.create_verifying_key_from_string(verifying_key)
  77             s = rsa.create_signing_key_from_string(signing_key)
  78             junk = os.urandom(42)
  79             sig = s.sign(junk)
  80             self.failUnless(v.verify(junk, sig))
  81         d.addCallback(check_key_works)
  82
  83         d.addCallback(p, 'checking pool exhaustion')
  84         # exhaust the pool
  85         for i in range(kgs.key_generator.pool_size):
  86             d.addCallback(get_key)
  87         d.addCallback(check_poolsize, 0)
  88
  89         # and check it still works (will gen key synchronously on demand)
  90         d.addCallback(get_key)
  91         d.addCallback(check_key_works)
  92
  93         d.addCallback(p, 'checking pool replenishment')
  94         # check that the pool will refill
  95         d.addCallback(lambda junk: self.poll(keypool_full))
  96
  97         return d