Fixes to handle retries for WantWriteError and WantReadError in SSL

As discussed in #245

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
@@ -31,7 +34,15 @@ def _flush_unlocked(self):
                 # so perhaps we should conditionally wrap this for perf?
                 n = self.raw.write(bytes(self._write_buf))
             except io.BlockingIOError as e:
-                n = e.characters_written
+                n = e.characters_writteni
+            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,15 @@ 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
+        while True:
+            try:
+                val = super().read(*args, **kwargs)
+                self.bytes_read += len(val)
+                return val
+            except SSL.WantReadError:
+                time.sleep(0.1)  # allow some retry delay
+            except SSL.WantWriteError:
+                time.sleep(0.1)
 
     def has_data(self):
         """Return true if there is buffered data to read."""

cheroot/server.pyi

@@ -1,5 +1,18 @@
 from typing import Any
 
+__all__ = (
+    'ChunkedRFile',
+    'DropUnderscoreHeaderReader',
+    'Gateway',
+    'HTTPConnection',
+    'HTTPRequest',
+    'HTTPServer',
+    'HeaderReader',
+    'KnownLengthRFile',
+    'SizeCheckWrapper',
+    'get_ssl_adapter_class',
+)
+
 class HeaderReader:
     def __call__(self, rfile, hdict: Any | None = ...): ...
 

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,47 @@ def test_ssl_env(  # noqa: C901  # FIXME
     )
 
 
+@pytest.fixture
+def ssl_writer(mocker):
+    """Return a BufferedWriter instance with a mocked raw socket."""
+    mock_raw = mocker.Mock()
+    mock_raw.closed = False
+    writer = BufferedWriter(mock_raw)
+
+    writer.mock_raw = mock_raw
+
+    return writer
+
+
+def test_want_write_error_retry(ssl_writer):
+    """Test that WantWriteError causes retry with same data."""
+    test_data = b'hello world'
+
+    # Access the mock object via the attribute we attached in the fixture
+    ssl_writer.mock_raw.write.side_effect = [
+        OpenSSL.SSL.WantWriteError(),
+        len(test_data),
+    ]
+
+    bytes_written = ssl_writer.write(test_data)
+    assert bytes_written == len(test_data)
+    assert ssl_writer.mock_raw.write.call_count == 2
+
+
+def test_want_read_error_retry(ssl_writer):
+    """Test that WantReadError causes retry with same data."""
+    test_data = b'test data'
+
+    # Access the mock object via the attribute we attached in the fixture
+    ssl_writer.mock_raw.write.side_effect = [
+        OpenSSL.SSL.WantReadError(),
+        len(test_data),
+    ]
+
+    bytes_written = ssl_writer.write(test_data)
+    assert bytes_written == len(test_data)
+
+
 @pytest.mark.parametrize(
     'ip_addr',
     (

cheroot/workers/threadpool.pyi

@@ -1,6 +1,8 @@
 import threading
 from typing import Any
 
+__all__ = ('ThreadPool', 'WorkerThread')
+
 class TrueyZero:
     def __add__(self, other): ...
     def __radd__(self, other): ...