open-swe/tests/test_http_security.py

from __future__ import annotations

import importlib
import socket as real_socket
import sys
import types
from typing import Any
from urllib.parse import urlparse

import httpx
import pytest

exa_py_stub = types.ModuleType("exa_py")
exa_py_stub.Exa = object
sys.modules.setdefault("exa_py", exa_py_stub)

importlib.import_module("agent.tools.fetch_url")
importlib.import_module("agent.tools.http_request")
fetch_url_tool = sys.modules["agent.tools.fetch_url"]
http_request_tool = sys.modules["agent.tools.http_request"]
# SSRF resolution now lives in the shared validator; patch DNS there.
url_safety = importlib.import_module("agent.utils.url_safety")
multimodal = importlib.import_module("agent.utils.multimodal")

_NO_JSON = object()


def _addr_info(ip: str, port: int | None = None) -> tuple:
    return (
        real_socket.AF_INET,
        real_socket.SOCK_STREAM,
        6,
        "",
        (ip, port or 0),
    )


class FakeResponse:
    def __init__(
        self,
        *,
        status_code: int,
        url: str,
        headers: dict[str, str] | None = None,
        text: str = "",
        json_data: object = _NO_JSON,
    ) -> None:
        self.status_code = status_code
        self.url = url
        self.headers = headers or {}
        self.text = text
        self._json_data = json_data

    def json(self) -> object:
        if self._json_data is _NO_JSON:
            raise ValueError("response is not json")
        return self._json_data

    def raise_for_status(self) -> None:
        if self.status_code >= 400:
            raise httpx.HTTPStatusError(f"{self.status_code} error", request=None, response=None)


class FakeAsyncClient:
    """Records each request and replays programmed responses.

    ``responder(method, url, **kwargs)`` returns a ``FakeResponse``. The class is
    installed in place of ``httpx.AsyncClient`` on the tool module under test.
    """

    last_instance: FakeAsyncClient | None = None

    def __init__(self, responder, *args: Any, **kwargs: Any) -> None:
        self._responder = responder
        self.calls: list[dict[str, Any]] = []
        FakeAsyncClient.last_instance = self

    async def __aenter__(self) -> FakeAsyncClient:
        return self

    async def __aexit__(self, *exc: Any) -> bool:
        return False

    async def request(self, method: str, url: str, **kwargs: Any) -> FakeResponse:
        self.calls.append({"method": method, "url": url, **kwargs})
        return self._responder(method, url, **kwargs)


def _install_client(monkeypatch, module, responder) -> type:
    def factory(*args: Any, **kwargs: Any) -> FakeAsyncClient:
        return FakeAsyncClient(responder, *args, **kwargs)

    fake_httpx = types.SimpleNamespace(
        AsyncClient=factory,
        HTTPError=httpx.HTTPError,
        TimeoutException=httpx.TimeoutException,
    )
    monkeypatch.setattr(module, "httpx", fake_httpx)
    return factory


# --- _resolve_and_validate (pure IP gating) ----------------------------------


def test_resolve_and_validate_rejects_unsupported_scheme() -> None:
    is_safe, reason, _, _ = http_request_tool._resolve_and_validate("ftp://example.com/x")
    assert is_safe is False
    assert "scheme" in reason.lower()


@pytest.mark.parametrize(
    "ip",
    ["127.0.0.1", "169.254.169.254", "10.0.0.5", "192.168.1.1"],
)
def test_resolve_and_validate_rejects_private_ranges(monkeypatch, ip: str) -> None:
    monkeypatch.setattr(
        url_safety.socket,
        "getaddrinfo",
        lambda host, port, *a, **k: [_addr_info(ip, port)],
    )
    is_safe, reason, hostname, _ = http_request_tool._resolve_and_validate("http://evil.test/")
    assert is_safe is False
    assert "blocked address" in reason
    assert hostname == "evil.test"


def test_resolve_and_validate_accepts_public_ip(monkeypatch) -> None:
    monkeypatch.setattr(
        url_safety.socket,
        "getaddrinfo",
        lambda host, port, *a, **k: [_addr_info("93.184.216.34", port)],
    )
    is_safe, reason, hostname, addr_infos = http_request_tool._resolve_and_validate(
        "https://example.com/path"
    )
    assert is_safe is True
    assert reason == ""
    assert hostname == "example.com"
    assert addr_infos[0][4][0] == "93.184.216.34"


