odysseus/tests/test_gallery_endpoint_hardening.py

"""Focused security tests for gallery endpoint URL hardening.

Covers:
- _is_openai_api_base: exact hostname matching (no substring bypass)
- _join_checked_gallery_endpoint: allowlist-only path construction
- No bare str(e) / f"...{e}" in gallery exception handlers
- harmonize validates _endpoint via check_outbound_url
- Target URL construction only appends constant paths to the validated base
"""
import ast
import re
from pathlib import Path

SRC = Path(__file__).resolve().parent.parent / "routes" / "gallery" / "gallery_routes.py"

import routes.gallery_routes as gallery_routes


# ---------------------------------------------------------------------------
# _is_openai_api_base — exact hostname, no substring tricks
# ---------------------------------------------------------------------------

def test_is_openai_api_base_accepts_exact_host():
    f = gallery_routes._is_openai_api_base
    assert f("https://api.openai.com") is True
    assert f("https://api.openai.com/v1") is True
    assert f("https://api.openai.com/") is True
    assert f("api.openai.com") is True


def test_is_openai_api_base_rejects_path_embed():
    # attacker hides api.openai.com in the path, not the hostname
    f = gallery_routes._is_openai_api_base
    assert f("https://evil.test/api.openai.com/v1") is False


def test_is_openai_api_base_rejects_subdomain_suffix():
    # hostname ends with .openai.com but isn't exactly api.openai.com
    f = gallery_routes._is_openai_api_base
    assert f("https://api.openai.com.evil.test/v1") is False
    assert f("https://evil-api.openai.com/v1") is False
    assert f("https://notapi.openai.com/v1") is False


def test_is_openai_api_base_rejects_malformed():
    f = gallery_routes._is_openai_api_base
    assert f("") is False
    assert f("not a url at all !!!") is False


# ---------------------------------------------------------------------------
# Source-level: gallery no longer uses substring "api.openai.com" in base
# ---------------------------------------------------------------------------

def test_gallery_does_not_use_openai_substring_check():
    src = SRC.read_text()
    assert '"api.openai.com" in base' not in src, (
        "Substring OpenAI check still present — use _is_openai_api_base instead"
    )
    assert "'api.openai.com' in base" not in src, (
        "Substring OpenAI check still present — use _is_openai_api_base instead"
    )


# ---------------------------------------------------------------------------
# _join_checked_gallery_endpoint — allowlist enforcement
# ---------------------------------------------------------------------------

def test_join_checked_accepts_known_paths():
    j = gallery_routes._join_checked_gallery_endpoint
    assert j("http://localhost:7860/v1", "/images/img2img") == "http://localhost:7860/v1/images/img2img"
    assert j("http://localhost:7860", "/sdapi/v1/img2img") == "http://localhost:7860/sdapi/v1/img2img"
    assert j("https://api.openai.com/v1", "/images/edits") == "https://api.openai.com/v1/images/edits"


def test_join_checked_rejects_unknown_path():
    import pytest
    j = gallery_routes._join_checked_gallery_endpoint
    with pytest.raises(ValueError):
        j("http://localhost/v1", "/arbitrary/user/path")
    with pytest.raises(ValueError):
        j("http://localhost/v1", "")
    with pytest.raises(ValueError):
        j("http://localhost/v1", "https://evil.test/steal")


# ---------------------------------------------------------------------------
# Source-level: no raw str(e) / f"...{e}" returned to API clients
# ---------------------------------------------------------------------------

def test_no_raw_exception_string_in_client_responses():
    src = SRC.read_text()
    # Patterns that indicate exception internals flowing into client-visible values.
    # We allow them only in logger calls (checked separately below).
    bad_patterns = [
        r'return \{"error": str\(e\)\}',
        r'return \{"error": f"[^"]*\{e\}[^"]*"\}',
        r'HTTPException\(\d+, str\(e\)\)',
        r'HTTPException\(\d+, f"[^"]*\{e\}[^"]*"\)',
    ]
    for pattern in bad_patterns:
        matches = re.findall(pattern, src)
        assert not matches, (
            f"Pattern {pattern!r} matched — raw exception string returned to client: {matches}"
        )


# ---------------------------------------------------------------------------
# harmonize: validates _endpoint via check_outbound_url before outbound request
# ---------------------------------------------------------------------------

