llamacppctl/tests/test_prompt_urls.py

import socket
from dataclasses import replace

import pytest

from llamacppctl.prompt_io import (
    PromptSourceError,
    _pin_dns,
    load_text_url,
    validate_url_target,
)


class FakeResponse:
    def __init__(self, status_code=200, headers=None, chunks=None, text=""):
        self.status_code = status_code
        self.headers = headers or {}
        self._chunks = chunks if chunks is not None else [text.encode("utf-8")]

    def raise_for_status(self):
        if self.status_code >= 400:
            raise RuntimeError(f"HTTP {self.status_code}")

    def iter_content(self, chunk_size=8192):
        for c in self._chunks:
            yield c


def _fake_getaddrinfo_factory(mapping):
    def fake_getaddrinfo(host, port, *args, **kwargs):
        ips = mapping.get(host, [])
        return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", (ip, port or 0)) for ip in ips]

    return fake_getaddrinfo


def test_reject_http_without_opt_in(policy):
    with pytest.raises(PromptSourceError):
        validate_url_target("http://example.org/x.txt", policy)


def test_reject_localhost(policy):
    with pytest.raises(PromptSourceError):
        validate_url_target("https://localhost/x.txt", policy)


def test_reject_ip_literal_without_opt_in(policy):
    with pytest.raises(PromptSourceError):
        validate_url_target("https://93.184.216.34/x.txt", policy)


def test_allow_ip_literal_with_opt_in(monkeypatch, policy):
    strict_policy = replace(policy, allow_ip_host=True)
    monkeypatch.setattr(
        socket, "getaddrinfo", _fake_getaddrinfo_factory({"93.184.216.34": ["93.184.216.34"]})
    )
    validate_url_target("https://93.184.216.34/x.txt", strict_policy)


def test_reject_private_ip_resolution(monkeypatch, policy):
    monkeypatch.setattr(
        socket, "getaddrinfo", _fake_getaddrinfo_factory({"internal.example.org": ["10.0.0.5"]})
    )
    with pytest.raises(PromptSourceError):
        validate_url_target("https://internal.example.org/x.txt", policy)


def test_reject_metadata_ip(monkeypatch, policy):
    monkeypatch.setattr(
        socket, "getaddrinfo", _fake_getaddrinfo_factory({"evil.example.org": ["169.254.169.254"]})
    )
    with pytest.raises(PromptSourceError):
        validate_url_target("https://evil.example.org/x.txt", policy)


def test_reject_embedded_credentials(policy):
    with pytest.raises(PromptSourceError):
        validate_url_target("https://user:pass@example.org/x.txt", policy)


def test_allowlist_blocks_other_hosts(monkeypatch, policy):
    strict_policy = replace(policy, url_allow_hosts=["good.example.org"])
    monkeypatch.setattr(
        socket,
        "getaddrinfo",
        _fake_getaddrinfo_factory({"bad.example.org": ["93.184.216.34"]}),
    )
    with pytest.raises(PromptSourceError):
        validate_url_target("https://bad.example.org/x.txt", strict_policy)


def test_accept_https_text_plain(monkeypatch, policy):
    monkeypatch.setattr(
        socket, "getaddrinfo", _fake_getaddrinfo_factory({"example.org": ["93.184.216.34"]})
    )

    def fake_get(url, **kwargs):
        return FakeResponse(
            status_code=200,
            headers={"Content-Type": "text/plain", "Content-Length": "5"},
            text="hello",
        )

    import llamacppctl.prompt_io as prompt_io

    monkeypatch.setattr(prompt_io.requests, "get", fake_get)
    text = load_text_url("https://example.org/x.txt", policy)
    assert text == "hello"


def test_reject_disallowed_content_type(monkeypatch, policy):
    monkeypatch.setattr(
        socket, "getaddrinfo", _fake_getaddrinfo_factory({"example.org": ["93.184.216.34"]})
    )

    def fake_get(url, **kwargs):
        return FakeResponse(status_code=200, headers={"Content-Type": "text/html"}, text="<h1>hi</h1>")

    import llamacppctl.prompt_io as prompt_io

    monkeypatch.setattr(prompt_io.requests, "get", fake_get)
    with pytest.raises(PromptSourceError):
        load_text_url("https://example.org/x.txt", policy)


