[Bugfix] Handle DoS Overflow Cases (#6598)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py

@@ -157,8 +157,8 @@ def test_speech_missing_required_fields(omni_server: OmniServer, online_client:
                 "ref_audio": None,
                 "ref_text": None,
             },
-            ("CustomVoice", "does not support"),
-            id="customvoice_task_not_supported",
+            ("Invalid voice", "vivian"),
+            id="customvoice_invalid_voice",
         ),
         pytest.param(
             {"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid"
@@ -181,13 +181,17 @@ def test_speech_missing_required_fields(omni_server: OmniServer, online_client:
             ("mutually exclusive", "speaker_embedding", "ref_audio"),
             id="speaker_embedding_with_ref_audio_mutually_exclusive",
         ),
-        pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"),
+        pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "greater_than_equal"), id="max_new_tokens_below_min"),
         pytest.param(
             {"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
             ("max_new_tokens", "exceed 4096"),
             id="max_new_tokens_above_max",
         ),
-        pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"),
+        pytest.param(
+            {"seed": -(2**63) - 1},
+            ("seed", "greater_than_equal", "-9223372036854775808"),
+            id="seed_below_min",
+        ),
         pytest.param({"seed": 2**63}, ("seed", "less_than_equal", "9223372036854775807"), id="seed_above_max"),
         pytest.param(
             {"initial_codec_chunk_frames": -1},
@@ -275,7 +279,9 @@ def test_speech_invalid_field_values(
         pytest.param({"input": "", "response_format": "wav"}, ("input", "empty"), id="input_empty"),
         pytest.param({"input": "   ", "response_format": "wav"}, ("input", "empty"), id="input_whitespace_only"),
         pytest.param(
-            {"input": "Hello world.", "max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"
+            {"input": "Hello world.", "max_new_tokens": 0},
+            ("max_new_tokens", "greater_than_equal"),
+            id="max_new_tokens_below_min",
         ),
         pytest.param(
             {
@@ -554,14 +560,14 @@ def test_speech_batch_missing_items(omni_server: OmniServer, online_client: Onli
         pytest.param(
             "batch",
             {"max_new_tokens": 0},
-            ("max_new_tokens", "least 1"),
+            ("max_new_tokens", "greater_than_equal"),
             id="batch_max_new_tokens_below_min",
             marks=_SKIP_ISSUE_3649,
         ),
         pytest.param(
             "item",
             {"max_new_tokens": 0},
-            ("max_new_tokens", "least 1"),
+            ("max_new_tokens", "greater_than_equal"),
             id="item_max_new_tokens_below_min",
             marks=_SKIP_ISSUE_3649,
         ),

tests/engine/test_orchestrator_error_handling.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for error propagation paths within the Orchestrator.
 
 Covers:
@@ -16,7 +16,9 @@ from types import SimpleNamespace
 
 import janus
 import pytest
+from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.exceptions import EngineDeadError
+from vllm.v1.serial_utils import MsgpackEncoder
 
 from vllm_omni.engine.messages import (
     AddCompanionRequestMessage,
@@ -46,8 +48,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def _sampling_params(max_tokens: int = 4):
-    from vllm.sampling_params import SamplingParams
-
     return SamplingParams(max_tokens=max_tokens)
 
 
@@ -188,6 +188,45 @@ class _UnavailableDiffusionStageClient(FakeStageClient):
         raise StageUnavailableError(f"stage {self.stage_id} has no live replica")
 
 
+class _OverflowOnEncodeStageClient(FakeStageClient):
+    """Msgpack-encodes the request's sampling params on dispatch, as the real
+    StageEngineCoreClient does. We use this to make sure overflow sampling params
+    do not break the orchestrator."""
+
+    async def add_request_async(self, request, *args, **kwargs) -> None:
+        MsgpackEncoder().encode(request.sampling_params)
+        await super().add_request_async(request, *args, **kwargs)  # unreachable
+
+
+@pytest.mark.asyncio
+async def test_encode_overflow_should_fail_request_not_loop(orchestrator_factory) -> None:
+    """Regression test for over-range seed handling."""
+    stage0 = _OverflowOnEncodeStageClient(stage_type="llm", final_output=True)
+    orchestrator_fixture = orchestrator_factory([stage0])
+
+    # seed > 2**63-1 cannot be msgpack-serialized -> the malicious input under test.
+    overflow_params = SamplingParams(seed=2**70)
+
+    await _enqueue_add_request(
+        orchestrator_fixture,
+        request_id="req-overflow",
+        prompt=SimpleNamespace(
+            request_id="req-overflow",
+            prompt_token_ids=[1, 2],
+            sampling_params=overflow_params,
+        ),
+        original_prompt={"prompt": "hi"},
+        sampling_params_list=[overflow_params],
+        final_stage_id=0,
+    )
+
+    error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-overflow")
+    assert error_msg.fatal is False
+    assert error_msg.status_code == 400
+    assert orchestrator_fixture.thread.is_alive()
+    assert "req-overflow" not in orchestrator_fixture.orchestrator.request_states
+
+
 @pytest.mark.asyncio
 async def test_engine_dead_error_fails_only_dead_replica_requests(orchestrator_factory) -> None:
     """Multi-replica stage: one replica dying must fail only the requests bound

tests/entrypoints/openai_api/test_serving_speech.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # tests/entrypoints/openai/test_serving_speech.py
 import asyncio
 import base64
@@ -4147,16 +4150,15 @@ class TestCosyVoice3Serving:
         error = cosyvoice3_server._validate_tts_request(request)
         assert error is None
 
-    def test_validate_cosyvoice3_max_new_tokens_range(self, cosyvoice3_server):
-        request = OpenAICreateSpeechRequest(
-            input="Hello",
-            ref_audio="data:audio/wav;base64,abc",
-            ref_text="hello",
-            max_new_tokens=0,
-        )
-        error = cosyvoice3_server._validate_tts_request(request)
-        assert error is not None
-        assert "max_new_tokens" in error
+    def test_validate_cosyvoice3_max_new_tokens_range(self):
+        """Ensure max_new_tokens below the minimum is rejected during request validation."""
+        with pytest.raises(ValidationError, match="max_new_tokens"):
+            OpenAICreateSpeechRequest(
+                input="Hello",
+                ref_audio="data:audio/wav;base64,abc",
+                ref_text="hello",
+                max_new_tokens=0,
+            )
 
     @pytest.mark.parametrize(
         ("max_new_tokens", "expected_min_tokens", "expected_max_tokens"),