1use std::collections::{BTreeSet, HashMap};
10use std::sync::{Arc, RwLock};
11
12use aion::Engine;
13use aion_core::{ScheduleId, SearchAttributeValue, WorkflowId, search_attributes_from_events};
14use aion_proto::WireError;
15use async_trait::async_trait;
16
17use crate::config::{NamespaceConfig, NamespaceMode};
18use crate::error::ServerError;
19
20use super::schedule_source::{HistoryScheduleNamespaceSource, ScheduleNamespaceSource};
21
22pub const NAMESPACE_ATTRIBUTE: &str = "aion.namespace";
25
26#[derive(Clone, Copy, Debug, Eq, PartialEq)]
30pub(crate) enum GrantSource {
31 NamespacesHeader,
34 TokenClaim,
36}
37
38impl GrantSource {
39 pub(crate) const fn label(self) -> &'static str {
41 match self {
42 Self::NamespacesHeader => "header",
43 Self::TokenClaim => "token_claim",
44 }
45 }
46}
47
48#[derive(Clone, Debug, Eq, PartialEq)]
50pub struct CallerIdentity {
51 subject: String,
52 namespaces: BTreeSet<String>,
53 denial_reason: Option<String>,
54 grant_source: GrantSource,
55 deploy: bool,
58}
59
60impl CallerIdentity {
61 #[must_use]
64 pub fn new(subject: impl Into<String>, namespaces: impl IntoIterator<Item = String>) -> Self {
65 Self {
66 subject: subject.into(),
67 namespaces: namespaces.into_iter().collect(),
68 denial_reason: None,
69 grant_source: GrantSource::NamespacesHeader,
70 deploy: false,
71 }
72 }
73
74 #[must_use]
78 pub fn from_token_claims(
79 subject: impl Into<String>,
80 namespaces: impl IntoIterator<Item = String>,
81 ) -> Self {
82 Self {
83 subject: subject.into(),
84 namespaces: namespaces.into_iter().collect(),
85 denial_reason: None,
86 grant_source: GrantSource::TokenClaim,
87 deploy: false,
88 }
89 }
90
91 #[must_use]
98 pub fn with_deploy(mut self, deploy: bool) -> Self {
99 self.deploy = deploy;
100 self
101 }
102
103 #[must_use]
105 pub const fn deploy_granted(&self) -> bool {
106 self.deploy
107 }
108
109 #[must_use]
111 pub fn denied(subject: impl Into<String>, reason: impl Into<String>) -> Self {
112 Self {
113 subject: subject.into(),
114 namespaces: BTreeSet::new(),
115 denial_reason: Some(reason.into()),
116 grant_source: GrantSource::NamespacesHeader,
117 deploy: false,
118 }
119 }
120
121 #[must_use]
123 pub fn subject(&self) -> &str {
124 &self.subject
125 }
126
127 fn can_access(&self, namespace: &str) -> bool {
128 self.namespaces.contains(namespace)
129 }
130
131 pub(crate) fn denial_reason(&self) -> Option<&str> {
132 self.denial_reason.as_deref()
133 }
134
135 pub(crate) const fn grant_source(&self) -> GrantSource {
138 self.grant_source
139 }
140}
141
142#[derive(Clone)]
144pub struct ScopedEngine {
145 namespace: String,
146 engine: Option<Arc<Engine>>,
147}
148
149impl ScopedEngine {
150 #[must_use]
152 pub fn namespace(&self) -> &str {
153 &self.namespace
154 }
155
156 pub fn engine(&self) -> Result<&Arc<Engine>, ServerError> {
163 self.engine.as_ref().ok_or_else(|| ServerError::Config {
164 message: "namespace resolver has no engine handle".to_owned(),
165 })
166 }
167}
168
169#[derive(Clone, Debug, Eq, PartialEq)]
176pub struct WorkflowAttribution {
177 pub namespace: String,
179 pub workflow_type: Option<String>,
182}
183
184#[async_trait]
190pub trait WorkflowNamespaceSource: Send + Sync {
191 async fn workflow_attribution(
198 &self,
199 workflow_id: &WorkflowId,
200 ) -> Result<Option<WorkflowAttribution>, ServerError>;
201}
202
203struct HistoryNamespaceSource {
207 engine: Arc<Engine>,
208}
209
210#[async_trait]
211impl WorkflowNamespaceSource for HistoryNamespaceSource {
212 async fn workflow_attribution(
213 &self,
214 workflow_id: &WorkflowId,
215 ) -> Result<Option<WorkflowAttribution>, ServerError> {
216 let history = self
217 .engine
218 .store()
219 .read_history(workflow_id)
220 .await
221 .map_err(ServerError::from)?;
222 let namespace = match search_attributes_from_events(&history).remove(NAMESPACE_ATTRIBUTE) {
223 Some(SearchAttributeValue::String(namespace)) => namespace,
224 Some(other) => {
225 return Err(ServerError::Config {
226 message: format!(
227 "workflow {workflow_id} recorded a non-string {NAMESPACE_ATTRIBUTE} search attribute: {other:?}"
228 ),
229 });
230 }
231 None => return Ok(None),
232 };
233 let workflow_type = history.iter().rev().find_map(|event| match event {
236 aion_core::Event::WorkflowStarted { workflow_type, .. } => Some(workflow_type.clone()),
237 _ => None,
238 });
239 Ok(Some(WorkflowAttribution {
240 namespace,
241 workflow_type,
242 }))
243 }
244}
245
246#[derive(Clone, Default)]
249pub struct StaticWorkflowNamespaces {
250 inner: Arc<RwLock<HashMap<WorkflowId, WorkflowAttribution>>>,
251}
252
253impl StaticWorkflowNamespaces {
254 pub fn record(&self, workflow_id: WorkflowId, namespace: &str) -> Result<(), ServerError> {
262 self.insert(
263 workflow_id,
264 WorkflowAttribution {
265 namespace: namespace.to_owned(),
266 workflow_type: None,
267 },
268 )
269 }
270
271 pub fn record_with_type(
278 &self,
279 workflow_id: WorkflowId,
280 namespace: &str,
281 workflow_type: &str,
282 ) -> Result<(), ServerError> {
283 self.insert(
284 workflow_id,
285 WorkflowAttribution {
286 namespace: namespace.to_owned(),
287 workflow_type: Some(workflow_type.to_owned()),
288 },
289 )
290 }
291
292 fn insert(
293 &self,
294 workflow_id: WorkflowId,
295 attribution: WorkflowAttribution,
296 ) -> Result<(), ServerError> {
297 let mut ownership = self
298 .inner
299 .write()
300 .map_err(|_| ServerError::lock_poisoned("namespace workflow ownership"))?;
301 ownership.insert(workflow_id, attribution);
302 Ok(())
303 }
304}
305
306#[async_trait]
307impl WorkflowNamespaceSource for StaticWorkflowNamespaces {
308 async fn workflow_attribution(
309 &self,
310 workflow_id: &WorkflowId,
311 ) -> Result<Option<WorkflowAttribution>, ServerError> {
312 let ownership = self
313 .inner
314 .read()
315 .map_err(|_| ServerError::lock_poisoned("namespace workflow ownership"))?;
316 Ok(ownership.get(workflow_id).cloned())
317 }
318}
319
320#[derive(Clone)]
322pub struct NamespaceResolver {
323 mode: NamespaceMode,
324 engine: Option<Arc<Engine>>,
325 ownership: Arc<dyn WorkflowNamespaceSource>,
326 schedule_ownership: Arc<dyn ScheduleNamespaceSource>,
327}
328
329impl NamespaceResolver {
330 #[must_use]
333 pub fn from_config(config: NamespaceConfig, engine: Arc<Engine>) -> Self {
334 Self {
335 mode: config.mode,
336 ownership: Arc::new(HistoryNamespaceSource {
337 engine: Arc::clone(&engine),
338 }),
339 schedule_ownership: Arc::new(HistoryScheduleNamespaceSource::new(Arc::clone(&engine))),
340 engine: Some(engine),
341 }
342 }
343
344 #[must_use]
346 pub fn from_parts(
347 mode: NamespaceMode,
348 engine: Option<Arc<Engine>>,
349 ownership: Arc<dyn WorkflowNamespaceSource>,
350 schedule_ownership: Arc<dyn ScheduleNamespaceSource>,
351 ) -> Self {
352 Self {
353 mode,
354 engine,
355 ownership,
356 schedule_ownership,
357 }
358 }
359
360 #[must_use]
365 pub fn authorization_only(
366 mode: NamespaceMode,
367 ownership: impl WorkflowNamespaceSource + 'static,
368 schedule_ownership: impl ScheduleNamespaceSource + 'static,
369 ) -> Self {
370 Self::from_parts(
371 mode,
372 None,
373 Arc::new(ownership),
374 Arc::new(schedule_ownership),
375 )
376 }
377
378 #[must_use]
380 pub const fn mode(&self) -> &NamespaceMode {
381 &self.mode
382 }
383
384 pub(crate) fn engine(&self) -> Result<&Arc<Engine>, ServerError> {
391 self.engine.as_ref().ok_or_else(|| ServerError::Config {
392 message: "namespace resolver has no engine handle".to_owned(),
393 })
394 }
395
396 pub fn shutdown_engine(&self) -> Result<(), ServerError> {
403 self.engine
404 .as_ref()
405 .ok_or_else(|| ServerError::Config {
406 message: "namespace resolver has no engine handle".to_owned(),
407 })?
408 .shutdown()
409 .map_err(ServerError::from)
410 }
411
412 pub(super) fn resolve(
420 &self,
421 caller: &CallerIdentity,
422 requested_namespace: &str,
423 ) -> Result<ScopedEngine, ServerError> {
424 if requested_namespace.is_empty() {
425 return Err(ServerError::namespace_denied(
426 "requested namespace must not be empty",
427 ));
428 }
429
430 if let Some(reason) = caller.denial_reason() {
431 return Err(ServerError::namespace_denied(reason));
432 }
433
434 match &self.mode {
435 NamespaceMode::SingleTenant { namespace } if namespace == requested_namespace => {
436 Ok(self.scoped(requested_namespace))
437 }
438 NamespaceMode::SharedEngine if caller.can_access(requested_namespace) => {
439 Ok(self.scoped(requested_namespace))
440 }
441 NamespaceMode::SingleTenant { .. } | NamespaceMode::SharedEngine => {
442 Err(namespace_denied(caller, requested_namespace))
443 }
444 }
445 }
446
447 pub async fn verify_workflow_ownership(
465 &self,
466 namespace: &str,
467 workflow_id: &WorkflowId,
468 ) -> Result<(), ServerError> {
469 match self.workflow_attribution(namespace, workflow_id).await? {
470 Some(_) => Ok(()),
471 None => Err(ServerError::Wire {
472 wire: WireError::not_found(format!("workflow not found in namespace {namespace}")),
473 }),
474 }
475 }
476
477 pub async fn workflow_attribution(
493 &self,
494 namespace: &str,
495 workflow_id: &WorkflowId,
496 ) -> Result<Option<WorkflowAttribution>, ServerError> {
497 Ok(self
498 .ownership
499 .workflow_attribution(workflow_id)
500 .await?
501 .filter(|attribution| attribution.namespace == namespace))
502 }
503
504 pub async fn verify_schedule_ownership(
522 &self,
523 namespace: &str,
524 schedule_id: &ScheduleId,
525 ) -> Result<(), ServerError> {
526 match self
527 .schedule_ownership
528 .schedule_namespace(schedule_id)
529 .await?
530 {
531 Some(owner) if owner == namespace => Ok(()),
532 Some(_) | None => Err(ServerError::Wire {
535 wire: WireError::not_found(format!("schedule not found in namespace {namespace}")),
536 }),
537 }
538 }
539
540 fn scoped(&self, namespace: &str) -> ScopedEngine {
541 ScopedEngine {
542 namespace: namespace.to_owned(),
543 engine: self.engine.clone(),
544 }
545 }
546}
547
548fn namespace_denied(caller: &CallerIdentity, requested_namespace: &str) -> ServerError {
549 let hint = match caller.grant_source {
550 GrantSource::NamespacesHeader => format!(
551 "add {requested_namespace} to x-aion-namespaces for subject `{}` or request a namespace listed in that header",
552 caller.subject()
553 ),
554 GrantSource::TokenClaim => format!(
555 "grant {requested_namespace} in the namespace claim of the token minted for subject `{}` or request a namespace the token grants",
556 caller.subject()
557 ),
558 };
559 ServerError::namespace_denied(format!(
560 "subject not authorized for namespace {requested_namespace}; {hint}"
561 ))
562}
563
564#[cfg(test)]
565mod tests {
566 use super::{
567 CallerIdentity, NamespaceResolver, StaticWorkflowNamespaces, WorkflowNamespaceSource,
568 };
569 use crate::config::NamespaceMode;
570 use crate::namespace::StaticScheduleNamespaces;
571 use aion_core::{ScheduleId, WorkflowId};
572
573 fn resolver(mode: NamespaceMode) -> NamespaceResolver {
574 NamespaceResolver::authorization_only(
575 mode,
576 StaticWorkflowNamespaces::default(),
577 StaticScheduleNamespaces::default(),
578 )
579 }
580
581 #[test]
582 fn shared_engine_authorizes_explicit_caller_grant() -> Result<(), Box<dyn std::error::Error>> {
583 let resolver = resolver(NamespaceMode::SharedEngine);
584 let caller = CallerIdentity::new("alice", [String::from("tenant-a")]);
585
586 let scoped = resolver.resolve(&caller, "tenant-a")?;
587
588 assert_eq!(scoped.namespace(), "tenant-a");
589 Ok(())
590 }
591
592 #[test]
593 fn shared_engine_denies_missing_caller_grant() {
594 let resolver = resolver(NamespaceMode::SharedEngine);
595 let caller = CallerIdentity::new("alice", [String::from("tenant-a")]);
596
597 let denied = resolver.resolve(&caller, "tenant-b");
598
599 assert!(denied.is_err());
600 }
601
602 #[test]
603 fn single_tenant_authorizes_only_configured_namespace() -> Result<(), Box<dyn std::error::Error>>
604 {
605 let resolver = resolver(NamespaceMode::SingleTenant {
606 namespace: String::from("tenant-a"),
607 });
608 let caller = CallerIdentity::new("alice", [String::from("tenant-b")]);
609
610 let scoped = resolver.resolve(&caller, "tenant-a")?;
611 let denied = resolver.resolve(&caller, "tenant-b");
612
613 assert_eq!(scoped.namespace(), "tenant-a");
614 assert!(denied.is_err());
615 Ok(())
616 }
617
618 #[test]
623 fn denial_hint_names_the_grant_source() -> Result<(), Box<dyn std::error::Error>> {
624 let resolver = resolver(NamespaceMode::SharedEngine);
625
626 let header_caller = CallerIdentity::new("alice", [String::from("tenant-a")]);
627 let header_denial = resolver
628 .resolve(&header_caller, "tenant-b")
629 .err()
630 .map(|error| error.to_wire_error())
631 .ok_or("expected header-sourced caller to be denied")?;
632 assert!(
633 header_denial.message.contains("x-aion-namespaces"),
634 "header-path denial must hint the dev header: {}",
635 header_denial.message
636 );
637 assert!(
638 !header_denial.message.contains("namespace claim"),
639 "header-path denial must not hint the token claim: {}",
640 header_denial.message
641 );
642
643 let token_caller = CallerIdentity::from_token_claims("alice", [String::from("tenant-a")]);
644 let token_denial = resolver
645 .resolve(&token_caller, "tenant-b")
646 .err()
647 .map(|error| error.to_wire_error())
648 .ok_or("expected token-sourced caller to be denied")?;
649 assert!(
650 token_denial.message.contains("namespace claim"),
651 "JWT-path denial must hint the token's namespace claim: {}",
652 token_denial.message
653 );
654 assert!(
655 !token_denial.message.contains("x-aion-namespaces"),
656 "JWT-path denial must not hint the dev header: {}",
657 token_denial.message
658 );
659 Ok(())
660 }
661
662 #[test]
663 fn empty_namespace_is_denied_before_scoping() {
664 let resolver = resolver(NamespaceMode::SharedEngine);
665 let caller = CallerIdentity::new("alice", [String::new()]);
666
667 let denied = resolver.resolve(&caller, "");
668
669 assert!(denied.is_err());
670 }
671
672 #[tokio::test]
673 async fn ownership_misses_are_indistinguishable_not_found()
674 -> Result<(), Box<dyn std::error::Error>> {
675 let ownership = StaticWorkflowNamespaces::default();
676 let owned = WorkflowId::new(uuid::Uuid::from_u128(1));
677 let unknown = WorkflowId::new(uuid::Uuid::from_u128(2));
678 ownership.record(owned.clone(), "tenant-a")?;
679 let resolver = NamespaceResolver::authorization_only(
680 NamespaceMode::SharedEngine,
681 ownership,
682 StaticScheduleNamespaces::default(),
683 );
684
685 resolver
686 .verify_workflow_ownership("tenant-a", &owned)
687 .await?;
688
689 let foreign = resolver
693 .verify_workflow_ownership("tenant-b", &owned)
694 .await
695 .err()
696 .map(|error| error.to_wire_error())
697 .ok_or("expected foreign-owned workflow to be rejected")?;
698 let absent = resolver
699 .verify_workflow_ownership("tenant-b", &unknown)
700 .await
701 .err()
702 .map(|error| error.to_wire_error())
703 .ok_or("expected unknown workflow to be rejected")?;
704
705 assert_eq!(foreign.code, aion_proto::WireErrorCode::NotFound);
706 assert_eq!(foreign, absent);
707 assert_eq!(foreign.message, "workflow not found in namespace tenant-b");
708
709 let absent_in_granted = resolver
710 .verify_workflow_ownership("tenant-a", &unknown)
711 .await
712 .err()
713 .map(|error| error.to_wire_error())
714 .ok_or("expected unknown workflow to be rejected in granted namespace")?;
715 assert_eq!(absent_in_granted.code, aion_proto::WireErrorCode::NotFound);
716 assert_eq!(
717 absent_in_granted.message,
718 "workflow not found in namespace tenant-a"
719 );
720 Ok(())
721 }
722
723 #[tokio::test]
724 async fn schedule_ownership_misses_are_indistinguishable_not_found()
725 -> Result<(), Box<dyn std::error::Error>> {
726 let schedule_ownership = StaticScheduleNamespaces::default();
727 let owned = ScheduleId::new(uuid::Uuid::from_u128(1));
728 let unknown = ScheduleId::new(uuid::Uuid::from_u128(2));
729 schedule_ownership.record(owned.clone(), "tenant-a")?;
730 let resolver = NamespaceResolver::authorization_only(
731 NamespaceMode::SharedEngine,
732 StaticWorkflowNamespaces::default(),
733 schedule_ownership,
734 );
735
736 resolver
737 .verify_schedule_ownership("tenant-a", &owned)
738 .await?;
739
740 let foreign = resolver
744 .verify_schedule_ownership("tenant-b", &owned)
745 .await
746 .err()
747 .map(|error| error.to_wire_error())
748 .ok_or("expected foreign-owned schedule to be rejected")?;
749 let absent = resolver
750 .verify_schedule_ownership("tenant-b", &unknown)
751 .await
752 .err()
753 .map(|error| error.to_wire_error())
754 .ok_or("expected unknown schedule to be rejected")?;
755
756 assert_eq!(foreign.code, aion_proto::WireErrorCode::NotFound);
757 assert_eq!(foreign, absent);
758 assert_eq!(foreign.message, "schedule not found in namespace tenant-b");
759
760 let absent_in_granted = resolver
761 .verify_schedule_ownership("tenant-a", &unknown)
762 .await
763 .err()
764 .map(|error| error.to_wire_error())
765 .ok_or("expected unknown schedule to be rejected in granted namespace")?;
766 assert_eq!(absent_in_granted.code, aion_proto::WireErrorCode::NotFound);
767 assert_eq!(
768 absent_in_granted.message,
769 "schedule not found in namespace tenant-a"
770 );
771 Ok(())
772 }
773
774 #[tokio::test]
775 async fn static_source_reports_recorded_namespace() -> Result<(), Box<dyn std::error::Error>> {
776 let ownership = StaticWorkflowNamespaces::default();
777 let workflow_id = WorkflowId::new(uuid::Uuid::from_u128(3));
778 ownership.record(workflow_id.clone(), "tenant-a")?;
779
780 assert_eq!(
781 ownership.workflow_attribution(&workflow_id).await?,
782 Some(super::WorkflowAttribution {
783 namespace: String::from("tenant-a"),
784 workflow_type: None,
785 })
786 );
787 Ok(())
788 }
789
790 #[tokio::test]
791 async fn static_source_reports_recorded_workflow_type() -> Result<(), Box<dyn std::error::Error>>
792 {
793 let ownership = StaticWorkflowNamespaces::default();
794 let workflow_id = WorkflowId::new(uuid::Uuid::from_u128(4));
795 ownership.record_with_type(workflow_id.clone(), "tenant-a", "checkout")?;
796
797 assert_eq!(
798 ownership.workflow_attribution(&workflow_id).await?,
799 Some(super::WorkflowAttribution {
800 namespace: String::from("tenant-a"),
801 workflow_type: Some(String::from("checkout")),
802 })
803 );
804 Ok(())
805 }
806
807 #[tokio::test]
811 async fn scoped_attribution_hides_foreign_and_unknown_identically()
812 -> Result<(), Box<dyn std::error::Error>> {
813 let ownership = StaticWorkflowNamespaces::default();
814 let owned = WorkflowId::new(uuid::Uuid::from_u128(5));
815 let foreign = WorkflowId::new(uuid::Uuid::from_u128(6));
816 let unknown = WorkflowId::new(uuid::Uuid::from_u128(7));
817 ownership.record_with_type(owned.clone(), "tenant-a", "checkout")?;
818 ownership.record_with_type(foreign.clone(), "tenant-b", "checkout")?;
819 let resolver = NamespaceResolver::authorization_only(
820 NamespaceMode::SharedEngine,
821 ownership,
822 StaticScheduleNamespaces::default(),
823 );
824
825 let visible = resolver
826 .workflow_attribution("tenant-a", &owned)
827 .await?
828 .ok_or("owned workflow attribution must be visible")?;
829 assert_eq!(visible.workflow_type.as_deref(), Some("checkout"));
830 assert_eq!(
831 resolver.workflow_attribution("tenant-a", &foreign).await?,
832 None
833 );
834 assert_eq!(
835 resolver.workflow_attribution("tenant-a", &unknown).await?,
836 None
837 );
838 Ok(())
839 }
840}