mirror of
https://github.com/rookiestar28/ComfyUI-OpenClaw.git
synced 2026-08-14 00:48:07 +00:00
security(connector): add rate limiting, LINE replay protection, and safer defaults for post-connectivity hardening
This commit is contained in:
@@ -7,6 +7,40 @@ ComfyUI-OpenClaw is a **ComfyUI custom node pack** that adds:
|
||||
- A secure-by-default HTTP API for automation (webhooks, triggers, schedules, approvals, presets)
|
||||
- And more exciting features being added continuously
|
||||
|
||||

|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Installation](#installation)
|
||||
- [Quick Start (Minimal)](#quick-start-minimal)
|
||||
- [Configure an LLM key](#1-configure-an-llm-key-for-plannerrefinervision-helpers)
|
||||
- [Configure webhook auth](#2-configure-webhook-auth-required-for-webhook)
|
||||
- [Set an Admin Token](#3-optional-recommended-set-an-admin-token)
|
||||
- [Nodes](#nodes)
|
||||
- [Extension UI](#extension-ui)
|
||||
- [API Overview](#api-overview)
|
||||
- [Observability](#observability-read-only)
|
||||
- [LLM config](#llm-config-non-secret)
|
||||
- [Webhooks](#webhooks)
|
||||
- [Triggers + approvals](#triggers--approvals-admin)
|
||||
- [Schedules](#schedules-admin)
|
||||
- [Presets](#presets-admin)
|
||||
- [Packs](#packs-admin)
|
||||
- [Bridge](#bridge-sidecar-optional)
|
||||
- [Templates](#templates)
|
||||
- [Execution Budgets](#execution-budgets)
|
||||
- [LLM Failover](#llm-failover)
|
||||
- [State Directory & Logs](#state-directory--logs)
|
||||
- [Troubleshooting](#troubleshooting)
|
||||
- [Tests](#tests)
|
||||
- [Updating](#updating)
|
||||
- [Remote Control (Connector)](#-remote-control-connector)
|
||||
- [Security](#security)
|
||||
|
||||
---
|
||||
|
||||
## Installation
|
||||
|
||||
- ComfyUI-Manager: install as a custom node (recommended for most users), then restart ComfyUI.
|
||||
@@ -223,6 +257,30 @@ Templates live in `data/templates/`.
|
||||
|
||||
For the full step-by-step guide (where to put exported workflow JSON, how to author `manifest.json`, how to verify `/openclaw/templates`, and how to use `/run`), see `tests/TEST_SOP.md`.
|
||||
|
||||
### Basic `/run` usage (chat)
|
||||
|
||||
**Free-text prompt mode (no `key=value` needed):**
|
||||
|
||||
```
|
||||
/run z 画面中央是一位年轻女性… seed=-1
|
||||
```
|
||||
|
||||
The connector will map the free text into a prompt field using:
|
||||
|
||||
- `allowed_inputs` if a single key is declared in `manifest.json`, or
|
||||
- fallback order: `positive_prompt` → `prompt` → `text` → `positive` → `caption`.
|
||||
|
||||
**Key=value mode (explicit mapping):**
|
||||
|
||||
```
|
||||
/run z positive_prompt="a cat" seed=-1
|
||||
```
|
||||
|
||||
Important:
|
||||
|
||||
- Ensure your workflow uses the same placeholder (e.g., `"text": "{{positive_prompt}}"`).
|
||||
- `seed=-1` gives random seeds; a fixed seed reproduces outputs.
|
||||
|
||||
## Execution Budgets
|
||||
|
||||
Queue submissions are protected by concurrency caps and render size budgets (`services/execution_budgets.py`).
|
||||
@@ -295,6 +353,13 @@ Notes:
|
||||
|
||||
Set webhook auth env vars (see “Quick Start”) and restart ComfyUI.
|
||||
|
||||
### Admin Token: server-side vs UI
|
||||
|
||||
`OPENCLAW_ADMIN_TOKEN` is a **server-side environment variable**.
|
||||
The Settings UI can **use** an Admin Token for authenticated requests, but **cannot set or persist** the server token.
|
||||
|
||||
Full setup steps: see `tests/TEST_SOP.md`.
|
||||
|
||||
## Tests
|
||||
|
||||
Run unit tests from the repo root:
|
||||
|
||||
@@ -634,3 +634,91 @@ async def llm_test_handler(request: web.Request) -> web.Response:
|
||||
},
|
||||
status=500,
|
||||
)
|
||||
|
||||
|
||||
async def llm_chat_handler(request: web.Request) -> web.Response:
|
||||
"""
|
||||
POST /openclaw/llm/chat (legacy: /moltbot/llm/chat)
|
||||
Run a simple chat completion using server-side LLM config + keys.
|
||||
This endpoint is intended for the connector; no prompt content is logged.
|
||||
"""
|
||||
if web is None:
|
||||
raise RuntimeError("aiohttp not available")
|
||||
try:
|
||||
from ..services.async_utils import run_in_thread
|
||||
except ImportError:
|
||||
from services.async_utils import run_in_thread
|
||||
|
||||
# S17: Rate Limit
|
||||
if not check_rate_limit(request, "admin"):
|
||||
return web.json_response(
|
||||
{"ok": False, "error": "Rate limit exceeded"}, status=429
|
||||
)
|
||||
|
||||
# NOTE: Keep this server-side. Connector cannot access UI-stored secrets directly.
|
||||
# This endpoint ensures keys are resolved via backend config + secret store.
|
||||
# S13: Validate admin boundary (or loopback if no admin token configured)
|
||||
allowed, err = require_admin_token(request)
|
||||
if not allowed:
|
||||
return web.json_response(
|
||||
{
|
||||
"ok": False,
|
||||
"error": err or "Unauthorized",
|
||||
},
|
||||
status=403,
|
||||
)
|
||||
|
||||
try:
|
||||
body = await request.json()
|
||||
except Exception:
|
||||
body = {}
|
||||
|
||||
if not isinstance(body, dict):
|
||||
return web.json_response(
|
||||
{"ok": False, "error": "Expected JSON object body"},
|
||||
status=400,
|
||||
)
|
||||
|
||||
system = body.get("system") if isinstance(body.get("system"), str) else ""
|
||||
user_message = (
|
||||
body.get("user_message")
|
||||
if isinstance(body.get("user_message"), str)
|
||||
else body.get("message") if isinstance(body.get("message"), str) else ""
|
||||
)
|
||||
temperature = body.get("temperature") if isinstance(body.get("temperature"), (int, float)) else 0.7
|
||||
max_tokens = body.get("max_tokens") if isinstance(body.get("max_tokens"), int) else 1024
|
||||
|
||||
if not user_message:
|
||||
return web.json_response(
|
||||
{"ok": False, "error": "missing_user_message"},
|
||||
status=400,
|
||||
)
|
||||
|
||||
try:
|
||||
client = LLMClient()
|
||||
|
||||
def _run():
|
||||
return client.complete(
|
||||
system=system,
|
||||
user_message=user_message,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
|
||||
result = await run_in_thread(_run)
|
||||
text = ""
|
||||
if isinstance(result, dict):
|
||||
text = result.get("text") or ""
|
||||
return web.json_response({"ok": True, "text": text})
|
||||
except ValueError as e:
|
||||
# Common: missing API key for selected provider
|
||||
return web.json_response(
|
||||
{"ok": False, "error": str(e)},
|
||||
status=400,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"LLM chat request failed: {type(e).__name__}")
|
||||
return web.json_response(
|
||||
{"ok": False, "error": "llm_request_failed"},
|
||||
status=500,
|
||||
)
|
||||
|
||||
+5
-1
@@ -20,7 +20,7 @@ except ModuleNotFoundError: # pragma: no cover (optional for unit tests)
|
||||
PACK_NAME = PACK_VERSION = PACK_START_TIME = LOG_FILE = get_api_key = None # type: ignore
|
||||
metrics = tail_log = require_observability_access = check_rate_limit = trace_store = None # type: ignore
|
||||
webhook_handler = webhook_submit_handler = webhook_validate_handler = capabilities_handler = preflight_handler = None # type: ignore
|
||||
config_get_handler = config_put_handler = llm_test_handler = llm_models_handler = None # type: ignore
|
||||
config_get_handler = config_put_handler = llm_test_handler = llm_models_handler = llm_chat_handler = None # type: ignore
|
||||
templates_list_handler = None # type: ignore
|
||||
secrets_status_handler = secrets_put_handler = secrets_delete_handler = None # type: ignore
|
||||
list_checkpoints_handler = create_checkpoint_handler = get_checkpoint_handler = delete_checkpoint_handler = None # type: ignore
|
||||
@@ -43,6 +43,7 @@ if web is not None:
|
||||
config_put_handler,
|
||||
llm_models_handler,
|
||||
llm_test_handler,
|
||||
llm_chat_handler,
|
||||
)
|
||||
from ..api.preflight_handler import inventory_handler, preflight_handler
|
||||
from ..api.secrets import (
|
||||
@@ -84,6 +85,7 @@ if web is not None:
|
||||
config_put_handler,
|
||||
llm_models_handler,
|
||||
llm_test_handler,
|
||||
llm_chat_handler,
|
||||
)
|
||||
from api.preflight_handler import inventory_handler, preflight_handler
|
||||
from api.secrets import (
|
||||
@@ -454,6 +456,8 @@ def register_routes(server) -> None:
|
||||
("GET", f"{prefix}/config", config_get_handler),
|
||||
("PUT", f"{prefix}/config", config_put_handler),
|
||||
("POST", f"{prefix}/llm/test", llm_test_handler),
|
||||
# NOTE: Connector uses this endpoint to avoid missing UI-stored keys.
|
||||
("POST", f"{prefix}/llm/chat", llm_chat_handler),
|
||||
(
|
||||
"GET",
|
||||
f"{prefix}/llm/models",
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 30 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 260 KiB |
@@ -24,6 +24,33 @@ logging.basicConfig(
|
||||
logger = logging.getLogger("connector")
|
||||
|
||||
|
||||
def _print_security_banner(config):
|
||||
"""
|
||||
F32 WP1: Print security warning when allowlists are empty.
|
||||
Fail-closed: empty allowlists = all users treated as untrusted.
|
||||
"""
|
||||
has_trusted_users = bool(
|
||||
config.telegram_allowed_users or
|
||||
config.telegram_allowed_chats or
|
||||
config.discord_allowed_users or
|
||||
config.discord_allowed_channels or
|
||||
config.line_allowed_users or
|
||||
config.line_allowed_groups
|
||||
)
|
||||
has_admins = bool(config.admin_users)
|
||||
|
||||
if not has_trusted_users:
|
||||
logger.warning("=" * 60)
|
||||
logger.warning("⚠️ SECURITY: No trusted users configured.")
|
||||
logger.warning("⚠️ All /run commands will require approval.")
|
||||
logger.warning("⚠️ Set OPENCLAW_CONNECTOR_*_ALLOWED_USERS to enable auto-exec.")
|
||||
logger.warning("=" * 60)
|
||||
|
||||
if not has_admins:
|
||||
logger.warning("⚠️ No admin users configured (OPENCLAW_CONNECTOR_ADMIN_USERS).")
|
||||
logger.warning("⚠️ Admin commands (/approve, /reject, etc.) will be unavailable.")
|
||||
|
||||
|
||||
async def main():
|
||||
logger.info("Initializing OpenClaw Connector (Phase 5)...")
|
||||
|
||||
@@ -39,6 +66,9 @@ async def main():
|
||||
logging.getLogger("connector").setLevel(logging.DEBUG)
|
||||
logger.debug("Debug mode enabled")
|
||||
|
||||
# F32 WP1: Security warning banner when no trusted users configured
|
||||
_print_security_banner(config)
|
||||
|
||||
# 2. Components
|
||||
client = OpenClawClient(config)
|
||||
await client.start() # Start session
|
||||
|
||||
@@ -43,6 +43,12 @@ class ConnectorConfig:
|
||||
# Privileged Access (ID match across platforms; Telegram Int vs Discord Str handled by router)
|
||||
admin_users: List[str] = field(default_factory=list)
|
||||
|
||||
# Security (F32)
|
||||
rate_limit_user_rpm: int = 10 # Requests per minute per user
|
||||
rate_limit_channel_rpm: int = 30 # Requests per minute per channel
|
||||
max_command_length: int = 4096 # Max characters in a single command
|
||||
llm_max_tokens_per_request: int = 1024 # LLM token budget
|
||||
|
||||
# Global
|
||||
debug: bool = False
|
||||
state_path: Optional[str] = None
|
||||
@@ -108,4 +114,15 @@ def load_config() -> ConnectorConfig:
|
||||
if admins := os.environ.get("OPENCLAW_CONNECTOR_ADMIN_USERS"):
|
||||
cfg.admin_users = [u.strip() for u in admins.split(",") if u.strip()]
|
||||
|
||||
# Security (F32)
|
||||
if rpm := os.environ.get("OPENCLAW_CONNECTOR_RATE_LIMIT_USER_RPM"):
|
||||
if rpm.isdigit():
|
||||
cfg.rate_limit_user_rpm = int(rpm)
|
||||
if rpm := os.environ.get("OPENCLAW_CONNECTOR_RATE_LIMIT_CHANNEL_RPM"):
|
||||
if rpm.isdigit():
|
||||
cfg.rate_limit_channel_rpm = int(rpm)
|
||||
if max_len := os.environ.get("OPENCLAW_CONNECTOR_MAX_COMMAND_LENGTH"):
|
||||
if max_len.isdigit():
|
||||
cfg.max_command_length = int(max_len)
|
||||
|
||||
return cfg
|
||||
|
||||
+27
-38
@@ -20,7 +20,7 @@ class LLMClient:
|
||||
|
||||
Security:
|
||||
- No conversation memory (stateless).
|
||||
- No user prompt logging (privacy).
|
||||
- Never logs user prompt content.
|
||||
- Never auto-executes commands.
|
||||
"""
|
||||
|
||||
@@ -42,7 +42,13 @@ class LLMClient:
|
||||
|
||||
res = await self._client.get_openclaw_config()
|
||||
if res.get("ok"):
|
||||
self._config_cache = res.get("data", {})
|
||||
data = res.get("data", {})
|
||||
# /openclaw/config returns { ok, config, sources, providers }
|
||||
# Keep only the effective config block.
|
||||
if isinstance(data, dict) and isinstance(data.get("config"), dict):
|
||||
self._config_cache = data.get("config", {})
|
||||
else:
|
||||
self._config_cache = data if isinstance(data, dict) else {}
|
||||
else:
|
||||
self._config_cache = {}
|
||||
return self._config_cache
|
||||
@@ -60,9 +66,13 @@ class LLMClient:
|
||||
self._configured = True
|
||||
return True
|
||||
|
||||
# Check if API key is configured (via secret store or env)
|
||||
# OpenClaw settings will include "api_key_configured" flag
|
||||
self._configured = config.get("api_key_configured", False) or bool(config.get("provider"))
|
||||
# If backend includes an explicit flag, honor it.
|
||||
if "api_key_configured" in config:
|
||||
self._configured = bool(config.get("api_key_configured"))
|
||||
return self._configured
|
||||
|
||||
# Best-effort: consider configured if provider is set (key lookup happens at call time).
|
||||
self._configured = bool(provider)
|
||||
return self._configured
|
||||
|
||||
async def chat(
|
||||
@@ -81,44 +91,23 @@ class LLMClient:
|
||||
if not await self.is_configured():
|
||||
return "[Error] LLM not configured. Configure in OpenClaw Settings."
|
||||
|
||||
config = await self._fetch_config()
|
||||
provider = config.get("provider", "openai")
|
||||
model = config.get("model", "gpt-4o-mini")
|
||||
base_url = config.get("base_url")
|
||||
|
||||
# Default base URLs per provider (matches OpenClaw catalog)
|
||||
if not base_url:
|
||||
base_url = self._get_default_base_url(provider)
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_message},
|
||||
]
|
||||
|
||||
# NOTE: Use backend chat endpoint so we don't bypass server-side key resolution.
|
||||
# Direct provider calls from the connector can miss UI-stored keys and produce 401 errors.
|
||||
try:
|
||||
# Try using services.providers.openai_compat if available
|
||||
from services.providers.openai_compat import make_request
|
||||
from services.providers.keys import get_api_key_for_provider
|
||||
|
||||
api_key = get_api_key_for_provider(provider)
|
||||
|
||||
result = make_request(
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
messages=messages,
|
||||
model=model,
|
||||
res = await self._client.chat_llm(
|
||||
system=system_prompt,
|
||||
user_message=user_message,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
timeout=60.0,
|
||||
)
|
||||
return result.get("text", "[No response]")
|
||||
except ImportError:
|
||||
# Fallback: use aiohttp directly
|
||||
return await self._fallback_chat(config, messages, temperature, max_tokens)
|
||||
except Exception as e:
|
||||
if res.get("ok"):
|
||||
data = res.get("text") or res.get("data", {}).get("text")
|
||||
return data or "[No response]"
|
||||
return f"[LLM Error] {res.get('error', 'Request failed')}"
|
||||
except Exception:
|
||||
# Log error without user content
|
||||
logger.error(f"LLM request failed: {type(e).__name__}")
|
||||
return f"[LLM Error] Request failed. Please try again."
|
||||
logger.error("LLM request failed")
|
||||
return "[LLM Error] Request failed. Please try again."
|
||||
|
||||
def _get_default_base_url(self, provider: str) -> str:
|
||||
"""Get default base URL for provider (matches OpenClaw catalog)."""
|
||||
|
||||
@@ -52,7 +52,9 @@ class OpenClawClient:
|
||||
if self.session:
|
||||
await self.session.close()
|
||||
|
||||
async def _request(self, method: str, path: str, json_data: dict = None) -> dict:
|
||||
async def _request(
|
||||
self, method: str, path: str, json_data: dict = None, timeout: int = 10
|
||||
) -> dict:
|
||||
url = f"{self.base_url}{path}"
|
||||
session = self.session
|
||||
|
||||
@@ -64,7 +66,7 @@ class OpenClawClient:
|
||||
|
||||
try:
|
||||
async with session.request(
|
||||
method, url, headers=self.headers, json=json_data, timeout=10
|
||||
method, url, headers=self.headers, json=json_data, timeout=timeout
|
||||
) as resp:
|
||||
result = {"ok": resp.status in (200, 201, 202)}
|
||||
|
||||
@@ -87,7 +89,9 @@ class OpenClawClient:
|
||||
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Request failed {method} {path}: {e}")
|
||||
logger.error(
|
||||
f"Request failed {method} {path}: {type(e).__name__}: {e}"
|
||||
)
|
||||
return {"ok": False, "error": str(e)}
|
||||
finally:
|
||||
if local_session:
|
||||
@@ -99,6 +103,30 @@ class OpenClawClient:
|
||||
"""Fetch OpenClaw runtime config (provider, model, base_url, etc.)."""
|
||||
return await self._request("GET", "/openclaw/config")
|
||||
|
||||
async def get_templates(self) -> dict:
|
||||
"""Fetch available templates (ids + metadata)."""
|
||||
return await self._request("GET", "/openclaw/templates")
|
||||
|
||||
async def chat_llm(
|
||||
self,
|
||||
system: str,
|
||||
user_message: str,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: int = 1024,
|
||||
) -> dict:
|
||||
"""Run a server-side LLM chat (uses backend config + keys)."""
|
||||
# NOTE: Must call backend so UI-stored secrets are available (connector has no access).
|
||||
payload = {
|
||||
"system": system,
|
||||
"user_message": user_message,
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
}
|
||||
# NOTE: LLM calls can exceed the default 10s HTTP timeout.
|
||||
return await self._request(
|
||||
"POST", "/openclaw/llm/chat", payload, timeout=120
|
||||
)
|
||||
|
||||
async def get_health(self) -> dict:
|
||||
res = await self._request("GET", "/openclaw/health")
|
||||
# Health endpoint might return nested structure, but we just want the wrapper
|
||||
|
||||
@@ -34,6 +34,10 @@ def _import_aiohttp_web():
|
||||
|
||||
|
||||
class LINEWebhookServer:
|
||||
# F32 WP2: Replay protection config
|
||||
REPLAY_WINDOW_SEC = 300 # 5 minutes
|
||||
NONCE_CACHE_SIZE = 1000
|
||||
|
||||
def __init__(self, config: ConnectorConfig, router: CommandRouter):
|
||||
self.config = config
|
||||
self.router = router
|
||||
@@ -41,6 +45,8 @@ class LINEWebhookServer:
|
||||
self.runner = None
|
||||
self.site = None
|
||||
self.session = None
|
||||
# F32 WP2: LRU nonce cache (event_id -> timestamp)
|
||||
self._nonce_cache: dict = {}
|
||||
|
||||
async def start(self):
|
||||
"""Start the webhook server."""
|
||||
@@ -96,6 +102,11 @@ class LINEWebhookServer:
|
||||
logger.warning("Invalid LINE Signature")
|
||||
return web.Response(status=401, text="Invalid Signature")
|
||||
|
||||
# F32 WP2: Replay protection (timestamp + nonce)
|
||||
if not self._check_replay_protection(body_text):
|
||||
logger.warning("Replay attack detected or stale request")
|
||||
return web.Response(status=403, text="Replay Rejected")
|
||||
|
||||
# 2. Parse Event
|
||||
try:
|
||||
payload = json.loads(body_text)
|
||||
@@ -124,6 +135,55 @@ class LINEWebhookServer:
|
||||
|
||||
return hmac.compare_digest(generated, signature)
|
||||
|
||||
def _check_replay_protection(self, body_text: str) -> bool:
|
||||
"""
|
||||
F32 WP2: Replay protection using timestamp + nonce.
|
||||
Returns False if request should be rejected.
|
||||
"""
|
||||
try:
|
||||
payload = json.loads(body_text)
|
||||
except json.JSONDecodeError:
|
||||
return False # Will be caught later as Bad JSON
|
||||
|
||||
events = payload.get("events", [])
|
||||
if not events:
|
||||
return True # No events to process
|
||||
|
||||
now = time.time() * 1000 # LINE timestamps are in ms
|
||||
|
||||
for event in events:
|
||||
# Check timestamp freshness
|
||||
ts = event.get("timestamp", 0)
|
||||
age_sec = (now - ts) / 1000
|
||||
if age_sec > self.REPLAY_WINDOW_SEC or age_sec < -60:
|
||||
# Allow 60s clock skew in the future
|
||||
logger.debug(f"Stale or future event: age={age_sec:.1f}s")
|
||||
return False
|
||||
|
||||
# Check nonce (use replyToken or webhookEventId as unique identifier)
|
||||
nonce = event.get("webhookEventId") or event.get("replyToken")
|
||||
if nonce:
|
||||
if nonce in self._nonce_cache:
|
||||
logger.debug(f"Duplicate nonce: {nonce}")
|
||||
return False
|
||||
# Add to cache with timestamp
|
||||
self._nonce_cache[nonce] = ts
|
||||
# Evict old entries if cache is full
|
||||
self._evict_old_nonces()
|
||||
|
||||
return True
|
||||
|
||||
def _evict_old_nonces(self):
|
||||
"""Remove old entries from nonce cache."""
|
||||
if len(self._nonce_cache) <= self.NONCE_CACHE_SIZE:
|
||||
return
|
||||
|
||||
now = time.time() * 1000
|
||||
cutoff = now - (self.REPLAY_WINDOW_SEC * 1000)
|
||||
self._nonce_cache = {
|
||||
k: v for k, v in self._nonce_cache.items() if v > cutoff
|
||||
}
|
||||
|
||||
async def _process_event(self, event: dict):
|
||||
"""Convert LINE event to CommandRequest and route."""
|
||||
source = event.get("source", {})
|
||||
|
||||
@@ -88,6 +88,8 @@ class TelegramPolling:
|
||||
return
|
||||
|
||||
updates = data.get("result", [])
|
||||
if self.config.debug and not updates:
|
||||
logger.debug("Telegram poll OK (no updates). offset=%s", self.offset)
|
||||
for update in updates:
|
||||
next_offset = update["update_id"] + 1
|
||||
if next_offset > self.offset:
|
||||
@@ -98,18 +100,35 @@ class TelegramPolling:
|
||||
await self._process_update(update)
|
||||
|
||||
async def _process_update(self, update: dict):
|
||||
message = update.get("message")
|
||||
# Telegram update shapes vary by chat type and sender mode.
|
||||
# - Normal groups/DMs: `message`
|
||||
# - Edited messages: `edited_message`
|
||||
# - Channels: `channel_post` / `edited_channel_post`
|
||||
#
|
||||
# IMPORTANT (recurring support issue):
|
||||
# If users say "DM works but group/channel does nothing" AND connector logs show no
|
||||
# `DEBUG raw message`, it's often because updates are arriving under `channel_post`
|
||||
# (or `sender_chat` anonymous posts) which older code ignored.
|
||||
message = (
|
||||
update.get("message")
|
||||
or update.get("edited_message")
|
||||
or update.get("channel_post")
|
||||
or update.get("edited_channel_post")
|
||||
)
|
||||
if not message or "text" not in message:
|
||||
return
|
||||
|
||||
chat_id = message["chat"]["id"]
|
||||
user_id = message["from"]["id"]
|
||||
username = message["from"].get("username", "unknown")
|
||||
# `from` may be missing for channel posts; `sender_chat` is used for anonymous admins.
|
||||
from_obj = message.get("from") or {}
|
||||
sender_chat = message.get("sender_chat") or {}
|
||||
user_id = from_obj.get("id")
|
||||
username = from_obj.get("username") or sender_chat.get("username") or "unknown"
|
||||
text = message["text"]
|
||||
|
||||
# Security Check
|
||||
is_allowed = False
|
||||
if user_id in self.config.telegram_allowed_users:
|
||||
if isinstance(user_id, int) and user_id in self.config.telegram_allowed_users:
|
||||
is_allowed = True
|
||||
if chat_id in self.config.telegram_allowed_chats:
|
||||
is_allowed = True
|
||||
@@ -122,7 +141,9 @@ class TelegramPolling:
|
||||
# Build Request
|
||||
req = CommandRequest(
|
||||
platform="telegram",
|
||||
sender_id=str(user_id),
|
||||
# If `user_id` is missing (channel posts), fall back to chat_id so allowlisting by
|
||||
# `OPENCLAW_CONNECTOR_TELEGRAM_ALLOWED_CHATS` still works deterministically.
|
||||
sender_id=str(user_id) if user_id is not None else str(chat_id),
|
||||
channel_id=str(chat_id),
|
||||
username=username,
|
||||
message_id=str(message["message_id"]),
|
||||
@@ -194,4 +215,3 @@ class TelegramPolling:
|
||||
logger.error(f"Telegram send_message failed: {r.status} {err}")
|
||||
except Exception as e:
|
||||
logger.error(f"Telegram send_message error: {e}")
|
||||
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
"""
|
||||
Rate Limiter for Connector (F32 WP2).
|
||||
Token bucket implementation for per-user and per-channel rate limiting.
|
||||
"""
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict
|
||||
|
||||
|
||||
@dataclass
|
||||
class TokenBucket:
|
||||
"""Simple token bucket for rate limiting."""
|
||||
capacity: float # Max tokens
|
||||
refill_rate: float # Tokens per second
|
||||
tokens: float = field(default=0.0)
|
||||
last_refill: float = field(default_factory=time.time)
|
||||
|
||||
def consume(self, tokens: int = 1) -> bool:
|
||||
"""
|
||||
Try to consume tokens. Returns True if allowed, False if rate limited.
|
||||
"""
|
||||
now = time.time()
|
||||
elapsed = now - self.last_refill
|
||||
self.tokens = min(self.capacity, self.tokens + elapsed * self.refill_rate)
|
||||
self.last_refill = now
|
||||
|
||||
if self.tokens >= tokens:
|
||||
self.tokens -= tokens
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
class RateLimiter:
|
||||
"""
|
||||
Per-user and per-channel rate limiter.
|
||||
|
||||
Default: 10 req/min per user, 30 req/min per channel.
|
||||
"""
|
||||
|
||||
def __init__(self, user_rpm: int = 10, channel_rpm: int = 30):
|
||||
self.user_rpm = user_rpm
|
||||
self.channel_rpm = channel_rpm
|
||||
self._user_buckets: Dict[str, TokenBucket] = {}
|
||||
self._channel_buckets: Dict[str, TokenBucket] = {}
|
||||
|
||||
def _get_user_bucket(self, user_id: str) -> TokenBucket:
|
||||
if user_id not in self._user_buckets:
|
||||
# capacity = rpm, refill_rate = rpm/60
|
||||
self._user_buckets[user_id] = TokenBucket(
|
||||
capacity=float(self.user_rpm),
|
||||
refill_rate=self.user_rpm / 60.0,
|
||||
tokens=float(self.user_rpm),
|
||||
)
|
||||
return self._user_buckets[user_id]
|
||||
|
||||
def _get_channel_bucket(self, channel_id: str) -> TokenBucket:
|
||||
if channel_id not in self._channel_buckets:
|
||||
self._channel_buckets[channel_id] = TokenBucket(
|
||||
capacity=float(self.channel_rpm),
|
||||
refill_rate=self.channel_rpm / 60.0,
|
||||
tokens=float(self.channel_rpm),
|
||||
)
|
||||
return self._channel_buckets[channel_id]
|
||||
|
||||
def is_allowed(self, user_id: str, channel_id: str) -> bool:
|
||||
"""
|
||||
Check if request is allowed. Returns False if rate limited.
|
||||
Both user and channel must have available tokens.
|
||||
"""
|
||||
user_bucket = self._get_user_bucket(user_id)
|
||||
channel_bucket = self._get_channel_bucket(channel_id)
|
||||
|
||||
# Check both limits
|
||||
user_ok = user_bucket.consume(1)
|
||||
channel_ok = channel_bucket.consume(1)
|
||||
|
||||
return user_ok and channel_ok
|
||||
|
||||
def cleanup(self, max_age_seconds: float = 3600.0):
|
||||
"""Remove stale buckets (optional, for memory management)."""
|
||||
now = time.time()
|
||||
cutoff = now - max_age_seconds
|
||||
|
||||
self._user_buckets = {
|
||||
k: v for k, v in self._user_buckets.items() if v.last_refill > cutoff
|
||||
}
|
||||
self._channel_buckets = {
|
||||
k: v for k, v in self._channel_buckets.items() if v.last_refill > cutoff
|
||||
}
|
||||
+105
-4
@@ -17,6 +17,7 @@ if False: # Type hinting only
|
||||
|
||||
from .llm_client import LLMClient
|
||||
from .prompts import CHAT_SYSTEM_PROMPT, CHAT_STATUS_PROMPT
|
||||
from .rate_limiter import RateLimiter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -27,10 +28,38 @@ class CommandRouter:
|
||||
self.client = client
|
||||
self.poller = poller
|
||||
self.state = ConnectorState(path=self.config.state_path)
|
||||
self._template_meta_cache: Dict[str, Dict[str, Any]] = {}
|
||||
# F32 WP2: Rate limiter
|
||||
self._rate_limiter = RateLimiter(
|
||||
user_rpm=self.config.rate_limit_user_rpm,
|
||||
channel_rpm=self.config.rate_limit_channel_rpm,
|
||||
)
|
||||
|
||||
async def handle(self, req: CommandRequest) -> CommandResponse:
|
||||
"""Main dispatch loop."""
|
||||
text = req.text.strip()
|
||||
# NOTE: Debug-only raw message logging for troubleshooting parsing issues.
|
||||
# Enable with OPENCLAW_CONNECTOR_DEBUG=1. May include sensitive user content.
|
||||
if self.config.debug:
|
||||
logger.info(
|
||||
"DEBUG raw message: platform=%s user=%s chat=%s text=%r",
|
||||
req.platform,
|
||||
req.sender_id,
|
||||
req.channel_id,
|
||||
text,
|
||||
)
|
||||
|
||||
# F32 WP2: Rate limiting
|
||||
if not self._rate_limiter.is_allowed(str(req.sender_id), str(req.channel_id)):
|
||||
return CommandResponse(
|
||||
text="[Rate Limited] Too many requests. Please wait a moment."
|
||||
)
|
||||
|
||||
# F32 WP5: Command length limit
|
||||
if len(text) > self.config.max_command_length:
|
||||
return CommandResponse(
|
||||
text=f"[Error] Command too long ({len(text)} chars). Max: {self.config.max_command_length}."
|
||||
)
|
||||
|
||||
try:
|
||||
parts = shlex.split(text)
|
||||
@@ -45,6 +74,17 @@ class CommandRouter:
|
||||
cmd = parts[0].lower()
|
||||
args = parts[1:]
|
||||
|
||||
# Telegram group commands often include the bot username suffix, e.g. `/help@mybot`.
|
||||
# If we don't strip it, the command won't match our dispatch table and appears "dead"
|
||||
# even though polling is working.
|
||||
if (req.platform or "").lower() == "telegram" and cmd.startswith("/") and "@" in cmd:
|
||||
cmd = cmd.split("@", 1)[0]
|
||||
|
||||
# Some users type `@bot /help` in group chats. Treat that as a command too.
|
||||
if cmd.startswith("@") and args and args[0].startswith("/"):
|
||||
cmd = args[0].lower()
|
||||
args = args[1:]
|
||||
|
||||
# Dispatch Table
|
||||
handlers = {
|
||||
("/status", "status"): (self._handle_status, False),
|
||||
@@ -179,7 +219,7 @@ class CommandRouter:
|
||||
) -> CommandResponse:
|
||||
if not args:
|
||||
return CommandResponse(
|
||||
text="Usage: /run <template_id> [key=value ...] [--approval]"
|
||||
text="Usage: /run <template_id> [prompt text] [key=value ...] [--approval]"
|
||||
)
|
||||
|
||||
# Parse flags
|
||||
@@ -195,11 +235,37 @@ class CommandRouter:
|
||||
return CommandResponse(text="Usage: /run <template_id> ...")
|
||||
|
||||
template_id = clean_args[0]
|
||||
inputs = {}
|
||||
inputs: Dict[str, str] = {}
|
||||
free_text_parts: List[str] = []
|
||||
for arg in clean_args[1:]:
|
||||
if "=" in arg:
|
||||
k, v = arg.split("=", 1)
|
||||
inputs[k.strip()] = v.strip()
|
||||
else:
|
||||
free_text_parts.append(arg)
|
||||
|
||||
# If user provided free text without key=value, treat it as the prompt.
|
||||
# We map it to a best-effort prompt key (prefers template metadata if available).
|
||||
if free_text_parts:
|
||||
prompt_key = await self._resolve_prompt_key(template_id)
|
||||
if prompt_key not in inputs:
|
||||
inputs[prompt_key] = " ".join(free_text_parts).strip()
|
||||
elif self.config.debug:
|
||||
logger.info(
|
||||
"DEBUG /run free-text ignored (prompt key already set): %s",
|
||||
prompt_key,
|
||||
)
|
||||
|
||||
# NOTE: Debug-only payload logging for troubleshooting prompt mismatches.
|
||||
# Enable with OPENCLAW_CONNECTOR_DEBUG=1 to log template_id + inputs.
|
||||
if self.config.debug:
|
||||
logger.info(
|
||||
"DEBUG /run payload: template=%s inputs=%s approval_flag=%s trusted=%s",
|
||||
template_id,
|
||||
inputs,
|
||||
explicit_approval,
|
||||
self._is_trusted(req),
|
||||
)
|
||||
|
||||
trusted = self._is_trusted(req)
|
||||
require_approval = explicit_approval or (not trusted)
|
||||
@@ -229,6 +295,42 @@ class CommandRouter:
|
||||
err = res.get("error", "Unknown error")
|
||||
return CommandResponse(text=f"[Submission Failed] Reason: {err}")
|
||||
|
||||
async def _resolve_prompt_key(self, template_id: str) -> str:
|
||||
"""
|
||||
Best-effort prompt key resolution.
|
||||
Prefer template metadata (allowed_inputs), then fall back to common names.
|
||||
"""
|
||||
meta = await self._get_template_meta(template_id)
|
||||
allowed = meta.get("allowed_inputs") or []
|
||||
|
||||
# If template explicitly declares a single input, use it.
|
||||
if isinstance(allowed, list) and len(allowed) == 1:
|
||||
return str(allowed[0])
|
||||
|
||||
preferred = ("positive_prompt", "prompt", "text", "positive", "caption")
|
||||
if isinstance(allowed, list):
|
||||
for key in preferred:
|
||||
if key in allowed:
|
||||
return key
|
||||
|
||||
# Default fallback
|
||||
return "positive_prompt"
|
||||
|
||||
async def _get_template_meta(self, template_id: str) -> Dict[str, Any]:
|
||||
if template_id in self._template_meta_cache:
|
||||
return self._template_meta_cache[template_id]
|
||||
try:
|
||||
res = await self.client.get_templates()
|
||||
if res.get("ok"):
|
||||
for item in res.get("templates", []) or []:
|
||||
if item.get("id") == template_id:
|
||||
self._template_meta_cache[template_id] = item
|
||||
return item
|
||||
except Exception as e:
|
||||
if self.config.debug:
|
||||
logger.info(f"DEBUG template meta fetch failed: {e}")
|
||||
return {}
|
||||
|
||||
async def _handle_interrupt(
|
||||
self, req: CommandRequest, args: List[str]
|
||||
) -> CommandResponse:
|
||||
@@ -361,7 +463,7 @@ class CommandRouter:
|
||||
text=(
|
||||
"OpenClaw Connector\n"
|
||||
"/status - Check system health and queue\n"
|
||||
"/run <template> [k=v] - Run a generation (trusted users auto-exec; others require approval)\n"
|
||||
"/run <template> [prompt] [k=v] - Run a generation (trusted users auto-exec; others require approval)\n"
|
||||
"/stop - Global Interrupt (Admin)\n"
|
||||
"/history <id> - Job details\n"
|
||||
"/jobs - Queue summary\n"
|
||||
@@ -534,4 +636,3 @@ Keep it minimal."""
|
||||
|
||||
response = await llm.chat(system_prompt, user_prompt)
|
||||
return CommandResponse(text=response)
|
||||
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
# OpenClaw Connector Security Checklist
|
||||
|
||||
> **Complete this checklist before enabling public ingress (tunnel, reverse proxy, or direct exposure).**
|
||||
|
||||
## ✅ Pre-Deployment Checklist
|
||||
|
||||
### 1. Authentication & Trust
|
||||
|
||||
- [ ] Set `OPENCLAW_CONNECTOR_ADMIN_USERS` with at least one admin ID.
|
||||
- [ ] Configure platform-specific allowlists:
|
||||
- Telegram: `OPENCLAW_CONNECTOR_TELEGRAM_ALLOWED_USERS` / `_ALLOWED_CHATS`
|
||||
- Discord: `OPENCLAW_CONNECTOR_DISCORD_ALLOWED_USERS` / `_ALLOWED_CHANNELS`
|
||||
- LINE: `OPENCLAW_CONNECTOR_LINE_ALLOWED_USERS` / `_ALLOWED_GROUPS`
|
||||
- [ ] Verify startup banner shows "No trusted users" warning if allowlists are empty.
|
||||
|
||||
### 2. Webhook Security (LINE)
|
||||
|
||||
- [ ] HTTPS only — never expose webhook over HTTP.
|
||||
- [ ] Verify `OPENCLAW_CONNECTOR_LINE_CHANNEL_SECRET` is set (signature verification).
|
||||
- [ ] Consider using a randomized webhook path (e.g., `/line/webhook-abc123def`).
|
||||
|
||||
### 3. Rate Limiting
|
||||
|
||||
- [ ] Review default limits: 10 req/min per user, 30 req/min per channel.
|
||||
- [ ] Adjust if needed: `OPENCLAW_CONNECTOR_RATE_LIMIT_USER_RPM`, `_CHANNEL_RPM`.
|
||||
|
||||
### 4. Payload Limits
|
||||
|
||||
- [ ] Default max command length: 4096 chars.
|
||||
- [ ] Adjust if needed: `OPENCLAW_CONNECTOR_MAX_COMMAND_LENGTH`.
|
||||
|
||||
### 5. Server API Access
|
||||
|
||||
- [ ] Keep ComfyUI on localhost (`--listen 127.0.0.1`) unless LAN access required.
|
||||
- [ ] If exposing to LAN/Internet: set `OPENCLAW_ADMIN_TOKEN` environment variable.
|
||||
- [ ] Never expose admin endpoints without token.
|
||||
|
||||
### 6. Debug Mode
|
||||
|
||||
- [ ] `OPENCLAW_CONNECTOR_DEBUG=1` logs sensitive data — **disable in production**.
|
||||
- [ ] Ensure no debug flags are set in production environment.
|
||||
|
||||
### 7. Tunnel / Reverse Proxy
|
||||
|
||||
- [ ] Use ngrok, Cloudflare Tunnel, or similar with TLS termination.
|
||||
- [ ] Restrict access by IP if possible.
|
||||
- [ ] Consider authentication layer (e.g., Cloudflare Access).
|
||||
|
||||
## ⚠️ Security Defaults
|
||||
|
||||
| Feature | Default | Effect |
|
||||
|---------|---------|--------|
|
||||
| Empty allowlists | Untrusted | All `/run` requires approval |
|
||||
| No admin users | Limited | Admin commands unavailable |
|
||||
| Rate limiting | Enabled | 10 req/min/user, 30 req/min/channel |
|
||||
| Debug mode | Disabled | No sensitive logging |
|
||||
| Replay protection | Enabled | LINE webhooks reject replays >5min old |
|
||||
|
||||
## 📞 Support
|
||||
|
||||
If you suspect a security issue, contact the maintainers via GitHub Issues (private for sensitive reports).
|
||||
@@ -67,6 +67,28 @@ async def submit_prompt(
|
||||
if extra_data:
|
||||
payload["extra_data"] = extra_data
|
||||
|
||||
# NOTE: Debug-only full payload logging for troubleshooting mismatched outputs.
|
||||
# Enable with OPENCLAW_DEBUG_PROMPT_PAYLOAD=1. This may include sensitive prompt content.
|
||||
if os.environ.get("OPENCLAW_DEBUG_PROMPT_PAYLOAD", "").strip().lower() in (
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
"on",
|
||||
):
|
||||
try:
|
||||
logger.warning(
|
||||
"DEBUG prompt payload (trace=%s source=%s): %s",
|
||||
trace_id,
|
||||
source,
|
||||
json.dumps(payload, ensure_ascii=False),
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"DEBUG prompt payload (trace=%s source=%s): <failed to serialize>",
|
||||
trace_id,
|
||||
source,
|
||||
)
|
||||
|
||||
# R33: Acquire concurrency budget
|
||||
limiter = get_limiter()
|
||||
async with limiter.acquire(source=source, trace_id=trace_id):
|
||||
|
||||
@@ -271,6 +271,11 @@ def validate_config_update(updates: Dict[str, Any]) -> Tuple[Dict[str, Any], lis
|
||||
if not isinstance(val, str):
|
||||
errors.append("base_url must be a string")
|
||||
continue
|
||||
# NOTE: Allow empty base_url (use provider default).
|
||||
# Without this, UI saves can fail with "Invalid scheme" on blank base_url.
|
||||
if val.strip() == "":
|
||||
sanitized[key] = ""
|
||||
continue
|
||||
# S16: Base URL policy
|
||||
|
||||
# 1. Allow if it matches the *default* base_url for the selected provider
|
||||
|
||||
@@ -217,6 +217,14 @@ Expected response:
|
||||
|
||||
### Use `/run` from chat
|
||||
|
||||
**Free-text prompt support (no `key=value` needed):**
|
||||
- `/run <template_id> <free text> seed=-1`
|
||||
- Connector maps free-text to a prompt key:
|
||||
- If `manifest.json` `allowed_inputs` has exactly one key → it uses that.
|
||||
- Otherwise prefers: `positive_prompt` → `prompt` → `text` → `positive` → `caption`.
|
||||
- If none match, defaults to `positive_prompt`.
|
||||
- Ensure the template uses the same placeholder (e.g., `"text": "{{positive_prompt}}"`).
|
||||
|
||||
Once the template appears in `/openclaw/templates`, you can run it via chat:
|
||||
- Run immediately:
|
||||
- `/run your_template_id positive_prompt="a cat" seed=123`
|
||||
@@ -225,6 +233,40 @@ Once the template appears in `/openclaw/templates`, you can run it via chat:
|
||||
|
||||
Unused keys have no effect unless the workflow contains a matching `{{key}}` placeholder.
|
||||
|
||||
## Admin Token & UI Usage (SOP)
|
||||
|
||||
**Key rule:** `OPENCLAW_ADMIN_TOKEN` is a **server-side environment variable**.
|
||||
The UI can **use** an Admin Token for authenticated requests, but **cannot set or persist** the server token.
|
||||
|
||||
### Recommended setup (local only)
|
||||
1) **Set server token (env)**
|
||||
```powershell
|
||||
$env:OPENCLAW_ADMIN_TOKEN="your_admin_token_here"
|
||||
```
|
||||
2) **Restart ComfyUI**
|
||||
3) **Enter the same token in the Settings UI**
|
||||
- This only stores it in the browser session for API calls.
|
||||
|
||||
### Windows CMD (per-session)
|
||||
```cmd
|
||||
set OPENCLAW_ADMIN_TOKEN=your_admin_token_here
|
||||
set OPENCLAW_LLM_API_KEY=your_api_key_here
|
||||
set OPENCLAW_LLM_PROVIDER=gemini
|
||||
```
|
||||
|
||||
### Windows CMD (persistent, user-level)
|
||||
```cmd
|
||||
setx OPENCLAW_ADMIN_TOKEN "your_admin_token_here"
|
||||
setx OPENCLAW_LLM_API_KEY "your_api_key_here"
|
||||
setx OPENCLAW_LLM_PROVIDER "gemini"
|
||||
```
|
||||
|
||||
> After `setx`, open a **new** terminal session before launching ComfyUI.
|
||||
|
||||
### Security Notes
|
||||
- Do **not** expose ComfyUI to the internet with UI-only tokens.
|
||||
- Admin token must remain server-side and protected by OS/environment.
|
||||
|
||||
## WSL / Restricted Environments
|
||||
If `pre-commit` fails due to cache permissions, run with a writable cache directory:
|
||||
```bash
|
||||
|
||||
@@ -171,27 +171,29 @@ class TestLLMClient(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
client = MagicMock()
|
||||
client.get_openclaw_config = AsyncMock(
|
||||
return_value={"ok": True, "data": {"provider": "openai", "model": "gpt-4o", "api_key_configured": True}}
|
||||
return_value={"ok": True, "data": {"config": {"provider": "openai", "model": "gpt-4o"}}}
|
||||
)
|
||||
client.chat_llm = AsyncMock(return_value={"ok": False, "error": "llm_request_failed"})
|
||||
|
||||
llm = LLMClient(client)
|
||||
|
||||
# Mock the services import to fail, then fallback to fail too
|
||||
with patch.dict("sys.modules", {"services": None, "services.providers": None}):
|
||||
with patch("connector.llm_client.LLMClient._fallback_chat", new_callable=AsyncMock) as mock_fallback:
|
||||
mock_fallback.return_value = "[LLM Error] Request failed."
|
||||
result = await llm.chat("system", "user message")
|
||||
# Should have attempted fallback
|
||||
# Note: actual behavior depends on import structure
|
||||
result = await llm.chat("system", "user message")
|
||||
self.assertIn("LLM Error", result)
|
||||
|
||||
async def test_no_prompt_logging(self):
|
||||
"""Verify user prompts are not logged."""
|
||||
import logging
|
||||
"""Verify user prompts are not logged (by design, not by level)."""
|
||||
from connector.llm_client import LLMClient
|
||||
|
||||
# The module should have WARNING level to avoid logging user prompts
|
||||
llm_logger = logging.getLogger("connector.llm_client")
|
||||
self.assertGreaterEqual(llm_logger.level, logging.WARNING)
|
||||
client = MagicMock()
|
||||
client.get_openclaw_config = AsyncMock(
|
||||
return_value={"ok": True, "data": {"config": {"provider": "openai"}}}
|
||||
)
|
||||
client.chat_llm = AsyncMock(return_value={"ok": True, "text": "response"})
|
||||
|
||||
llm = LLMClient(client)
|
||||
|
||||
# Verify the LLM client docstring mentions no prompt logging
|
||||
# This is a design verification, not a runtime check
|
||||
self.assertIn("Never logs user prompt content", LLMClient.__doc__ or "")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -0,0 +1,151 @@
|
||||
"""
|
||||
Unit tests for F32 Security Hardening.
|
||||
Tests rate limiting, command length limits, and replay protection.
|
||||
"""
|
||||
|
||||
import time
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from connector.config import ConnectorConfig
|
||||
from connector.contract import CommandRequest
|
||||
from connector.rate_limiter import RateLimiter, TokenBucket
|
||||
from connector.router import CommandRouter
|
||||
|
||||
|
||||
def make_request(sender_id: str, text: str, channel_id: str = "123") -> CommandRequest:
|
||||
"""Helper to create CommandRequest."""
|
||||
return CommandRequest(
|
||||
platform="telegram",
|
||||
channel_id=channel_id,
|
||||
sender_id=sender_id,
|
||||
username="testuser",
|
||||
message_id="msg-001",
|
||||
text=text,
|
||||
timestamp=time.time(),
|
||||
)
|
||||
|
||||
|
||||
class TestRateLimiter(unittest.TestCase):
|
||||
"""Test rate limiter token bucket."""
|
||||
|
||||
def test_token_bucket_allows_initial(self):
|
||||
"""Should allow requests up to capacity."""
|
||||
bucket = TokenBucket(capacity=5.0, refill_rate=1.0, tokens=5.0)
|
||||
for _ in range(5):
|
||||
self.assertTrue(bucket.consume())
|
||||
# 6th should fail
|
||||
self.assertFalse(bucket.consume())
|
||||
|
||||
def test_token_bucket_refills(self):
|
||||
"""Should refill tokens over time."""
|
||||
bucket = TokenBucket(capacity=5.0, refill_rate=5.0, tokens=0.0)
|
||||
bucket.last_refill = time.time() - 1.0 # 1 second ago
|
||||
# Should have refilled 5 tokens
|
||||
self.assertTrue(bucket.consume())
|
||||
|
||||
def test_rate_limiter_per_user(self):
|
||||
"""Should track per-user limits."""
|
||||
limiter = RateLimiter(user_rpm=2, channel_rpm=100)
|
||||
# User 1 can make 2 requests
|
||||
self.assertTrue(limiter.is_allowed("user1", "channel1"))
|
||||
self.assertTrue(limiter.is_allowed("user1", "channel1"))
|
||||
# User 1 blocked
|
||||
self.assertFalse(limiter.is_allowed("user1", "channel1"))
|
||||
# User 2 can still make requests
|
||||
self.assertTrue(limiter.is_allowed("user2", "channel1"))
|
||||
|
||||
def test_rate_limiter_per_channel(self):
|
||||
"""Should track per-channel limits."""
|
||||
limiter = RateLimiter(user_rpm=100, channel_rpm=2)
|
||||
# Channel 1 can handle 2 requests
|
||||
self.assertTrue(limiter.is_allowed("user1", "channel1"))
|
||||
self.assertTrue(limiter.is_allowed("user2", "channel1"))
|
||||
# Channel 1 blocked
|
||||
self.assertFalse(limiter.is_allowed("user3", "channel1"))
|
||||
# Channel 2 still works
|
||||
self.assertTrue(limiter.is_allowed("user1", "channel2"))
|
||||
|
||||
|
||||
class TestRouterSecurityChecks(unittest.IsolatedAsyncioTestCase):
|
||||
"""Test router security checks."""
|
||||
|
||||
def setUp(self):
|
||||
self.config = ConnectorConfig()
|
||||
self.config.max_command_length = 100
|
||||
self.config.rate_limit_user_rpm = 5
|
||||
|
||||
async def test_command_length_rejected(self):
|
||||
"""Should reject commands exceeding max length."""
|
||||
client = MagicMock()
|
||||
client.get_health = AsyncMock(return_value={"ok": True})
|
||||
client.get_prompt_queue = AsyncMock(return_value={"ok": True})
|
||||
router = CommandRouter(self.config, client)
|
||||
|
||||
long_command = "/status " + "x" * 200
|
||||
req = make_request("user1", long_command)
|
||||
resp = await router.handle(req)
|
||||
self.assertIn("too long", resp.text.lower())
|
||||
|
||||
async def test_rate_limit_response(self):
|
||||
"""Should return rate limit message when exceeded."""
|
||||
client = MagicMock()
|
||||
router = CommandRouter(self.config, client)
|
||||
|
||||
# Exhaust rate limit
|
||||
for _ in range(6):
|
||||
req = make_request("user1", "/status")
|
||||
resp = await router.handle(req)
|
||||
|
||||
# Last response should be rate limit
|
||||
self.assertIn("rate limit", resp.text.lower())
|
||||
|
||||
|
||||
class TestLineReplayProtection(unittest.TestCase):
|
||||
"""Test LINE webhook replay protection."""
|
||||
|
||||
def test_stale_timestamp_rejected(self):
|
||||
"""Should reject events with timestamps > 5 min old."""
|
||||
from connector.platforms.line_webhook import LINEWebhookServer
|
||||
|
||||
config = ConnectorConfig()
|
||||
config.line_channel_secret = "test_secret"
|
||||
server = LINEWebhookServer(config, MagicMock())
|
||||
|
||||
# Event from 10 minutes ago
|
||||
old_ts = int((time.time() - 600) * 1000)
|
||||
body = f'{{"events": [{{"timestamp": {old_ts}, "webhookEventId": "evt1"}}]}}'
|
||||
self.assertFalse(server._check_replay_protection(body))
|
||||
|
||||
def test_fresh_timestamp_accepted(self):
|
||||
"""Should accept events with recent timestamps."""
|
||||
from connector.platforms.line_webhook import LINEWebhookServer
|
||||
|
||||
config = ConnectorConfig()
|
||||
config.line_channel_secret = "test_secret"
|
||||
server = LINEWebhookServer(config, MagicMock())
|
||||
|
||||
# Event from 1 minute ago
|
||||
recent_ts = int((time.time() - 60) * 1000)
|
||||
body = f'{{"events": [{{"timestamp": {recent_ts}, "webhookEventId": "evt2"}}]}}'
|
||||
self.assertTrue(server._check_replay_protection(body))
|
||||
|
||||
def test_duplicate_nonce_rejected(self):
|
||||
"""Should reject duplicate webhook event IDs."""
|
||||
from connector.platforms.line_webhook import LINEWebhookServer
|
||||
|
||||
config = ConnectorConfig()
|
||||
config.line_channel_secret = "test_secret"
|
||||
server = LINEWebhookServer(config, MagicMock())
|
||||
|
||||
recent_ts = int((time.time() - 10) * 1000)
|
||||
body = f'{{"events": [{{"timestamp": {recent_ts}, "webhookEventId": "dup_evt"}}]}}'
|
||||
|
||||
# First request accepted
|
||||
self.assertTrue(server._check_replay_protection(body))
|
||||
# Second request (replay) rejected
|
||||
self.assertFalse(server._check_replay_protection(body))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user