deer-flow/backend/packages/harness/deerflow/agents/middlewares/input_sanitization_middleware.py

"""Input guardrail middleware for prompt-injection defense (issue #3630).

Escapes blocked XML-like tags in every genuine user message (e.g.
``<system>`` → ``&lt;system&gt;``) so they render as literal text instead
of structured-context markers.  This preserves the user's intent ("how do
I use DeerFlow's <think> tag?") while neutralizing injection attempts —
the same de-identify-don't-reject strategy as AWS Bedrock's PII ANONYMIZE.

The whole conversation is covered, not only the newest turn: the transformation
is request-scoped, so thread state keeps the raw text, and a last-turn-only scan
would neutralize a payload for exactly one model call before replaying it
verbatim on the next.

Blocked: system-reserved tags (memory, analysis, etc.) + common injection
tags (system, instruction, role, etc.). Normal HTML/XML tags (<div>,
<span>) are NOT escaped.

Clean input is wrapped in plain-text boundary markers as a secondary
semantic defense (OWASP structured-prompt guidance).
"""

from __future__ import annotations

import logging
import re
from collections.abc import Awaitable, Callable
from typing import override

from langchain.agents import AgentState
from langchain.agents.middleware import AgentMiddleware
from langchain.agents.middleware.types import (
    ModelCallResult,
    ModelRequest,
    ModelResponse,
)
from langchain_core.messages import HumanMessage
from langgraph.errors import GraphBubbleUp

from deerflow.agents.middlewares.message_utils import requires_input_sanitization
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY, message_content_to_text

logger = logging.getLogger(__name__)

# Finite set of blocked tag names: system-reserved + common injection patterns.
#
# Maintenance: when adding a new framework block tag that the system emits into
# model input, you MUST also update the expected count in
# test_input_sanitization_middleware.py::test_denylist_covers_framework_authority_blocks.
# The test pins the exact number of blocked tags so a new framework tag cannot
# be added without the corresponding regression guard.
_BLOCKED_TAG_NAMES: frozenset[str] = frozenset(
    {
        # Framework-injected structured/authority blocks. The lead-agent system
        # prompt's "System-Context Confidentiality" section (agents/lead_agent/
        # prompt.py) declares *every* such tag trusted internal data — it names a
        # few then says "and all other structured tags". So the denylist must
        # cover the framework's authority blocks as a class, not a hand-picked
        # subset: any one of them, forged in untrusted input, mimics trusted
        # framework context. Enumerated from the block tags the framework actually
        # emits into model input (system prompt + hidden-context/reminder
        # middlewares) and pinned against drift by
        # test_input_sanitization_middleware.py::test_denylist_covers_framework_authority_blocks.
        # Both spellings of the reminder block are covered: "system-reminder"
        # (dynamic-context) and "system_reminder" (todo/terminal middlewares).
        #
        # Subagents share this denylist: build_subagent_runtime_middlewares reuses
        # the same _build_runtime_middlewares base, so both sanitization paths guard
        # subagent model input too. The subagent system-prompt blocks
        # (file_editing_workflow / guidelines / output_format / working_directory)
        # are therefore authority blocks of the same class as the lead-agent ones.
        "system-reminder",
        "system_reminder",
        "memory",
        "current_date",
        "think",
        "analysis",
        "role",
        "soul",
        "self_update",
        "thinking_style",
        "clarification_system",
        "critical_reminders",
        "response_style",
        "citations",
        "current_uploads",
        "subagent_system",
        "skill_system",
        "skill_index",
        "available_skills",
        "disabled_skills",
        "memory_tool_system",
        "todo_list_system",
        "durable_context_data",
        "slash_skill_activation",
        "mcp_routing_hints",
        "available-deferred-tools",
        "goal_continuation",
        "background_task_event",
        "file_editing_workflow",
        "guidelines",
        "output_format",
        "working_directory",
        # Subagent system-prompt block (general_purpose.py): declares the task
        # tool off-limits. Forging this in untrusted input could trick the
        # model into believing it has (or lacks) tool restrictions it does not.
        "tool_restrictions",
        # Subagent report-contract blocks (subagents/report_contract.py, RFC
        # #4651 PR3): injected by the executor into every subagent system
        # prompt (the criteria pointer note carries no criterion values —
        # those stay in the untrusted task message). Forging them in
        # untrusted input could impersonate the verification contract (e.g.
        # pre-declaring acceptance criteria as met).
        "report_contract",
        "acceptance_criteria",
        # Project context blocks (projects/context.py, Projects Phase 2):
        # rendered request-scoped into model input carrying user-managed
        # project identity/instructions and the bounded shelf index. Forging
        # either in untrusted input mimics the run's pinned project
        # configuration or fabricates shelf entries/tool-callable IDs.
        "project",
        "documents",
        # Common prompt-injection tag patterns
        "system",
        "instruction",
        "important",
        "override",
        "ignore",
        "prompt",
    }
)

