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 #[must_use]
132 pub fn namespaces(&self) -> Vec<String> {
133 self.namespaces.iter().cloned().collect()
134 }
135
136 fn can_access(&self, namespace: &str) -> bool {
137 self.namespaces.contains(namespace)
138 }
139
140 pub(crate) fn denial_reason(&self) -> Option<&str> {
141 self.denial_reason.as_deref()
142 }
143
144 pub(crate) const fn grant_source(&self) -> GrantSource {
147 self.grant_source
148 }
149}
150
151#[derive(Clone)]
153pub struct ScopedEngine {
154 namespace: String,
155 engine: Option<Arc<Engine>>,
156}
157
158impl ScopedEngine {
159 #[must_use]
161 pub fn namespace(&self) -> &str {
162 &self.namespace
163 }
164
165 pub fn engine(&self) -> Result<&Arc<Engine>, ServerError> {
172 self.engine.as_ref().ok_or_else(|| ServerError::Config {
173 message: "namespace resolver has no engine handle".to_owned(),
174 })
175 }
176}
177
178#[derive(Clone, Debug, Eq, PartialEq)]
185pub struct WorkflowAttribution {
186 pub namespace: String,
188 pub workflow_type: Option<String>,
191}
192
193#[async_trait]
199pub trait WorkflowNamespaceSource: Send + Sync {
200 async fn workflow_attribution(
207 &self,
208 workflow_id: &WorkflowId,
209 ) -> Result<Option<WorkflowAttribution>, ServerError>;
210}
211
212struct HistoryNamespaceSource {
216 engine: Arc<Engine>,
217}
218
219#[async_trait]
220impl WorkflowNamespaceSource for HistoryNamespaceSource {
221 async fn workflow_attribution(
222 &self,
223 workflow_id: &WorkflowId,
224 ) -> Result<Option<WorkflowAttribution>, ServerError> {
225 let history = self
226 .engine
227 .store()
228 .read_history(workflow_id)
229 .await
230 .map_err(ServerError::from)?;
231 let namespace = match search_attributes_from_events(&history).remove(NAMESPACE_ATTRIBUTE) {
232 Some(SearchAttributeValue::String(namespace)) => namespace,
233 Some(other) => {
234 return Err(ServerError::Config {
235 message: format!(
236 "workflow {workflow_id} recorded a non-string {NAMESPACE_ATTRIBUTE} search attribute: {other:?}"
237 ),
238 });
239 }
240 None => return Ok(None),
241 };
242 let workflow_type = history.iter().rev().find_map(|event| match event {
245 aion_core::Event::WorkflowStarted { workflow_type, .. } => Some(workflow_type.clone()),
246 _ => None,
247 });
248 Ok(Some(WorkflowAttribution {
249 namespace,
250 workflow_type,
251 }))
252 }
253}
254
255#[derive(Clone, Default)]
258pub struct StaticWorkflowNamespaces {
259 inner: Arc<RwLock<HashMap<WorkflowId, WorkflowAttribution>>>,
260}
261
262impl StaticWorkflowNamespaces {
263 pub fn record(&self, workflow_id: WorkflowId, namespace: &str) -> Result<(), ServerError> {
271 self.insert(
272 workflow_id,
273 WorkflowAttribution {
274 namespace: namespace.to_owned(),
275 workflow_type: None,
276 },
277 )
278 }
279
280 pub fn record_with_type(
287 &self,
288 workflow_id: WorkflowId,
289 namespace: &str,
290 workflow_type: &str,
291 ) -> Result<(), ServerError> {
292 self.insert(
293 workflow_id,
294 WorkflowAttribution {
295 namespace: namespace.to_owned(),
296 workflow_type: Some(workflow_type.to_owned()),
297 },
298 )
299 }
300
301 fn insert(
302 &self,
303 workflow_id: WorkflowId,
304 attribution: WorkflowAttribution,
305 ) -> Result<(), ServerError> {
306 let mut ownership = self
307 .inner
308 .write()
309 .map_err(|_| ServerError::lock_poisoned("namespace workflow ownership"))?;
310 ownership.insert(workflow_id, attribution);
311 Ok(())
312 }
313}
314
315#[async_trait]
316impl WorkflowNamespaceSource for StaticWorkflowNamespaces {
317 async fn workflow_attribution(
318 &self,
319 workflow_id: &WorkflowId,
320 ) -> Result<Option<WorkflowAttribution>, ServerError> {
321 let ownership = self
322 .inner
323 .read()
324 .map_err(|_| ServerError::lock_poisoned("namespace workflow ownership"))?;
325 Ok(ownership.get(workflow_id).cloned())
326 }
327}
328
329#[derive(Clone)]
331pub struct NamespaceResolver {
332 mode: NamespaceMode,
333 engine: Option<Arc<Engine>>,
334 ownership: Arc<dyn WorkflowNamespaceSource>,
335 schedule_ownership: Arc<dyn ScheduleNamespaceSource>,
336}
337
338impl NamespaceResolver {
339 #[must_use]
342 pub fn from_config(config: NamespaceConfig, engine: Arc<Engine>) -> Self {
343 Self {
344 mode: config.mode,
345 ownership: Arc::new(HistoryNamespaceSource {
346 engine: Arc::clone(&engine),
347 }),
348 schedule_ownership: Arc::new(HistoryScheduleNamespaceSource::new(Arc::clone(&engine))),
349 engine: Some(engine),
350 }
351 }
352
353 #[must_use]
355 pub fn from_parts(
356 mode: NamespaceMode,
357 engine: Option<Arc<Engine>>,
358 ownership: Arc<dyn WorkflowNamespaceSource>,
359 schedule_ownership: Arc<dyn ScheduleNamespaceSource>,
360 ) -> Self {
361 Self {
362 mode,
363 engine,
364 ownership,
365 schedule_ownership,
366 }
367 }
368
369 #[must_use]
374 pub fn authorization_only(
375 mode: NamespaceMode,
376 ownership: impl WorkflowNamespaceSource + 'static,
377 schedule_ownership: impl ScheduleNamespaceSource + 'static,
378 ) -> Self {
379 Self::from_parts(
380 mode,
381 None,
382 Arc::new(ownership),
383 Arc::new(schedule_ownership),
384 )
385 }
386
387 #[must_use]
389 pub const fn mode(&self) -> &NamespaceMode {
390 &self.mode
391 }
392
393 pub(crate) fn engine(&self) -> Result<&Arc<Engine>, ServerError> {
400 self.engine.as_ref().ok_or_else(|| ServerError::Config {
401 message: "namespace resolver has no engine handle".to_owned(),
402 })
403 }
404
405 pub fn shutdown_engine(&self) -> Result<(), ServerError> {
412 self.engine
413 .as_ref()
414 .ok_or_else(|| ServerError::Config {
415 message: "namespace resolver has no engine handle".to_owned(),
416 })?
417 .shutdown()
418 .map_err(ServerError::from)
419 }
420
421 pub(super) fn resolve(
429 &self,
430 caller: &CallerIdentity,
431 requested_namespace: &str,
432 ) -> Result<ScopedEngine, ServerError> {
433 if requested_namespace.is_empty() {
434 return Err(ServerError::namespace_denied(
435 "requested namespace must not be empty",
436 ));
437 }
438
439 if let Some(reason) = caller.denial_reason() {
440 return Err(ServerError::namespace_denied(reason));
441 }
442
443 match &self.mode {
444 NamespaceMode::SingleTenant { namespace } if namespace == requested_namespace => {
445 Ok(self.scoped(requested_namespace))
446 }
447 NamespaceMode::SharedEngine if caller.can_access(requested_namespace) => {
448 Ok(self.scoped(requested_namespace))
449 }
450 NamespaceMode::SingleTenant { .. } | NamespaceMode::SharedEngine => {
451 Err(namespace_denied(caller, requested_namespace))
452 }
453 }
454 }
455
456 pub async fn verify_workflow_ownership(
474 &self,
475 namespace: &str,
476 workflow_id: &WorkflowId,
477 ) -> Result<(), ServerError> {
478 match self.workflow_attribution(namespace, workflow_id).await? {
479 Some(_) => Ok(()),
480 None => Err(ServerError::Wire {
481 wire: WireError::not_found(format!("workflow not found in namespace {namespace}")),
482 }),
483 }
484 }
485
486 pub async fn workflow_attribution(
502 &self,
503 namespace: &str,
504 workflow_id: &WorkflowId,
505 ) -> Result<Option<WorkflowAttribution>, ServerError> {
506 Ok(self
507 .ownership
508 .workflow_attribution(workflow_id)
509 .await?
510 .filter(|attribution| attribution.namespace == namespace))
511 }
512
513 pub async fn verify_schedule_ownership(
531 &self,
532 namespace: &str,
533 schedule_id: &ScheduleId,
534 ) -> Result<(), ServerError> {
535 match self
536 .schedule_ownership
537 .schedule_namespace(schedule_id)
538 .await?
539 {
540 Some(owner) if owner == namespace => Ok(()),
541 Some(_) | None => Err(ServerError::Wire {
544 wire: WireError::not_found(format!("schedule not found in namespace {namespace}")),
545 }),
546 }
547 }
548
549 fn scoped(&self, namespace: &str) -> ScopedEngine {
550 ScopedEngine {
551 namespace: namespace.to_owned(),
552 engine: self.engine.clone(),
553 }
554 }
555}
556
557fn namespace_denied(caller: &CallerIdentity, requested_namespace: &str) -> ServerError {
558 let hint = match caller.grant_source {
559 GrantSource::NamespacesHeader => format!(
560 "add {requested_namespace} to x-aion-namespaces for subject `{}` or request a namespace listed in that header",
561 caller.subject()
562 ),
563 GrantSource::TokenClaim => format!(
564 "grant {requested_namespace} in the namespace claim of the token minted for subject `{}` or request a namespace the token grants",
565 caller.subject()
566 ),
567 };
568 ServerError::namespace_denied(format!(
569 "subject not authorized for namespace {requested_namespace}; {hint}"
570 ))
571}
572
573#[cfg(test)]
574mod tests {
575 use super::{
576 CallerIdentity, NamespaceResolver, StaticWorkflowNamespaces, WorkflowNamespaceSource,
577 };
578 use crate::config::NamespaceMode;
579 use crate::namespace::StaticScheduleNamespaces;
580 use aion_core::{ScheduleId, WorkflowId};
581
582 fn resolver(mode: NamespaceMode) -> NamespaceResolver {
583 NamespaceResolver::authorization_only(
584 mode,
585 StaticWorkflowNamespaces::default(),
586 StaticScheduleNamespaces::default(),
587 )
588 }
589
590 #[test]
591 fn shared_engine_authorizes_explicit_caller_grant() -> Result<(), Box<dyn std::error::Error>> {
592 let resolver = resolver(NamespaceMode::SharedEngine);
593 let caller = CallerIdentity::new("alice", [String::from("tenant-a")]);
594
595 let scoped = resolver.resolve(&caller, "tenant-a")?;
596
597 assert_eq!(scoped.namespace(), "tenant-a");
598 Ok(())
599 }
600
601 #[test]
602 fn shared_engine_denies_missing_caller_grant() {
603 let resolver = resolver(NamespaceMode::SharedEngine);
604 let caller = CallerIdentity::new("alice", [String::from("tenant-a")]);
605
606 let denied = resolver.resolve(&caller, "tenant-b");
607
608 assert!(denied.is_err());
609 }
610
611 #[test]
612 fn single_tenant_authorizes_only_configured_namespace() -> Result<(), Box<dyn std::error::Error>>
613 {
614 let resolver = resolver(NamespaceMode::SingleTenant {
615 namespace: String::from("tenant-a"),
616 });
617 let caller = CallerIdentity::new("alice", [String::from("tenant-b")]);
618
619 let scoped = resolver.resolve(&caller, "tenant-a")?;
620 let denied = resolver.resolve(&caller, "tenant-b");
621
622 assert_eq!(scoped.namespace(), "tenant-a");
623 assert!(denied.is_err());
624 Ok(())
625 }
626
627 #[test]
632 fn denial_hint_names_the_grant_source() -> Result<(), Box<dyn std::error::Error>> {
633 let resolver = resolver(NamespaceMode::SharedEngine);
634
635 let header_caller = CallerIdentity::new("alice", [String::from("tenant-a")]);
636 let header_denial = resolver
637 .resolve(&header_caller, "tenant-b")
638 .err()
639 .map(|error| error.to_wire_error())
640 .ok_or("expected header-sourced caller to be denied")?;
641 assert!(
642 header_denial.message.contains("x-aion-namespaces"),
643 "header-path denial must hint the dev header: {}",
644 header_denial.message
645 );
646 assert!(
647 !header_denial.message.contains("namespace claim"),
648 "header-path denial must not hint the token claim: {}",
649 header_denial.message
650 );
651
652 let token_caller = CallerIdentity::from_token_claims("alice", [String::from("tenant-a")]);
653 let token_denial = resolver
654 .resolve(&token_caller, "tenant-b")
655 .err()
656 .map(|error| error.to_wire_error())
657 .ok_or("expected token-sourced caller to be denied")?;
658 assert!(
659 token_denial.message.contains("namespace claim"),
660 "JWT-path denial must hint the token's namespace claim: {}",
661 token_denial.message
662 );
663 assert!(
664 !token_denial.message.contains("x-aion-namespaces"),
665 "JWT-path denial must not hint the dev header: {}",
666 token_denial.message
667 );
668 Ok(())
669 }
670
671 #[test]
672 fn empty_namespace_is_denied_before_scoping() {
673 let resolver = resolver(NamespaceMode::SharedEngine);
674 let caller = CallerIdentity::new("alice", [String::new()]);
675
676 let denied = resolver.resolve(&caller, "");
677
678 assert!(denied.is_err());
679 }
680
681 #[tokio::test]
682 async fn ownership_misses_are_indistinguishable_not_found()
683 -> Result<(), Box<dyn std::error::Error>> {
684 let ownership = StaticWorkflowNamespaces::default();
685 let owned = WorkflowId::new(uuid::Uuid::from_u128(1));
686 let unknown = WorkflowId::new(uuid::Uuid::from_u128(2));
687 ownership.record(owned.clone(), "tenant-a")?;
688 let resolver = NamespaceResolver::authorization_only(
689 NamespaceMode::SharedEngine,
690 ownership,
691 StaticScheduleNamespaces::default(),
692 );
693
694 resolver
695 .verify_workflow_ownership("tenant-a", &owned)
696 .await?;
697
698 let foreign = resolver
702 .verify_workflow_ownership("tenant-b", &owned)
703 .await
704 .err()
705 .map(|error| error.to_wire_error())
706 .ok_or("expected foreign-owned workflow to be rejected")?;
707 let absent = resolver
708 .verify_workflow_ownership("tenant-b", &unknown)
709 .await
710 .err()
711 .map(|error| error.to_wire_error())
712 .ok_or("expected unknown workflow to be rejected")?;
713
714 assert_eq!(foreign.code, aion_proto::WireErrorCode::NotFound);
715 assert_eq!(foreign, absent);
716 assert_eq!(foreign.message, "workflow not found in namespace tenant-b");
717
718 let absent_in_granted = resolver
719 .verify_workflow_ownership("tenant-a", &unknown)
720 .await
721 .err()
722 .map(|error| error.to_wire_error())
723 .ok_or("expected unknown workflow to be rejected in granted namespace")?;
724 assert_eq!(absent_in_granted.code, aion_proto::WireErrorCode::NotFound);
725 assert_eq!(
726 absent_in_granted.message,
727 "workflow not found in namespace tenant-a"
728 );
729 Ok(())
730 }
731
732 #[tokio::test]
733 async fn schedule_ownership_misses_are_indistinguishable_not_found()
734 -> Result<(), Box<dyn std::error::Error>> {
735 let schedule_ownership = StaticScheduleNamespaces::default();
736 let owned = ScheduleId::new(uuid::Uuid::from_u128(1));
737 let unknown = ScheduleId::new(uuid::Uuid::from_u128(2));
738 schedule_ownership.record(owned.clone(), "tenant-a")?;
739 let resolver = NamespaceResolver::authorization_only(
740 NamespaceMode::SharedEngine,
741 StaticWorkflowNamespaces::default(),
742 schedule_ownership,
743 );
744
745 resolver
746 .verify_schedule_ownership("tenant-a", &owned)
747 .await?;
748
749 let foreign = resolver
753 .verify_schedule_ownership("tenant-b", &owned)
754 .await
755 .err()
756 .map(|error| error.to_wire_error())
757 .ok_or("expected foreign-owned schedule to be rejected")?;
758 let absent = resolver
759 .verify_schedule_ownership("tenant-b", &unknown)
760 .await
761 .err()
762 .map(|error| error.to_wire_error())
763 .ok_or("expected unknown schedule to be rejected")?;
764
765 assert_eq!(foreign.code, aion_proto::WireErrorCode::NotFound);
766 assert_eq!(foreign, absent);
767 assert_eq!(foreign.message, "schedule not found in namespace tenant-b");
768
769 let absent_in_granted = resolver
770 .verify_schedule_ownership("tenant-a", &unknown)
771 .await
772 .err()
773 .map(|error| error.to_wire_error())
774 .ok_or("expected unknown schedule to be rejected in granted namespace")?;
775 assert_eq!(absent_in_granted.code, aion_proto::WireErrorCode::NotFound);
776 assert_eq!(
777 absent_in_granted.message,
778 "schedule not found in namespace tenant-a"
779 );
780 Ok(())
781 }
782
783 #[tokio::test]
784 async fn static_source_reports_recorded_namespace() -> Result<(), Box<dyn std::error::Error>> {
785 let ownership = StaticWorkflowNamespaces::default();
786 let workflow_id = WorkflowId::new(uuid::Uuid::from_u128(3));
787 ownership.record(workflow_id.clone(), "tenant-a")?;
788
789 assert_eq!(
790 ownership.workflow_attribution(&workflow_id).await?,
791 Some(super::WorkflowAttribution {
792 namespace: String::from("tenant-a"),
793 workflow_type: None,
794 })
795 );
796 Ok(())
797 }
798
799 #[tokio::test]
800 async fn static_source_reports_recorded_workflow_type() -> Result<(), Box<dyn std::error::Error>>
801 {
802 let ownership = StaticWorkflowNamespaces::default();
803 let workflow_id = WorkflowId::new(uuid::Uuid::from_u128(4));
804 ownership.record_with_type(workflow_id.clone(), "tenant-a", "checkout")?;
805
806 assert_eq!(
807 ownership.workflow_attribution(&workflow_id).await?,
808 Some(super::WorkflowAttribution {
809 namespace: String::from("tenant-a"),
810 workflow_type: Some(String::from("checkout")),
811 })
812 );
813 Ok(())
814 }
815
816 #[tokio::test]
820 async fn scoped_attribution_hides_foreign_and_unknown_identically()
821 -> Result<(), Box<dyn std::error::Error>> {
822 let ownership = StaticWorkflowNamespaces::default();
823 let owned = WorkflowId::new(uuid::Uuid::from_u128(5));
824 let foreign = WorkflowId::new(uuid::Uuid::from_u128(6));
825 let unknown = WorkflowId::new(uuid::Uuid::from_u128(7));
826 ownership.record_with_type(owned.clone(), "tenant-a", "checkout")?;
827 ownership.record_with_type(foreign.clone(), "tenant-b", "checkout")?;
828 let resolver = NamespaceResolver::authorization_only(
829 NamespaceMode::SharedEngine,
830 ownership,
831 StaticScheduleNamespaces::default(),
832 );
833
834 let visible = resolver
835 .workflow_attribution("tenant-a", &owned)
836 .await?
837 .ok_or("owned workflow attribution must be visible")?;
838 assert_eq!(visible.workflow_type.as_deref(), Some("checkout"));
839 assert_eq!(
840 resolver.workflow_attribution("tenant-a", &foreign).await?,
841 None
842 );
843 assert_eq!(
844 resolver.workflow_attribution("tenant-a", &unknown).await?,
845 None
846 );
847 Ok(())
848 }
849}