def test_reject_large_content_length(monkeypatch, small_policy):
    monkeypatch.setattr(
        socket, "getaddrinfo", _fake_getaddrinfo_factory({"example.org": ["93.184.216.34"]})
    )

    def fake_get(url, **kwargs):
        return FakeResponse(
            status_code=200,
            headers={"Content-Type": "text/plain", "Content-Length": "1000"},
            text="x" * 1000,
        )

    import llamacppctl.prompt_io as prompt_io

    monkeypatch.setattr(prompt_io.requests, "get", fake_get)
    with pytest.raises(PromptSourceError):
        load_text_url("https://example.org/x.txt", small_policy)


def test_reject_large_streamed_response(monkeypatch, small_policy):
    monkeypatch.setattr(
        socket, "getaddrinfo", _fake_getaddrinfo_factory({"example.org": ["93.184.216.34"]})
    )

    def fake_get(url, **kwargs):
        return FakeResponse(
            status_code=200,
            headers={"Content-Type": "text/plain"},
            chunks=[b"x" * 10, b"x" * 10, b"x" * 10],
        )

    import llamacppctl.prompt_io as prompt_io

    monkeypatch.setattr(prompt_io.requests, "get", fake_get)
    with pytest.raises(PromptSourceError):
        load_text_url("https://example.org/x.txt", small_policy)


def test_reject_redirect_by_default(monkeypatch, policy):
    monkeypatch.setattr(
        socket, "getaddrinfo", _fake_getaddrinfo_factory({"example.org": ["93.184.216.34"]})
    )

    def fake_get(url, **kwargs):
        return FakeResponse(status_code=302, headers={"Location": "https://example.org/other.txt"})

    import llamacppctl.prompt_io as prompt_io

    monkeypatch.setattr(prompt_io.requests, "get", fake_get)
    with pytest.raises(PromptSourceError):
        load_text_url("https://example.org/x.txt", policy)


def test_pin_dns_allows_validated_ip(monkeypatch):
    monkeypatch.setattr(
        socket, "getaddrinfo", _fake_getaddrinfo_factory({"host.example": ["1.2.3.4"]})
    )
    with _pin_dns("host.example", ["1.2.3.4"]):
        assert socket.getaddrinfo("host.example", 443)  # not filtered away


def test_pin_dns_rejects_rebound_ip(monkeypatch):
    # Simulate DNS rebinding: after validation to 1.2.3.4, DNS now returns a
    # private IP; the pin must refuse it during the actual connect.
    monkeypatch.setattr(
        socket, "getaddrinfo", _fake_getaddrinfo_factory({"host.example": ["10.0.0.9"]})
    )
    with _pin_dns("host.example", ["1.2.3.4"]):
        with pytest.raises(socket.gaierror):
            socket.getaddrinfo("host.example", 443)


def test_pin_dns_does_not_touch_other_hosts(monkeypatch):
    monkeypatch.setattr(
        socket, "getaddrinfo", _fake_getaddrinfo_factory({"other.example": ["9.9.9.9"]})
    )
    with _pin_dns("host.example", ["1.2.3.4"]):
        assert socket.getaddrinfo("other.example", 443)  # unrelated host passes


def test_pin_dns_restores_resolver(monkeypatch):
    sentinel = _fake_getaddrinfo_factory({"host.example": ["1.2.3.4"]})
    monkeypatch.setattr(socket, "getaddrinfo", sentinel)
    with _pin_dns("host.example", ["1.2.3.4"]):
        pass
    assert socket.getaddrinfo is sentinel  # restored after the context


def test_revalidate_redirect_target(monkeypatch, policy):
    strict_policy = replace(policy, follow_redirects=True)
    monkeypatch.setattr(
        socket,
        "getaddrinfo",
        _fake_getaddrinfo_factory(
            {"example.org": ["93.184.216.34"], "internal.example.org": ["10.0.0.9"]}
        ),
    )

    calls = {"n": 0}

    def fake_get(url, **kwargs):
        calls["n"] += 1
        if "internal" not in url:
            return FakeResponse(
                status_code=302, headers={"Location": "https://internal.example.org/secret.txt"}
            )
        return FakeResponse(status_code=200, headers={"Content-Type": "text/plain"}, text="leak")

    import llamacppctl.prompt_io as prompt_io

    monkeypatch.setattr(prompt_io.requests, "get", fake_get)
    with pytest.raises(PromptSourceError):
        load_text_url("https://example.org/x.txt", strict_policy)
    assert calls["n"] == 1  # must fail validation before the second request