[#20986][fix] Cap decoded audio samples in AudioMediaIO

tensorrt_llm/inputs/media_io.py

@@ -176,6 +176,17 @@ async def _close_aiohttp_session() -> None:
 # Maximum number of redirects allowed for remote fetches.
 _MAX_REDIRECTS = 5
 
+# Maximum number of decoded audio samples (frames x channels) per audio input.
+# Compressed formats such as FLAC can expand a few hundred kilobytes into
+# gigabytes of samples, so decoded audio needs its own limit. The value is the
+# sample count of _MAX_RESPONSE_BYTES of 16-bit PCM, about 100 Mi samples, or
+# 800 MiB as float64.
+_MAX_AUDIO_SAMPLES = _MAX_RESPONSE_BYTES // 2
+
+# Frame count that libsndfile reports when a file does not state its length,
+# such as a FLAC stream whose STREAMINFO total is 0.
+_SF_UNKNOWN_FRAMES = 2**63 - 1
+
 _REDIRECT_STATUSES = (301, 302, 303, 307, 308)
 
 # Opt-in switch relaxing the requirement that media URLs resolve to global
@@ -1092,6 +1103,38 @@ def _decode_base64_with_limit(data: str) -> bytes:
     return decoded
 
 
+class _BoundedAudioSoundFile(soundfile.SoundFile):
+    def seekable(self) -> bool:
+        # libsndfile reports unknown-length FLAC as a huge sentinel. soundfile
+        # seeks after each read on seekable files, including a short final read;
+        # that seek fails when the header has no frame count. Read sequentially.
+        return self.frames != _SF_UNKNOWN_FRAMES and super().seekable()
+
+
+def _read_audio_with_limit(source: Any) -> Tuple[np.ndarray, int]:
+    """Decode audio like ``soundfile.read`` with a cap on decoded samples.
+
+    A file whose header reports more than ``_MAX_AUDIO_SAMPLES`` samples is
+    rejected before any samples are read. For a file of unknown length, the
+    read stops at one frame beyond the limit, so an oversized payload is
+    rejected before it is fully decoded.
+    """
+    too_many_samples = (
+        f"Audio input decodes to more than {_MAX_AUDIO_SAMPLES} samples "
+        f"(frames x channels), the maximum allowed."
+    )
+    with _BoundedAudioSoundFile(source) as audio_file:
+        # libsndfile refuses to open a file with fewer than one channel.
+        max_frames = _MAX_AUDIO_SAMPLES // audio_file.channels
+        if audio_file.frames != _SF_UNKNOWN_FRAMES and audio_file.frames > max_frames:
+            raise ValueError(too_many_samples)
+        data = audio_file.read(frames=max_frames + 1)
+        samplerate = audio_file.samplerate
+    if len(data) > max_frames:
+        raise ValueError(too_many_samples)
+    return data, samplerate
+
+
 def _safe_read_local_file(location: str, *, max_bytes: int | None = None) -> bytes:
     """Read a local file, refusing anything that has no end.
 
@@ -1339,7 +1382,7 @@ def load_bytes(self, data: bytes) -> Tuple[np.ndarray, int]:
         return self._read(BytesIO(data))
 
     def load_base64(self, media_type: str, data: str) -> Tuple[np.ndarray, int]:
-        return self.load_bytes(base64.b64decode(data))
+        return self.load_bytes(_decode_base64_with_limit(data))
 
     def load_file(self, url: str) -> Tuple[np.ndarray, int]:
         # Strip `file://` and unquote the path for `file:` URIs; pass
@@ -1350,16 +1393,23 @@ def load_file(self, url: str) -> Tuple[np.ndarray, int]:
 
     def _read(self, source: Union[BytesIO, str]) -> Tuple[np.ndarray, int]:
         if self._decoded_memory_reserver is None:
-            return soundfile.read(source)
+            return _read_audio_with_limit(source)
 
         reservation = None
         try:
             # Read metadata and samples from the same handle. Reserve the float64
             # output before allocation; libsndfile's internal buffers are excluded.
-            with soundfile.SoundFile(source) as audio:
+            with _BoundedAudioSoundFile(source) as audio:
                 frames, channels = audio.frames, audio.channels
                 if frames < 0 or channels < 1:
                     raise ValueError("invalid audio frame or channel count")
+                if frames == _SF_UNKNOWN_FRAMES:
+                    return self._read_unknown_length(audio, channels)
+                if frames * channels > _MAX_AUDIO_SAMPLES:
+                    raise ValueError(
+                        f"Audio input decodes to more than {_MAX_AUDIO_SAMPLES} samples "
+                        f"(frames x channels), the maximum allowed."
+                    )
                 decoded_bytes = frames * channels * np.dtype(np.float64).itemsize
                 reservation = self._decoded_memory_reserver(decoded_bytes, decoded_bytes)
                 shape = (frames,) if channels == 1 else (frames, channels)
@@ -1379,6 +1429,64 @@ def _read(self, source: Union[BytesIO, str]) -> Tuple[np.ndarray, int]:
             if reservation is not None:
                 reservation.release()
 
+    def _read_unknown_length(
+        self, audio: soundfile.SoundFile, channels: int
+    ) -> Tuple[np.ndarray, int]:
+        """Read an unknown-length stream in bounded, budgeted chunks."""
+        reserver = self._decoded_memory_reserver
+        assert reserver is not None
+        max_frames = _MAX_AUDIO_SAMPLES // channels
+        too_many_samples = (
+            f"Audio input decodes to more than {_MAX_AUDIO_SAMPLES} samples "
+            f"(frames x channels), the maximum allowed."
+        )
+        chunks: List[np.ndarray] = []
+        reservations: List[Tuple[_DecodedMemoryReservation, int]] = []
+        concatenation_reservation: Optional[_DecodedMemoryReservation] = None
+        total_frames = 0
+        try:
+            while True:
+                chunk_frames = min(16384, max_frames + 1 - total_frames)
+                chunk_bytes = chunk_frames * channels * np.dtype(np.float64).itemsize
+                reservation = reserver(chunk_bytes, chunk_bytes)
+                reservations.append((reservation, 0))
+                shape = (chunk_frames,) if channels == 1 else (chunk_frames, channels)
+                samples = np.empty(shape, dtype=np.float64)
+                decoded = audio.read(frames=chunk_frames, out=samples)
+                read_frames = len(decoded)
+                if total_frames + read_frames > max_frames:
+                    raise ValueError(too_many_samples)
+                if read_frames != chunk_frames:
+                    decoded = None
+                    samples.resize((read_frames,) if channels == 1 else (read_frames, channels))
+                    reservation.shrink(samples.nbytes)
+                if read_frames:
+                    chunks.append(samples)
+                    reservations[-1] = (reservation, samples.nbytes)
+                total_frames += read_frames
+                if read_frames < chunk_frames:
+                    break
+
+            if chunks:
+                total_bytes = sum(size for _, size in reservations)
+                concatenation_reservation = reserver(total_bytes, 0)
+                result = np.concatenate(chunks)
+                chunks.clear()
+                concatenation_reservation.release()
+                concatenation_reservation = None
+            else:
+                shape = (0,) if channels == 1 else (0, channels)
+                result = np.empty(shape, dtype=np.float64)
+            for reservation, size in reservations:
+                reservation.commit(size)
+            return result, audio.samplerate
+        finally:
+            chunks.clear()
+            if concatenation_reservation is not None:
+                concatenation_reservation.release()
+            for reservation, _ in reservations:
+                reservation.release()
+
 
 class VideoMediaIO(BaseMediaIO[VideoData]):
     """Video I/O with a 200 MiB encoded-input limit and coupled sampling kwargs."""

tests/integration/test_lists/test-db/l0_cpu.yml

@@ -243,6 +243,7 @@ l0_cpu:
   - unittest/llmapi/apps/test_harmony_channel_validation.py
   - unittest/llmapi/apps/test_harmony_parsing.py::TestStripIncompleteMessagesReporting
   - unittest/llmapi/apps/test_kimi_serve_extensions.py
+  - unittest/llmapi/apps/test_media_io.py
   - unittest/llmapi/apps/test_openai_chat_disagg_multimodal.py
   - unittest/llmapi/apps/test_openai_chat_multimodal_prompt.py
   - unittest/llmapi/apps/test_openai_protocol_chat_store.py

tests/unittest/llmapi/apps/test_media_io.py

@@ -1,9 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-from unittest.mock import AsyncMock, patch
+import base64
+from io import BytesIO
+from unittest.mock import AsyncMock, Mock, patch
 
+import numpy as np
 import pytest
+import soundfile
 
+import tensorrt_llm.inputs.media_io as media_io_module
 from tensorrt_llm.inputs import MultimodalDataTracker
 from tensorrt_llm.inputs.media_io import AudioMediaIO, BaseMediaIO, ImageMediaIO, VideoMediaIO
 from tensorrt_llm.serve.chat_utils import parse_chat_message_content_part
@@ -91,3 +96,157 @@ def test_non_video_classes_use_plain_shallow_merge(self, media_io_cls):
             {"num_frames": 32},
         )
         assert merged == {"num_frames": 32, "fps": 1}
+
+
+def _flac_bytes(frames: int, channels: int = 2) -> bytes:
+    # Nonzero samples that differ per channel, so a loader that silences or
+    # reorders channels fails the comparison with soundfile.read.
+    t = np.arange(frames)[:, None]
+    c = np.arange(channels)[None, :]
+    pcm = ((t * (2 * c + 3)) % 20000 + 1) * (1 - 2 * (c % 2))
+    buffer = BytesIO()
+    soundfile.write(buffer, pcm.astype(np.int16), 16000, format="FLAC", subtype="PCM_16")
+    return buffer.getvalue()
+
+
+def _without_frame_count(flac: bytes) -> bytes:
+    # Zero the 36-bit total-samples field of STREAMINFO, which starts in the
+    # low nibble of byte 21, so the stream no longer states its length.
+    data = bytearray(flac)
+    data[21] &= 0xF0
+    data[22:26] = bytes(4)
+    return bytes(data)
+
+
+class TestAudioDecodedSizeLimit:
+    """Audio decoding stops at `_MAX_AUDIO_SAMPLES` decoded samples."""
+
+    LIMIT = 4000
+
+    @pytest.fixture(autouse=True)
+    def small_limit(self, monkeypatch):
+        monkeypatch.setattr(media_io_module, "_MAX_AUDIO_SAMPLES", self.LIMIT)
+
+    @pytest.fixture
+    def audio_reads(self, monkeypatch):
+        """Frame counts passed to each ``SoundFile.read`` call."""
+        reads = []
+        original_read = soundfile.SoundFile.read
+
+        def recording_read(audio_file, *args, **kwargs):
+            reads.append(kwargs.get("frames"))
+            return original_read(audio_file, *args, **kwargs)
+
+        monkeypatch.setattr(soundfile.SoundFile, "read", recording_read)
+        return reads
+
+    def test_audio_within_limit_matches_soundfile(self, tmp_path):
+        encoded = _flac_bytes(frames=self.LIMIT // 2, channels=2)
+        path = tmp_path / "ok.flac"
+        path.write_bytes(encoded)
+        expected, expected_rate = soundfile.read(BytesIO(encoded))
+        loader = AudioMediaIO()
+        for data, rate in (
+            loader.load_bytes(encoded),
+            loader.load_base64("audio/flac", base64.b64encode(encoded).decode()),
+            loader.load_file(str(path)),
+            loader.load_file(path.as_uri()),
+        ):
+            assert rate == expected_rate
+            np.testing.assert_array_equal(data, expected, strict=True)
+
+    def test_audio_over_limit_is_rejected(self, tmp_path):
+        # Two channels, so LIMIT // 2 + 1 frames is just over the sample limit.
+        encoded = _flac_bytes(frames=self.LIMIT // 2 + 1, channels=2)
+        path = tmp_path / "big.flac"
+        path.write_bytes(encoded)
+        loader = AudioMediaIO()
+        with pytest.raises(ValueError, match="more than 4000 samples"):
+            loader.load_bytes(encoded)
+        with pytest.raises(ValueError, match="more than 4000 samples"):
+            loader.load_base64("audio/flac", base64.b64encode(encoded).decode())
+        with pytest.raises(ValueError, match="more than 4000 samples"):
+            loader.load_file(str(path))
+
+    def test_audio_over_limit_in_header_is_rejected_before_reading(self, audio_reads):
+        encoded = _flac_bytes(frames=self.LIMIT // 2 + 1, channels=2)
+        with pytest.raises(ValueError, match="more than 4000 samples"):
+            AudioMediaIO().load_bytes(encoded)
+        assert audio_reads == []
+
+    def test_audio_with_memory_reservation_obeys_limit(self, audio_reads):
+        reservation = Mock()
+        reserver = Mock(return_value=reservation)
+        loader = AudioMediaIO(decoded_memory_reserver=reserver)
+        encoded = _flac_bytes(frames=self.LIMIT // 2, channels=2)
+        data, rate = loader.load_base64("audio/flac", base64.b64encode(encoded).decode())
+        assert data.shape == (self.LIMIT // 2, 2)
+        assert rate == 16000
+        reserver.assert_called_once_with(self.LIMIT * 8, self.LIMIT * 8)
+        reservation.commit.assert_called_once_with(data.nbytes)
+        reservation.release.assert_called_once()
+        assert audio_reads == [self.LIMIT // 2]
+
+        oversized = _flac_bytes(frames=self.LIMIT // 2 + 1, channels=2)
+        with pytest.raises(ValueError, match="more than 4000 samples"):
+            loader.load_bytes(oversized)
+        assert audio_reads == [self.LIMIT // 2]
+        reserver.assert_called_once()
+
+    def test_audio_of_unknown_length_over_limit_is_rejected(self, audio_reads):
+        encoded = _without_frame_count(_flac_bytes(frames=self.LIMIT, channels=2))
+        with soundfile.SoundFile(BytesIO(encoded)) as audio_file:
+            assert audio_file.frames == media_io_module._SF_UNKNOWN_FRAMES
+        with pytest.raises(ValueError, match="more than 4000 samples"):
+            AudioMediaIO().load_bytes(encoded)
+        assert audio_reads == [self.LIMIT // 2 + 1]
+
+    def test_unknown_length_audio_with_reservation(self, audio_reads):
+        original = _flac_bytes(frames=1000, channels=2)
+        encoded = _without_frame_count(original)
+        expected, expected_rate = soundfile.read(BytesIO(original))
+        audio_reads.clear()
+        reservation = Mock()
+        reserver = Mock(return_value=reservation)
+
+        data, rate = AudioMediaIO(decoded_memory_reserver=reserver).load_bytes(encoded)
+
+        assert rate == expected_rate
+        np.testing.assert_array_equal(data, expected, strict=True)
+        assert audio_reads == [self.LIMIT // 2 + 1]
+        assert reserver.call_count == 2  # Bounded read and final concatenation.
+        reservation.shrink.assert_called_once_with(data.nbytes)
+        reservation.commit.assert_any_call(data.nbytes)
+
+        unreserved, unreserved_rate = AudioMediaIO().load_bytes(encoded)
+        assert unreserved_rate == expected_rate
+        np.testing.assert_array_equal(unreserved, expected, strict=True)
+
+    def test_unknown_length_audio_over_limit_with_reservation(self, audio_reads):
+        encoded = _without_frame_count(_flac_bytes(frames=self.LIMIT, channels=2))
+        reservation = Mock()
+        reserver = Mock(return_value=reservation)
+
+        with pytest.raises(ValueError, match="more than 4000 samples"):
+            AudioMediaIO(decoded_memory_reserver=reserver).load_bytes(encoded)
+
+        assert audio_reads == [self.LIMIT // 2 + 1]
+        reservation.commit.assert_not_called()
+        reservation.release.assert_called_once()
+
+    def test_audio_base64_over_encoded_limit_is_rejected(self, monkeypatch):
+        encoded = _flac_bytes(frames=self.LIMIT // 2, channels=2)
+        monkeypatch.setattr(media_io_module, "_MAX_RESPONSE_BYTES", len(encoded) // 2)
+        with pytest.raises(ValueError, match="Base64 media input exceeds"):
+            AudioMediaIO().load_base64("audio/flac", base64.b64encode(encoded).decode())
+
+    @pytest.mark.asyncio
+    async def test_chat_audio_part_over_limit_is_rejected(self):
+        encoded = _flac_bytes(frames=self.LIMIT + 1, channels=1)
+        part = {
+            "type": "input_audio",
+            "input_audio": {"data": base64.b64encode(encoded).decode(), "format": "flac"},
+        }
+        result = parse_chat_message_content_part(part, MultimodalDataTracker(model_type="dummy"))
+        with pytest.raises(ValueError, match="more than 4000 samples"):
+            await result["data"]