feat(ssl): support custom SSLContext

redis/asyncio/client.py

@@ -97,8 +97,9 @@
 )
 
 if TYPE_CHECKING and SSL_AVAILABLE:
-    from ssl import TLSVersion, VerifyFlags, VerifyMode
+    from ssl import SSLContext, TLSVersion, VerifyFlags, VerifyMode
 else:
+    SSLContext = None
     TLSVersion = None
     VerifyMode = None
     VerifyFlags = None
@@ -293,6 +294,7 @@ def __init__(
         ssl_min_version: "TLSVersion | None" = None,
         ssl_ciphers: str | None = None,
         ssl_password: str | None = None,
+        ssl_context: "SSLContext | None" = None,
         max_connections: int | None = None,
         single_connection_client: bool = False,
         health_check_interval: int = 0,
@@ -449,6 +451,7 @@ def __init__(
                             "ssl_min_version": ssl_min_version,
                             "ssl_ciphers": ssl_ciphers,
                             "ssl_password": ssl_password,
+                            "ssl_context": ssl_context,
                         }
                     )
             maint_notifications_enabled = (

redis/asyncio/cluster.py

@@ -133,8 +133,9 @@
 )
 
 if SSL_AVAILABLE:
-    from ssl import TLSVersion, VerifyFlags, VerifyMode
+    from ssl import SSLContext, TLSVersion, VerifyFlags, VerifyMode
 else:
+    SSLContext = None
     TLSVersion = None
     VerifyMode = None
     VerifyFlags = None
@@ -443,6 +444,7 @@ def __init__(
         ssl_keyfile: str | None = None,
         ssl_min_version: "TLSVersion | None" = None,
         ssl_ciphers: str | None = None,
+        ssl_context: "SSLContext | None" = None,
         protocol: int | None = None,
         legacy_responses: bool = True,
         address_remap: Callable[[Tuple[str, int]], Tuple[str, int]] | None = None,
@@ -510,6 +512,7 @@ def __init__(
                     "ssl_keyfile": ssl_keyfile,
                     "ssl_min_version": ssl_min_version,
                     "ssl_ciphers": ssl_ciphers,
+                    "ssl_context": ssl_context,
                 }
             )
 

redis/asyncio/connection.py

@@ -1507,6 +1507,7 @@ def __init__(
         ssl_min_version: Optional[TLSVersion] = None,
         ssl_ciphers: Optional[str] = None,
         ssl_password: Optional[str] = None,
+        ssl_context: Optional[SSLContext] = None,
         **kwargs,
     ):
         if not SSL_AVAILABLE:
@@ -1525,6 +1526,7 @@ def __init__(
             min_version=ssl_min_version,
             ciphers=ssl_ciphers,
             password=ssl_password,
+            context=ssl_context,
         )
         super().__init__(**kwargs)
 
@@ -1601,6 +1603,7 @@ def __init__(
         min_version: Optional[TLSVersion] = None,
         ciphers: Optional[str] = None,
         password: Optional[str] = None,
+        context: Optional[SSLContext] = None,
     ):
         if not SSL_AVAILABLE:
             raise RedisError("Python wasn't built with SSL support")
