1use 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
15pub const POLICY_STORE_ANNOTATION: &str = "treetop_store";
17
18#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
20pub struct PolicyStoreId(Arc<str>);
21
22impl PolicyStoreId {
23 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 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#[derive(Debug, Clone, PartialEq, Eq)]
59pub struct PolicyStoreConfig {
60 id: PolicyStoreId,
61 namespace: Arc<str>,
62}
63
64impl PolicyStoreConfig {
65 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 pub fn id(&self) -> &PolicyStoreId {
85 &self.id
86 }
87
88 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#[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 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 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 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}