def test_pinned_url_rewrites_host_to_ip_keeping_path_and_port() -> None:
    assert (
        http_request_tool._pinned_url("https://example.com:8443/a/b?q=1", "93.184.216.34")
        == "https://93.184.216.34:8443/a/b?q=1"
    )
    # IPv6 literal is bracketed
    assert http_request_tool._pinned_url("http://h/x", "::1").startswith("http://[::1]/x")


# --- fetch_url ---------------------------------------------------------------


async def test_fetch_url_blocks_private_ip_without_issuing_a_request(monkeypatch) -> None:
    def fail_responder(*args: Any, **kwargs: Any) -> FakeResponse:
        raise AssertionError("request should not be issued for blocked URLs")

    _install_client(monkeypatch, fetch_url_tool, fail_responder)
    # Real DNS resolution of the metadata IP literal yields the private IP itself.

    result = await fetch_url_tool.fetch_url(
        "http://169.254.169.254/latest/meta-data/iam/security-credentials/"
    )

    assert result["status_code"] == 0
    assert "Request blocked" in result["error"]
    assert result["url"].startswith("http://169.254.169.254/")


async def test_fetch_url_blocks_redirects_to_private_ips(monkeypatch) -> None:
    def fake_getaddrinfo(host, port, *args, **kwargs):  # type: ignore[no-untyped-def]
        ip = "93.184.216.34" if host == "example.com" else host
        return [_addr_info(ip, port)]

    monkeypatch.setattr(url_safety.socket, "getaddrinfo", fake_getaddrinfo)

    def responder(method: str, url: str, **kwargs: Any) -> FakeResponse:
        return FakeResponse(
            status_code=302,
            url=url,
            headers={"Location": "http://169.254.169.254/latest/meta-data"},
        )

    _install_client(monkeypatch, fetch_url_tool, responder)

    result = await fetch_url_tool.fetch_url("https://example.com/start")

    # First hop targets the validated public IP, with Host preserved.
    client = FakeAsyncClient.last_instance
    assert client is not None
    assert len(client.calls) == 1
    first = client.calls[0]
    assert urlparse(first["url"]).hostname == "93.184.216.34"
    assert first["headers"]["Host"] == "example.com"
    assert first["extensions"]["sni_hostname"] == "example.com"
    # The redirect to a private IP was blocked before a second request was issued.
    assert result["status_code"] == 0
    assert result["url"] == "http://169.254.169.254/latest/meta-data"
    assert "Request blocked" in result["error"]


# --- http_request ------------------------------------------------------------


async def test_http_request_pins_connection_to_validated_public_ip(monkeypatch) -> None:
    """Validation sees a public IP and the connection must target that exact IP.

    A resolver that later flips to a private address cannot rebind because the
    request URL is pinned to the validated IP (with Host + SNI preserved).
    """
    hostname = "rebind.example.com"
    public_addr = "93.184.216.34"
    private_addr = "127.0.0.1"

    call_count = {"n": 0}

    def fake_getaddrinfo(host, port, *args, **kwargs):  # type: ignore[no-untyped-def]
        call_count["n"] += 1
        ip = public_addr if call_count["n"] == 1 else private_addr
        return [_addr_info(ip, port)]

    monkeypatch.setattr(url_safety.socket, "getaddrinfo", fake_getaddrinfo)

    def responder(method: str, url: str, **kwargs: Any) -> FakeResponse:
        return FakeResponse(status_code=200, url=url, text="ok", json_data="ok")

    _install_client(monkeypatch, http_request_tool, responder)

    result = await http_request_tool.http_request(f"http://{hostname}/probe")

    client = FakeAsyncClient.last_instance
    assert client is not None
    assert len(client.calls) == 1
    call = client.calls[0]
    assert urlparse(call["url"]).hostname == public_addr, (
        f"connection must target pinned public IP, got {call['url']}"
    )
    assert call["headers"]["Host"] == hostname
    assert call["extensions"]["sni_hostname"] == hostname
    assert result["status_code"] == 200


async def test_http_request_blocks_when_only_private_ips(monkeypatch) -> None:
    """If the first resolution returns a private IP, no request is issued."""
    hostname = "evil.example.com"
    private_addr = "169.254.169.254"

    monkeypatch.setattr(
        url_safety.socket,
        "getaddrinfo",
        lambda host, port, *a, **k: [_addr_info(private_addr, port)],
    )

    def fail_responder(*args: Any, **kwargs: Any) -> FakeResponse:
        raise AssertionError("request should not be issued for blocked URLs")

    _install_client(monkeypatch, http_request_tool, fail_responder)

    result = await http_request_tool.http_request(f"http://{hostname}/")

    assert result["status_code"] == 0
    assert "Request blocked" in result["content"]