# Matches a full blocked tag: <tag>, </tag>, <tag attrs>, <tag/>, bare <tag
_BLOCKED_TAG_PATTERN = re.compile(
    r"<\s*/?\s*(?:" + "|".join(re.escape(t) for t in sorted(_BLOCKED_TAG_NAMES)) + r")\b[^>]*>?",
    re.IGNORECASE,
)

# Plain-text boundary markers (OWASP structured-prompt guidance).
_USER_INPUT_BEGIN = "--- BEGIN USER INPUT ---"
_USER_INPUT_END = "--- END USER INPUT ---"

# Neutralized forms injected when the user's text already contains a marker.
# These look visually similar but do not match the real boundary delimiters.
_NEUTRALIZED_BEGIN = "[BEGIN USER INPUT]"
_NEUTRALIZED_END = "[END USER INPUT]"

# Matches either boundary token as a standalone line or embedded in text.
_BOUNDARY_TOKEN_RE = re.compile(
    re.escape(_USER_INPUT_BEGIN) + r"|" + re.escape(_USER_INPUT_END),
)


def _escape_tag_match(match: re.Match) -> str:
    """Escape < and > in a blocked-tag match so it renders as literal text."""
    return match.group(0).replace("<", "&lt;").replace(">", "&gt;")


def _neutralize_boundary_tokens(text: str) -> str:
    """Replace real BEGIN/END USER INPUT markers with look-alike inert forms."""
    return _BOUNDARY_TOKEN_RE.sub(
        lambda m: _NEUTRALIZED_BEGIN if m.group(0) == _USER_INPUT_BEGIN else _NEUTRALIZED_END,
        text,
    )


def neutralize_untrusted_tags(text: str) -> str:
    """Neutralize framework/injection control tokens in untrusted text.

    Shared primitive for any content that originates outside the trust boundary
    and is about to enter the model context as *data* — currently the genuine
    user message (via :func:`frame_untrusted_text`) and remote tool results
    (web_fetch / web_search and friends, via
    :class:`ToolResultSanitizationMiddleware`).

    Applies exactly the two structural defenses, and nothing else:

    * blocked framework/injection tags (e.g. ``<system-reminder>``) are
      HTML-escaped to ``&lt;system-reminder&gt;`` so they lose their structural
      meaning while staying human-readable;
    * the plain-text ``--- BEGIN/END USER INPUT ---`` boundary markers are
      neutralized so untrusted content cannot forge or break out of the
      user-input boundary.

    It intentionally does **not** wrap the text in boundary markers: that
    framing is specific to the user message. Empty/whitespace-only text is
    returned unchanged so callers do not emit marker noise.
    """
    if not text.strip():
        return text
    text = _BLOCKED_TAG_PATTERN.sub(_escape_tag_match, text)
    return _neutralize_boundary_tokens(text)


