HERMES-GIELIE/tests/test_issue1105_ssrf_custom_providers.py

"""Tests for #1105 — SSRF check allows user-configured custom_providers hostnames.

The SSRF check blocks requests to private IPs unless the hostname is in a
hardcoded allowlist. This fix extracts hostnames from custom_providers config
and adds them to the trusted set, so user-explicitly configured local endpoints
(ollama, llama.cpp, vLLM, TabbyAPI, etc.) are not blocked.
"""
import os
import pytest


# ---------- Source-code analysis tests ----------

def test_ssrf_trusted_hosts_variable_exists():
    """The _ssrf_trusted_hosts set must be built from custom_providers config."""
    with open("api/config.py") as f:
        src = f.read()
    assert "_ssrf_trusted_hosts" in src
    assert "_ssrf_trusted_hosts: set[str] = set()" in src


def test_ssrf_trusted_hosts_populated_from_custom_providers():
    """Trusted hosts are extracted by iterating custom_providers[].base_url."""
    with open("api/config.py") as f:
        src = f.read()
    # Must read custom_providers from cfg
    assert 'cfg.get("custom_providers"' in src
    # Must extract base_url from each entry
    assert '_cp.get("base_url")' in src
    # Must parse hostname with urlparse
    assert "_cp_parsed.hostname" in src
    # Must add to trusted set
    assert "_ssrf_trusted_hosts.add" in src


def test_ssrf_check_uses_trusted_hosts():
    """The SSRF check must consult _ssrf_trusted_hosts before blocking."""
    with open("api/config.py") as f:
        src = f.read()
    # The is_known_local check must include _ssrf_trusted_hosts
    assert "in _ssrf_trusted_hosts" in src


def test_ssrf_known_local_still_present():
    """Original hardcoded allowlist must still be present (no regression)."""
    with open("api/config.py") as f:
        src = f.read()
    for keyword in ("ollama", "localhost", "127.0.0.1", "lmstudio", "lm-studio"):
        assert keyword in src, f"Missing hardcoded allowlist entry: {keyword}"


def test_ssrf_block_still_present():
    """SSRF ValueError must still be raised for unknown private IPs."""
    with open("api/config.py") as f:
        src = f.read()
    assert 'SSRF: resolved hostname to private IP' in src


# ---------- Functional tests (mocked socket) ----------

def test_custom_provider_hostname_added_to_trusted():
    """A hostname from custom_providers base_url is added to trusted set."""
    import api.config as config
    import socket
    from unittest.mock import patch, MagicMock

    old_cfg = dict(config.cfg)
    try:
        config.cfg.update({
            "model": {"model": "my-model", "base_url": "http://my-llama-server:8080/v1"},
            "custom_providers": [
                {"name": "my-llama", "base_url": "http://my-llama-server:8080/v1", "model": "llama-3"}
            ],
            "providers": {},
        })
        config.invalidate_models_cache()

        # Mock socket.getaddrinfo to return a private IP for my-llama-server
        private_addr = ("192.168.1.100", None)
        mock_getaddrinfo = MagicMock(return_value=[
            (socket.AF_INET, socket.SOCK_STREAM, 6, "", private_addr)
        ])

        # Mock urllib to prevent actual HTTP call
        mock_urlopen = MagicMock()
        mock_urlopen.read.return_value = b'{"data": [{"id": "test-model"}]}'
        mock_urlopen.__enter__ = MagicMock(return_value=mock_urlopen)
        mock_urlopen.__exit__ = MagicMock(return_value=False)

        with patch("socket.getaddrinfo", mock_getaddrinfo), \
             patch("urllib.request.urlopen", mock_urlopen):
            # Should NOT raise ValueError (SSRF) because hostname is in trusted set
            result = config.get_available_models()

        # Verify models were returned (auto-detection succeeded)
        assert result is not None
        assert "groups" in result

    finally:
        config.cfg.clear()
        config.cfg.update(old_cfg)
        config.invalidate_models_cache()


def test_unknown_private_ip_still_blocked():
    """A private IP from a hostname NOT in custom_providers is still blocked.

    The SSRF ValueError is caught by the broad `except Exception` around
    the custom endpoint fetch (line ~1571), so get_available_models() doesn't
    crash — but no models are auto-detected from that endpoint.
    """
    import api.config as config
    import socket
    from unittest.mock import patch, MagicMock

    old_cfg = dict(config.cfg)
    try:
        config.cfg.update({
            "model": {"model": "test", "base_url": "http://unknown-local-server:9999/v1"},
            "custom_providers": [
                {"name": "other", "base_url": "http://other-server:8080/v1", "model": "x"}
            ],
            "providers": {},
        })
        config.invalidate_models_cache()

        # Mock socket.getaddrinfo to return a private IP for unknown-local-server
        private_addr = ("10.0.0.50", None)
        mock_getaddrinfo = MagicMock(return_value=[
            (socket.AF_INET, socket.SOCK_STREAM, 6, "", private_addr)
        ])

        with patch("socket.getaddrinfo", mock_getaddrinfo):
            # Should NOT crash (ValueError is caught internally)
            result = config.get_available_models()

        # But no models should be auto-detected from the blocked endpoint
        assert result is not None
        assert "groups" in result
        # Verify no group with "unknown-local-server" models exists
        for group in result["groups"]:
            provider_name = group.get("provider", "")
            assert "unknown-local-server" not in provider_name

    finally:
        config.cfg.clear()
        config.cfg.update(old_cfg)
        config.invalidate_models_cache()