Skip to content

Commit 78f1e14

Browse files
vuanhphungIsaac
andcommitted
Stop urllib3 from rewriting shared SSL contexts per connection
urllib3 calls set_alpn_protocols() and sets verify_mode on its context for every connection. On a shared context, replacing the ALPN buffer races with SSL_new() in other threads (which runs without the GIL) and can corrupt the ClientHello: 1-4 corrupted ALPN lists per ~500k handshake starts in a 32-thread stress test. Cached contexts now ignore same-value writes, with ALPN set once before the context is shared. Contexts that are not plain ssl.SSLContext objects are not cached. Signed-off-by: Vu Anh Phung <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com>
1 parent d4f71e8 commit 78f1e14

2 files changed

Lines changed: 118 additions & 6 deletions

File tree

‎src/databricks/sql/common/ssl_context_cache.py‎

Lines changed: 46 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616
from collections import OrderedDict
1717
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple
1818

19+
from urllib3.util.ssl_ import ALPN_PROTOCOLS
20+
1921
if TYPE_CHECKING:
2022
from databricks.sql.types import SSLOptions
2123

@@ -26,6 +28,41 @@
2628
_contexts: "OrderedDict[tuple, Tuple[ssl.SSLContext, float]]" = OrderedDict()
2729
_build_locks: Dict[tuple, Any] = {}
2830

31+
_BASE_CONTEXT_CLASS = ssl.SSLContext
32+
33+
34+
class _SharedSSLContext(_BASE_CONTEXT_CLASS): # type: ignore[valid-type,misc]
35+
"""A cached context that ignores same-value writes.
36+
37+
urllib3 re-applies ALPN and verify_mode to its context on every connection.
38+
SSL_new() copies the context's ALPN buffer without the GIL, so replacing
39+
that buffer while another thread connects can corrupt the ClientHello.
40+
"""
41+
42+
def set_alpn_protocols(self, alpn_protocols):
43+
protocols = list(alpn_protocols)
44+
if protocols != getattr(self, "_shared_alpn_protocols", None):
45+
super().set_alpn_protocols(protocols)
46+
self._shared_alpn_protocols = protocols
47+
48+
@property
49+
def verify_mode(self):
50+
return super().verify_mode
51+
52+
@verify_mode.setter
53+
def verify_mode(self, value):
54+
if super().verify_mode != value:
55+
_BASE_CONTEXT_CLASS.verify_mode.__set__(self, value)
56+
57+
58+
def _make_shareable(context: ssl.SSLContext) -> bool:
59+
# Only plain stdlib contexts are shared; others (e.g. pyOpenSSL) stay per caller.
60+
if type(context) is not _BASE_CONTEXT_CLASS:
61+
return False
62+
context.__class__ = _SharedSSLContext
63+
context.set_alpn_protocols(ALPN_PROTOCOLS)
64+
return True
65+
2966

