Sem-Ver: api-break Optimise the standard requests session based public key retrieval by obtaining proxy information from the environment once per a public key server and setting session.trust_env to False.

As a result of this change retrieving the same public key, with caching enabled, 10,000 times takes ~ 6 seconds instead of ~ 14 seconds.

Signed-off-by: David Black <dblack@atlassian.com>

Author: dbaxa

atlassian_jwt_auth/key.py

@@ -9,6 +9,7 @@
 import cryptography.hazmat.backends
 import jwt
 import requests
+import requests.utils
 from cryptography.hazmat.primitives import serialization
 from requests.exceptions import RequestException, ConnectionError
 
@@ -84,18 +85,21 @@ def __init__(self, base_url):
             base_url += '/'
         self.base_url = base_url
         self._session = self._get_session()
+        self._proxies = requests.utils.get_environ_proxies(self.base_url)
 
     def _get_session(self):
         if HTTPSPublicKeyRetriever._class_session is None:
             session = cachecontrol.CacheControl(requests.Session())
+            session.trust_env = False
             HTTPSPublicKeyRetriever._class_session = session
         return HTTPSPublicKeyRetriever._class_session
 
     def retrieve(self, key_identifier, **requests_kwargs):
         """ returns the public key for given key_identifier. """
         if not isinstance(key_identifier, KeyIdentifier):
             key_identifier = KeyIdentifier(key_identifier)
-
+        if self._proxies and 'proxies' not in requests_kwargs:
+            requests_kwargs['proxies'] = self._proxies
         url = self.base_url + key_identifier.key_id
         try:
             return self._retrieve(url, requests_kwargs)

atlassian_jwt_auth/tests/test_public_key_provider.py

@@ -1,3 +1,4 @@
+import os
 import re
 import unittest
 
@@ -8,10 +9,26 @@
 from atlassian_jwt_auth.key import (
     HTTPSPublicKeyRetriever,
     HTTPSMultiRepositoryPublicKeyRetriever,
+    PEM_FILE_TYPE,
 )
 from atlassian_jwt_auth.tests import utils
 
 
+def get_expected_and_os_proxies_dict(proxy_location):
+    """ returns expected proxy & environmental
+        proxy dictionary based upon the provided proxy location.
+    """
+    expected_proxies = {
+        'http': proxy_location,
+        'https': proxy_location,
+    }
+    os_proxy_dict = {
+        'HTTP_PROXY': proxy_location,
+        'HTTPS_PROXY': proxy_location
+    }
+    return expected_proxies, os_proxy_dict
+
+
 class BaseHTTPSPublicKeyRetrieverTest(object):
     """ tests for the HTTPSPublicKeyRetriever class. """
 
@@ -39,6 +56,21 @@ def test_https_public_key_retriever_does_not_support_none_url(self):
         with self.assertRaises(ValueError):
             self.create_retriever(None)
 
+    def test_https_public_key_retriever_session_uses_env_proxy(self):
+        """ tests that the underlying session makes use of environmental
+            proxy configured.
+        """
+        proxy_location = 'https://example.proxy'
+        expected_proxies, proxy_dict = get_expected_and_os_proxies_dict(
+            proxy_location)
+        with mock.patch.dict(os.environ, proxy_dict, clear=True):
+            retriever = self.create_retriever(self.base_url)
+            key_retrievers = [retriever]
+            if isinstance(retriever, HTTPSMultiRepositoryPublicKeyRetriever):
+                key_retrievers = retriever._retrievers
+            for key_retriever in key_retrievers:
+                self.assertEqual(key_retriever._proxies, expected_proxies)
+
     def test_https_public_key_retriever_supports_https_url(self):
         """ tests that HTTPSPublicKeyRetriever supports https://
             base urls.
@@ -55,6 +87,48 @@ def test_retrieve(self, mock_get_method):
             retriever.retrieve('example/eg'),
             self._public_key_pem)
 
+    @mock.patch.object(requests.Session, 'get')
+    def test_retrieve_with_proxy(self, mock_get_method):
+        """ tests that the retrieve method works as expected when a proxy
+            should be used.
+        """
+        proxy_location = 'https://example.proxy'
+        key_id = 'example/eg'
+        expected_proxies, proxy_dict = get_expected_and_os_proxies_dict(
+            proxy_location)
+        _setup_mock_response_for_retriever(
+            mock_get_method, self._public_key_pem)
+        with mock.patch.dict(os.environ, proxy_dict, clear=True):
+            retriever = self.create_retriever(self.base_url)
+            retriever.retrieve(key_id)
+            mock_get_method.assert_called_once_with(
+                '%s/%s' % (self.base_url, key_id),
+                headers={'accept': PEM_FILE_TYPE},
+                proxies=expected_proxies
+            )
+
+    @mock.patch.object(requests.Session, 'get')
+    def test_retrieve_with_proxy_explicitly_set(self, mock_get_method):
+        """ tests that the retrieve method works as expected when a proxy
+            should be used and has been explicitly provided.
+        """
+        proxy_location = 'https://example.proxy'
+        explicit_proxy_location = 'https://explicit.proxy'
+        key_id = 'example/eg'
+        _, proxy_dict = get_expected_and_os_proxies_dict(proxy_location)
+        expected_proxies, _ = get_expected_and_os_proxies_dict(
+            explicit_proxy_location)
+        _setup_mock_response_for_retriever(
+            mock_get_method, self._public_key_pem)
+        with mock.patch.dict(os.environ, proxy_dict, clear=True):
+            retriever = self.create_retriever(self.base_url)
+            retriever.retrieve(key_id, proxies=expected_proxies)
+            mock_get_method.assert_called_once_with(
+                '%s/%s' % (self.base_url, key_id),
+                headers={'accept': PEM_FILE_TYPE},
+                proxies=expected_proxies
+            )
+
     @mock.patch.object(requests.Session, 'get')
     def test_retrieve_with_charset_in_content_type_h(self, mock_get_method):
         """ tests that the retrieve method works expected when there is