@@ -1632,10 +1635,10 @@ def __init__(
         self.min_version = min_version
         self.ciphers = ciphers
         self.password = password
-        self.context: Optional[SSLContext] = None
+        self.context: Optional[SSLContext] = context
 
     def get(self) -> SSLContext:
-        if not self.context:
+        if self.context is None:
             context = ssl.create_default_context()
             context.check_hostname = self.check_hostname
             context.verify_mode = self.cert_reqs

redis/client.py

@@ -315,6 +315,7 @@ def __init__(
         maint_notifications_config: MaintNotificationsConfig | None = None,
         oss_cluster_maint_notifications_handler: OSSMaintNotificationsHandler
         | None = None,
+        ssl_context: "ssl.SSLContext | None" = None,
     ) -> None:
         """
         Initialize a new Redis client.
@@ -464,6 +465,7 @@ def __init__(
                             "ssl_ocsp_expected_cert": ssl_ocsp_expected_cert,
                             "ssl_min_version": ssl_min_version,
                             "ssl_ciphers": ssl_ciphers,
+                            "ssl_context": ssl_context,
                         }
                     )
                 if (cache_config or cache) and check_protocol_version(protocol, 3):

redis/cluster.py

@@ -302,6 +302,7 @@ def parse_cluster_myshardid(resp, **options):
     "ssl_exclude_verify_flags",
     "ssl_keyfile",
     "ssl_password",
+    "ssl_context",
     "ssl_check_hostname",
     "unix_socket_path",
     "username",

redis/connection.py

@@ -2035,6 +2035,7 @@ def __init__(
         ssl_ocsp_expected_cert=None,
         ssl_min_version=None,
         ssl_ciphers=None,
+        ssl_context=None,
         **kwargs,
     ):
         """Constructor
@@ -2058,6 +2059,9 @@ def __init__(
             ssl_ocsp_expected_cert: A PEM armoured string containing the expected certificate to be returned from the ocsp verification service.
             ssl_min_version: The lowest supported SSL version. It affects the supported SSL versions of the SSLContext. None leaves the default provided by ssl module.
             ssl_ciphers: A string listing the ciphers that are allowed to be used. Defaults to None, which means that the default ciphers are used. See https://docs.python.org/3/library/ssl.html#ssl.SSLContext.set_ciphers for more information.
+            ssl_context: A pre-configured ``ssl.SSLContext`` to use for the
+                connection. If provided, it takes precedence over the other
+                SSL configuration options.
 
         Raises:
             RedisError
@@ -2096,6 +2100,7 @@ def __init__(
         self.ssl_ocsp_expected_cert = ssl_ocsp_expected_cert
         self.ssl_min_version = ssl_min_version
         self.ssl_ciphers = ssl_ciphers
+        self.ssl_context = ssl_context
         super().__init__(**kwargs)
 
     def _connect(self):
@@ -2119,33 +2124,36 @@ def _wrap_socket_with_ssl(self, sock):
         Returns:
             An SSL wrapped socket.
         """
-        context = ssl.create_default_context()
-        context.check_hostname = self.check_hostname
-        context.verify_mode = self.cert_reqs
-        if self.ssl_include_verify_flags:
-            for flag in self.ssl_include_verify_flags:
-                context.verify_flags |= flag
-        if self.ssl_exclude_verify_flags:
-            for flag in self.ssl_exclude_verify_flags:
-                context.verify_flags &= ~flag
-        if self.certfile or self.keyfile:
-            context.load_cert_chain(
-                certfile=self.certfile,
-                keyfile=self.keyfile,
-                password=self.certificate_password,
-            )
-        if (
-            self.ca_certs is not None
-            or self.ca_path is not None
-            or self.ca_data is not None
-        ):
-            context.load_verify_locations(
-                cafile=self.ca_certs, capath=self.ca_path, cadata=self.ca_data
-            )
-        if self.ssl_min_version is not None:
-            context.minimum_version = self.ssl_min_version
-        if self.ssl_ciphers:
-            context.set_ciphers(self.ssl_ciphers)
+        if self.ssl_context is None:
+            context = ssl.create_default_context()
+            context.check_hostname = self.check_hostname
+            context.verify_mode = self.cert_reqs
+            if self.ssl_include_verify_flags:
+                for flag in self.ssl_include_verify_flags:
+                    context.verify_flags |= flag
+            if self.ssl_exclude_verify_flags:
+                for flag in self.ssl_exclude_verify_flags:
+                    context.verify_flags &= ~flag
+            if self.certfile or self.keyfile:
+                context.load_cert_chain(
+                    certfile=self.certfile,
+                    keyfile=self.keyfile,
+                    password=self.certificate_password,
+                )
+            if (
+                self.ca_certs is not None
+                or self.ca_path is not None
+                or self.ca_data is not None
+            ):
+                context.load_verify_locations(
+                    cafile=self.ca_certs, capath=self.ca_path, cadata=self.ca_data
+                )
+            if self.ssl_min_version is not None:
+                context.minimum_version = self.ssl_min_version
+            if self.ssl_ciphers:
+                context.set_ciphers(self.ssl_ciphers)
+        else:
+            context = self.ssl_context
         if self.ssl_validate_ocsp is True and CRYPTOGRAPHY_AVAILABLE is False:
             raise RedisError("cryptography is not installed.")
 

tests/test_asyncio/test_cluster.py

@@ -77,6 +77,17 @@
 ]
 
 
+def test_cluster_preserves_custom_ssl_context():
+    context = ssl.create_default_context()
+    cluster = RedisCluster(
+        startup_nodes=[ClusterNode("localhost", 6379)],
+        ssl=True,
+        ssl_context=context,
+    )
+
+    assert cluster.connection_kwargs["ssl_context"] is context
+
+
 class NodeProxy:
     """A class to proxy a node connection to a different port"""
 

tests/test_asyncio/test_connect.py

@@ -4,6 +4,7 @@
 import ssl
 
 import pytest
+import redis.asyncio as redis
 from redis.asyncio.connection import (
     Connection,
     SSLConnection,
@@ -53,6 +54,26 @@ async def test_uds_connect(uds_address):
     await _assert_connect(conn, path)
 
 
+@pytest.mark.ssl
+async def test_tcp_ssl_uses_custom_context(tcp_address):
+    context = ssl.create_default_context()
+    conn = SSLConnection(host="localhost", port=tcp_address[1], ssl_context=context)
+
+    assert conn.ssl_context.get() is context
+
+
+@pytest.mark.ssl
+async def test_redis_passes_custom_context_to_ssl_connection():
+    context = ssl.create_default_context()
+    client = redis.Redis(ssl=True, ssl_context=context)
+
+    try:
+        connection = client.connection_pool.make_connection()
+        assert connection.ssl_context.get() is context
+    finally:
+        await client.aclose()
+
+
 @pytest.mark.ssl
 @pytest.mark.parametrize(
     "ssl_ciphers",

tests/test_asyncio/test_connection_pool.py

@@ -753,6 +753,19 @@ def test_host(self):
         assert pool.connection_class == redis.SSLConnection
         assert_kwargs_subset(pool.connection_kwargs, {"host": "my.host"})
 
+    def test_custom_ssl_context(self):
+        import ssl
+
+        context = ssl.create_default_context()
+
+        class DummyConnectionPool(redis.ConnectionPool):
+            def get_connection(self):
+                return self.make_connection()
+
+        pool = DummyConnectionPool.from_url("rediss://my.host", ssl_context=context)
+
+        assert pool.get_connection().ssl_context.get() is context
+
     def test_cert_reqs_options(self):
         import ssl
 

tests/test_cluster.py

@@ -3,6 +3,7 @@
 import select
 import socket
 import socketserver
+import ssl
 import threading
 from typing import List
 import warnings
@@ -75,6 +76,22 @@
 ]
 
 
+def test_cluster_preserves_custom_ssl_context():
+    context = ssl.create_default_context()
+
+    with (
+        patch.object(NodesManager, "initialize"),
+        patch.object(CommandsParser, "initialize"),
+    ):
+        cluster = RedisCluster(
+            startup_nodes=[ClusterNode("localhost", 6379)],
+            ssl=True,
+            ssl_context=context,
+        )
+
+    assert cluster.get_connection_kwargs()["ssl_context"] is context
+
+
 class ProxyRequestHandler(socketserver.BaseRequestHandler):
     def recv(self, sock):
         """A recv with a timeout"""

tests/test_connect.py

@@ -3,8 +3,10 @@
 import socketserver
 import ssl
 import threading
+import unittest.mock
 
 import pytest
+import redis
 from redis.connection import Connection, SSLConnection, UnixDomainSocketConnection
 from redis.exceptions import RedisError
 
@@ -48,6 +50,29 @@ def test_uds_connect(uds_address):
     _assert_connect(conn, path)
 
 
+@pytest.mark.ssl
+def test_tcp_ssl_uses_custom_context(tcp_address):
+    context = unittest.mock.Mock(spec=ssl.SSLContext)
+    wrapped_socket = unittest.mock.Mock()
+    context.wrap_socket.return_value = wrapped_socket
+    conn = SSLConnection(host="localhost", port=tcp_address[1], ssl_context=context)
+
+    assert conn._wrap_socket_with_ssl(unittest.mock.Mock()) is wrapped_socket
+    context.wrap_socket.assert_called_once()
+
+
+@pytest.mark.ssl
+def test_redis_passes_custom_context_to_ssl_connection():
+    context = unittest.mock.Mock(spec=ssl.SSLContext)
+    client = redis.Redis(ssl=True, ssl_context=context)
+
+    try:
+        connection = client.connection_pool.make_connection()
+        assert connection.ssl_context is context
+    finally:
+        client.close()
+
+
 @pytest.mark.ssl
 @pytest.mark.parametrize(
     "ssl_min_version",

tests/test_connection_pool.py

@@ -849,6 +849,17 @@ def test_host(self):
         assert pool.connection_class == redis.SSLConnection
         assert_kwargs_subset(pool.connection_kwargs, {"host": "my.host"})
 
+    def test_custom_ssl_context(self):
+        context = ssl.create_default_context()
+
+        class DummyConnectionPool(redis.ConnectionPool):
+            def get_connection(self):
+                return self.make_connection()
+
+        pool = DummyConnectionPool.from_url("rediss://my.host", ssl_context=context)
+
+        assert pool.get_connection().ssl_context is context
+
     def test_connection_class_override(self):
         class MyConnection(redis.SSLConnection):
             pass