如何在已连接的Python异步IO流上启用TLS？

Question

我使用高级Streams API编写了一个Python异步IO服务器。我想在已经建立的连接上启用TLS，如SMTP和IMAP协议中的STARTTLS。asyncio事件循环有一个start_tls()函数（在Python 3.7中添加），但它采用协议和传输，而不是流。streams API确实允许您通过StreamWriter.transport获得传输。但是我没有看到改变传输的方法，调用start_tls（）后需要改变传输。是否可以将start_tls（）与streams API一起使用

Answer 1

我需要为Python 3.8的异步IO流实现代理支持，并提出了以下解决方案：

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 2

查看streams API的the code，您会注意到StreamReader和StreamWriter都将其传输存储在内部_transport变量中。事实证明，如果调用start_tls（），然后将新传输存储在这些变量中，它就可以正常工作。使用内部API的所有常见注意事项当然都适用。以下是服务器的外观。在客户机上，我认为您可以删除load_cert_chain和server_side位

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