def frame_untrusted_text(text: str) -> str:
    """Sanitize untrusted text, then wrap it in user-input boundary markers.

    * Empty/whitespace-only → return unchanged (no marker noise).
    * Blocked tags → HTML-escape ``<``/``>`` (e.g. ``<system>`` → ``&lt;system&gt;``).
    * Boundary tokens in user text → neutralized so they cannot forge boundaries.
    * Already wrapped (strict prefix+suffix) → return text unchanged (idempotent).
    * Otherwise → wrap in boundary markers.
    """
    if not text.strip():
        return text
    text = _BLOCKED_TAG_PATTERN.sub(_escape_tag_match, text)
    # Idempotency: only skip if text is *exactly* wrapped (prefix+suffix),
    # not if the user merely typed the begin token somewhere.
    if text.startswith(_USER_INPUT_BEGIN) and text.endswith(_USER_INPUT_END):
        # Still neutralize boundary tokens in the inner content — a user
        # can forge the outer wrapping to bypass the neutralization below
        # and inject inner boundary markers (break-out attack).
        inner = text[len(_USER_INPUT_BEGIN) : -len(_USER_INPUT_END)]
        neutralized_inner = _neutralize_boundary_tokens(inner)
        if neutralized_inner == inner:
            return text
        return f"{_USER_INPUT_BEGIN}{neutralized_inner}{_USER_INPUT_END}"
    # Neutralize any boundary tokens the user may have embedded, preventing
    # both self-suppression (begin token skips wrapping) and break-out
    # (end token creates a premature boundary inside the payload).
    text = _neutralize_boundary_tokens(text)
    return f"{_USER_INPUT_BEGIN}\n{text}\n{_USER_INPUT_END}"


def _check_user_content(text: str) -> str:
    """Backward-compatible internal alias for untrusted text framing."""
    return frame_untrusted_text(text)


