gh-155027: Flush close_notify before closing SSL transport

When SSL shutdown failed, SSLProtocol copied close_notify to the raw
transport and immediately force-closed it. Selector transports discarded
buffered output and proactor transports cancelled the active write.

Close the raw transport normally after SSL write flow resumes, keep the
shutdown timeout active while it drains, and preserve the SSL exception for
the application protocol unless the raw transport reports its own error.

Author: matthiasgoergens

Lib/asyncio/sslproto.py

@@ -323,6 +323,8 @@ def __init__(self, loop, app_protocol, sslcontext, waiter,
         self._outgoing = ssl.MemoryBIO()
         self._state = SSLProtocolState.UNWRAPPED
         self._conn_lost = 0  # Set when connection_lost called
+        self._shutdown_exc = None
+        self._shutdown_close_pending = False
         if call_connection_made:
             self._app_state = AppProtocolState.STATE_INIT
         else:
@@ -397,6 +399,11 @@ def connection_lost(self, exc):
         meaning a regular EOF is received or the connection was
         aborted or closed).
         """
+        if exc is None and self._shutdown_exc is not None:
+            exc = self._shutdown_exc
+        self._shutdown_exc = None
+        self._shutdown_close_pending = False
+
         self._write_backlog.clear()
         self._outgoing.read()
         self._conn_lost += 1
@@ -669,14 +676,21 @@ def _do_shutdown(self):
             self._on_shutdown_complete(None)
 
     def _on_shutdown_complete(self, shutdown_exc):
-        if self._shutdown_timeout_handle is not None:
-            self._shutdown_timeout_handle.cancel()
-            self._shutdown_timeout_handle = None
+        # close() lets the raw transport flush data queued by
+        # _process_outgoing().  _fatal_error() would force-close it and
+        # discard the close_notify that shutdown just produced. Keep the
+        # shutdown timeout active until connection_lost() bounds the drain.
+        if shutdown_exc is not None:
+            self._shutdown_exc = shutdown_exc
+        self._shutdown_close_pending = True
+        if not self._ssl_writing_paused:
+            self._loop.call_soon(self._close_transport)
 
-        if shutdown_exc:
-            self._fatal_error(shutdown_exc)
-        else:
-            self._loop.call_soon(self._transport.close)
+    def _close_transport(self):
+        if self._shutdown_close_pending:
+            self._shutdown_close_pending = False
+            if self._transport is not None:
+                self._transport.close()
 
     def _abort(self, exc):
         self._set_state(SSLProtocolState.UNWRAPPED)
@@ -927,6 +941,8 @@ def resume_writing(self):
         assert self._ssl_writing_paused
         self._ssl_writing_paused = False
         self._process_outgoing()
+        if self._shutdown_close_pending:
+            self._loop.call_soon(self._close_transport)
 
     def _fatal_error(self, exc, message='Fatal error on transport'):
         if self._transport:

Lib/test/test_asyncio/test_sslproto.py

@@ -96,6 +96,77 @@ def test_fatal_error_no_name_error(self):
             # Restore error logging.
             log.logger.setLevel(log_level)
 
+    def test_shutdown_error_closes_after_flushing(self):
+        app_proto = mock.Mock(spec=asyncio.Protocol)
+        app_proto.eof_received.return_value = False
+        ssl_proto = self.ssl_protocol(proto=app_proto)
+        ssl_proto._state = sslproto.SSLProtocolState.SHUTDOWN
+        ssl_proto._app_state = sslproto.AppProtocolState.STATE_CON_MADE
+        transport = mock.Mock()
+        ssl_proto._transport = transport
+        ssl_proto._sslobj = mock.Mock()
+        shutdown_exc = ssl.SSLError(ssl.SSL_ERROR_SSL, 'shutdown failed')
+        ssl_proto._sslobj.unwrap.side_effect = shutdown_exc
+        ssl_proto._outgoing = mock.Mock()
+        ssl_proto._outgoing.read.side_effect = [b'close notify', b'', b'']
+        ssl_proto._outgoing.pending = 0
+        timeout_handle = mock.Mock()
+        ssl_proto._shutdown_timeout_handle = timeout_handle
+
+        ssl_proto._do_shutdown()
+        ssl_proto.eof_received()
+
+        transport.write.assert_called_once_with(b'close notify')
+        transport._force_close.assert_not_called()
+        test_utils.run_briefly(self.loop)
+        transport.close.assert_called_once_with()
+        timeout_handle.cancel.assert_not_called()
+
+        ssl_proto.connection_lost(None)
+        test_utils.run_briefly(self.loop)
+        app_proto.connection_lost.assert_called_once_with(shutdown_exc)
+        timeout_handle.cancel.assert_called_once_with()
+
+    def test_shutdown_waits_for_resume_writing_before_close(self):
+        app_proto = mock.Mock(spec=asyncio.Protocol)
+        ssl_proto = self.ssl_protocol(proto=app_proto)
+        ssl_proto._state = sslproto.SSLProtocolState.SHUTDOWN
+        ssl_proto._app_state = sslproto.AppProtocolState.STATE_CON_MADE
+        transport = mock.Mock()
+        ssl_proto._transport = transport
+        ssl_proto._sslobj = mock.Mock()
+        shutdown_exc = ssl.SSLError(ssl.SSL_ERROR_SSL, 'shutdown failed')
+        ssl_proto._sslobj.unwrap.side_effect = shutdown_exc
+        ssl_proto._outgoing = mock.Mock()
+        ssl_proto._outgoing.read.return_value = b'close notify'
+        ssl_proto._outgoing.pending = 0
+        ssl_proto._ssl_writing_paused = True
+
+        ssl_proto._do_shutdown()
+        test_utils.run_briefly(self.loop)
+
+        transport.write.assert_not_called()
+        transport.close.assert_not_called()
+
+        ssl_proto.resume_writing()
+        transport.write.assert_called_once_with(b'close notify')
+        test_utils.run_briefly(self.loop)
+        transport.close.assert_called_once_with()
+
+    def test_shutdown_raw_error_takes_precedence(self):
+        app_proto = mock.Mock(spec=asyncio.Protocol)
+        ssl_proto = self.ssl_protocol(proto=app_proto)
+        ssl_proto._state = sslproto.SSLProtocolState.SHUTDOWN
+        ssl_proto._app_state = sslproto.AppProtocolState.STATE_CON_MADE
+        ssl_proto._shutdown_exc = ssl.SSLError(
+            ssl.SSL_ERROR_SSL, 'shutdown failed')
+        raw_exc = ConnectionResetError('raw write failed')
+
+        ssl_proto.connection_lost(raw_exc)
+        test_utils.run_briefly(self.loop)
+
+        app_proto.connection_lost.assert_called_once_with(raw_exc)
+
     def test_connection_lost(self):
         # From issue #472.
         # yield from waiter hang if lost_connection was called.