Changed eof_received() to reflect close_notify

Author: fantix

tests/test_tcp.py

@@ -1573,6 +1573,7 @@ def serve(sock):
             data = sock.recv_all(len(HELLO_MSG))
             self.assertEqual(len(data), len(HELLO_MSG))
 
+            sock.unwrap()
             sock.shutdown(socket.SHUT_RDWR)
             sock.close()
 
@@ -1643,6 +1644,7 @@ def serve(sock):
             data = sock.recv_all(len(HELLO_MSG))
             self.assertEqual(len(data), len(HELLO_MSG))
 
+            sock.unwrap()
             sock.shutdown(socket.SHUT_RDWR)
             sock.close()
 
@@ -1798,6 +1800,7 @@ def client(sock, addr):
             sock.starttls(client_context)
             sock.sendall(HELLO_MSG)
 
+            sock.unwrap()
             sock.shutdown(socket.SHUT_RDWR)
             sock.close()
 
@@ -2303,9 +2306,11 @@ def test_write_to_closed_transport(self):
 
         sslctx = self._create_server_ssl_context(self.ONLYCERT, self.ONLYKEY)
         client_sslctx = self._create_client_ssl_context()
+        future = None
 
         def server(sock):
             sock.starttls(sslctx, server_side=True)
+            sock.shutdown(socket.SHUT_RDWR)
             sock.close()
 
         def unwrap_server(sock):
@@ -2320,6 +2325,9 @@ def unwrap_server(sock):
             sock.close()
 
         async def client(addr):
