iAs/src/tools/approval.rs

//! ## 审批管理器 —— 高风险工具的用户确认流程
//!
//! 当 LLM 调用高风险工具时，需要用户通过微信输入确认码来批准。
//!
//! ### 完整流程
//! 1. `create(user_id, skill_name)` → 生成 6 位随机确认码，SHA-256 哈希后存储
//! 2. 向用户发送 "⚠️ 操作确认\n\n技能：{name}\n确认码：{code}"
//! 3. 用户在微信聊天中输入确认码
//! 4. `handle_reply(user_id, reply)` → 验证确认码哈希
//!    - 匹配 → `Approved`
//!    - 不匹配 → 减少尝试次数（最多 3 次）
//!    - 输入 0 或"取消" → `Rejected`
//!    - 超时 → `Expired`（每 60 秒清理一次）
//!
//! ### 安全设计
//! - 确认码用 SHA-256 哈希后存储，原始码不持久化
//! - 每次生成随机 6 位数字（100000-999999）
//! - 最多 3 次尝试，超时 5 分钟
//! - 同时支持内存和数据库双写

use rand::Rng;
use sha2::{Digest, Sha256};
use sqlx::PgPool;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::{Mutex, oneshot};

/// ## 审批决策
///
/// * `Approved` — 用户输入了正确的确认码
/// * `Rejected` — 用户拒绝了（输入 0 或"取消"）
/// * `Expired` — 超时未确认（5 分钟有效）
#[derive(Debug, Clone, PartialEq)]
pub enum ApprovalDecision {
    Approved,
    Rejected,
    Expired,
}

/// 内存中待审批的条目
struct PendingEntry {
    code_hash: String,
    skill_name: String,
    attempts_left: i32,
    expires_at: Instant,
    /// 通知等待方
    sender: oneshot::Sender<ApprovalDecision>,
}

/// ## 审批管理器
///
/// 管理高风险工具的用户确认流程。
///
/// ### 流程
/// 1. `create(user_id, skill_name)` → 生成 6 位随机确认码，哈希后存储
/// 2. 向用户发送"⚠️ 操作确认\n\n技能：{name}\n确认码：{code}"
/// 3. 用户在微信聊天中输入确认码
/// 4. `handle_reply(user_id, reply)` → 验证确认码哈希
///    - 匹配 → Approved
///    - 不匹配 → 减少尝试次数（最多 3 次）
///    - 输入 0 或"取消" → Rejected
///    - 超时 → Expired（每 60 秒清理一次）
///
/// ### 存储
/// * 内存 — `HashMap<user_id, Vec<PendingEntry>>`（快速访问）
/// * 数据库 — `pending_approvals` 表（持久化审计）
pub struct ApprovalManager {
    pool: Option<Arc<PgPool>>,
    /// user_id → 该用户的待审批队列
    pending: Arc<Mutex<HashMap<String, Vec<PendingEntry>>>>,
}

