From d84643a63da0357645b023c90a32c40c1f07550f Mon Sep 17 00:00:00 2001
From: Rodrigo Barbosa <rodrigo@roboflow.com>
Date: Wed, 9 Sep 2026 09:14:44 -0300
Subject: [PATCH] Bound WebRTC video uploads and clean up owned files

---
 .../interfaces/webrtc_worker/sources/file.py  | 107 +++++--
 .../core/interfaces/webrtc_worker/webrtc.py   |  47 ++-
 .../webrtc_worker/test_video_upload.py        | 272 ++++++++++++++++++
 3 files changed, 391 insertions(+), 35 deletions(-)
 create mode 100644 tests/inference/unit_tests/core/interfaces/webrtc_worker/test_video_upload.py

diff --git a/inference/core/interfaces/webrtc_worker/sources/file.py b/inference/core/interfaces/webrtc_worker/sources/file.py
index 9eb8987031..6943a3da3d 100644
--- a/inference/core/interfaces/webrtc_worker/sources/file.py
+++ b/inference/core/interfaces/webrtc_worker/sources/file.py
@@ -1,7 +1,9 @@
 """Video file source for WebRTC - handles uploaded video files."""
 
 import asyncio
+import os
 import queue
+import tempfile
 import threading
 from typing import Dict, Optional
 
@@ -10,6 +12,7 @@
 from av import VideoFrame
 
 from inference.core import logger
+from inference.core.env import MAX_VIDEO_DOWNLOAD_SIZE_MB
 from inference.core.interfaces.webrtc_worker.entities import VideoFileUploadState
 
 
