Fix thread-safety issue while configuring SSLContext

Author: sethmlarson

noxfile.py

@@ -1,4 +1,5 @@
 import os
+import sys
 
 import nox
 
@@ -15,6 +16,11 @@ def test(session):
     session.env["YARL_NO_EXTENSIONS"] = "1"
     session.env["FROZENLIST_NO_EXTENSIONS"] = "1"
 
+    # This would need to be updated if we added PyPy
+    # to our default Python versions above. Right now
+    # PyPy is only used on CI.
+    pypy = session.python is None and sys.implementation.name == "pypy"
+
     session.install("-rdev-requirements.txt", ".")
     session.run("pip", "freeze")
     session.run(
@@ -23,6 +29,7 @@ def test(session):
         "-s",
         "-rs",
         "--no-flaky-report",
+        *(("--config-file=pypy-pytest.ini",) if pypy else ()),
         "--max-runs=3",
         *(session.posargs or ("tests/",)),
     )

pypy-pytest.ini

@@ -0,0 +1,14 @@
+# This pytest.ini is only for PyPy
+# to suppress 'ResourceWarning' from aiohttp.
+
+[pytest]
+asyncio_mode = strict
+asyncio_default_fixture_loop_scope = function
+filterwarnings =
+  error
+  # See: aio-libs/aiohttp#7545
+  ignore:.*datetime.utcfromtimestamp().*:DeprecationWarning
+  # PyPy and aiohttp don't play nice.
+  ignore:.*:ResourceWarning
+markers =
+  internet: test requires Internet access

src/truststore/_api.py

@@ -1,8 +1,10 @@
+import contextlib
 import os
 import platform
 import socket
 import ssl
 import sys
+import threading
 import typing
 
 import _ssl
@@ -84,6 +86,7 @@ def __class__(self) -> type:
 
     def __init__(self, protocol: int = None) -> None:  # type: ignore[assignment]
         self._ctx = _original_SSLContext(protocol)
+        self._ctx_lock = threading.Lock()
 
         class TruststoreSSLObject(ssl.SSLObject):
             # This object exists because wrap_bio() doesn't
@@ -106,10 +109,15 @@ def wrap_socket(
         server_hostname: str | None = None,
         session: ssl.SSLSession | None = None,
     ) -> ssl.SSLSocket:
-        # Use a context manager here because the
-        # inner SSLContext holds on to our state
-        # but also does the actual handshake.
-        with _configure_context(self._ctx):
+
+        # We need to lock around the .__enter__()
+        # but we don't need to lock within the
+        # context manager, so we need to expand the
+        # syntactic sugar of the `with` statement.
+        with contextlib.ExitStack() as stack:
+            with self._ctx_lock:
+                stack.enter_context(_configure_context(self._ctx))
+
             ssl_sock = self._ctx.wrap_socket(
                 sock,
                 server_side=server_side,

tests/conftest.py

@@ -172,4 +172,6 @@ async def handler(request: web.Request) -> web.Response:
         yield Server(host="localhost", port=port)
     finally:
         await site.stop()
+        # Wait for the server to actually close.
+        await site._server.wait_closed()
         await runner.cleanup()

tests/test_threading.py

@@ -0,0 +1,38 @@
+import asyncio
+import socket
+import ssl
+import threading
+
+import pytest
+
+import truststore
+
+
+def wrap_and_close_sockets(ctx: truststore.SSLContext, host: str, port: int) -> None:
+    for _ in range(100):
+        sock = None
+        try:
+            sock = socket.create_connection((host, port))
+            sock = ctx.wrap_socket(sock, server_hostname=host)
+        finally:
+            if sock:
+                sock.close()
+
+
+@pytest.mark.asyncio
+async def test_threading(server):
+    def run_threads():
+        ctx = truststore.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+        threads = [
+            threading.Thread(
+                target=wrap_and_close_sockets, args=(ctx, server.host, server.port)
+            )
+            for _ in range(16)
+        ]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+    thread = asyncio.to_thread(run_threads)
+    await thread