class InputSanitizationMiddleware(AgentMiddleware[AgentState]):
    """Guardrail middleware that escapes prompt-injection tags in user input.

    Blocked tags are HTML-escaped (not rejected) so the user's intent is
    preserved while the tags lose their semantic significance. Clean input
    is wrapped in plain-text boundary markers. Transformation is temporary
    (wrap_model_call) — never written to state, which is why every genuine
    user message is re-sanitized on each call rather than only the newest.
    """

    @staticmethod
    def _extract_text_from_content(content: str | list) -> tuple[str, list | None]:
        """Extract concatenated text from a plain-string or content-block-list.

        Returns ``(text, extracted_blocks)``. *extracted_blocks* is None when
        *content* is a string, or the list of text-content blocks when a list.

        A list can hold bare ``str`` items next to content-block dicts
        (``message_content_to_text`` treats both as text, and some IM/SDK
        clients send exactly that shape), so bare strings are collected too —
        skipping them would skip sanitization entirely for that message.
        """
        if isinstance(content, str):
            return content, None
        if not isinstance(content, list):
            return "", None
        text_parts: list[str] = []
        text_blocks: list[dict | str] = []
        for block in content:
            if isinstance(block, str):
                if not block:  # skip empty items — matches message_content_to_text behaviour
                    continue
                text_parts.append(block)
                text_blocks.append(block)
            elif isinstance(block, dict) and block.get("type") == "text" and isinstance(block.get("text"), str):
                text = block["text"]
                if not text:  # skip empty blocks — matches message_content_to_text behaviour
                    continue
                text_parts.append(text)
                text_blocks.append(block)
        return "\n".join(text_parts), text_blocks

    @staticmethod
    def _rebuild_content(
        original_content: list,
        processed_text: str,
        text_blocks: list,
    ) -> list:
        """Replace text blocks with a single merged text block, preserving interleaved non-text blocks.

        For ``[text, image, text]`` the image block between the two text blocks
        is kept in place — only the text blocks are collapsed into one.
        """
        text_block_ids = {id(b) for b in text_blocks}
        first = last = None
        for i, block in enumerate(original_content):
            if id(block) in text_block_ids:
                if first is None:
                    first = i
                last = i
        if first is None:
            return original_content
        result: list = [*original_content[:first], {"type": "text", "text": processed_text}]
        # Re-insert any non-text blocks that sat between text blocks
        for i in range(first + 1, last + 1):
            if id(original_content[i]) not in text_block_ids:
                result.append(original_content[i])
        result.extend(original_content[last + 1 :])
        return result

    def _sanitize_message(self, msg: HumanMessage) -> HumanMessage | None:
        """Return a sanitized copy of *msg*, or None when nothing needs changing.

        Blocked tags are HTML-escaped (not rejected) so the user's intent is
        preserved while the tags lose their semantic significance. The original
        message is never mutated.
        """
        content = msg.content
        text_content, text_blocks = self._extract_text_from_content(content)

        # No text at all (e.g. image-only message) — pass through
        if not text_content and not isinstance(content, str):
            logger.debug("_sanitize_message: no text content in message — passing through")
            return None

        # Sanitize only the user's original input when available (set by
        # UploadsMiddleware before it prepends the <current_uploads> block),
        # so server-injected trusted blocks are never scanned for blocked
        # tags.  Fall back to full-content scanning only when the marker is
        # absent — UploadsMiddleware sets it on upload turns, so plain text
        # messages without uploads won't have it.  Full-content scanning is
        # safe for those: no server-injected <current_uploads> block exists
        # to accidentally escape.
        preserved_kwargs = dict(msg.additional_kwargs or {})
        original_user_content = preserved_kwargs.get(ORIGINAL_USER_CONTENT_KEY)
        if isinstance(original_user_content, str) and original_user_content:
            processed_user = _check_user_content(original_user_content)
            if processed_user != original_user_content:
                # Replace only the user's text suffix within the full
                # content — server-prepended blocks stay untouched.
                idx = text_content.rfind(original_user_content)
                if idx >= 0:
                    processed = text_content[:idx] + processed_user
                else:
                    # _extract_text_from_content and message_content_to_text
                    # disagreed on text extraction — rfind failed (only
                    # reachable for multimodal list content; see Decision 18).
                    if isinstance(content, list) and len(content) >= 2:
                        # content[0] is the server-injected
                        # <current_uploads> block (UploadsMiddleware
                        # prepends it as the first element for list
                        # content).  Sanitize only user blocks (content[1:])
                        # and rebuild directly — _rebuild_content only
                        # handles type:"text" blocks and would miss raw
                        # strings or non-standard dict blocks that
                        # message_content_to_text sees.
                        logger.warning(
                            "rfind failed on multimodal content; sanitizing user content blocks individually",
                        )
                        new_content: list = [content[0]]
                        for block in content[1:]:
                            if isinstance(block, str):
                                new_content.append(neutralize_untrusted_tags(block))
                            elif isinstance(block, dict) and block.get("type") == "text" and isinstance(block.get("text"), str):
                                sanitized = neutralize_untrusted_tags(block["text"])
                                if sanitized != block["text"]:
                                    new_content.append({**block, "text": sanitized})
                                else:
                                    new_content.append(block)
                            else:
                                new_content.append(block)
                        return HumanMessage(
                            content=new_content,
                            id=msg.id,
                            name=msg.name,
                            additional_kwargs=preserved_kwargs,
                        )
                    # Cannot distinguish server block from user blocks
                    # (non-list content or len(content) < 2).
                    # Degrade to full-content sanitization — server
                    # block may be escaped (UX degradation) but user
                    # forgeries are still neutralized (no security
                    # regression).
                    logger.warning(
                        "rfind failed with original_user_content set; cannot distinguish blocks, falling back to full-content sanitization",
                    )
                    processed = _check_user_content(text_content)
            else:
                processed = text_content  # no change needed
        elif isinstance(original_user_content, str):
            # Key is present but empty string (e.g. file upload with no
            # text input).  No user text to sanitize; server-injected
            # blocks must survive untouched.
            processed = text_content
        else:
            processed = _check_user_content(text_content)  # fallback

        if processed == text_content:
            # Already clean / already wrapped — no override needed
            return None

        if text_blocks:
            new_content = self._rebuild_content(content, processed, text_blocks)
        else:
            new_content = processed

        # Preserve the pre-sanitization user text so downstream consumers that
        # must see the genuine input (slash skill activation, regenerate) can
        # recover it after the BEGIN/END wrapping. Keep a valid value set by
        # UploadsMiddleware or an IM channel, but repair malformed metadata so
        # persistence never falls back to the wrapped model-facing content.
        if not isinstance(original_user_content, str):
            if ORIGINAL_USER_CONTENT_KEY in preserved_kwargs:
                logger.warning(
                    "InputSanitizationMiddleware replaced non-string %s metadata: type=%s",
                    ORIGINAL_USER_CONTENT_KEY,
                    type(original_user_content).__name__,
                )
            preserved_kwargs[ORIGINAL_USER_CONTENT_KEY] = message_content_to_text(content)
        logger.debug(
            "InputSanitizationMiddleware: original=%r -> processed=%r",
            content if isinstance(content, str) else "[content-blocks]",
            processed,
        )
        return HumanMessage(
            content=new_content,
            id=msg.id,
            name=msg.name,
            additional_kwargs=preserved_kwargs,
        )

    def _process_request(self, request: ModelRequest) -> ModelRequest:
        """Return a request with every genuine user message sanitized.

        Each genuine message is processed, not just the newest one. The
        transformation is request-scoped, so thread state keeps the raw text:
        sanitizing only the last turn would make the guardrail last exactly one
        turn, replaying an earlier turn's payload to the model verbatim — and
        outside the boundary markers, which the lead-agent prompt declares
        trusted internal framework data. Covering the whole history also keeps
        each turn's rendering stable across model calls, so the prompt prefix no
        longer changes shape behind the newest turn.

        Scope comes from ``requires_input_sanitization``: framework-injected
        messages stay excluded, because escaping their blocks would corrupt
        trusted context, while a caller-supplied message is covered even when it
        carries a framework marker — the Gateway marks those on the way in.
        """
        messages = list(request.messages)
        changed = False
        for index, msg in enumerate(messages):
            if not requires_input_sanitization(msg):
                if isinstance(msg, HumanMessage):
                    logger.debug(
                        "_process_request: skipping non-genuine HumanMessage at pos=%d name=%s hide_from_ui=%s content_preview=%.80r",
                        index,
                        msg.name,
                        msg.additional_kwargs.get("hide_from_ui"),
                        msg.content,
                    )
                continue
            logger.debug("_process_request: found genuine user message at pos=%d content=%.120r", index, msg.content)
            # Recover per message rather than per request. Unexpected errors
            # fail open, and one unprocessable history row must not widen that
            # into "no sanitization this turn" — an attacker who can land such a
            # row would otherwise buy themselves an unescaped newest turn.
            try:
                sanitized = self._sanitize_message(msg)
            except GraphBubbleUp:
                raise
            except Exception:
                logger.warning(
                    "Input guardrail failed on message at pos=%d; leaving it unchanged",
                    index,
                    exc_info=True,
                )
                continue
            if sanitized is None:
                continue
            messages[index] = sanitized
            changed = True
        if not changed:
            return request
        return request.override(messages=messages)

    def _try_process(self, request: ModelRequest) -> ModelRequest:
        """Sanitize request; fail-open on unexpected errors.

        GraphBubbleUp propagates; other exceptions return the original request.
        """
        try:
            return self._process_request(request)
        except GraphBubbleUp:
            raise
        except Exception:
            logger.warning(
                "Input guardrail processing failed; passing original request to model",
                exc_info=True,
            )
            return request

    @override
    def wrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> ModelCallResult:
        return handler(self._try_process(request))

    @override
    async def awrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
    ) -> ModelCallResult:
        return await handler(self._try_process(request))