fix: 4 个安全/正确性问题

1. (High) 定时任务 FOR UPDATE 无事务保护 → BEGIN/COMMIT 事务包裹 fetch + update 2. (Medium) 审批 DB 状态未更新 → handle_reply 更新 status/consumed_at 3. (Medium) 过期审批未清理 → 每 30s tokio::spawn 调用 clean_expired 4. (Low) 非零退出丢弃 stdout → 退时 stderr 空则用 stdout 作为错误信息
2026-06-02 14:32:18 +08:00
parent 3f447668b8
commit a32593bc25
4 changed files with 95 additions and 44 deletions
@@ -438,6 +438,14 @@ async fn cmd_listen(
                .await;
        });
        info!("定时任务调度器已启动");
+        // 审批定期清理
+        let approval_clean = approval_manager.clone();
+        tokio::spawn(async move {
+            loop {
+                tokio::time::sleep(std::time::Duration::from_secs(30)).await;
+                approval_clean.clean_expired().await;
+            }
+        });
    }

    info!("开始监听消息 (llm={}, echo={})...", enable_llm, echo);
@@ -70,6 +70,7 @@ struct DueTask {
 }

 async fn fetch_due_tasks(pool: &PgPool) -> Result<Vec<DueTask>, String> {
+    let mut tx = pool.begin().await.map_err(|e| format!("开始事务失败: {}", e))?;
    let rows = sqlx::query_as::<_, (uuid::Uuid, String, String, String, String, String)>(
        r#"
        SELECT id, name, user_id, command, shell, cwd
@@ -80,10 +81,20 @@ async fn fetch_due_tasks(pool: &PgPool) -> Result<Vec<DueTask>, String> {
        FOR UPDATE SKIP LOCKED
        "#,
    )
-    .fetch_all(pool)
+    .fetch_all(&mut *tx)
    .await
    .map_err(|e| format!("查询到期任务失败: {}", e))?;

+    // 立即标记为已调度（防止其他实例重复获取）
+    for (id, _, _, _, _, _) in &rows {
+        sqlx::query("UPDATE scheduled_tasks SET next_run_at = NOW() + (interval_seconds * INTERVAL '1 second') WHERE id = $1")
+            .bind(id)
+            .execute(&mut *tx)
+            .await
+            .ok();
+    }
+    tx.commit().await.map_err(|e| format!("提交事务失败: {}", e))?;
+
    Ok(rows
        .into_iter()
        .map(|(id, name, user_id, command, shell, cwd)| DueTask {
@@ -89,61 +89,91 @@ impl ApprovalManager {
    /// 处理用户回复（由消息循环调用）
    /// 匹配成功时返回 Some(技能名)，匹配失败返回 None
    pub async fn handle_reply(&self, user_id: &str, reply: &str) -> Option<String> {
-        let mut map = self.pending.lock().await;
-        let entries = map.get_mut(user_id)?;
-        if entries.is_empty() {
-            return None;
-        }
+        let (result, name) = {
+            let mut map = self.pending.lock().await;
+            let entries = map.get_mut(user_id)?;
+            if entries.is_empty() { return None; }

-        let reply_hash = format!("{:x}", Sha256::digest(reply.as_bytes()));
-        let reply_trimmed = reply.trim();
+            let reply_hash = format!("{:x}", Sha256::digest(reply.as_bytes()));
+            let reply_trimmed = reply.trim();

-        // 检查取消
-        if reply_trimmed == "0" || reply_trimmed == "取消" {
-            let entry = entries.remove(0);
-            let name = entry.skill_name.clone();
-            let _ = entry.sender.send(ApprovalDecision::Rejected);
-            return Some(name);
-        }
-
-        // 检查确认码
-        let pos = entries.iter().position(|e| e.code_hash == reply_hash);
-        if let Some(idx) = pos {
-            let entry = entries.remove(idx);
-            let name = entry.skill_name.clone();
-            let _ = entry.sender.send(ApprovalDecision::Approved);
-            return Some(name);
-        }
-
-        // 确认码错误
-        if let Some(entry) = entries.first_mut() {
-            entry.attempts_left -= 1;
-            if entry.attempts_left <= 0 {
+            if reply_trimmed == "0" || reply_trimmed == "取消" {
                let entry = entries.remove(0);
-                let name = entry.skill_name.clone();
                let _ = entry.sender.send(ApprovalDecision::Rejected);
-                return Some(name);
+                (ApprovalDecision::Rejected, entry.skill_name.clone())
+            } else if let Some(idx) = entries.iter().position(|e| e.code_hash == reply_hash) {
+                let entry = entries.remove(idx);
+                let _ = entry.sender.send(ApprovalDecision::Approved);
+                (ApprovalDecision::Approved, entry.skill_name.clone())
+            } else if let Some(entry) = entries.first_mut() {
+                entry.attempts_left -= 1;
+                if entry.attempts_left <= 0 {
+                    let entry = entries.remove(0);
+                    let n = entry.skill_name.clone();
+                    let _ = entry.sender.send(ApprovalDecision::Rejected);
+                    (ApprovalDecision::Rejected, n)
+                } else {
+                    return None;
+                }
+            } else {
+                return None;
            }
-        }
-
-        None
-    }
-
-    /// 清理过期（由定时任务调用）
-    pub async fn clean_expired(&self) {
-        let mut map = self.pending.lock().await;
-        for entries in map.values_mut() {
-            entries.retain(|e| !e.sender.is_closed());
-        }
-        map.retain(|_, v| !v.is_empty());
+        };

+        // 更新数据库状态
        if let Some(ref pool) = self.pool {
+            let status = match result {
+                ApprovalDecision::Approved => "approved",
+                _ => "rejected",
+            };
            let _ = sqlx::query(
-                "UPDATE pending_approvals SET status = 'expired' WHERE status = 'pending' AND expires_at < NOW()",
+                "UPDATE pending_approvals SET status = $1, consumed_at = NOW() WHERE user_id = $2 AND code_hash = $3 AND status = 'pending'",
            )
+            .bind(status)
+            .bind(user_id)
+            .bind(&format!("{:x}", Sha256::digest(reply.as_bytes())))
            .execute(pool.as_ref())
            .await;
        }
+
+        Some(name)
+    }
+
+    /// 清理过期审批（通知等待方 + 更新 DB）
+    pub async fn clean_expired(&self) {
+        let mut expired = Vec::new();
+        {
+            let map = self.pending.lock().await;
+            for (uid, entries) in map.iter() {
+                for (i, e) in entries.iter().enumerate() {
+                    if e.sender.is_closed() {
+                        expired.push((uid.clone(), i, e.code_hash.clone()));
+                    }
+                }
+            }
+        }
+
+        if let Some(ref pool) = self.pool {
+            for (_uid, _i, _code_hash) in &expired {
+                let _ = sqlx::query(
+                    "UPDATE pending_approvals SET status = 'expired' WHERE code_hash = $1 AND status = 'pending'",
+                )
+                .bind(_code_hash)
+                .execute(pool.as_ref())
+                .await;
+            }
+        }
+
+        let mut map = self.pending.lock().await;
+        for (uid, i, _) in expired.iter().rev() {
+            if let Some(entries) = map.get_mut(uid) {
+                if *i < entries.len() {
+                    let entry = entries.remove(*i);
+                    let _ = entry.sender.send(ApprovalDecision::Expired);
+                }
+            }
+        }
+        map.retain(|_, v| !v.is_empty());
    }

    /// 获取某个用户的待审批条数
@@ -267,6 +267,8 @@ async fn run_skill(skill: &Skill, params: serde_json::Value, _ctx: &ExecutionCon
            } else {
                let msg = if !stderr.is_empty() {
                    format!("退出码 {}: {}", output.status.code().unwrap_or(-1), stderr)
+                } else if !stdout.is_empty() {
+                    format!("退出码 {}: {}", output.status.code().unwrap_or(-1), stdout)
                } else {
                    format!("退出码 {}", output.status.code().unwrap_or(-1))
                };