Fixes to handle retries for WantWriteError and WantReadError in SSL

Author: julianz-

cheroot/makefile.py

@@ -3,6 +3,9 @@
 # prefer slower Python-based io module
 import _pyio as io
 import socket
+import time
+
+from OpenSSL import SSL
 
 
 # Write only 16K at a time to sockets
@@ -32,6 +35,14 @@ def _flush_unlocked(self):
                 n = self.raw.write(bytes(self._write_buf))
             except io.BlockingIOError as e:
                 n = e.characters_written
+            except (
+                SSL.WantReadError,
+                SSL.WantWriteError,
+                SSL.WantX509LookupError,
+            ):
+                # these errors require retries with the same data
+                # if some data has already been written
+                n = 0
             del self._write_buf[:n]
 
 
@@ -45,9 +56,22 @@ def __init__(self, sock, mode='r', bufsize=io.DEFAULT_BUFFER_SIZE):
 
     def read(self, *args, **kwargs):
         """Capture bytes read."""
-        val = super().read(*args, **kwargs)
-        self.bytes_read += len(val)
-        return val
+        MAX_ATTEMPTS = 10
+        attempts = 0
+        while True:
+            try:
+                val = super().read(*args, **kwargs)
+            except (SSL.WantReadError, SSL.WantWriteError):
+                attempts += 1
+                if attempts >= MAX_ATTEMPTS:
+                    # Raise an error if max retries reached
+                    raise TimeoutError(
+                        'Max retries exceeded while waiting for data.',
+                    )
+                time.sleep(0.1)
+            else:
+                self.bytes_read += len(val)
+                return val
 
     def has_data(self):
         """Return true if there is buffered data to read."""

cheroot/test/test_ssl.py

@@ -17,6 +17,8 @@
 import requests
 import trustme
 
+from cheroot.makefile import BufferedWriter
+
 from .._compat import (
     IS_ABOVE_OPENSSL10,
     IS_ABOVE_OPENSSL31,
@@ -625,6 +627,143 @@ def test_ssl_env(  # noqa: C901  # FIXME
     )
 
 
+@pytest.fixture
+def mock_raw_open(mocker):
+    """Return a mocked raw socket prepared for writing (closed=False)."""
+    # This fixture sets the state on the injected object
+    mock_raw = mocker.Mock()
+    mock_raw.closed = False
+    return mock_raw
+
+
+@pytest.fixture
+def ssl_writer(mock_raw_open):
+    """Return a BufferedWriter instance with a mocked raw socket."""
+    return BufferedWriter(mock_raw_open)
+
+
+def test_want_write_error_retry(ssl_writer, mock_raw_open):
+    """Test that WantWriteError causes retry with same data."""
+    test_data = b'hello world'
+
+    # set up mock socket so that when its write() method is called,
+    # we get WantWriteError first, then success on the second call
+    # indicated by returning the number of bytes written
+    mock_raw_open.write.side_effect = [
+        OpenSSL.SSL.WantWriteError(),
+        len(test_data),
+    ]
+
+    bytes_written = ssl_writer.write(test_data)
+    assert bytes_written == len(test_data)
+
+    # Assert against the injected mock object
+    assert mock_raw_open.write.call_count == 2
+
+
+def test_want_read_error_retry(ssl_writer, mock_raw_open):
+    """Test that WantReadError causes retry with same data."""
+    test_data = b'test data'
+
+    # set up mock socket so that when its write() method is called,
+    # we get WantReadError first, then success on the second call
+    # indicated by returning the number of bytes written
+    mock_raw_open.write.side_effect = [
+        OpenSSL.SSL.WantReadError(),
+        len(test_data),
+    ]
+
+    bytes_written = ssl_writer.write(test_data)
+    assert bytes_written == len(test_data)
+
+
+@pytest.fixture(
+    params=['builtin', 'pyopenssl'],
+)
+def adapter_type(request):
+    """Fixture that yields the name of the SSL adapter."""
+    return request.param
+
+
+@pytest.fixture
+def ssl_writer_integration(
+    mocker,
+    adapter_type,
+    tls_certificate_chain_pem_path,
+    tls_certificate_private_key_pem_path,
+):
+    """
+    Set up mock SSL writer for integration test.
+
+    Mocks the lowest-level write/send method to simulate a
+    transient WantWriteError.
+    """
+    # Set up SSL adapter
+    tls_adapter_cls = get_ssl_adapter_class(name=adapter_type)
+    tls_adapter = tls_adapter_cls(
+        tls_certificate_chain_pem_path,
+        tls_certificate_private_key_pem_path,
+    )
+
+    # Ensure context is initialized if needed
+    if adapter_type == 'pyopenssl':
+        # --- PYOPENSSL SETUP
+        tls_adapter.context = tls_adapter.get_context()
+        mock_raw_socket = mocker.Mock(name='mock_raw_socket')
+        mock_raw_socket.fileno.return_value = 1  # need to mock a dummy fd
+
+        # Create the real OpenSSL.SSL.Connection object
+        ssl_conn = OpenSSL.SSL.Connection(tls_adapter.context, mock_raw_socket)
+        ssl_conn.set_connect_state()
+        ssl_conn.closed = False
+
+        # Return the BufferedWriter and the specific mock for assertions
+        raw_io_object = ssl_conn
+        raw_io_object.write = mocker.Mock(name='ssl_conn_write_mock')
+    else:
+        # adapter_type == 'builtin'
+        # --- BUILTIN ADAPTER SETUP (Requires different mocking) ---
+        # Mock the adapter's own low-level write method
+        raw_io_object = tls_adapter
+        raw_io_object.write = mocker.Mock(
+            name='builtin_adapter_write',
+            autospec=True,
+        )
+        raw_io_object.closed = False
+        raw_io_object.writable = mocker.Mock(return_value=True)
+
+    # Return both the writer and the specific mock assertion target
+    return BufferedWriter(raw_io_object), raw_io_object.write
+
+
+def test_want_write_error_integration(ssl_writer_integration):
+    """Integration test for SSL writer handling of WantWriteError."""
+    writer, mock_write = ssl_writer_integration
+    test_data = b'integration test data'
+    successful_write_length = len(test_data)
+
+    # Determine the failure mechanism
+    failure_error = (
+        OpenSSL.SSL.WantWriteError() if adapter_type == 'pyopenssl' else 0
+    )
+
+    # Configure the mock's side effect with the first error
+    # and then the calculated buffer length for success
+    mock_write.side_effect = [
+        failure_error,
+        successful_write_length,
+    ]
+
+    # write data and then flush
+    # with the way the mock_write is set up this should fail once,
+    # and then succeed on the retry.
+    bytes_written = writer.write(test_data)
+    writer.flush()
+
+    assert bytes_written == successful_write_length
+    assert mock_write.call_count == 2
+
+
 @pytest.mark.parametrize(
     'ip_addr',
     (