From b90ecda752fa1d8ffa363ba26c08e1f2a5094901 Mon Sep 17 00:00:00 2001
From: Vlad Ciobanu <95963142+vl3c@users.noreply.github.com>
Date: Sat, 11 Jul 2026 22:37:07 +0300
Subject: [PATCH] Fix chat XSS and harden login, request size, and secret key

XSS (client):
- markdown_parser.py now HTML-escapes source text before any markdown
  transformation, so raw <script>/<img onerror> in AI output can never
  reach the innerHTML sinks; only parser-emitted tags survive
- Added URL scheme allowlist (http/https/mailto + relative) for links;
  javascript:/data:/vbscript: and control-char tricks are neutralized to
  plain text
- chat_ui_manager.py escapes sender/content in the error-fallback sink

Server hardening:
- Login rate limit keys on request.remote_addr instead of the spoofable
  X-Forwarded-For; added a global failed-attempt ceiling (50 / 5 min)
  that blocks brute force from distributed/spoofed sources
- Set MAX_CONTENT_LENGTH (80 MB) with a JSON 413 handler; enforce
  server-side image caps (max 5, 20 MB base64 each) on both send routes
- Warn when SECRET_KEY is unset in a deployed environment

Tests:
- 13 markdown XSS regression tests; new test_login_and_hardening.py
  covering login, per-IP + global rate limiting, XFF spoof resistance,
  and image validation
- Server: 1293 passed; client: 2733 passed; mypy and ruff clean
---
 server_tests/test_login_and_hardening.py | 192 +++++++++++++++++++++++
 server_tests/test_markdown_parser.py     |  76 ++++++++-
 static/app_manager.py                    |  30 +++-
 static/client/chat_ui_manager.py         |   8 +-
 static/client/markdown_parser.py         |  61 ++++++-
 static/config.py                         |  19 +++
 static/routes.py                         | 144 +++++++++++++----
 7 files changed, 487 insertions(+), 43 deletions(-)
 create mode 100644 server_tests/test_login_and_hardening.py

