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
26pub const TASK_QUEUE_ATTRIBUTE: &str = aion_core::START_TIME_TASK_QUEUE_ATTRIBUTE;
37
38#[derive(Clone, Copy, Debug, Eq, PartialEq)]
42pub(crate) enum GrantSource {
43 NamespacesHeader,
46 TokenClaim,
48 Operator,
54}
55
56impl GrantSource {
57 pub(crate) const fn label(self) -> &'static str {
59 match self {
60 Self::NamespacesHeader => "header",
61 Self::TokenClaim => "token_claim",
62 Self::Operator => "operator",
63 }
64 }
65}
66
67#[derive(Clone, Debug, Eq, PartialEq)]
69pub struct CallerIdentity {
70 subject: String,
71 namespaces: BTreeSet<String>,
72 denial_reason: Option<String>,
73 grant_source: GrantSource,
74 deploy: bool,
77 all_namespaces: bool,
82}
83
84impl CallerIdentity {
85 #[must_use]
88 pub fn new(subject: impl Into<String>, namespaces: impl IntoIterator<Item = String>) -> Self {
89 Self {
90 subject: subject.into(),
91 namespaces: namespaces.into_iter().collect(),
92 denial_reason: None,
93 grant_source: GrantSource::NamespacesHeader,
94 deploy: false,
95 all_namespaces: false,
96 }
97 }
98
99 #[must_use]
103 pub fn from_token_claims(
104 subject: impl Into<String>,
105 namespaces: impl IntoIterator<Item = String>,
106 ) -> Self {
107 Self {
108 subject: subject.into(),
109 namespaces: namespaces.into_iter().collect(),
110 denial_reason: None,
111 grant_source: GrantSource::TokenClaim,
112 deploy: false,
113 all_namespaces: false,
114 }
115 }
116
117 #[must_use]
127 pub fn operator(subject: impl Into<String>) -> Self {
128 Self {
129 subject: subject.into(),
130 namespaces: BTreeSet::new(),
131 denial_reason: None,
132 grant_source: GrantSource::Operator,
133 deploy: true,
134 all_namespaces: true,
135 }
136 }
137
138 #[must_use]
145 pub fn with_deploy(mut self, deploy: bool) -> Self {
146 self.deploy = deploy;
147 self
148 }
149
150 #[must_use]
152 pub const fn deploy_granted(&self) -> bool {
153 self.deploy
154 }
155
156 #[must_use]
158 pub fn denied(subject: impl Into<String>, reason: impl Into<String>) -> Self {
159 Self {
160 subject: subject.into(),
161 namespaces: BTreeSet::new(),
162 denial_reason: Some(reason.into()),
163 grant_source: GrantSource::NamespacesHeader,
164 deploy: false,
165 all_namespaces: false,
166 }
167 }
168
169 #[must_use]
171 pub fn subject(&self) -> &str {
172 &self.subject
173 }
174
175 #[must_use]
180 pub fn namespaces(&self) -> Vec<String> {
181 self.namespaces.iter().cloned().collect()
182 }
183
184 #[must_use]
189 pub const fn all_namespaces(&self) -> bool {
190 self.all_namespaces
191 }
192
193 pub(crate) fn can_access(&self, namespace: &str) -> bool {
202 self.all_namespaces || self.namespaces.contains(namespace)
203 }
204
205 pub(crate) fn denial_reason(&self) -> Option<&str> {
206 self.denial_reason.as_deref()
207 }
208
209 pub(crate) const fn grant_source(&self) -> GrantSource {
212 self.grant_source
213 }
214}
215
216#[derive(Clone)]
218pub struct ScopedEngine {
219 namespace: String,
220 engine: Option<Arc<Engine>>,
221}
222
223impl ScopedEngine {
224 #[must_use]
226 pub fn namespace(&self) -> &str {
227 &self.namespace
228 }
229
230 pub fn engine(&self) -> Result<&Arc<Engine>, ServerError> {
237 self.engine.as_ref().ok_or_else(|| ServerError::Config {
238 message: "namespace resolver has no engine handle".to_owned(),
239 })
240 }
241}
242
243#[derive(Clone, Debug, Eq, PartialEq)]
250pub struct WorkflowAttribution {
251 pub namespace: String,
253 pub workflow_type: Option<String>,
256}
257
258#[async_trait]
264pub trait WorkflowNamespaceSource: Send + Sync {
265 async fn workflow_attribution(
272 &self,
273 workflow_id: &WorkflowId,
274 ) -> Result<Option<WorkflowAttribution>, ServerError>;
275}
276
277struct HistoryNamespaceSource {
281 engine: Arc<Engine>,
282}
283
284#[async_trait]
285impl WorkflowNamespaceSource for HistoryNamespaceSource {
286 async fn workflow_attribution(
287 &self,
288 workflow_id: &WorkflowId,
289 ) -> Result<Option<WorkflowAttribution>, ServerError> {
290 let history = self
291 .engine
292 .store()
293 .read_history(workflow_id)
294 .await
295 .map_err(ServerError::from)?;
296 let namespace = match search_attributes_from_events(&history).remove(NAMESPACE_ATTRIBUTE) {
297 Some(SearchAttributeValue::String(namespace)) => namespace,
298 Some(other) => {
299 return Err(ServerError::Config {
300 message: format!(
301 "workflow {workflow_id} recorded a non-string {NAMESPACE_ATTRIBUTE} search attribute: {other:?}"
302 ),
303 });
304 }
305 None => return Ok(None),
306 };
307 let workflow_type = history.iter().rev().find_map(|event| match event {
310 aion_core::Event::WorkflowStarted { workflow_type, .. } => Some(workflow_type.clone()),
311 _ => None,
312 });
313 Ok(Some(WorkflowAttribution {
314 namespace,
315 workflow_type,
316 }))
317 }
318}
319
320#[derive(Clone, Default)]
323pub struct StaticWorkflowNamespaces {
324 inner: Arc<RwLock<HashMap<WorkflowId, WorkflowAttribution>>>,
325}
326
327impl StaticWorkflowNamespaces {
328 pub fn record(&self, workflow_id: WorkflowId, namespace: &str) -> Result<(), ServerError> {
336 self.insert(
337 workflow_id,
338 WorkflowAttribution {
339 namespace: namespace.to_owned(),
340 workflow_type: None,
341 },
342 )
343 }
344
345 pub fn record_with_type(
352 &self,
353 workflow_id: WorkflowId,
354 namespace: &str,
355 workflow_type: &str,
356 ) -> Result<(), ServerError> {
357 self.insert(
358 workflow_id,
359 WorkflowAttribution {
360 namespace: namespace.to_owned(),
361 workflow_type: Some(workflow_type.to_owned()),
362 },
363 )
364 }
365
366 fn insert(
367 &self,
368 workflow_id: WorkflowId,
369 attribution: WorkflowAttribution,
370 ) -> Result<(), ServerError> {
371 let mut ownership = self
372 .inner
373 .write()
374 .map_err(|_| ServerError::lock_poisoned("namespace workflow ownership"))?;
375 ownership.insert(workflow_id, attribution);
376 Ok(())
377 }
378}
379
380#[async_trait]
381impl WorkflowNamespaceSource for StaticWorkflowNamespaces {
382 async fn workflow_attribution(
383 &self,
384 workflow_id: &WorkflowId,
385 ) -> Result<Option<WorkflowAttribution>, ServerError> {
386 let ownership = self
387 .inner
388 .read()
389 .map_err(|_| ServerError::lock_poisoned("namespace workflow ownership"))?;
390 Ok(ownership.get(workflow_id).cloned())
391 }
392}
393
394#[derive(Clone)]
396pub struct NamespaceResolver {
397 mode: NamespaceMode,
398 engine: Option<Arc<Engine>>,
399 ownership: Arc<dyn WorkflowNamespaceSource>,
400 schedule_ownership: Arc<dyn ScheduleNamespaceSource>,
401}
402
403impl NamespaceResolver {
404 #[must_use]
407 pub fn from_config(config: NamespaceConfig, engine: Arc<Engine>) -> Self {
408 Self {
409 mode: config.mode,
410 ownership: Arc::new(HistoryNamespaceSource {
411 engine: Arc::clone(&engine),
412 }),
413 schedule_ownership: Arc::new(HistoryScheduleNamespaceSource::new(Arc::clone(&engine))),
414 engine: Some(engine),
415 }
416 }
417
418 #[must_use]
420 pub fn from_parts(
421 mode: NamespaceMode,
422 engine: Option<Arc<Engine>>,
423 ownership: Arc<dyn WorkflowNamespaceSource>,
424 schedule_ownership: Arc<dyn ScheduleNamespaceSource>,
425 ) -> Self {
426 Self {
427 mode,
428 engine,
429 ownership,
430 schedule_ownership,
431 }
432 }
433
434 #[must_use]
439 pub fn authorization_only(
440 mode: NamespaceMode,
441 ownership: impl WorkflowNamespaceSource + 'static,
442 schedule_ownership: impl ScheduleNamespaceSource + 'static,
443 ) -> Self {
444 Self::from_parts(
445 mode,
446 None,
447 Arc::new(ownership),
448 Arc::new(schedule_ownership),
449 )
450 }
451
452 #[must_use]
454 pub const fn mode(&self) -> &NamespaceMode {
455 &self.mode
456 }
457
458 pub(crate) fn engine(&self) -> Result<&Arc<Engine>, ServerError> {
465 self.engine.as_ref().ok_or_else(|| ServerError::Config {
466 message: "namespace resolver has no engine handle".to_owned(),
467 })
468 }
469
470 pub fn shutdown_engine(&self) -> Result<(), ServerError> {
477 self.engine
478 .as_ref()
479 .ok_or_else(|| ServerError::Config {
480 message: "namespace resolver has no engine handle".to_owned(),
481 })?
482 .shutdown()
483 .map_err(ServerError::from)
484 }
485
486 pub(super) fn resolve(
494 &self,
495 caller: &CallerIdentity,
496 requested_namespace: &str,
497 ) -> Result<ScopedEngine, ServerError> {
498 if requested_namespace.is_empty() {
499 return Err(ServerError::namespace_denied(
500 "requested namespace must not be empty",
501 ));
502 }
503
504 if let Some(reason) = caller.denial_reason() {
505 return Err(ServerError::namespace_denied(reason));
506 }
507
508 match &self.mode {
509 NamespaceMode::SingleTenant { namespace } if namespace == requested_namespace => {
510 Ok(self.scoped(requested_namespace))
511 }
512 NamespaceMode::SharedEngine if caller.can_access(requested_namespace) => {
513 Ok(self.scoped(requested_namespace))
514 }
515 NamespaceMode::SingleTenant { .. } | NamespaceMode::SharedEngine => {
516 Err(namespace_denied(caller, requested_namespace))
517 }
518 }
519 }
520
521 pub async fn verify_workflow_ownership(
539 &self,
540 namespace: &str,
541 workflow_id: &WorkflowId,
542 ) -> Result<(), ServerError> {
543 match self.workflow_attribution(namespace, workflow_id).await? {
544 Some(_) => Ok(()),
545 None => Err(ServerError::Wire {
546 wire: WireError::not_found(format!("workflow not found in namespace {namespace}")),
547 }),
548 }
549 }
550
551 pub async fn workflow_attribution(
567 &self,
568 namespace: &str,
569 workflow_id: &WorkflowId,
570 ) -> Result<Option<WorkflowAttribution>, ServerError> {
571 Ok(self
572 .ownership
573 .workflow_attribution(workflow_id)
574 .await?
575 .filter(|attribution| attribution.namespace == namespace))
576 }
577
578 pub async fn verify_schedule_ownership(
596 &self,
597 namespace: &str,
598 schedule_id: &ScheduleId,
599 ) -> Result<(), ServerError> {
600 match self
601 .schedule_ownership
602 .schedule_namespace(schedule_id)
603 .await?
604 {
605 Some(owner) if owner == namespace => Ok(()),
606 Some(_) | None => Err(ServerError::Wire {
609 wire: WireError::not_found(format!("schedule not found in namespace {namespace}")),
610 }),
611 }
612 }
613
614 fn scoped(&self, namespace: &str) -> ScopedEngine {
615 ScopedEngine {
616 namespace: namespace.to_owned(),
617 engine: self.engine.clone(),
618 }
619 }
620}
621
622fn namespace_denied(caller: &CallerIdentity, requested_namespace: &str) -> ServerError {
623 let hint = match caller.grant_source {
624 GrantSource::NamespacesHeader => format!(
625 "add {requested_namespace} to x-aion-namespaces for subject `{}` or request a namespace listed in that header",
626 caller.subject()
627 ),
628 GrantSource::TokenClaim => format!(
629 "grant {requested_namespace} in the namespace claim of the token minted for subject `{}` or request a namespace the token grants",
630 caller.subject()
631 ),
632 GrantSource::Operator => format!(
636 "subject `{}` is the operator and already holds every namespace",
637 caller.subject()
638 ),
639 };
640 ServerError::namespace_denied(format!(
641 "subject not authorized for namespace {requested_namespace}; {hint}"
642 ))
643}
644
645#[cfg(test)]
646mod tests {
647 use super::{
648 CallerIdentity, NamespaceResolver, StaticWorkflowNamespaces, WorkflowNamespaceSource,
649 };
650 use crate::config::NamespaceMode;
651 use crate::namespace::StaticScheduleNamespaces;
652 use aion_core::{ScheduleId, WorkflowId};
653
654 fn resolver(mode: NamespaceMode) -> NamespaceResolver {
655 NamespaceResolver::authorization_only(
656 mode,
657 StaticWorkflowNamespaces::default(),
658 StaticScheduleNamespaces::default(),
659 )
660 }
661
662 #[test]
663 fn shared_engine_authorizes_explicit_caller_grant() -> Result<(), Box<dyn std::error::Error>> {
664 let resolver = resolver(NamespaceMode::SharedEngine);
665 let caller = CallerIdentity::new("alice", [String::from("tenant-a")]);
666
667 let scoped = resolver.resolve(&caller, "tenant-a")?;
668
669 assert_eq!(scoped.namespace(), "tenant-a");
670 Ok(())
671 }
672
673 #[test]
678 fn operator_is_authorized_for_any_namespace() -> Result<(), Box<dyn std::error::Error>> {
679 let resolver = resolver(NamespaceMode::SharedEngine);
680 let operator = CallerIdentity::operator("operator");
681
682 assert!(operator.all_namespaces());
683 assert!(operator.deploy_granted());
684 assert!(operator.namespaces().is_empty());
685
686 assert_eq!(
687 resolver.resolve(&operator, "tenant-a")?.namespace(),
688 "tenant-a"
689 );
690 assert_eq!(
691 resolver.resolve(&operator, "tenant-z")?.namespace(),
692 "tenant-z"
693 );
694 Ok(())
695 }
696
697 #[test]
698 fn shared_engine_denies_missing_caller_grant() {
699 let resolver = resolver(NamespaceMode::SharedEngine);
700 let caller = CallerIdentity::new("alice", [String::from("tenant-a")]);
701
702 let denied = resolver.resolve(&caller, "tenant-b");
703
704 assert!(denied.is_err());
705 }
706
707 #[test]
708 fn single_tenant_authorizes_only_configured_namespace() -> Result<(), Box<dyn std::error::Error>>
709 {
710 let resolver = resolver(NamespaceMode::SingleTenant {
711 namespace: String::from("tenant-a"),
712 });
713 let caller = CallerIdentity::new("alice", [String::from("tenant-b")]);
714
715 let scoped = resolver.resolve(&caller, "tenant-a")?;
716 let denied = resolver.resolve(&caller, "tenant-b");
717
718 assert_eq!(scoped.namespace(), "tenant-a");
719 assert!(denied.is_err());
720 Ok(())
721 }
722
723 #[test]
728 fn denial_hint_names_the_grant_source() -> Result<(), Box<dyn std::error::Error>> {
729 let resolver = resolver(NamespaceMode::SharedEngine);
730
731 let header_caller = CallerIdentity::new("alice", [String::from("tenant-a")]);
732 let header_denial = resolver
733 .resolve(&header_caller, "tenant-b")
734 .err()
735 .map(|error| error.to_wire_error())
736 .ok_or("expected header-sourced caller to be denied")?;
737 assert!(
738 header_denial.message.contains("x-aion-namespaces"),
739 "header-path denial must hint the dev header: {}",
740 header_denial.message
741 );
742 assert!(
743 !header_denial.message.contains("namespace claim"),
744 "header-path denial must not hint the token claim: {}",
745 header_denial.message
746 );
747
748 let token_caller = CallerIdentity::from_token_claims("alice", [String::from("tenant-a")]);
749 let token_denial = resolver
750 .resolve(&token_caller, "tenant-b")
751 .err()
752 .map(|error| error.to_wire_error())
753 .ok_or("expected token-sourced caller to be denied")?;
754 assert!(
755 token_denial.message.contains("namespace claim"),
756 "JWT-path denial must hint the token's namespace claim: {}",
757 token_denial.message
758 );
759 assert!(
760 !token_denial.message.contains("x-aion-namespaces"),
761 "JWT-path denial must not hint the dev header: {}",
762 token_denial.message
763 );
764 Ok(())
765 }
766
767 #[test]
768 fn empty_namespace_is_denied_before_scoping() {
769 let resolver = resolver(NamespaceMode::SharedEngine);
770 let caller = CallerIdentity::new("alice", [String::new()]);
771
772 let denied = resolver.resolve(&caller, "");
773
774 assert!(denied.is_err());
775 }
776
777 #[tokio::test]
778 async fn ownership_misses_are_indistinguishable_not_found()
779 -> Result<(), Box<dyn std::error::Error>> {
780 let ownership = StaticWorkflowNamespaces::default();
781 let owned = WorkflowId::new(uuid::Uuid::from_u128(1));
782 let unknown = WorkflowId::new(uuid::Uuid::from_u128(2));
783 ownership.record(owned.clone(), "tenant-a")?;
784 let resolver = NamespaceResolver::authorization_only(
785 NamespaceMode::SharedEngine,
786 ownership,
787 StaticScheduleNamespaces::default(),
788 );
789
790 resolver
791 .verify_workflow_ownership("tenant-a", &owned)
792 .await?;
793
794 let foreign = resolver
798 .verify_workflow_ownership("tenant-b", &owned)
799 .await
800 .err()
801 .map(|error| error.to_wire_error())
802 .ok_or("expected foreign-owned workflow to be rejected")?;
803 let absent = resolver
804 .verify_workflow_ownership("tenant-b", &unknown)
805 .await
806 .err()
807 .map(|error| error.to_wire_error())
808 .ok_or("expected unknown workflow to be rejected")?;
809
810 assert_eq!(foreign.code, aion_proto::WireErrorCode::NotFound);
811 assert_eq!(foreign, absent);
812 assert_eq!(foreign.message, "workflow not found in namespace tenant-b");
813
814 let absent_in_granted = resolver
815 .verify_workflow_ownership("tenant-a", &unknown)
816 .await
817 .err()
818 .map(|error| error.to_wire_error())
819 .ok_or("expected unknown workflow to be rejected in granted namespace")?;
820 assert_eq!(absent_in_granted.code, aion_proto::WireErrorCode::NotFound);
821 assert_eq!(
822 absent_in_granted.message,
823 "workflow not found in namespace tenant-a"
824 );
825 Ok(())
826 }
827
828 #[tokio::test]
829 async fn schedule_ownership_misses_are_indistinguishable_not_found()
830 -> Result<(), Box<dyn std::error::Error>> {
831 let schedule_ownership = StaticScheduleNamespaces::default();
832 let owned = ScheduleId::new(uuid::Uuid::from_u128(1));
833 let unknown = ScheduleId::new(uuid::Uuid::from_u128(2));
834 schedule_ownership.record(owned.clone(), "tenant-a")?;
835 let resolver = NamespaceResolver::authorization_only(
836 NamespaceMode::SharedEngine,
837 StaticWorkflowNamespaces::default(),
838 schedule_ownership,
839 );
840
841 resolver
842 .verify_schedule_ownership("tenant-a", &owned)
843 .await?;
844
845 let foreign = resolver
849 .verify_schedule_ownership("tenant-b", &owned)
850 .await
851 .err()
852 .map(|error| error.to_wire_error())
853 .ok_or("expected foreign-owned schedule to be rejected")?;
854 let absent = resolver
855 .verify_schedule_ownership("tenant-b", &unknown)
856 .await
857 .err()
858 .map(|error| error.to_wire_error())
859 .ok_or("expected unknown schedule to be rejected")?;
860
861 assert_eq!(foreign.code, aion_proto::WireErrorCode::NotFound);
862 assert_eq!(foreign, absent);
863 assert_eq!(foreign.message, "schedule not found in namespace tenant-b");
864
865 let absent_in_granted = resolver
866 .verify_schedule_ownership("tenant-a", &unknown)
867 .await
868 .err()
869 .map(|error| error.to_wire_error())
870 .ok_or("expected unknown schedule to be rejected in granted namespace")?;
871 assert_eq!(absent_in_granted.code, aion_proto::WireErrorCode::NotFound);
872 assert_eq!(
873 absent_in_granted.message,
874 "schedule not found in namespace tenant-a"
875 );
876 Ok(())
877 }
878
879 #[tokio::test]
880 async fn static_source_reports_recorded_namespace() -> Result<(), Box<dyn std::error::Error>> {
881 let ownership = StaticWorkflowNamespaces::default();
882 let workflow_id = WorkflowId::new(uuid::Uuid::from_u128(3));
883 ownership.record(workflow_id.clone(), "tenant-a")?;
884
885 assert_eq!(
886 ownership.workflow_attribution(&workflow_id).await?,
887 Some(super::WorkflowAttribution {
888 namespace: String::from("tenant-a"),
889 workflow_type: None,
890 })
891 );
892 Ok(())
893 }
894
895 #[tokio::test]
896 async fn static_source_reports_recorded_workflow_type() -> Result<(), Box<dyn std::error::Error>>
897 {
898 let ownership = StaticWorkflowNamespaces::default();
899 let workflow_id = WorkflowId::new(uuid::Uuid::from_u128(4));
900 ownership.record_with_type(workflow_id.clone(), "tenant-a", "checkout")?;
901
902 assert_eq!(
903 ownership.workflow_attribution(&workflow_id).await?,
904 Some(super::WorkflowAttribution {
905 namespace: String::from("tenant-a"),
906 workflow_type: Some(String::from("checkout")),
907 })
908 );
909 Ok(())
910 }
911
912 #[tokio::test]
916 async fn scoped_attribution_hides_foreign_and_unknown_identically()
917 -> Result<(), Box<dyn std::error::Error>> {
918 let ownership = StaticWorkflowNamespaces::default();
919 let owned = WorkflowId::new(uuid::Uuid::from_u128(5));
920 let foreign = WorkflowId::new(uuid::Uuid::from_u128(6));
921 let unknown = WorkflowId::new(uuid::Uuid::from_u128(7));
922 ownership.record_with_type(owned.clone(), "tenant-a", "checkout")?;
923 ownership.record_with_type(foreign.clone(), "tenant-b", "checkout")?;
924 let resolver = NamespaceResolver::authorization_only(
925 NamespaceMode::SharedEngine,
926 ownership,
927 StaticScheduleNamespaces::default(),
928 );
929
930 let visible = resolver
931 .workflow_attribution("tenant-a", &owned)
932 .await?
933 .ok_or("owned workflow attribution must be visible")?;
934 assert_eq!(visible.workflow_type.as_deref(), Some("checkout"));
935 assert_eq!(
936 resolver.workflow_attribution("tenant-a", &foreign).await?,
937 None
938 );
939 assert_eq!(
940 resolver.workflow_attribution("tenant-a", &unknown).await?,
941 None
942 );
943 Ok(())
944 }
945}