Bound WebRTC video uploads and clean up owned files

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

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(