impl ApprovalManager {
    pub fn new(pool: Option<Arc<PgPool>>) -> Self {
        Self {
            pool,
            pending: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// 创建审批请求
    /// 返回 (确认码明文, 接收器) — 调用方 await 接收器等待用户确认
    pub async fn create(
        &self,
        user_id: &str,
        skill_name: &str,
    ) -> Result<(String, oneshot::Receiver<ApprovalDecision>), String> {
        let code: u32 = rand::rng().random_range(100000..999999);
        let code_str = code.to_string();
        let code_hash = format!("{:x}", Sha256::digest(code_str.as_bytes()));

        let (tx, rx) = oneshot::channel();

        // 入库
        if let Some(ref pool) = self.pool {
            let id = uuid::Uuid::new_v4();
            let expires_at = chrono::Utc::now() + chrono::Duration::minutes(5);
            let params = serde_json::json!({"name": skill_name});

            let _ = sqlx::query(
                r#"
                INSERT INTO pending_approvals (id, expires_at, user_id, skill_name, params, code_hash, status)
                VALUES ($1, $2, $3, $4, $5, $6, 'pending')
                "#,
            )
            .bind(id)
            .bind(expires_at)
            .bind(user_id)
            .bind(skill_name)
            .bind(&params)
            .bind(&code_hash)
            .execute(pool.as_ref())
            .await;
        }

        // 注册到内存
        let mut map = self.pending.lock().await;
        map.entry(user_id.to_string())
            .or_default()
            .push(PendingEntry {
                code_hash,
                skill_name: skill_name.to_string(),
                attempts_left: 3,
                expires_at: Instant::now() + std::time::Duration::from_secs(300),
                sender: tx,
            });

        Ok((code_str, rx))
    }

    /// 测试专用：创建指定 TTL（秒）的审批请求
    #[cfg(test)]
    async fn create_with_ttl(
        &self,
        user_id: &str,
        skill_name: &str,
        ttl_secs: u64,
    ) -> Result<(String, oneshot::Receiver<ApprovalDecision>), String> {
        let code: u32 = rand::rng().random_range(100000..999999);
        let code_str = code.to_string();
        let code_hash = format!("{:x}", Sha256::digest(code_str.as_bytes()));
        let (tx, rx) = oneshot::channel();

        let mut map = self.pending.lock().await;
        map.entry(user_id.to_string())
            .or_default()
            .push(PendingEntry {
                code_hash,
                skill_name: skill_name.to_string(),
                attempts_left: 3,
                expires_at: Instant::now() + std::time::Duration::from_secs(ttl_secs),
                sender: tx,
            });

        Ok((code_str, rx))
    }

    /// 处理用户回复（由消息循环调用）
    /// 返回 (技能名, 决策)；匹配失败返回 None
    pub async fn handle_reply(
        &self,
        user_id: &str,
        reply: &str,
    ) -> Option<(String, ApprovalDecision)> {
        let (result, name, code_hash) = {
            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();

            if reply_trimmed == "0" || reply_trimmed == "取消" {
                let entry = entries.remove(0);
                let hash = entry.code_hash.clone();
                let _ = entry.sender.send(ApprovalDecision::Rejected);
                (ApprovalDecision::Rejected, entry.skill_name.clone(), hash)
            } else if let Some(idx) = entries.iter().position(|e| e.code_hash == reply_hash) {
                let entry = entries.remove(idx);
                let hash = entry.code_hash.clone();
                let _ = entry.sender.send(ApprovalDecision::Approved);
                (ApprovalDecision::Approved, entry.skill_name.clone(), hash)
            } 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 hash = entry.code_hash.clone();
                    let _ = entry.sender.send(ApprovalDecision::Rejected);
                    (ApprovalDecision::Rejected, n, hash)
                } else {
                    return None;
                }
            } else {
                return None;
            }
        };

        // 更新数据库状态
        if let Some(ref pool) = self.pool {
            let status = match result {
                ApprovalDecision::Approved => "approved",
                _ => "rejected",
            };
            // 按 user_id + 最近 pending 记录更新
            let _ = sqlx::query(
                "UPDATE pending_approvals SET status = $1, consumed_at = NOW() WHERE code_hash = $2 AND status = 'pending'",
            )
            .bind(status)
            .bind(&code_hash)
            .execute(pool.as_ref())
            .await;
        }

        Some((name, result))
    }

