From 0883bb28ff430ac061a783961141aad2c838afde Mon Sep 17 00:00:00 2001
From: iven <iven_h@qq.com>
Date: Thu, 9 Apr 2026 17:24:36 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20validation=20hardening=20=E2=80=94=20age?=
 =?UTF-8?q?nt=20import=20prompt=20limit,=20relay=20retry=20tracking,=20hea?=
 =?UTF-8?q?rtbeat=20validation?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- agent_import: add system_prompt length validation (max 50K chars)
  to prevent excessive token consumption from imported configs
- relay retry_task: wrap JoinHandle to log abort on server shutdown
- device_heartbeat: validate device_id length (1-64 chars) matching
  register endpoint constraints
---
 crates/zclaw-saas/src/account/handlers.rs      |  5 +++++
 crates/zclaw-saas/src/relay/handlers.rs        |  9 ++++++++-
 desktop/src-tauri/src/kernel_commands/agent.rs | 12 ++++++++++++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/crates/zclaw-saas/src/account/handlers.rs b/crates/zclaw-saas/src/account/handlers.rs
index 329f84f..d36e25b 100644
--- a/crates/zclaw-saas/src/account/handlers.rs
+++ b/crates/zclaw-saas/src/account/handlers.rs
@@ -283,6 +283,11 @@ pub async fn device_heartbeat(
         .and_then(|v| v.as_str())
         .ok_or_else(|| SaasError::InvalidInput("缺少 device_id".into()))?;
 
+    // Validate device_id length (must match register endpoint constraints)
+    if device_id.is_empty() || device_id.len() > 64 {
+        return Err(SaasError::InvalidInput("device_id 长度必须在 1-64 个字符之间".into()));
+    }
+
     let now = chrono::Utc::now();
 
     // Also update platform/app_version if provided (supports client upgrades)
diff --git a/crates/zclaw-saas/src/relay/handlers.rs b/crates/zclaw-saas/src/relay/handlers.rs
index d6a0976..b9f06ef 100644
--- a/crates/zclaw-saas/src/relay/handlers.rs
+++ b/crates/zclaw-saas/src/relay/handlers.rs
@@ -550,7 +550,7 @@ pub async fn retry_task(
     // 异步执行重试 — 根据解析结果选择执行路径
     let db = state.db.clone();
     let task_id = id.clone();
-    tokio::spawn(async move {
+    let handle = tokio::spawn(async move {
         let result = match model_resolution {
             ModelResolution::Direct(ref candidate) => {
                 service::execute_relay(
@@ -575,6 +575,13 @@ pub async fn retry_task(
             Err(e) => tracing::warn!("Relay task {} 重试失败: {}", task_id, e),
         }
     });
+    // Detach with warning — if server shuts down mid-retry, the task is lost.
+    // The DB status is already reset to 'queued', so a future restart can pick it up.
+    tokio::spawn(async move {
+        if let Err(e) = handle.await {
+            tracing::warn!("Relay retry task aborted (server shutdown?): {}", e);
+        }
+    });
 
     // 异步派发操作日志
     state.dispatch_log_operation(
diff --git a/desktop/src-tauri/src/kernel_commands/agent.rs b/desktop/src-tauri/src/kernel_commands/agent.rs
index 243cc3a..abe3675 100644
--- a/desktop/src-tauri/src/kernel_commands/agent.rs
+++ b/desktop/src-tauri/src/kernel_commands/agent.rs
@@ -295,6 +295,18 @@ pub async fn agent_import(
     let mut config: AgentConfig = serde_json::from_str(&config_json)
         .map_err(|e| format!("Invalid agent config JSON: {}", e))?;
 
+    // Validate system_prompt length to prevent excessive token consumption
+    const MAX_SYSTEM_PROMPT_LEN: usize = 50_000;
+    if let Some(ref prompt) = config.system_prompt {
+        if prompt.len() > MAX_SYSTEM_PROMPT_LEN {
+            return Err(format!(
+                "system_prompt too long: {} chars (max {})",
+                prompt.len(),
+                MAX_SYSTEM_PROMPT_LEN
+            ));
+        }
+    }
+
     // Regenerate ID to avoid collisions
     config.id = AgentId::new();