PYTHON-2132 cache OCSP responses

Author: prashantmital

pymongo/ocsp_cache.py

@@ -0,0 +1,87 @@
+# Copyright 2020-present MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for caching OCSP responses."""
+
+from collections import namedtuple
+from datetime import datetime as _datetime
+from threading import Lock
+
+
+class _OCSPCache(object):
+    """A cache for OCSP responses."""
+    CACHE_KEY_TYPE = namedtuple('OcspResponseCacheKey',
+                                ['hash_algorithm', 'issuer_name_hash',
+                                 'issuer_key_hash', 'serial_number'])
+
+    def __init__(self):
+        self._data = {}
+        # Hold this lock when accessing _data.
+        self._lock = Lock()
+
+    def _get_cache_key(self, ocsp_request):
+        return self.CACHE_KEY_TYPE(
+            hash_algorithm=ocsp_request.hash_algorithm.name.lower(),
+            issuer_name_hash=ocsp_request.issuer_name_hash,
+            issuer_key_hash=ocsp_request.issuer_key_hash,
+            serial_number=ocsp_request.serial_number)
+
+    def __setitem__(self, key, value):
+        """Add/update a cache entry.
+
+        'key' is of type cryptography.x509.ocsp.OCSPRequest
+        'value' is of type cryptography.x509.ocsp.OCSPResponse
+
+        Validity of the OCSP response must be checked by caller.
+        """
+        with self._lock:
+            cache_key = self._get_cache_key(key)
+
+            # As per the OCSP protocol, if the response's nextUpdate field is
+            # not set, the responder is indicating that newer revocation
+            # information is available all the time.
+            if value.next_update is None:
+                self._data.pop(cache_key, None)
+                return
+
+            # Do nothing if the response is invalid.
+            if not (value.this_update <= _datetime.utcnow()
+                    < value.next_update):
+                return
+
+            # Cache new response OR update cached response if new response
+            # has longer validity.
+            cached_value = self._data.get(cache_key, None)
+            if (cached_value is None or
+                    cached_value.next_update < value.next_update):
+                self._data[cache_key] = value
+
+    def __getitem__(self, item):
+        """Get a cache entry if it exists.
+
+        'item' is of type cryptography.x509.ocsp.OCSPRequest
+
+        Raises KeyError if the item is not in the cache.
+        """
+        with self._lock:
+            cache_key = self._get_cache_key(item)
+            value = self._data[cache_key]
+
+            # Return cached response if it is still valid.
+            if (value.this_update <= _datetime.utcnow() <
+                    value.next_update):
+                return value
+
+            self._data.pop(cache_key, None)
+            raise KeyError(cache_key)

pymongo/ocsp_support.py

@@ -207,37 +207,11 @@ def _verify_response_signature(issuer, response):
     return ret
 
 
-def _request_ocsp(cert, issuer, uri):
+def _build_ocsp_request(cert, issuer):
     # https://cryptography.io/en/latest/x509/ocsp/#creating-requests
     builder = _OCSPRequestBuilder()
-    # add_certificate returns a new instance
     builder = builder.add_certificate(cert, issuer, _SHA1())
-    ocsp_request = builder.build()
-    try:
-        response = _post(
-            uri,
-            data=ocsp_request.public_bytes(_Encoding.DER),
-            headers={'Content-Type': 'application/ocsp-request'},
-            timeout=5)
-    except _RequestException:
-        _LOGGER.debug("HTTP request failed")
-        return None
-    if response.status_code != 200:
-        _LOGGER.debug("HTTP request returned %d", response.status_code)
-        return None
-    ocsp_response = _load_der_ocsp_response(response.content)
-    _LOGGER.debug(
-        "OCSP response status: %r", ocsp_response.response_status)
-    if ocsp_response.response_status != _OCSPResponseStatus.SUCCESSFUL:
-        return None
-    # RFC6960, Section 3.2, Number 1. Only relevant if we need to
-    # talk to the responder directly.
-    # Accessing response.serial_number raises if response status is not
-    # SUCCESSFUL.
-    if ocsp_response.serial_number != ocsp_request.serial_number:
-        _LOGGER.debug("Response serial number does not match request")
-        return None
-    return ocsp_response
+    return builder.build()
 
 
 def _verify_response(issuer, response):
