learningcircuit--local-deep-research/tests/web/test_app_factory_middleware.py

"""
Tests for web-layer middleware: SecureCookieMiddleware, ServerHeaderMiddleware,
DiskSpoolingRequest, and the is_private_ip helper they depend on.
"""

import ipaddress

from flask import Flask, Request

from local_deep_research.security.network_utils import is_private_ip
from local_deep_research.security.web_middleware import (
    SecureCookieMiddleware,
    ServerHeaderMiddleware,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_simple_app(**config_overrides):
    """Create a minimal Flask app with a cookie-setting route."""
    app = Flask(__name__)
    app.config["SECRET_KEY"] = "test-secret"
    app.config["LDR_TESTING_MODE"] = False
    app.config.update(config_overrides)

    @app.route("/set-cookie")
    def set_cookie():
        from flask import make_response

        resp = make_response("ok")
        resp.set_cookie("session_id", "abc123")
        return resp

    @app.route("/no-cookie")
    def no_cookie():
        return "ok"

    return app


# ===================================================================
# 1. Tests for the is_private_ip helper used by SecureCookieMiddleware
# ===================================================================


class TestIsPrivateIp:
    """Unit tests for security.network_utils.is_private_ip."""

    def test_loopback_ipv4(self):
        assert is_private_ip("127.0.0.1") is True

    def test_loopback_ipv6(self):
        assert is_private_ip("::1") is True

    def test_class_a_private(self):
        assert is_private_ip("10.0.0.1") is True

    def test_class_b_private(self):
        assert is_private_ip("172.16.0.1") is True

    def test_class_c_private(self):
        assert is_private_ip("192.168.1.1") is True

    def test_public_google_dns(self):
        assert is_private_ip("8.8.8.8") is False

    def test_public_cloudflare_dns(self):
        assert is_private_ip("1.1.1.1") is False

    def test_public_documentation_range(self):
        """203.0.113.0/24 (TEST-NET-3) is documentation/reserved, not routable.
        Python's ipaddress module treats it as private."""

        # This is actually considered private by Python's ipaddress module
        # because it is reserved. Verify against the real implementation.
        result = is_private_ip("203.0.113.1")
        ip = ipaddress.ip_address("203.0.113.1")
        assert result == (ip.is_private or ip.is_loopback)

    def test_invalid_string(self):
        assert is_private_ip("not-an-ip") is False

    def test_empty_string(self):
        assert is_private_ip("") is False

    def test_invalid_octet_values(self):
        assert is_private_ip("256.256.256.256") is False

    def test_ipv6_private(self):
        # fc00::/7 is unique local address (private)
        assert is_private_ip("fc00::1") is True

    def test_ipv6_public(self):
        # 2001:4860:4860::8888 is Google Public DNS IPv6
        assert is_private_ip("2001:4860:4860::8888") is False


# ===================================================================
# 2. Tests for DiskSpoolingRequest (module-level, directly importable)
# ===================================================================


class TestDiskSpoolingRequest:
    """Unit tests for the DiskSpoolingRequest custom Request class."""

    def test_max_form_memory_size_is_5mb(self):
        from local_deep_research.web.app_factory import DiskSpoolingRequest

        assert DiskSpoolingRequest.max_form_memory_size == 5 * 1024 * 1024

    def test_is_subclass_of_flask_request(self):
        from local_deep_research.web.app_factory import DiskSpoolingRequest

        assert issubclass(DiskSpoolingRequest, Request)

    def test_inherits_request_methods(self):
        """DiskSpoolingRequest should inherit all standard Request class attributes."""
        from local_deep_research.web.app_factory import DiskSpoolingRequest

        # Spot-check class-level attributes and methods inherited from Request
        assert hasattr(DiskSpoolingRequest, "from_values")
        assert hasattr(DiskSpoolingRequest, "application")
        assert hasattr(DiskSpoolingRequest, "max_content_length")
        assert hasattr(DiskSpoolingRequest, "max_form_memory_size")

    def test_can_be_assigned_as_request_class(self):
        """Flask app should accept DiskSpoolingRequest as its request_class."""
        from local_deep_research.web.app_factory import DiskSpoolingRequest

        app = Flask(__name__)
        app.request_class = DiskSpoolingRequest
        assert app.request_class is DiskSpoolingRequest


# ===================================================================
# 3. Tests for SecureCookieMiddleware (_should_add_secure_flag logic)
# ===================================================================


class TestSecureCookieMiddlewareShouldAddSecure:
    """Tests for SecureCookieMiddleware._should_add_secure_flag decision logic.

    The new rule is: add Secure iff wsgi.url_scheme == 'https' (and not in
    testing mode). Source IP no longer affects the decision.
    """

    def _make_middleware(self, **config):
        flask_app = Flask(__name__)
        flask_app.config["LDR_TESTING_MODE"] = False
        flask_app.config.update(config)
        return SecureCookieMiddleware(wsgi_app=None, flask_app=flask_app)

    def test_testing_mode_returns_false(self):
        mw = self._make_middleware(LDR_TESTING_MODE=True)
        environ = {"REMOTE_ADDR": "8.8.8.8", "wsgi.url_scheme": "https"}
        assert mw._should_add_secure_flag(environ) is False

    def test_http_returns_false_regardless_of_source_ip(self):
        """HTTP requests never get Secure, even from public IPs (#3849)."""
        mw = self._make_middleware()
        for remote_addr in [
            "127.0.0.1",
            "::1",
            "192.168.1.1",
            "10.0.0.1",
            "172.16.0.1",
            "172.17.0.2",  # Default Docker bridge
            "172.67.130.145",  # Docker Desktop NAT (the #3849 trigger)
            "8.8.8.8",
            "1.1.1.1",
            "",
        ]:
            environ = {
                "REMOTE_ADDR": remote_addr,
                "wsgi.url_scheme": "http",
            }
            assert mw._should_add_secure_flag(environ) is False, (
                f"HTTP from {remote_addr!r} should not get Secure"
            )

    def test_https_returns_true_regardless_of_source_ip(self):
        mw = self._make_middleware()
        for remote_addr in ["127.0.0.1", "192.168.1.1", "8.8.8.8"]:
            environ = {
                "REMOTE_ADDR": remote_addr,
                "wsgi.url_scheme": "https",
            }
            assert mw._should_add_secure_flag(environ) is True


class TestSecureCookieMiddlewareCall:
    """Tests for SecureCookieMiddleware.__call__ cookie header modification."""

    def _capture_headers(self, wrapped, environ):
        captured = []

        def start_response(status, headers, exc_info=None):
            captured.extend(headers)

        list(wrapped(environ, start_response))
        return captured

    def test_no_secure_appended_for_public_ip_http(self):
        """Public IP over HTTP must NOT get Secure - core fix for #3849."""
        app = _make_simple_app()
        wrapped = SecureCookieMiddleware(app.wsgi_app, app)
        app.wsgi_app = wrapped

        with app.test_request_context():
            captured = self._capture_headers(
                wrapped,
                {
                    "REQUEST_METHOD": "GET",
                    "PATH_INFO": "/set-cookie",
                    "SERVER_NAME": "localhost",
                    "SERVER_PORT": "5000",
                    "REMOTE_ADDR": "8.8.8.8",
                    "wsgi.url_scheme": "http",
                    "wsgi.input": b"",
                },
            )
            cookies = [v for n, v in captured if n.lower() == "set-cookie"]
            for cookie in cookies:
                assert "Secure" not in cookie

    def test_no_secure_appended_for_private_ip_http(self):
        """Private IP HTTP requests are unmodified."""
        app = _make_simple_app()
        wrapped = SecureCookieMiddleware(app.wsgi_app, app)
        app.wsgi_app = wrapped

        with app.test_request_context():
            captured = self._capture_headers(
                wrapped,
                {
                    "REQUEST_METHOD": "GET",
                    "PATH_INFO": "/set-cookie",
                    "SERVER_NAME": "localhost",
                    "SERVER_PORT": "5000",
                    "REMOTE_ADDR": "127.0.0.1",
                    "wsgi.url_scheme": "http",
                    "wsgi.input": b"",
                },
            )
            cookies = [v for n, v in captured if n.lower() == "set-cookie"]
            for cookie in cookies:
                assert "Secure" not in cookie

    def test_secure_appended_for_https(self):
        """HTTPS requests get Secure flag added."""
        app = _make_simple_app()
        wrapped = SecureCookieMiddleware(app.wsgi_app, app)
        app.wsgi_app = wrapped

        with app.test_request_context():
            captured = self._capture_headers(
                wrapped,
                {
                    "REQUEST_METHOD": "GET",
                    "PATH_INFO": "/set-cookie",
                    "SERVER_NAME": "localhost",
                    "SERVER_PORT": "5000",
                    "REMOTE_ADDR": "127.0.0.1",
                    "wsgi.url_scheme": "https",
                    "wsgi.input": b"",
                },
            )
            cookies = [v for n, v in captured if n.lower() == "set-cookie"]
            assert cookies, "expected at least one Set-Cookie header"
            for cookie in cookies:
                assert "Secure" in cookie

    def test_no_duplicate_secure_flag(self):
        """If cookie already has '; Secure', don't add it again."""

        def inner_app(environ, start_response):
            headers = [
                ("Content-Type", "text/plain"),
                ("Set-Cookie", "token=xyz; HttpOnly; Secure"),
            ]
            start_response("200 OK", headers)
            return [b"ok"]

        flask_app = Flask(__name__)
        flask_app.config["LDR_TESTING_MODE"] = False
        wrapped = SecureCookieMiddleware(inner_app, flask_app)

        captured = self._capture_headers(
            wrapped,
            {"REMOTE_ADDR": "8.8.8.8", "wsgi.url_scheme": "https"},
        )
        cookies = [v for n, v in captured if n.lower() == "set-cookie"]
        assert len(cookies) == 1
        assert cookies[0].count("; Secure") == 1

    def test_non_cookie_headers_unchanged(self):
        """Non-Set-Cookie headers pass through unmodified."""

        def inner_app(environ, start_response):
            headers = [
                ("Content-Type", "text/html"),
                ("X-Custom", "value123"),
                ("Set-Cookie", "foo=bar"),
            ]
            start_response("200 OK", headers)
            return [b"ok"]

        flask_app = Flask(__name__)
        flask_app.config["LDR_TESTING_MODE"] = False
        wrapped = SecureCookieMiddleware(inner_app, flask_app)

        captured = self._capture_headers(
            wrapped,
            {"REMOTE_ADDR": "8.8.8.8", "wsgi.url_scheme": "http"},
        )
        content_type = [v for n, v in captured if n == "Content-Type"]
        assert content_type == ["text/html"]
        x_custom = [v for n, v in captured if n == "X-Custom"]
        assert x_custom == ["value123"]

    def test_testing_mode_leaves_cookies_alone(self):
        """When LDR_TESTING_MODE is True, cookies are never modified."""

        def inner_app(environ, start_response):
            headers = [("Set-Cookie", "test=value")]
            start_response("200 OK", headers)
            return [b"ok"]

        flask_app = Flask(__name__)
        flask_app.config["LDR_TESTING_MODE"] = True
        wrapped = SecureCookieMiddleware(inner_app, flask_app)

        captured = self._capture_headers(
            wrapped,
            {"REMOTE_ADDR": "8.8.8.8", "wsgi.url_scheme": "https"},
        )
        cookies = [v for n, v in captured if n.lower() == "set-cookie"]
        assert cookies == ["test=value"]
        assert "Secure" not in cookies[0]


# ===================================================================
# 4. Tests for ServerHeaderMiddleware
# ===================================================================


class TestServerHeaderMiddleware:
    """Tests for ServerHeaderMiddleware WSGI middleware."""

    def _make_inner_app(self, headers):
        """Create a simple WSGI app that returns the given headers."""

        def inner_app(environ, start_response):
            start_response("200 OK", list(headers))
            return [b"ok"]

        return inner_app

    def test_removes_server_header_title_case(self):
        inner = self._make_inner_app(
            [("Content-Type", "text/plain"), ("Server", "Werkzeug/2.3.0")]
        )
        wrapped = ServerHeaderMiddleware(inner)

        captured_headers = []

        def mock_start_response(status, headers, exc_info=None):
            captured_headers.extend(headers)

        list(wrapped({}, mock_start_response))

        header_names = [n for n, v in captured_headers]
        assert "Server" not in header_names
        assert "Content-Type" in header_names

    def test_removes_server_header_lowercase(self):
        inner = self._make_inner_app(
            [("content-type", "text/plain"), ("server", "nginx")]
        )
        wrapped = ServerHeaderMiddleware(inner)

        captured_headers = []

        def mock_start_response(status, headers, exc_info=None):
            captured_headers.extend(headers)

        list(wrapped({}, mock_start_response))

        header_names_lower = [n.lower() for n, v in captured_headers]
        assert "server" not in header_names_lower

    def test_removes_server_header_uppercase(self):
        inner = self._make_inner_app(
            [("Content-Type", "text/plain"), ("SERVER", "Apache")]
        )
        wrapped = ServerHeaderMiddleware(inner)

        captured_headers = []

        def mock_start_response(status, headers, exc_info=None):
            captured_headers.extend(headers)

        list(wrapped({}, mock_start_response))

        header_names_lower = [n.lower() for n, v in captured_headers]
        assert "server" not in header_names_lower

    def test_passes_through_other_headers(self):
        inner = self._make_inner_app(
            [
                ("Content-Type", "text/html"),
                ("X-Custom-Header", "foobar"),
                ("Set-Cookie", "id=123"),
                ("Server", "should-be-removed"),
            ]
        )
        wrapped = ServerHeaderMiddleware(inner)

        captured_headers = []

        def mock_start_response(status, headers, exc_info=None):
            captured_headers.extend(headers)

        list(wrapped({}, mock_start_response))

        header_dict = dict(captured_headers)
        assert header_dict["Content-Type"] == "text/html"
        assert header_dict["X-Custom-Header"] == "foobar"
        assert header_dict["Set-Cookie"] == "id=123"
        assert "Server" not in header_dict

    def test_no_server_header_present_is_noop(self):
        """When there is no Server header, all headers pass through."""
        original_headers = [
            ("Content-Type", "text/plain"),
            ("X-Request-Id", "abc"),
        ]
        inner = self._make_inner_app(original_headers)
        wrapped = ServerHeaderMiddleware(inner)

        captured_headers = []

        def mock_start_response(status, headers, exc_info=None):
            captured_headers.extend(headers)

        list(wrapped({}, mock_start_response))

        assert captured_headers == original_headers

    def test_empty_headers(self):
        """Middleware handles empty header list gracefully."""
        inner = self._make_inner_app([])
        wrapped = ServerHeaderMiddleware(inner)

        captured_headers = []

        def mock_start_response(status, headers, exc_info=None):
            captured_headers.extend(headers)

        list(wrapped({}, mock_start_response))

        assert captured_headers == []

    def test_status_and_body_pass_through(self):
        """Status code and response body are unmodified."""

        def inner_app(environ, start_response):
            start_response("404 Not Found", [("Server", "x")])
            return [b"not found"]

        wrapped = ServerHeaderMiddleware(inner_app)

        captured_status = []

        def mock_start_response(status, headers, exc_info=None):
            captured_status.append(status)

        body_parts = list(wrapped({}, mock_start_response))

        assert captured_status == ["404 Not Found"]
        assert body_parts == [b"not found"]


# ===================================================================
# 5. Integration: both middlewares composed together
# ===================================================================


class TestMiddlewareComposition:
    """Test that SecureCookieMiddleware and ServerHeaderMiddleware
    compose correctly when stacked (as they are in production)."""

    def _run_stack(self, environ, set_cookies, with_server_header=True):
        """Run inner -> SecureCookie -> ProxyFix -> ServerHeader stack."""
        from werkzeug.middleware.proxy_fix import ProxyFix

        def inner_app(environ, start_response):
            headers = [("Content-Type", "text/plain")] + [
                ("Set-Cookie", c) for c in set_cookies
            ]
            if with_server_header:
                headers.append(("Server", "Werkzeug/2.3.0"))
            start_response("200 OK", headers)
            return [b"ok"]

        flask_app = Flask(__name__)
        flask_app.config["LDR_TESTING_MODE"] = False

        # Match create_app's wrap order:
        #   inner -> SecureCookie (innermost) -> ProxyFix -> ServerHeader (outer)
        wsgi = SecureCookieMiddleware(inner_app, flask_app)
        wsgi = ProxyFix(wsgi, x_for=1, x_proto=1)
        wsgi = ServerHeaderMiddleware(wsgi)

        captured = []

        def start_response(status, headers, exc_info=None):
            captured.extend(headers)

        list(wsgi(environ, start_response))
        return captured

    def test_stacked_http_public_ip_no_secure(self):
        """Public IP over HTTP: no Secure flag, Server header stripped."""
        captured = self._run_stack(
            {"REMOTE_ADDR": "8.8.8.8", "wsgi.url_scheme": "http"},
            set_cookies=["sid=abc123"],
        )
        assert "Server" not in [n for n, _ in captured]
        cookies = [v for n, v in captured if n.lower() == "set-cookie"]
        assert len(cookies) == 1
        assert "Secure" not in cookies[0]
        assert ("Content-Type", "text/plain") in captured

    def test_stacked_https_adds_secure(self):
        """HTTPS request gets Secure flag added."""
        captured = self._run_stack(
            {"REMOTE_ADDR": "8.8.8.8", "wsgi.url_scheme": "https"},
            set_cookies=["sid=abc123"],
        )
        cookies = [v for n, v in captured if n.lower() == "set-cookie"]
        assert len(cookies) == 1
        assert "Secure" in cookies[0]

    def test_stacked_private_ip_no_secure(self):
        """Private IP over HTTP: no Secure flag, Server header stripped."""
        captured = self._run_stack(
            {"REMOTE_ADDR": "192.168.1.100", "wsgi.url_scheme": "http"},
            set_cookies=["sid=abc123"],
            with_server_header=True,
        )
        assert "Server" not in [n for n, _ in captured]
        cookies = [v for n, v in captured if n.lower() == "set-cookie"]
        assert len(cookies) == 1
        assert "Secure" not in cookies[0]

    def test_stacked_proxyfix_translates_xfp_https(self):
        """ProxyFix should rewrite scheme from X-Forwarded-Proto, after
        which SecureCookieMiddleware adds Secure even though the raw
        wsgi.url_scheme is http."""
        captured = self._run_stack(
            {
                "REMOTE_ADDR": "10.0.0.1",  # trusted proxy
                "wsgi.url_scheme": "http",
                "HTTP_X_FORWARDED_PROTO": "https",
                "HTTP_X_FORWARDED_FOR": "203.0.113.50",
            },
            set_cookies=["sid=abc123"],
        )
        cookies = [v for n, v in captured if n.lower() == "set-cookie"]
        assert len(cookies) == 1
        assert "Secure" in cookies[0]