3067
def _file_signature(path) -> Optional[Tuple]:
3168
# Rewriting or replacing a CA file yields a new key, so rotation is picked up.
@@ -74,9 +111,11 @@ def get_ssl_context(
74111
75112
``kind`` separates builders that configure contexts differently. urllib3
76113
writes a pool's ``cert_reqs`` and ALPN protocols back onto its context on
77-
every connection, so every pool sharing a context must use TLS pool kwargs
78-
derived only from ``ssl_options`` (and no ``assert_hostname``). Do not
79-
otherwise modify the returned context.
114+
every connection; the shared context skips those writes only when nothing
115+
changes, so every pool sharing a context must use TLS pool kwargs derived
116+
only from ``ssl_options`` (and no ``assert_hostname``). Do not otherwise
117+
modify the returned context. Contexts other than plain ``ssl.SSLContext``
118+
are returned without caching.
80119
"""
81120
if not is_cacheable(ssl_options):
82121
return build()
@@ -96,9 +135,13 @@ def get_ssl_context(
96135
return context
97136
try:
98137
context = build()
138+
shareable = _make_shareable(context)
99139
except BaseException:
100140
_forget_build_lock(key, key_lock)
101141
raise
142+
if not shareable:
143+
_forget_build_lock(key, key_lock)
144+
return context
102145
with _lock:
103146
_contexts[key] = (context, time.monotonic())
104147
_contexts.move_to_end(key)

‎tests/unit/test_ssl_context_cache.py‎

Lines changed: 72 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,16 @@
1818
from databricks.sql.types import SSLOptions
1919

2020

21+
def bare_context():
22+
# Real but cheap: no CA bundle is loaded.
23+
return ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
24+
25+
2126
def counting_builder():
2227
built = []
2328

2429
def build():
25-
context = Mock(spec=ssl.SSLContext)
30+
context = bare_context()
2631
built.append(context)
2732
return context
2833

@@ -144,7 +149,7 @@ def test_concurrent_first_use_builds_once(self):
144149

145150
def slow_build():
146151
time.sleep(0.05)
147-
context = Mock(spec=ssl.SSLContext)
152+
context = bare_context()
148153
built.append(context)
149154
return context
150155

@@ -175,7 +180,7 @@ def test_slow_build_does_not_block_other_keys(self):
175180

176181
def stuck_build():
177182
release.wait(10)
178-
return Mock(spec=ssl.SSLContext)
183+
return bare_context()
179184

180185
stuck = threading.Thread(
181186
target=ssl_context_cache.get_ssl_context,
@@ -233,6 +238,70 @@ def test_cache_is_bounded(self, monkeypatch):
233238
built[0]
234239
)
235240

241+
def test_non_stdlib_contexts_are_not_cached(self):
242+
build = Mock(side_effect=lambda: Mock(spec=ssl.SSLContext))
243+
244+
first = ssl_context_cache.get_ssl_context("k", SSLOptions(), build)
245+
second = ssl_context_cache.get_ssl_context("k", SSLOptions(), build)
246+
247+
assert first is not second
248+
assert not ssl_context_cache._contexts
249+
250+
def test_shared_context_skips_same_value_writes(self):
251+
context = ssl_context_cache.get_ssl_context("k", SSLOptions(), bare_context)
252+
assert isinstance(context, ssl.SSLContext)
253+
254+
with patch.object(ssl.SSLContext, "_set_alpn_protocols") as set_alpn:
255+
context.set_alpn_protocols(["http/1.1"]) # what urllib3 sends
256+
set_alpn.assert_not_called()
257+
context.set_alpn_protocols(["h2"])
258+
set_alpn.assert_called_once()
259+
260+
real = ssl.SSLContext.verify_mode
261+
writes = []
262+
263+
def spy_set(self, value):
264+
writes.append(value)
265+
real.fset(self, value)
266+
267+
context.check_hostname = False
268+
with patch.object(ssl.SSLContext, "verify_mode", property(real.fget, spy_set)):
269+
context.verify_mode = ssl.CERT_REQUIRED
270+
assert writes == []
271+
context.verify_mode = ssl.CERT_NONE
272+
assert writes == [ssl.CERT_NONE]
273+
assert context.verify_mode == ssl.CERT_NONE
274+
275+
def test_reset_after_fork_releases_locks(self):
276+
held = threading.Event()
277+
release = threading.Event()
278+
279+
def hold_locks():
280+
with ssl_context_cache._lock:
281+
held.set()
282+
release.wait(10)
283+
284+
holder = threading.Thread(target=hold_locks)
285+
holder.start()
286+
held.wait(5)
287+
stuck_build_lock = threading.Lock()
288+
stuck_build_lock.acquire()
289+
key = ssl_context_cache._cache_key("k", SSLOptions())
290+
ssl_context_cache._build_locks[key] = stuck_build_lock
291+
try:
292+
ssl_context_cache._reset_after_fork() # as in a forked child
293+
294+
done = threading.Thread(
295+
target=ssl_context_cache.get_ssl_context,
296+
args=("k", SSLOptions(), bare_context),
297+
)
298+
done.start()
299+
done.join(timeout=2)
300+
assert not done.is_alive()
301+
finally:
302+
release.set()
303+
holder.join()
304+
236305

237306
class TestThriftTransportContext:
238307
@staticmethod

0 commit comments

Comments
 (0)