mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-24 13:49:17 +00:00
2e9ee329ce
InputSanitizationMiddleware's text extraction only collected dict blocks with type == "text", so a HumanMessage whose content list carried a bare str item (a shape message_content_to_text treats as text and some IM/SDK clients send) yielded no text at all — the request passed through unwrapped and unescaped, letting forged framework tags (<system-reminder> etc.) reach the model untouched. The sibling rfind-fallback path in _process_request already neutralized bare strings individually, and both ToolResultSanitizationMiddleware and ToolOutputBudgetMiddleware treat bare strings as text; the extraction helper was the odd one out. Collect bare string blocks alongside text-block dicts (skipping empty items, matching message_content_to_text), merging them into the single sanitized text block on rebuild while interleaved non-text blocks keep their positions.
461 lines
21 KiB
Python
461 lines
21 KiB
Python
"""Input guardrail middleware for prompt-injection defense (issue #3630).
|
|
|
|
Escapes blocked XML-like tags in the last genuine user message (e.g.
|
|
``<system>`` → ``<system>``) 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.
|
|
|
|
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.human_input import read_human_input_response
|
|
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY, message_content_to_text
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_SUMMARY_MESSAGE_NAME = "summary"
|
|
|
|
# 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",
|
|
"uploaded_files", # old uploads tag — still processed by deermem for backward-compat
|
|
"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",
|
|
"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",
|
|
# 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("<", "<").replace(">", ">")
|
|
|
|
|
|
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:`_check_user_content`) 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 ``<system-reminder>`` 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 _is_genuine_user_message(message: object) -> bool:
|
|
"""Return True for real user messages, excluding system-injected HumanMessages.
|
|
|
|
``hide_from_ui`` is also used by hidden UI replies from HumanInputCard, so
|
|
only skip hidden HumanMessages that do not carry a valid user response.
|
|
"""
|
|
if not isinstance(message, HumanMessage):
|
|
return False
|
|
if message.name == _SUMMARY_MESSAGE_NAME:
|
|
return False
|
|
if message.additional_kwargs.get("hide_from_ui") and read_human_input_response(message.additional_kwargs) is None:
|
|
return False
|
|
return True
|
|
|
|
|
|
def _check_user_content(text: str) -> str:
|
|
"""Sanitize user content: escape blocked tags, then wrap in boundary markers.
|
|
|
|
* Empty/whitespace-only → return unchanged (no marker noise).
|
|
* Blocked tags → HTML-escape ``<``/``>`` (e.g. ``<system>`` → ``<system>``).
|
|
* 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}"
|
|
|
|
|
|
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.
|
|
"""
|
|
|
|
@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 _process_request(self, request: ModelRequest) -> ModelRequest:
|
|
"""Return a request with the last genuine user message sanitized.
|
|
|
|
Blocked tags are HTML-escaped (not rejected) so the user's intent is
|
|
preserved while the tags lose their semantic significance. Transformation
|
|
is temporary — the original request is never mutated.
|
|
"""
|
|
messages = list(request.messages)
|
|
for i in range(len(messages) - 1, -1, -1):
|
|
msg = messages[i]
|
|
if not _is_genuine_user_message(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",
|
|
i,
|
|
msg.name,
|
|
msg.additional_kwargs.get("hide_from_ui"),
|
|
msg.content,
|
|
)
|
|
continue
|
|
content = msg.content
|
|
logger.debug("_process_request: found genuine user message at pos=%d content=%.120r", i, 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("_process_request: no text content in message — passing through")
|
|
return request
|
|
|
|
# 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)
|
|
messages[i] = HumanMessage(
|
|
content=new_content,
|
|
id=msg.id,
|
|
name=msg.name,
|
|
additional_kwargs=preserved_kwargs,
|
|
)
|
|
return request.override(messages=messages)
|
|
else:
|
|
# 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 request
|
|
|
|
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)
|
|
messages[i] = HumanMessage(
|
|
content=new_content,
|
|
id=msg.id,
|
|
name=msg.name,
|
|
additional_kwargs=preserved_kwargs,
|
|
)
|
|
logger.debug(
|
|
"InputSanitizationMiddleware: original=%r -> processed=%r",
|
|
content if isinstance(content, str) else "[content-blocks]",
|
|
processed,
|
|
)
|
|
return request.override(messages=messages)
|
|
return request
|
|
|
|
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))
|