| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests verify that malicious sparse tensors are rejected before they can trigger |
| 4 | out-of-bounds memory writes, or an unbounded allocation, during to_dense() |
| 5 | operations. |
| 6 | |
| 7 | A sparse tensor carries its own declared shape, so the two failure modes are |
| 8 | independent: invalid *indices* are what the invariant checks catch, while a |
| 9 | valid tensor with an enormous declared *shape* passes those checks and is |
| 10 | bounded by `VLLM_MAX_EMBED_DECODE_BYTES` instead. |
| 11 | """ |
| 12 | |
| 13 | import io |
| 14 | |
| 15 | import numpy as np |
| 16 | import pybase64 as base64 |
| 17 | import pytest |
| 18 | import torch |
| 19 | |
| 20 | import vllm.envs as envs |
| 21 | from vllm.exceptions import VLLMValidationError |
| 22 | from vllm.multimodal.media import AudioEmbeddingMediaIO, ImageEmbeddingMediaIO |
| 23 | from vllm.multimodal.media.video import VideoEmbeddingMediaIO |
| 24 | from vllm.renderers.embed_utils import safe_load_prompt_embeds |
| 25 | from vllm.utils.sparse_utils import safe_to_dense |
| 26 | |
| 27 | |
| 28 | @pytest.fixture |
| 29 | def model_config(): |
| 30 | """Mock ModelConfig for testing.""" |
| 31 | from vllm.config import ModelConfig |
| 32 | |
| 33 | return ModelConfig( |
| 34 | model="facebook/opt-125m", |
| 35 | tokenizer="facebook/opt-125m", |
| 36 | tokenizer_mode="auto", |
| 37 | trust_remote_code=False, |
| 38 | dtype="float32", |
| 39 | seed=0, |
| 40 | enable_prompt_embeds=True, # Required for prompt embeds tests |
| 41 | ) |
| 42 | |
| 43 | |
| 44 | def _encode_tensor(tensor: torch.Tensor) -> bytes: |
| 45 | """Helper to encode a tensor as base64 bytes.""" |
| 46 | buffer = io.BytesIO() |
| 47 | torch.save(tensor, buffer) |
| 48 | buffer.seek(0) |
| 49 | return base64.b64encode(buffer.read()) |
| 50 | |
| 51 | |
| 52 | def _create_malicious_sparse_tensor() -> torch.Tensor: |
| 53 | """Create a malicious sparse COO tensor with out-of-bounds indices. |
| 54 | |
| 55 | This tensor has indices that point beyond the declared shape, which would |
| 56 | cause an out-of-bounds write when converted to dense format without |
| 57 | validation. |
| 58 | """ |
| 59 | # Create a 3x3 sparse tensor but with indices pointing to (10, 10) |
| 60 | indices = torch.tensor([[10], [10]]) # Out of bounds for 3x3 shape |
| 61 | values = torch.tensor([1.0]) |
| 62 | shape = (3, 3) |
| 63 | |
| 64 | # Create sparse tensor (this will be invalid). Pass `check_invariants=False` |
| 65 | # explicitly so this fixture is robust to process-wide invariant-check state |
| 66 | # left enabled by other tests (the global flag isn't thread-local, and |
| 67 | # concurrent users of the `check_sparse_tensor_invariants` context manager |
| 68 | # can leak the "enabled" state across tests). |
| 69 | sparse_tensor = torch.sparse_coo_tensor( |
| 70 | indices, values, shape, dtype=torch.float32, check_invariants=False |
| 71 | ) |
| 72 | return sparse_tensor |
| 73 | |
| 74 | |
| 75 | def _create_valid_sparse_tensor() -> torch.Tensor: |
| 76 | """Create a valid sparse COO tensor for baseline testing.""" |
| 77 | indices = torch.tensor([[0, 1, 2], [0, 1, 2]]) |
| 78 | values = torch.tensor([1.0, 2.0, 3.0]) |
| 79 | shape = (3, 3) |
| 80 | |
| 81 | sparse_tensor = torch.sparse_coo_tensor(indices, values, shape, dtype=torch.float32) |
| 82 | return sparse_tensor |
| 83 | |
| 84 | |
| 85 | def _create_valid_dense_tensor() -> torch.Tensor: |
| 86 | """Create a valid dense tensor for baseline testing.""" |
| 87 | return torch.randn(10, 768, dtype=torch.float32) # (seq_len, hidden_size) |
| 88 | |
| 89 | |
| 90 | class TestPromptEmbedsValidation: |
| 91 | """Test sparse tensor validation in prompt embeddings (Completions API).""" |
| 92 | |
| 93 | def test_valid_dense_tensor_accepted(self, model_config): |
| 94 | """Baseline: Valid dense tensors should work normally.""" |
| 95 | valid_tensor = _create_valid_dense_tensor() |
| 96 | encoded = _encode_tensor(valid_tensor) |
| 97 | |
| 98 | # Should not raise any exception |
| 99 | result = safe_load_prompt_embeds(model_config, encoded) |
| 100 | assert result.shape == valid_tensor.shape |
| 101 | |
| 102 | def test_valid_sparse_tensor_accepted(self): |
| 103 | """Baseline: Valid sparse tensors should load successfully.""" |
| 104 | io_handler = ImageEmbeddingMediaIO() |
| 105 | |
| 106 | valid_sparse = _create_valid_sparse_tensor() |
| 107 | encoded = _encode_tensor(valid_sparse) |
| 108 | |
| 109 | # Should not raise any exception (sparse tensors remain sparse) |
| 110 | result = io_handler.load_base64("", encoded.decode("utf-8")) |
| 111 | assert result.shape == valid_sparse.shape |
| 112 | |
| 113 | def test_malicious_sparse_tensor_rejected(self, model_config): |
| 114 | """Security: Malicious sparse tensors should be rejected.""" |
| 115 | malicious_tensor = _create_malicious_sparse_tensor() |
| 116 | encoded = _encode_tensor(malicious_tensor) |
| 117 | |
| 118 | # The invariant check raises RuntimeError from inside torch.load; |
| 119 | # safe_load_prompt_embeds now reports that as a client error so the |
| 120 | # caller gets a 400 rather than the 500 a bare RuntimeError produced. |
| 121 | with pytest.raises(VLLMValidationError) as exc_info: |
| 122 | safe_load_prompt_embeds(model_config, encoded) |
| 123 | |
| 124 | # Error should indicate sparse tensor validation failure |
| 125 | error_msg = str(exc_info.value).lower() |
| 126 | assert "sparse" in error_msg or "index" in error_msg or "bounds" in error_msg |
| 127 | |
| 128 | def test_extremely_large_indices_rejected(self, model_config): |
| 129 | """Security: Sparse tensors with extremely large indices should be rejected.""" |
| 130 | # Create tensor with indices far beyond reasonable bounds |
| 131 | indices = torch.tensor([[999999], [999999]]) |
| 132 | values = torch.tensor([1.0]) |
| 133 | shape = (10, 10) |
| 134 | |
| 135 | malicious_tensor = torch.sparse_coo_tensor( |
| 136 | indices, values, shape, dtype=torch.float32, check_invariants=False |
| 137 | ) |
| 138 | encoded = _encode_tensor(malicious_tensor) |
| 139 | |
| 140 | with pytest.raises(VLLMValidationError): |
| 141 | safe_load_prompt_embeds(model_config, encoded) |
| 142 | |
| 143 | def test_negative_indices_rejected(self, model_config): |
| 144 | """Security: Sparse tensors with negative indices should be rejected.""" |
| 145 | # Create tensor with negative indices |
| 146 | indices = torch.tensor([[-1], [-1]]) |
| 147 | values = torch.tensor([1.0]) |
| 148 | shape = (10, 10) |
| 149 | |
| 150 | malicious_tensor = torch.sparse_coo_tensor( |
| 151 | indices, values, shape, dtype=torch.float32, check_invariants=False |
| 152 | ) |
| 153 | encoded = _encode_tensor(malicious_tensor) |
| 154 | |
| 155 | with pytest.raises(VLLMValidationError): |
| 156 | safe_load_prompt_embeds(model_config, encoded) |
| 157 | |
| 158 | def test_hidden_size_mismatch_rejected(self, model_config): |
| 159 | """Tensors whose trailing dim doesn't match the model's hidden_size |
| 160 | must be rejected at parse time.""" |
| 161 | # opt-125m has hidden_size=768, passing 512 triggers the check. |
| 162 | wrong_hidden = torch.randn(10, 512, dtype=torch.float32) |
| 163 | encoded = _encode_tensor(wrong_hidden) |
| 164 | |
| 165 | with pytest.raises(VLLMValidationError, match="hidden_size"): |
| 166 | safe_load_prompt_embeds(model_config, encoded) |
| 167 | |
| 168 | def test_float_dtype_mismatch_cast_to_model_dtype(self, model_config): |
| 169 | """Tensors whose dtype doesn't match the model's dtype but are still |
| 170 | floating-point are cast, since API clients generally can't know the |
| 171 | server's `--dtype` setting ahead of time.""" |
| 172 | # Fixture pins model dtype to float32, upload a bfloat16 tensor. |
| 173 | mismatched_float = torch.randn(10, 768, dtype=torch.bfloat16) |
| 174 | encoded = _encode_tensor(mismatched_float) |
| 175 | |
| 176 | result = safe_load_prompt_embeds(model_config, encoded) |
| 177 | |
| 178 | assert result.dtype == torch.float32 |
| 179 | assert result.shape == mismatched_float.shape |
| 180 | |
| 181 | def test_non_float_dtype_rejected(self, model_config): |
| 182 | """Non-floating-point dtypes cannot be safely cast for embeddings |
| 183 | (e.g. integer tensors almost certainly indicate caller confusion), |
| 184 | so they are rejected at parse time.""" |
| 185 | non_float = torch.randint(0, 100, (10, 768), dtype=torch.int32) |
| 186 | encoded = _encode_tensor(non_float) |
| 187 | |
| 188 | with pytest.raises(VLLMValidationError, match="floating-point"): |
| 189 | safe_load_prompt_embeds(model_config, encoded) |
| 190 | |
| 191 | def test_non_2d_tensor_rejected(self, model_config): |
| 192 | """Tensors that aren't 2D (even after squeezing a leading dim) |
| 193 | must be rejected with a clear error.""" |
| 194 | # A 1D tensor cannot be interpreted as (num_tokens, hidden_size). |
| 195 | bad = torch.randn(768, dtype=torch.float32) |
| 196 | encoded = _encode_tensor(bad) |
| 197 | |
| 198 | with pytest.raises(VLLMValidationError, match="2D tensor"): |
| 199 | safe_load_prompt_embeds(model_config, encoded) |
| 200 | |
| 201 | def test_non_tensor_payload_rejected(self, model_config): |
| 202 | """Deserializing to a non-Tensor object must raise a clear error |
| 203 | instead of propagating an AssertionError.""" |
| 204 | # `torch.save` will serialize a plain dict; `weights_only=True` allows |
| 205 | # loading built-in containers, so this exercises the isinstance check. |
| 206 | buffer = io.BytesIO() |
| 207 | torch.save({"not": "a tensor"}, buffer) |
| 208 | buffer.seek(0) |
| 209 | encoded = base64.b64encode(buffer.read()) |
| 210 | |
| 211 | with pytest.raises(VLLMValidationError, match="torch.Tensor"): |
| 212 | safe_load_prompt_embeds(model_config, encoded) |
| 213 | |
| 214 | |
| 215 | class TestImageEmbedsValidation: |
| 216 | """Test sparse tensor validation in image embeddings (Chat API).""" |
| 217 | |
| 218 | def test_valid_dense_tensor_accepted(self): |
| 219 | """Baseline: Valid dense tensors should work normally.""" |
| 220 | io_handler = ImageEmbeddingMediaIO() |
| 221 | |
| 222 | valid_tensor = _create_valid_dense_tensor() |
| 223 | encoded = _encode_tensor(valid_tensor) |
| 224 | |
| 225 | # Should not raise any exception |
| 226 | result = io_handler.load_base64("", encoded.decode("utf-8")) |
| 227 | assert result.shape == valid_tensor.shape |
| 228 | |
| 229 | def test_valid_sparse_tensor_accepted(self): |
| 230 | """Baseline: Valid sparse tensors should load successfully.""" |
| 231 | io_handler = AudioEmbeddingMediaIO() |
| 232 | |
| 233 | valid_sparse = _create_valid_sparse_tensor() |
| 234 | encoded = _encode_tensor(valid_sparse) |
| 235 | |
| 236 | # Should not raise any exception (sparse tensors remain sparse) |
| 237 | result = io_handler.load_base64("", encoded.decode("utf-8")) |
| 238 | assert result.shape == valid_sparse.shape |
| 239 | |
| 240 | def test_malicious_sparse_tensor_rejected(self): |
| 241 | """Security: Malicious sparse tensors should be rejected.""" |
| 242 | io_handler = ImageEmbeddingMediaIO() |
| 243 | |
| 244 | malicious_tensor = _create_malicious_sparse_tensor() |
| 245 | encoded = _encode_tensor(malicious_tensor) |
| 246 | |
| 247 | # Should raise RuntimeError due to invalid sparse tensor |
| 248 | with pytest.raises((RuntimeError, ValueError)) as exc_info: |
| 249 | io_handler.load_base64("", encoded.decode("utf-8")) |
| 250 | |
| 251 | error_msg = str(exc_info.value).lower() |
| 252 | assert "sparse" in error_msg or "index" in error_msg or "bounds" in error_msg |
| 253 | |
| 254 | def test_load_bytes_validates(self): |
| 255 | """Security: Validation should also work for load_bytes method.""" |
| 256 | io_handler = ImageEmbeddingMediaIO() |
| 257 | |
| 258 | malicious_tensor = _create_malicious_sparse_tensor() |
| 259 | buffer = io.BytesIO() |
| 260 | torch.save(malicious_tensor, buffer) |
| 261 | buffer.seek(0) |
| 262 | |
| 263 | with pytest.raises((RuntimeError, ValueError)): |
| 264 | io_handler.load_bytes(buffer.read()) |
| 265 | |
| 266 | def test_valid_numpy_tensor_accepted(self): |
| 267 | """Numpy .npy format should load and return correct tensor.""" |
| 268 | io_handler = ImageEmbeddingMediaIO() |
| 269 | |
| 270 | arr = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32) |
| 271 | buf = io.BytesIO() |
| 272 | np.save(buf, arr) |
| 273 | encoded = base64.b64encode(buf.getvalue()).decode("utf-8") |
| 274 | |
| 275 | result = io_handler.load_base64("", encoded) |
| 276 | assert isinstance(result, torch.Tensor) |
| 277 | assert result.shape == torch.Size([2, 3]) |
| 278 | assert result.dtype == torch.float32 |
| 279 | assert torch.allclose(result, torch.from_numpy(arr)) |
| 280 | |
| 281 | def test_numpy_int32_tensor_accepted(self): |
| 282 | """Numpy int32 arrays should round-trip correctly.""" |
| 283 | io_handler = ImageEmbeddingMediaIO() |
| 284 | |
| 285 | arr = np.arange(280, dtype=np.int32) |
| 286 | buf = io.BytesIO() |
| 287 | np.save(buf, arr) |
| 288 | encoded = base64.b64encode(buf.getvalue()).decode("utf-8") |
| 289 | |
| 290 | result = io_handler.load_base64("", encoded) |
| 291 | assert result.dtype == torch.int32 |
| 292 | assert result.shape == torch.Size([280]) |
| 293 | assert (result == torch.from_numpy(arr)).all() |
| 294 | |
| 295 | def test_load_file_numpy_tensor_accepted(self, tmp_path): |
| 296 | """Numpy .npy files should load correctly via load_file.""" |
| 297 | io_handler = ImageEmbeddingMediaIO() |
| 298 | |
| 299 | arr = np.array([[1.5, 2.5], [3.5, 4.5]], dtype=np.float32) |
| 300 | npy_path = tmp_path / "image_embeds.npy" |
| 301 | np.save(npy_path, arr) |
| 302 | |
| 303 | result = io_handler.load_file(npy_path) |
| 304 | assert isinstance(result, torch.Tensor) |
| 305 | assert result.shape == torch.Size([2, 2]) |
| 306 | assert result.dtype == torch.float32 |
| 307 | assert torch.allclose(result, torch.from_numpy(arr)) |
| 308 | |
| 309 | |
| 310 | class TestAudioEmbedsValidation: |
| 311 | """Test sparse tensor validation in audio embeddings (Chat API).""" |
| 312 | |
| 313 | def test_valid_dense_tensor_accepted(self): |
| 314 | """Baseline: Valid dense tensors should work normally.""" |
| 315 | io_handler = AudioEmbeddingMediaIO() |
| 316 | |
| 317 | valid_tensor = _create_valid_dense_tensor() |
| 318 | encoded = _encode_tensor(valid_tensor) |
| 319 | |
| 320 | # Should not raise any exception |
| 321 | result = io_handler.load_base64("", encoded.decode("utf-8")) |
| 322 | assert result.shape == valid_tensor.shape |
| 323 | |
| 324 | def test_valid_sparse_tensor_accepted(self): |
| 325 | """Baseline: Valid sparse tensors should be converted successfully.""" |
| 326 | io_handler = AudioEmbeddingMediaIO() |
| 327 | |
| 328 | valid_sparse = _create_valid_sparse_tensor() |
| 329 | encoded = _encode_tensor(valid_sparse) |
| 330 | |
| 331 | # Should not raise any exception |
| 332 | result = io_handler.load_base64("", encoded.decode("utf-8")) |
| 333 | assert result.is_sparse is False |
| 334 | |
| 335 | def test_malicious_sparse_tensor_rejected(self): |
| 336 | """Security: Malicious sparse tensors should be rejected.""" |
| 337 | io_handler = AudioEmbeddingMediaIO() |
| 338 | |
| 339 | malicious_tensor = _create_malicious_sparse_tensor() |
| 340 | encoded = _encode_tensor(malicious_tensor) |
| 341 | |
| 342 | # Should raise RuntimeError due to invalid sparse tensor |
| 343 | with pytest.raises((RuntimeError, ValueError)) as exc_info: |
| 344 | io_handler.load_base64("", encoded.decode("utf-8")) |
| 345 | |
| 346 | error_msg = str(exc_info.value).lower() |
| 347 | assert "sparse" in error_msg or "index" in error_msg or "bounds" in error_msg |
| 348 | |
| 349 | def test_load_bytes_validates(self): |
| 350 | """Security: Validation should also work for load_bytes method.""" |
| 351 | io_handler = AudioEmbeddingMediaIO() |
| 352 | |
| 353 | malicious_tensor = _create_malicious_sparse_tensor() |
| 354 | buffer = io.BytesIO() |
| 355 | torch.save(malicious_tensor, buffer) |
| 356 | buffer.seek(0) |
| 357 | |
| 358 | with pytest.raises((RuntimeError, ValueError)): |
| 359 | io_handler.load_bytes(buffer.read()) |
| 360 | |
| 361 | |
| 362 | class TestSparseTensorValidationIntegration: |
| 363 | """These tests verify the complete attack chain is blocked at all entry points.""" |
| 364 | |
| 365 | def test_attack_scenario_completions_api(self, model_config): |
| 366 | """Simulate a complete attack through the Completions API. |
| 367 | |
| 368 | Attack scenario: |
| 369 | 1. Attacker crafts malicious sparse tensor |
| 370 | 2. Encodes it as base64 |
| 371 | 3. Sends to /v1/completions with prompt_embeds parameter |
| 372 | 4. Server should reject before memory corruption occurs |
| 373 | """ |
| 374 | # Step 1-2: Attacker creates malicious payload |
| 375 | attack_payload = _encode_tensor(_create_malicious_sparse_tensor()) |
| 376 | |
| 377 | # Step 3-4: Server processes and should reject, as a client error |
| 378 | with pytest.raises(VLLMValidationError): |
| 379 | safe_load_prompt_embeds(model_config, attack_payload) |
| 380 | |
| 381 | def test_attack_scenario_chat_api_image(self): |
| 382 | """Simulate attack through Chat API with image_embeds. |
| 383 | |
| 384 | Verifies the image embeddings path is protected. |
| 385 | """ |
| 386 | io_handler = ImageEmbeddingMediaIO() |
| 387 | attack_payload = _encode_tensor(_create_malicious_sparse_tensor()) |
| 388 | |
| 389 | with pytest.raises((RuntimeError, ValueError)): |
| 390 | io_handler.load_base64("", attack_payload.decode("utf-8")) |
| 391 | |
| 392 | def test_attack_scenario_chat_api_audio(self): |
| 393 | """Simulate attack through Chat API with audio_embeds. |
| 394 | |
| 395 | Verifies the audio embeddings path is protected. |
| 396 | """ |
| 397 | io_handler = AudioEmbeddingMediaIO() |
| 398 | attack_payload = _encode_tensor(_create_malicious_sparse_tensor()) |
| 399 | |
| 400 | with pytest.raises((RuntimeError, ValueError)): |
| 401 | io_handler.load_base64("", attack_payload.decode("utf-8")) |
| 402 | |
| 403 | |
| 404 | # The real attack payload: 4 TiB of float32 once densified, from a single |
| 405 | # stored element. Only ever measured, never decoded. |
| 406 | BOMB_SHAPE = (2**20, 2**20) |
| 407 | |
| 408 | # What the size-limit tests below actually build. Each of them also lowers |
| 409 | # `VLLM_MAX_EMBED_DECODE_BYTES` to `TEST_LIMIT_BYTES`, so the guard still has |
| 410 | # something to reject while nothing here can allocate more than 4 MiB even if |
| 411 | # the guard regressed. Keeping the 4 TiB shape in these tests would make the |
| 412 | # outcome depend on the runner: with `vm.overcommit_memory=0` the allocator |
| 413 | # refuses it outright, but with `vm.overcommit_memory=1` the malloc succeeds |
| 414 | # and the zero-fill is what fails, which would OOM the test process. |
| 415 | OVERSIZED_SHAPE = (1024, 1024) # 4 MiB dense |
| 416 | TEST_LIMIT_BYTES = 4096 |
| 417 | |
| 418 | EMBEDDING_MEDIA_IO = [ |
| 419 | (ImageEmbeddingMediaIO, "image_embeds"), |
| 420 | (AudioEmbeddingMediaIO, "audio_embeds"), |
| 421 | (VideoEmbeddingMediaIO, "video_embeds"), |
| 422 | ] |
| 423 | |
| 424 | |
| 425 | def _create_oversized_sparse_tensor(shape: tuple[int, int]) -> torch.Tensor: |
| 426 | """A *valid* sparse tensor with one stored element and a declared `shape`. |
| 427 | |
| 428 | Indices are in bounds, so `check_sparse_tensor_invariants()` accepts it. |
| 429 | That is the point: the invariant checks are not what stops this one. |
| 430 | """ |
| 431 | indices = torch.tensor([[0], [0]]) |
| 432 | values = torch.tensor([1.0]) |
| 433 | return torch.sparse_coo_tensor(indices, values, shape, dtype=torch.float32) |
| 434 | |
| 435 | |
| 436 | @pytest.fixture |
| 437 | def small_decode_limit(monkeypatch: pytest.MonkeyPatch) -> None: |
| 438 | monkeypatch.setenv("VLLM_MAX_EMBED_DECODE_BYTES", str(TEST_LIMIT_BYTES)) |
| 439 | |
| 440 | |
| 441 | class TestEmbeddingDecodeSizeLimit: |
| 442 | """Bound the dense allocation a payload can declare. |
| 443 | |
| 444 | `weights_only=True` stops code execution and the invariant checks stop |
| 445 | out-of-bounds indices, but neither bounds the declared shape. |
| 446 | """ |
| 447 | |
| 448 | def test_oversized_sparse_tensor_rejected(self, small_decode_limit): |
| 449 | with pytest.raises(VLLMValidationError) as exc_info: |
| 450 | safe_to_dense( |
| 451 | _create_oversized_sparse_tensor(OVERSIZED_SHAPE), |
| 452 | parameter="prompt_embeds", |
| 453 | ) |
| 454 | |
| 455 | assert "VLLM_MAX_EMBED_DECODE_BYTES" in str(exc_info.value) |
| 456 | assert "prompt_embeds" in str(exc_info.value) |
| 457 | |
| 458 | def test_default_limit_would_reject_the_real_payload(self): |
| 459 | """The shipped default has to bound the actual attack, not just the |
| 460 | shrunken shape used above. Checked arithmetically so no test ever asks |
| 461 | the allocator for 4 TiB.""" |
| 462 | dense_bytes = BOMB_SHAPE[0] * BOMB_SHAPE[1] * 4 |
| 463 | assert dense_bytes > envs.VLLM_MAX_EMBED_DECODE_BYTES > 0 |
| 464 | |
| 465 | def test_attack_payload_is_tiny(self): |
| 466 | """The payload really is small - this is amplification, not an upload.""" |
| 467 | assert len(_encode_tensor(_create_oversized_sparse_tensor(BOMB_SHAPE))) < 4096 |
| 468 | |
| 469 | def test_valid_tensors_still_densify(self): |
| 470 | assert safe_to_dense( |
| 471 | _create_valid_dense_tensor(), parameter="prompt_embeds" |
| 472 | ).shape == (10, 768) |
| 473 | |
| 474 | dense = safe_to_dense(_create_valid_sparse_tensor(), parameter="image_embeds") |
| 475 | assert dense.shape == (3, 3) |
| 476 | assert not dense.is_sparse |
| 477 | |
| 478 | def test_limit_is_configurable(self, monkeypatch: pytest.MonkeyPatch): |
| 479 | # 32 x 32 float32 == 4096 bytes, exactly at the limit. |
| 480 | monkeypatch.setenv("VLLM_MAX_EMBED_DECODE_BYTES", str(TEST_LIMIT_BYTES)) |
| 481 | |
| 482 | at_limit = _create_oversized_sparse_tensor((32, 32)) |
| 483 | assert safe_to_dense(at_limit, parameter="image_embeds").shape == (32, 32) |
| 484 | |
| 485 | over_limit = _create_oversized_sparse_tensor((33, 32)) |
| 486 | with pytest.raises(VLLMValidationError): |
| 487 | safe_to_dense(over_limit, parameter="image_embeds") |
| 488 | |
| 489 | def test_limit_can_be_disabled(self, monkeypatch: pytest.MonkeyPatch): |
| 490 | monkeypatch.setenv("VLLM_MAX_EMBED_DECODE_BYTES", "0") |
| 491 | tensor = _create_oversized_sparse_tensor(OVERSIZED_SHAPE) |
| 492 | |
| 493 | assert safe_to_dense(tensor, parameter="image_embeds").shape == OVERSIZED_SHAPE |
| 494 | |
| 495 | @pytest.mark.parametrize("io_cls,parameter", EMBEDDING_MEDIA_IO) |
| 496 | def test_media_io_rejects_oversized_base64( |
| 497 | self, io_cls, parameter, small_decode_limit |
| 498 | ): |
| 499 | encoded = _encode_tensor(_create_oversized_sparse_tensor(OVERSIZED_SHAPE)) |
| 500 | |
| 501 | with pytest.raises(VLLMValidationError) as exc_info: |
| 502 | io_cls().load_base64("", encoded.decode("utf-8")) |
| 503 | |
| 504 | assert parameter in str(exc_info.value) |
| 505 | |
| 506 | @pytest.mark.parametrize("io_cls,parameter", EMBEDDING_MEDIA_IO) |
| 507 | def test_media_io_rejects_oversized_file( |
| 508 | self, io_cls, parameter, tmp_path, small_decode_limit |
| 509 | ): |
| 510 | path = tmp_path / "embeds.pt" |
| 511 | torch.save(_create_oversized_sparse_tensor(OVERSIZED_SHAPE), path) |
| 512 | |
| 513 | with pytest.raises(VLLMValidationError): |
| 514 | io_cls().load_file(path) |
| 515 | |
| 516 | @pytest.mark.parametrize("io_cls,parameter", EMBEDDING_MEDIA_IO) |
| 517 | def test_media_io_rejects_non_tensor_payload(self, io_cls, parameter): |
| 518 | """These three had no `isinstance` check, so a non-tensor payload |
| 519 | surfaced as an `AttributeError` (500) rather than a 4xx.""" |
| 520 | buffer = io.BytesIO() |
| 521 | torch.save({"not": "a tensor"}, buffer) |
| 522 | encoded = base64.b64encode(buffer.getvalue()).decode("utf-8") |
| 523 | |
| 524 | with pytest.raises(VLLMValidationError) as exc_info: |
| 525 | io_cls().load_base64("", encoded) |
| 526 | |
| 527 | assert "torch.Tensor" in str(exc_info.value) |
| 528 | |
| 529 | @pytest.mark.parametrize("io_cls,parameter", EMBEDDING_MEDIA_IO) |
| 530 | def test_media_io_still_accepts_valid_embeddings(self, io_cls, parameter): |
| 531 | tensor = torch.randn(4, 16, dtype=torch.float32) |
| 532 | encoded = _encode_tensor(tensor) |
| 533 | |
| 534 | assert torch.equal(io_cls().load_base64("", encoded.decode("utf-8")), tensor) |