async def test_http_request_downgrades_method_on_303(monkeypatch) -> None:
    """A 303 redirect must switch the follow-up request to GET and drop the body."""

    def fake_getaddrinfo(host, port, *args, **kwargs):  # type: ignore[no-untyped-def]
        return [_addr_info("93.184.216.34", port)]

    monkeypatch.setattr(url_safety.socket, "getaddrinfo", fake_getaddrinfo)

    def responder(method: str, url: str, **kwargs: Any) -> FakeResponse:
        if "start" in url:
            return FakeResponse(
                status_code=303,
                url=url,
                headers={"Location": "https://example.com/done"},
            )
        return FakeResponse(status_code=200, url=url, json_data={"ok": True})

    _install_client(monkeypatch, http_request_tool, responder)

    result = await http_request_tool.http_request(
        "https://example.com/start", method="POST", data={"x": 1}
    )

    client = FakeAsyncClient.last_instance
    assert client is not None
    assert len(client.calls) == 2
    assert client.calls[0]["method"] == "POST"
    assert client.calls[1]["method"] == "GET"
    assert "json" not in client.calls[1] and "content" not in client.calls[1]
    assert result["status_code"] == 200
    assert result["content"] == {"ok": True}


async def test_http_request_returns_timeout_result(monkeypatch) -> None:
    def fake_getaddrinfo(host, port, *args, **kwargs):  # type: ignore[no-untyped-def]
        return [_addr_info("93.184.216.34", port)]

    monkeypatch.setattr(url_safety.socket, "getaddrinfo", fake_getaddrinfo)

    def responder(method: str, url: str, **kwargs: Any) -> FakeResponse:
        raise httpx.TimeoutException("timed out")

    _install_client(monkeypatch, http_request_tool, responder)

    result = await http_request_tool.http_request("https://example.com/", timeout=7)

    assert result["success"] is False
    assert result["status_code"] == 0
    assert "timed out after 7 seconds" in result["content"]


# --- fetch_image_block SSRF (INJ-01 regression) ------------------------------


def _image_dns(monkeypatch, ip_map: dict[str, str]) -> None:
    monkeypatch.setattr(
        url_safety.socket,
        "getaddrinfo",
        lambda host, port, *a, **k: [_addr_info(ip_map[host], port)],
    )


async def test_fetch_image_block_blocks_redirect_to_internal_host(monkeypatch) -> None:
    """A public image URL that 302-redirects to the metadata endpoint must not be
    followed: the redirect target is re-validated and the internal hop blocked."""
    _image_dns(
        monkeypatch,
        {"public.example": "93.184.216.34", "internal.evil": "169.254.169.254"},
    )

    def responder(method, url, **kwargs):
        return FakeResponse(
            status_code=302,
            url=url,
            headers={"Location": "http://internal.evil/latest/meta-data/"},
        )

    client = FakeAsyncClient(responder)
    result = await multimodal.fetch_image_block("http://public.example/x.png", client)

    assert result is None
    # Exactly one hop was issued (to the validated public host); the redirect to
    # the internal host was blocked before any request left the process.
    assert len(client.calls) == 1
    assert all("169.254.169.254" not in call["url"] for call in client.calls)


async def test_fetch_image_block_strips_auth_header_on_redirect(monkeypatch) -> None:
    """The per-host Slack bearer token must not be replayed to a redirect target."""
    monkeypatch.setenv("SLACK_BOT_TOKEN", "xoxb-secret")
    _image_dns(
        monkeypatch,
        {"files.slack.com": "93.184.216.34", "cdn.public": "93.184.216.35"},
    )

    def responder(method, url, **kwargs):
        if len(client.calls) == 1:
            return FakeResponse(
                status_code=302, url=url, headers={"Location": "https://cdn.public/x.png"}
            )
        resp = FakeResponse(status_code=200, url=url, headers={"Content-Type": "image/png"})
        resp.content = b"\x89PNG\r\n"
        return resp

    client = FakeAsyncClient(responder)
    result = await multimodal.fetch_image_block("https://files.slack.com/x.png", client)

    assert result is not None
    assert len(client.calls) == 2
    first_auth = client.calls[0]["headers"].get("Authorization")
    second_auth = client.calls[1]["headers"].get("Authorization")
    assert first_auth == "Bearer xoxb-secret"
    assert second_auth is None