bpo-31386: Custom wrap_bio and wrap_socket type (#3426)

SSLSocket.wrap_bio() and SSLSocket.wrap_socket() hard-code SSLObject and
SSLSocket as return types. In the light of future deprecation of
ssl.wrap_socket() module function and direct instantiation of SSLSocket,
it is desirable to make the return type of SSLSocket.wrap_bio() and
SSLSocket.wrap_socket() customizable.

Signed-off-by: Christian Heimes <christian@python.org>

2 files changed, 34 insertions, 8 deletions

diff --git a/Lib/ssl.py b/Lib/ssl.py
index 062e802118..2849deee07 100644
--- a/Lib/ssl.py
+++ b/Lib/ssl.py
@@ -383,10 +383,11 @@ class Purpose(_ASN1Object, _Enum):
 class SSLContext(_SSLContext):
     """An SSLContext holds various SSL-related configuration options and
     data, such as certificates and possibly a private key."""
-
-    __slots__ = ('protocol', '__weakref__')
     _windows_cert_stores = ("CA", "ROOT")
 
+    sslsocket_class = None  # SSLSocket is assigned later.
+    sslobject_class = None  # SSLObject is assigned later.
+
     def __new__(cls, protocol=PROTOCOL_TLS, *args, **kwargs):
         self = _SSLContext.__new__(cls, protocol)
         if protocol != _SSLv2_IF_EXISTS:
@@ -400,17 +401,21 @@ class SSLContext(_SSLContext):
                     do_handshake_on_connect=True,
                     suppress_ragged_eofs=True,
                     server_hostname=None, session=None):
-        return SSLSocket(sock=sock, server_side=server_side,
-                         do_handshake_on_connect=do_handshake_on_connect,
-                         suppress_ragged_eofs=suppress_ragged_eofs,
-                         server_hostname=server_hostname,
-                         _context=self, _session=session)
+        return self.sslsocket_class(
+            sock=sock,
+            server_side=server_side,
+            do_handshake_on_connect=do_handshake_on_connect,
+            suppress_ragged_eofs=suppress_ragged_eofs,
+            server_hostname=server_hostname,
+            _context=self,
+            _session=session
+        )
 
     def wrap_bio(self, incoming, outgoing, server_side=False,
                  server_hostname=None, session=None):
         sslobj = self._wrap_bio(incoming, outgoing, server_side=server_side,
                                 server_hostname=server_hostname)
-        return SSLObject(sslobj, session=session)
+        return self.sslobject_class(sslobj, session=session)
 
     def set_npn_protocols(self, npn_protocols):
         protos = bytearray()
@@ -1135,6 +1140,11 @@ class SSLSocket(socket):
         return self._sslobj.version()
 
 
+# Python does not support forward declaration of types.
+SSLContext.sslsocket_class = SSLSocket
+SSLContext.sslobject_class = SSLObject
+
+
 def wrap_socket(sock, keyfile=None, certfile=None,
                 server_side=False, cert_reqs=CERT_NONE,
                 ssl_version=PROTOCOL_TLS, ca_certs=None,
diff --git a/Lib/test/test_ssl.py b/Lib/test/test_ssl.py
index 523322da2f..fb5958f1a5 100644
--- a/Lib/test/test_ssl.py
+++ b/Lib/test/test_ssl.py
@@ -1359,6 +1359,22 @@ class ContextTests(unittest.TestCase):
         self.assertFalse(ctx.check_hostname)
         self.assertEqual(ctx.verify_mode, ssl.CERT_NONE)
 
+    def test_context_custom_class(self):
+        class MySSLSocket(ssl.SSLSocket):
+            pass
+
+        class MySSLObject(ssl.SSLObject):
+            pass
+
+        ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+        ctx.sslsocket_class = MySSLSocket
+        ctx.sslobject_class = MySSLObject
+
+        with ctx.wrap_socket(socket.socket(), server_side=True) as sock:
+            self.assertIsInstance(sock, MySSLSocket)
+        obj = ctx.wrap_bio(ssl.MemoryBIO(), ssl.MemoryBIO())
+        self.assertIsInstance(obj, MySSLObject)
+
 
 class SSLErrorTests(unittest.TestCase):