+            nonlocal future
+            future = self.loop.create_future()
+
             reader, writer = await asyncio.open_connection(
                 *addr,
                 ssl=client_sslctx,
@@ -2330,7 +2338,7 @@ async def client(addr):
             try:
                 data = await reader.read()
                 self.assertEqual(data, b'')
-            except ConnectionResetError:
+            except (ConnectionResetError, BrokenPipeError):
                 pass
 
             for i in range(25):
@@ -2339,11 +2347,23 @@ async def client(addr):
             self.assertEqual(
                 len(writer.transport._ssl_protocol._write_backlog), 0)
 
+            await future
+
+        def run(meth):
+            def wrapper(sock):
+                try:
+                    meth(sock)
+                except Exception as ex:
+                    self.loop.call_soon_threadsafe(future.set_exception, ex)
+                else:
+                    self.loop.call_soon_threadsafe(future.set_result, None)
+            return wrapper
+
         with self._silence_eof_received_warning():
-            with self.tcp_server(server) as srv:
+            with self.tcp_server(run(server)) as srv:
                 self.loop.run_until_complete(client(srv.addr))
 
-            with self.tcp_server(unwrap_server) as srv:
+            with self.tcp_server(run(unwrap_server)) as srv:
                 self.loop.run_until_complete(client(srv.addr))
 
     def test_flush_before_shutdown(self):
@@ -2438,6 +2458,98 @@ async def client(addr):
         with self.tcp_server(run(openssl_server)) as srv:
             self.loop.run_until_complete(client(srv.addr))
 
+    def test_remote_shutdown_receives_trailing_data(self):
+        if self.implementation == 'asyncio':
+            raise unittest.SkipTest()
+
+        CHUNK = 1024 * 128
+        SIZE = 32
+
+        sslctx = self._create_server_ssl_context(self.ONLYCERT, self.ONLYKEY)
+        client_sslctx = self._create_client_ssl_context()
+        future = None
+
+        def server(sock):
+            sock.starttls(sslctx, server_side=True)
+            self.assertEqual(sock.recv_all(4), b'ping')
+            sock.send(b'pong')
+
+            time.sleep(0.2)  # wait for the peer to fill its backlog
+
+            # send close_notify but don't wait for response
+            sock.setblocking(0)
+            with self.assertRaises(ssl.SSLWantReadError):
+                sock.unwrap()
+            sock.setblocking(1)
+
+            # should receive all data
+            data = sock.recv_all(CHUNK * SIZE)
+            self.assertEqual(len(data), CHUNK * SIZE)
+
+            # wait for close_notify
+            sock.unwrap()
+
+            sock.close()
+
+        def eof_server(sock):
+            sock.starttls(sslctx, server_side=True)
+            self.assertEqual(sock.recv_all(4), b'ping')
+            sock.send(b'pong')
+
+            time.sleep(0.2)  # wait for the peer to fill its backlog
+
+            # send EOF
+            sock.shutdown(socket.SHUT_WR)
+
+            # should receive all data
+            data = sock.recv_all(CHUNK * SIZE)
+            self.assertEqual(len(data), CHUNK * SIZE)
+
+            sock.close()
+
+        async def client(addr):
+            nonlocal future
+            future = self.loop.create_future()
+
+            reader, writer = await asyncio.open_connection(
+                *addr,
+                ssl=client_sslctx,
+                server_hostname='',
+                loop=self.loop)
+            writer.write(b'ping')
+            data = await reader.readexactly(4)
+            self.assertEqual(data, b'pong')
+
+            # fill write backlog in a hacky way - renegotiation won't help
+            ssl_protocol = writer.transport._ssl_protocol
+            for _ in range(SIZE):
+                ssl_protocol._write_backlog.append(b'x' * CHUNK)
+                ssl_protocol._write_buffer_size += CHUNK
+
+            try:
+                data = await reader.read()
+                self.assertEqual(data, b'')
+            except (BrokenPipeError, ConnectionResetError):
+                pass
+
+            await future
+
+        def run(meth):
+            def wrapper(sock):
+                try:
+                    meth(sock)
+                except Exception as ex:
+                    self.loop.call_soon_threadsafe(future.set_exception, ex)
+                else:
+                    self.loop.call_soon_threadsafe(future.set_result, None)
+            return wrapper
+
+        with self.tcp_server(run(server)) as srv:
+            self.loop.run_until_complete(client(srv.addr))
+
+        with self.tcp_server(run(eof_server)) as srv:
+            self.loop.run_until_complete(client(srv.addr))
+
 
 class Test_UV_TCPSSL(_TestSSL, tb.UVTestCase):
     pass

uvloop/sslproto.pyx

@@ -262,6 +262,7 @@ class SSLProtocol(object):
         self._ssl_buffer = bytearray(256 * 1024)
         self._state = _UNWRAPPED
         self._conn_lost = 0  # Set when connection_lost called
+        self._eof_received = False
 
         # Flow Control
 
@@ -362,13 +363,23 @@ class SSLProtocol(object):
             if self._loop.get_debug():
                 aio_logger.debug("%r received EOF", self)
 
-            self._wakeup_waiter(ConnectionResetError)
+            if self._state == _DO_HANDSHAKE:
+                self._on_handshake_complete(ConnectionResetError)
+
+            elif self._state == _WRAPPED:
+                self._set_state(_FLUSHING)
+                self._do_write()
+                self._set_state(_SHUTDOWN)
+                self._do_shutdown()
+
+            elif self._state == _FLUSHING:
+                self._do_write()
+                self._set_state(_SHUTDOWN)
+                self._do_shutdown()
+
+            elif self._state == _SHUTDOWN:
+                self._do_shutdown()
 
-            if self._state != _DO_HANDSHAKE:
-                keep_open = self._app_protocol.eof_received()
-                if keep_open:
-                    aio_logger.warning('returning true from eof_received() '
-                                       'has no effect when using ssl')
         finally:
             self._transport.close()
 
@@ -529,6 +540,7 @@ class SSLProtocol(object):
                 self._on_shutdown_complete(exc)
         else:
             self._process_outgoing()
+            self._call_eof_received()
             self._on_shutdown_complete(None)
 
     def _on_shutdown_complete(self, shutdown_exc):
@@ -627,6 +639,7 @@ class SSLProtocol(object):
             self._app_protocol.buffer_updated(offset)
         if not count:
             # close_notify
+            self._call_eof_received()
             self._start_shutdown()
 
     def _do_read__copied(self):
@@ -647,8 +660,20 @@ class SSLProtocol(object):
             self._app_protocol.data_received(b''.join(data))
         if not chunk:
             # close_notify
+            self._call_eof_received()
             self._start_shutdown()
 
+    def _call_eof_received(self):
+        try:
+            if not self._eof_received:
+                self._eof_received = True
+                keep_open = self._app_protocol.eof_received()
+                if keep_open:
+                    aio_logger.warning('returning true from eof_received() '
+                                       'has no effect when using ssl')
+        except Exception as ex:
+            self._fatal_error(ex, 'Error calling eof_received()')
+
     # Flow control for writes from APP socket
 
     def _control_app_writing(self):