learningcircuit--local-deep-research/tests/security/test_ssrf_response_size.py

"""Tests for SSRF response size check restructuring (PR #1942).

Verifies try/except/else pattern ensures ValueError from int() parsing
does not mask ValueError from legitimate size-exceeded checks.
"""

import requests
import pytest
from unittest.mock import patch, MagicMock

from io import BytesIO

from local_deep_research.security import ssrf_validator
from local_deep_research.security.safe_requests import (
    safe_get,
    safe_post,
    SafeSession,
    _check_response_size,
    _install_body_guard,
    MAX_RESPONSE_SIZE,
)


@pytest.fixture
def mock_validate_url():
    with patch.object(
        ssrf_validator,
        "validate_url",
        return_value=True,
    ) as m:
        yield m


class TestSafeGetResponseSize:
    """Verify safe_get Content-Length handling with try/except/else."""

    def test_large_content_length_raises(self, mock_validate_url):
        """A Content-Length exceeding MAX_RESPONSE_SIZE should raise ValueError."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": str(MAX_RESPONSE_SIZE + 1)}

        with patch(
            "local_deep_research.security.safe_requests.requests.get",
            return_value=mock_resp,
        ):
            with pytest.raises(ValueError, match="Response too large"):
                safe_get("https://example.com")
            mock_resp.close.assert_called_once()

    def test_valid_content_length_returns_response(self, mock_validate_url):
        """A Content-Length within limit should return the response."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": "1024"}

        with patch(
            "local_deep_research.security.safe_requests.requests.get",
            return_value=mock_resp,
        ):
            result = safe_get("https://example.com")
            assert result is mock_resp

    def test_invalid_content_length_ignored(self, mock_validate_url):
        """Non-numeric Content-Length should be silently ignored."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": "not-a-number"}

        with patch(
            "local_deep_research.security.safe_requests.requests.get",
            return_value=mock_resp,
        ):
            result = safe_get("https://example.com")
            assert result is mock_resp

    def test_missing_content_length_returns_response(self, mock_validate_url):
        """Missing Content-Length header should not affect response."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {}

        with patch(
            "local_deep_research.security.safe_requests.requests.get",
            return_value=mock_resp,
        ):
            result = safe_get("https://example.com")
            assert result is mock_resp

    def test_exact_max_size_allowed(self, mock_validate_url):
        """Content-Length exactly at MAX_RESPONSE_SIZE should be allowed."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": str(MAX_RESPONSE_SIZE)}

        with patch(
            "local_deep_research.security.safe_requests.requests.get",
            return_value=mock_resp,
        ):
            result = safe_get("https://example.com")
            assert result is mock_resp

    def test_comma_separated_identical_content_length(self, mock_validate_url):
        """Identical comma-separated Content-Length values should be accepted."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": "1024, 1024"}

        with patch(
            "local_deep_research.security.safe_requests.requests.get",
            return_value=mock_resp,
        ):
            result = safe_get("https://example.com")
            assert result is mock_resp

    def test_comma_separated_different_content_length_rejected(
        self, mock_validate_url
    ):
        """Differing comma-separated Content-Length values should be rejected."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": "100, 200"}

        with patch(
            "local_deep_research.security.safe_requests.requests.get",
            return_value=mock_resp,
        ):
            with pytest.raises(ValueError, match="Conflicting"):
                safe_get("https://example.com")
            mock_resp.close.assert_called_once()

    def test_comma_separated_oversized_content_length(self, mock_validate_url):
        """Identical but oversized comma-separated values should be rejected."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {
            "Content-Length": f"{MAX_RESPONSE_SIZE + 1}, {MAX_RESPONSE_SIZE + 1}"
        }

        with patch(
            "local_deep_research.security.safe_requests.requests.get",
            return_value=mock_resp,
        ):
            with pytest.raises(ValueError, match="Response too large"):
                safe_get("https://example.com")
            mock_resp.close.assert_called_once()

    def test_zero_content_length_accepted(self, mock_validate_url):
        """Zero Content-Length should be accepted."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": "0"}

        with patch(
            "local_deep_research.security.safe_requests.requests.get",
            return_value=mock_resp,
        ):
            result = safe_get("https://example.com")
            assert result is mock_resp

    def test_negative_content_length_ignored(self, mock_validate_url):
        """Negative Content-Length should be treated as absent."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": "-1"}

        with patch(
            "local_deep_research.security.safe_requests.requests.get",
            return_value=mock_resp,
        ):
            result = safe_get("https://example.com")
            assert result is mock_resp

    def test_trailing_comma_within_limit_accepted(self, mock_validate_url):
        """Trailing comma with valid size should be accepted."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": "100,"}

        with patch(
            "local_deep_research.security.safe_requests.requests.get",
            return_value=mock_resp,
        ):
            result = safe_get("https://example.com")
            assert result is mock_resp

    def test_trailing_comma_oversized_raises(self, mock_validate_url):
        """Trailing comma must not bypass size check for oversized values."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": f"{MAX_RESPONSE_SIZE + 1},"}

        with patch(
            "local_deep_research.security.safe_requests.requests.get",
            return_value=mock_resp,
        ):
            with pytest.raises(ValueError, match="Response too large"):
                safe_get("https://example.com")
            mock_resp.close.assert_called_once()

    def test_double_comma_conflicting_values_rejected(self, mock_validate_url):
        """Double comma between values must not mask a conflict."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": "100,,200"}

        with patch(
            "local_deep_research.security.safe_requests.requests.get",
            return_value=mock_resp,
        ):
            with pytest.raises(ValueError, match="Conflicting"):
                safe_get("https://example.com")
            mock_resp.close.assert_called_once()

    def test_only_commas_treated_as_absent(self, mock_validate_url):
        """Header containing only commas should be treated as absent."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": ","}

        with patch(
            "local_deep_research.security.safe_requests.requests.get",
            return_value=mock_resp,
        ):
            result = safe_get("https://example.com")
            assert result is mock_resp

    def test_trailing_comma_identical_values_accepted(self, mock_validate_url):
        """Identical values with trailing comma should be accepted."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": "100, 100,"}

        with patch(
            "local_deep_research.security.safe_requests.requests.get",
            return_value=mock_resp,
        ):
            result = safe_get("https://example.com")
            assert result is mock_resp


class TestSafePostResponseSize:
    """Verify safe_post Content-Length handling with try/except/else."""

    def test_large_content_length_raises(self, mock_validate_url):
        """A Content-Length exceeding MAX_RESPONSE_SIZE should raise ValueError."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": str(MAX_RESPONSE_SIZE + 1)}

        with patch(
            "local_deep_research.security.safe_requests.requests.post",
            return_value=mock_resp,
        ):
            with pytest.raises(ValueError, match="Response too large"):
                safe_post("https://example.com")
            mock_resp.close.assert_called_once()

    def test_valid_content_length_returns_response(self, mock_validate_url):
        """A Content-Length within limit should return the response."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": "512"}

        with patch(
            "local_deep_research.security.safe_requests.requests.post",
            return_value=mock_resp,
        ):
            result = safe_post("https://example.com")
            assert result is mock_resp

    def test_invalid_content_length_ignored(self, mock_validate_url):
        """Non-numeric Content-Length should be silently ignored."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": "abc"}

        with patch(
            "local_deep_research.security.safe_requests.requests.post",
            return_value=mock_resp,
        ):
            result = safe_post("https://example.com")
            assert result is mock_resp

    def test_trailing_comma_oversized_raises(self, mock_validate_url):
        """Trailing comma must not bypass size check for oversized values."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": f"{MAX_RESPONSE_SIZE + 1},"}

        with patch(
            "local_deep_research.security.safe_requests.requests.post",
            return_value=mock_resp,
        ):
            with pytest.raises(ValueError, match="Response too large"):
                safe_post("https://example.com")
            mock_resp.close.assert_called_once()


class TestSafeSessionResponseSize:
    """Verify SafeSession.send() size checks on final and redirect responses."""

    def test_session_large_content_length_raises(self, mock_validate_url):
        """Oversized Content-Length on final response raises ValueError."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": str(MAX_RESPONSE_SIZE + 1)}

        with patch(
            "local_deep_research.security.safe_requests.requests.Session.send",
            return_value=mock_resp,
        ):
            with SafeSession() as session:
                with pytest.raises(ValueError, match="Response too large"):
                    session.get("https://example.com")
                mock_resp.close.assert_called_once()

    def test_session_valid_content_length_returns(self, mock_validate_url):
        """Normal Content-Length returns the response."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {"Content-Length": "1024"}

        with patch(
            "local_deep_research.security.safe_requests.requests.Session.send",
            return_value=mock_resp,
        ):
            with SafeSession() as session:
                result = session.get("https://example.com")
                assert result is mock_resp

    def test_session_missing_content_length_returns(self, mock_validate_url):
        """No Content-Length header returns the response."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {}

        with patch(
            "local_deep_research.security.safe_requests.requests.Session.send",
            return_value=mock_resp,
        ):
            with SafeSession() as session:
                result = session.get("https://example.com")
                assert result is mock_resp

    def test_session_redirect_with_oversized_final(self, mock_validate_url):
        """302→200: size check on redirect hop is harmless, oversized final raises."""
        redirect_resp = MagicMock()
        redirect_resp.status_code = 302
        redirect_resp.headers = {"Location": "https://example.com/final"}

        final_resp = MagicMock()
        final_resp.status_code = 200
        final_resp.headers = {"Content-Length": str(MAX_RESPONSE_SIZE + 1)}

        session = SafeSession()
        prep = requests.Request("GET", "https://example.com").prepare()

        with patch(
            "local_deep_research.security.safe_requests.requests.Session.send",
            side_effect=[redirect_resp, final_resp],
        ):
            # Intermediate redirect: no CL, size check passes harmlessly
            result = session.send(prep)
            assert result is redirect_resp

            # Final response: oversized CL triggers ValueError
            with pytest.raises(ValueError, match="Response too large"):
                session.send(prep)
            final_resp.close.assert_called_once()

    def test_session_redirect_with_valid_final(self, mock_validate_url):
        """302→200: redirect hop passes, valid final returns normally."""
        redirect_resp = MagicMock()
        redirect_resp.status_code = 302
        redirect_resp.headers = {"Location": "https://example.com/final"}

        final_resp = MagicMock()
        final_resp.status_code = 200
        final_resp.headers = {"Content-Length": "512"}

        session = SafeSession()
        prep = requests.Request("GET", "https://example.com").prepare()

        with patch(
            "local_deep_research.security.safe_requests.requests.Session.send",
            side_effect=[redirect_resp, final_resp],
        ):
            result = session.send(prep)
            assert result is redirect_resp

            result = session.send(prep)
            assert result is final_resp

    def test_session_comma_separated_content_length(self, mock_validate_url):
        """Multi-value Content-Length handled correctly via session path."""
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.headers = {
            "Content-Length": f"{MAX_RESPONSE_SIZE + 1}, {MAX_RESPONSE_SIZE + 1}"
        }

        with patch(
            "local_deep_research.security.safe_requests.requests.Session.send",
            return_value=mock_resp,
        ):
            with SafeSession() as session:
                with pytest.raises(ValueError, match="Response too large"):
                    session.get("https://example.com")
                mock_resp.close.assert_called_once()


