fix(sglang): validate diffusion input_reference and bound media fetches

components/src/dynamo/sglang/request_handlers/image_diffusion/image_diffusion_handler.py

@@ -9,11 +9,17 @@
 import time
 import uuid
 from typing import Any, AsyncGenerator, Optional
+from urllib.parse import urlparse
 
 import torch
 from PIL import Image
 
 from dynamo._core import Context
+from dynamo.common.http.url_validator import (
+    UrlValidationPolicy,
+    validate_local_path,
+    validate_url,
+)
 from dynamo.common.protocols.image_protocol import ImageNvExt
 from dynamo.common.storage import upload_to_fs
 from dynamo.sglang.args import Config
@@ -23,6 +29,20 @@
 
 logger = logging.getLogger(__name__)
 
+
+async def _validate_input_reference(reference: str) -> str:
+    """Validate a client-supplied I2I image reference before it becomes the
+    generator's ``image_path``. Local paths are confined to ``DYN_MM_LOCAL_PATH``
+    (traversal rejected); URLs are blocked from resolving to internal addresses.
+    """
+    policy = UrlValidationPolicy.from_env()
+    parsed = urlparse(reference)
+    if parsed.scheme.lower() in ("", "file"):
+        raw = parsed.path if parsed.scheme else reference
+        return str(validate_local_path(raw, policy))
+    return await validate_url(reference, policy)
+
+
 MAX_NUM_INFERENCE_STEPS = 50
 DEFAULT_NUM_INFERENCE_STEPS = 50
 DEFAULT_GUIDANCE_SCALE = 7.5
@@ -179,7 +199,7 @@ async def _generate_images(
         if input_reference is not None:
             if not input_reference.strip():
                 raise ValueError("input_reference must be a non-empty string")
-            args["image_path"] = input_reference
+            args["image_path"] = await _validate_input_reference(input_reference)
 
         result = await asyncio.to_thread(
             self.generator.generate,

components/src/dynamo/sglang/tests/test_sglang_image_diffusion_handler.py

@@ -392,32 +392,54 @@ async def test_generate_with_nvext(self, handler, mock_context):
 
     @pytest.mark.asyncio
     async def test_generate_i2i_passes_image_path(
-        self, handler, mock_context, tmp_path
+        self, handler, mock_context, tmp_path, monkeypatch
     ):
-        """Test that input_reference is passed as image_path to the generator."""
+        """input_reference inside DYN_MM_LOCAL_PATH is resolved and passed as image_path."""
         test_image = Image.new("RGB", (256, 256), color="green")
 
         handler.generator.generate = Mock(
             return_value=SimpleNamespace(frames=[test_image])
         )
 
-        input_ref = str(tmp_path / "test_input.png")
+        monkeypatch.setenv("DYN_MM_LOCAL_PATH", str(tmp_path))
+        input_ref = tmp_path / "test_input.png"
+        input_ref.write_bytes(b"x")
         request = {
             "prompt": "Transform this image",
             "model": "test-model",
             "size": "256x256",
             "response_format": "b64_json",
-            "input_reference": input_ref,
+            "input_reference": str(input_ref),
         }
 
         results = []
         async for result in handler.generate(request, mock_context):
             results.append(result)
 
-        # Verify image_path was passed to the generator
-        call_args = handler.generator.generate.call_args
-        sampling_params = call_args[1]["sampling_params_kwargs"]
-        assert sampling_params["image_path"] == input_ref
+        sampling_params = handler.generator.generate.call_args[1][
+            "sampling_params_kwargs"
+        ]
+        assert sampling_params["image_path"] == str(input_ref.resolve())
+
+    @pytest.mark.asyncio
+    async def test_generate_i2i_rejects_path_outside_allowed_dir(
+        self, handler, mock_context, tmp_path, monkeypatch
+    ):
+        """A path outside DYN_MM_LOCAL_PATH is rejected before the generator runs."""
+        handler.generator.generate = Mock()
+        monkeypatch.setenv("DYN_MM_LOCAL_PATH", str(tmp_path))
+        request = {
+            "prompt": "x",
+            "model": "test-model",
+            "size": "256x256",
+            "response_format": "b64_json",
+            "input_reference": "/etc/passwd",
+        }
+
+        results = [r async for r in handler.generate(request, mock_context)]
+
+        assert any("error" in r for r in results)
+        handler.generator.generate.assert_not_called()
 
     @pytest.mark.asyncio
     async def test_generate_t2i_no_image_path(self, handler, mock_context):