fix(streaming): hold sync client TLS reads until the handshake request is written

assemblyai/__version__.py

@@ -1 +1 @@
-__version__ = "1.6.1"
+__version__ = "1.6.2"

assemblyai/streaming/v3/_tls.py

@@ -0,0 +1,103 @@
+"""TLS context for the thread-based streaming client.
+
+websockets' sync client starts the thread that reads the connection before the
+calling thread writes the opening handshake request. With servers that send
+TLS 1.3 post-handshake messages, such as session tickets, a read already
+blocked on the socket while that request is written occasionally leaves the
+connection without a handshake response, and the connect times out. Sockets
+created from ``_create_tls_context()`` hold reads until the first write has
+returned, so the request is written before anything is read.
+
+Reads are held only on sockets the context builds from its
+``sslsocket_class``. An ``SSLContext`` replacement whose ``wrap_socket()``
+builds sockets some other way keeps its own sockets, and reads on them are not
+held.
+"""
+
+from __future__ import annotations
+
+import ssl
+import threading
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from _typeshed import ReadableBuffer, WriteableBuffer
+
+# Upper bound on how long a read is held, so a socket that is never written to
+# can still be read.
+_FIRST_WRITE_WAIT_SEC = 30.0
+
+# Serializes the creation of each socket's first-write event.
+_FIRST_WRITE_LOCK = threading.Lock()
+
+
+class _WriteFirstSSLSocket(ssl.SSLSocket):
+    """``SSLSocket`` whose reads wait until the first write has returned.
+
+    The first ``send()`` or ``sendall()`` releases held reads whether or not it
+    succeeds, and so does ``shutdown()``, which lets closing the socket
+    interrupt a held read. Once released, reads and writes proceed
+    concurrently as with a plain ``SSLSocket``.
+    """
+
+    _first_write_event: Optional[threading.Event] = None
+
+    def _first_write_done(self) -> threading.Event:
+        # SSLContext.wrap_socket() builds instances without calling __init__,
+        # so the event is created on first use.
+        done = self._first_write_event
+        if done is None:
+            with _FIRST_WRITE_LOCK:
+                done = self._first_write_event
+                if done is None:
+                    done = self._first_write_event = threading.Event()
+        return done
+
+    def send(self, data: ReadableBuffer, flags: int = 0) -> int:
+        try:
+            return super().send(data, flags)
+        finally:
+            self._first_write_done().set()
+
+    def sendall(self, data: ReadableBuffer, flags: int = 0) -> None:
+        try:
+            super().sendall(data, flags)
+        finally:
+            self._first_write_done().set()
+
+    def recv(self, buflen: int = 1024, flags: int = 0) -> bytes:
+        self._first_write_done().wait(_FIRST_WRITE_WAIT_SEC)
+        return super().recv(buflen, flags)
+
+    def recv_into(
+        self,
+        buffer: WriteableBuffer,
+        nbytes: Optional[int] = None,
+        flags: int = 0,
+    ) -> int:
+        self._first_write_done().wait(_FIRST_WRITE_WAIT_SEC)
+        return super().recv_into(buffer, nbytes, flags)
+
+    def shutdown(self, how: int) -> None:
+        self._first_write_done().set()
+        super().shutdown(how)
+
+
+def _create_tls_context() -> ssl.SSLContext:
+    """Return a default client context whose sockets are
+    ``_WriteFirstSSLSocket``.
+
+    Certificate and hostname verification are those of
+    ``ssl.create_default_context()``, which is also what websockets uses when
+    no context is given.
+
+    ``_WriteFirstSSLSocket`` extends the ``ssl.SSLSocket`` present when this
+    module was imported. If the default context uses a socket class that
+    ``_WriteFirstSSLSocket`` does not extend, for example because gevent
+    patched ``ssl`` afterwards, the context keeps that class: its
+    ``wrap_socket()`` may not be able to construct ``_WriteFirstSSLSocket``.
+    """
+    context = ssl.create_default_context()
+    if issubclass(_WriteFirstSSLSocket, context.sslsocket_class):
+        context.sslsocket_class = _WriteFirstSSLSocket
+    return context

assemblyai/streaming/v3/client.py

@@ -1,14 +1,17 @@
+import inspect
 import json
 import logging
 import queue
 import threading
 import time
 from typing import Any, Dict, Generator, Iterable, Optional, Union
+from urllib.parse import urlsplit
 
 import httpx
 import websockets
 from pydantic import BaseModel
-from websockets.sync.client import connect as websocket_connect
+from websockets.sync.client import ClientConnection
+from websockets.sync.client import connect as _ws_sync_connect
 
 from ._base import (
     _BaseStreamingClient,
@@ -21,6 +24,7 @@
     _resolve_options,
     _user_agent,
 )
+from ._tls import _create_tls_context
 from .models import (
     ErrorEvent,
     EventMessage,
@@ -40,6 +44,23 @@
 
 logger = logging.getLogger(__name__)
 
+# The sync ``connect()`` TLS parameter is ``ssl`` from websockets 13 on and
+# ``ssl_context`` before (``setup.py`` allows ``websockets>=11.0``).
+_WS_SSL_KW = (
+    "ssl" if "ssl" in inspect.signature(_ws_sync_connect).parameters else "ssl_context"
+)
+
+
+def websocket_connect(uri: str, **kwargs: Any) -> ClientConnection:
+    """Open a sync websocket connection with ``websockets.sync.client.connect``.
+
+    ``wss://`` URIs use the context from ``_create_tls_context()``. Module-level
+    indirection so tests can patch a single attribute.
+    """
+    if urlsplit(uri).scheme == "wss":
+        kwargs[_WS_SSL_KW] = _create_tls_context()
+    return _ws_sync_connect(uri, **kwargs)
+
 
 class RealTimeTranscriber(_BaseStreamingClient):
     def __init__(