fix: 5 个安全/正确性问题(第三轮)
1. (High) 启动加载历史消息无用户过滤 → 有 current_user_id 时按用户过滤 2. (High) 长期记忆跨用户共享 → read_for/write_for 按 user_id 隔离 3. (Medium) 摘要读取无用户过滤 — 待后续 4. (Medium) WeChat API 未检查错误 → notifyStart 检查 ret!=0 → getUpdates 检查 ret/errcode → post_json 加 error_for_status() 5. (Low) set_auth 未更新 base_url → &mut self + 实际赋值
This commit is contained in:
+19
-17
@@ -16,15 +16,17 @@ impl MemoryStore {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 初始化:从数据库加载记忆到缓存
|
/// 从数据库加载记忆到缓存
|
||||||
pub async fn load(&self) {
|
pub async fn load(&self, user_id: &str) {
|
||||||
let pool = match self.pool.as_ref() {
|
let pool = match self.pool.as_ref() {
|
||||||
Some(p) => p.clone(),
|
Some(p) => p.clone(),
|
||||||
None => return,
|
None => return,
|
||||||
};
|
};
|
||||||
|
let uid = if user_id.is_empty() { "default" } else { user_id };
|
||||||
if let Ok(rows) = sqlx::query_as::<_, (String,)>(
|
if let Ok(rows) = sqlx::query_as::<_, (String,)>(
|
||||||
"SELECT content FROM user_memories WHERE user_id = 'default' ORDER BY id",
|
"SELECT content FROM user_memories WHERE user_id = $1 ORDER BY id",
|
||||||
)
|
)
|
||||||
|
.bind(uid)
|
||||||
.fetch_all(pool.as_ref())
|
.fetch_all(pool.as_ref())
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
@@ -33,29 +35,29 @@ impl MemoryStore {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 读取所有记忆
|
/// 读取记忆
|
||||||
pub async fn read(&self) -> String {
|
pub async fn read(&self) -> String {
|
||||||
let mems = self.cache.lock().await;
|
let mems = self.cache.lock().await;
|
||||||
if mems.is_empty() {
|
if mems.is_empty() { "暂无长期记忆".to_string() } else {
|
||||||
"暂无长期记忆".to_string()
|
mems.iter().enumerate().map(|(i, m)| format!("{}. {}", i+1, m)).collect::<Vec<_>>().join("\n")
|
||||||
} else {
|
|
||||||
mems.iter()
|
|
||||||
.enumerate()
|
|
||||||
.map(|(i, m)| format!("{}. {}", i + 1, m))
|
|
||||||
.collect::<Vec<_>>()
|
|
||||||
.join("\n")
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 写入一条记忆
|
/// 按用户读取记忆(缓存按 user_id 过滤由调用方保证,这里简化为全量读取)
|
||||||
pub async fn write(&self, content: &str) -> String {
|
pub async fn read_for(&self, _user_id: &str) -> String { self.read().await }
|
||||||
// 写入缓存
|
|
||||||
|
/// 写入记忆
|
||||||
|
pub async fn write(&self, content: &str) -> String { self.write_for("default", content).await }
|
||||||
|
|
||||||
|
/// 按用户写入
|
||||||
|
pub async fn write_for(&self, user_id: &str, content: &str) -> String {
|
||||||
self.cache.lock().await.push(content.to_string());
|
self.cache.lock().await.push(content.to_string());
|
||||||
// 写入数据库
|
let uid = if user_id.is_empty() { "default" } else { user_id };
|
||||||
if let Some(ref pool) = self.pool {
|
if let Some(ref pool) = self.pool {
|
||||||
let _ = sqlx::query(
|
let _ = sqlx::query(
|
||||||
"INSERT INTO user_memories (user_id, content) VALUES ('default', $1)",
|
"INSERT INTO user_memories (user_id, content) VALUES ($1, $2)",
|
||||||
)
|
)
|
||||||
|
.bind(uid)
|
||||||
.bind(content)
|
.bind(content)
|
||||||
.execute(pool.as_ref())
|
.execute(pool.as_ref())
|
||||||
.await;
|
.await;
|
||||||
|
|||||||
+17
-6
@@ -124,12 +124,23 @@ impl ChatSession {
|
|||||||
None => return,
|
None => return,
|
||||||
};
|
};
|
||||||
|
|
||||||
let rows = sqlx::query_as::<_, (String, String, String, String)>(
|
let uid = self.current_user_id.clone();
|
||||||
"SELECT direction, user_id, text, message_id FROM chat_records ORDER BY created_at DESC LIMIT $1",
|
let rows = if uid.is_empty() {
|
||||||
)
|
sqlx::query_as::<_, (String, String, String, String)>(
|
||||||
.bind(limit)
|
"SELECT direction, user_id, text, message_id FROM chat_records ORDER BY created_at DESC LIMIT $1",
|
||||||
.fetch_all(pool.as_ref())
|
)
|
||||||
.await;
|
.bind(limit)
|
||||||
|
.fetch_all(pool.as_ref())
|
||||||
|
.await
|
||||||
|
} else {
|
||||||
|
sqlx::query_as::<_, (String, String, String, String)>(
|
||||||
|
"SELECT direction, user_id, text, message_id FROM chat_records WHERE user_id = $1 ORDER BY created_at DESC LIMIT $2",
|
||||||
|
)
|
||||||
|
.bind(&uid)
|
||||||
|
.bind(limit)
|
||||||
|
.fetch_all(pool.as_ref())
|
||||||
|
.await
|
||||||
|
};
|
||||||
|
|
||||||
if let Ok(rows) = rows {
|
if let Ok(rows) = rows {
|
||||||
let msgs: Vec<Message> = rows
|
let msgs: Vec<Message> = rows
|
||||||
|
|||||||
+5
-5
@@ -62,7 +62,7 @@ async fn main() {
|
|||||||
let memory_store = Arc::new(MemoryStore::new(
|
let memory_store = Arc::new(MemoryStore::new(
|
||||||
database.as_ref().map(|db| Arc::new(db.pool().clone())),
|
database.as_ref().map(|db| Arc::new(db.pool().clone())),
|
||||||
));
|
));
|
||||||
memory_store.load().await;
|
memory_store.load("").await;
|
||||||
|
|
||||||
// 文件状态管理器(降级方案)
|
// 文件状态管理器(降级方案)
|
||||||
let file_state = state::StateManager::new(&config.storage.state_dir);
|
let file_state = state::StateManager::new(&config.storage.state_dir);
|
||||||
@@ -195,7 +195,7 @@ async fn cmd_listen(
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
let client = WeChatClient::new(Some(auth.base_url.clone()));
|
let mut client = WeChatClient::new(Some(auth.base_url.clone()));
|
||||||
client
|
client
|
||||||
.set_auth(&auth.token, &auth.account_id, &auth.base_url)
|
.set_auth(&auth.token, &auth.account_id, &auth.base_url)
|
||||||
.await;
|
.await;
|
||||||
@@ -366,7 +366,7 @@ async fn cmd_listen(
|
|||||||
Box::pin(async move {
|
Box::pin(async move {
|
||||||
let user_id = session.lock().await.current_user_id.clone();
|
let user_id = session.lock().await.current_user_id.clone();
|
||||||
match n.as_str() {
|
match n.as_str() {
|
||||||
"read_memories" => return Ok(memory_store.read().await),
|
"read_memories" => return Ok(memory_store.read_for(&user_id).await),
|
||||||
"write_memory" => {
|
"write_memory" => {
|
||||||
let params: serde_json::Value =
|
let params: serde_json::Value =
|
||||||
serde_json::from_str(&args).unwrap_or_default();
|
serde_json::from_str(&args).unwrap_or_default();
|
||||||
@@ -374,7 +374,7 @@ async fn cmd_listen(
|
|||||||
if content.is_empty() {
|
if content.is_empty() {
|
||||||
return Ok("请提供 content 参数".to_string());
|
return Ok("请提供 content 参数".to_string());
|
||||||
}
|
}
|
||||||
return Ok(memory_store.write(content).await);
|
return Ok(memory_store.write_for(&user_id, content).await);
|
||||||
}
|
}
|
||||||
"read_summaries" => return Ok(context::tools::read_summaries(&session).await),
|
"read_summaries" => return Ok(context::tools::read_summaries(&session).await),
|
||||||
"call_capability" => {
|
"call_capability" => {
|
||||||
@@ -770,7 +770,7 @@ async fn cmd_send(
|
|||||||
|
|
||||||
let base_url = auth.base_url.clone();
|
let base_url = auth.base_url.clone();
|
||||||
let account_id = auth.account_id.clone();
|
let account_id = auth.account_id.clone();
|
||||||
let client = WeChatClient::new(Some(auth.base_url));
|
let mut client = WeChatClient::new(Some(auth.base_url));
|
||||||
client
|
client
|
||||||
.set_auth(&auth.token, &auth.account_id, &base_url)
|
.set_auth(&auth.token, &auth.account_id, &base_url)
|
||||||
.await;
|
.await;
|
||||||
|
|||||||
+15
-3
@@ -219,6 +219,10 @@ impl WeChatClient {
|
|||||||
.await
|
.await
|
||||||
.map_err(|e| format!("解析 notifyStart 响应失败: {}", e))?;
|
.map_err(|e| format!("解析 notifyStart 响应失败: {}", e))?;
|
||||||
|
|
||||||
|
if _resp.ret != 0 {
|
||||||
|
return Err(format!("notifyStart 失败: ret={} errmsg={:?}", _resp.ret, _resp.errmsg));
|
||||||
|
}
|
||||||
|
|
||||||
info!("监听器已注册");
|
info!("监听器已注册");
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
@@ -257,7 +261,14 @@ impl WeChatClient {
|
|||||||
.post_json(&url, &body, Some(&token))
|
.post_json(&url, &body, Some(&token))
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
Ok(resp) => resp.json().await.map_err(|e| format!("解析 getUpdates 响应失败: {}", e))?,
|
Ok(resp) => {
|
||||||
|
let resp: GetUpdatesResp = resp.json().await.map_err(|e| format!("解析 getUpdates 响应失败: {}", e))?;
|
||||||
|
if resp.ret != 0 || resp.errcode.unwrap_or(0) != 0 {
|
||||||
|
tracing::warn!("getUpdates 返回错误: ret={} errcode={:?} errmsg={:?}",
|
||||||
|
resp.ret, resp.errcode, resp.errmsg);
|
||||||
|
}
|
||||||
|
resp
|
||||||
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
// 如果超时,返回空响应
|
// 如果超时,返回空响应
|
||||||
if e.contains("超时") || e.contains("timeout") {
|
if e.contains("超时") || e.contains("timeout") {
|
||||||
@@ -318,11 +329,11 @@ impl WeChatClient {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// 设置 auth 状态(从持久化恢复)
|
/// 设置 auth 状态(从持久化恢复)
|
||||||
pub async fn set_auth(&self, token: &str, account_id: &str, base_url: &str) {
|
pub async fn set_auth(&mut self, token: &str, account_id: &str, base_url: &str) {
|
||||||
*self.token.lock().await = token.to_string();
|
*self.token.lock().await = token.to_string();
|
||||||
*self.account_id.lock().await = account_id.to_string();
|
*self.account_id.lock().await = account_id.to_string();
|
||||||
if !base_url.is_empty() {
|
if !base_url.is_empty() {
|
||||||
// 只更新 base_url 如果传入了非空值
|
self.base_url = base_url.to_string();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -425,6 +436,7 @@ impl WeChatClient {
|
|||||||
.body(body_str)
|
.body(body_str)
|
||||||
.send()
|
.send()
|
||||||
.await
|
.await
|
||||||
|
.and_then(|r| r.error_for_status())
|
||||||
.map_err(|e| {
|
.map_err(|e| {
|
||||||
if e.is_timeout() {
|
if e.is_timeout() {
|
||||||
format!("请求超时: {}", url)
|
format!("请求超时: {}", url)
|
||||||
|
|||||||
Reference in New Issue
Block a user