1#![forbid(unsafe_code)]
43
44use std::sync::Arc;
45use std::{error::Error, fmt};
46
47use crate::orm::repository::{
48 EntityAttributes, Repository, RepositoryError, RepositoryResult, WhereCondition, WhereOp,
49};
50use crate::orm::Value;
51
52thread_local! {
57 static TENANT_ID: std::cell::Cell<Option<i64>> = const { std::cell::Cell::new(None) };
58}
59
60pub struct TenantContext;
68
69impl TenantContext {
70 pub fn set_current(tenant_id: i64) {
72 TENANT_ID.with(|cell| cell.set(Some(tenant_id)));
73 }
74
75 pub fn clear() {
77 TENANT_ID.with(|cell| cell.set(None));
78 }
79
80 pub fn current() -> Option<i64> {
84 TENANT_ID.with(|cell| cell.get())
85 }
86
87 pub fn require_current() -> Result<i64, TenantError> {
89 Self::current().ok_or(TenantError::TenantNotSet)
90 }
91
92 pub fn is_set() -> bool {
94 Self::current().is_some()
95 }
96
97 pub fn guard() -> Result<TenantGuard, TenantError> {
128 Self::require_current().map(TenantGuard::new)
129 }
130}
131
132#[derive(Debug, Clone, Copy, PartialEq, Eq)]
158pub struct TenantGuard {
159 tenant_id: i64,
160}
161
162impl TenantGuard {
163 fn new(tenant_id: i64) -> Self {
165 Self { tenant_id }
166 }
167
168 pub fn tenant_id(self) -> i64 {
172 self.tenant_id
173 }
174
175 pub fn assert_current(&self) -> Result<(), TenantError> {
186 match TenantContext::current() {
187 Some(current) if current == self.tenant_id => Ok(()),
188 Some(current) => Err(TenantError::TenantMismatch {
189 entity_tenant: self.tenant_id,
190 current_tenant: current,
191 }),
192 None => Err(TenantError::TenantNotSet),
193 }
194 }
195}
196
197pub trait TenantAware: Clone + Send + Sync + 'static {
205 fn tenant_id_field() -> &'static str {
207 "tenant_id"
208 }
209
210 fn tenant_id(&self) -> i64;
212
213 fn set_tenant_id(&mut self, tenant_id: i64);
215}
216
217#[derive(Debug, Clone, PartialEq)]
223pub enum TenantError {
224 TenantNotSet,
226 TenantMismatch {
228 entity_tenant: i64,
230 current_tenant: i64,
232 },
233}
234
235impl fmt::Display for TenantError {
236 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
237 match self {
238 TenantError::TenantNotSet => {
239 write!(f, "未设置租户上下文,请先调用 TenantContext::set_current()")
240 }
241 TenantError::TenantMismatch {
242 entity_tenant,
243 current_tenant,
244 } => {
245 write!(
246 f,
247 "租户不匹配:实体 tenant_id={},当前租户={}",
248 entity_tenant, current_tenant
249 )
250 }
251 }
252 }
253}
254
255impl Error for TenantError {}
256
257impl From<TenantError> for RepositoryError {
258 fn from(err: TenantError) -> Self {
259 RepositoryError::Other(err.to_string())
260 }
261}
262
263pub struct TenantRepository<E, R> {
282 inner: Arc<R>,
283 _marker: std::marker::PhantomData<E>,
284}
285
286impl<E: TenantAware, R> TenantRepository<E, R> {
287 pub fn new(inner: Arc<R>) -> Self {
289 Self {
290 inner,
291 _marker: std::marker::PhantomData,
292 }
293 }
294
295 fn tenant_condition() -> Result<WhereCondition, TenantError> {
297 let tid = TenantContext::require_current()?;
298 Ok(WhereCondition::new(
299 E::tenant_id_field(),
300 WhereOp::Eq,
301 Value::I64(tid),
302 ))
303 }
304
305 fn with_tenant_filter(
307 conditions: &[WhereCondition],
308 ) -> Result<Vec<WhereCondition>, TenantError> {
309 let mut all = conditions.to_vec();
310 all.push(Self::tenant_condition()?);
311 Ok(all)
312 }
313
314 fn validate_tenant(&self, entity: &mut E) -> Result<(), TenantError> {
316 let current = TenantContext::require_current()?;
317 let entity_tid = entity.tenant_id();
318 if entity_tid == 0 {
319 entity.set_tenant_id(current);
320 Ok(())
321 } else if entity_tid == current {
322 Ok(())
323 } else {
324 Err(TenantError::TenantMismatch {
325 entity_tenant: entity_tid,
326 current_tenant: current,
327 })
328 }
329 }
330}
331
332impl<E, R> Repository<E> for TenantRepository<E, R>
333where
334 E: TenantAware + EntityAttributes,
335 R: Repository<E>,
336{
337 type Key = R::Key;
338
339 fn key_of(&self, entity: &E) -> Self::Key {
340 self.inner.key_of(entity)
341 }
342
343 fn find_by_id(&self, key: &Self::Key) -> RepositoryResult<Option<E>> {
344 let entity = self.inner.find_by_id(key)?;
345 match entity {
346 Some(e) => {
347 let current = match TenantContext::current() {
348 Some(t) => t,
349 None => return Err(TenantError::TenantNotSet.into()),
350 };
351 if e.tenant_id() == current {
352 Ok(Some(e))
353 } else {
354 Ok(None)
356 }
357 }
358 None => Ok(None),
359 }
360 }
361
362 fn find_all(&self) -> RepositoryResult<Vec<E>> {
363 let cond = Self::tenant_condition()?;
364 self.inner.find_by(&[cond])
365 }
366
367 fn find_by(&self, conditions: &[WhereCondition]) -> RepositoryResult<Vec<E>> {
368 let all = Self::with_tenant_filter(conditions)?;
369 self.inner.find_by(&all)
370 }
371
372 fn find_one_by(&self, conditions: &[WhereCondition]) -> RepositoryResult<Option<E>> {
373 let all = Self::with_tenant_filter(conditions)?;
374 self.inner.find_one_by(&all)
375 }
376
377 fn save(&self, mut entity: E) -> RepositoryResult<E> {
378 self.validate_tenant(&mut entity)?;
379 self.inner.save(entity)
380 }
381
382 fn save_many(&self, mut entities: Vec<E>) -> RepositoryResult<Vec<E>> {
383 for e in &mut entities {
384 self.validate_tenant(e)?;
385 }
386 self.inner.save_many(entities)
387 }
388
389 fn delete(&self, key: &Self::Key) -> RepositoryResult<usize> {
390 match self.find_by_id(key)? {
391 Some(_) => self.inner.delete(key),
392 None => Ok(0),
393 }
394 }
395
396 fn delete_by(&self, conditions: &[WhereCondition]) -> RepositoryResult<usize> {
397 let all = Self::with_tenant_filter(conditions)?;
398 self.inner.delete_by(&all)
399 }
400
401 fn count(&self) -> RepositoryResult<u64> {
402 let cond = Self::tenant_condition()?;
403 self.inner.count_by(&[cond])
404 }
405
406 fn count_by(&self, conditions: &[WhereCondition]) -> RepositoryResult<u64> {
407 let all = Self::with_tenant_filter(conditions)?;
408 self.inner.count_by(&all)
409 }
410}
411
412use axum::{
417 body::Body,
418 http::{Request, StatusCode},
419 middleware::Next,
420 response::Response,
421};
422
423pub async fn tenant_middleware(
435 req: Request<Body>,
436 next: Next,
437) -> Result<Response, (StatusCode, String)> {
438 let tenant_id_str = req
439 .headers()
440 .get("X-Tenant-Id")
441 .and_then(|v| v.to_str().ok())
442 .ok_or_else(|| {
443 (
444 StatusCode::BAD_REQUEST,
445 "Missing X-Tenant-Id header".to_string(),
446 )
447 })?;
448
449 let tenant_id: i64 = tenant_id_str.parse().map_err(|_| {
450 (
451 StatusCode::BAD_REQUEST,
452 "X-Tenant-Id must be a valid integer".to_string(),
453 )
454 })?;
455
456 TenantContext::set_current(tenant_id);
457
458 let response = next.run(req).await;
459 TenantContext::clear();
460
461 Ok(response)
462}
463
464#[cfg(test)]
469mod tests {
470 use super::*;
471 use crate::orm::repository::InMemoryRepository;
472
473 #[derive(Clone, Debug, PartialEq)]
476 struct TenantOrder {
477 id: i64,
478 tenant_id: i64,
479 order_no: String,
480 }
481
482 impl EntityAttributes for TenantOrder {
483 fn get_attribute(&self, field: &str) -> Option<Value> {
484 match field {
485 "id" => Some(Value::I64(self.id)),
486 "tenant_id" => Some(Value::I64(self.tenant_id)),
487 "order_no" => Some(Value::String(self.order_no.clone())),
488 _ => None,
489 }
490 }
491 }
492
493 impl TenantAware for TenantOrder {
494 fn tenant_id_field() -> &'static str {
495 "tenant_id"
496 }
497 fn tenant_id(&self) -> i64 {
498 self.tenant_id
499 }
500 fn set_tenant_id(&mut self, tid: i64) {
501 self.tenant_id = tid;
502 }
503 }
504
505 fn make_order(id: i64, tenant_id: i64, no: &str) -> TenantOrder {
506 TenantOrder {
507 id,
508 tenant_id,
509 order_no: no.to_string(),
510 }
511 }
512
513 fn repo() -> TenantRepository<TenantOrder, InMemoryRepository<TenantOrder>> {
514 TenantRepository::new(Arc::new(InMemoryRepository::new()))
515 }
516
517 #[test]
520 fn test_tenant_context_set_and_get() {
521 TenantContext::clear();
522 assert!(!TenantContext::is_set());
523 TenantContext::set_current(1001);
524 assert!(TenantContext::is_set());
525 assert_eq!(TenantContext::current(), Some(1001));
526 assert_eq!(TenantContext::require_current(), Ok(1001));
527 TenantContext::clear();
528 }
529
530 #[test]
531 fn test_tenant_context_require_current_fails_when_unset() {
532 TenantContext::clear();
533 assert!(matches!(
534 TenantContext::require_current(),
535 Err(TenantError::TenantNotSet)
536 ));
537 }
538
539 #[test]
542 fn test_find_by_auto_filters_tenant() {
543 TenantContext::clear();
544 let r = repo();
545
546 r.inner.save(make_order(1, 1001, "ORD-001")).unwrap();
548 r.inner.save(make_order(2, 1001, "ORD-002")).unwrap();
549 r.inner.save(make_order(3, 2002, "ORD-003")).unwrap();
550
551 TenantContext::clear();
553 assert!(r.find_by(&[]).is_err());
554
555 TenantContext::set_current(1001);
557 let orders = r.find_by(&[]).unwrap();
558 assert_eq!(orders.len(), 2);
559 assert!(orders.iter().all(|o| o.tenant_id == 1001));
560
561 TenantContext::set_current(2002);
563 let orders = r.find_by(&[]).unwrap();
564 assert_eq!(orders.len(), 1);
565 assert_eq!(orders[0].order_no, "ORD-003");
566 }
567
568 #[test]
569 fn test_find_by_with_additional_conditions() {
570 TenantContext::clear();
571 let r = repo();
572 r.inner.save(make_order(1, 1001, "ORD-001")).unwrap();
573 r.inner.save(make_order(2, 1001, "ORD-002")).unwrap();
574 r.inner.save(make_order(3, 1001, "ORD-003")).unwrap();
575
576 TenantContext::set_current(1001);
577 let orders = r
578 .find_by(&[WhereCondition::new("id", WhereOp::Ge, Value::I64(2))])
579 .unwrap();
580 assert_eq!(orders.len(), 2);
581 }
582
583 #[test]
586 fn test_find_by_id_hides_other_tenant_data() {
587 TenantContext::clear();
588 let r = repo();
589 r.inner.save(make_order(42, 2002, "ORD-042")).unwrap();
590
591 TenantContext::set_current(1001);
592 let result = r.find_by_id(&Value::I64(42)).unwrap();
593 assert!(result.is_none(), "跨租户数据应被隐藏");
594 }
595
596 #[test]
597 fn test_find_by_id_returns_own_data() {
598 TenantContext::clear();
599 let r = repo();
600 r.inner.save(make_order(42, 1001, "ORD-042")).unwrap();
601
602 TenantContext::set_current(1001);
603 let result = r.find_by_id(&Value::I64(42)).unwrap();
604 assert!(result.is_some());
605 assert_eq!(result.unwrap().order_no, "ORD-042");
606 }
607
608 #[test]
611 fn test_save_auto_injects_tenant_when_zero() {
612 TenantContext::clear();
613 let r = repo();
614 TenantContext::set_current(1001);
615
616 let order = make_order(0, 0, "ORD-NEW");
617 let saved = r.save(order).unwrap();
618 assert_eq!(saved.tenant_id, 1001, "tenant_id 应自动注入为当前租户");
619 }
620
621 #[test]
622 fn test_save_rejects_cross_tenant_write() {
623 TenantContext::clear();
624 let r = repo();
625 TenantContext::set_current(1001);
626
627 let order = make_order(0, 2002, "ORD-BAD");
628 let result = r.save(order);
629 assert!(matches!(result, Err(RepositoryError::Other(_))));
630 let err_msg = result.unwrap_err().to_string();
631 assert!(
632 err_msg.contains("租户不匹配"),
633 "错误信息应包含租户不匹配: {}",
634 err_msg
635 );
636 }
637
638 #[test]
639 fn test_save_many_all_must_match_tenant() {
640 TenantContext::clear();
641 let r = repo();
642 TenantContext::set_current(1001);
643
644 let orders = vec![make_order(0, 0, "ORD-A"), make_order(0, 2002, "ORD-B")];
645 let result = r.save_many(orders);
646 assert!(result.is_err(), "批量保存中存在跨租户数据应整体失败");
647 }
648
649 #[test]
652 fn test_delete_only_deletes_own_tenant() {
653 TenantContext::clear();
654 let r = repo();
655 r.inner.save(make_order(1, 2002, "ORD-001")).unwrap();
656
657 TenantContext::set_current(1001);
658 let count = r.delete(&Value::I64(1)).unwrap();
659 assert_eq!(count, 0, "跨租户删除应返回 0");
660
661 TenantContext::set_current(2002);
662 let found = r.find_by_id(&Value::I64(1)).unwrap();
663 assert!(found.is_some());
664 }
665
666 #[test]
669 fn test_delete_by_auto_filters_tenant() {
670 TenantContext::clear();
671 let r = repo();
672 r.inner.save(make_order(1, 1001, "ORD-001")).unwrap();
673 r.inner.save(make_order(2, 2002, "ORD-002")).unwrap();
674
675 TenantContext::set_current(1001);
676 let count = r.delete_by(&[]).unwrap();
677 assert_eq!(count, 1);
678
679 TenantContext::set_current(2002);
680 let remaining = r.find_by(&[]).unwrap();
681 assert_eq!(remaining.len(), 1);
682 assert_eq!(remaining[0].id, 2);
683 }
684
685 #[test]
688 fn test_count_by_auto_filters_tenant() {
689 TenantContext::clear();
690 let r = repo();
691 r.inner.save(make_order(1, 1001, "ORD-001")).unwrap();
692 r.inner.save(make_order(2, 1001, "ORD-002")).unwrap();
693 r.inner.save(make_order(3, 2002, "ORD-003")).unwrap();
694
695 TenantContext::set_current(1001);
696 let count = r.count_by(&[]).unwrap();
697 assert_eq!(count, 2);
698 }
699
700 #[test]
701 fn test_count_auto_filters_tenant() {
702 TenantContext::clear();
703 let r = repo();
704 r.inner.save(make_order(1, 1001, "ORD-001")).unwrap();
705 r.inner.save(make_order(2, 2002, "ORD-002")).unwrap();
706
707 TenantContext::set_current(1001);
708 let count = r.count().unwrap();
709 assert_eq!(count, 1);
710 }
711
712 #[test]
715 fn test_tenant_error_display() {
716 let e = TenantError::TenantNotSet;
717 assert!(e.to_string().contains("未设置租户上下文"));
718
719 let e = TenantError::TenantMismatch {
720 entity_tenant: 2002,
721 current_tenant: 1001,
722 };
723 let msg = e.to_string();
724 assert!(msg.contains("2002"));
725 assert!(msg.contains("1001"));
726 assert!(msg.contains("租户不匹配"));
727 }
728
729 #[test]
732 fn test_tenant_guard_captures_current_tenant() {
733 TenantContext::clear();
734 TenantContext::set_current(1001);
735
736 let guard = TenantContext::guard().expect("应成功创建 guard");
737 assert_eq!(guard.tenant_id(), 1001);
738
739 TenantContext::set_current(2002);
741 assert_eq!(guard.tenant_id(), 1001, "guard 值不应随 thread_local 改变");
742
743 TenantContext::clear();
744 }
745
746 #[test]
747 fn test_tenant_guard_fails_when_unset() {
748 TenantContext::clear();
749 let result = TenantContext::guard();
750 assert!(matches!(result, Err(TenantError::TenantNotSet)));
751 }
752
753 #[test]
754 fn test_tenant_guard_assert_current_matches() {
755 TenantContext::clear();
756 TenantContext::set_current(1001);
757
758 let guard = TenantContext::guard().unwrap();
759 assert!(guard.assert_current().is_ok());
760
761 TenantContext::clear();
762 }
763
764 #[test]
765 fn test_tenant_guard_assert_current_mismatch() {
766 TenantContext::clear();
767 TenantContext::set_current(1001);
768
769 let guard = TenantContext::guard().unwrap();
770
771 TenantContext::set_current(2002);
773 let result = guard.assert_current();
774 assert!(matches!(result, Err(TenantError::TenantMismatch { .. })));
775
776 TenantContext::clear();
777 }
778
779 #[test]
780 fn test_tenant_guard_assert_current_after_clear() {
781 TenantContext::clear();
782 TenantContext::set_current(1001);
783
784 let guard = TenantContext::guard().unwrap();
785
786 TenantContext::clear();
788 let result = guard.assert_current();
789 assert!(matches!(result, Err(TenantError::TenantNotSet)));
790 }
791
792 #[test]
793 fn test_tenant_guard_is_copy() {
794 TenantContext::clear();
795 TenantContext::set_current(1001);
796
797 let guard = TenantContext::guard().unwrap();
798 let guard_copy = guard; assert_eq!(guard.tenant_id(), 1001);
800 assert_eq!(guard_copy.tenant_id(), 1001);
801
802 TenantContext::clear();
803 }
804
805 #[test]
808 fn test_find_all_auto_filters_tenant() {
809 TenantContext::clear();
810 let r = repo();
811 r.inner.save(make_order(1, 1001, "ORD-001")).unwrap();
812 r.inner.save(make_order(2, 2002, "ORD-002")).unwrap();
813
814 TenantContext::set_current(1001);
815 let all = r.find_all().unwrap();
816 assert_eq!(all.len(), 1);
817 assert_eq!(all[0].tenant_id, 1001);
818 }
819
820 #[test]
823 fn test_save_fails_without_tenant_context() {
824 TenantContext::clear();
825 let r = repo();
826 let order = make_order(0, 0, "ORD-NEW");
827 let result = r.save(order);
828 assert!(matches!(
829 result.unwrap_err().to_string().as_str(),
830 s if s.contains("未设置租户上下文")
831 ));
832 }
833}