def _function_source(src_text: str, func_name: str) -> str:
    tree = ast.parse(src_text)
    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name:
            return ast.get_source_segment(src_text, node) or ""
    raise AssertionError(f"{func_name} not found in {SRC}")


def test_harmonize_validates_endpoint_before_fetch():
    src = SRC.read_text()
    body = _function_source(src, "harmonize_image")
    assert "check_outbound_url" in body, (
        "harmonize_image must validate _endpoint via check_outbound_url before outbound requests"
    )


# ---------------------------------------------------------------------------
# harmonize: target URL only appends constant allowed paths
# ---------------------------------------------------------------------------

def test_harmonize_uses_join_checked_for_target_construction():
    src = SRC.read_text()
    body = _function_source(src, "harmonize_image")
    assert "_join_checked_gallery_endpoint" in body, (
        "harmonize_image must use _join_checked_gallery_endpoint to build target URLs"
    )
    # Raw concatenation patterns that bypass the allowlist must not appear in harmonize
    assert "base_root + path" not in body, (
        "harmonize_image must not concatenate base_root + path directly"
    )
    assert "base + path" not in body, (
        "harmonize_image must not concatenate base + path directly"
    )


def test_gallery_endpoint_paths_allowlist_covers_all_harmonize_candidates():
    # Every path in the candidates list must be in the pre-approved allowlist.
    src = SRC.read_text()
    body = _function_source(src, "harmonize_image")
    # Extract string literals that look like route paths from candidates
    candidate_paths = re.findall(r'"/(?:images|sdapi)/[^"]*"', body)
    allowed = gallery_routes._GALLERY_ENDPOINT_PATHS
    for p in candidate_paths:
        p = p.strip('"')
        assert p in allowed, (
            f"Path {p!r} used in harmonize candidates but not in _GALLERY_ENDPOINT_PATHS allowlist"
        )


# ---------------------------------------------------------------------------
# _is_openai_api_base — userinfo bypass
# ---------------------------------------------------------------------------

def test_is_openai_api_base_rejects_userinfo_bypass():
    # userinfo trick: user = api.openai.com, host = evil.test
    f = gallery_routes._is_openai_api_base
    assert f("https://api.openai.com@evil.test/v1") is False


# ---------------------------------------------------------------------------
# Source-level: no client-visible error leaks upstream body fragments
# ---------------------------------------------------------------------------

def _extract_httpexception_call(src: str, pos: int) -> str:
    """Paren-match from the opening '(' of an HTTPException call."""
    start = src.index("(", pos)
    depth = 0
    for k, ch in enumerate(src[start:]):
        if ch == "(":
            depth += 1
        elif ch == ")":
            depth -= 1
            if depth == 0:
                return src[start : start + k + 1]
    return src[start:]


def test_no_upstream_data_in_client_responses():
    """No raise HTTPException or return {"error": ...} may expose upstream body data."""
    src = SRC.read_text()
    forbidden = [
        "r.text",
        "body[:",
        'data["error"]',
        "data['error']",
        "last_err",
        "{base}",
    ]

    for m in re.finditer(r"\braise\s+HTTPException\s*\(", src):
        line_start = src.rfind("\n", 0, m.start()) + 1
        if "logger." in src[line_start : m.start()]:
            continue
        call_text = _extract_httpexception_call(src, m.start())
        for frag in forbidden:
            assert frag not in call_text, (
                f"HTTPException raise at byte {m.start()} exposes {frag!r} to client:\n{call_text[:300]}"
            )

    for m in re.finditer(r'return\s*\{"error":', src):
        line_end = src.find("\n", m.start())
        line = src[m.start() : line_end if line_end != -1 else len(src)]
        for frag in forbidden:
            assert frag not in line, (
                f"Error return at byte {m.start()} exposes {frag!r} to client:\n{line}"
            )


# ---------------------------------------------------------------------------
# inpaint_proxy: endpoint construction via _join_checked_gallery_endpoint
# ---------------------------------------------------------------------------