    /// 清理过期审批（按时间判定，通知等待方 + 更新 DB）
    ///
    /// 两阶段加锁：收集过期 code_hash → 更新 DB → 重新加锁按 code_hash 定位删除。
    /// 不依赖快照 index，避免并发回复导致 index 错位误删。
    pub async fn clean_expired(&self) {
        let now = Instant::now();

        // 第一阶段：持锁收集过期条目的 code_hash
        let mut expired_hashes: Vec<String> = Vec::new();
        {
            let map = self.pending.lock().await;
            for entries in map.values() {
                for e in entries {
                    if now > e.expires_at {
                        expired_hashes.push(e.code_hash.clone());
                    }
                }
            }
        }

        if expired_hashes.is_empty() {
            return;
        }

        // DB 更新（按 code_hash 精确匹配，无错位问题）
        if let Some(ref pool) = self.pool {
            for code_hash in &expired_hashes {
                let _ = sqlx::query(
                    "UPDATE pending_approvals SET status = 'expired' WHERE code_hash = $1 AND status = 'pending'",
                )
                .bind(code_hash)
                .execute(pool.as_ref())
                .await;
            }
        }

        // 第二阶段：重新持锁，按 code_hash 定位并移除
        // （即使两阶段间 handle_reply 增删了条目，按 hash 查找也能安全跳过）
        let mut map = self.pending.lock().await;
        for code_hash in &expired_hashes {
            for entries in map.values_mut() {
                if let Some(idx) = entries.iter().position(|e| &e.code_hash == code_hash) {
                    let entry = entries.remove(idx);
                    let _ = entry.sender.send(ApprovalDecision::Expired);
                }
            }
        }
        map.retain(|_, v| !v.is_empty());
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn clean_expired_only_removes_expired_entries() {
        // 无 DB，纯内存
        let mgr = ApprovalManager::new(None);

        // 两条审批：一条立即过期（ttl=0），一条 300s 不过期
        let (_expired_code, expired_rx) = mgr.create_with_ttl("u1", "tool_a", 0).await.unwrap();
        let (live_code, live_rx) = mgr.create_with_ttl("u1", "tool_b", 300).await.unwrap();

        // 立即过期的条目 expires_at 已成过去，sleep 确保时间推进
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;

        mgr.clean_expired().await;

        // 过期条目应收到 Expired
        assert_eq!(expired_rx.await.unwrap(), ApprovalDecision::Expired);
        // 未过期条目不应被误删：handle_reply 用正确确认码应能匹配
        let (name, decision) = mgr.handle_reply("u1", &live_code).await.unwrap();
        assert_eq!(name, "tool_b");
        assert_eq!(decision, ApprovalDecision::Approved);
        // live_rx 应收到 Approved（未被 clean_expired 误标记）
        assert_eq!(live_rx.await.unwrap(), ApprovalDecision::Approved);
    }

    #[tokio::test]
    async fn clean_expired_safe_when_no_pending() {
        // 空 pending 不应 panic
        let mgr = ApprovalManager::new(None);
        mgr.clean_expired().await;
        // 重复调用也安全
        mgr.clean_expired().await;
    }

    #[tokio::test]
    async fn clean_expired_skips_entries_already_removed_by_reply() {
        // 回归保护：clean_expired 第一阶段收集过期 hash 后、第二阶段删除前，
        // 若该条目已被 handle_reply 提前处理，clean_expired 应安全跳过，不误删其他条目。
        // （原 bug：按快照 index 删除会导致误删；现按 code_hash 删除应安全跳过）
        let mgr = ApprovalManager::new(None);

        // 两条同用户审批，都设为立即过期
        let (_code_a, rx_a) = mgr.create_with_ttl("u1", "tool_a", 0).await.unwrap();
        let (_code_b, rx_b) = mgr.create_with_ttl("u1", "tool_b", 0).await.unwrap();

        tokio::time::sleep(std::time::Duration::from_millis(10)).await;

        // 先用 handle_reply “0” 取消最早的条目（tool_a），模拟并发回复
        let (name, decision) = mgr.handle_reply("u1", "0").await.unwrap();
        assert_eq!(name, "tool_a");
        assert_eq!(decision, ApprovalDecision::Rejected);
        assert_eq!(rx_a.await.unwrap(), ApprovalDecision::Rejected);

        // 再 clean_expired：tool_a 已被 remove，按 hash 查找应跳过；
        // tool_b 仍在且过期，应被正确标记 Expired（不会被误当成 tool_a 跳过）
        mgr.clean_expired().await;
        assert_eq!(rx_b.await.unwrap(), ApprovalDecision::Expired);
    }
}