class TestBodySizeEnforcement:
    """Verify _install_body_guard enforces size limits on responses without CL."""

    def _make_response(self, body: bytes, headers: dict | None = None):
        """Create a mock response with a real readable raw body."""
        resp = MagicMock(spec=requests.Response)
        resp.headers = headers or {}
        raw = MagicMock()
        stream = BytesIO(body)
        raw.read = stream.read
        resp.raw = raw
        return resp

    def test_chunked_response_within_limit(self):
        """Small body without Content-Length succeeds."""
        body = b"x" * 1024
        resp = self._make_response(body)
        _install_body_guard(resp)

        data = resp.raw.read()
        assert data == body

    def test_chunked_response_exceeds_limit(self):
        """Large body without Content-Length raises ValueError on read."""
        body = b"x" * (MAX_RESPONSE_SIZE + 1)
        resp = self._make_response(body)
        _install_body_guard(resp)

        with pytest.raises(ValueError, match="Response body too large"):
            resp.raw.read()
        resp.close.assert_called_once()

    def test_missing_cl_within_limit(self):
        """No Content-Length, small body succeeds."""
        body = b"hello world"
        resp = self._make_response(body)
        _install_body_guard(resp)

        data = resp.raw.read()
        assert data == body

    def test_missing_cl_exceeds_limit(self):
        """No Content-Length, oversized body raises ValueError."""
        body = b"x" * (MAX_RESPONSE_SIZE + 100)
        resp = self._make_response(body)
        _install_body_guard(resp)

        with pytest.raises(ValueError, match="Response body too large"):
            resp.raw.read()
        resp.close.assert_called_once()

    def test_guard_installed_even_when_cl_present(self):
        """Guard always installs — callers decide when to call it."""
        body = b"x" * 100
        resp = self._make_response(body, headers={"Content-Length": "100"})
        original_read = resp.raw.read
        _install_body_guard(resp)

        # read function should be wrapped
        assert resp.raw.read is not original_read
        # but small body still reads fine
        data = resp.raw.read()
        assert data == body

    def test_invalid_cl_gets_body_guard(self):
        """Invalid Content-Length must still get body guard via _check_response_size."""
        body = b"x" * (MAX_RESPONSE_SIZE + 1)
        resp = self._make_response(body, headers={"Content-Length": "abc"})
        _check_response_size(resp)

        with pytest.raises(ValueError, match="Response body too large"):
            resp.raw.read()
        resp.close.assert_called_once()

    def test_negative_cl_gets_body_guard(self):
        """Negative Content-Length must still get body guard via _check_response_size."""
        body = b"x" * (MAX_RESPONSE_SIZE + 1)
        resp = self._make_response(body, headers={"Content-Length": "-1"})
        _check_response_size(resp)

        with pytest.raises(ValueError, match="Response body too large"):
            resp.raw.read()
        resp.close.assert_called_once()

    def test_streamed_access_exceeds_limit(self):
        """Chunked reads that cumulatively exceed limit raise ValueError."""
        chunk_size = MAX_RESPONSE_SIZE // 2
        body = b"x" * (MAX_RESPONSE_SIZE + 1)
        resp = self._make_response(body)
        _install_body_guard(resp)

        # First chunk: within limit
        data1 = resp.raw.read(chunk_size)
        assert len(data1) == chunk_size

        # Second chunk: pushes over the limit
        with pytest.raises(ValueError, match="Response body too large"):
            resp.raw.read(chunk_size + 1)
        resp.close.assert_called_once()