[Bugfix][Security] Fix concurrent sparse invariant race bypassing CVE remediation (#48583)
Signed-off-by: jperezde <jperezde@redhat.com>
Juan Pérez de Algaba committed
Jul 14, 2026 at 13:08 UTC
793cf79c89d4049124e756915468ac30318f2e50
5 files changed
+298
-13
tests/renderers/test_sparse_tensor_concurrent_race.py
new
+263
@@ -0,0 +1,263 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import contextlib
5
+import io
6
+import threading
7
+from concurrent.futures import ThreadPoolExecutor
8
+
9
+import pybase64
10
+import pytest
11
+import torch
12
+
13
+from vllm.config import ModelConfig
14
+from vllm.multimodal.media import AudioEmbeddingMediaIO, ImageEmbeddingMediaIO
15
+from vllm.renderers.embed_utils import safe_load_prompt_embeds
16
+from vllm.utils.sparse_utils import check_sparse_tensor_invariants_threadsafe
17
+
18
+
19
+def _encode_tensor(tensor: torch.Tensor) -> bytes:
20
+ buffer = io.BytesIO()
21
+ torch.save(tensor, buffer)
22
+ buffer.seek(0)
23
+ return pybase64.b64encode(buffer.read())
24
+
25
+
26
+def _create_malicious_sparse_tensor() -> torch.Tensor:
27
+ indices = torch.tensor([[10], [10]])
28
+ values = torch.tensor([1.0])
29
+ shape = (3, 3)
30
+ return torch.sparse_coo_tensor(
31
+ indices, values, shape, dtype=torch.float32, check_invariants=False
32
+ )
33
+
34
+
35
+def _create_valid_dense_tensor(hidden_size: int = 768) -> torch.Tensor:
36
+ return torch.randn(4, hidden_size, dtype=torch.float32)
37
+
38
+
39
+@pytest.fixture
40
+def model_config():
41
+ return ModelConfig(
42
+ model="facebook/opt-125m",
43
+ tokenizer="facebook/opt-125m",
44
+ tokenizer_mode="auto",
45
+ trust_remote_code=False,
46
+ dtype="float32",
47
+ seed=0,
48
+ enable_prompt_embeds=True,
49
+ )
50
+
51
+
52
+class TestNegativeControlWithoutRace:
53
+ """The invalid payload must be rejected even without concurrency."""
54
+
55
+ def test_malicious_sparse_rejected_by_prompt_loader(self, model_config):
56
+ encoded = _encode_tensor(_create_malicious_sparse_tensor())
57
+ with pytest.raises((RuntimeError, ValueError)):
58
+ safe_load_prompt_embeds(model_config, encoded)
59
+
60
+ def test_malicious_sparse_rejected_by_image_loader(self):
61
+ io_handler = ImageEmbeddingMediaIO()
62
+ encoded = _encode_tensor(_create_malicious_sparse_tensor())
63
+ with pytest.raises((RuntimeError, ValueError)):
64
+ io_handler.load_base64("", encoded.decode("utf-8"))
65
+
66
+ def test_malicious_sparse_rejected_by_audio_loader(self):
67
+ io_handler = AudioEmbeddingMediaIO()
68
+ encoded = _encode_tensor(_create_malicious_sparse_tensor())
69
+ with pytest.raises((RuntimeError, ValueError)):
70
+ io_handler.load_base64("", encoded.decode("utf-8"))
71
+
72
+
73
+class TestConcurrentRaceProtection:
74
+ """Verify the lock prevents the A-enter, B-enter, A-exit, B-load race."""
75
+
76
+ def test_forced_interleaving_still_rejects_invalid(self, model_config):
77
+ """Force the exact interleaving that bypasses the guard without a lock.
78
+
79
+ Thread A (benign): enters context, loads valid tensor, exits context.
80
+ Thread B (malicious): enters context after A, but loads after A exits.
81
+
82
+ Without _SPARSE_LOAD_LOCK the flag would be False when B loads.
83
+ With the lock, B cannot enter until A fully completes.
84
+ """
85
+ barrier = threading.Barrier(2, timeout=5)
86
+ results: dict[str, object] = {}
87
+
88
+ valid_encoded = _encode_tensor(_create_valid_dense_tensor())
89
+ malicious_encoded = _encode_tensor(_create_malicious_sparse_tensor())
90
+
91
+ def thread_a_benign():
92
+ """Enter context, signal B, load, exit."""
93
+ try:
94
+ with check_sparse_tensor_invariants_threadsafe():
95
+ barrier.wait() # signal B that A holds the lock
96
+ tensor = torch.load(
97
+ io.BytesIO(pybase64.b64decode(valid_encoded, validate=True)),
98
+ weights_only=True,
99
+ map_location=torch.device("cpu"),
100
+ )
101
+ results["a_loaded"] = True
102
+ results["a_tensor_valid"] = isinstance(tensor, torch.Tensor)
103
+ except threading.BrokenBarrierError:
104
+ results["a_barrier_broken"] = True
105
+
106
+ def thread_b_malicious():
107
+ """Wait for A to hold the lock, then try to acquire it."""
108
+ try:
109
+ barrier.wait() # wait until A is inside the lock
110
+ with check_sparse_tensor_invariants_threadsafe():
111
+ tensor = torch.load(
112
+ io.BytesIO(
113
+ pybase64.b64decode(malicious_encoded, validate=True)
114
+ ),
115
+ weights_only=True,
116
+ map_location=torch.device("cpu"),
117
+ )
118
+ tensor.to_dense()
119
+ results["b_bypass"] = True
120
+ except (RuntimeError, ValueError):
121
+ results["b_rejected"] = True
122
+ except threading.BrokenBarrierError:
123
+ results["b_barrier_broken"] = True
124
+
125
+ t_a = threading.Thread(target=thread_a_benign)
126
+ t_b = threading.Thread(target=thread_b_malicious)
127
+ t_a.start()
128
+ t_b.start()
129
+ t_a.join(timeout=10)
130
+ t_b.join(timeout=10)
131
+
132
+ assert results.get("a_loaded") is True
133
+ assert results.get("a_tensor_valid") is True
134
+ assert results.get("b_rejected") is True, (
135
+ "Malicious tensor was NOT rejected under concurrency"
136
+ )
137
+ assert results.get("b_bypass") is not True
138
+
139
+ def test_concurrent_loads_all_reject_invalid(self, model_config):
140
+ """Multiple concurrent malicious loads must all be rejected."""
141
+ malicious_encoded = _encode_tensor(_create_malicious_sparse_tensor())
142
+ num_workers = 4
143
+ rejected = []
144
+ bypassed = []
145
+
146
+ def attempt_load(_):
147
+ try:
148
+ safe_load_prompt_embeds(model_config, malicious_encoded)
149
+ bypassed.append(True)
150
+ except (RuntimeError, ValueError):
151
+ rejected.append(True)
152
+
153
+ with ThreadPoolExecutor(max_workers=num_workers) as pool:
154
+ list(pool.map(attempt_load, range(num_workers)))
155
+
156
+ assert len(rejected) == num_workers
157
+ assert len(bypassed) == 0
158
+
159
+ def test_concurrent_valid_loads_succeed(self, model_config):
160
+ """Concurrent valid loads must all succeed (no false rejections)."""
161
+ valid_encoded = _encode_tensor(_create_valid_dense_tensor())
162
+ num_workers = 4
163
+ successes = []
164
+ failures = []
165
+
166
+ def attempt_load(_):
167
+ try:
168
+ result = safe_load_prompt_embeds(model_config, valid_encoded)
169
+ successes.append(result.shape)
170
+ except Exception as e:
171
+ failures.append(str(e))
172
+
173
+ with ThreadPoolExecutor(max_workers=num_workers) as pool:
174
+ list(pool.map(attempt_load, range(num_workers)))
175
+
176
+ assert len(successes) == num_workers
177
+ assert len(failures) == 0
178
+
179
+
180
+class TestGlobalFlagRestoration:
181
+ """The global invariant flag must be correctly restored after use."""
182
+
183
+ def test_flag_restored_after_success(self, model_config):
184
+ initial = torch.sparse.check_sparse_tensor_invariants.is_enabled()
185
+ valid_encoded = _encode_tensor(_create_valid_dense_tensor())
186
+
187
+ safe_load_prompt_embeds(model_config, valid_encoded)
188
+
189
+ assert torch.sparse.check_sparse_tensor_invariants.is_enabled() == initial
190
+
191
+ def test_flag_restored_after_exception(self, model_config):
192
+ initial = torch.sparse.check_sparse_tensor_invariants.is_enabled()
193
+ malicious_encoded = _encode_tensor(_create_malicious_sparse_tensor())
194
+
195
+ with pytest.raises((RuntimeError, ValueError)):
196
+ safe_load_prompt_embeds(model_config, malicious_encoded)
197
+
198
+ assert torch.sparse.check_sparse_tensor_invariants.is_enabled() == initial
199
+
200
+ def test_flag_restored_after_concurrent_exceptions(self, model_config):
201
+ """Flag must be restored even when multiple threads raise."""
202
+ initial = torch.sparse.check_sparse_tensor_invariants.is_enabled()
203
+ malicious_encoded = _encode_tensor(_create_malicious_sparse_tensor())
204
+
205
+ def attempt_load(_):
206
+ with contextlib.suppress(RuntimeError, ValueError):
207
+ safe_load_prompt_embeds(model_config, malicious_encoded)
208
+
209
+ with ThreadPoolExecutor(max_workers=4) as pool:
210
+ list(pool.map(attempt_load, range(4)))
211
+
212
+ assert torch.sparse.check_sparse_tensor_invariants.is_enabled() == initial
213
+
214
+
215
+class TestCrossLoaderLockSharing:
216
+ """All loaders must share the same lock to prevent cross-loader races."""
217
+
218
+ def test_prompt_and_image_share_lock(self, model_config):
219
+ """Prompt and image loaders cannot run their sparse guards
220
+ concurrently."""
221
+ valid_prompt = _encode_tensor(_create_valid_dense_tensor())
222
+ valid_image = _encode_tensor(_create_valid_dense_tensor(hidden_size=10))
223
+
224
+ io_handler = ImageEmbeddingMediaIO()
225
+
226
+ def load_prompt():
227
+ safe_load_prompt_embeds(model_config, valid_prompt)
228
+
229
+ def load_image():
230
+ io_handler.load_base64("", valid_image.decode("utf-8"))
231
+
232
+ # If both threads could hold the lock simultaneously, the barrier
233
+ # would succeed. Since the lock serializes them, the barrier will
234
+ # time out for the second thread (proving mutual exclusion).
235
+ t1 = threading.Thread(target=load_prompt)
236
+ t2 = threading.Thread(target=load_image)
237
+
238
+ t1.start()
239
+ t2.start()
240
+ t1.join(timeout=10)
241
+ t2.join(timeout=10)
242
+
243
+ # The key assertion: both completed (barrier didn't deadlock due to
244
+ # serialization — one finishes before the other starts), AND invalid
245
+ # tensors are still rejected.
246
+ malicious = _encode_tensor(_create_malicious_sparse_tensor())
247
+ with pytest.raises((RuntimeError, ValueError)):
248
+ safe_load_prompt_embeds(model_config, malicious)
249
+ with pytest.raises((RuntimeError, ValueError)):
250
+ io_handler.load_base64("", malicious.decode("utf-8"))
251
+
252
+ def test_all_loaders_use_same_context_manager(self):
253
+ """Verify all modules reference the same context manager function."""
254
+ import vllm.multimodal.media.audio as audio_mod
255
+ import vllm.multimodal.media.image as image_mod
256
+ import vllm.renderers.embed_utils as embed_mod
257
+ from vllm.utils.sparse_utils import (
258
+ check_sparse_tensor_invariants_threadsafe as cm_from_utils,
259
+ )
260
+
261
+ assert embed_mod.check_sparse_tensor_invariants_threadsafe is cm_from_utils
262
+ assert image_mod.check_sparse_tensor_invariants_threadsafe is cm_from_utils
263
+ assert audio_mod.check_sparse_tensor_invariants_threadsafe is cm_from_utils
vllm/multimodal/media/audio.py
+3
-6
@@ -14,6 +14,7 @@ from vllm.logger import init_logger
14
from vllm.multimodal.audio import resample_audio_pyav
15
from vllm.utils.import_utils import PlaceholderModule
16
from vllm.utils.serial_utils import tensor2base64
17
+from vllm.utils.sparse_utils import check_sparse_tensor_invariants_threadsafe
18
19
from .base import MediaIO
20
@@ -282,9 +283,7 @@ class AudioEmbeddingMediaIO(MediaIO[torch.Tensor]):
283
284
def load_bytes(self, data: bytes) -> torch.Tensor:
285
buffer = BytesIO(data)
285
- # Enable sparse tensor integrity checks to prevent out-of-bounds
286
- # writes from maliciously crafted tensors
287
- with torch.sparse.check_sparse_tensor_invariants():
286
+ with check_sparse_tensor_invariants_threadsafe():
287
tensor = torch.load(buffer, weights_only=True)
288
return tensor.to_dense()
289
@@ -292,9 +291,7 @@ class AudioEmbeddingMediaIO(MediaIO[torch.Tensor]):
291
return self.load_bytes(pybase64.b64decode(data, validate=True))
292
293
def load_file(self, filepath: Path) -> torch.Tensor:
295
- # Enable sparse tensor integrity checks to prevent out-of-bounds
296
- # writes from maliciously crafted tensors
297
- with torch.sparse.check_sparse_tensor_invariants():
294
+ with check_sparse_tensor_invariants_threadsafe():
295
tensor = torch.load(filepath, weights_only=True)
296
return tensor.to_dense()
297
vllm/multimodal/media/image.py
+3
-4
@@ -11,6 +11,7 @@ from PIL import Image
11
12
import vllm.envs as envs
13
from vllm.utils.serial_utils import tensor2base64
14
+from vllm.utils.sparse_utils import check_sparse_tensor_invariants_threadsafe
15
16
from ..image import convert_image_mode, normalize_image, rgba_to_rgb
17
from .base import MediaIO, MediaWithBytes
@@ -123,9 +124,7 @@ class ImageEmbeddingMediaIO(MediaIO[torch.Tensor]):
124
125
def _load_pickled_torch(self, data: bytes) -> torch.Tensor:
126
buffer = BytesIO(data)
126
- # Enable sparse tensor integrity checks to prevent out-of-bounds
127
- # writes from maliciously crafted tensors
128
- with torch.sparse.check_sparse_tensor_invariants():
127
+ with check_sparse_tensor_invariants_threadsafe():
128
tensor = torch.load(buffer, weights_only=True)
129
return tensor.to_dense()
130
@@ -146,7 +145,7 @@ class ImageEmbeddingMediaIO(MediaIO[torch.Tensor]):
145
if filepath.suffix == ".npy":
146
return torch.from_numpy(np.load(filepath))
147
149
- with torch.sparse.check_sparse_tensor_invariants():
148
+ with check_sparse_tensor_invariants_threadsafe():
149
tensor = torch.load(filepath, weights_only=True)
150
return tensor.to_dense()
151
vllm/renderers/embed_utils.py
+2
-3
@@ -8,6 +8,7 @@ import torch
8
9
from vllm.exceptions import VLLMValidationError
10
from vllm.utils.async_utils import make_async
11
+from vllm.utils.sparse_utils import check_sparse_tensor_invariants_threadsafe
12
13
if TYPE_CHECKING:
14
from vllm.config import ModelConfig
@@ -23,9 +24,7 @@ def safe_load_prompt_embeds(
24
parameter="prompt_embeds",
25
)
26
26
- # Enable sparse tensor integrity checks to prevent out-of-bounds
27
- # writes from maliciously crafted tensors
28
- with torch.sparse.check_sparse_tensor_invariants():
27
+ with check_sparse_tensor_invariants_threadsafe():
28
tensor = torch.load(
29
BytesIO(pybase64.b64decode(embed, validate=True)),
30
weights_only=True,
vllm/utils/sparse_utils.py
new
+27
@@ -0,0 +1,27 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Thread-safe wrapper for sparse tensor invariant validation.
5
+
6
+PyTorch's `torch.sparse.check_sparse_tensor_invariants()` context manager
7
+manipulates a **process-global** flag (save/enable/restore). When multiple
8
+embedding-load operations run concurrently on a thread-pool executor, one
9
+context can restore the flag to `False` while another thread is still inside
10
+its guard, bypassing the invariant check.
11
+
12
+All call sites MUST use `check_sparse_tensor_invariants_threadsafe()`
13
+which serializes access behind a lock.
14
+"""
15
+
16
+import contextlib
17
+import threading
18
+
19
+import torch
20
+
21
+_SPARSE_LOAD_LOCK = threading.Lock()
22
+
23
+
24
+@contextlib.contextmanager
25
+def check_sparse_tensor_invariants_threadsafe():
26
+ with _SPARSE_LOAD_LOCK, torch.sparse.check_sparse_tensor_invariants():
27
+ yield