如何在已连接的 Python asyncio stream 上启用 TLS？

Question

I have a Python asyncio server written using the high level Streams API . I want to enable TLS on an already established connection, as in STARTTLS in the SMTP and IMAP protocols. The asyncio event loop has a start_tls() function (added in Python 3.7), but it takes a protocol and a transport rather than a stream. The streams API does let you get the transport via StreamWriter.transport . But I don't see a way to change the transport, which would be required after calling start_tls(). Is it possible to use start_tls() with the streams API?

Answer 1

Looking at the code for the streams API you'll notice that StreamReader and StreamWriter both store their transport in an internal _transport variable. It turns out that if you call start_tls() and then store the new transport in those variables it works just fine. All the usual caveats with using an internal API apply of course. Here's what this looks like for a server. On a client I think you can just drop the load_cert_chain and server_side bits.

transport = writer.transport
protocol = transport.get_protocol()
loop = asyncio.get_event_loop()
ssl_context = ssl.SSLContext()
ssl_context.load_cert_chain("/path/to/certchain", "/path/to/key")
new_transport = await loop.start_tls(
    transport, protocol, ssl_context, server_side=True)
writer._transport = new_transport
reader._transport = new_transport

Answer 2

I needed to implement proxy support for asyncio streams of Python 3.8 and came up with following solution:

import socket
import weakref
import asyncio
import typing as t
from ssl import create_default_context, Purpose, SSLContext


class TLSStreamReaderProtocol(asyncio.StreamReaderProtocol):

    def upgrade_reader(self):
        if self._stream_reader is not None:
            self._stream_reader.set_exception(Exception('upgraded connection to TLS, this reader is obsolete now.'))
        self._stream_reader_wr = weakref.ref(reader)
        self._source_traceback = reader._source_traceback


async def open_tls_stream(host: str, port: int, ssl: t.Union[SSLContext, bool]=False):
    # this does the same as loop.open_connection(), but TLS upgrade is done
    # manually after connection be established.
    loop = asyncio.get_running_loop()
    reader = asyncio.StreamReader(limit=2**64, loop=loop)
    protocol = TLSStreamReaderProtocol(reader, loop=loop)
    transport, _ = await loop.create_connection(
        lambda: protocol, host, port, family=socket.AF_INET
    )
    writer = asyncio.StreamWriter(transport, protocol, reader, loop)
    # here you can use reader and writer for whatever you want, for example
    # start a proxy connection and start TLS to target host later...
    # now perform TLS upgrade
    if ssl:
        transport = await loop.start_tls(
            transport,
            protocol,
            sslcontext=create_default_context(Purpose.SERVER_AUTH) if isinstance(ssl, bool) else ssl,
            server_side=False,
            server_hostname=host
        )
        reader = asyncio.StreamReader(limit=2**64, loop=loop)
        protocol.upgrade_reader(reader) # update reader
        protocol.connection_made(transport) # update transport
        writer = asyncio.StreamWriter(transport, protocol, reader, loop) # update writer
    return reader, writer     

Answer 3

I'm using the following:

import asyncio
import ssl

async def tls_handshake(
    reader: asyncio.StreamReader,
    writer: asyncio.StreamWriter,
    ssl_context: ssl.SSLContext,
):
    """
    Upgrades the client connection to TLS/SSL.

    Args:
        reader: The reader of the client connection.
        writer: The writer of the client connection.
        ssl_context: The SSL context to use for the TLS/SSL handshake.

    Notes:
        For Python 3.6 to 3.9 you can use ``ssl.PROTOCOL_TLS`` for the SSL context. For
        Python 3.10+ you need to either use ``ssl.PROTOCOL_TLS_CLIENT`` or
        ``ssl.PROTOCOL_TLS_SERVER`` depending on the role of the reader/writer.
    """

    transport = writer.transport
    protocol = transport.get_protocol()

    loop = asyncio.get_event_loop()
    new_transport = await loop.start_tls(
        transport,
        protocol,
        ssl_context,
        server_side=True,
    )

    reader._transport = new_transport
    writer._transport = new_transport