Skip to main content

treetop_core/
policy_store.rs

1//! Policy-store configuration and namespace-based routing.
2
3use std::collections::{BTreeSet, HashMap, HashSet};
4use std::fmt::{Display, Formatter, Result as FmtResult};
5use std::sync::Arc;
6
7use cedar_policy::pst::{
8    ActionConstraint, Clause, EntityOrSlot, Expr, Literal, PrincipalConstraint, ResourceConstraint,
9};
10use cedar_policy::{EntityTypeName, Policy};
11
12use crate::error::PolicyError;
13use crate::types::{Action, Resource};
14
15/// Reserved policy annotation used for explicit local or global assignment.
16pub const POLICY_STORE_ANNOTATION: &str = "treetop_store";
17
18/// A stable application-defined policy-store identifier.
19#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
20pub struct PolicyStoreId(Arc<str>);
21
22impl PolicyStoreId {
23    /// Construct a non-empty policy-store identifier.
24    pub fn new(value: impl AsRef<str>) -> Result<Self, PolicyError> {
25        let value = value.as_ref();
26        if value.trim().is_empty() {
27            return Err(PolicyError::PolicyStoreConfigError(
28                "policy-store ID must not be empty".to_string(),
29            ));
30        }
31        if value == "*" {
32            return Err(PolicyError::PolicyStoreConfigError(
33                "policy-store ID '*' is reserved for global policies".to_string(),
34            ));
35        }
36        Ok(Self(value.into()))
37    }
38
39    /// Borrow the configured identifier.
40    pub fn as_str(&self) -> &str {
41        &self.0
42    }
43}
44
45impl Display for PolicyStoreId {
46    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
47        f.write_str(self.as_str())
48    }
49}
50
51impl AsRef<str> for PolicyStoreId {
52    fn as_ref(&self) -> &str {
53        self.as_str()
54    }
55}
56
57/// One declared policy store and the Cedar namespace subtree it owns.
58#[derive(Debug, Clone, PartialEq, Eq)]
59pub struct PolicyStoreConfig {
60    id: PolicyStoreId,
61    namespace: Arc<str>,
62}
63
64impl PolicyStoreConfig {
65    /// Declare a store whose requests and ordinary policies use `namespace`.
66    ///
67    /// The namespace is a Cedar-qualified name such as `ExampleCo::DNS`. The
68    /// store also owns nested namespaces such as `ExampleCo::DNS::Admin`.
69    pub fn new(id: impl AsRef<str>, namespace: impl AsRef<str>) -> Result<Self, PolicyError> {
70        let id = PolicyStoreId::new(id)?;
71        let namespace = namespace.as_ref();
72        let _: EntityTypeName = namespace.parse().map_err(|error| {
73            PolicyError::PolicyStoreConfigError(format!(
74                "policy store '{id}' has invalid Cedar namespace '{namespace}': {error}"
75            ))
76        })?;
77        Ok(Self {
78            id,
79            namespace: namespace.into(),
80        })
81    }
82
83    /// Return the stable store identifier.
84    pub fn id(&self) -> &PolicyStoreId {
85        &self.id
86    }
87
88    /// Return the namespace subtree owned by this store.
89    pub fn namespace(&self) -> &str {
90        &self.namespace
91    }
92}
93
94#[derive(Debug, Default)]
95struct NamespaceRouter {
96    root: NamespaceNode,
97}
98
99#[derive(Debug, Default)]
100struct NamespaceNode {
101    store_index: Option<usize>,
102    children: HashMap<String, NamespaceNode>,
103}
104
105impl NamespaceRouter {
106    fn insert(&mut self, namespace: &str, store_index: usize) -> Option<usize> {
107        let mut node = &mut self.root;
108        for component in namespace.split("::") {
109            if let Some(existing_store) = node.store_index {
110                return Some(existing_store);
111            }
112            node = node.children.entry(component.to_string()).or_default();
113        }
114        if let Some(existing_store) = node.store_index.or_else(|| node.first_store_index()) {
115            return Some(existing_store);
116        }
117        node.store_index = Some(store_index);
118        None
119    }
120
121    fn resolve<'a>(&self, components: impl IntoIterator<Item = &'a str>) -> Option<usize> {
122        let mut node = &self.root;
123        for component in components {
124            node = node.children.get(component)?;
125            if let Some(store_index) = node.store_index {
126                return Some(store_index);
127            }
128        }
129        None
130    }
131
132    fn resolve_qualified_name(&self, name: &str) -> Option<usize> {
133        self.resolve(name.split("::"))
134    }
135}
136
137impl NamespaceNode {
138    fn first_store_index(&self) -> Option<usize> {
139        self.store_index.or_else(|| {
140            self.children
141                .values()
142                .find_map(NamespaceNode::first_store_index)
143        })
144    }
145}
146
147/// Validated configuration for namespace-partitioned policy stores.
148///
149/// Policies are assigned from the configured namespace references in their
150/// scope constraints and conditions. Policies whose `@id` value is registered
151/// as global, or which carry
152/// `@treetop_store("*")`, are installed in every store. An otherwise unscoped
153/// policy can use `@treetop_store("store-id")` for an explicit assignment.
154/// Entity namespaces outside every configured store may be shared across
155/// stores; they do not participate in assignment or request routing.
156#[derive(Debug, Clone)]
157pub struct PolicyStoreLayout {
158    stores: Arc<[PolicyStoreConfig]>,
159    global_policy_ids: Arc<HashSet<String>>,
160    namespace_router: Arc<NamespaceRouter>,
161}
162
163impl PolicyStoreLayout {
164    /// Create a layout containing non-overlapping store namespaces.
165    pub fn new(stores: impl IntoIterator<Item = PolicyStoreConfig>) -> Result<Self, PolicyError> {
166        let stores = stores.into_iter().collect::<Vec<_>>();
167        if stores.is_empty() {
168            return Err(PolicyError::PolicyStoreConfigError(
169                "a policy-store layout must declare at least one store".to_string(),
170            ));
171        }
172
173        let mut ids = HashSet::with_capacity(stores.len());
174        for store in &stores {
175            if !ids.insert(store.id().clone()) {
176                return Err(PolicyError::PolicyStoreConfigError(format!(
177                    "duplicate policy-store ID '{}'",
178                    store.id()
179                )));
180            }
181        }
182        let mut namespace_router = NamespaceRouter::default();
183        for (index, store) in stores.iter().enumerate() {
184            if let Some(existing_index) = namespace_router.insert(store.namespace(), index) {
185                return Err(PolicyError::PolicyStoreConfigError(format!(
186                    "policy-store namespaces '{}' and '{}' overlap",
187                    stores[existing_index].namespace(),
188                    store.namespace()
189                )));
190            }
191        }
192
193        Ok(Self {
194            stores: stores.into(),
195            global_policy_ids: Arc::new(HashSet::new()),
196            namespace_router: Arc::new(namespace_router),
197        })
198    }
199
200    /// Register policy `@id` values that must be installed in every store.
201    ///
202    /// Every registered ID must exist in the loaded policy source. This catches
203    /// stale or misspelled global-policy assignments during snapshot creation.
204    pub fn with_global_policy_ids<I, S>(self, policy_ids: I) -> Result<Self, PolicyError>
205    where
206        I: IntoIterator<Item = S>,
207        S: AsRef<str>,
208    {
209        let mut global_policy_ids = self.global_policy_ids.as_ref().clone();
210        for policy_id in policy_ids {
211            let policy_id = policy_id.as_ref();
212            if policy_id.trim().is_empty() {
213                return Err(PolicyError::PolicyStoreConfigError(
214                    "global policy ID must not be empty".to_string(),
215                ));
216            }
217            global_policy_ids.insert(policy_id.to_string());
218        }
219        Ok(Self {
220            stores: self.stores,
221            global_policy_ids: Arc::new(global_policy_ids),
222            namespace_router: self.namespace_router,
223        })
224    }
225
226    /// Return the configured stores in declaration order.
227    pub fn stores(&self) -> &[PolicyStoreConfig] {
228        &self.stores
229    }
230
231    pub(crate) fn global_policy_ids(&self) -> &HashSet<String> {
232        &self.global_policy_ids
233    }
234
235    pub(crate) fn explicit_policy_store(
236        &self,
237        policy: &Policy,
238    ) -> Result<Option<ExplicitPolicyStore>, PolicyError> {
239        let Some(value) = policy.annotation(POLICY_STORE_ANNOTATION) else {
240            return Ok(None);
241        };
242        if value == "*" {
243            return Ok(Some(ExplicitPolicyStore::Global));
244        }
245        let store_index = self.store_index_by_id(value).ok_or_else(|| {
246            PolicyError::PolicyStoreConfigError(format!(
247                "policy '{}' names unknown policy store '{value}' in @{POLICY_STORE_ANNOTATION}",
248                display_policy_id(policy)
249            ))
250        })?;
251        Ok(Some(ExplicitPolicyStore::Store(store_index)))
252    }
253
254    pub(crate) fn policy_candidates(&self, policy: &Policy) -> Result<Vec<usize>, PolicyError> {
255        let pst = policy.to_pst().map_err(|error| {
256            PolicyError::PolicyStoreConfigError(format!(
257                "policy '{}' cannot be inspected for policy-store assignment: {error}",
258                display_policy_id(policy)
259            ))
260        })?;
261        let body = pst.body();
262        let mut references = Vec::new();
263
264        match &body.principal {
265            PrincipalConstraint::Any => {}
266            PrincipalConstraint::Eq(value) | PrincipalConstraint::In(value) => {
267                collect_entity_or_slot(value, &mut references);
268            }
269            PrincipalConstraint::Is(entity_type) => {
270                references.push(entity_type.to_string());
271            }
272            PrincipalConstraint::IsIn(entity_type, value) => {
273                references.push(entity_type.to_string());
274                collect_entity_or_slot(value, &mut references);
275            }
276        }
277        let mut candidates = BTreeSet::new();
278        match &body.action {
279            ActionConstraint::Any => {}
280            ActionConstraint::Eq(uid) => {
281                references.push(uid.ty.to_string());
282            }
283            ActionConstraint::In(uids) => {
284                for uid in uids {
285                    references.push(uid.ty.to_string());
286                }
287            }
288        }
289        match &body.resource {
290            ResourceConstraint::Any => {}
291            ResourceConstraint::Eq(value) | ResourceConstraint::In(value) => {
292                collect_entity_or_slot(value, &mut references);
293            }
294            ResourceConstraint::Is(entity_type) => {
295                references.push(entity_type.to_string());
296            }
297            ResourceConstraint::IsIn(entity_type, value) => {
298                references.push(entity_type.to_string());
299                collect_entity_or_slot(value, &mut references);
300            }
301        }
302        for clause in body.clauses() {
303            let expression = match clause {
304                Clause::When(expression) | Clause::Unless(expression) => expression,
305            };
306            collect_expr_references(expression, &mut references);
307        }
308        for reference in references {
309            self.insert_name_candidate(&reference, &mut candidates);
310        }
311
312        if candidates.len() > 1 {
313            let names = candidates
314                .iter()
315                .map(|index| self.stores[*index].id().as_str())
316                .collect::<Vec<_>>()
317                .join(", ");
318            return Err(PolicyError::PolicyStoreConfigError(format!(
319                "policy '{}' spans multiple policy stores ({names}); split it or mark it global",
320                display_policy_id(policy)
321            )));
322        }
323        Ok(candidates.into_iter().collect())
324    }
325
326    pub(crate) fn resolve_request(
327        &self,
328        action: &Action,
329        resource: &Resource,
330    ) -> Result<usize, PolicyError> {
331        let action_store = self
332            .namespace_router
333            .resolve(action.namespace().iter().map(String::as_str));
334        let resource_store = self
335            .namespace_router
336            .resolve_qualified_name(resource.kind());
337
338        match (action_store, resource_store) {
339            (Some(action_store), Some(resource_store)) if action_store != resource_store => {
340                Err(PolicyError::PolicyStoreRoutingError(format!(
341                    "request action namespace '{}' and resource type '{}' resolve to different policy stores ({}, {})",
342                    display_action_namespace(action),
343                    resource.kind(),
344                    self.stores[action_store].id(),
345                    self.stores[resource_store].id()
346                )))
347            }
348            (Some(index), _) | (_, Some(index)) => Ok(index),
349            (None, None) => Err(PolicyError::PolicyStoreRoutingError(format!(
350                "request action namespace '{}' and resource type '{}' do not belong to a configured policy store",
351                display_action_namespace(action),
352                resource.kind()
353            ))),
354        }
355    }
356
357    fn store_index_by_id(&self, id: &str) -> Option<usize> {
358        self.stores
359            .iter()
360            .position(|store| store.id().as_str() == id)
361    }
362
363    fn insert_name_candidate(&self, name: &str, candidates: &mut BTreeSet<usize>) {
364        if let Some(index) = self.namespace_router.resolve_qualified_name(name) {
365            candidates.insert(index);
366        }
367    }
368}
369
370fn display_action_namespace(action: &Action) -> String {
371    if action.namespace().is_empty() {
372        "<none>".to_string()
373    } else {
374        action.namespace().join("::")
375    }
376}
377
378fn collect_entity_or_slot(value: &EntityOrSlot, references: &mut Vec<String>) {
379    if let EntityOrSlot::Entity(uid) = value {
380        references.push(uid.ty.to_string());
381    }
382}
383
384fn collect_expr_references(expression: &Expr, references: &mut Vec<String>) {
385    match expression {
386        Expr::Literal(Literal::EntityUID(uid)) => references.push(uid.ty.to_string()),
387        Expr::UnaryOp { expr, .. }
388        | Expr::GetAttr { expr, .. }
389        | Expr::HasAttr { expr, .. }
390        | Expr::Like { expr, .. } => collect_expr_references(expr, references),
391        Expr::BinaryOp { left, right, .. } => {
392            collect_expr_references(left, references);
393            collect_expr_references(right, references);
394        }
395        Expr::Is {
396            expr,
397            entity_type,
398            in_expr,
399        } => {
400            references.push(entity_type.to_string());
401            collect_expr_references(expr, references);
402            if let Some(in_expr) = in_expr {
403                collect_expr_references(in_expr, references);
404            }
405        }
406        Expr::IfThenElse {
407            cond,
408            then_expr,
409            else_expr,
410        } => {
411            collect_expr_references(cond, references);
412            collect_expr_references(then_expr, references);
413            collect_expr_references(else_expr, references);
414        }
415        Expr::Set(expressions) => {
416            for expression in expressions {
417                collect_expr_references(expression, references);
418            }
419        }
420        Expr::Record(expressions) => {
421            for expression in expressions.values() {
422                collect_expr_references(expression, references);
423            }
424        }
425        _ => {}
426    }
427}
428
429pub(crate) enum ExplicitPolicyStore {
430    Global,
431    Store(usize),
432}
433
434pub(crate) fn display_policy_id(policy: &Policy) -> &str {
435    policy
436        .annotation("id")
437        .filter(|id| !id.is_empty())
438        .unwrap_or_else(|| policy.id().as_ref())
439}
440
441#[cfg(test)]
442mod tests {
443    use super::*;
444    use crate::types::{Action, Resource};
445
446    fn layout() -> PolicyStoreLayout {
447        PolicyStoreLayout::new([
448            PolicyStoreConfig::new("dns", "ExampleCo::DNS").unwrap(),
449            PolicyStoreConfig::new("www", "ExampleCo::WWW").unwrap(),
450        ])
451        .unwrap()
452    }
453
454    #[test]
455    fn rejects_overlapping_namespaces() {
456        let result = PolicyStoreLayout::new([
457            PolicyStoreConfig::new("parent", "ExampleCo").unwrap(),
458            PolicyStoreConfig::new("child", "ExampleCo::DNS").unwrap(),
459        ]);
460        assert!(matches!(
461            result,
462            Err(PolicyError::PolicyStoreConfigError(_))
463        ));
464
465        let reverse = PolicyStoreLayout::new([
466            PolicyStoreConfig::new("child", "ExampleCo::DNS").unwrap(),
467            PolicyStoreConfig::new("parent", "ExampleCo").unwrap(),
468        ]);
469        assert!(matches!(
470            reverse,
471            Err(PolicyError::PolicyStoreConfigError(_))
472        ));
473    }
474
475    #[test]
476    fn rejects_invalid_layout_entries() {
477        assert!(matches!(
478            PolicyStoreConfig::new("dns", "not a namespace"),
479            Err(PolicyError::PolicyStoreConfigError(_))
480        ));
481        assert!(matches!(
482            PolicyStoreLayout::new([
483                PolicyStoreConfig::new("duplicate", "ExampleCo::DNS").unwrap(),
484                PolicyStoreConfig::new("duplicate", "ExampleCo::WWW").unwrap(),
485            ]),
486            Err(PolicyError::PolicyStoreConfigError(_))
487        ));
488    }
489
490    #[test]
491    fn routes_namespaced_requests() {
492        let layout = layout();
493        let store = layout
494            .resolve_request(
495                &Action::new("read", Some(vec!["ExampleCo".into(), "DNS".into()])).unwrap(),
496                &Resource::new("ExampleCo::DNS::Host", "host-1").unwrap(),
497            )
498            .unwrap();
499        assert_eq!(layout.stores()[store].id().as_str(), "dns");
500    }
501
502    #[test]
503    fn rejects_cross_store_requests() {
504        let result = layout().resolve_request(
505            &Action::new("read", Some(vec!["ExampleCo".into(), "DNS".into()])).unwrap(),
506            &Resource::new("ExampleCo::WWW::Page", "page-1").unwrap(),
507        );
508        assert!(matches!(
509            result,
510            Err(PolicyError::PolicyStoreRoutingError(_))
511        ));
512    }
513
514    #[test]
515    fn namespace_matching_is_segment_aware() {
516        let result = layout().resolve_request(
517            &Action::new("read", Some(vec!["ExampleCo".into(), "DNSAdmin".into()])).unwrap(),
518            &Resource::new("ExampleCo::DNSAdmin::Host", "host-1").unwrap(),
519        );
520        assert!(matches!(
521            result,
522            Err(PolicyError::PolicyStoreRoutingError(_))
523        ));
524    }
525}