fix entity_facade test for aggregate id

tutao · Nov 14, 2024 · c54b003 · c54b003
1 parent 97eb93a
commit c54b003
Show file tree

Hide file tree

Showing 2 changed files with 326 additions and 280 deletions.
diff --git a/tuta-sdk/rust/sdk/src/entities/entity_facade.rs b/tuta-sdk/rust/sdk/src/entities/entity_facade.rs
@@ -192,11 +192,9 @@ impl EntityFacadeImpl {
             let id_key = String::from("_id");
             match encrypted.get(&id_key) {
                 Some(ElementValue::Null) => {
-                    let new_id = self.randomizer_facade.generate_random_array::<4>();
-
                     encrypted.insert(
                         id_key,
-                        ElementValue::String(BASE64_URL_SAFE_NO_PAD.encode(new_id)),
+                        make_random_aggregate_id(&self.randomizer_facade),
                     );
                 }
                 Some(_) => {
@@ -626,6 +624,13 @@ impl EntityFacade for EntityFacadeImpl {
     }
 }
 
+fn make_random_aggregate_id(random: &RandomizerFacade) -> ElementValue {
+    let new_id_bytes = random.generate_random_array::<4>();
+    let new_id_string = BASE64_URL_SAFE_NO_PAD.encode(new_id_bytes);
+    let new_id = crate::CustomId(new_id_string);
+    ElementValue::IdCustomId(new_id)
+}
+
 #[cfg(test)]
 mod lz4_compressed_string_compatibility_tests {
     use crate::crypto::compatibility_test_utils::{
@@ -710,9 +715,7 @@ mod tests {
     use crate::crypto::{aes::Iv, Aes256Key};
     use crate::date::DateTime;
     use crate::element_value::{ElementValue, ParsedEntity};
-    use crate::entities::entity_facade::{
-        EntityFacade, EntityFacadeImpl, MappedValue, MAX_UNCOMPRESSED_INPUT_LZ4,
-    };
+    use crate::entities::entity_facade::{make_random_aggregate_id, EntityFacade, EntityFacadeImpl, MappedValue, MAX_UNCOMPRESSED_INPUT_LZ4};
     use crate::entities::generated::sys::CustomerAccountTerminationRequest;
     use crate::entities::generated::tutanota::Mail;
     use crate::entities::Entity;
@@ -1075,7 +1078,12 @@ mod tests {
         let owner_enc_session_key = [0, 1, 2];
 
         let deterministic_rng = DeterministicRng(20);
-        let iv = Iv::generate(&RandomizerFacade::from_core(deterministic_rng.clone()));
+        let random = RandomizerFacade::from_core(deterministic_rng.clone());
+        // we're kind of using the implementation to test itself here, but
+        // this is not about the actual value, but that it is set at all
+        let expected_aggregate_id = make_random_aggregate_id(&random);
+
+        let iv = Iv::generate(&random);
         let type_model_provider = Arc::new(init_type_model_provider());
 
         let type_ref = Mail::type_ref();
@@ -1088,7 +1096,7 @@ mod tests {
             RandomizerFacade::from_core(deterministic_rng),
         );
 
-        let (mut expected_encrypted_mail, raw_mail) = generate_email_entity(
+        let (mut expected_encrypted_mail, mut raw_mail) = generate_email_entity(
             &sk,
             &iv,
             true,
@@ -1097,7 +1105,7 @@ mod tests {
             String::from("Munich"),
         );
 
-        // remove finalIvs for easy comparision
+        // removes finalIvs for easy comparison as well as setting the aggregate _id fields
         {
             expected_encrypted_mail.remove("_finalIvs").unwrap();
             expected_encrypted_mail
@@ -1106,22 +1114,59 @@ mod tests {
                 .assert_dict_mut_ref()
                 .remove("_finalIvs")
                 .unwrap();
+            expected_encrypted_mail
+                .get_mut("sender")
+                .unwrap()
+                .assert_dict_mut_ref()
+                .insert("_id".to_string(), expected_aggregate_id.clone());
             expected_encrypted_mail
                 .get_mut("firstRecipient")
                 .unwrap()
                 .assert_dict_mut_ref()
                 .remove("_finalIvs")
                 .unwrap();
+            expected_encrypted_mail
+                .get_mut("firstRecipient")
+                .unwrap()
+                .assert_dict_mut_ref()
+                .insert("_id".to_string(), expected_aggregate_id.clone());
         }
 
-        let encrypted_mail = entity_facade.encrypt_and_map_inner(type_model, &raw_mail, &sk);
+        {
+            // generate_email_entity generates aggregate ids for us, but the entity_facade is supposed to set
+            // them on the fly if they're missing. by removing them here we test that they're re-created
+            raw_mail
+                .get_mut("sender")
+                .unwrap()
+                .assert_dict_mut_ref()
+                .insert(String::from("_id"), ElementValue::Null);
+            raw_mail
+                .get_mut("firstRecipient")
+                .unwrap()
+                .assert_dict_mut_ref()
+                .insert(String::from("_id"), ElementValue::Null);
+        }
+
+        let encrypted_mail = entity_facade.encrypt_and_map_inner(type_model, &raw_mail, &sk).unwrap();
 
-        assert_eq!(Ok(expected_encrypted_mail), encrypted_mail);
+        assert_eq!(expected_encrypted_mail, encrypted_mail);
 
         // verify every data is preserved as is after decryption
         {
-            let original_mail = raw_mail;
-            let encrypted_mail = encrypted_mail.unwrap();
+            let mut original_mail = raw_mail;
+            {
+                original_mail
+                    .get_mut("sender")
+                    .unwrap()
+                    .assert_dict_mut_ref()
+                    .insert(String::from("_id"), expected_aggregate_id.clone());
+                original_mail
+                    .get_mut("firstRecipient")
+                    .unwrap()
+                    .assert_dict_mut_ref()
+                    .insert(String::from("_id"), expected_aggregate_id.clone());
+            }
+            let encrypted_mail = encrypted_mail;
 
             let mut decrypted_mail = entity_facade
                 .decrypt_and_map(