Merge pull request #5100 from freedomofpress/speedy-getkey

Cache CryptoUtil.getkey (redshiftzero's idea)

securedrop/crypto_util.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 
+import collections
 from distutils.version import StrictVersion
 import pretty_bad_protocol as gnupg
 import os
@@ -54,6 +55,32 @@ def monkey_patch_delete_handle_status(self, key, value):
 gnupg._parsers.DeleteResult._handle_status = monkey_patch_delete_handle_status
 
 
+class FIFOCache():
+    """
+    We implemented this simple cache instead of using functools.lru_cache
+    (this uses a different cache replacement policy (FIFO), but either
+    FIFO or LRU works for our key fingerprint cache)
+    due to the inability to remove an item from its cache.
+
+    See: https://bugs.python.org/issue28178
+    """
+    def __init__(self, maxsize: int):
+        self.maxsize = maxsize
+        self.cache = collections.OrderedDict()  # type: collections.OrderedDict
+
+    def get(self, item):
+        if item in self.cache:
+            return self.cache[item]
+
+    def put(self, item, value):
+        self.cache[item] = value
+        if len(self.cache) > self.maxsize:
+            self.cache.popitem(last=False)
+
+    def delete(self, item):
+        del self.cache[item]
+
+
 class CryptoException(Exception):
     pass
 
@@ -72,6 +99,9 @@ class CryptoUtil:
     # to set an expiration date.
     DEFAULT_KEY_EXPIRATION_DATE = '0'
 
+    keycache_limit = 1000
+    keycache = FIFOCache(keycache_limit)
+
     def __init__(self,
                  scrypt_params,
                  scrypt_id_pepper,
@@ -224,12 +254,20 @@ def delete_reply_keypair(self, source_filesystem_id):
         temp_gpg = gnupg.GPG(binary='gpg2', homedir=self.gpg_key_dir)
         # The subkeys keyword argument deletes both secret and public keys.
         temp_gpg.delete_keys(key, secret=True, subkeys=True)
+        self.keycache.delete(source_filesystem_id)
 
     def getkey(self, name):
+        fingerprint = self.keycache.get(name)
+        if fingerprint:  # cache hit
+            return fingerprint
+
+        # cache miss
         for key in self.gpg.list_keys():
             for uid in key['uids']:
                 if name in uid:
+                    self.keycache.put(name, key['fingerprint'])
                     return key['fingerprint']
+
         return None
 
     def export_pubkey(self, name):

securedrop/journalist.py

@@ -3,10 +3,25 @@
 from sdconfig import config
 
 from journalist_app import create_app
+from models import Source
+from source_app.utils import asynchronous
 
 app = create_app(config)
 
 
+@asynchronous
+def prime_keycache():
+    """
+    Preloads CryptoUtil.keycache.
+    """
+    with app.app_context():
+        for source in Source.query.filter_by(pending=False).all():
+            app.crypto_util.getkey(source.filesystem_id)
+
+
+prime_keycache()
+
+
 if __name__ == "__main__":  # pragma: no cover
     debug = getattr(config, 'env', 'prod') != 'prod'
     app.run(debug=debug, host='0.0.0.0', port=8081)  # nosec

securedrop/tests/test_crypto_util.py

@@ -11,7 +11,7 @@
 import crypto_util
 import models
 
-from crypto_util import CryptoUtil, CryptoException
+from crypto_util import CryptoUtil, CryptoException, FIFOCache
 from db import db
 
 
@@ -343,3 +343,25 @@ def test_encrypt_then_decrypt_gives_same_result(
     decrypted_text = crypto.decrypt(secret, ciphertext)
 
     assert decrypted_text == message
+
+
+def test_fifo_cache():
+    cache = FIFOCache(3)
+
+    cache.put('item 1', 1)
+    cache.put('item 2', 2)
+    cache.put('item 3', 3)
+
+    assert cache.get('item 1') == 1
+    assert cache.get('item 2') == 2
+    assert cache.get('item 3') == 3
+
+    cache.put('item 4', 4)
+    # Maxsize is 3, so adding item 4 should kick out item 1
+    assert not cache.get('item 1')
+    assert cache.get('item 2') == 2
+    assert cache.get('item 3') == 3
+    assert cache.get('item 4') == 4
+
+    cache.delete('item 2')
+    assert not cache.get('item 2')