#![forbid(unsafe_code)]
use std::sync::Arc;
use std::{error::Error, fmt};
use crate::orm::repository::{
EntityAttributes, Repository, RepositoryError, RepositoryResult, WhereCondition, WhereOp,
};
use crate::orm::Value;
thread_local! {
static TENANT_ID: std::cell::Cell<Option<i64>> = const { std::cell::Cell::new(None) };
}
pub struct TenantContext;
impl TenantContext {
pub fn set_current(tenant_id: i64) {
TENANT_ID.with(|cell| cell.set(Some(tenant_id)));
}
pub fn clear() {
TENANT_ID.with(|cell| cell.set(None));
}
pub fn current() -> Option<i64> {
TENANT_ID.with(|cell| cell.get())
}
pub fn require_current() -> Result<i64, TenantError> {
Self::current().ok_or(TenantError::TenantNotSet)
}
pub fn is_set() -> bool {
Self::current().is_some()
}
pub fn guard() -> Result<TenantGuard, TenantError> {
Self::require_current().map(TenantGuard::new)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TenantGuard {
tenant_id: i64,
}
impl TenantGuard {
fn new(tenant_id: i64) -> Self {
Self { tenant_id }
}
pub fn tenant_id(self) -> i64 {
self.tenant_id
}
pub fn assert_current(&self) -> Result<(), TenantError> {
match TenantContext::current() {
Some(current) if current == self.tenant_id => Ok(()),
Some(current) => Err(TenantError::TenantMismatch {
entity_tenant: self.tenant_id,
current_tenant: current,
}),
None => Err(TenantError::TenantNotSet),
}
}
}
pub trait TenantAware: Clone + Send + Sync + 'static {
fn tenant_id_field() -> &'static str {
"tenant_id"
}
fn tenant_id(&self) -> i64;
fn set_tenant_id(&mut self, tenant_id: i64);
}
#[derive(Debug, Clone, PartialEq)]
pub enum TenantError {
TenantNotSet,
TenantMismatch {
entity_tenant: i64,
current_tenant: i64,
},
}
impl fmt::Display for TenantError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
TenantError::TenantNotSet => {
write!(f, "未设置租户上下文,请先调用 TenantContext::set_current()")
}
TenantError::TenantMismatch {
entity_tenant,
current_tenant,
} => {
write!(
f,
"租户不匹配:实体 tenant_id={},当前租户={}",
entity_tenant, current_tenant
)
}
}
}
}
impl Error for TenantError {}
impl From<TenantError> for RepositoryError {
fn from(err: TenantError) -> Self {
RepositoryError::Other(err.to_string())
}
}
pub struct TenantRepository<E, R> {
inner: Arc<R>,
_marker: std::marker::PhantomData<E>,
}
impl<E: TenantAware, R> TenantRepository<E, R> {
pub fn new(inner: Arc<R>) -> Self {
Self {
inner,
_marker: std::marker::PhantomData,
}
}
fn tenant_condition() -> Result<WhereCondition, TenantError> {
let tid = TenantContext::require_current()?;
Ok(WhereCondition::new(
E::tenant_id_field(),
WhereOp::Eq,
Value::I64(tid),
))
}
fn with_tenant_filter(
conditions: &[WhereCondition],
) -> Result<Vec<WhereCondition>, TenantError> {
let mut all = conditions.to_vec();
all.push(Self::tenant_condition()?);
Ok(all)
}
fn validate_tenant(&self, entity: &mut E) -> Result<(), TenantError> {
let current = TenantContext::require_current()?;
let entity_tid = entity.tenant_id();
if entity_tid == 0 {
entity.set_tenant_id(current);
Ok(())
} else if entity_tid == current {
Ok(())
} else {
Err(TenantError::TenantMismatch {
entity_tenant: entity_tid,
current_tenant: current,
})
}
}
}
impl<E, R> Repository<E> for TenantRepository<E, R>
where
E: TenantAware + EntityAttributes,
R: Repository<E>,
{
type Key = R::Key;
fn key_of(&self, entity: &E) -> Self::Key {
self.inner.key_of(entity)
}
fn find_by_id(&self, key: &Self::Key) -> RepositoryResult<Option<E>> {
let entity = self.inner.find_by_id(key)?;
match entity {
Some(e) => {
let current = match TenantContext::current() {
Some(t) => t,
None => return Err(TenantError::TenantNotSet.into()),
};
if e.tenant_id() == current {
Ok(Some(e))
} else {
Ok(None)
}
}
None => Ok(None),
}
}
fn find_all(&self) -> RepositoryResult<Vec<E>> {
let cond = Self::tenant_condition()?;
self.inner.find_by(&[cond])
}
fn find_by(&self, conditions: &[WhereCondition]) -> RepositoryResult<Vec<E>> {
let all = Self::with_tenant_filter(conditions)?;
self.inner.find_by(&all)
}
fn find_one_by(&self, conditions: &[WhereCondition]) -> RepositoryResult<Option<E>> {
let all = Self::with_tenant_filter(conditions)?;
self.inner.find_one_by(&all)
}
fn save(&self, mut entity: E) -> RepositoryResult<E> {
self.validate_tenant(&mut entity)?;
self.inner.save(entity)
}
fn save_many(&self, mut entities: Vec<E>) -> RepositoryResult<Vec<E>> {
for e in &mut entities {
self.validate_tenant(e)?;
}
self.inner.save_many(entities)
}
fn delete(&self, key: &Self::Key) -> RepositoryResult<usize> {
match self.find_by_id(key)? {
Some(_) => self.inner.delete(key),
None => Ok(0),
}
}
fn delete_by(&self, conditions: &[WhereCondition]) -> RepositoryResult<usize> {
let all = Self::with_tenant_filter(conditions)?;
self.inner.delete_by(&all)
}
fn count(&self) -> RepositoryResult<u64> {
let cond = Self::tenant_condition()?;
self.inner.count_by(&[cond])
}
fn count_by(&self, conditions: &[WhereCondition]) -> RepositoryResult<u64> {
let all = Self::with_tenant_filter(conditions)?;
self.inner.count_by(&all)
}
}
use axum::{
body::Body,
http::{Request, StatusCode},
middleware::Next,
response::Response,
};
pub async fn tenant_middleware(
req: Request<Body>,
next: Next,
) -> Result<Response, (StatusCode, String)> {
let tenant_id_str = req
.headers()
.get("X-Tenant-Id")
.and_then(|v| v.to_str().ok())
.ok_or_else(|| {
(
StatusCode::BAD_REQUEST,
"Missing X-Tenant-Id header".to_string(),
)
})?;
let tenant_id: i64 = tenant_id_str.parse().map_err(|_| {
(
StatusCode::BAD_REQUEST,
"X-Tenant-Id must be a valid integer".to_string(),
)
})?;
TenantContext::set_current(tenant_id);
let response = next.run(req).await;
TenantContext::clear();
Ok(response)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::orm::repository::InMemoryRepository;
#[derive(Clone, Debug, PartialEq)]
struct TenantOrder {
id: i64,
tenant_id: i64,
order_no: String,
}
impl EntityAttributes for TenantOrder {
fn get_attribute(&self, field: &str) -> Option<Value> {
match field {
"id" => Some(Value::I64(self.id)),
"tenant_id" => Some(Value::I64(self.tenant_id)),
"order_no" => Some(Value::String(self.order_no.clone())),
_ => None,
}
}
}
impl TenantAware for TenantOrder {
fn tenant_id_field() -> &'static str {
"tenant_id"
}
fn tenant_id(&self) -> i64 {
self.tenant_id
}
fn set_tenant_id(&mut self, tid: i64) {
self.tenant_id = tid;
}
}
fn make_order(id: i64, tenant_id: i64, no: &str) -> TenantOrder {
TenantOrder {
id,
tenant_id,
order_no: no.to_string(),
}
}
fn repo() -> TenantRepository<TenantOrder, InMemoryRepository<TenantOrder>> {
TenantRepository::new(Arc::new(InMemoryRepository::new()))
}
#[test]
fn test_tenant_context_set_and_get() {
TenantContext::clear();
assert!(!TenantContext::is_set());
TenantContext::set_current(1001);
assert!(TenantContext::is_set());
assert_eq!(TenantContext::current(), Some(1001));
assert_eq!(TenantContext::require_current(), Ok(1001));
TenantContext::clear();
}
#[test]
fn test_tenant_context_require_current_fails_when_unset() {
TenantContext::clear();
assert!(matches!(
TenantContext::require_current(),
Err(TenantError::TenantNotSet)
));
}
#[test]
fn test_find_by_auto_filters_tenant() {
TenantContext::clear();
let r = repo();
r.inner.save(make_order(1, 1001, "ORD-001")).unwrap();
r.inner.save(make_order(2, 1001, "ORD-002")).unwrap();
r.inner.save(make_order(3, 2002, "ORD-003")).unwrap();
TenantContext::clear();
assert!(r.find_by(&[]).is_err());
TenantContext::set_current(1001);
let orders = r.find_by(&[]).unwrap();
assert_eq!(orders.len(), 2);
assert!(orders.iter().all(|o| o.tenant_id == 1001));
TenantContext::set_current(2002);
let orders = r.find_by(&[]).unwrap();
assert_eq!(orders.len(), 1);
assert_eq!(orders[0].order_no, "ORD-003");
}
#[test]
fn test_find_by_with_additional_conditions() {
TenantContext::clear();
let r = repo();
r.inner.save(make_order(1, 1001, "ORD-001")).unwrap();
r.inner.save(make_order(2, 1001, "ORD-002")).unwrap();
r.inner.save(make_order(3, 1001, "ORD-003")).unwrap();
TenantContext::set_current(1001);
let orders = r
.find_by(&[WhereCondition::new("id", WhereOp::Ge, Value::I64(2))])
.unwrap();
assert_eq!(orders.len(), 2);
}
#[test]
fn test_find_by_id_hides_other_tenant_data() {
TenantContext::clear();
let r = repo();
r.inner.save(make_order(42, 2002, "ORD-042")).unwrap();
TenantContext::set_current(1001);
let result = r.find_by_id(&Value::I64(42)).unwrap();
assert!(result.is_none(), "跨租户数据应被隐藏");
}
#[test]
fn test_find_by_id_returns_own_data() {
TenantContext::clear();
let r = repo();
r.inner.save(make_order(42, 1001, "ORD-042")).unwrap();
TenantContext::set_current(1001);
let result = r.find_by_id(&Value::I64(42)).unwrap();
assert!(result.is_some());
assert_eq!(result.unwrap().order_no, "ORD-042");
}
#[test]
fn test_save_auto_injects_tenant_when_zero() {
TenantContext::clear();
let r = repo();
TenantContext::set_current(1001);
let order = make_order(0, 0, "ORD-NEW");
let saved = r.save(order).unwrap();
assert_eq!(saved.tenant_id, 1001, "tenant_id 应自动注入为当前租户");
}
#[test]
fn test_save_rejects_cross_tenant_write() {
TenantContext::clear();
let r = repo();
TenantContext::set_current(1001);
let order = make_order(0, 2002, "ORD-BAD");
let result = r.save(order);
assert!(matches!(result, Err(RepositoryError::Other(_))));
let err_msg = result.unwrap_err().to_string();
assert!(
err_msg.contains("租户不匹配"),
"错误信息应包含租户不匹配: {}",
err_msg
);
}
#[test]
fn test_save_many_all_must_match_tenant() {
TenantContext::clear();
let r = repo();
TenantContext::set_current(1001);
let orders = vec![make_order(0, 0, "ORD-A"), make_order(0, 2002, "ORD-B")];
let result = r.save_many(orders);
assert!(result.is_err(), "批量保存中存在跨租户数据应整体失败");
}
#[test]
fn test_delete_only_deletes_own_tenant() {
TenantContext::clear();
let r = repo();
r.inner.save(make_order(1, 2002, "ORD-001")).unwrap();
TenantContext::set_current(1001);
let count = r.delete(&Value::I64(1)).unwrap();
assert_eq!(count, 0, "跨租户删除应返回 0");
TenantContext::set_current(2002);
let found = r.find_by_id(&Value::I64(1)).unwrap();
assert!(found.is_some());
}
#[test]
fn test_delete_by_auto_filters_tenant() {
TenantContext::clear();
let r = repo();
r.inner.save(make_order(1, 1001, "ORD-001")).unwrap();
r.inner.save(make_order(2, 2002, "ORD-002")).unwrap();
TenantContext::set_current(1001);
let count = r.delete_by(&[]).unwrap();
assert_eq!(count, 1);
TenantContext::set_current(2002);
let remaining = r.find_by(&[]).unwrap();
assert_eq!(remaining.len(), 1);
assert_eq!(remaining[0].id, 2);
}
#[test]
fn test_count_by_auto_filters_tenant() {
TenantContext::clear();
let r = repo();
r.inner.save(make_order(1, 1001, "ORD-001")).unwrap();
r.inner.save(make_order(2, 1001, "ORD-002")).unwrap();
r.inner.save(make_order(3, 2002, "ORD-003")).unwrap();
TenantContext::set_current(1001);
let count = r.count_by(&[]).unwrap();
assert_eq!(count, 2);
}
#[test]
fn test_count_auto_filters_tenant() {
TenantContext::clear();
let r = repo();
r.inner.save(make_order(1, 1001, "ORD-001")).unwrap();
r.inner.save(make_order(2, 2002, "ORD-002")).unwrap();
TenantContext::set_current(1001);
let count = r.count().unwrap();
assert_eq!(count, 1);
}
#[test]
fn test_tenant_error_display() {
let e = TenantError::TenantNotSet;
assert!(e.to_string().contains("未设置租户上下文"));
let e = TenantError::TenantMismatch {
entity_tenant: 2002,
current_tenant: 1001,
};
let msg = e.to_string();
assert!(msg.contains("2002"));
assert!(msg.contains("1001"));
assert!(msg.contains("租户不匹配"));
}
#[test]
fn test_tenant_guard_captures_current_tenant() {
TenantContext::clear();
TenantContext::set_current(1001);
let guard = TenantContext::guard().expect("应成功创建 guard");
assert_eq!(guard.tenant_id(), 1001);
TenantContext::set_current(2002);
assert_eq!(guard.tenant_id(), 1001, "guard 值不应随 thread_local 改变");
TenantContext::clear();
}
#[test]
fn test_tenant_guard_fails_when_unset() {
TenantContext::clear();
let result = TenantContext::guard();
assert!(matches!(result, Err(TenantError::TenantNotSet)));
}
#[test]
fn test_tenant_guard_assert_current_matches() {
TenantContext::clear();
TenantContext::set_current(1001);
let guard = TenantContext::guard().unwrap();
assert!(guard.assert_current().is_ok());
TenantContext::clear();
}
#[test]
fn test_tenant_guard_assert_current_mismatch() {
TenantContext::clear();
TenantContext::set_current(1001);
let guard = TenantContext::guard().unwrap();
TenantContext::set_current(2002);
let result = guard.assert_current();
assert!(matches!(result, Err(TenantError::TenantMismatch { .. })));
TenantContext::clear();
}
#[test]
fn test_tenant_guard_assert_current_after_clear() {
TenantContext::clear();
TenantContext::set_current(1001);
let guard = TenantContext::guard().unwrap();
TenantContext::clear();
let result = guard.assert_current();
assert!(matches!(result, Err(TenantError::TenantNotSet)));
}
#[test]
fn test_tenant_guard_is_copy() {
TenantContext::clear();
TenantContext::set_current(1001);
let guard = TenantContext::guard().unwrap();
let guard_copy = guard; assert_eq!(guard.tenant_id(), 1001);
assert_eq!(guard_copy.tenant_id(), 1001);
TenantContext::clear();
}
#[test]
fn test_find_all_auto_filters_tenant() {
TenantContext::clear();
let r = repo();
r.inner.save(make_order(1, 1001, "ORD-001")).unwrap();
r.inner.save(make_order(2, 2002, "ORD-002")).unwrap();
TenantContext::set_current(1001);
let all = r.find_all().unwrap();
assert_eq!(all.len(), 1);
assert_eq!(all[0].tenant_id, 1001);
}
#[test]
fn test_save_fails_without_tenant_context() {
TenantContext::clear();
let r = repo();
let order = make_order(0, 0, "ORD-NEW");
let result = r.save(order);
assert!(matches!(
result.unwrap_err().to_string().as_str(),
s if s.contains("未设置租户上下文")
));
}
}