""" 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]