Merge pull request #3316 from amadeuppereira/fix-ssl-context-performance-issue

Support requests>=2.32.5, reimplement the fix previously there for only loading ssl certificates once

Author: cyberw

locust/clients.py

@@ -2,6 +2,7 @@
 
 from locust.event import EventHook
 
+import os
 import re
 import sys
 import time
@@ -10,11 +11,15 @@
 from urllib.parse import urlparse, urlunparse
 
 import requests
+from packaging.version import Version
 from requests import Response
 from requests.adapters import HTTPAdapter
 from requests.auth import HTTPBasicAuth
+from requests.compat import basestring
 from requests.exceptions import InvalidSchema, InvalidURL, MissingSchema, RequestException
+from requests.utils import DEFAULT_CA_BUNDLE_PATH, extract_zipped_paths
 from urllib3 import PoolManager
+from urllib3.util import create_urllib3_context
 
 from .exception import CatchResponseError, LocustError, ResponseError
 
@@ -23,6 +28,7 @@
 else:
     from typing_extensions import override
 
+requests_version = Version(requests.__version__).release
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Generator, Iterable, Mapping, MutableMapping
@@ -60,6 +66,10 @@ class RESTKwargs(RequestKwargs, total=False):
 
 absolute_http_url_regexp = re.compile(r"^https?://", re.IGNORECASE)
 
+if requests_version >= (2, 32, 5):
+    _preloaded_ssl_context = create_urllib3_context()
+    _preloaded_ssl_context.load_verify_locations(extract_zipped_paths(DEFAULT_CA_BUNDLE_PATH))
+
 
 class HttpSession(requests.Session):
     """
@@ -482,6 +492,52 @@ def init_poolmanager(self, *args, **kwargs):
         if self.poolmanager is None:
             super().init_poolmanager(*args, **kwargs)
 
+    # In python requests version 2.32.5 they reverted
+    # https://github.com/psf/requests/pull/6667
+    # Without this change the root CA certificates are loaded on every request
+    # We re-implement this change to increase the performance
+    def cert_verify(self, conn, url, verify, cert):
+        if requests_version < (2, 32, 5):
+            return super().cert_verify(conn, url, verify, cert)
+
+        if url.lower().startswith("https") and verify:
+            conn.cert_reqs = "CERT_REQUIRED"
+
+            if verify is not True:
+                cert_loc = verify
+
+                if not os.path.exists(cert_loc):
+                    raise OSError(f"Could not find a suitable TLS CA certificate bundle, invalid path: {cert_loc}")
+
+                if not os.path.isdir(cert_loc):
+                    conn.ca_certs = cert_loc
+                else:
+                    conn.ca_cert_dir = cert_loc
+        else:
+            conn.cert_reqs = "CERT_NONE"
+            conn.ca_certs = None
+            conn.ca_cert_dir = None
+
+        if cert:
+            if not isinstance(cert, basestring):
+                conn.cert_file = cert[0]
+                conn.key_file = cert[1]
+            else:
+                conn.cert_file = cert
+                conn.key_file = None
+            if conn.cert_file and not os.path.exists(conn.cert_file):
+                raise OSError(f"Could not find the TLS certificate file, invalid path: {conn.cert_file}")
+            if conn.key_file and not os.path.exists(conn.key_file):
+                raise OSError(f"Could not find the TLS key file, invalid path: {conn.key_file}")
+
+    def build_connection_pool_key_attributes(self, request, verify, cert=None):
+        host_params, pool_kwargs = super().build_connection_pool_key_attributes(request, verify, cert)
+
+        if requests_version >= (2, 32, 5) and verify is True:
+            pool_kwargs["ssl_context"] = _preloaded_ssl_context
+
+        return host_params, pool_kwargs
+
 
 # Monkey patch Response class to give some guidance
 def _missing_catch_response_True(self, *_args, **_kwargs):

locust/test/test_http.py

@@ -4,10 +4,14 @@
 
 import time
 
+import urllib3
 from requests.exceptions import InvalidSchema, InvalidURL, MissingSchema, RequestException
+from urllib3.exceptions import SSLError
 
 from .testcases import WebserverTestCase
 
+urllib3.disable_warnings()
+
 
 class TestHttpSession(WebserverTestCase):
     def get_client(self, base_url=None):
@@ -315,3 +319,14 @@ def on_request(**kw):
         self.assertEqual("GET", kwargs["response"].text)
         user.client.request("get", "/request_method", context={"user": "bar"})  # override User context
         self.assertDictEqual({"user": "bar"}, kwargs["context"])
+
+    def test_verify_true_fails_with_bad_cert(self):
+        s = self.get_client("https://expired.badssl.com")
+        r = s.get("/", verify=True)
+        self.assertTrue("exception" in r.request_meta)
+        self.assertIsInstance(r.request_meta["exception"], SSLError)
+
+    def test_verify_false_succeeds_with_bad_cert(self):
+        s = self.get_client("https://expired.badssl.com")
+        r = s.get("/", verify=False)
+        self.assertEqual(r.status_code, 200)

pyproject.toml

@@ -35,7 +35,7 @@ classifiers = [
 dependencies = [
     "flask>=2.0.0",
     "Werkzeug>=2.0.0",
-    "requests>=2.32.2,<2.32.5",
+    "requests>=2.32.2",
     "msgpack>=1.0.0",
     "pyzmq>=25.0.0",
     "geventhttpclient>=2.3.1",