bpo-33654: Support BufferedProtocol in set_protocol() and start_tls()

Author: 1st1

Doc/library/asyncio-protocol.rst

@@ -469,10 +469,16 @@ the buffer only once at creation time.
 The following callbacks are called on :class:`BufferedProtocol`
 instances:
 
-.. method:: BufferedProtocol.get_buffer()
+.. method:: BufferedProtocol.get_buffer(sizehint)
 
-   Called to allocate a new receive buffer.  Must return an object
-   that implements the :ref:`buffer protocol <bufferobjects>`.
+   Called to allocate a new receive buffer.
+
+   *sizehint* is a recommended minimal size for the returned
+   buffer.  When set to -1, the buffer size can be arbitrary.
+
+   Must return an object that implements the
+   :ref:`buffer protocol <bufferobjects>`.
+   It is an error to return a zero-sized buffer.
 
 .. method:: BufferedProtocol.buffer_updated(nbytes)
 

Lib/asyncio/proactor_events.py

@@ -30,7 +30,7 @@ def __init__(self, loop, sock, protocol, waiter=None,
         super().__init__(extra, loop)
         self._set_extra(sock)
         self._sock = sock
-        self._protocol = protocol
+        self.set_protocol(protocol)
         self._server = server
         self._buffer = None  # None or bytearray.
         self._read_fut = None
@@ -159,16 +159,27 @@ class _ProactorReadPipeTransport(_ProactorBasePipeTransport,
 
     def __init__(self, loop, sock, protocol, waiter=None,
                  extra=None, server=None):
+        self._loop_reading_cb = None
+        self._paused = True
         super().__init__(loop, sock, protocol, waiter, extra, server)
-        self._paused = False
+
         self._reschedule_on_resume = False
+        self._loop.call_soon(self._loop_reading)
+        self._paused = False
 
-        if protocols._is_buffered_protocol(protocol):
-            self._loop_reading = self._loop_reading__get_buffer
+    def set_protocol(self, protocol):
+        if isinstance(protocol, protocols.BufferedProtocol):
+            self._loop_reading_cb = self._loop_reading__get_buffer
         else:
-            self._loop_reading = self._loop_reading__data_received
+            self._loop_reading_cb = self._loop_reading__data_received
 
-        self._loop.call_soon(self._loop_reading)
+        super().set_protocol(protocol)
+
+        if not self._paused:
+            # reset reading callback / buffers / self._read_fut
+            self.pause_reading()
+            self._reschedule_on_resume = True
+            self.resume_reading()
 
     def is_reading(self):
         return not self._paused and not self._closing
@@ -210,7 +221,10 @@ def _loop_reading__on_eof(self):
         if not keep_open:
             self.close()
 
-    def _loop_reading__data_received(self, fut=None):
+    def _loop_reading(self, fut=None):
+        self._loop_reading_cb(fut)
+
+    def _loop_reading__data_received(self, fut):
         if self._paused:
             self._reschedule_on_resume = True
             return
@@ -260,7 +274,7 @@ def _loop_reading__data_received(self, fut=None):
             elif data == b'':
                 self._loop_reading__on_eof()
 
-    def _loop_reading__get_buffer(self, fut=None):
+    def _loop_reading__get_buffer(self, fut):
         if self._paused:
             self._reschedule_on_resume = True
             return
@@ -310,7 +324,9 @@ def _loop_reading__get_buffer(self, fut=None):
             return
 
         try:
-            buf = self._protocol.get_buffer()
+            buf = self._protocol.get_buffer(-1)
+            if not len(buf):
+                raise RuntimeError('get_buffer() returned an empty buffer')
         except Exception as exc:
             self._fatal_error(
                 exc, 'Fatal error: protocol.get_buffer() call failed.')

Lib/asyncio/protocols.py

@@ -130,11 +130,15 @@ class BufferedProtocol(BaseProtocol):
     * CL: connection_lost()
     """
 
-    def get_buffer(self):
+    def get_buffer(self, sizehint):
         """Called to allocate a new receive buffer.
 
+        *sizehint* is a recommended minimal size for the returned
+        buffer.  When set to -1, the buffer size can be arbitrary.
+
         Must return an object that implements the
         :ref:`buffer protocol <bufferobjects>`.
+        It is an error to return a zero-sized buffer.
         """
 
     def buffer_updated(self, nbytes):
@@ -185,7 +189,3 @@ def pipe_connection_lost(self, fd, exc):
 
     def process_exited(self):
         """Called when subprocess has exited."""
-
-
-def _is_buffered_protocol(proto):
-    return hasattr(proto, 'get_buffer') and not hasattr(proto, 'data_received')

Lib/asyncio/selector_events.py

@@ -597,8 +597,10 @@ def __init__(self, loop, sock, protocol, extra=None, server=None):
                 self._extra['peername'] = None
         self._sock = sock
         self._sock_fd = sock.fileno()
-        self._protocol = protocol
-        self._protocol_connected = True
+
+        self._protocol_connected = False
+        self.set_protocol(protocol)
+
         self._server = server
         self._buffer = self._buffer_factory()
         self._conn_lost = 0  # Set when call to connection_lost scheduled.
@@ -640,6 +642,7 @@ def abort(self):
 
     def set_protocol(self, protocol):
         self._protocol = protocol
+        self._protocol_connected = True
 
     def get_protocol(self):
         return self._protocol
@@ -721,11 +724,7 @@ class _SelectorSocketTransport(_SelectorTransport):
     def __init__(self, loop, sock, protocol, waiter=None,
                  extra=None, server=None):
 
-        if protocols._is_buffered_protocol(protocol):
-            self._read_ready = self._read_ready__get_buffer
-        else:
-            self._read_ready = self._read_ready__data_received
-
+        self._read_ready_cb = None
         super().__init__(loop, sock, protocol, extra, server)
         self._eof = False
         self._paused = False
@@ -745,6 +744,14 @@ def __init__(self, loop, sock, protocol, waiter=None,
             self._loop.call_soon(futures._set_result_unless_cancelled,
                                  waiter, None)
 
+    def set_protocol(self, protocol):
+        if isinstance(protocol, protocols.BufferedProtocol):
+            self._read_ready_cb = self._read_ready__get_buffer
+        else:
+            self._read_ready_cb = self._read_ready__data_received
+
+        super().set_protocol(protocol)
+
     def is_reading(self):
         return not self._paused and not self._closing
 
@@ -764,12 +771,17 @@ def resume_reading(self):
         if self._loop.get_debug():
             logger.debug("%r resumes reading", self)
 
+    def _read_ready(self):
+        self._read_ready_cb()
+
     def _read_ready__get_buffer(self):
         if self._conn_lost:
             return
 
         try:
-            buf = self._protocol.get_buffer()
+            buf = self._protocol.get_buffer(-1)
+            if not len(buf):
+                raise RuntimeError('get_buffer() returned an empty buffer')
         except Exception as exc:
             self._fatal_error(
                 exc, 'Fatal error: protocol.get_buffer() call failed.')

Lib/asyncio/sslproto.py

@@ -441,6 +441,8 @@ def __init__(self, loop, app_protocol, sslcontext, waiter,
         self._waiter = waiter
         self._loop = loop
         self._app_protocol = app_protocol
+        self._app_protocol_is_buffer = \
+            isinstance(app_protocol, protocols.BufferedProtocol)
         self._app_transport = _SSLProtocolTransport(self._loop, self)
         # _SSLPipe instance (None until the connection is made)
         self._sslpipe = None
@@ -522,7 +524,16 @@ def data_received(self, data):
 
         for chunk in appdata:
             if chunk:
-                self._app_protocol.data_received(chunk)
+                try:
+                    if self._app_protocol_is_buffer:
+                        _feed_data_to_bufferred_proto(
+                            self._app_protocol, chunk)
+                    else:
+                        self._app_protocol.data_received(chunk)
+                except Exception as ex:
+                    self._fatal_error(
+                        ex, 'application protocol failed to receive SSL data')
+                    return
             else:
                 self._start_shutdown()
                 break
@@ -709,3 +720,22 @@ def _abort(self):
                 self._transport.abort()
         finally:
             self._finalize()
+
+
+def _feed_data_to_bufferred_proto(proto, data):
+    data_len = len(data)
+    while data_len:
+        buf = proto.get_buffer(data_len)
+        buf_len = len(buf)
+        if not buf_len:
+            raise RuntimeError('get_buffer() returned an empty buffer')
+
+        if buf_len >= data_len:
+            buf[:data_len] = data
+            proto.buffer_updated(data_len)
+            return
+        else:
+            buf[:buf_len] = data[:buf_len]
+            proto.buffer_updated(buf_len)
+            data = data[buf_len:]
+            data_len = len(data)

Lib/test/test_asyncio/test_buffered_proto.py

@@ -9,7 +9,7 @@ def __init__(self, cb, con_lost_fut):
         self.cb = cb
         self.con_lost_fut = con_lost_fut
 
-    def get_buffer(self):
+    def get_buffer(self, sizehint):
         self.buffer = bytearray(100)
         return self.buffer
 

Lib/test/test_asyncio/test_proactor_events.py

@@ -465,8 +465,8 @@ def setUp(self):
         self.loop._proactor = self.proactor
 
         self.protocol = test_utils.make_test_protocol(asyncio.BufferedProtocol)
-        self.buf = mock.Mock()
-        self.protocol.get_buffer.side_effect = lambda: self.buf
+        self.buf = bytearray(1)
+        self.protocol.get_buffer.side_effect = lambda hint: self.buf
 
         self.sock = mock.Mock(socket.socket)
 
@@ -505,6 +505,62 @@ def test_get_buffer_error(self):
         self.assertTrue(self.protocol.get_buffer.called)
         self.assertFalse(self.protocol.buffer_updated.called)
 
+    def test_get_buffer_zerosized(self):
+        transport = self.socket_transport()
+        transport._fatal_error = mock.Mock()
+
+        self.loop.call_exception_handler = mock.Mock()
+        self.protocol.get_buffer.side_effect = lambda hint: bytearray(0)
+
+        transport._loop_reading()
+
+        self.assertTrue(transport._fatal_error.called)
+        self.assertTrue(self.protocol.get_buffer.called)
+        self.assertFalse(self.protocol.buffer_updated.called)
+
+    def test_proto_type_switch(self):
+        self.protocol = test_utils.make_test_protocol(asyncio.Protocol)
+        tr = self.socket_transport()
+
+        res = asyncio.Future(loop=self.loop)
+        res.set_result(b'data')
+
+        tr = self.socket_transport()
+        tr._read_fut = res
+        tr._loop_reading(res)
+        self.loop._proactor.recv.assert_called_with(self.sock, 32768)
+        self.protocol.data_received.assert_called_with(b'data')
+
+        # switch protocol to a BufferedProtocol
+
+        buf_proto = test_utils.make_test_protocol(asyncio.BufferedProtocol)
+        buf = bytearray(4)
+        buf_proto.get_buffer.side_effect = lambda hint: buf
+
+        tr.set_protocol(buf_proto)
+        res = asyncio.Future(loop=self.loop)
+        res.set_result(4)
+
+        tr._read_fut = res
+        tr._loop_reading(res)
+        self.loop._proactor.recv_into.assert_called_with(self.sock, buf)
+        buf_proto.buffer_updated.assert_called_with(4)
+
+    def test_proto_buf_switch(self):
+        tr = self.socket_transport()
+        test_utils.run_briefly(self.loop)
+        self.protocol.get_buffer.assert_called_with(-1)
+
+        # switch protocol to *another* BufferedProtocol
+
+        buf_proto = test_utils.make_test_protocol(asyncio.BufferedProtocol)
+        buf = bytearray(4)
+        buf_proto.get_buffer.side_effect = lambda hint: buf
+        tr.set_protocol(buf_proto)
+        self.assertFalse(buf_proto.get_buffer.called)
+        test_utils.run_briefly(self.loop)
+        buf_proto.get_buffer.assert_called_with(-1)
+
     def test_buffer_updated_error(self):
         transport = self.socket_transport()
         transport._fatal_error = mock.Mock()

Lib/test/test_asyncio/test_selector_events.py

@@ -1279,8 +1279,8 @@ def setUp(self):
         self.loop = self.new_test_loop()
 
         self.protocol = test_utils.make_test_protocol(asyncio.BufferedProtocol)
-        self.buf = mock.Mock()
-        self.protocol.get_buffer.side_effect = lambda: self.buf
+        self.buf = bytearray(1)
+        self.protocol.get_buffer.side_effect = lambda hint: self.buf
 
         self.sock = mock.Mock(socket.socket)
         self.sock_fd = self.sock.fileno.return_value = 7
@@ -1313,6 +1313,42 @@ def test_get_buffer_error(self):
         self.assertTrue(self.protocol.get_buffer.called)
         self.assertFalse(self.protocol.buffer_updated.called)
 
+    def test_get_buffer_zerosized(self):
+        transport = self.socket_transport()
+        transport._fatal_error = mock.Mock()
+
+        self.loop.call_exception_handler = mock.Mock()
+        self.protocol.get_buffer.side_effect = lambda hint: bytearray(0)
+
+        transport._read_ready()
+
+        self.assertTrue(transport._fatal_error.called)
+        self.assertTrue(self.protocol.get_buffer.called)
+        self.assertFalse(self.protocol.buffer_updated.called)
+
+    def test_proto_type_switch(self):
+        self.protocol = test_utils.make_test_protocol(asyncio.Protocol)
+        transport = self.socket_transport()
+
+        self.sock.recv.return_value = b'data'
+        transport._read_ready()
+
+        self.protocol.data_received.assert_called_with(b'data')
+
+        # switch protocol to a BufferedProtocol
+
+        buf_proto = test_utils.make_test_protocol(asyncio.BufferedProtocol)
+        buf = bytearray(4)
+        buf_proto.get_buffer.side_effect = lambda hint: buf
+
+        transport.set_protocol(buf_proto)
+
+        self.sock.recv_into.return_value = 10
+        transport._read_ready()
+
+        buf_proto.get_buffer.assert_called_with(-1)
+        buf_proto.buffer_updated.assert_called_with(10)
+
     def test_buffer_updated_error(self):
         transport = self.socket_transport()
         transport._fatal_error = mock.Mock()
@@ -1348,7 +1384,7 @@ def test_read_ready(self):
         self.sock.recv_into.return_value = 10
         transport._read_ready()
 
-        self.protocol.get_buffer.assert_called_with()
+        self.protocol.get_buffer.assert_called_with(-1)
         self.protocol.buffer_updated.assert_called_with(10)
 
     def test_read_ready_eof(self):