sz_orm_core/
tenant_context.rs

1//! 多租户上下文与隔离策略
2//!
3//! 本模块在 `multi-tenant-enhanced` feature gate 下导出，提供：
4//! - [`IsolationStrategy`] — 隔离策略枚举（行级 / Schema 隔离）
5//! - [`TenantContext`] + [`TenantContextGuard`] — 租户上下文 + RAII 守卫（task-local 异步隔离）
6//! - [`SchemaIsolationRouter`] — Schema 隔离路由器（表名重写 `tenant_{id}_{table}`）
7//! - [`TenantPoolRegistry`] — 租户连接池注册表（按 tenant_id 维护独立 Pool）
8
9use crate::pool::{Pool, PoolConfig};
10use crate::tenant_security::{ColumnMaskingRule, RowLevelSecurityPolicy};
11use crate::PoolError;
12use std::collections::HashMap;
13use std::sync::Arc;
14
15// ─── M1-T2：租户上下文与 RAII 守卫 ─────────────────────────────────
16
17/// 隔离策略枚举
18///
19/// - `RowLevel`：行级隔离，追加 `WHERE tenant_id = ?`
20/// - `SchemaIsolation`：Schema 隔离，路由到 `tenant_{id}_{table}`
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum IsolationStrategy {
23    /// 行级隔离（追加 `WHERE tenant_id = ?`）
24    RowLevel,
25    /// Schema 隔离（路由到 `tenant_{id}_{table}`）
26    SchemaIsolation,
27}
28
29/// 租户权限（行级安全策略 + 列级脱敏规则 + 角色列表）
30#[derive(Debug, Clone, Default)]
31pub struct TenantPermissions {
32    /// 行级安全策略列表
33    pub row_level_policies: Vec<RowLevelSecurityPolicy>,
34    /// 列级脱敏规则列表
35    pub column_masking_rules: Vec<ColumnMaskingRule>,
36    /// 角色列表
37    pub roles: Vec<String>,
38}
39
40impl TenantPermissions {
41    /// 创建空的权限集合
42    pub fn new() -> Self {
43        Self::default()
44    }
45
46    /// 添加行级安全策略
47    pub fn with_row_level_policy(mut self, policy: RowLevelSecurityPolicy) -> Self {
48        self.row_level_policies.push(policy);
49        self
50    }
51
52    /// 添加列级脱敏规则
53    pub fn with_column_masking_rule(mut self, rule: ColumnMaskingRule) -> Self {
54        self.column_masking_rules.push(rule);
55        self
56    }
57
58    /// 设置角色列表
59    pub fn with_roles(mut self, roles: Vec<String>) -> Self {
60        self.roles = roles;
61        self
62    }
63}
64
65// 线程局部存储：用于 RAII guard 模式（同步上下文安全）
66std::thread_local! {
67    static TENANT_CONTEXT_THREAD: std::cell::RefCell<Option<TenantContext>> = const { std::cell::RefCell::new(None) };
68}
69
70// task-local 存储：用于 scope 模式（异步上下文安全）
71tokio::task_local! {
72    static TENANT_CONTEXT_TASK: std::cell::RefCell<Option<TenantContext>>;
73}
74
75/// 租户上下文（运行时自动注入）
76///
77/// 由可信路径（中间件/网关）设置，不可被客户端篡改。
78/// `tenant_id` 必填 `i64`（禁止字符串避免注入）。
79#[derive(Debug, Clone)]
80pub struct TenantContext {
81    /// 租户 ID（必填，与既有 `with_tenant_id(tenant_id: i64)` 类型一致）
82    pub tenant_id: i64,
83    /// 隔离策略（行级 / Schema 隔离）
84    pub isolation_strategy: IsolationStrategy,
85    /// 权限（行级安全策略 + 列级脱敏规则 + 角色列表）
86    pub permissions: TenantPermissions,
87}
88
89impl TenantContext {
90    /// 创建新的租户上下文
91    pub fn new(tenant_id: i64, isolation_strategy: IsolationStrategy) -> Self {
92        Self {
93            tenant_id,
94            isolation_strategy,
95            permissions: TenantPermissions::new(),
96        }
97    }
98
99    /// 设置权限
100    pub fn with_permissions(mut self, permissions: TenantPermissions) -> Self {
101        self.permissions = permissions;
102        self
103    }
104
105    /// 进入上下文作用域，返回 RAII 守卫（线程局部存储）
106    ///
107    /// 守卫在 Drop 时自动清理线程局部上下文。
108    /// **注意**：此方法使用 `thread_local` 存储，适用于同步上下文或
109    /// 单线程 tokio 运行时。对于多线程 tokio 运行时中的异步任务，
110    /// 请使用 [`TenantContext::scope`] 代替。
111    pub fn enter(self) -> TenantContextGuard {
112        TenantContextGuard::enter(self)
113    }
114
115    /// 在指定异步块的作用域内设置上下文（task-local 存储）
116    ///
117    /// 这是异步安全的上下文设置方式，使用 `tokio::task_local!` 实现
118    /// 异步任务边界隔离。在 scope 内部，[`TenantContext::current`]
119    /// 返回 `Some`；scope 结束后返回 `None`。
120    ///
121    /// # 示例
122    ///
123    /// ```ignore
124    /// let ctx = TenantContext::new(42, IsolationStrategy::RowLevel);
125    /// ctx.scope(async {
126    ///     // 在此作用域内，TenantContext::current() 返回 Some
127    ///     assert_eq!(TenantContext::current().unwrap().tenant_id, 42);
128    /// }).await;
129    /// ```
130    pub async fn scope<F, R>(self, f: F) -> R
131    where
132        F: std::future::Future<Output = R>,
133    {
134        TENANT_CONTEXT_TASK
135            .scope(std::cell::RefCell::new(Some(self)), f)
136            .await
137    }
138
139    /// 读取当前上下文
140    ///
141    /// 优先从 task-local（异步 scope）读取，其次从 thread-local（RAII guard）读取。
142    /// 返回 `None` 表示未设置上下文。
143    pub fn current() -> Option<TenantContext> {
144        // 先尝试 task-local
145        if let Some(ctx) = TENANT_CONTEXT_TASK
146            .try_with(|cell| cell.borrow().clone())
147            .ok()
148            .flatten()
149        {
150            return Some(ctx);
151        }
152        // 再尝试 thread-local
153        TENANT_CONTEXT_THREAD.with(|cell| cell.borrow().clone())
154    }
155
156    /// 检查当前是否在上下文作用域内
157    pub fn is_set() -> bool {
158        Self::current().is_some()
159    }
160}
161
162/// RAII 上下文守卫（线程局部存储）
163///
164/// 在作用域结束时自动清理线程局部上下文。
165/// 守卫作用域内上下文不变，保证租户切换原子。
166pub struct TenantContextGuard {
167    _prev: Option<TenantContext>,
168}
169
170impl TenantContextGuard {
171    fn enter(context: TenantContext) -> Self {
172        let prev = TENANT_CONTEXT_THREAD.with(|cell| cell.borrow().clone());
173        TENANT_CONTEXT_THREAD.with(|cell| {
174            *cell.borrow_mut() = Some(context);
175        });
176        Self { _prev: prev }
177    }
178}
179
180impl Drop for TenantContextGuard {
181    fn drop(&mut self) {
182        TENANT_CONTEXT_THREAD.with(|cell| {
183            *cell.borrow_mut() = self._prev.take();
184        });
185    }
186}
187
188// ─── M1-T3：Schema 隔离路由器 ──────────────────────────────────────
189
190/// Schema 隔离路由器
191///
192/// 将表名重写为 `tenant_{tenant_id}_{table}` 格式。
193/// Schema 命名遵循固定格式，禁止用户自定义避免冲突。
194pub struct SchemaIsolationRouter;
195
196impl SchemaIsolationRouter {
197    /// 重写表名：`table` → `tenant_{tenant_id}_{table}`
198    ///
199    /// # 示例
200    ///
201    /// ```
202    /// use sz_orm_core::tenant_context::SchemaIsolationRouter;
203    /// assert_eq!(SchemaIsolationRouter::rewrite_table("users", 42), "tenant_42_users");
204    /// ```
205    pub fn rewrite_table(table: &str, tenant_id: i64) -> String {
206        format!("tenant_{}_{}", tenant_id, table)
207    }
208
209    /// 重写多条表名
210    pub fn rewrite_tables(tables: &[&str], tenant_id: i64) -> Vec<String> {
211        tables
212            .iter()
213            .map(|t| Self::rewrite_table(t, tenant_id))
214            .collect()
215    }
216}
217
218// ─── M1-T5：租户连接池注册表 ───────────────────────────────────────
219
220/// 租户连接池注册表
221///
222/// 按 `tenant_id` 维护独立的 `Pool`，各租户池共享 `PoolConfig`。
223/// 租户切换原子（RAII 守卫），路由开销 ≤ 50μs（HashMap 查找 + Arc clone）。
224pub struct TenantPoolRegistry {
225    pools: parking_lot::RwLock<HashMap<i64, Arc<Pool>>>,
226    pool_config: PoolConfig,
227}
228
229impl TenantPoolRegistry {
230    /// 创建新的租户连接池注册表
231    pub fn new(pool_config: PoolConfig) -> Self {
232        Self {
233            pools: parking_lot::RwLock::new(HashMap::new()),
234            pool_config,
235        }
236    }
237
238    /// 获取或创建租户连接池
239    ///
240    /// Pool 已存在时返回既有 Pool；不存在时创建新 Pool 并插入 HashMap。
241    /// 使用读锁优先 + 写锁降级策略避免不必要的写锁竞争。
242    pub fn get_or_create(
243        &self,
244        tenant_id: i64,
245        factory: &Arc<dyn crate::pool::ConnectionFactory>,
246    ) -> Result<Arc<Pool>, PoolError> {
247        // 先尝试读锁快速路径
248        {
249            let pools = self.pools.read();
250            if let Some(pool) = pools.get(&tenant_id) {
251                return Ok(Arc::clone(pool));
252            }
253        }
254
255        // 写锁慢速路径：创建新 Pool
256        let mut pools = self.pools.write();
257        // 双检查：可能在等写锁时其他线程已创建
258        if let Some(pool) = pools.get(&tenant_id) {
259            return Ok(Arc::clone(pool));
260        }
261
262        let pool = Arc::new(Pool::new(self.pool_config.clone(), Arc::clone(factory))?);
263        pools.insert(tenant_id, Arc::clone(&pool));
264        Ok(pool)
265    }
266
267    /// 原子切换到指定租户的连接池
268    ///
269    /// 返回 RAII 守卫，在 Drop 时切换回原租户。
270    pub fn switch(
271        &self,
272        tenant_id: i64,
273        factory: &Arc<dyn crate::pool::ConnectionFactory>,
274    ) -> Result<TenantPoolGuard<'_>, PoolError> {
275        let new_pool = self.get_or_create(tenant_id, factory)?;
276        Ok(TenantPoolGuard {
277            registry: self,
278            new_pool,
279        })
280    }
281
282    /// 获取已注册的租户数量
283    pub fn tenant_count(&self) -> usize {
284        self.pools.read().len()
285    }
286
287    /// 获取连接池配置
288    pub fn config(&self) -> &PoolConfig {
289        &self.pool_config
290    }
291}
292
293/// 租户连接池 RAII 守卫
294///
295/// 在 Drop 时不需要显式切换回原租户（各租户 Pool 独立，无全局"当前池"状态）。
296/// 守卫持有新租户 Pool 的 Arc 引用，Drop 时自动释放。
297pub struct TenantPoolGuard<'a> {
298    registry: &'a TenantPoolRegistry,
299    new_pool: Arc<Pool>,
300}
301
302impl<'a> TenantPoolGuard<'a> {
303    /// 获取当前租户的连接池
304    pub fn pool(&self) -> &Arc<Pool> {
305        &self.new_pool
306    }
307
308    /// 获取注册表引用
309    pub fn registry(&self) -> &TenantPoolRegistry {
310        self.registry
311    }
312}
313
314#[cfg(test)]
315mod tests {
316    use super::*;
317
318    // ─── M1-T9.1：TenantContext + TenantContextGuard 测试 ──────────
319
320    #[tokio::test]
321    async fn test_tenant_context_enter_and_current() {
322        let ctx = TenantContext::new(42, IsolationStrategy::RowLevel);
323        let guard = ctx.enter();
324        let current = TenantContext::current();
325        assert!(current.is_some());
326        assert_eq!(current.unwrap().tenant_id, 42);
327        drop(guard);
328        // 守卫 Drop 后上下文应清理
329        assert!(TenantContext::current().is_none());
330    }
331
332    #[tokio::test]
333    async fn test_tenant_context_not_set() {
334        // 未设置上下文时 current() 返回 None
335        // 注意：thread_local 可能残留其他测试的值，所以仅测试 scope 外
336        let ctx = TenantContext::new(99, IsolationStrategy::RowLevel);
337        ctx.scope(async {
338            assert!(TenantContext::is_set());
339        })
340        .await;
341        // scope 结束后 task-local 清理，但 thread-local 可能仍残留
342    }
343
344    #[tokio::test]
345    async fn test_tenant_context_is_set() {
346        let ctx = TenantContext::new(7, IsolationStrategy::SchemaIsolation);
347        ctx.scope(async {
348            assert!(TenantContext::is_set());
349            assert_eq!(TenantContext::current().unwrap().tenant_id, 7);
350        })
351        .await;
352    }
353
354    #[tokio::test]
355    async fn test_tenant_context_nested_scope() {
356        let ctx_a = TenantContext::new(1, IsolationStrategy::RowLevel);
357        ctx_a
358            .scope(async {
359                assert_eq!(TenantContext::current().unwrap().tenant_id, 1);
360
361                let ctx_b = TenantContext::new(2, IsolationStrategy::RowLevel);
362                ctx_b
363                    .scope(async {
364                        assert_eq!(TenantContext::current().unwrap().tenant_id, 2);
365                    })
366                    .await;
367
368                // 恢复回 ctx_a
369                assert_eq!(TenantContext::current().unwrap().tenant_id, 1);
370            })
371            .await;
372    }
373
374    #[tokio::test]
375    async fn test_tenant_context_async_isolation() {
376        // 不同异步任务有独立的 task-local 上下文
377        let ctx_a = TenantContext::new(1, IsolationStrategy::RowLevel);
378        let ctx_b = TenantContext::new(2, IsolationStrategy::RowLevel);
379
380        let handle_a = tokio::spawn(async move {
381            ctx_a
382                .scope(async {
383                    tokio::task::yield_now().await;
384                    TenantContext::current().unwrap().tenant_id
385                })
386                .await
387        });
388
389        let handle_b = tokio::spawn(async move {
390            ctx_b
391                .scope(async {
392                    tokio::task::yield_now().await;
393                    TenantContext::current().unwrap().tenant_id
394                })
395                .await
396        });
397
398        let (id_a, id_b) = tokio::join!(handle_a, handle_b);
399        assert_eq!(id_a.unwrap(), 1);
400        assert_eq!(id_b.unwrap(), 2);
401    }
402
403    // ─── M1-T9.2：SchemaIsolationRouter 测试 ──────────────────────
404
405    #[test]
406    fn test_schema_isolation_router_rewrite() {
407        assert_eq!(
408            SchemaIsolationRouter::rewrite_table("users", 42),
409            "tenant_42_users"
410        );
411        assert_eq!(
412            SchemaIsolationRouter::rewrite_table("orders", 1),
413            "tenant_1_orders"
414        );
415    }
416
417    #[test]
418    fn test_schema_isolation_router_different_tenants() {
419        let table_a = SchemaIsolationRouter::rewrite_table("users", 1);
420        let table_b = SchemaIsolationRouter::rewrite_table("users", 2);
421        assert_ne!(table_a, table_b);
422        assert_eq!(table_a, "tenant_1_users");
423        assert_eq!(table_b, "tenant_2_users");
424    }
425
426    #[test]
427    fn test_schema_isolation_router_rewrite_tables() {
428        let rewritten = SchemaIsolationRouter::rewrite_tables(&["users", "orders"], 42);
429        assert_eq!(rewritten, vec!["tenant_42_users", "tenant_42_orders"]);
430    }
431
432    // ─── TenantPermissions 测试 ───────────────────────────────────
433
434    #[test]
435    fn test_tenant_permissions_default() {
436        let perms = TenantPermissions::new();
437        assert!(perms.row_level_policies.is_empty());
438        assert!(perms.column_masking_rules.is_empty());
439        assert!(perms.roles.is_empty());
440    }
441
442    #[test]
443    fn test_isolation_strategy_copy() {
444        let strategy = IsolationStrategy::RowLevel;
445        let copied = strategy;
446        assert_eq!(strategy, copied);
447    }
448}
sz_orm_core/tenant_context.rs

sz_orm_core/
tenant_context.rs