diff --git a/server_tests/test_login_and_hardening.py b/server_tests/test_login_and_hardening.py
new file mode 100644
index 00000000..27c2a5ba
--- /dev/null
+++ b/server_tests/test_login_and_hardening.py
@@ -0,0 +1,192 @@
+"""Tests for login rate limiting, request-size limits, and image validation.
+
+Covers the server-side security hardening added to the authentication flow and
+the AI message endpoints:
+
+    - Correct/incorrect PIN handling with constant-time comparison.
+    - Per-IP failed-attempt cooldown keyed on the real socket peer (X-Forwarded-For
+      spoofing must not reset or bypass the limit).
+    - A global failed-attempt ceiling that locks out even fresh IPs.
+    - Server-side attached-image validation (count and per-image size caps).
+    - MAX_CONTENT_LENGTH configuration and the JSON 413 error handler.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import unittest
+from typing import List, Optional
+from unittest.mock import Mock, patch
+
+from werkzeug.test import TestResponse
+
+from static.app_manager import AppManager, MatHudFlask
+from static.config import MAX_ATTACHED_IMAGES, MAX_CONTENT_LENGTH_BYTES, MAX_IMAGE_BASE64_BYTES
+from static.routes import GLOBAL_FAILED_ATTEMPTS_LIMIT, reset_login_rate_limit_state
+
+
+TEST_PIN = "123456"
+
+
+class TestLoginRateLimiting(unittest.TestCase):
+    """Exercise the /login handler's authentication and rate-limiting logic."""
+
+    def setUp(self) -> None:
+        self._original_require_auth: Optional[str] = os.environ.get("REQUIRE_AUTH")
+        self._original_auth_pin: Optional[str] = os.environ.get("AUTH_PIN")
+        os.environ["REQUIRE_AUTH"] = "true"
+        os.environ["AUTH_PIN"] = TEST_PIN
+
+        self.app: MatHudFlask = AppManager.create_app()
+        self.app.config["TESTING"] = True
+        self.client = self.app.test_client()
+
+        # Ensure no leaked rate-limit state from other tests.
+        reset_login_rate_limit_state()
+
+    def tearDown(self) -> None:
+        reset_login_rate_limit_state()
+        self._restore_env("REQUIRE_AUTH", self._original_require_auth)
+        self._restore_env("AUTH_PIN", self._original_auth_pin)
+
+    @staticmethod
+    def _restore_env(name: str, value: Optional[str]) -> None:
+        if value is not None:
+            os.environ[name] = value
+        else:
+            os.environ.pop(name, None)
+
+    def _post_login(self, pin: str, remote_addr: str = "127.0.0.1", xff: Optional[str] = None) -> TestResponse:
+        headers = {"X-Forwarded-For": xff} if xff is not None else None
+        return self.client.post(
+            "/login",
+            data={"pin": pin},
+            headers=headers,
+            environ_base={"REMOTE_ADDR": remote_addr},
+        )
+
+    def test_correct_pin_succeeds(self) -> None:
+        response = self._post_login(TEST_PIN)
+        self.assertEqual(response.status_code, 302)  # redirect to index
+        with self.client.session_transaction() as sess:
+            self.assertTrue(sess.get("authenticated"))
+
+    def test_wrong_pin_fails(self) -> None:
+        response = self._post_login("000000")
+        self.assertEqual(response.status_code, 200)  # re-render login, not a redirect
+        with self.client.session_transaction() as sess:
+            self.assertFalse(sess.get("authenticated"))
+
+    def test_per_ip_cooldown_triggers_after_failure(self) -> None:
+        # First failure records a timestamp and re-renders the login page.
+        first = self._post_login("000000")
+        self.assertEqual(first.status_code, 200)
+
+        # Second failure within the cooldown window is rejected with 429.
+        second = self._post_login("000000")
+        self.assertEqual(second.status_code, 429)
+
+    def test_x_forwarded_for_does_not_bypass_per_ip_limit(self) -> None:
+        # Same socket peer (REMOTE_ADDR), different spoofed XFF values on each request.
+        self.assertEqual(self._post_login("000000", xff="10.0.0.1").status_code, 200)
+
+        # Varying the client-controlled header must NOT reset the per-IP cooldown,
+        # because limiting keys on REMOTE_ADDR, not X-Forwarded-For.
+        self.assertEqual(self._post_login("000000", xff="10.0.0.2").status_code, 429)
+        self.assertEqual(self._post_login("000000", xff="203.0.113.9").status_code, 429)
+
+    def test_global_ceiling_locks_out_fresh_ip(self) -> None:
+        # Drive the global counter to its ceiling using a distinct IP per request so
+        # the per-IP cooldown never engages.
+        for i in range(GLOBAL_FAILED_ATTEMPTS_LIMIT):
+            response = self._post_login("000000", remote_addr=f"198.51.100.{i}")
+            self.assertEqual(response.status_code, 200)
+
+        # A brand-new, never-seen IP is now locked out purely by the global ceiling.
+        locked = self._post_login("000000", remote_addr="192.0.2.55")
+        self.assertEqual(locked.status_code, 429)
+
+        # Even a correct PIN is rejected while the global lockout is active.
+        locked_correct = self._post_login(TEST_PIN, remote_addr="192.0.2.77")
+        self.assertEqual(locked_correct.status_code, 429)
+
+
+class TestImageAndSizeHardening(unittest.TestCase):
+    """Exercise server-side image validation and request-size configuration."""
+
+    SAMPLE_PNG_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/w8AAwMB/aqVw0sAAAAASUVORK5CYII="
+
+    def setUp(self) -> None:
+        self._original_require_auth: Optional[str] = os.environ.get("REQUIRE_AUTH")
+        os.environ["REQUIRE_AUTH"] = "false"
+
+        self.app = AppManager.create_app()
+        self.app.config["TESTING"] = True
+        self.client = self.app.test_client()
+
+    def tearDown(self) -> None:
+        if self._original_require_auth is not None:
+            os.environ["REQUIRE_AUTH"] = self._original_require_auth
+        else:
+            os.environ.pop("REQUIRE_AUTH", None)
+
+    def _small_image(self) -> str:
+        return f"data:image/png;base64,{self.SAMPLE_PNG_BASE64}"
+
+    def _post_message(self, images: List[str]) -> TestResponse:
+        payload = {"message": json.dumps({"user_message": "hi", "attached_images": images})}
+        return self.client.post("/send_message", json=payload)
+
+    def test_too_many_images_rejected(self) -> None:
+        images = [self._small_image()] * (MAX_ATTACHED_IMAGES + 1)
+        response = self._post_message(images)
+        self.assertEqual(response.status_code, 400)
+        data = json.loads(response.data)
+        self.assertEqual(data["status"], "error")
+        self.assertIn("Too many attached images", data["message"])
+
+    def test_oversized_single_image_rejected(self) -> None:
+        oversized = "data:image/png;base64," + ("A" * (MAX_IMAGE_BASE64_BYTES + 1))
+        response = self._post_message([oversized])
+        self.assertEqual(response.status_code, 400)
+        data = json.loads(response.data)
+        self.assertEqual(data["status"], "error")
+        self.assertIn("maximum size", data["message"])
+
+    @patch("static.openai_completions_api.OpenAIChatCompletionsAPI.create_chat_completion")
+    def test_valid_image_payload_passes_validation(self, mock_chat: Mock) -> None:
+        class MockMessage:
+            content = "ok"
+            tool_calls = None
+
+        class MockResponse:
+            message = MockMessage()
+            finish_reason = "stop"
+
+        mock_chat.return_value = MockResponse()
+
+        response = self._post_message([self._small_image()] * MAX_ATTACHED_IMAGES)
+        self.assertEqual(response.status_code, 200)
+        data = json.loads(response.data)
+        self.assertEqual(data["status"], "success")
+
+    def test_max_content_length_configured(self) -> None:
+        self.assertEqual(self.app.config["MAX_CONTENT_LENGTH"], MAX_CONTENT_LENGTH_BYTES)
+
+    def test_413_handler_returns_json(self) -> None:
+        # Shrink the limit for this request so we don't have to send 80 MB.
+        self.app.config["MAX_CONTENT_LENGTH"] = 64
+        response = self.client.post(
+            "/send_message",
+            data="x" * 512,
+            content_type="application/json",
+        )
+        self.assertEqual(response.status_code, 413)
+        self.assertIn("application/json", response.content_type)
+        data = json.loads(response.data)
+        self.assertIn("error", data)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/server_tests/test_markdown_parser.py b/server_tests/test_markdown_parser.py
index cd428f6e..ec7b5262 100644
--- a/server_tests/test_markdown_parser.py
+++ b/server_tests/test_markdown_parser.py
@@ -89,10 +89,84 @@ def hello():
         self.assertIn("</code></pre>", result)
 
     def test_links(self) -> None:
