From 47aca595089d8403d01e367e4306068fca06dd41 Mon Sep 17 00:00:00 2001 From: yanghailan Date: Wed, 22 Jul 2026 18:11:22 +0800 Subject: [PATCH] fix(demohouse/computer_use): prevent press key command injection --- demohouse/computer_use/tool_server/README.md | 2 +- .../computer_use/tool_server/config.toml | 1 + demohouse/computer_use/tool_server/main.py | 6 ++ .../tool_server/middleware/auth.py | 3 - .../tests/test_presskey_security.py | 85 +++++++++++++++++++ .../tests/test_startup_security.py | 23 +++++ .../tool_server/tools/computer_xdotool.py | 56 ++++++++++-- 7 files changed, 164 insertions(+), 12 deletions(-) create mode 100644 demohouse/computer_use/tool_server/tests/test_presskey_security.py create mode 100644 demohouse/computer_use/tool_server/tests/test_startup_security.py diff --git a/demohouse/computer_use/tool_server/README.md b/demohouse/computer_use/tool_server/README.md index 139f456f..b0eef110 100644 --- a/demohouse/computer_use/tool_server/README.md +++ b/demohouse/computer_use/tool_server/README.md @@ -120,7 +120,7 @@ The tool server supports a simple API-key based authentication and TLS: ```toml # Shared API key. Clients must send it via the `X-API-Key` (or `Authorization`) header. -# Leave empty to disable authentication. +# Required. Set AUTH_KEY or inject a non-empty value through deployment configuration. auth_key = "your-secret-api-key-here" [plugins] diff --git a/demohouse/computer_use/tool_server/config.toml b/demohouse/computer_use/tool_server/config.toml index b9d608e0..907b42d5 100644 --- a/demohouse/computer_use/tool_server/config.toml +++ b/demohouse/computer_use/tool_server/config.toml @@ -1,5 +1,6 @@ port = 8102 display = ":5" +# Required at startup. Set a non-empty value through AUTH_KEY or deployment configuration. auth_key = "" [log] diff --git a/demohouse/computer_use/tool_server/main.py b/demohouse/computer_use/tool_server/main.py index bc2dfb3a..8e4ba7df 100644 --- a/demohouse/computer_use/tool_server/main.py +++ b/demohouse/computer_use/tool_server/main.py @@ -28,10 +28,16 @@ configure_logging() +def validate_security_settings(settings) -> None: + if not settings.auth_key: + raise RuntimeError("auth_key must be configured before starting tool_server") + + def main(): import uvicorn settings = get_settings() + validate_security_settings(settings) uvicorn_kwargs = { "host": "0.0.0.0", diff --git a/demohouse/computer_use/tool_server/middleware/auth.py b/demohouse/computer_use/tool_server/middleware/auth.py index d8776bd4..938e927f 100644 --- a/demohouse/computer_use/tool_server/middleware/auth.py +++ b/demohouse/computer_use/tool_server/middleware/auth.py @@ -26,9 +26,6 @@ def __init__(self, app): super().__init__(app) async def dispatch(self, request: Request, call_next) -> Response: - if get_settings().auth_key == "": - return await call_next(request) - auth_header = request.headers.get("X-API-Key") if not auth_header: auth_header = request.headers.get("Authorization") diff --git a/demohouse/computer_use/tool_server/tests/test_presskey_security.py b/demohouse/computer_use/tool_server/tests/test_presskey_security.py new file mode 100644 index 00000000..e690763d --- /dev/null +++ b/demohouse/computer_use/tool_server/tests/test_presskey_security.py @@ -0,0 +1,85 @@ +import asyncio +import sys +import unittest +from unittest.mock import AsyncMock, Mock, patch +from pathlib import Path + +from fastapi import HTTPException + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(PROJECT_ROOT)) + +from tools.computer import PressKeyRequest +from tools.computer_xdotool import XDOComputerTool, normalize_key + + +class NormalizeKeyTests(unittest.TestCase): + def test_maps_existing_named_key_combination(self): + self.assertEqual(normalize_key("ctrl c"), "Control_L+c") + self.assertEqual(normalize_key("alt+tab"), "Alt_L+Tab") + + def test_accepts_single_alphanumeric_key(self): + self.assertEqual(normalize_key("a"), "a") + self.assertEqual(normalize_key("7"), "7") + + def test_rejects_shell_metacharacters(self): + with self.assertRaises(HTTPException) as context: + normalize_key("a;id>/tmp/poc") + + self.assertEqual(context.exception.status_code, 400) + + +class PressKeyExecutionTests(unittest.IsolatedAsyncioTestCase): + async def test_uses_exec_with_key_as_a_separate_argument(self): + process = type("Process", (), {"communicate": AsyncMock(return_value=(b"", b""))})() + tool = XDOComputerTool(display=":99") + + with patch( + "tools.computer_xdotool.asyncio.create_subprocess_exec", + new=AsyncMock(return_value=process), + ) as create_process: + result = await tool.press_key(PressKeyRequest(key="ctrl c")) + + create_process.assert_awaited_once() + args, kwargs = create_process.call_args + self.assertEqual(args, ("xdotool", "key", "--", "Control_L+c")) + self.assertEqual(kwargs["env"]["DISPLAY"], ":99") + self.assertEqual(result.output, "") + self.assertEqual(result.error, "") + + async def test_kills_subprocess_when_press_key_times_out(self): + process = type( + "Process", + (), + { + "communicate": AsyncMock(side_effect=asyncio.TimeoutError), + "kill": Mock(), + }, + )() + tool = XDOComputerTool(display=":99") + + with patch( + "tools.computer_xdotool.asyncio.create_subprocess_exec", + new=AsyncMock(return_value=process), + ): + with self.assertRaises(TimeoutError): + await tool.press_key(PressKeyRequest(key="enter")) + + process.kill.assert_called_once() + + async def test_returns_error_when_xdotool_is_unavailable(self): + tool = XDOComputerTool(display=":99") + + with patch( + "tools.computer_xdotool.asyncio.create_subprocess_exec", + new=AsyncMock(side_effect=FileNotFoundError("xdotool not found")), + ): + result = await tool.press_key(PressKeyRequest(key="enter")) + + self.assertEqual(result.output, "") + self.assertIn("xdotool not found", result.error) + + +if __name__ == "__main__": + unittest.main() diff --git a/demohouse/computer_use/tool_server/tests/test_startup_security.py b/demohouse/computer_use/tool_server/tests/test_startup_security.py new file mode 100644 index 00000000..ce3b1b62 --- /dev/null +++ b/demohouse/computer_use/tool_server/tests/test_startup_security.py @@ -0,0 +1,23 @@ +import sys +import unittest +from pathlib import Path +from types import SimpleNamespace + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(PROJECT_ROOT)) + +from main import validate_security_settings + + +class StartupSecurityTests(unittest.TestCase): + def test_rejects_empty_auth_key(self): + with self.assertRaisesRegex(RuntimeError, "auth_key"): + validate_security_settings(SimpleNamespace(auth_key="")) + + def test_accepts_configured_auth_key(self): + validate_security_settings(SimpleNamespace(auth_key="configured-secret")) + + +if __name__ == "__main__": + unittest.main() diff --git a/demohouse/computer_use/tool_server/tools/computer_xdotool.py b/demohouse/computer_use/tool_server/tools/computer_xdotool.py index 4f2a7a09..13880991 100644 --- a/demohouse/computer_use/tool_server/tools/computer_xdotool.py +++ b/demohouse/computer_use/tool_server/tools/computer_xdotool.py @@ -11,6 +11,7 @@ import base64 import asyncio +import os from typing import Optional from pathlib import Path from fastapi import HTTPException @@ -21,6 +22,30 @@ from .computer import * XDOTOOL_DELAY = 50 + + +def normalize_key(raw_key: str) -> str: + """Convert a supported client key expression to an xdotool key expression.""" + if not raw_key or not raw_key.strip(): + raise HTTPException(status_code=400, detail="Invalid key") + + tokens = [token for token in raw_key.replace("+", " ").split() if token] + if not tokens: + raise HTTPException(status_code=400, detail="Invalid key") + + normalized = [] + for token in tokens: + key_name = token.lower() + if key_name in KEYS: + normalized.append(KEYS[key_name]) + elif len(token) == 1 and token.isascii() and token.isalnum(): + normalized.append(key_name) + else: + raise HTTPException(status_code=400, detail="Invalid key") + + return "+".join(normalized) + + class XDOComputerTool(IComputerTool): def __init__( self, @@ -110,15 +135,30 @@ async def scroll(self, r: ScrollRequest): async def press_key(self, r: PressKeyRequest): self.logger.debug(f"press_key, {r}") - keys = [x for x in r.key.split(' ') if x] - if isinstance(keys, list): - key = "+".join( - KEYS[k.lower()] if k.lower() in KEYS else k.lower() for k in keys + key = normalize_key(r.key) + try: + process = await asyncio.create_subprocess_exec( + "xdotool", + "key", + "--", + key, + env={**os.environ, "DISPLAY": self._display}, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, ) - else: - key = KEYS[r.key.lower()] if r.key.lower() in KEYS else r.key.lower() - command_parts = [self._xdotool, f"key -- {key}"] - return await self.shell(" ".join(command_parts)) + except FileNotFoundError as exc: + return BaseResult(output="", error=str(exc)) + + try: + stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=10) + except asyncio.TimeoutError as exc: + try: + process.kill() + except ProcessLookupError: + pass + raise TimeoutError("xdotool key command timed out after 10 seconds") from exc + + return BaseResult(output=stdout.decode(), error=stderr.decode()) async def type_text(self, r: TypeTextRequest): results: list[BaseResult] = []