@@ -261,6 +235,45 @@ def _verify_response(issuer, response):
     return 1
 
 
+def _get_ocsp_response(cert, issuer, uri, ocsp_response_cache):
+    ocsp_request = _build_ocsp_request(cert, issuer)
+    try:
+        ocsp_response = ocsp_response_cache[ocsp_request]
+        _LOGGER.debug("Using cached OCSP response.")
+    except KeyError:
+        try:
+            response = _post(
+                uri,
+                data=ocsp_request.public_bytes(_Encoding.DER),
+                headers={'Content-Type': 'application/ocsp-request'},
+                timeout=5)
+        except _RequestException:
+            _LOGGER.debug("HTTP request failed")
+            return None
+        if response.status_code != 200:
+            _LOGGER.debug("HTTP request returned %d", response.status_code)
+            return None
+        ocsp_response = _load_der_ocsp_response(response.content)
+        _LOGGER.debug(
+            "OCSP response status: %r", ocsp_response.response_status)
+        if ocsp_response.response_status != _OCSPResponseStatus.SUCCESSFUL:
+            return None
+        # RFC6960, Section 3.2, Number 1. Only relevant if we need to
+        # talk to the responder directly.
+        # Accessing response.serial_number raises if response status is not
+        # SUCCESSFUL.
+        if ocsp_response.serial_number != ocsp_request.serial_number:
+            _LOGGER.debug("Response serial number does not match request")
+            return None
+        if not _verify_response(issuer, ocsp_response):
+            # The response failed verification.
+            return None
+        _LOGGER.debug("Caching OCSP response.")
+        ocsp_response_cache[ocsp_request] = ocsp_response
+
+    return ocsp_response
+
+
 def _ocsp_callback(conn, ocsp_bytes, user_data):
     """Callback for use with OpenSSL.SSL.Context.set_ocsp_client_callback."""
     cert = conn.get_peer_certificate()
@@ -283,6 +296,8 @@ def _ocsp_callback(conn, ocsp_bytes, user_data):
                 _LOGGER.debug("Peer presented a must-staple cert")
                 must_staple = True
                 break
+    ocsp_response_cache = user_data.ocsp_response_cache
+
     # No stapled OCSP response
     if ocsp_bytes == b'':
         _LOGGER.debug("Peer did not staple an OCSP response")
@@ -314,13 +329,12 @@ def _ocsp_callback(conn, ocsp_bytes, user_data):
         # successful, valid responses with a certificate status of REVOKED.
         for uri in uris:
             _LOGGER.debug("Trying %s", uri)
-            response = _request_ocsp(cert, issuer, uri)
+            response = _get_ocsp_response(
+                cert, issuer, uri, ocsp_response_cache)
             if response is None:
                 # The endpoint didn't respond in time, or the response was
-                # unsuccessful or didn't match the request.
-                continue
-            if not _verify_response(issuer, response):
-                # The response failed verification.
+                # unsuccessful or didn't match the request, or the response
+                # failed verification.
                 continue
             _LOGGER.debug("OCSP cert status: %r", response.certificate_status)
             if response.certificate_status == _OCSPCertStatus.GOOD:
@@ -344,6 +358,8 @@ def _ocsp_callback(conn, ocsp_bytes, user_data):
         return 0
     if not _verify_response(issuer, response):
         return 0
+    # Cache the verified, stapled response.
+    ocsp_response_cache[_build_ocsp_request(cert, issuer)] = response
     _LOGGER.debug("OCSP cert status: %r", response.certificate_status)
     if response.certificate_status == _OCSPCertStatus.REVOKED:
         return 0

pymongo/pyopenssl_context.py

@@ -40,6 +40,7 @@
 from pymongo.ocsp_support import (
     _load_trusted_ca_certs,
     _ocsp_callback)
+from pymongo.ocsp_cache import _OCSPCache
 from pymongo.socket_checker import (
     _errno_from_exception, SocketChecker as _SocketChecker)
 
@@ -142,6 +143,7 @@ class _CallbackData(object):
     def __init__(self):
         self.trusted_ca_certs = None
         self.check_ocsp_endpoint = None
+        self.ocsp_response_cache = _OCSPCache()
 
 
 class SSLContext(object):

test/test_ocsp_cache.py