-        """Test link formatting."""
+        """Test link formatting for a safe (allowed-scheme) URL."""
         result = self.parser.parse("This is a [link](https://example.com) text")
         self.assertIn('<a href="https://example.com">link</a>', result)
 
+    def test_raw_script_tag_is_escaped(self) -> None:
+        """Raw <script> in AI output must be escaped, not executed."""
+        result = self.parser.parse("<script>alert(1)</script>")
+        self.assertNotIn("<script>", result)
+        self.assertIn("&lt;script&gt;", result)
+        self.assertIn("&lt;/script&gt;", result)
+
+    def test_img_onerror_is_escaped(self) -> None:
+        """An <img onerror=...> payload must be escaped, not rendered."""
+        result = self.parser.parse("<img src=x onerror=alert(1)>")
+        self.assertNotIn("<img", result)
+        self.assertIn("&lt;img", result)
+
+    def test_javascript_link_is_neutralized(self) -> None:
+        """javascript: links must not produce a live href."""
+        result = self.parser.parse("[x](javascript:alert(1))")
+        self.assertNotIn("href=\"javascript", result)
+        self.assertNotIn("javascript:alert", result)
+        # Link text is preserved as plain text.
+        self.assertIn("x", result)
+
+    def test_javascript_link_mixed_case_is_neutralized(self) -> None:
+        """Mixed-case JaVaScRiPt: scheme must be rejected case-insensitively."""
+        result = self.parser.parse("[x](JaVaScRiPt:alert(1))")
+        self.assertNotIn("<a href", result)
+        self.assertNotIn("alert(1)", result)
+
+    def test_javascript_link_with_control_char_is_neutralized(self) -> None:
+        """Embedded control chars (java\\tscript:) must not bypass sanitization."""
+        result = self.parser.parse("[x](java\tscript:alert(1))")
+        self.assertNotIn("<a href", result)
+        self.assertNotIn("alert(1)", result)
+
+    def test_data_uri_link_is_neutralized(self) -> None:
+        """data: URIs must be rejected."""
+        result = self.parser.parse("[x](data:text/html,<script>alert(1)</script>)")
+        self.assertNotIn("<a href", result)
+
+    def test_safe_https_link_still_renders(self) -> None:
+        """Allowed https links must still render as real anchors."""
+        result = self.parser.parse("[ok](https://example.com)")
+        self.assertIn('<a href="https://example.com">ok</a>', result)
+
+    def test_mailto_link_still_renders(self) -> None:
+        """mailto: links are allowed."""
+        result = self.parser.parse("[mail](mailto:a@b.com)")
+        self.assertIn('<a href="mailto:a@b.com">mail</a>', result)
+
+    def test_relative_link_still_renders(self) -> None:
+        """Relative URLs (no scheme) are allowed."""
+        result = self.parser.parse("[rel](/path/page)")
+        self.assertIn('<a href="/path/page">rel</a>', result)
+
+    def test_normal_markdown_still_renders_after_escaping(self) -> None:
+        """Bold, code, and headers must still render once escaping is in place."""
+        result = self.parser.parse("# Title\n\nSome **bold** and `code` here")
+        self.assertIn("<h1>Title</h1>", result)
+        self.assertIn("<strong>bold</strong>", result)
+        self.assertIn("<code>code</code>", result)
+
+    def test_angle_brackets_in_prose_are_escaped(self) -> None:
+        """Comparison operators in prose render as escaped entities."""
+        result = self.parser.parse("a < b and c > d")
+        self.assertIn("a &lt; b and c &gt; d", result)
+        self.assertNotIn("<b", result)
+
+    def test_code_block_with_script_is_escaped(self) -> None:
+        """A code block containing <script> must be escaped inside <pre><code>."""
+        code_block = "```\n<script>alert(1)</script>\n```"
+        result = self.parser.parse(code_block)
+        self.assertIn("<pre><code>", result)
+        self.assertIn("&lt;script&gt;", result)
+        self.assertNotIn("<script>", result)
+
     def test_unordered_lists(self) -> None:
         """Test unordered list formatting."""
         markdown = """- Item 1
diff --git a/static/app_manager.py b/static/app_manager.py
index f5880b13..5dd37a93 100644
--- a/static/app_manager.py
+++ b/static/app_manager.py
@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import logging
 import os
 import secrets
 from typing import TYPE_CHECKING, Dict, Optional, Tuple, TypedDict, Union
@@ -22,6 +23,7 @@
 from flask import Flask, Response, jsonify
 from flask_session import Session as FlaskSession
 
+from static.config import MAX_CONTENT_LENGTH_BYTES
 from static.env_config import load_env_files
 from static.log_manager import LogManager
 from static.openai_completions_api import OpenAIChatCompletionsAPI
@@ -35,6 +37,9 @@
     from static.webdriver_manager import WebDriverManager
 
 
+_logger = logging.getLogger(__name__)
+
+
 JsonValue = Union[str, int, float, bool, None, Dict[str, "JsonValue"], list["JsonValue"]]
 
 
@@ -156,8 +161,29 @@ def create_app() -> MatHudFlask:
         # Load environment variables from project .env and parent .env (API keys)
         AppManager._load_env()
 
-        # Configure session management for authentication using modern CacheLib backend
-        app.secret_key = os.getenv("SECRET_KEY", secrets.token_hex(32))
+        # Configure session management for authentication using modern CacheLib backend.
+        # A persistent SECRET_KEY is required for sessions to survive process restarts.
+        # In deployed mode an ephemeral fallback silently logs everyone out on every
+        # restart, so warn prominently (but do not hard-fail).
+        secret_key_env = os.getenv("SECRET_KEY")
+        if secret_key_env:
+            app.secret_key = secret_key_env
+        else:
+            app.secret_key = secrets.token_hex(32)
+            if AppManager.is_deployed():
+                _logger.warning(
+                    "SECRET_KEY is not set in a deployed environment. A random key was "
+                    "generated for this process, so ALL user sessions will be invalidated "
+                    "on every restart. Set SECRET_KEY to a stable secret to persist sessions."
+                )
+
+        # Cap total request body size and return clean JSON on overflow so clients
+        # do not receive Flask's default HTML 413 page.
+        app.config["MAX_CONTENT_LENGTH"] = MAX_CONTENT_LENGTH_BYTES
+
+        @app.errorhandler(413)
+        def _handle_request_entity_too_large(_error: Exception) -> Tuple[Response, int]:
+            return jsonify({"error": "Request payload too large"}), 413
 
         # Create session directory if it doesn't exist
         session_dir = os.path.join(os.getcwd(), "flask_session")
diff --git a/static/client/chat_ui_manager.py b/static/client/chat_ui_manager.py
index 5fb26ed4..8b7ef860 100644
--- a/static/client/chat_ui_manager.py
+++ b/static/client/chat_ui_manager.py
@@ -187,10 +187,12 @@ def handler(event: Any) -> None:
 
         except Exception as e:
             print(f"Error creating message element: {e}")
-            # Fall back to simple paragraph
+            # Fall back to simple paragraph. Escape sender/content before
+            # interpolating into innerHTML so raw HTML never executes here.
             if sender == "AI":
-                content = message.replace("\n", "<br>")
-                return html.P(f"<strong>{sender}:</strong> {content}", innerHTML=True)
+                content = self._escape_html(message).replace("\n", "<br>")
+                safe_sender = self._escape_html(sender)
+                return html.P(f"<strong>{safe_sender}:</strong> {content}", innerHTML=True)
             else:
                 return html.P(f"<strong>{sender}:</strong> {message}")
 
diff --git a/static/client/markdown_parser.py b/static/client/markdown_parser.py
index d5ab5891..a7a4fcf8 100644
--- a/static/client/markdown_parser.py
+++ b/static/client/markdown_parser.py
@@ -25,6 +25,10 @@
 class MarkdownParser:
     """Custom markdown parser optimized for chat interface display."""
 
+    # URL schemes considered safe for link ``href`` attributes. Anything else
+    # (javascript:, data:, vbscript:, ...) is rejected by ``_sanitize_url``.
+    _ALLOWED_URL_SCHEMES = ("http", "https", "mailto")
+
     def parse(self, text: str) -> str:
         """Parse markdown text to HTML."""
         try:
@@ -35,12 +39,51 @@ def parse(self, text: str) -> str:
 
         except Exception as e:
             print(f"Error in custom markdown parsing: {e}")
-            # Ultimate fallback
-            return text.replace("\n", "<br>")
+            # Ultimate fallback — escape so raw HTML never reaches innerHTML.
+            return self._escape_html(text).replace("\n", "<br>")
+
+    def _escape_html(self, text: str) -> str:
+        """Escape HTML special characters before any markdown transformation.
+
+        ``&`` is escaped first so later replacements do not double-escape, then
+        ``<``/``>`` neutralise raw tags, and the quotes protect attribute
+        contexts (e.g. sanitized link URLs interpolated into ``href``). This
+        runs up-front so the only tags in the output are ones this parser emits.
+        """
+        return (
+            text.replace("&", "&amp;")
+            .replace("<", "&lt;")
+            .replace(">", "&gt;")
+            .replace('"', "&quot;")
+            .replace("'", "&#x27;")
+        )
+
+    def _sanitize_url(self, url: str) -> Optional[str]:
+        """Return ``url`` if safe for an ``href``, otherwise ``None``.
+
+        Allows http/https/mailto plus scheme-relative (``//host``) and relative
+        URLs. Rejects javascript:, data:, vbscript: and any other explicit
+        scheme. Whitespace and control characters are stripped before the
+        scheme check so tricks like ``java\tscript:`` cannot slip through.
+        """
+        cleaned = re.sub(r"[\x00-\x20\x7f]", "", url)
+        scheme_match = re.match(r"^([a-zA-Z][a-zA-Z0-9+.\-]*):", cleaned)
+        if scheme_match:
+            scheme = scheme_match.group(1).lower()
+            if scheme not in self._ALLOWED_URL_SCHEMES:
+                return None
+        return url
 
     def _simple_markdown_parse(self, text: str) -> str:
         """Simple markdown parser for basic formatting using string operations."""
         try:
+            # Escape HTML entities up-front so raw tags in the source (e.g.
+            # <script> or <img onerror=...>) can never become live HTML at the
+            # innerHTML sink. Structural markdown markers (#, -, *, |, digits,
+            # backticks) are ASCII that escaping leaves intact; the only
+            # collision is the blockquote ">", handled as "&gt;" in the loop.
+            text = self._escape_html(text)
+
             # First handle tables
             text = self._process_tables(text)
 
@@ -83,9 +126,9 @@ def _simple_markdown_parse(self, text: str) -> str:
                 # Lists - handle ordered and unordered with indentation
                 elif self._is_list_item(processed_line):
                     processed_line = self._process_list_item(processed_line)
-                # Blockquotes
-                elif processed_line.startswith("> "):
-                    processed_line = f"<blockquote>{processed_line[2:]}</blockquote>"
+                # Blockquotes (the ">" marker is escaped to "&gt;" up-front)
+                elif processed_line.startswith("&gt; "):
+                    processed_line = f"<blockquote>{processed_line[len('&gt; '):]}</blockquote>"
                 # Horizontal rules
                 elif processed_line.strip() == "---":
                     processed_line = "<hr>"
@@ -648,7 +691,13 @@ def _process_inline_markdown(self, text: str) -> str:
                 link_text = text[start + 1 : middle]
                 link_url = text[middle + 2 : end]
                 after = text[end + 1 :]
-                text = before + f'<a href="{link_url}">{link_text}</a>' + after
+                safe_url = self._sanitize_url(link_url)
+                if safe_url is None:
+                    # Unsafe scheme (e.g. javascript:) — drop the link, keep text.
+                    replacement = link_text
+                else:
+                    replacement = f'<a href="{safe_url}">{link_text}</a>'
+                text = before + replacement + after
 
             return text
 
diff --git a/static/config.py b/static/config.py
index ec920e25..0b9d3027 100644
--- a/static/config.py
+++ b/static/config.py
@@ -25,3 +25,22 @@
 # Paths for Selenium-captured canvas screenshots
 CANVAS_SNAPSHOT_DIR: str = "canvas_snapshots"
 CANVAS_SNAPSHOT_PATH: str = os.path.join(CANVAS_SNAPSHOT_DIR, "canvas.png")
+
+# ===== REQUEST SIZE & IMAGE VALIDATION CONSTANTS =====
+# Server-side hard caps enforced independently of the (advisory) client-side
+# limits in static/client/constants.py. These defend the backend against
+# oversized or abusive uploads regardless of what the client claims to enforce.
+
+# Maximum size of an entire request body (bytes). Sized to comfortably hold the
+# legitimate worst case: MAX_ATTACHED_IMAGES base64 data URLs (~13.4 MB each for
+# a 10 MB binary image) plus message text and JSON overhead. 80 MB leaves ample
+# headroom while still bounding memory use per request.
+MAX_CONTENT_LENGTH_BYTES: int = 80 * 1024 * 1024
+
+# Maximum number of attached images accepted per request.
+MAX_ATTACHED_IMAGES: int = 5
+
+# Per-image base64 payload cap (bytes). ~20 MB base64 ≈ 15 MB of binary image
+# data, comfortably above the client's 10 MB warning threshold so legitimate
+# uploads pass, while still bounding a single image's footprint.
+MAX_IMAGE_BASE64_BYTES: int = 20 * 1024 * 1024
diff --git a/static/routes.py b/static/routes.py
index 8413bb52..d328d66a 100644
--- a/static/routes.py
+++ b/static/routes.py
@@ -21,7 +21,7 @@
 import os
 import time
 from collections.abc import Callable, Iterator, Set as AbstractSet
-from typing import Any, Dict, List, Optional, TypeVar, Union, cast
+from typing import Any, Dict, List, Optional, Tuple, TypeVar, Union, cast
 
 from flask import Response, flash, jsonify, redirect, render_template, request, session, stream_with_context, url_for
 from flask.typing import ResponseReturnValue
@@ -29,7 +29,12 @@
 from static.ai_model import AIModel, PROVIDER_OPENAI, PROVIDER_ANTHROPIC, PROVIDER_OPENROUTER, PROVIDER_OLLAMA
 from static.app_manager import AppManager, MatHudFlask
 from static.canvas_state_summarizer import compare_canvas_states
-from static.config import CANVAS_SNAPSHOT_DIR, CANVAS_SNAPSHOT_PATH
+from static.config import (
+    CANVAS_SNAPSHOT_DIR,
+    CANVAS_SNAPSHOT_PATH,
+    MAX_ATTACHED_IMAGES,
+    MAX_IMAGE_BASE64_BYTES,
+)
 from static.openai_api_base import OpenAIAPIBase
 from static.providers import ProviderRegistry, create_provider_instance
 from static.route_helpers import get_active_provider, reset_tools_for_all_providers
@@ -41,16 +46,82 @@
 
 _logger = logging.getLogger(__name__)
 
-# Global dictionary to track login attempts by IP address
-# Format: {ip_address: last_attempt_timestamp}
 JsonValue = Union[str, int, float, bool, None, Dict[str, "JsonValue"], List["JsonValue"]]
 JsonObject = Dict[str, JsonValue]
 StreamEventDict = Dict[str, JsonValue]
 
+# ===== LOGIN RATE-LIMIT CONSTANTS =====
+# Per-IP cooldown between consecutive failed login attempts (seconds).
+LOGIN_ATTEMPT_COOLDOWN_SECONDS: float = 5.0
+# Prune the per-IP tracking dict once it grows past this many entries.
+LOGIN_ATTEMPTS_CLEANUP_THRESHOLD: int = 1000
+# Entries older than this (seconds) are removed during cleanup.
+LOGIN_ATTEMPTS_CLEANUP_AGE_SECONDS: float = 3600.0
+# Global failed-attempt ceiling, independent of source IP. Because per-IP limiting
+# keys on request.remote_addr, an attacker behind spoofed X-Forwarded-For headers or
+# a distributed set of sources could otherwise brute-force the 6-digit PIN. Once this
+# many failures occur across ALL IPs within the rolling window, every login attempt is
+# rejected with HTTP 429 until the window drains.
+GLOBAL_FAILED_ATTEMPTS_WINDOW_SECONDS: float = 300.0  # 5-minute rolling window
+GLOBAL_FAILED_ATTEMPTS_LIMIT: int = 50
+
+# Per-IP login attempt tracking. Format: {ip_address: last_failed_attempt_timestamp}
 login_attempts: Dict[str, float] = {}
+# Timestamps of recent failed login attempts across all IPs (for the global ceiling).
+global_failed_attempts: List[float] = []
 ToolCallList = List[ProcessedToolCall]
 
 
+def _prune_global_failed_attempts(current_time: float) -> None:
+    """Drop failed-attempt timestamps that have aged out of the global window."""
+    cutoff = current_time - GLOBAL_FAILED_ATTEMPTS_WINDOW_SECONDS
+    while global_failed_attempts and global_failed_attempts[0] < cutoff:
+        global_failed_attempts.pop(0)
+
+
+def _cleanup_login_attempts(current_time: float) -> None:
+    """Prune stale per-IP entries once the tracking dict grows too large."""
+    if len(login_attempts) <= LOGIN_ATTEMPTS_CLEANUP_THRESHOLD:
+        return
+    cutoff = current_time - LOGIN_ATTEMPTS_CLEANUP_AGE_SECONDS
+    for ip, timestamp in list(login_attempts.items()):
+        if timestamp < cutoff:
+            del login_attempts[ip]
+
+
+def reset_login_rate_limit_state() -> None:
+    """Reset all module-level login rate-limit state (used by tests)."""
+    login_attempts.clear()
+    global_failed_attempts.clear()
+
+
+def validate_attached_images(images: Optional[List[str]]) -> Optional[Tuple[Response, int]]:
+    """Enforce server-side limits on attached images.
+
+    Args:
+        images: List of base64 data-URL image strings, or None.
+
+    Returns:
+        An error response tuple if a limit is violated, otherwise None.
+    """
+    if not images:
+        return None
+    if len(images) > MAX_ATTACHED_IMAGES:
+        return AppManager.make_response(
+            message=f"Too many attached images (maximum {MAX_ATTACHED_IMAGES})",
+            status="error",
+            code=400,
+        )
+    for image in images:
+        if len(image) > MAX_IMAGE_BASE64_BYTES:
+            return AppManager.make_response(
+                message=f"Attached image exceeds the maximum size of {MAX_IMAGE_BASE64_BYTES} bytes",
+                status="error",
+                code=400,
+            )
+    return None
+
+
 def get_provider_for_model(app: MatHudFlask, model_id: str) -> OpenAIAPIBase:
     """Get or create the appropriate provider instance for a model.
 
@@ -394,9 +465,20 @@ def login() -> ResponseReturnValue:
             return redirect(url_for("get_index"))
 
         if request.method == "POST":
-            client_ip = request.environ.get("HTTP_X_FORWARDED_FOR", request.remote_addr)
+            # Key per-IP limiting on the real socket peer. The X-Forwarded-For header
+            # is client-controlled and MUST NOT be trusted for security decisions;
+            # trusting it would let an attacker reset the limiter on every request.
+            client_ip = request.remote_addr or "unknown"
             current_time = time.time()
 
+            # Global ceiling: independent of IP so it cannot be bypassed by spoofed
+            # X-Forwarded-For values or distributed sources. Once tripped, ALL logins
+            # (even from otherwise-clean IPs) are rejected until the window drains.
+            _prune_global_failed_attempts(current_time)
+            if len(global_failed_attempts) >= GLOBAL_FAILED_ATTEMPTS_LIMIT:
+                flash("Too many failed login attempts across all users. Please try again later.")
+                return render_template("login.html"), 429
+
             pin_submitted = request.form.get("pin", "")
             auth_pin = AppManager.get_auth_pin()
 
@@ -406,37 +488,29 @@ def login() -> ResponseReturnValue:
 
             is_pin_correct = hmac.compare_digest(pin_submitted, auth_pin)
 
-            # --- NEW, CORRECT LOGIC ---
-
             if is_pin_correct:
-                # 1. PIN is correct. Login succeeds immediately.
+                # PIN is correct. Login succeeds immediately.
                 session["authenticated"] = True
                 login_attempts.pop(client_ip, None)  # Clear any old rate limit.
                 return redirect(url_for("get_index"))
-            else:
-                # 2. PIN is incorrect. Now we handle rate limiting.
-                if client_ip in login_attempts:
-                    time_since_last_failed = current_time - login_attempts[client_ip]
-
-                    if time_since_last_failed < 5.0:
-                        # 2a. Cooldown is ACTIVE. Block and show countdown.
-                        remaining_cooldown = 5.0 - time_since_last_failed
-                        display_time = math.ceil(remaining_cooldown)
-                        flash(f"Too many attempts. Please wait {display_time} seconds.")
-                        return render_template("login.html")
-
-                # 2b. PIN was wrong, but no active cooldown. Start a new one.
-                login_attempts[client_ip] = current_time
-                flash("Invalid access code")
-
-                # Cleanup logic (runs only after a failed attempt)
-                if len(login_attempts) > 1000:
-                    cutoff = current_time - 3600  # 1 hour
-                    for ip, timestamp in list(login_attempts.items()):
-                        if timestamp < cutoff:
-                            del login_attempts[ip]
 
-                return render_template("login.html")
+            # PIN is incorrect: record the failure globally, then enforce per-IP cooldown.
+            global_failed_attempts.append(current_time)
+
+            if client_ip in login_attempts:
+                time_since_last_failed = current_time - login_attempts[client_ip]
+                if time_since_last_failed < LOGIN_ATTEMPT_COOLDOWN_SECONDS:
+                    # Cooldown is ACTIVE. Block and show countdown.
+                    remaining_cooldown = LOGIN_ATTEMPT_COOLDOWN_SECONDS - time_since_last_failed
+                    display_time = math.ceil(remaining_cooldown)
+                    flash(f"Too many attempts. Please wait {display_time} seconds.")
+                    return render_template("login.html"), 429
+
+            # No active cooldown. Record this failure and start a new cooldown window.
+            login_attempts[client_ip] = current_time
+            flash("Invalid access code")
+            _cleanup_login_attempts(current_time)
+            return render_template("login.html")
 
         # For GET request
         return render_template("login.html")
@@ -810,6 +884,10 @@ def send_message_stream() -> ResponseReturnValue:
         if isinstance(attached_images_raw, list):
             attached_images = [img for img in attached_images_raw if isinstance(img, str)]
 
+        image_error = validate_attached_images(attached_images)
+        if image_error is not None:
+            return image_error
+
         # Get the provider for this model and update all relevant APIs
         provider = get_active_provider(app, ai_model)
 
@@ -1082,6 +1160,10 @@ def send_message() -> ResponseReturnValue:
         if isinstance(attached_images_raw, list):
             attached_images = [img for img in attached_images_raw if isinstance(img, str)]
 
+        image_error = validate_attached_images(attached_images)
+        if image_error is not None:
+            return image_error
+
         # Get the provider for this model and update all relevant APIs
         provider = get_active_provider(app, ai_model)