Problem 118 · Agents

Prompt-Injection-Safe Tool Router with an Allowlist

MediumAI engineering

01Problem

Model output proposes tool calls like `{"name": ..., "arguments": {...}}`. Write `ToolRouter(tools, allowed_domains).route(call, session_allowlist, tainted)` returning a `Decision(allowed, reason, needs_approval)`. Reject unknown tools (exact name match only), tools not in this session's allowlist, missing or extra arguments, and wrong argument types (exact type: a bool is not an int). For URL parameters require `https` and a hostname equal to, or a subdomain of, an allowed domain. If the context is `tainted` (it contains untrusted content such as a fetched web page) a `write` tool is not executed but returned with `needs_approval=True`.

02Key insight

You cannot prompt your way out of prompt injection; the model will sometimes follow instructions hidden in a web page or email. So the security boundary is deterministic code between the model and the tools: an allowlist scoped per session, strict argument validation, and egress control on URLs, all of which assume the model's output is attacker-controlled. Parse URLs with a real parser and compare hostnames, because substring checks fall to `docs.example.com.evil.io` and `user@host` tricks. The taint rule encodes the 'untrusted input + private data + an exfiltration channel' combination: once untrusted text is in context, side-effecting calls need a human.

03Solution

solution.pypython59 lines

from dataclasses import dataclass
from typing import Any
from urllib.parse import urlparse


@dataclass(frozen=True)
class Tool:
    name: str
    params: dict[str, type]
    effect: str = "read"  # "read" | "write"
    url_params: tuple[str, ...] = ()


@dataclass(frozen=True)
class Decision:
    allowed: bool
    reason: str
    needs_approval: bool = False


class ToolRouter:
    def __init__(self, tools: list[Tool], allowed_domains: list[str]) -> None:
        self.tools = {t.name: t for t in tools}
        self.allowed_domains = frozenset(d.lower() for d in allowed_domains)

    def _url_ok(self, value: Any) -> bool:
        if not isinstance(value, str):
            return False
        parsed = urlparse(value)
        host = (parsed.hostname or "").lower()
        if parsed.scheme != "https" or not host:
            return False
        return any(host == d or host.endswith("." + d) for d in self.allowed_domains)

    def route(self, call: Any, session_allowlist: set[str], tainted: bool) -> Decision:
        if not isinstance(call, dict):
            return Decision(False, "malformed call")
        name = call.get("name")
        if not isinstance(name, str) or name not in self.tools:
            return Decision(False, "unknown tool")
        if name not in session_allowlist:
            return Decision(False, "tool not allowed in this session")
        tool = self.tools[name]
        args = call.get("arguments")
        if not isinstance(args, dict):
            return Decision(False, "arguments must be an object")
        missing = sorted(set(tool.params) - set(args))
        extra = sorted(set(args) - set(tool.params))
        if missing or extra:
            return Decision(False, f"bad arguments: missing={missing} extra={extra}")
        for key, typ in tool.params.items():
            if type(args[key]) is not typ:
                return Decision(False, f"argument {key} must be {typ.__name__}")
        for key in tool.url_params:
            if not self._url_ok(args[key]):
                return Decision(False, f"url in {key} is not an allowed https host")
        if tool.effect == "write" and tainted:
            return Decision(False, "write after untrusted input needs human approval", needs_approval=True)
        return Decision(True, "ok")

Try it first.

TimeO(number of arguments + allowed domains) per call

SpaceO(number of tools)

tests.py · 12 passing

tests.pypython23 lines

router = ToolRouter([
    Tool("fetch", {"url": str}, "read", ("url",)),
    Tool("send_email", {"to": str, "body": str}, "write"),
    Tool("set_limit", {"n": int}, "write"),
], ["docs.example.com", "example.org"])
allow = {"fetch", "send_email", "set_limit"}
ok = lambda d: d.allowed
assert ok(router.route({"name": "fetch", "arguments": {"url": "https://docs.example.com/a"}}, allow, False))
assert ok(router.route({"name": "fetch", "arguments": {"url": "https://api.example.org/x"}}, allow, True))
for url in ["https://docs.example.com.evil.io/", "https://evil.io/?u=docs.example.com",
            "http://docs.example.com/", "https://docs.example.com@evil.io/", "https://notexample.org/",
            "javascript:alert(1)"]:
    assert not ok(router.route({"name": "fetch", "arguments": {"url": url}}, allow, False)), url
assert router.route({"name": "Send_Email", "arguments": {}}, allow, False).reason == "unknown tool"
assert router.route({"name": "send_email", "arguments": {"to": "a", "body": "b"}}, {"fetch"}, False).reason.startswith("tool not allowed")
assert not ok(router.route({"name": "send_email", "arguments": {"to": "a"}}, allow, False))
assert not ok(router.route({"name": "send_email", "arguments": {"to": "a", "body": "b", "bcc": "x"}}, allow, False))
assert not ok(router.route({"name": "set_limit", "arguments": {"n": True}}, allow, False))
assert ok(router.route({"name": "set_limit", "arguments": {"n": 5}}, allow, False))
d = router.route({"name": "send_email", "arguments": {"to": "a", "body": "b"}}, allow, True)
assert not d.allowed and d.needs_approval
assert ok(router.route({"name": "send_email", "arguments": {"to": "a", "body": "b"}}, allow, False))
assert not ok(router.route("send_email()", allow, False))

04Staff-level follow-up

How do you know when context is tainted? Track provenance per message (system, user, tool output from a trusted internal API, tool output from the open web) and propagate taint through summaries the model writes. Then go beyond allowlists: per-tool rate limits, argument-level policies (send only to the user's own domain), dry-run previews for writes, and a separate quarantined model that reads untrusted content and returns only structured fields.