fix: 5 个安全/正确性问题(第二轮)
1. (High) user_id NOT NULL 导致 add 任务失败 → ALTER TABLE 加 DEFAULT '', DROP NOT NULL 2. (High) 调度器任务卡住阻塞后续 → timeout(120s) + kill_on_drop 3. (Medium) 多实例重复执行 → SELECT ... FOR UPDATE SKIP LOCKED 4. (Medium) SQL 注入(字符串拼接) → psql -v 参数化 (:name, :cmd, :interval) 5. (Medium) 非工具对话 usage 未记录 → consume() 改为 &mut self,消费后再读 usage
This commit is contained in:
@@ -0,0 +1,3 @@
|
|||||||
|
-- scheduled_tasks.user_id 默认值修复
|
||||||
|
ALTER TABLE scheduled_tasks ALTER COLUMN user_id SET DEFAULT '';
|
||||||
|
ALTER TABLE scheduled_tasks ALTER COLUMN user_id DROP NOT NULL;
|
||||||
@@ -1,5 +1,5 @@
|
|||||||
#!/bin/bash
|
#!/bin/bash
|
||||||
# 管理定时任务 → JSON 输出
|
# 管理定时任务 → JSON 输出(参数化查询)
|
||||||
read -r input
|
read -r input
|
||||||
ACTION=$(echo "$input" | python3 -c "import sys,json; print(json.load(sys.stdin).get('action',''))" 2>/dev/null)
|
ACTION=$(echo "$input" | python3 -c "import sys,json; print(json.load(sys.stdin).get('action',''))" 2>/dev/null)
|
||||||
NAME=$(echo "$input" | python3 -c "import sys,json; print(json.load(sys.stdin).get('name',''))" 2>/dev/null)
|
NAME=$(echo "$input" | python3 -c "import sys,json; print(json.load(sys.stdin).get('name',''))" 2>/dev/null)
|
||||||
@@ -14,22 +14,23 @@ DB_URL="${DATABASE_URL}"
|
|||||||
case "$ACTION" in
|
case "$ACTION" in
|
||||||
add)
|
add)
|
||||||
[ -z "$CMD" ] && echo '{"error":"add 需要 command 参数"}' && exit 0
|
[ -z "$CMD" ] && echo '{"error":"add 需要 command 参数"}' && exit 0
|
||||||
psql "$DB_URL" -c "
|
psql "$DB_URL" -v name="$NAME" -v desc="$DESC" -v cmd="$CMD" -v interval="$INTERVAL" << 'SQL'
|
||||||
INSERT INTO scheduled_tasks (id, name, description, command, interval_seconds, enabled, next_run_at)
|
INSERT INTO scheduled_tasks (id, name, description, command, interval_seconds, enabled, next_run_at)
|
||||||
VALUES (gen_random_uuid(), '$NAME', '$DESC', '$CMD', $INTERVAL, true, NOW() + interval '$INTERVAL seconds')
|
VALUES (gen_random_uuid(), :'name', :'desc', :'cmd', :interval, true, NOW() + interval ':interval seconds')
|
||||||
ON CONFLICT (name) DO UPDATE SET command=EXCLUDED.command, interval_seconds=EXCLUDED.interval_seconds, enabled=true, description=EXCLUDED.description
|
ON CONFLICT (name) DO UPDATE SET command=EXCLUDED.command, interval_seconds=EXCLUDED.interval_seconds, enabled=true, description=EXCLUDED.description
|
||||||
" 2>/dev/null && echo '{"ok":true,"action":"add","name":"'"$NAME"'"}' || echo '{"error":"添加失败"}'
|
SQL
|
||||||
|
echo '{"ok":true,"action":"add","name":"'"$NAME"'"}'
|
||||||
;;
|
;;
|
||||||
delete)
|
delete)
|
||||||
psql "$DB_URL" -c "DELETE FROM scheduled_tasks WHERE name = '$NAME'" 2>/dev/null
|
psql "$DB_URL" -v name="$NAME" -c "DELETE FROM scheduled_tasks WHERE name = :'name'" 2>/dev/null
|
||||||
echo '{"ok":true,"action":"delete","name":"'"$NAME"'"}'
|
echo '{"ok":true,"action":"delete","name":"'"$NAME"'"}'
|
||||||
;;
|
;;
|
||||||
enable)
|
enable)
|
||||||
psql "$DB_URL" -c "UPDATE scheduled_tasks SET enabled = true WHERE name = '$NAME'" 2>/dev/null
|
psql "$DB_URL" -v name="$NAME" -c "UPDATE scheduled_tasks SET enabled = true WHERE name = :'name'" 2>/dev/null
|
||||||
echo '{"ok":true,"action":"enable","name":"'"$NAME"'"}'
|
echo '{"ok":true,"action":"enable","name":"'"$NAME"'"}'
|
||||||
;;
|
;;
|
||||||
disable)
|
disable)
|
||||||
psql "$DB_URL" -c "UPDATE scheduled_tasks SET enabled = false WHERE name = '$NAME'" 2>/dev/null
|
psql "$DB_URL" -v name="$NAME" -c "UPDATE scheduled_tasks SET enabled = false WHERE name = :'name'" 2>/dev/null
|
||||||
echo '{"ok":true,"action":"disable","name":"'"$NAME"'"}'
|
echo '{"ok":true,"action":"disable","name":"'"$NAME"'"}'
|
||||||
;;
|
;;
|
||||||
*)
|
*)
|
||||||
|
|||||||
@@ -222,7 +222,8 @@ impl ChatHandle {
|
|||||||
chunk
|
chunk
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn consume(mut self) -> Result<String, String> {
|
/// 消费所有流式块,返回完整文本
|
||||||
|
pub async fn consume(&mut self) -> Result<String, String> {
|
||||||
while let Some(chunk) = self.next_chunk().await {
|
while let Some(chunk) = self.next_chunk().await {
|
||||||
if let StreamChunk::Error(e) = chunk { return Err(e); }
|
if let StreamChunk::Error(e) = chunk { return Err(e); }
|
||||||
}
|
}
|
||||||
|
|||||||
+4
-2
@@ -598,7 +598,9 @@ async fn generate_reply(
|
|||||||
user_text: String,
|
user_text: String,
|
||||||
db: &Option<Arc<Database>>,
|
db: &Option<Arc<Database>>,
|
||||||
) -> Result<String, String> {
|
) -> Result<String, String> {
|
||||||
let handle = conv.chat(user_text).await?;
|
let mut handle = conv.chat(user_text).await?;
|
||||||
|
// 先消费流(usage 在 Done chunk 中才可用)
|
||||||
|
let reply = handle.consume().await?;
|
||||||
if let Some(u) = handle.get_usage() {
|
if let Some(u) = handle.get_usage() {
|
||||||
info!(
|
info!(
|
||||||
"📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}",
|
"📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}",
|
||||||
@@ -610,7 +612,7 @@ async fn generate_reply(
|
|||||||
);
|
);
|
||||||
store_usage(db, "", conv.model(), conv.provider_name(), u).await;
|
store_usage(db, "", conv.model(), conv.provider_name(), u).await;
|
||||||
}
|
}
|
||||||
handle.consume().await
|
Ok(reply)
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn generate_reply_with_tools(
|
async fn generate_reply_with_tools(
|
||||||
|
|||||||
+16
-11
@@ -76,7 +76,8 @@ async fn fetch_due_tasks(pool: &PgPool) -> Result<Vec<DueTask>, String> {
|
|||||||
FROM scheduled_tasks
|
FROM scheduled_tasks
|
||||||
WHERE enabled = true AND next_run_at <= NOW()
|
WHERE enabled = true AND next_run_at <= NOW()
|
||||||
ORDER BY next_run_at ASC
|
ORDER BY next_run_at ASC
|
||||||
LIMIT 10
|
LIMIT 5
|
||||||
|
FOR UPDATE SKIP LOCKED
|
||||||
"#,
|
"#,
|
||||||
)
|
)
|
||||||
.fetch_all(pool)
|
.fetch_all(pool)
|
||||||
@@ -110,17 +111,20 @@ async fn execute_task(task: &DueTask) -> String {
|
|||||||
task.cwd.clone()
|
task.cwd.clone()
|
||||||
};
|
};
|
||||||
|
|
||||||
let result = tokio::process::Command::new(shell)
|
let result = tokio::time::timeout(
|
||||||
.arg("-c")
|
std::time::Duration::from_secs(120),
|
||||||
.arg(&task.command)
|
tokio::process::Command::new(shell)
|
||||||
.current_dir(&cwd)
|
.arg("-c")
|
||||||
.stdout(std::process::Stdio::piped())
|
.arg(&task.command)
|
||||||
.stderr(std::process::Stdio::piped())
|
.current_dir(&cwd)
|
||||||
.output()
|
.stdout(std::process::Stdio::piped())
|
||||||
.await;
|
.stderr(std::process::Stdio::piped())
|
||||||
|
.kill_on_drop(true)
|
||||||
|
.output()
|
||||||
|
).await;
|
||||||
|
|
||||||
match result {
|
match result {
|
||||||
Ok(output) => {
|
Ok(Ok(output)) => {
|
||||||
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
|
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
|
||||||
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
|
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
|
||||||
|
|
||||||
@@ -138,7 +142,8 @@ async fn execute_task(task: &DueTask) -> String {
|
|||||||
)
|
)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
Err(e) => format!("执行失败: {}", e),
|
Ok(Err(e)) => format!("进程错误: {}", e),
|
||||||
|
Err(_timeout) => "执行超时 (120s)".to_string(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user