@@ -114,8 +117,22 @@ class VideoFileUploadHandler:
     Auto-completes when all chunks received.
     """
 
-    def __init__(self):
+    def __init__(self, chunk_size: int):
+        self.chunk_size = chunk_size
+        # Use the same clip budget as URL/base64 video input, including its opt-out.
+        self.max_bytes = (
+            MAX_VIDEO_DOWNLOAD_SIZE_MB * 1024 * 1024
+            if MAX_VIDEO_DOWNLOAD_SIZE_MB >= 0
+            else None
+        )
+        self.max_chunks = (
+            (self.max_bytes + chunk_size - 1) // chunk_size
+            if self.max_bytes is not None
+            else None
+        )
+        self._lock = threading.Lock()
         self._chunks: Dict[int, bytes] = {}
+        self._received_bytes = 0
         self._total_chunks: Optional[int] = None
         self._temp_file_path: Optional[str] = None
         self._state = VideoFileUploadState.IDLE
@@ -125,48 +142,78 @@ def __init__(self):
     def temp_file_path(self) -> Optional[str]:
         return self._temp_file_path
 
-    def handle_chunk(self, chunk_index: int, total_chunks: int, data: bytes) -> None:
-        """Handle a chunk. Auto-completes when all chunks received."""
-        # TODO: we need to refactor this...
-        if self._total_chunks is None:
+    def handle_chunk(self, chunk_index: int, total_chunks: int, data: bytes) -> bool:
+        """Accept one chunk; return whether it adds new upload data."""
+        with self._lock:
+            if self._state not in (
+                VideoFileUploadState.IDLE,
+                VideoFileUploadState.UPLOADING,
+            ):
+                raise ValueError("Video upload is no longer accepting chunks")
+            if total_chunks <= 0 or (
+                self.max_chunks is not None and total_chunks > self.max_chunks
+            ):
+                raise ValueError("Invalid video upload chunk count")
+            if not 0 <= chunk_index < total_chunks:
+                raise ValueError("Invalid video upload chunk index")
+            if not data or len(data) > self.chunk_size:
+                raise ValueError("Invalid video upload chunk size")
+            if self._total_chunks is not None and total_chunks != self._total_chunks:
+                raise ValueError("Video upload chunk count changed")
+            if chunk_index in self._chunks:
+                if self._chunks[chunk_index] != data:
+                    raise ValueError("Conflicting video upload chunk")
+                return False
+            if (
+                self.max_bytes is not None
+                and self._received_bytes + len(data) > self.max_bytes
+            ):
+                raise ValueError("Video upload exceeds the server video size limit")
+
             self._total_chunks = total_chunks
             self._state = VideoFileUploadState.UPLOADING
-
-        self._chunks[chunk_index] = data
-
-        if len(self._chunks) == total_chunks:
-            self._write_to_temp_file()
-            self._state = VideoFileUploadState.COMPLETE
-            self.upload_complete_event.set()
+            self._chunks[chunk_index] = data
+            self._received_bytes += len(data)
+            if len(self._chunks) == self._total_chunks:
+                self._write_to_temp_file()
+                self._state = VideoFileUploadState.COMPLETE
+                self.upload_complete_event.set()
+            return True
 
     def _write_to_temp_file(self) -> None:
         """Reassemble chunks and write to temp file."""
-        import tempfile
-
-        # TODO: we need to refactor this...
         with tempfile.NamedTemporaryFile(mode="wb", suffix=".mp4", delete=False) as f:
+            # Retain ownership even if writing or closing the file fails.
+            self._temp_file_path = f.name
             for i in range(self._total_chunks):
                 f.write(self._chunks[i])
-            self._temp_file_path = f.name
 
         self._chunks.clear()
+        self._received_bytes = 0
 
     def try_start_processing(self) -> Optional[str]:
         """Check if upload complete and transition to PROCESSING. Returns path or None."""
-        if self._state == VideoFileUploadState.COMPLETE:
-            self._state = VideoFileUploadState.PROCESSING
-            return self._temp_file_path
+        with self._lock:
+            if self._state == VideoFileUploadState.COMPLETE:
+                self._state = VideoFileUploadState.PROCESSING
+                return self._temp_file_path
         return None
 
     async def cleanup(self) -> None:
-        """Clean up temp file."""
-        # TODO: we need to refactor this...
-        if self._temp_file_path:
-            import os
-
-            path = self._temp_file_path
-            self._temp_file_path = None
-            try:
-                await asyncio.to_thread(os.unlink, path)
-            except Exception:
-                pass
+        """Wait for pending writes, release chunks and remove the owned file."""
+        await asyncio.to_thread(self._cleanup)
+
+    def _cleanup(self) -> None:
+        with self._lock:
+            self._state = VideoFileUploadState.ERROR
+            self._chunks.clear()
+            self._received_bytes = 0
+            if self._temp_file_path:
+                try:
+                    os.unlink(self._temp_file_path)
+                except FileNotFoundError:
+                    pass
+                except OSError:
+                    logger.warning("Could not remove uploaded video", exc_info=True)
+                    return
+                self._temp_file_path = None
diff --git a/inference/core/interfaces/webrtc_worker/webrtc.py b/inference/core/interfaces/webrtc_worker/webrtc.py
index 93625b3911..7cf763e2df 100644
--- a/inference/core/interfaces/webrtc_worker/webrtc.py
+++ b/inference/core/interfaces/webrtc_worker/webrtc.py
@@ -55,6 +55,7 @@
     WorkflowConfiguration,
 )
 from inference.core.interfaces.webrtc_worker.entities import (
+    VIDEO_FILE_HEADER_SIZE,
     DataOutputMode,
     StreamOutputMode,
     WebRTCOutput,
@@ -1238,12 +1239,19 @@ def on_datachannel(channel: RTCDataChannel):
         logger.info("Data channel '%s' received", channel.label)
         # Handle video file upload channel
         if channel.label == "video_upload":
+            if video_processor.video_upload_handler is not None:
+                channel.close()
+                return
             logger.info("Video upload channel established")
 
-            video_processor.video_upload_handler = VideoFileUploadHandler()
+            upload_handler = VideoFileUploadHandler(chunk_size=CHUNK_SIZE)
+            video_processor.video_upload_handler = upload_handler
+            pending_chunks = 0
+            pending_bytes = 0
 
             @channel.on("message")
             async def on_upload_message(message):
+                nonlocal pending_chunks, pending_bytes
                 # Keep watchdog alive during upload and keepalive pings
                 if video_processor.heartbeat_callback:
                     video_processor.heartbeat_callback()
@@ -1252,10 +1260,39 @@ async def on_upload_message(message):
                 if len(message) <= 1:
                     channel.send(message)
                     return
-                loop = asyncio.get_running_loop()
-                video_path = await loop.run_in_executor(
-                    None, process_video_upload_message, message, video_processor
-                )
+                payload_size = len(message) - VIDEO_FILE_HEADER_SIZE
+                if (
+                    not isinstance(message, bytes)
+                    or not 0 < payload_size <= CHUNK_SIZE
+                    or (
+                        upload_handler.max_chunks is not None
+                        and pending_chunks >= upload_handler.max_chunks
+                    )
+                    or (
+                        upload_handler.max_bytes is not None
+                        and pending_bytes + payload_size > upload_handler.max_bytes
+                    )
+                ):
+                    terminate_event.set()
+                    channel.close()
+                    await upload_handler.cleanup()
+                    return
+                pending_chunks += 1
+                pending_bytes += payload_size
+                try:
+                    loop = asyncio.get_running_loop()
+                    video_path = await loop.run_in_executor(
+                        None, process_video_upload_message, message, video_processor
+                    )
+                except (ValueError, OSError):
+                    logger.warning("Video upload rejected", exc_info=True)
+                    terminate_event.set()
+                    channel.close()
+                    await upload_handler.cleanup()
+                    return
+                finally:
+                    pending_chunks -= 1
+                    pending_bytes -= payload_size
                 if video_path:
                     video_processor._file_processing = True
                     logger.info(
diff --git a/tests/inference/unit_tests/core/interfaces/webrtc_worker/test_video_upload.py b/tests/inference/unit_tests/core/interfaces/webrtc_worker/test_video_upload.py
new file mode 100644
index 0000000000..1f9709eb0b
--- /dev/null
+++ b/tests/inference/unit_tests/core/interfaces/webrtc_worker/test_video_upload.py
@@ -0,0 +1,272 @@
+import asyncio
+import struct
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from inference.core.interfaces.webrtc_worker import webrtc
+from inference.core.interfaces.webrtc_worker.entities import (
+    VideoFileUploadState,
+    WebRTCWorkerRequest,
+)
+from inference.core.interfaces.webrtc_worker.sources import file as video_file
+
+
+@pytest.fixture
+def handler(monkeypatch, tmp_path):
+    monkeypatch.setattr(video_file.tempfile, "tempdir", str(tmp_path))
+    instance = video_file.VideoFileUploadHandler(chunk_size=4)
+    instance.max_bytes = 10
+    instance.max_chunks = 3
+    yield instance
+    instance._cleanup()
+    assert list(tmp_path.iterdir()) == []
+
+
+def test_reuses_video_input_budget(monkeypatch):
+    monkeypatch.setattr(video_file, "MAX_VIDEO_DOWNLOAD_SIZE_MB", 1)
+    instance = video_file.VideoFileUploadHandler(chunk_size=webrtc.CHUNK_SIZE)
+    assert instance.max_bytes == 1024 * 1024
+    assert instance.max_chunks == 22
+
+    monkeypatch.setattr(video_file, "MAX_VIDEO_DOWNLOAD_SIZE_MB", -1)
+    instance = video_file.VideoFileUploadHandler(chunk_size=webrtc.CHUNK_SIZE)
+    assert instance.max_bytes is None
+    assert instance.max_chunks is None
+
+
+@pytest.mark.parametrize(
+    "index,total,data",
+    [
+        (0, 0, b"x"),
+        (0, 4, b"x"),
+        (-1, 3, b"x"),
+        (3, 3, b"x"),
+        (0, 3, b""),
+        (0, 3, b"12345"),
+    ],
+)
+def test_rejects_invalid_chunks_before_retaining_data(handler, index, total, data):
+    with pytest.raises(ValueError):
+        handler.handle_chunk(index, total, data)
+    assert handler._chunks == {}
+    assert handler.temp_file_path is None
+
+
+def test_checks_total_duplicates_and_cumulative_bytes(handler):
+    assert handler.handle_chunk(0, 3, b"1234")
+    assert not handler.handle_chunk(0, 3, b"1234")
+    with pytest.raises(ValueError, match="Conflicting"):
+        handler.handle_chunk(0, 3, b"5678")
+    with pytest.raises(ValueError, match="count changed"):
+        handler.handle_chunk(1, 2, b"x")
+    handler.handle_chunk(1, 3, b"5678")
+    with pytest.raises(ValueError, match="size limit"):
+        handler.handle_chunk(2, 3, b"901")
+    assert handler._received_bytes == 8
+    assert len(handler._chunks) == 2
+
+
+def test_out_of_order_upload_at_limit_only_processes_once(handler):
+    handler.handle_chunk(2, 3, b"90")
+    handler.handle_chunk(0, 3, b"1234")
+    handler.handle_chunk(1, 3, b"5678")
+    assert handler._chunks == {}
+    assert handler._state == VideoFileUploadState.COMPLETE
+    with pytest.raises(ValueError, match="no longer accepting"):
+        handler.handle_chunk(0, 3, b"xxxx")
+    path = handler.try_start_processing()
+    assert Path(path).read_bytes() == b"1234567890"
+    assert handler.try_start_processing() is None
+    with pytest.raises(ValueError, match="no longer accepting"):
+        handler.handle_chunk(0, 1, b"new")
+
+
+@pytest.mark.asyncio
+async def test_cleanup_releases_partial_upload_and_is_terminal(handler):
+    handler.handle_chunk(0, 3, b"1234")
+    await handler.cleanup()
+    await handler.cleanup()
+    assert handler._received_bytes == 0
+    assert handler._chunks == {}
+    with pytest.raises(ValueError, match="no longer accepting"):
+        handler.handle_chunk(1, 3, b"x")
+
+
+@pytest.mark.asyncio
+async def test_write_error_retains_file_for_cleanup(handler, monkeypatch, tmp_path):
+    real_temporary_file = video_file.tempfile.NamedTemporaryFile
+
+    def failing_file(*args, **kwargs):
+        file = real_temporary_file(*args, **kwargs)
+        file.write = MagicMock(side_effect=OSError("disk full"))
+        return file
+
+    monkeypatch.setattr(video_file.tempfile, "NamedTemporaryFile", failing_file)
+    with pytest.raises(OSError, match="disk full"):
+        handler.handle_chunk(0, 1, b"x")
+    assert Path(handler.temp_file_path).exists()
+    await handler.cleanup()
+    assert handler.temp_file_path is None
+    assert list(tmp_path.iterdir()) == []
+
+
+@pytest.mark.asyncio
+async def test_unlink_failure_can_be_retried(handler, monkeypatch):
+    handler.handle_chunk(0, 1, b"x")
+    path = handler.temp_file_path
+    with monkeypatch.context() as patch:
+        patch.setattr(video_file.os, "unlink", MagicMock(side_effect=PermissionError()))
+        await handler.cleanup()
+    assert handler.temp_file_path == path
+    await handler.cleanup()
+    assert not Path(path).exists()
+
+
+def test_concurrent_duplicate_completion_creates_one_file(handler, tmp_path):
+    def complete():
+        try:
+            handler.handle_chunk(0, 1, b"x")
+        except ValueError:
+            return None
+        return handler.try_start_processing()
+
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        paths = list(pool.map(lambda _: complete(), range(2)))
+    assert sum(path is not None for path in paths) == 1
+    assert len(list(tmp_path.iterdir())) == 1
+
+
+@pytest.mark.asyncio
+async def test_cleanup_waits_for_inflight_write(handler, monkeypatch, tmp_path):
+    writing = threading.Event()
+    release = threading.Event()
+    write = handler._write_to_temp_file
+
+    def delayed_write():
+        writing.set()
+        assert release.wait(timeout=5)
+        write()
+
+    monkeypatch.setattr(handler, "_write_to_temp_file", delayed_write)
+    upload = asyncio.create_task(asyncio.to_thread(handler.handle_chunk, 0, 1, b"x"))
+    try:
+        assert await asyncio.to_thread(writing.wait, 5)
+        cleanup = asyncio.create_task(handler.cleanup())
+        await asyncio.sleep(0)
+    finally:
+        release.set()
+    await upload
+    await cleanup
+    assert handler._state == VideoFileUploadState.ERROR
+    assert list(tmp_path.iterdir()) == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "message",
+    [
+        b"bad",
+        "not binary",
+        struct.pack("<II", 2, 1) + b"x",
+        b"x" * (webrtc.CHUNK_SIZE + 9),
+        "queued_bytes",
+        "queued_chunks",
+    ],
+)
+async def test_upload_channel_rejects_invalid_messages_and_extra_channel(
+    monkeypatch, tmp_path, message
+):
+    monkeypatch.setattr(video_file.tempfile, "tempdir", str(tmp_path))
+    events = {}
+    peer = MagicMock(connectionState="connected")
+    peer.on.side_effect = lambda event: lambda callback: events.setdefault(
+        event, callback
+    )
+    peer.setRemoteDescription = AsyncMock()
+    peer.createAnswer = AsyncMock()
+    peer.setLocalDescription = AsyncMock()
+    peer.close = AsyncMock()
+    peer.localDescription.type = "answer"
+    peer.localDescription.sdp = ""
+    processor = MagicMock(
+        video_upload_handler=None, track=None, heartbeat_callback=None
+    )
+
+    async def close():
+        if processor.video_upload_handler:
+            await processor.video_upload_handler.cleanup()
+
+    processor.close = close
+    monkeypatch.setattr(
+        webrtc, "RTCPeerConnectionWithLoop", MagicMock(return_value=peer)
+    )
+    monkeypatch.setattr(
+        webrtc, "VideoFrameProcessor", MagicMock(return_value=processor)
+    )
+    monkeypatch.setattr(webrtc, "_wait_ice_complete", AsyncMock())
+    monkeypatch.setattr(
+        webrtc.usage_collector, "async_push_usage_payloads", AsyncMock()
+    )
+    ready = asyncio.Event()
+    request = WebRTCWorkerRequest(
+        workflow_configuration={
+            "type": "WorkflowConfiguration",
+            "workflow_specification": {},
+        },
+        webrtc_offer={"sdp": "", "type": "offer"},
+        stream_output=[],
+        processing_timeout=None,
+    )
+    session = asyncio.create_task(
+        webrtc.init_rtc_peer_connection_with_loop(request, lambda _: ready.set())
+    )
+    try:
+        await asyncio.wait_for(ready.wait(), timeout=5)
+        messages = {}
+        channel = MagicMock(label="video_upload")
+        channel.on.side_effect = lambda event: lambda callback: messages.setdefault(
+            event, callback
+        )
+        events["datachannel"](channel)
+        first_handler = processor.video_upload_handler
+        extra = MagicMock(label="video_upload")
+        events["datachannel"](extra)
+        extra.close.assert_called_once()
+        assert processor.video_upload_handler is first_handler
+        if message in ("queued_bytes", "queued_chunks"):
+            working = threading.Event()
+            release = threading.Event()
+
+            def delayed_chunk(*args):
+                working.set()
+                assert release.wait(timeout=5)
+                return False
+
+            monkeypatch.setattr(first_handler, "handle_chunk", delayed_chunk)
+            if message == "queued_bytes":
+                first_handler.max_bytes = 1
+            else:
+                first_handler.max_chunks = 1
+            chunk = struct.pack("<II", 0, 1) + b"x"
+            pending = asyncio.create_task(messages["message"](chunk))
+            try:
+                assert await asyncio.to_thread(working.wait, 5)
+                await messages["message"](chunk)
+            finally:
+                release.set()
+                await pending
+        else:
+            await messages["message"](message)
+        await asyncio.wait_for(session, timeout=5)
+        channel.close.assert_called_once()
+        assert first_handler._state == VideoFileUploadState.ERROR
+        assert first_handler._chunks == {}
+        assert list(tmp_path.iterdir()) == []
+    finally:
+        if not session.done():
+            session.cancel()
+            await asyncio.gather(session, return_exceptions=True)