diff --git a/src/soup_cli/eval/gate.py b/src/soup_cli/eval/gate.py index 4dedbce..a4e338f 100644 --- a/src/soup_cli/eval/gate.py +++ b/src/soup_cli/eval/gate.py @@ -11,6 +11,7 @@ import json from dataclasses import dataclass from pathlib import Path from typing import Callable, Literal, Optional +from urllib.parse import urlparse import yaml from pydantic import BaseModel, Field, field_validator @@ -81,13 +82,18 @@ class GateTask(BaseModel): if value is None: return None # Allowlist of schemes — SSRF hardening consistent with the project. - allowed = ("ollama://", "https://", "http://localhost", "http://127.0.0.1") - if not value.startswith(allowed): - raise ValueError( - f"judge_model URL '{value}' uses disallowed scheme - " - "use ollama://, https://, or http://localhost" - ) - return value + parsed = urlparse(value) + + if parsed.scheme in ("ollama","https"): + return value + + if (parsed.scheme == "http" and parsed.hostname in {"localhost", "127.0.0.1"}): + return value + + raise ValueError( + f"judge_model URL '{value}' uses disallowed scheme - " + "use ollama://, https://, or http://localhost" + ) class EvalSuite(BaseModel): @@ -168,31 +174,37 @@ def _parse_judge_url(judge_model: str) -> tuple[str, str, Optional[str]]: ``http://localhost:8000/Qwen2.5`` -> ("server", "Qwen2.5", "http://localhost:8000") ``https://api.openai.com/gpt-4o-mini`` -> ("openai", "gpt-4o-mini", "https://api.openai.com") """ - if judge_model.startswith("ollama://"): - return ("ollama", judge_model[len("ollama://"):], None) - # http(s):// — last path segment is the model id; the rest is api_base. - # Defence-in-depth: GateTask._valid_judge_url already restricts the - # scheme to ollama:// / https:// / http://localhost / http://127.0.0.1. - # We match prefixes in the same order here. We do NOT include a bare - # ``http://`` catch-all — if validation is ever bypassed and a non-loopback - # http URL reaches us, the trailing ``raise ValueError`` will fire. - for prefix, default_provider in ( - ("http://localhost", "server"), - ("http://127.0.0.1", "server"), - ("https://", "openai"), - ): - if judge_model.startswith(prefix): - try: - base, model = judge_model.rsplit("/", 1) - except ValueError as exc: - raise ValueError( - f"judge_model '{judge_model}' missing model id" - ) from exc - if not model: - raise ValueError(f"judge_model '{judge_model}' missing model id") - return (default_provider, model, base) - raise ValueError(f"judge_model '{judge_model}' uses unsupported scheme") + parsed = urlparse(judge_model) + + if parsed.scheme == "ollama": + return ("ollama", judge_model[len("ollama://"):], None) + + if parsed.scheme == "https": + default_provider = "openai" + elif ( + parsed.scheme == "http" + and parsed.hostname in ("localhost", "127.0.0.1") + ): + default_provider = "server" + else: + raise ValueError( + f"judge_model '{judge_model}' uses unsupported scheme" + ) + + try: + base, model = judge_model.rsplit("/", 1) + except ValueError as exc: + raise ValueError( + f"judge_model '{judge_model}' missing model id" + ) from exc + + if not model: + raise ValueError( + f"judge_model '{judge_model}' missing model id" + ) + + return (default_provider, model, base) def _run_judge_task( task: GateTask, generate_fn: Callable[[str], str], diff --git a/tests/test_eval_gate.py b/tests/test_eval_gate.py index a16e45f..e2560d2 100644 --- a/tests/test_eval_gate.py +++ b/tests/test_eval_gate.py @@ -498,3 +498,15 @@ class TestTrainGateFlag: cleaned = re.sub(r"\x1b\[[0-9;]*m", "", result.output) assert "--gate" in cleaned assert "eval-gated" in cleaned +class TestJudgeModelValidation: + def test_rejects_localhost_prefix_bypass(self): + from soup_cli.eval.gate import GateTask + + with pytest.raises(ValueError, match="judge_model"): + GateTask( + type="judge", + name="judge", + threshold=0.5, + prompts="prompts.jsonl", + judge_model="http://localhost.attacker.com/model", + ) diff --git a/tests/test_part_a_wave1.py b/tests/test_part_a_wave1.py index 7f72150..d81a2bf 100644 --- a/tests/test_part_a_wave1.py +++ b/tests/test_part_a_wave1.py @@ -53,7 +53,14 @@ class TestParseJudgeURL: assert provider == "server" assert model == "Qwen2.5" assert base == "http://localhost:8000" + def test_rejects_localhost_prefix_bypass(self): + from soup_cli.eval.gate import _parse_judge_url + with pytest.raises(ValueError, match="unsupported scheme"): + _parse_judge_url("http://localhost.attacker.com/model") + + with pytest.raises(ValueError, match="unsupported scheme"): + _parse_judge_url("http://127.0.0.1.evil/model") def test_rejects_unsupported_scheme(self): from soup_cli.eval.gate import _parse_judge_url