F1 — Modell-Whitelist-Bypass via version/deployment Argument

Datei: ~/source/mcps/mcp-replicate-hosted/src/mcp_replicate_hosted/main.py

Patch-Skizze

Erweitere die Modell-Whitelist-Logik in GuardMiddleware.on_call_tool, so dass auch die version- und deployment-Pfade von replicate_create_prediction gepruft werden. Praktischer Ansatz: version + deployment komplett ablehnen (keine Layer-Tool nutzt diese Pfade, alle laufen ueber model="owner/name").

# main.py — Tools die alternative Target-Refs erlauben (Replicate-API spezifisch)
_PREDICTION_TOOLS_WITH_ALT_REFS: frozenset[str] = frozenset({
    "replicate_create_prediction",
})
 
# ...
 
class GuardMiddleware(Middleware):
    # ...
    async def on_call_tool(
        self,
        context: MiddlewareContext[Any],
        call_next: Callable[[MiddlewareContext[Any]], Awaitable[Any]],
    ) -> Any:
        if get_settings().emergency_disable:
            audit_log("emergency_disable_blocked", tool=_tool_name(context))
            raise RuntimeError("Service temporarily unavailable")
 
        tool = _tool_name(context)
 
        # NEU: version/deployment-Bypass-Block VOR dem Standard-Model-Check
        if tool in _PREDICTION_TOOLS_WITH_ALT_REFS:
            args = _tool_args(context)
            if args.get("version") or args.get("deployment"):
                subject = _extract_subject(context)
                audit_log(
                    "model_whitelist_bypass_attempt",
                    tool=tool,
                    subject=subject,
                    reason=f"alt_ref={'version' if args.get('version') else 'deployment'}",
                )
                raise RuntimeError(
                    "version/deployment Pfade sind blockiert. Nutze `model='owner/name'` mit "
                    "einem Modell aus der Whitelist (z.B. 'black-forest-labs/flux-2-pro'). "
                    "Whitelist ist in config.py — Erweiterung via PR mit Begruendung + Cost-Estimate."
                )
 
        # Standard-Model-Check (unveraendert)
        if tool in _MODEL_GATED_TOOLS:
            requested_model = _extract_model_arg(context)
            if requested_model and not is_model_allowed(requested_model, self._model_whitelist):
                # ... (existing rejection logic)
                pass
 
        # Rate-Limit, Audit, Result-Return — unveraendert

Test

tests/test_model_whitelist.py ergaenzen:

import pytest
from unittest.mock import MagicMock, AsyncMock
from mcp_replicate_hosted.main import GuardMiddleware, _PREDICTION_TOOLS_WITH_ALT_REFS
from mcp_replicate_hosted.config import DEFAULT_MODEL_WHITELIST
from mcp_replicate_hosted.ratelimit import RateLimiter
 
 
@pytest.mark.asyncio
async def test_version_path_blocked() -> None:
    """F1 regression: version-Arg darf Modell-Whitelist nicht bypassen."""
    middleware = GuardMiddleware(RateLimiter(60, 1000), DEFAULT_MODEL_WHITELIST)
    ctx = MagicMock()
    ctx.message = MagicMock(name="replicate_create_prediction", arguments={
        "version": "abc123def456",
        "input": {"prompt": "..."},
    })
    ctx.fastmcp_context = None  # anonymous
    call_next = AsyncMock(return_value="should_not_be_called")
    with pytest.raises(RuntimeError, match="version/deployment Pfade sind blockiert"):
        await middleware.on_call_tool(ctx, call_next)
    call_next.assert_not_called()
 
 
@pytest.mark.asyncio
async def test_deployment_path_blocked() -> None:
    """F1 regression: deployment-Arg darf Modell-Whitelist nicht bypassen."""
    middleware = GuardMiddleware(RateLimiter(60, 1000), DEFAULT_MODEL_WHITELIST)
    ctx = MagicMock()
    ctx.message = MagicMock(name="replicate_create_prediction", arguments={
        "deployment": "evil/expensive-model",
        "input": {"prompt": "..."},
    })
    ctx.fastmcp_context = None
    call_next = AsyncMock(return_value="should_not_be_called")
    with pytest.raises(RuntimeError, match="version/deployment Pfade sind blockiert"):
        await middleware.on_call_tool(ctx, call_next)
    call_next.assert_not_called()

Aufwand

~10 Zeilen Code + ~25 Zeilen Tests. <15 Min inkl. pytest-Lauf.

Side-Effects

Keine — keine Layer-Tool und keine Slash-Prompts nutzen version= oder deployment=. Marvin-Workflows ueber model= bleiben unveraendert.

Falls Marvin spaeter Fine-Tunes auf eigenen Replicate-Account hostet und ueber version=<hash> ansprechen will: dann pro-User-Whitelist-Mechanik oder Allow-Liste fuer version-Hashes ergaenzen. Heute nicht im Scope.