@@ -0,0 +1,133 @@
+# Copyright 2020-present MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Test the pymongo ocsp_support module."""
+
+from collections import namedtuple
+from datetime import datetime, timedelta
+from os import urandom
+import random
+import sys
+from time import sleep
+
+sys.path[0:0] = [""]
+
+from pymongo.ocsp_cache import _OCSPCache
+from test import unittest
+
+
+class TestOcspCache(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.MockHashAlgorithm = namedtuple(
+            "MockHashAlgorithm", ['name'])
+        cls.MockOcspRequest = namedtuple(
+            "MockOcspRequest", ['hash_algorithm', 'issuer_name_hash',
+                                'issuer_key_hash', 'serial_number'])
+        cls.MockOcspResponse = namedtuple(
+            "MockOcspResponse", ["this_update", "next_update"])
+
+    def setUp(self):
+        self.cache = _OCSPCache()
+
+    def _create_mock_request(self):
+        hash_algorithm = self.MockHashAlgorithm(
+            random.choice(['sha1', 'md5', 'sha256']))
+        issuer_name_hash = urandom(8)
+        issuer_key_hash = urandom(8)
+        serial_number = random.randint(0, 10**10)
+        return self.MockOcspRequest(
+            hash_algorithm=hash_algorithm,
+            issuer_name_hash=issuer_name_hash,
+            issuer_key_hash=issuer_key_hash,
+            serial_number=serial_number)
+
+    def _create_mock_response(self, this_update_delta_seconds,
+                              next_update_delta_seconds):
+        now = datetime.utcnow()
+        this_update = now + timedelta(seconds=this_update_delta_seconds)
+        if next_update_delta_seconds is not None:
+            next_update = now + timedelta(seconds=next_update_delta_seconds)
+        else:
+            next_update = None
+        return self.MockOcspResponse(
+            this_update=this_update,
+            next_update=next_update)
+
+    def _add_mock_cache_entry(self, mock_request, mock_response):
+        key = self.cache._get_cache_key(mock_request)
+        self.cache._data[key] = mock_response
+
+    def test_simple(self):
+        # Start with 1 valid entry in the cache.
+        request = self._create_mock_request()
+        response = self._create_mock_response(-10, +3600)
+        self._add_mock_cache_entry(request, response)
+
+        # Ensure entry can be retrieved.
+        self.assertEqual(self.cache[request], response)
+
+        # Valid entries with an earlier next_update have no effect.
+        response_1 = self._create_mock_response(-20, +1800)
+        self.cache[request] = response_1
+        self.assertEqual(self.cache[request], response)
+
+        # Invalid entries with a later this_update have no effect.
+        response_2 = self._create_mock_response(+20, +1800)
+        self.cache[request] = response_2
+        self.assertEqual(self.cache[request], response)
+
+        # Invalid entries with passed next_update have no effect.
+        response_3 = self._create_mock_response(-10, -5)
+        self.cache[request] = response_3
+        self.assertEqual(self.cache[request], response)
+
+        # Valid entries with a later next_update update the cache.
+        response_new = self._create_mock_response(-5, +7200)
+        self.cache[request] = response_new
+        self.assertEqual(self.cache[request], response_new)
+
+        # Entries with an unset next_update purge the cache.
+        response_notset = self._create_mock_response(-5, None)
+        self.cache[request] = response_notset
+        with self.assertRaises(KeyError):
+            _ = self.cache[request]
+
+    def test_invalidate(self):
+        # Start with 1 valid entry in the cache.
+        request = self._create_mock_request()
+        response = self._create_mock_response(-10, +0.25)
+        self._add_mock_cache_entry(request, response)
+
+        # Ensure entry can be retrieved.
+        self.assertEqual(self.cache[request], response)
+
+        # Wait for entry to become invalid and ensure KeyError is raised.
+        sleep(0.5)
+        with self.assertRaises(KeyError):
+            _ = self.cache[request]
+
+    def test_non_existent(self):
+        # Start with 1 valid entry in the cache.
+        request = self._create_mock_request()
+        response = self._create_mock_response(-10, +10)
+        self._add_mock_cache_entry(request, response)
+
+        # Attempt to retrieve non-existent entry must raise KeyError.
+        with self.assertRaises(KeyError):
+            _ = self.cache[self._create_mock_request()]
+
+
+if __name__ == "__main__":
+    unittest.main()