refacto: move refresh logic to an iterator struct

benches/benches.rs

@@ -310,57 +310,58 @@ fn bench_header_encryption(c: &mut Criterion) {
 fn bench_header_decryption(c: &mut Criterion) {
     let policy = policy().expect("cannot generate policy");
     let authenticated_data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9];
-    let (user_access_policy, access_policies) = get_access_policies();
+    let (user_access_policies, access_policies) = get_access_policies();
     let cover_crypt = Covercrypt::default();
     let (msk, mpk) = cover_crypt
         .generate_master_keys(&policy)
         .expect("cannot generate master keys");
-    let user_decryption_keys: Vec<_> = user_access_policy
-        .iter()
-        .map(|ap| {
-            cover_crypt
-                .generate_user_secret_key(&msk, ap, &policy)
-                .expect("cannot generate user private key")
-        })
-        .collect();
     let mut group = c.benchmark_group("Header encryption and decryption");
-    for (n_partitions_usk, usk) in user_decryption_keys.iter().enumerate() {
+    for (n_user, user_access_policy) in user_access_policies.iter().enumerate() {
         for (n_partition_ct, access_policy) in access_policies.iter().enumerate() {
             group.bench_function(
                 &format!(
                     "ciphertexts with {} partition(s), usk with {} partitions",
                     n_partition_ct + 1,
-                    n_partitions_usk + 1
+                    n_user + 1
                 ),
                 |b| {
-                    b.iter(|| {
-                        let (_, encrypted_header) = EncryptedHeader::generate(
-                            &cover_crypt,
-                            &policy,
-                            &mpk,
-                            access_policy,
-                            None,
-                            Some(&authenticated_data),
-                        )
-                        .unwrap_or_else(|_| {
-                            panic!(
-                                "cannot encrypt header for {} ciphertext partition(s), {} usk \
-                                 partition(s)",
-                                n_partition_ct + 1,
-                                n_partitions_usk
+                    b.iter_batched(
+                        || {
+                            let usk = cover_crypt
+                                .generate_user_secret_key(&msk, user_access_policy, &policy)
+                                .expect("cannot generate user private key");
+                            let (_, encrypted_header) = EncryptedHeader::generate(
+                                &cover_crypt,
+                                &policy,
+                                &mpk,
+                                access_policy,
+                                None,
+                                Some(&authenticated_data),
                             )
-                        });
-                        encrypted_header
-                            .decrypt(&cover_crypt, usk, Some(&authenticated_data))
                             .unwrap_or_else(|_| {
                                 panic!(
-                                    "cannot decrypt header for {} ciphertext partition(s), {} usk \
+                                    "cannot encrypt header for {} ciphertext partition(s), {} usk \
                                      partition(s)",
                                     n_partition_ct + 1,
-                                    n_partitions_usk
+                                    n_user
                                 )
                             });
-                    });
+                            (usk, encrypted_header)
+                        },
+                        |(usk, encrypted_header)| {
+                            encrypted_header
+                                .decrypt(&cover_crypt, &usk, Some(&authenticated_data))
+                                .unwrap_or_else(|_| {
+                                    panic!(
+                                        "cannot decrypt header for {} ciphertext partition(s), {} \
+                                         usk partition(s)",
+                                        n_partition_ct + 1,
+                                        n_user
+                                    )
+                                });
+                        },
+                        BatchSize::SmallInput,
+                    );
                 },
             );
         }

src/core/primitives.rs

@@ -25,7 +25,7 @@ use crate::{
         EncryptionHint, Partition,
     },
     core::{Encapsulation, KeyEncapsulation, MasterPublicKey, MasterSecretKey, UserSecretKey},
-    data_struct::{Cursor, RevisionMap, RevisionVec},
+    data_struct::{List, RevisionMap, RevisionVec},
     Error,
 };
 
@@ -444,6 +444,54 @@ pub fn prune(msk: &mut MasterSecretKey, coordinates: &HashSet<Partition>) -> Res
     Ok(())
 }
 
+struct Merge<'a, M, U> {
+    msk_subkeys: M,
+    usk_subkeys: U,
+    usk_first_subkey: &'a SecretSubkey,
+    prepend: bool,
+}
+
+impl<'a, M, U> Merge<'a, M, U>
+where
+    M: Iterator<Item = &'a SecretSubkey>,
+    U: Iterator<Item = &'a SecretSubkey>,
+{
+    pub fn new(msk_subkeys: M, mut usk_subkeys: U) -> Self {
+        let usk_first_subkey = usk_subkeys.next().expect("first key");
+        Self {
+            msk_subkeys,
+            usk_subkeys,
+            usk_first_subkey,
+            prepend: true,
+        }
+    }
+}
+
+impl<'a, M, U> Iterator for Merge<'a, M, U>
+where
+    M: Iterator<Item = &'a SecretSubkey>,
+    U: Iterator<Item = &'a SecretSubkey>,
+{
+    type Item = SecretSubkey;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let msk_subkey = self.msk_subkeys.next()?;
+        if self.prepend {
+            if msk_subkey == self.usk_first_subkey {
+                self.prepend = false;
+            }
+            Some(msk_subkey.clone())
+        } else {
+            let usk_subkey = self.usk_subkeys.next()?;
+            if msk_subkey == usk_subkey {
+                Some(usk_subkey.clone())
+            } else {
+                None
+            }
+        }
+    }
+}
+
 /// Refresh a user key from the master secret key and the given decryption set.
 ///
 /// If `keep_old_rights` is set to false, old sub-keys are removed.
