bpo-23749: Implement loop.start_tls() (#5039)

4 files changed, 59 insertions, 1 deletions

diff --git a/Lib/asyncio/base_events.py b/Lib/asyncio/base_events.py
index 96cc4f0..00831b3 100644
--- a/Lib/asyncio/base_events.py
+++ b/Lib/asyncio/base_events.py
@@ -29,9 +29,15 @@ import sys
 import warnings
 import weakref
 
+try:
+    import ssl
+except ImportError:  # pragma: no cover
+    ssl = None
+
 from . import coroutines
 from . import events
 from . import futures
+from . import sslproto
 from . import tasks
 from .log import logger
 
@@ -279,7 +285,8 @@ class BaseEventLoop(events.AbstractEventLoop):
             self, rawsock, protocol, sslcontext, waiter=None,
             *, server_side=False, server_hostname=None,
             extra=None, server=None,
-            ssl_handshake_timeout=None):
+            ssl_handshake_timeout=None,
+            call_connection_made=True):
         """Create SSL transport."""
         raise NotImplementedError
 
@@ -795,6 +802,42 @@ class BaseEventLoop(events.AbstractEventLoop):
 
         return transport, protocol
 
+    async def start_tls(self, transport, protocol, sslcontext, *,
+                        server_side=False,
+                        server_hostname=None,
+                        ssl_handshake_timeout=None):
+        """Upgrade transport to TLS.
+
+        Return a new transport that *protocol* should start using
+        immediately.
+        """
+        if ssl is None:
+            raise RuntimeError('Python ssl module is not available')
+
+        if not isinstance(sslcontext, ssl.SSLContext):
+            raise TypeError(
+                f'sslcontext is expected to be an instance of ssl.SSLContext, '
+                f'got {sslcontext!r}')
+
+        if not getattr(transport, '_start_tls_compatible', False):
+            raise TypeError(
+                f'transport {self!r} is not supported by start_tls()')
+
+        waiter = self.create_future()
+        ssl_protocol = sslproto.SSLProtocol(
+            self, protocol, sslcontext, waiter,
+            server_side, server_hostname,
+            ssl_handshake_timeout=ssl_handshake_timeout,
+            call_connection_made=False)
+
+        transport.set_protocol(ssl_protocol)
+        self.call_soon(ssl_protocol.connection_made, transport)
+        if not transport.is_reading():
+            self.call_soon(transport.resume_reading)
+
+        await waiter
+        return ssl_protocol._app_transport
+
     async def create_datagram_endpoint(self, protocol_factory,
                                        local_addr=None, remote_addr=None, *,
                                        family=0, proto=0, flags=0,
diff --git a/Lib/asyncio/events.py b/Lib/asyncio/events.py
index 3a5dbad..9496d5c 100644
--- a/Lib/asyncio/events.py
+++ b/Lib/asyncio/events.py
@@ -305,6 +305,17 @@ class AbstractEventLoop:
         """
         raise NotImplementedError
 
+    async def start_tls(self, transport, protocol, sslcontext, *,
+                        server_side=False,
+                        server_hostname=None,
+                        ssl_handshake_timeout=None):
+        """Upgrade a transport to TLS.
+
+        Return a new transport that *protocol* should start using
+        immediately.
+        """
+        raise NotImplementedError
+
     async def create_unix_connection(
             self, protocol_factory, path=None, *,
             ssl=None, sock=None,
diff --git a/Lib/asyncio/proactor_events.py b/Lib/asyncio/proactor_events.py
index 2661cdd..ab1285b 100644
--- a/Lib/asyncio/proactor_events.py
+++ b/Lib/asyncio/proactor_events.py
@@ -223,6 +223,8 @@ class _ProactorBaseWritePipeTransport(_ProactorBasePipeTransport,
                                       transports.WriteTransport):
     """Transport for write pipes."""
 
+    _start_tls_compatible = True
+
     def write(self, data):
         if not isinstance(data, (bytes, bytearray, memoryview)):
             raise TypeError(
diff --git a/Lib/asyncio/selector_events.py b/Lib/asyncio/selector_events.py
index 1e4bd83..5692e38 100644
--- a/Lib/asyncio/selector_events.py
+++ b/Lib/asyncio/selector_events.py
@@ -694,6 +694,8 @@ class _SelectorTransport(transports._FlowControlMixin,
 
 class _SelectorSocketTransport(_SelectorTransport):
 
+    _start_tls_compatible = True
+
     def __init__(self, loop, sock, protocol, waiter=None,
                  extra=None, server=None):
         super().__init__(loop, sock, protocol, extra, server)