refacto: apply review suggestions

Cosmian · Dec 22, 2023 · 31e9f5f · 31e9f5f
1 parent 69a7c37
commit 31e9f5f
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 31 deletions.
diff --git a/benches/benches.rs b/benches/benches.rs
@@ -65,10 +65,10 @@ fn policy() -> Result<Policy, Error> {
 fn bench_policy_editing(c: &mut Criterion) {
     let cover_crypt = Covercrypt::default();
     let new_dep_attr = Attribute::new("Department", "Tech");
-    let new_dep_name = "IT";
+    let new_dep_name = "IT".to_string();
     let remove_dep_attr = Attribute::new("Department", "FIN");
     let old_sl_attr = Attribute::new("Security Level", "Protected");
-    let new_sl_name = "Open";
+    let new_sl_name = "Open".to_string();
     let disable_sl_attr = Attribute::new("Security Level", "Confidential");
 
     let mut group = c.benchmark_group("Edit Policy");
@@ -88,11 +88,13 @@ fn bench_policy_editing(c: &mut Criterion) {
                     .add_attribute(new_dep_attr.clone(), EncryptionHint::Classic)
                     .unwrap();
                 policy
-                    .rename_attribute(&new_dep_attr, new_dep_name)
+                    .rename_attribute(&new_dep_attr, new_dep_name.clone())
                     .unwrap();
                 policy.remove_attribute(&remove_dep_attr).unwrap();
 
-                policy.rename_attribute(&old_sl_attr, new_sl_name).unwrap();
+                policy
+                    .rename_attribute(&old_sl_attr, new_sl_name.clone())
+                    .unwrap();
                 policy.disable_attribute(&disable_sl_attr).unwrap();
 
                 cover_crypt

diff --git a/src/abe_policy/dimension.rs b/src/abe_policy/dimension.rs
@@ -1,4 +1,7 @@
-use std::{collections::HashMap, fmt::Debug};
+use std::{
+    collections::{hash_map::Entry, HashMap},
+    fmt::Debug,
+};
 
 use serde::{Deserialize, Serialize};
 
@@ -121,10 +124,10 @@ impl Dimension {
     /// * `dim` - The `DimensionBuilder` to base the dimension on.
     /// * `seed_id` - A mutable reference to a seed ID used for generating
     ///   unique values for attributes.
-    pub fn new(dim: &DimensionBuilder, seed_id: &mut u32) -> Self {
-        let attributes_mapping = dim.attributes_properties.iter().map(|attr| {
+    pub fn new(dim: DimensionBuilder, seed_id: &mut u32) -> Self {
+        let attributes_mapping = dim.attributes_properties.into_iter().map(|attr| {
             (
-                attr.name.clone(),
+                attr.name,
                 AttributeParameters::new(attr.encryption_hint, seed_id),
             )
         });
@@ -182,22 +185,19 @@ impl Dimension {
     /// Returns an error if the operation is not permitted.
     pub fn add_attribute(
         &mut self,
-        attr_name: &AttributeName,
+        attr_name: AttributeName,
         encryption_hint: EncryptionHint,
         seed_id: &mut u32,
     ) -> Result<(), Error> {
         match self {
             Self::Unordered(attributes) => {
-                if attributes.contains_key(attr_name) {
+                if let Entry::Vacant(entry) = attributes.entry(attr_name) {
+                    entry.insert(AttributeParameters::new(encryption_hint, seed_id));
+                    Ok(())
+                } else {
                     Err(Error::OperationNotPermitted(
                         "Attribute already in dimension".to_string(),
                     ))
-                } else {
-                    attributes.insert(
-                        attr_name.clone(),
-                        AttributeParameters::new(encryption_hint, seed_id),
-                    );
-                    Ok(())
                 }
             }
             Self::Ordered(_) => Err(Error::OperationNotPermitted(
@@ -264,25 +264,25 @@ impl Dimension {
     pub fn rename_attribute(
         &mut self,
         attr_name: &AttributeName,
-        new_name: &str,
+        new_name: String,
     ) -> Result<(), Error> {
         match self {
             Self::Unordered(attributes) => {
-                if attributes.contains_key(new_name) {
+                if attributes.contains_key(&new_name) {
                     return Err(Error::OperationNotPermitted(
                         "New attribute name is already used in the same dimension".to_string(),
                     ));
                 }
                 match attributes.remove(attr_name) {
                     Some(attr_params) => {
-                        attributes.insert(new_name.to_string(), attr_params);
+                        attributes.insert(new_name, attr_params);
                         Ok(())
                     }
                     None => Err(Error::AttributeNotFound(attr_name.to_string())),
                 }
             }
             Self::Ordered(attributes) => attributes
-                .update_key(attr_name, new_name.to_string())
+                .update_key(attr_name, new_name)
                 .map_err(|e| Error::OperationNotPermitted(e.to_string())),
         }
     }

diff --git a/src/abe_policy/policy.rs b/src/abe_policy/policy.rs
@@ -63,7 +63,7 @@ impl Policy {
 
         self.dimensions.insert(
             dim.name.clone(),
-            Dimension::new(&dim, &mut self.last_attribute_value),
+            Dimension::new(dim, &mut self.last_attribute_value),
         );
 
         Ok(())
@@ -95,11 +95,9 @@ impl Policy {
         encryption_hint: EncryptionHint,
     ) -> Result<(), Error> {
         match self.dimensions.get_mut(&attr.dimension) {
-            Some(policy_dim) => policy_dim.add_attribute(
-                &attr.name,
-                encryption_hint,
-                &mut self.last_attribute_value,
-            ),
+            Some(policy_dim) => {
+                policy_dim.add_attribute(attr.name, encryption_hint, &mut self.last_attribute_value)
+            }
             None => Err(Error::DimensionNotFound(attr.dimension)),
         }
     }
@@ -130,7 +128,7 @@ impl Policy {
     }
 
     /// Changes the name of an attribute.
-    pub fn rename_attribute(&mut self, attr: &Attribute, new_name: &str) -> Result<(), Error> {
+    pub fn rename_attribute(&mut self, attr: &Attribute, new_name: String) -> Result<(), Error> {
         match self.dimensions.get_mut(&attr.dimension) {
             Some(policy_dim) => policy_dim.rename_attribute(&attr.name, new_name),
             None => Err(Error::DimensionNotFound(attr.dimension.to_string())),

diff --git a/src/abe_policy/tests.rs b/src/abe_policy/tests.rs
@@ -73,17 +73,20 @@ fn test_edit_policy_attributes() -> Result<(), Error> {
 
     // Try renaming Research to already used name MKG
     assert!(policy
-        .rename_attribute(&Attribute::new("Department", "R&D"), "MKG",)
+        .rename_attribute(&Attribute::new("Department", "R&D"), "MKG".to_string(),)
         .is_err());
 
     // Rename R&D to Research
     assert!(policy
-        .rename_attribute(&Attribute::new("Department", "R&D"), "Research",)
+        .rename_attribute(&Attribute::new("Department", "R&D"), "Research".to_string(),)
         .is_ok());
 
     // Rename ordered dimension
     assert!(policy
-        .rename_attribute(&Attribute::new("Security Level", "Protected"), "Open",)
+        .rename_attribute(
+            &Attribute::new("Security Level", "Protected"),
+            "Open".to_string(),
+        )
         .is_ok());
     let order: Vec<_> = policy
         .dimensions

diff --git a/src/test_utils/mod.rs b/src/test_utils/mod.rs
@@ -398,7 +398,7 @@ mod tests {
             EncryptedHeader::generate(&cover_crypt, &policy, &mpk, &top_secret_ap, None, None)?;
 
         // remove the FIN department
-        policy.rename_attribute(&Attribute::new("Department", "FIN"), "Finance")?;
+        policy.rename_attribute(&Attribute::new("Department", "FIN"), "Finance".to_string())?;
 
         // update the master keys
         cover_crypt.update_master_keys(&policy, &mut msk, &mut mpk)?;