Adds ability to Merge policy sets with the option to rename policy id's upon duplicate ids.

cedar-policy-core/src/ast/policy_set.rs

@@ -264,6 +264,103 @@ impl PolicySet {
         Ok(())
     }
 
+    /// Helper function for `merge_policyset` to check if the `PolicyID` pid
+    /// appears in this `PolicySet`'s links or templates.
+    fn policy_id_is_bound(&self, pid: &PolicyID) -> bool {
+        self.templates.contains_key(pid) || self.links.contains_key(pid)
+    }
+
+    /// Helper function for `merge_policyset` to construct a renaming
+    /// that would resolve any conflicting `PolicyID`s. We use the type parameter `T`
+    /// to allow this code to be applied to both Templates and Policies.
+    fn update_renaming<T>(
+        &self,
+        this_contents: &HashMap<PolicyID, T>,
+        other: &Self,
+        other_contents: &HashMap<PolicyID, T>,
+        renaming: &mut HashMap<PolicyID, PolicyID>,
+        start_ind: &mut u32,
+    ) where
+        T: PartialEq + Clone,
+    {
+        for (pid, ot) in other_contents {
+            if let Some(tt) = this_contents.get(pid) {
+                if tt != ot {
+                    let mut new_pid = PolicyID::from_string(format!("policy{}", start_ind));
+                    *start_ind += 1;
+                    while self.policy_id_is_bound(&new_pid) || other.policy_id_is_bound(&new_pid) {
+                        new_pid = PolicyID::from_string(format!("policy{}", start_ind));
+                        *start_ind += 1;
+                    }
+                    renaming.insert(pid.clone(), new_pid);
+                }
+            }
+        }
+    }
+
+    /// Merges this `PolicySet` with another `PolicySet`.
+    /// This `PolicySet` is modified while the other `PolicySet`
+    /// remains unchanged.
+    ///
+    /// The flag `rename_duplicates` controls the expected behavior
+    /// when a `PolicyID` in this and the other `PolicySet` conflict.
+    ///
+    /// When `rename_duplicates` is false, conflicting `PolicyID`s result
+    /// in a occupied `PolicySetError`.
+    ///
+    /// Otherwise, when `rename_duplicates` is true, conflicting `PolicyID`s from
+    /// the other `PolicySet` are automatically renamed to avoid conflict.
+    /// This renaming is returned as a Hashmap from the old `PolicyID` to the
+    /// renamed `PolicyID`.
+    pub fn merge_policyset(
+        &mut self,
+        other: &PolicySet,
+        rename_duplicates: bool,
+    ) -> Result<HashMap<PolicyID, PolicyID>, PolicySetError> {
+        // Check for conflicting policy ids. If there is a conflict either
+        // throw an error or construct a renaming (if `rename_duplicates` is true)
+        let mut min_id = 0;
+        let mut renaming = HashMap::new();
+        self.update_renaming(
+            &self.templates,
+            other,
+            &other.templates,
+            &mut renaming,
+            &mut min_id,
+        );
+        self.update_renaming(&self.links, other, &other.links, &mut renaming, &mut min_id);
+        // If `rename_dupilicates` is false, then throw an error if any renaming should happen
+        if !rename_duplicates {
+            if let Some(pid) = renaming.keys().next() {
+                return Err(PolicySetError::Occupied { id: pid.clone() });
+            }
+        }
+        // either there are no conflicting policy ids
+        // or we should rename conflicting policy ids (using renaming) to avoid conflicting policy ids
+        for (pid, other_template) in &other.templates {
+            let pid = renaming.get(pid).unwrap_or(pid);
+            self.templates.insert(pid.clone(), other_template.clone());
+        }
+        for (pid, other_policy) in &other.links {
+            let pid = renaming.get(pid).unwrap_or(pid);
+            self.links.insert(pid.clone(), other_policy.clone());
+        }
+        for (tid, other_template_link_set) in &other.template_to_links_map {
+            let tid = renaming.get(tid).unwrap_or(tid);
+            let mut this_template_link_set = self
+                .template_to_links_map
+                .remove(tid)
+                .unwrap_or(HashSet::new());
+            for pid in other_template_link_set {
+                let pid = renaming.get(pid).unwrap_or(pid);
+                this_template_link_set.insert(pid.clone());
+            }
+            self.template_to_links_map
+                .insert(tid.clone(), this_template_link_set);
+        }
+        Ok(renaming)
+    }
+
     /// Remove a static `Policy`` from the `PolicySet`.
     pub fn remove_static(
         &mut self,
@@ -670,6 +767,110 @@ mod test {
         }
     }
 
