feat(auth): keep unsaved passwords in session memory
This commit is contained in:
1 parent
bcec60ebef
commit
1ff24aaeb4
23 files changed
+1265
-65
No files matched your search
@@ -29,6 +29,7 @@ use crate::models::connection::{
|
||||
use crate::path_utils::expand_tilde;
|
||||
use crate::plugins::{PluginDriverSession, PluginRegistry, PluginRuntimeEnv};
|
||||
use crate::query_cancel::RunningQueries;
|
||||
use crate::session_credentials::SessionCredentialStore;
|
||||
use crate::storage::{normalize_duckdb_worker_max_processes, Storage, DUCKDB_WORKER_MAX_PROCESSES_DEFAULT};
|
||||
use crate::task_supervisor::TaskSupervisor;
|
||||
|
||||
@@ -224,6 +225,11 @@ pub struct AppState {
|
||||
/// Used to reconstruct a TLS connector compatible with the original connection when cancelling.
|
||||
postgres_cancel_contexts: Arc<RwLock<HashMap<String, db::postgres::PostgresCancelContext>>>,
|
||||
pub transaction_sessions: Arc<RwLock<HashMap<String, TransactionSession>>>,
|
||||
/// `save_password=false` 连接本次运行期的临时密码(内存,进程退出即丢,
|
||||
/// 绝不落盘)。键为 `(owner_scope, connection_id)`:桌面端 owner 为空串,
|
||||
/// Web 端 owner 为已认证会话 token,不同登录会话互不可见。建池/池重建/
|
||||
/// AI/元数据从它读取,前端通过状态接口查询。
|
||||
pub session_credentials: SessionCredentialStore,
|
||||
#[cfg(feature = "mq-admin")]
|
||||
pub mq_registry: crate::mq::MqAdminRegistry,
|
||||
}
|
||||
@@ -960,6 +966,43 @@ impl AppState {
|
||||
}
|
||||
}
|
||||
|
||||
/// save_password=false 连接:持久化/运行态 config 的 password 恒为空,从本次
|
||||
/// 运行期会话凭据仓库补主密码,使手动/编辑器/AI/元数据/池重建复用首次输入,
|
||||
/// 不再反复弹窗。会话凭据只在内存中,进程退出即丢;"断开并忘记"后仓库为空,
|
||||
/// 此处自然回到空密码(需重新输入)。
|
||||
///
|
||||
/// owner 作用域取自 [`crate::session_credentials::current_credential_owner`]:
|
||||
/// Web 请求由鉴权中间件注入会话 token,桌面端/后台任务为空串。这样不同登录
|
||||
/// 会话只会读取各自输入的密码,不会跨会话复用。
|
||||
pub fn apply_session_credential(
|
||||
&self,
|
||||
config: &ConnectionConfig,
|
||||
db_config: &mut ConnectionConfig,
|
||||
connection_id: &str,
|
||||
) {
|
||||
if !config.save_password && db_config.password.is_empty() {
|
||||
let owner = crate::session_credentials::current_credential_owner().unwrap_or_default();
|
||||
if let Some(session_password) = self.session_credentials.get(&owner, connection_id) {
|
||||
db_config.password = session_password;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// no-save 连接复用已有池时,若该池由"另一个 owner 会话"创建,则不应复用
|
||||
/// (否则会话 B 会直接使用会话 A 输入的临时密码建立的连接)。返回 `true`
|
||||
/// 表示需要销毁旧池、以当前 owner 的凭据重建。
|
||||
async fn pool_credential_owner_mismatch(&self, config: &ConnectionConfig, pool_key: &str) -> bool {
|
||||
if config.save_password {
|
||||
return false;
|
||||
}
|
||||
let owner = crate::session_credentials::current_credential_owner().unwrap_or_default();
|
||||
if owner.is_empty() {
|
||||
// 桌面端/未注入 owner:单用户场景,按既有逻辑复用。
|
||||
return false;
|
||||
}
|
||||
self.session_credentials.pool_owner_mismatch(pool_key, &owner)
|
||||
}
|
||||
|
||||
async fn handle_shared_connection_open_error(
|
||||
&self,
|
||||
error: crate::agent_runtime::SharedConnectionOpenError,
|
||||
@@ -1041,6 +1084,7 @@ impl AppState {
|
||||
duckdb_worker_max_processes: AtomicUsize::new(DUCKDB_WORKER_MAX_PROCESSES_DEFAULT),
|
||||
postgres_cancel_contexts: Arc::new(RwLock::new(HashMap::new())),
|
||||
transaction_sessions: Arc::new(RwLock::new(HashMap::new())),
|
||||
session_credentials: SessionCredentialStore::new(),
|
||||
#[cfg(feature = "mq-admin")]
|
||||
mq_registry: crate::mq::MqAdminRegistry::new(),
|
||||
}
|
||||
@@ -1345,6 +1389,14 @@ impl AppState {
|
||||
routing.detach_pool_by_key(&pool_key, true).await;
|
||||
return Err("Agent runtime is unavailable while publishing the connection pool".to_string());
|
||||
}
|
||||
// 记录 no-save 池由哪个 owner 会话创建,供多会话复用判定使用
|
||||
// (见 pool_credential_owner_mismatch)。已保存密码连接始终共享池,不记录。
|
||||
if !config.save_password {
|
||||
let owner = crate::session_credentials::current_credential_owner().unwrap_or_default();
|
||||
if !owner.is_empty() {
|
||||
self.session_credentials.record_pool_owner(&pool_key, &owner);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -1735,6 +1787,11 @@ impl AppState {
|
||||
drop(conns);
|
||||
if self.remove_pool_if_duckdb_isolation_mismatch(&pool_key).await {
|
||||
// Recreate below using the current DuckDB isolation mode.
|
||||
} else if self.pool_credential_owner_mismatch(&config, &pool_key).await {
|
||||
// 创建该 no-save 池的是另一个登录会话:销毁旧池,用当前会话的
|
||||
// 凭据重建,避免跨会话复用他人输入的临时密码。
|
||||
self.remove_stale_connection_pool(&pool_key).await;
|
||||
break;
|
||||
} else if !validate_existing_pool || !self.remove_stale_connection_pool(&pool_key).await {
|
||||
self.touch_pool_activity(&pool_key).await;
|
||||
return Ok(pool_key);
|
||||
@@ -1752,7 +1809,8 @@ impl AppState {
|
||||
break;
|
||||
}
|
||||
|
||||
let db_config = database_connection_config_with_catalog(&config, database, catalog);
|
||||
let mut db_config = database_connection_config_with_catalog(&config, database, catalog);
|
||||
self.apply_session_credential(&config, &mut db_config, connection_id);
|
||||
|
||||
validate_h2_file_connection(&db_config)?;
|
||||
self.ensure_current_connection_attempt(connection_id, connection_attempt).await?;
|
||||
@@ -4192,6 +4250,7 @@ impl AppState {
|
||||
let pool_keys = self.connections.read().await.keys().cloned().collect::<Vec<_>>();
|
||||
self.stop_keepalive_tasks(&pool_keys).await;
|
||||
self.pool_activity.write().await.clear();
|
||||
self.session_credentials.clear_pool_owners();
|
||||
self.postgres_cancel_contexts.write().await.clear();
|
||||
self.draining_pools.lock().unwrap_or_else(|error| error.into_inner()).clear();
|
||||
self.connections.write().await.drain().collect()
|
||||
@@ -4230,6 +4289,7 @@ impl AppState {
|
||||
cancel_contexts.remove(key);
|
||||
}
|
||||
}
|
||||
self.session_credentials.remove_pool_owners(&keys_to_remove);
|
||||
let mut conns = self.connections.write().await;
|
||||
let mut removed = Vec::with_capacity(keys_to_remove.len());
|
||||
for key in keys_to_remove {
|
||||
@@ -4718,6 +4778,26 @@ fn session_scoped_pool_key_for(
|
||||
session_scoped_pool_key(base_pool_key, client_session_id)
|
||||
}
|
||||
|
||||
/// 两个运行态连接配置是否视为同一连接(用于决定是否销毁连接池)。
|
||||
///
|
||||
/// 仅当双方都是 `save_password=false` 时才忽略 `password` 字段的差异:这类连接
|
||||
/// 在 connect 时运行态配置可能携带会话密码,持久化同步后为空,这种空值差异不应
|
||||
/// 触发池重建。若任一方 `save_password=true`,密码是真实的连接参数,任何密码变更
|
||||
/// (包括用户保存了新密码)都必须销毁旧池,否则旧池会继续用旧密码认证。
|
||||
pub fn connection_configs_pool_equivalent(a: &ConnectionConfig, b: &ConnectionConfig) -> bool {
|
||||
if a == b {
|
||||
return true;
|
||||
}
|
||||
if !a.save_password && !b.save_password {
|
||||
let mut a = a.clone();
|
||||
a.password.clear();
|
||||
let mut b = b.clone();
|
||||
b.password.clear();
|
||||
return a == b;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
fn pool_key_for_session_role(
|
||||
config: Option<&ConnectionConfig>,
|
||||
base_pool_key: String,
|
||||
@@ -5161,15 +5241,16 @@ async fn detect_ob_oracle_mode(config: &ConnectionConfig, pool: &db::mysql::MySq
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
agent_connect_timeout, connection_probe_endpoints, connection_remote_endpoint, connection_url_for_endpoint,
|
||||
database_connection_config, database_connection_config_with_catalog,
|
||||
gaussdb_identifier_quote_from_query_result, gaussdb_m_jdbc_config_for_endpoint, gaussdb_uses_m_jdbc_driver,
|
||||
kafka_single_loopback_bootstrap_endpoint, metadata_connection_config, mysql_metadata_fallback_url,
|
||||
mysql_pool_setup_queries, oceanbase_mysql_query_timeout_sql, oceanbase_mysql_setup_queries,
|
||||
prestosql_jdbc_config_for_endpoint, redacted_connection_url_for_endpoint, redis_sentinel_transport_id,
|
||||
redis_sentinel_transport_prefix, sqlserver_legacy_agent_config, sqlserver_legacy_driver_error,
|
||||
sqlserver_uses_legacy_driver, task_client_session_id, upsert_connection_url_param, uses_bare_mysql_pool,
|
||||
uses_tcp_probe, validate_connection_url_params, validate_h2_database_path, AppState, MysqlMode, PoolKind,
|
||||
agent_connect_timeout, connection_configs_pool_equivalent, connection_probe_endpoints,
|
||||
connection_remote_endpoint, connection_url_for_endpoint, database_connection_config,
|
||||
database_connection_config_with_catalog, gaussdb_identifier_quote_from_query_result,
|
||||
gaussdb_m_jdbc_config_for_endpoint, gaussdb_uses_m_jdbc_driver, kafka_single_loopback_bootstrap_endpoint,
|
||||
metadata_connection_config, mysql_metadata_fallback_url, mysql_pool_setup_queries,
|
||||
oceanbase_mysql_query_timeout_sql, oceanbase_mysql_setup_queries, prestosql_jdbc_config_for_endpoint,
|
||||
redacted_connection_url_for_endpoint, redis_sentinel_transport_id, redis_sentinel_transport_prefix,
|
||||
sqlserver_legacy_agent_config, sqlserver_legacy_driver_error, sqlserver_uses_legacy_driver,
|
||||
task_client_session_id, upsert_connection_url_param, uses_bare_mysql_pool, uses_tcp_probe,
|
||||
validate_connection_url_params, validate_h2_database_path, AppState, MysqlMode, PoolKind,
|
||||
GAUSSDB_M_JDBC_DRIVER_CLASS, GAUSSDB_M_JDBC_DRIVER_PROFILE, PRESTOSQL_JDBC_DRIVER_CLASS,
|
||||
};
|
||||
use crate::agent_connection::{
|
||||
@@ -5280,6 +5361,172 @@ mod tests {
|
||||
assert_eq!(metadata_connection_config(&config).database, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn connection_configs_pool_equivalent_ignores_password_differences() {
|
||||
let mut a = mysql_config(None);
|
||||
a.save_password = false;
|
||||
a.password = "session-secret".to_string();
|
||||
// 持久化同步后 password 为空(save_password=false):视为同一连接,不销毁池。
|
||||
let mut b = a.clone();
|
||||
b.password.clear();
|
||||
assert!(connection_configs_pool_equivalent(&a, &b));
|
||||
assert!(connection_configs_pool_equivalent(&b, &a));
|
||||
|
||||
// 两个非空但不同的密码:同样忽略(密码不应触发池重建)。
|
||||
let mut c = a.clone();
|
||||
c.password = "changed-secret".to_string();
|
||||
assert!(connection_configs_pool_equivalent(&a, &c));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn connection_configs_pool_equivalent_detects_saved_password_change() {
|
||||
let mut a = mysql_config(None);
|
||||
a.save_password = true;
|
||||
a.password = "old-secret".to_string();
|
||||
let mut b = a.clone();
|
||||
b.password = "new-secret".to_string();
|
||||
// save_password=true:密码是真实连接参数,保存新密码后旧池不得继续用旧密码认证。
|
||||
assert!(!connection_configs_pool_equivalent(&a, &b));
|
||||
assert!(!connection_configs_pool_equivalent(&b, &a));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn connection_configs_pool_equivalent_detects_real_parameter_changes() {
|
||||
let mut a = mysql_config(None);
|
||||
a.password = "secret".to_string();
|
||||
// host / port / username / database 等真实连接参数变化应视为不同连接(销毁池)。
|
||||
let mut host = a.clone();
|
||||
host.host = "other-host".to_string();
|
||||
assert!(!connection_configs_pool_equivalent(&a, &host));
|
||||
|
||||
let mut port = a.clone();
|
||||
port.port = 5433;
|
||||
assert!(!connection_configs_pool_equivalent(&a, &port));
|
||||
|
||||
let mut user = a.clone();
|
||||
user.username = "other-user".to_string();
|
||||
assert!(!connection_configs_pool_equivalent(&a, &user));
|
||||
|
||||
let mut ssl = a.clone();
|
||||
ssl.ssl = true;
|
||||
assert!(!connection_configs_pool_equivalent(&a, &ssl));
|
||||
|
||||
// 完全相等(含相同密码)→ true。
|
||||
assert!(connection_configs_pool_equivalent(&a, &a.clone()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn apply_session_credential_injects_saved_password_only_for_no_save_connections() {
|
||||
let dir = std::env::temp_dir().join(format!("dbx-core-session-cred-{}", std::process::id()));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
|
||||
let state = AppState::new_with_plugin_dir(storage, dir.join("plugins"));
|
||||
let _ = state.session_credentials.set("", "conn-a", "s3cret");
|
||||
|
||||
let mut config = mysql_config(None);
|
||||
config.id = "conn-a".to_string();
|
||||
config.save_password = false;
|
||||
config.password.clear();
|
||||
|
||||
// save_password=false + db_config 密码为空 → 从会话凭据仓库注入。
|
||||
let mut db_config = metadata_connection_config(&config);
|
||||
state.apply_session_credential(&config, &mut db_config, &config.id);
|
||||
assert_eq!(db_config.password, "s3cret");
|
||||
|
||||
// save_password=true → 不注入(走持久化水合的密码,若为空则保持空)。
|
||||
config.save_password = true;
|
||||
let mut db_config = metadata_connection_config(&config);
|
||||
state.apply_session_credential(&config, &mut db_config, &config.id);
|
||||
assert_eq!(db_config.password, "");
|
||||
|
||||
// 无会话凭据 → 保持空密码("断开并忘记"后重新输入)。
|
||||
state.session_credentials.remove("", "conn-a");
|
||||
config.save_password = false;
|
||||
let mut db_config = metadata_connection_config(&config);
|
||||
state.apply_session_credential(&config, &mut db_config, &config.id);
|
||||
assert_eq!(db_config.password, "");
|
||||
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn apply_session_credential_reads_owner_scoped_credentials_only() {
|
||||
let dir = std::env::temp_dir().join(format!("dbx-core-session-cred-owner-{}", std::process::id()));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
|
||||
let state = AppState::new_with_plugin_dir(storage, dir.join("plugins"));
|
||||
|
||||
let mut config = mysql_config(None);
|
||||
config.id = "conn-a".to_string();
|
||||
config.save_password = false;
|
||||
config.password.clear();
|
||||
|
||||
// 会话 X 输入了密码,会话 Y 未输入:Y 的请求(owner=token-y)不得注入 X 的密码。
|
||||
let _ = state.session_credentials.set("token-x", "conn-a", "x-secret");
|
||||
crate::session_credentials::with_credential_owner(Some("token-y".to_string()), async {
|
||||
let mut db_config = metadata_connection_config(&config);
|
||||
state.apply_session_credential(&config, &mut db_config, &config.id);
|
||||
assert_eq!(db_config.password, "");
|
||||
})
|
||||
.await;
|
||||
|
||||
// 会话 X 自身能读到自己的密码。
|
||||
crate::session_credentials::with_credential_owner(Some("token-x".to_string()), async {
|
||||
let mut db_config = metadata_connection_config(&config);
|
||||
state.apply_session_credential(&config, &mut db_config, &config.id);
|
||||
assert_eq!(db_config.password, "x-secret");
|
||||
})
|
||||
.await;
|
||||
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_credential_owner_mismatch_prevents_cross_session_pool_reuse() {
|
||||
let dir = std::env::temp_dir().join(format!("dbx-core-pool-owner-{}", std::process::id()));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
|
||||
let state = AppState::new_with_plugin_dir(storage, dir.join("plugins"));
|
||||
|
||||
let mut config = mysql_config(None);
|
||||
config.id = "conn-a".to_string();
|
||||
config.save_password = false;
|
||||
|
||||
// 模拟会话 X 创建的 no-save 池。
|
||||
state.session_credentials.record_pool_owner("conn-a", "token-x");
|
||||
|
||||
// 同一会话 X 复用 → 无冲突。
|
||||
crate::session_credentials::with_credential_owner(Some("token-x".to_string()), async {
|
||||
assert!(!state.pool_credential_owner_mismatch(&config, "conn-a").await);
|
||||
})
|
||||
.await;
|
||||
|
||||
// 另一会话 Y 请求同一 no-save 池 → 冲突,需以 Y 的凭据重建,避免复用 X 的密码。
|
||||
crate::session_credentials::with_credential_owner(Some("token-y".to_string()), async {
|
||||
assert!(state.pool_credential_owner_mismatch(&config, "conn-a").await);
|
||||
})
|
||||
.await;
|
||||
|
||||
// owner 标记缺失时按不可信处理,避免全局配置失效与异步移除旧池之间复用旧池。
|
||||
state.session_credentials.clear_connection("conn-a");
|
||||
crate::session_credentials::with_credential_owner(Some("token-x".to_string()), async {
|
||||
assert!(state.pool_credential_owner_mismatch(&config, "conn-a").await);
|
||||
})
|
||||
.await;
|
||||
|
||||
// 桌面端(无 owner)单用户 → 不冲突,按既有逻辑复用。
|
||||
assert!(!state.pool_credential_owner_mismatch(&config, "conn-a").await);
|
||||
|
||||
// 已保存密码连接始终共享池 → 不冲突。
|
||||
config.save_password = true;
|
||||
crate::session_credentials::with_credential_owner(Some("token-y".to_string()), async {
|
||||
assert!(!state.pool_credential_owner_mismatch(&config, "conn-a").await);
|
||||
})
|
||||
.await;
|
||||
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn task_client_session_ids_are_stable_and_isolated() {
|
||||
assert_eq!(task_client_session_id("table-export", "job-1"), "table-export:job-1");
|
||||
|
||||
@@ -74,6 +74,7 @@ pub mod saved_sql;
|
||||
pub mod schema;
|
||||
pub mod schema_diff;
|
||||
pub mod script_generator;
|
||||
pub mod session_credentials;
|
||||
pub mod sql;
|
||||
pub mod sql_analysis;
|
||||
pub mod sql_diagnostics;
|
||||
|
||||
@@ -0,0 +1,348 @@
|
||||
//! 运行期会话凭据仓库。
|
||||
//!
|
||||
//! 为 `save_password=false` 的连接在本次进程内临时保留主密码,使手动操作、
|
||||
//! 编辑器 SQL、AI 工具、元数据请求以及池重建都能复用首次输入的密码,
|
||||
//! 无需反复弹窗。
|
||||
//!
|
||||
//! 与持久化 secret store([`crate::connection_secrets::FileSecretStore`])职责分离:
|
||||
//! 持久层保存"已保存密码"(save_password=true);本仓库只保存"本次运行期临时密码"
|
||||
//! (save_password=false,进程退出即丢,绝不落盘)。
|
||||
//!
|
||||
//! 约束:
|
||||
//! - 只在 Rust 进程内存中存在,进程退出自然丢失;
|
||||
//! - 不写 SQLite、云同步导出、日志或任何磁盘存储;
|
||||
//! - [`fmt::Debug`] 只暴露凭据数量,不输出 owner、连接 ID 或密码值,避免调试日志与异常文本泄露。
|
||||
//!
|
||||
//! # 多会话隔离(Web)
|
||||
//!
|
||||
//! 凭据按 `(owner_scope, connection_id)` 双键存储:桌面端(Tauri)以空字符串作为
|
||||
//! owner;Web 端以已认证会话 token 作为 owner,使不同登录会话无法复用彼此的临时
|
||||
//! 密码,登出也只清除当前会话的凭据([`SessionCredentialStore::clear_owner`])。
|
||||
//!
|
||||
//! owner 通过 [`tokio::task_local`] 在请求边界注入(见 [`with_credential_owner`] /
|
||||
//! [`current_credential_owner`]),因此 dbx-core 深层的池创建无需在每个函数签名上
|
||||
//! 逐层透传 owner。未注入 owner 的调用(桌面端、后台任务)按空 owner 处理,天然
|
||||
//! 隔离于任何 Web 会话凭据;即使 owner 因后台任务丢失,也会因键不匹配而"失败闭合"
|
||||
//! (读不到任何会话密码),不会跨会话泄露。
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::fmt;
|
||||
use std::future::Future;
|
||||
use std::sync::RwLock;
|
||||
|
||||
/// 桌面端(Tauri)使用的 owner 作用域:无认证会话概念,单用户。
|
||||
pub const DESKTOP_OWNER: &str = "";
|
||||
|
||||
#[derive(Clone, Eq, Hash, PartialEq)]
|
||||
struct CredentialKey {
|
||||
owner_scope: String,
|
||||
connection_id: String,
|
||||
}
|
||||
|
||||
impl CredentialKey {
|
||||
fn new(owner_scope: &str, connection_id: &str) -> Self {
|
||||
Self { owner_scope: owner_scope.to_string(), connection_id: connection_id.to_string() }
|
||||
}
|
||||
}
|
||||
|
||||
struct SessionCredential {
|
||||
password: String,
|
||||
generation: u64,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct SessionCredentialState {
|
||||
credentials: HashMap<CredentialKey, SessionCredential>,
|
||||
pool_credential_owners: HashMap<String, String>,
|
||||
next_generation: u64,
|
||||
}
|
||||
|
||||
pub struct SessionCredentialWriteToken {
|
||||
key: CredentialKey,
|
||||
generation: u64,
|
||||
}
|
||||
|
||||
/// 当前请求的认证会话作用域(Web 会话 token / 桌面端空串)。
|
||||
///
|
||||
/// 由 Web 鉴权中间件在请求边界通过 [`with_credential_owner`] 注入;未注入(桌面端、
|
||||
/// 后台任务、中间件未覆盖的路径)返回 `None`,调用方按空 owner 处理。
|
||||
pub fn current_credential_owner() -> Option<String> {
|
||||
CREDENTIAL_OWNER.try_get().ok().flatten()
|
||||
}
|
||||
|
||||
/// 在一个 future 的整个执行期间设置凭据 owner 作用域。
|
||||
///
|
||||
/// Web 鉴权中间件用它包裹下游处理器,使请求处理任务内(含其 await 到的池创建)
|
||||
/// 都能读到当前会话 owner。未被此函数包裹的调用(桌面端、独立后台任务)等价于
|
||||
/// 空 owner。
|
||||
pub async fn with_credential_owner<F, T>(owner_scope: Option<String>, future: F) -> T
|
||||
where
|
||||
F: Future<Output = T>,
|
||||
{
|
||||
CREDENTIAL_OWNER.scope(owner_scope, future).await
|
||||
}
|
||||
|
||||
tokio::task_local! {
|
||||
static CREDENTIAL_OWNER: Option<String>;
|
||||
}
|
||||
|
||||
/// 内存会话凭据仓库:`(owner_scope, connection_id) -> password`。
|
||||
#[derive(Default)]
|
||||
pub struct SessionCredentialStore {
|
||||
state: RwLock<SessionCredentialState>,
|
||||
}
|
||||
|
||||
impl SessionCredentialStore {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// 记录一个 no-save 连接在本次运行期输入的密码。
|
||||
///
|
||||
/// 空密码是 no-op(不覆盖已有凭据,也不删除),删除只能通过 [`Self::remove`]
|
||||
/// 显式触发("断开并忘记本次密码")。这样连接成功后以空密码 config 建池
|
||||
/// 不会误清已记录的凭据。
|
||||
pub fn set(&self, owner_scope: &str, connection_id: &str, password: &str) -> Option<SessionCredentialWriteToken> {
|
||||
if password.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let mut state = self.state.write().unwrap_or_else(|error| error.into_inner());
|
||||
state.next_generation = state.next_generation.checked_add(1).expect("session credential generation overflow");
|
||||
let key = CredentialKey::new(owner_scope, connection_id);
|
||||
let generation = state.next_generation;
|
||||
state.credentials.insert(key.clone(), SessionCredential { password: password.to_string(), generation });
|
||||
Some(SessionCredentialWriteToken { key, generation })
|
||||
}
|
||||
|
||||
/// 读取某个 owner 作用域下的本次运行期临时密码;不存在则返回 `None`。
|
||||
pub fn get(&self, owner_scope: &str, connection_id: &str) -> Option<String> {
|
||||
let state = self.state.read().unwrap_or_else(|error| error.into_inner());
|
||||
state.credentials.get(&CredentialKey::new(owner_scope, connection_id)).map(|entry| entry.password.clone())
|
||||
}
|
||||
|
||||
/// 某个 owner 作用域下是否存在本次运行期临时密码。
|
||||
pub fn has(&self, owner_scope: &str, connection_id: &str) -> bool {
|
||||
self.get(owner_scope, connection_id).is_some()
|
||||
}
|
||||
|
||||
/// 清除某个 owner 作用域下的临时密码(连接删除 / "断开并忘记本次密码")。
|
||||
pub fn remove(&self, owner_scope: &str, connection_id: &str) {
|
||||
let mut state = self.state.write().unwrap_or_else(|error| error.into_inner());
|
||||
state.credentials.remove(&CredentialKey::new(owner_scope, connection_id));
|
||||
}
|
||||
|
||||
/// 仅当指定写入仍是当前值时删除,用于连接失败回滚,避免旧 attempt 删除更新密码。
|
||||
pub fn remove_if_current(&self, token: &SessionCredentialWriteToken) -> bool {
|
||||
let mut state = self.state.write().unwrap_or_else(|error| error.into_inner());
|
||||
let is_current =
|
||||
state.credentials.get(&token.key).is_some_and(|credential| credential.generation == token.generation);
|
||||
if is_current {
|
||||
state.credentials.remove(&token.key);
|
||||
}
|
||||
is_current
|
||||
}
|
||||
|
||||
/// 清除某个 owner 作用域下的全部凭据(Web 登出 / 登录会话失效),不影响其他
|
||||
/// 登录会话的凭据。
|
||||
pub fn clear_owner(&self, owner_scope: &str) {
|
||||
let mut state = self.state.write().unwrap_or_else(|error| error.into_inner());
|
||||
state.credentials.retain(|key, _| key.owner_scope != owner_scope);
|
||||
}
|
||||
|
||||
/// 全局连接配置被编辑、删除或复用 ID 时,原子清除所有 owner 的临时凭据和
|
||||
/// 该连接的全部池 owner 标记,防止旧密码或旧池流入新的连接定义。
|
||||
pub fn clear_connection(&self, connection_id: &str) {
|
||||
let mut state = self.state.write().unwrap_or_else(|error| error.into_inner());
|
||||
state.credentials.retain(|key, _| key.connection_id != connection_id);
|
||||
let pool_prefix = format!("{connection_id}:");
|
||||
state
|
||||
.pool_credential_owners
|
||||
.retain(|pool_key, _| pool_key != connection_id && !pool_key.starts_with(&pool_prefix));
|
||||
}
|
||||
|
||||
/// 清空全部会话凭据(桌面端退出前兜底;Web 全实例重置)。
|
||||
pub fn clear(&self) {
|
||||
let mut state = self.state.write().unwrap_or_else(|error| error.into_inner());
|
||||
state.credentials.clear();
|
||||
}
|
||||
|
||||
pub fn record_pool_owner(&self, pool_key: &str, owner_scope: &str) {
|
||||
let mut state = self.state.write().unwrap_or_else(|error| error.into_inner());
|
||||
state.pool_credential_owners.insert(pool_key.to_string(), owner_scope.to_string());
|
||||
}
|
||||
|
||||
pub fn pool_owner_mismatch(&self, pool_key: &str, owner_scope: &str) -> bool {
|
||||
let state = self.state.read().unwrap_or_else(|error| error.into_inner());
|
||||
state.pool_credential_owners.get(pool_key).is_none_or(|recorded| recorded != owner_scope)
|
||||
}
|
||||
|
||||
pub fn has_pool_owner(&self, pool_key: &str) -> bool {
|
||||
let state = self.state.read().unwrap_or_else(|error| error.into_inner());
|
||||
state.pool_credential_owners.contains_key(pool_key)
|
||||
}
|
||||
|
||||
pub fn remove_pool_owners(&self, pool_keys: &[String]) {
|
||||
let mut state = self.state.write().unwrap_or_else(|error| error.into_inner());
|
||||
for pool_key in pool_keys {
|
||||
state.pool_credential_owners.remove(pool_key);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn clear_pool_owners(&self) {
|
||||
let mut state = self.state.write().unwrap_or_else(|error| error.into_inner());
|
||||
state.pool_credential_owners.clear();
|
||||
}
|
||||
|
||||
fn credential_count(&self) -> usize {
|
||||
let state = self.state.read().unwrap_or_else(|error| error.into_inner());
|
||||
state.credentials.len()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for SessionCredentialStore {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("SessionCredentialStore").field("credential_count", &self.credential_count()).finish()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SessionCredentialStore;
|
||||
use std::sync::{Arc, Barrier};
|
||||
|
||||
#[test]
|
||||
fn set_get_has_round_trip() {
|
||||
let store = SessionCredentialStore::new();
|
||||
assert!(!store.has("", "conn-a"));
|
||||
assert_eq!(store.get("", "conn-a"), None);
|
||||
|
||||
let _ = store.set("", "conn-a", "s3cret");
|
||||
assert!(store.has("", "conn-a"));
|
||||
assert_eq!(store.get("", "conn-a").as_deref(), Some("s3cret"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn set_with_empty_password_is_noop_and_keeps_existing() {
|
||||
let store = SessionCredentialStore::new();
|
||||
let _ = store.set("", "conn-a", "s3cret");
|
||||
let empty_write = store.set("", "conn-a", "");
|
||||
assert!(empty_write.is_none());
|
||||
assert_eq!(store.get("", "conn-a").as_deref(), Some("s3cret"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_clears_credential() {
|
||||
let store = SessionCredentialStore::new();
|
||||
let _ = store.set("", "conn-a", "s3cret");
|
||||
store.remove("", "conn-a");
|
||||
assert!(!store.has("", "conn-a"));
|
||||
assert_eq!(store.get("", "conn-a"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clear_removes_all_credentials() {
|
||||
let store = SessionCredentialStore::new();
|
||||
let _ = store.set("", "conn-a", "secret-a");
|
||||
let _ = store.set("", "conn-b", "secret-b");
|
||||
store.clear();
|
||||
assert!(!store.has("", "conn-a"));
|
||||
assert!(!store.has("", "conn-b"));
|
||||
assert_eq!(store.get("", "conn-a"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn credentials_are_isolated_by_connection_id() {
|
||||
let store = SessionCredentialStore::new();
|
||||
let _ = store.set("", "conn-a", "secret-a");
|
||||
let _ = store.set("", "conn-b", "secret-b");
|
||||
assert_eq!(store.get("", "conn-a").as_deref(), Some("secret-a"));
|
||||
assert_eq!(store.get("", "conn-b").as_deref(), Some("secret-b"));
|
||||
store.remove("", "conn-a");
|
||||
assert_eq!(store.get("", "conn-a"), None);
|
||||
assert_eq!(store.get("", "conn-b").as_deref(), Some("secret-b"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn credentials_are_isolated_by_owner_scope() {
|
||||
let store = SessionCredentialStore::new();
|
||||
let _ = store.set("token-x", "conn-a", "x-secret");
|
||||
let _ = store.set("token-y", "conn-a", "y-secret");
|
||||
// 同一连接在不同登录会话下互不可见。
|
||||
assert_eq!(store.get("token-x", "conn-a").as_deref(), Some("x-secret"));
|
||||
assert_eq!(store.get("token-y", "conn-a").as_deref(), Some("y-secret"));
|
||||
// 桌面端(空 owner)与任何会话 token 均隔离。
|
||||
assert!(!store.has("", "conn-a"));
|
||||
assert!(!store.has("token-z", "conn-a"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clear_owner_only_removes_that_sessions_credentials() {
|
||||
let store = SessionCredentialStore::new();
|
||||
let _ = store.set("token-x", "conn-a", "x-secret");
|
||||
let _ = store.set("token-y", "conn-a", "y-secret");
|
||||
let _ = store.set("token-y", "conn-b", "y2-secret");
|
||||
store.clear_owner("token-y");
|
||||
// Y 登出只清除 Y 的凭据,X 与桌面端不受影响。
|
||||
assert!(!store.has("token-y", "conn-a"));
|
||||
assert!(!store.has("token-y", "conn-b"));
|
||||
assert!(store.has("token-x", "conn-a"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn debug_output_hides_password_values() {
|
||||
let store = SessionCredentialStore::new();
|
||||
let _ = store.set("session-token-sensitive", "connection-sensitive", "super-secret-value");
|
||||
let debug = format!("{store:?}");
|
||||
assert!(debug.contains("credential_count: 1"));
|
||||
assert!(!debug.contains("session-token-sensitive"));
|
||||
assert!(!debug.contains("connection-sensitive"));
|
||||
assert!(!debug.contains("super-secret-value"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clear_connection_removes_all_owners_and_pool_owner_state() {
|
||||
let store = SessionCredentialStore::new();
|
||||
let _ = store.set("token-x", "conn-a", "x-secret");
|
||||
let _ = store.set("token-y", "conn-a", "y-secret");
|
||||
let _ = store.set("token-y", "conn-b", "other-secret");
|
||||
store.record_pool_owner("conn-a", "token-x");
|
||||
store.record_pool_owner("conn-a:analytics", "token-y");
|
||||
store.record_pool_owner("conn-b", "token-y");
|
||||
|
||||
store.clear_connection("conn-a");
|
||||
|
||||
assert!(!store.has("token-x", "conn-a"));
|
||||
assert!(!store.has("token-y", "conn-a"));
|
||||
assert!(store.has("token-y", "conn-b"));
|
||||
assert!(!store.has_pool_owner("conn-a"));
|
||||
assert!(!store.has_pool_owner("conn-a:analytics"));
|
||||
assert!(store.pool_owner_mismatch("conn-a", "token-z"));
|
||||
assert!(store.pool_owner_mismatch("conn-a:analytics", "token-z"));
|
||||
assert!(store.has_pool_owner("conn-b"));
|
||||
assert!(store.pool_owner_mismatch("conn-b", "token-z"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stale_write_token_cannot_remove_newer_credential() {
|
||||
let store = Arc::new(SessionCredentialStore::new());
|
||||
let first_write_done = Arc::new(Barrier::new(2));
|
||||
let second_write_done = Arc::new(Barrier::new(2));
|
||||
let attempt_store = store.clone();
|
||||
let attempt_first_write_done = first_write_done.clone();
|
||||
let attempt_second_write_done = second_write_done.clone();
|
||||
|
||||
let attempt_a = std::thread::spawn(move || {
|
||||
let token = attempt_store.set("token-x", "conn-a", "attempt-a-password").unwrap();
|
||||
attempt_first_write_done.wait();
|
||||
attempt_second_write_done.wait();
|
||||
assert!(!attempt_store.remove_if_current(&token));
|
||||
});
|
||||
|
||||
first_write_done.wait();
|
||||
let _attempt_b_token = store.set("token-x", "conn-a", "attempt-b-password").unwrap();
|
||||
second_write_done.wait();
|
||||
attempt_a.join().unwrap();
|
||||
|
||||
assert_eq!(store.get("token-x", "conn-a").as_deref(), Some("attempt-b-password"));
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user