@@ -463,32 +511,24 @@ pub fn refresh(
 ) -> Result<(), Error> {
     verify_user_key_kmac(msk, usk)?;
 
-    usk.subkeys
-        .retain(|coordinate| msk.subkeys.contains_key(coordinate));
-
-    for (partition, user_chain) in usk.subkeys.iter_mut() {
-        let master_chain = msk.subkeys.get(partition).expect("at least one key");
-        // compare against all master subkeys or the last one to remove old rights
-        let mut master_chain = match keep_old_rights {
-            true => master_chain.iter().take(master_chain.len()).peekable(),
-            false => master_chain.iter().take(1).peekable(),
-        };
-
-        // 1 - add new master subkeys in user key if any
-        let user_first_key = user_chain.front().expect("have one key").clone();
-        let mut cursor = Cursor::new(user_chain);
-        while let Some(new_master_subkey) =
-            master_chain.next_if(|&master_subkey| master_subkey != &user_first_key)
-        {
-            cursor = cursor.prepend(new_master_subkey.clone());
-        }
-
-        // 2 - go through the remaining matching keys between the master and user chain
-        cursor = cursor.skip_while(|user_subkey| Some(user_subkey) == master_chain.next());
-
-        // 3 - old keys in USK not present in the MSK should be removed
-        cursor.cutoff();
-    }
+    usk.subkeys = usk
+        .subkeys
+        .iter()
+        .filter_map(|(coordinate, user_chain)| {
+            msk.subkeys.get(coordinate).map(|msk_chain| {
+                let msk_subkeys = if keep_old_rights {
+                    msk_chain.iter().take(msk_chain.len())
+                } else {
+                    msk_chain.iter().take(1)
+                };
+                let usk_subkeys = user_chain.iter();
+                (
+                    coordinate.clone(),
+                    Merge::new(msk_subkeys, usk_subkeys).collect::<List<_>>(),
+                )
+            })
+        })
+        .collect::<RevisionVec<_, _>>();
 
     // Update user key KMAC
     usk.kmac = compute_user_key_kmac(msk, usk);

src/core/serialization.rs

@@ -15,7 +15,7 @@ use crate::{
         Encapsulation, KeyEncapsulation, MasterPublicKey, MasterSecretKey, UserSecretKey,
         SYM_KEY_LENGTH,
     },
-    data_struct::{RevisionList, RevisionMap, RevisionVec},
+    data_struct::{List, RevisionMap, RevisionVec},
     CleartextHeader, EncryptedHeader, Error,
 };
 
@@ -151,7 +151,7 @@ impl Serializable for MasterSecretKey {
         for _ in 0..n_partitions {
             let partition = Partition::from(de.read_vec()?);
             let n_keys = <usize>::try_from(de.read_leb128_u64()?)?;
-            let chain: Result<RevisionList<_>, Self::Error> = (0..n_keys)
+            let chain: Result<List<_>, Self::Error> = (0..n_keys)
                 .map(|_| {
                     let sk_i = deserialize_option!(de, KyberSecretKey(de.read_array()?));
                     let x_i = de.read_array::<{ R25519PrivateKey::LENGTH }>()?;
@@ -224,7 +224,7 @@ impl Serializable for UserSecretKey {
             let partition = Partition::from(de.read_vec()?);
             // read all keys forming a chain and inserting them all at once.
             let n_keys = <usize>::try_from(de.read_leb128_u64()?)?;
-            let new_chain: Result<RevisionList<_>, _> = (0..n_keys)
+            let new_chain: Result<List<_>, _> = (0..n_keys)
                 .map(|_| {
                     let sk_i = deserialize_option!(de, KyberSecretKey(de.read_array()?));
                     let x_i = de.read_array::<{ R25519PrivateKey::LENGTH }>()?;