+    #[test]
+    fn policy_merge_no_conflicts() {
+        let p1 = parser::parse_policy(
+            Some(PolicyID::from_string("policy0")),
+            "permit(principal,action,resource);",
+        )
+        .expect("Failed to parse");
+        let p2 = parser::parse_policy(
+            Some(PolicyID::from_string("policy1")),
+            "permit(principal,action,resource) when { false };",
+        )
+        .expect("Failed to parse");
+        let p3 = parser::parse_policy(
+            Some(PolicyID::from_string("policy0")),
+            "permit(principal,action,resource);",
+        )
+        .expect("Failed to parse");
+        let p4 = parser::parse_policy(
+            Some(PolicyID::from_string("policy2")),
+            "permit(principal,action,resource) when { true };",
+        )
+        .expect("Failed to parse");
+        let mut pset1 = PolicySet::new();
+        let mut pset2 = PolicySet::new();
+        pset1.add_static(p1).expect("Failed to add!");
+        pset1.add_static(p2).expect("Failed to add!");
+        pset2.add_static(p3).expect("Failed to add!");
+        pset2.add_static(p4).expect("Failed to add!");
+        // should not conflict because p1 == p3
+        match pset1.merge_policyset(&pset2, false) {
+            Ok(_) => (),
+            Err(PolicySetError::Occupied { id }) => panic!(
+                "There should not have been an error! Unexpected conflict for id {}",
+                id
+            ),
+        }
+    }
+
+    #[test]
+    fn policy_merge_with_conflicts() {
+        let pid0 = PolicyID::from_string("policy0");
+        let pid1 = PolicyID::from_string("policy1");
+        let pid2 = PolicyID::from_string("policy2");
+        let p1 = parser::parse_policy(Some(pid0.clone()), "permit(principal,action,resource);")
+            .expect("Failed to parse");
+        let p2 = parser::parse_policy(
+            Some(pid1.clone()),
+            "permit(principal,action,resource) when { false };",
+        )
+        .expect("Failed to parse");
+        let p3 = parser::parse_policy(Some(pid1.clone()), "permit(principal,action,resource);")
+            .expect("Failed to parse");
+        let p4 = parser::parse_policy(
+            Some(pid2.clone()),
+            "permit(principal,action,resource) when { true };",
+        )
+        .expect("Failed to parse");
+        let mut pset1 = PolicySet::new();
+        let mut pset2 = PolicySet::new();
+        pset1.add_static(p1.clone()).expect("Failed to add!");
+        pset1.add_static(p2.clone()).expect("Failed to add!");
+        pset2.add_static(p3.clone()).expect("Failed to add!");
+        pset2.add_static(p4.clone()).expect("Failed to add!");
+        // should conclict on pid "policy1"
+        match pset1.merge_policyset(&pset2, false) {
+            Ok(_) => panic!("`pset1` and `pset2` should conflict for PolicyID `policy1`"),
+            Err(PolicySetError::Occupied { id }) => {
+                assert_eq!(id, PolicyID::from_string("policy1"));
+            }
+        }
+        // should not conflict because of auto-renaming of conflicting policies
+        match pset1.merge_policyset(&pset2, true) {
+            Ok(renaming) => {
+                // ensure `policy1` was renamed
+                let new_pid1 = match renaming.get(&pid1) {
+                    Some(new_pid1) => new_pid1,
+                    None => panic!("Error: `policy1` is a conflict and should be renamed"),
+                };
+                // ensure no other policy was renamed
+                assert_eq!(renaming.keys().len(), 1);
+                match pset1.get(&pid0) {
+                    Some(new_p1) => assert_eq!(Policy::from(p1), new_p1.clone()),
+                    None => (),
+                }
+                match pset1.get(&pid1) {
+                    Some(new_p2) => assert_eq!(Policy::from(p2), new_p2.clone()),
+                    None => (),
+                }
+                match pset1.get(new_pid1) {
+                    Some(new_p3) => assert_eq!(Policy::from(p3), new_p3.clone()),
+                    None => (),
+                }
+                match pset1.get(&pid2) {
+                    Some(new_p4) => assert_eq!(Policy::from(p4), new_p4.clone()),
+                    None => (),
+                }
+            }
+            Err(PolicySetError::Occupied { id }) => panic!(
+                "There should not have been an error! Unexpected conflict for id {}",
+                id
+            ),
+        }
+    }
+
     #[test]
     fn policy_conflicts() {
         let mut pset = PolicySet::new();

cedar-policy/CHANGELOG.md

@@ -20,6 +20,7 @@ Cedar Language Version: TBD
 ### Added
 
 - Added `Entities::remove_entities()` to remove `Entity`s from an `Entities` struct (resolving #701)
+- Added `PolicySet::merge_policyset()` to merge a `PolicySet` into another `PolicySet` struct (resolving #610)
 - Implemented [RFC 53 (enumerated entity types)](https://github.com/cedar-policy/rfcs/blob/main/text/0053-enum-entities.md)  (#1377)
 
 ### Fixed

cedar-policy/src/api.rs

@@ -2237,6 +2237,73 @@ impl PolicySet {
         Ok(set)
     }
 
+    /// Helper function for `merge_policyset`
+    /// Merges two sets and avoids name clashes by using the provided
+    /// renaming. The type parameter `T` allows this code to be used for
+    /// both Templates and Policies.
+    fn merge_sets<T>(
+        this: &mut HashMap<PolicyId, T>,
+        other: &HashMap<PolicyId, T>,
+        renaming: &HashMap<PolicyId, PolicyId>,
+    ) where
+        T: PartialEq + Clone,
+    {
+        for (pid, ot) in other {
+            match renaming.get(&pid) {
+                Some(new_pid) => {
+                    this.insert(new_pid.clone(), ot.clone());
+                }
+                None => {
+                    if this.get(pid).is_none() {
+                        this.insert(pid.clone(), ot.clone());
+                    }
+                    // If pid is not in the renaming but is in both
+                    // this and other, then by assumption
+                    // the element at pid in this and other are equal
+                    // i.e., the renaming is expected to track all
+                    // conflicting pids.
+                }
+            }
+        }
+    }
+
+    /// Merges this `PolicySet` with another `PolicySet`.
+    /// This `PolicySet` is modified while the other `PolicySet`
+    /// remains unchanged.
+    ///
+    /// The flag `rename_duplicates` controls the expected behavior
+    /// when a `PolicyId` in this and the other `PolicySet` conflict.
+    ///
+    /// When `rename_duplicates` is false, conflicting `PolicyId`s result
+    /// in a `PolicySetError::AlreadyDefined` error.
+    ///
+    /// Otherwise, when `rename_duplicates` is true, conflicting `PolicyId`s from
+    /// the other `PolicySet` are automatically renamed to avoid conflict.
+    /// This renaming is returned as a Hashmap from the old `PolicyId` to the
+    /// renamed `PolicyId`.
+    pub fn merge_policyset(
+        &mut self,
+        other: &PolicySet,
+        rename_duplicates: bool,
+    ) -> Result<HashMap<PolicyId, PolicyId>, PolicySetError> {
+        match self.ast.merge_policyset(&other.ast, rename_duplicates) {
+            Ok(renaming) => {
+                let renaming: HashMap<PolicyId, PolicyId> = renaming
+                    .into_iter()
+                    .map(|(old_pid, new_pid)| (PolicyId::new(old_pid), PolicyId::new(new_pid)))
+                    .collect();
+                Self::merge_sets(&mut self.templates, &other.templates, &renaming);
+                Self::merge_sets(&mut self.policies, &other.policies, &renaming);
+                Ok(renaming)
+            }
+            Err(ast::PolicySetError::Occupied { id }) => Err(PolicySetError::AlreadyDefined(
+                policy_set_errors::AlreadyDefined {
+                    id: PolicyId::new(id),
+                },
+            )),
+        }
+    }
+
     /// Add an static policy to the `PolicySet`. To add a template instance, use
     /// `link` instead. This function will return an error (and not modify
     /// the `PolicySet`) if a template-linked policy is passed in.