use rand::Rng; use sha2::{Digest, Sha256}; use sqlx::PgPool; use std::collections::HashMap; use std::sync::Arc; use tokio::sync::{Mutex, oneshot}; /// 审批决策 #[derive(Debug, Clone, PartialEq)] pub enum ApprovalDecision { Approved, Rejected, Expired, } /// 内存中待审批的条目 struct PendingEntry { code_hash: String, skill_name: String, attempts_left: i32, /// 通知等待方 sender: oneshot::Sender, } /// 审批管理器 pub struct ApprovalManager { pool: Option>, /// user_id → 该用户的待审批队列 pending: Arc>>>, } impl ApprovalManager { pub fn new(pool: Option>) -> 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), String> { let code: u32 = rand::thread_rng().gen_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(¶ms) .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, sender: tx, }); Ok((code_str, rx)) } /// 处理用户回复(由消息循环调用) /// 匹配成功时返回 Some(技能名),匹配失败返回 None pub async fn handle_reply(&self, user_id: &str, reply: &str) -> Option { 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(); if reply_trimmed == "0" || reply_trimmed == "取消" { let entry = entries.remove(0); let _ = entry.sender.send(ApprovalDecision::Rejected); (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; } }; // 更新数据库状态 if let Some(ref pool) = self.pool { let status = match result { ApprovalDecision::Approved => "approved", _ => "rejected", }; let _ = sqlx::query( "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()); } /// 获取某个用户的待审批条数 pub async fn pending_count(&self, user_id: &str) -> usize { let map = self.pending.lock().await; map.get(user_id).map(|v| v.len()).unwrap_or(0) } }