def test_inpaint_uses_join_checked_endpoint():
    src = SRC.read_text()
    body = _function_source(src, "inpaint_proxy")
    assert 'f"{base}/images/edits"' not in body, (
        "inpaint_proxy must not build /images/edits via raw f-string"
    )
    assert 'f"{base}/images/inpaint"' not in body, (
        "inpaint_proxy must not build /images/inpaint via raw f-string"
    )
    assert '_join_checked_gallery_endpoint(base, "/images/edits")' in body, (
        "inpaint_proxy must use _join_checked_gallery_endpoint for /images/edits"
    )
    assert '_join_checked_gallery_endpoint(base, "/images/inpaint")' in body, (
        "inpaint_proxy must use _join_checked_gallery_endpoint for /images/inpaint"
    )


# ---------------------------------------------------------------------------
# harmonize final 502: no base URL or last_err in client message
# ---------------------------------------------------------------------------

def test_harmonize_final_502_omits_base_and_last_err():
    src = SRC.read_text()
    body = _function_source(src, "harmonize_image")
    # Collect all HTTPException raises in harmonize and check the last one (final 502)
    raises = list(re.finditer(r"\braise\s+HTTPException\s*\(", body))
    assert raises, "harmonize_image must contain at least one raise HTTPException"
    last_call = _extract_httpexception_call(body, raises[-1].start())
    for forbidden in ("last_err", "{base}", "r.text"):
        assert forbidden not in last_call, (
            f"harmonize final raise exposes {forbidden!r} to client:\n{last_call}"
        )


# ---------------------------------------------------------------------------
# inpaint/harmonize: _endpoint must resolve via DB; no raw admin bypass
# ---------------------------------------------------------------------------

def test_inpaint_endpoint_resolved_via_db_not_raw_input():
    """inpaint_proxy must not use the raw request-body value as the outbound base.
    The user-supplied value is stored as requested_base; outbound base comes from DB."""
    src = SRC.read_text()
    body = _function_source(src, "inpaint_proxy")
    # requested_base holds the user input; base is only set from ep.base_url
    assert "requested_base" in body, (
        "inpaint_proxy must use 'requested_base' for the user-supplied value"
    )
    # The admin bypass (not _current_user_is_admin) must not appear in inpaint
    assert "_current_user_is_admin" not in body, (
        "inpaint_proxy must not have an admin bypass for raw endpoint resolution"
    )
    # If no matching endpoint is found, a 403 must be raised unconditionally
    assert 'raise HTTPException(403, "Choose a registered image endpoint")' in body, (
        "inpaint_proxy must raise 403 when _endpoint doesn't match a registered endpoint"
    )


def test_inpaint_outbound_base_not_from_request_body():
    """Confirm _join_checked_gallery_endpoint is never called with the raw
    request-body variable (requested_base) — only with the DB-derived base."""
    src = SRC.read_text()
    body = _function_source(src, "inpaint_proxy")
    assert "_join_checked_gallery_endpoint(requested_base," not in body, (
        "inpaint_proxy must not pass requested_base to _join_checked_gallery_endpoint"
    )


def test_harmonize_endpoint_resolved_via_db_not_raw_input():
    """harmonize_image must not use the raw request-body value as the outbound base."""
    src = SRC.read_text()
    body = _function_source(src, "harmonize_image")
    assert "requested_base" in body, (
        "harmonize_image must use 'requested_base' for the user-supplied value"
    )
    assert "_current_user_is_admin" not in body, (
        "harmonize_image must not have an admin bypass for raw endpoint resolution"
    )
    assert 'raise HTTPException(403, "Choose a registered image endpoint")' in body, (
        "harmonize_image must raise 403 when _endpoint doesn't match a registered endpoint"
    )


def test_harmonize_outbound_base_not_from_request_body():
    """Confirm _join_checked_gallery_endpoint is never called with requested_base."""
    src = SRC.read_text()
    body = _function_source(src, "harmonize_image")
    assert "_join_checked_gallery_endpoint(requested_base," not in body, (
        "harmonize_image must not pass requested_base to _join_checked_gallery_endpoint"
    )


def test_inpaint_and_harmonize_no_base_equals_endpoint():
    """Neither function should assign `base = endpoint` or `base = requested_base`
    — the outbound base must come exclusively from DB (ep.base_url)."""
    src = SRC.read_text()
    for func_name in ("inpaint_proxy", "harmonize_image"):
        body = _function_source(src, func_name)
        assert "base = endpoint" not in body, (
            f"{func_name}: 'base = endpoint' carries request-body input into outbound request"
        )
        assert "base = requested_base" not in body, (
            f"{func_name}: 'base = requested_base' carries request-body input into outbound request"
        )