1
"""
2
These are top level entropy patches for packages.
3
4
We'll need to keep these for backwards compatibility but our goal is to stop including these and
5
move them into workers-runtime-sdk.
6
"""
7
 
8
import sys
9
from contextlib import contextmanager
10
 
11
from .allow_entropy import (
12
    allow_bad_entropy_calls,
13
)
14
from .import_patch_manager import (
15
    block_calls,
16
    register_after_snapshot,
17
    register_before_first_request,
18
    register_create_patch,
19
    register_exec_patch,
20
)
21
 
22
IMPORTED_RUST_PACKAGE = False
23
 
24
 
25
@register_create_patch("tiktoken._tiktoken")
26
@register_exec_patch("cryptography.exceptions")
27
@register_exec_patch("jiter")
28
@contextmanager
29
def rust_package_context(module):
30
    """Rust packages need one entropy call if they create a rust hash map at
31
    init time.
32
33
    For reasons I don't entirely understand, in Pyodide 0.28 only the first Rust package to be
34
    imported makes the get_entropy call. See gen_rust_import_tests() which tests that importing
35
    four rust packages in different permutations works correctly.
36
    """
37
    global IMPORTED_RUST_PACKAGE
38
    if sys.version_info >= (3, 13) and IMPORTED_RUST_PACKAGE:
39
        yield
40
        return
41
    IMPORTED_RUST_PACKAGE = True
42
    with allow_bad_entropy_calls(1):
43
        yield
44
 
45
 
46
NUMPY_RANDOM_STATE = None
47
 
48
 
49
@register_exec_patch("numpy.random")
50
@contextmanager
51
def numpy_random_context(numpy_random):
52
    """numpy.random doesn't call getentropy() itself, but we want to block calls
53
    that might use the bad seed.
54
55
    TODO: Maybe there are more calls we can whitelist?
56
    TODO: Is it not enough to just block numpy.random.mtrand calls?
57
    """
58
    yield
59
    # Calling default_rng() with a given seed is fine, calling it without a seed
60
    # will call getentropy() and fail.
61
    block_calls(numpy_random, allowlist=("default_rng", "RandomState"))
62
 
63
 
64
@register_after_snapshot("numpy.random")
65
def numpy_random_after_snapshot(numpy_random):
66
    r1 = numpy_random.random()
67
    numpy_random.set_state(NUMPY_RANDOM_STATE)
68
    r2 = numpy_random.random()
69
    if r1 != r2:
70
        raise RuntimeError("random seed in bad state")
71
 
72
 
73
@register_before_first_request("numpy.random")
74
def numpy_random_before_first_request(numpy_random):
75
    numpy_random.seed()
76
 
77
 
78
@register_exec_patch("numpy.random.mtrand")
79
@contextmanager
80
def numpy_random_mtrand_context(module):
81
    # numpy.random.mtrand calls secrets.randbits at top level to seed itself.
82
    # This will fail if we don't let it through.
83
    with allow_bad_entropy_calls(1):
84
        yield
85
    # Block calls until we get a chance to replace the bad random seed.
86
    global NUMPY_RANDOM_STATE
87
    NUMPY_RANDOM_STATE = module.get_state()
88
    block_calls(module, allowlist=("RandomState",))
89
 
90
 
91
@register_exec_patch("pydantic_core")
92
@contextmanager
93
def pydantic_core_context(module):
94
    try:
95
        # Initial import needs one entropy call to initialize
96
        # std::collections::HashMap hash seed
97
        with allow_bad_entropy_calls(1):
98
            yield
99
    finally:
100
        try:
101
            with allow_bad_entropy_calls(1):
102
                # validate_core_schema makes an ahash::AHashMap which makes
103
                # another entropy call for its hash seed. It will throw an error
104
                # but only after making the needed entropy call.
105
                module.validate_core_schema(None)
106
        except module.SchemaError:
107
            pass
108
 
109
 
110
@register_exec_patch("aiohttp.http_websocket")
111
@contextmanager
112
def aiohttp_http_websocket_context(module):
113
    import random
114
 
115
    Random = random.Random
116
 
117
    def patched_Random():
118
        return random
119
 
120
    random.Random = patched_Random
121
    try:
122
        yield
123
    finally:
124
        random.Random = Random
125
 
126
 
127
class NoSslFinder:
128
    def find_spec(self, fullname, path, target):
129
        if fullname == "ssl":
130
            raise ModuleNotFoundError(
131
                f"No module named {fullname!r}", name=fullname
132
            ) from None
133
 
134
 
135
@contextmanager
136
def no_ssl():
137
    """
138
    Various packages will call ssl.create_default_context() at top level which uses entropy if they
139
    can import ssl. By temporarily making importing ssl raise an import error, we exercise the
140
    workaround code and so avoid the entropy calls. After, we put the ssl module back to the normal
141
    value.
142
    """
143
    try:
144
        f = NoSslFinder()
145
        ssl = sys.modules.pop("ssl", None)
146
        sys.meta_path.insert(0, f)
147
        yield
148
    finally:
149
        sys.meta_path.remove(f)
150
        if ssl:
151
            sys.modules["ssl"] = ssl
152
 
153
 
154
@register_exec_patch("aiohttp.connector")
155
@contextmanager
156
def aiohttp_connector_context(module):
157
    with no_ssl():
158
        yield
159
 
160
 
161
@register_exec_patch("requests.adapters")
162
@contextmanager
163
def requests_adapters_context(module):
164
    with no_ssl():
165
        yield
166
 
167
 
168
@register_exec_patch("urllib3.util.ssl_")
169
@contextmanager
170
def urllib3_util_ssl__context(module):
171
    with no_ssl():
172
        yield
173
 
174
 
175
@register_exec_patch("langsmith._internal._constants")
176
@contextmanager
177
def langsmith__internal__constants_context(module):
178
    # Langsmith uses a UUID to communicate with a background thread. This obviously won't work so we
179
    # might as well allow it to make a UUID.
180
    with allow_bad_entropy_calls(1):
181
        yield
182
 
183
 
184
@register_exec_patch("langchain_openai.chat_models.base")
185
@contextmanager
186
def langchain_openai_chat_models_base_context(module):
187
    if sys.version_info >= (3, 13):
188
        # Creates an ssl context in the version used with 3.13
189
        with allow_bad_entropy_calls(1):
190
            yield
191
    else:
192
        yield