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 作为错误信息
This commit is contained in:
@@ -438,6 +438,14 @@ async fn cmd_listen(
|
|||||||
.await;
|
.await;
|
||||||
});
|
});
|
||||||
info!("定时任务调度器已启动");
|
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);
|
info!("开始监听消息 (llm={}, echo={})...", enable_llm, echo);
|
||||||
|
|||||||
+12
-1
@@ -70,6 +70,7 @@ struct DueTask {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn fetch_due_tasks(pool: &PgPool) -> Result<Vec<DueTask>, String> {
|
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)>(
|
let rows = sqlx::query_as::<_, (uuid::Uuid, String, String, String, String, String)>(
|
||||||
r#"
|
r#"
|
||||||
SELECT id, name, user_id, command, shell, cwd
|
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
|
FOR UPDATE SKIP LOCKED
|
||||||
"#,
|
"#,
|
||||||
)
|
)
|
||||||
.fetch_all(pool)
|
.fetch_all(&mut *tx)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| format!("查询到期任务失败: {}", e))?;
|
.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
|
Ok(rows
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|(id, name, user_id, command, shell, cwd)| DueTask {
|
.map(|(id, name, user_id, command, shell, cwd)| DueTask {
|
||||||
|
|||||||
+61
-31
@@ -89,61 +89,91 @@ impl ApprovalManager {
|
|||||||
/// 处理用户回复(由消息循环调用)
|
/// 处理用户回复(由消息循环调用)
|
||||||
/// 匹配成功时返回 Some(技能名),匹配失败返回 None
|
/// 匹配成功时返回 Some(技能名),匹配失败返回 None
|
||||||
pub async fn handle_reply(&self, user_id: &str, reply: &str) -> Option<String> {
|
pub async fn handle_reply(&self, user_id: &str, reply: &str) -> Option<String> {
|
||||||
|
let (result, name) = {
|
||||||
let mut map = self.pending.lock().await;
|
let mut map = self.pending.lock().await;
|
||||||
let entries = map.get_mut(user_id)?;
|
let entries = map.get_mut(user_id)?;
|
||||||
if entries.is_empty() {
|
if entries.is_empty() { return None; }
|
||||||
return None;
|
|
||||||
}
|
|
||||||
|
|
||||||
let reply_hash = format!("{:x}", Sha256::digest(reply.as_bytes()));
|
let reply_hash = format!("{:x}", Sha256::digest(reply.as_bytes()));
|
||||||
let reply_trimmed = reply.trim();
|
let reply_trimmed = reply.trim();
|
||||||
|
|
||||||
// 检查取消
|
|
||||||
if reply_trimmed == "0" || reply_trimmed == "取消" {
|
if reply_trimmed == "0" || reply_trimmed == "取消" {
|
||||||
let entry = entries.remove(0);
|
let entry = entries.remove(0);
|
||||||
let name = entry.skill_name.clone();
|
|
||||||
let _ = entry.sender.send(ApprovalDecision::Rejected);
|
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 pos = entries.iter().position(|e| e.code_hash == reply_hash);
|
|
||||||
if let Some(idx) = pos {
|
|
||||||
let entry = entries.remove(idx);
|
let entry = entries.remove(idx);
|
||||||
let name = entry.skill_name.clone();
|
|
||||||
let _ = entry.sender.send(ApprovalDecision::Approved);
|
let _ = entry.sender.send(ApprovalDecision::Approved);
|
||||||
return Some(name);
|
(ApprovalDecision::Approved, entry.skill_name.clone())
|
||||||
}
|
} else if let Some(entry) = entries.first_mut() {
|
||||||
|
|
||||||
// 确认码错误
|
|
||||||
if let Some(entry) = entries.first_mut() {
|
|
||||||
entry.attempts_left -= 1;
|
entry.attempts_left -= 1;
|
||||||
if entry.attempts_left <= 0 {
|
if entry.attempts_left <= 0 {
|
||||||
let entry = entries.remove(0);
|
let entry = entries.remove(0);
|
||||||
let name = entry.skill_name.clone();
|
let n = entry.skill_name.clone();
|
||||||
let _ = entry.sender.send(ApprovalDecision::Rejected);
|
let _ = entry.sender.send(ApprovalDecision::Rejected);
|
||||||
return Some(name);
|
(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 {
|
if let Some(ref pool) = self.pool {
|
||||||
|
let status = match result {
|
||||||
|
ApprovalDecision::Approved => "approved",
|
||||||
|
_ => "rejected",
|
||||||
|
};
|
||||||
let _ = sqlx::query(
|
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())
|
.execute(pool.as_ref())
|
||||||
.await;
|
.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 {
|
} else {
|
||||||
let msg = if !stderr.is_empty() {
|
let msg = if !stderr.is_empty() {
|
||||||
format!("退出码 {}: {}", output.status.code().unwrap_or(-1), stderr)
|
format!("退出码 {}: {}", output.status.code().unwrap_or(-1), stderr)
|
||||||
|
} else if !stdout.is_empty() {
|
||||||
|
format!("退出码 {}: {}", output.status.code().unwrap_or(-1), stdout)
|
||||||
} else {
|
} else {
|
||||||
format!("退出码 {}", output.status.code().unwrap_or(-1))
|
format!("退出码 {}", output.status.code().unwrap_or(-1))
|
||||||
};
|
};
|
||||||
|
|||||||
Reference in New Issue
Block a user