fix: allow partial rekeying of disabled hybrid partitions

src/core/primitives.rs

@@ -293,20 +293,14 @@ pub fn update(
             let h_i = &h * x_i;
             // Set the correct hybridization property.
             let (sk_i, pk_i) = if is_hybridized == EncryptionHint::Hybridized {
-                let (pk_i, _) = mpk.subkeys.get(partition).ok_or_else(|| {
-                    Error::KeyError(
-                        "Kyber public key cannot be computed from the secret key.".to_string(),
-                    )
-                })?;
+                let pk_i = mpk
+                    .subkeys
+                    .get(partition)
+                    .map(|(pk_i, _)| pk_i)
+                    .unwrap_or(&None);
 
                 if sk_i.is_some() {
-                    if pk_i.is_some() {
-                        (sk_i.clone(), pk_i.clone())
-                    } else {
-                        return Err(Error::KeyError(
-                            "Kyber public key cannot be computed from the secret key.".to_string(),
-                        ));
-                    }
+                    (sk_i.clone(), pk_i.clone())
                 } else {
                     let (mut sk_i, mut pk_i) = (
                         KyberSecretKey([0; KYBER_INDCPA_SECRETKEYBYTES]),
@@ -318,11 +312,17 @@ pub fn update(
             } else {
                 (None, None)
             };
-            new_sub_sk.insert(partition.clone(), (sk_i, x_i.clone()));
+
             if write_status == AttributeStatus::EncryptDecrypt {
                 // Only add non read only partition to the public key
+                if sk_i.is_some() && pk_i.is_none() {
+                    return Err(Error::KeyError(
+                        "Kyber public key cannot be computed from the secret key.".to_string(),
+                    ));
+                }
                 new_sub_pk.insert(partition.clone(), (pk_i, h_i));
             }
+            new_sub_sk.insert(partition.clone(), (sk_i, x_i.clone()));
         } else {
             // Create new entry.
             let x_i = R25519PrivateKey::new(rng);
@@ -338,7 +338,10 @@ pub fn update(
                 (None, None)
             };
             new_sub_sk.insert(partition.clone(), (sk_pq, x_i));
-            new_sub_pk.insert(partition.clone(), (pk_pq, h_i));
+            if write_status == AttributeStatus::EncryptDecrypt {
+                // Only add non read only partition to the public key
+                new_sub_pk.insert(partition.clone(), (pk_pq, h_i));
+            }
         }
     }
 

src/test_utils/mod.rs

@@ -315,8 +315,9 @@ mod tests {
         cover_crypt.update_master_keys(&policy, &mut msk, &mut mpk)?;
         let new_partitions_msk: Vec<Partition> = msk.subkeys.clone().into_keys().collect();
         let new_partitions_mpk: Vec<Partition> = mpk.subkeys.clone().into_keys().collect();
-        // 5 is the size of the security level dimension
+        // the disabled partition have been removed from mpk
         assert_eq!(new_partitions_msk.len() - 5, new_partitions_mpk.len());
+        // msk hasn't changed
         assert_eq!(new_partitions_msk.len(), partitions_msk.len());
 
         assert!(encrypted_header
@@ -358,6 +359,16 @@ mod tests {
             .decrypt(&cover_crypt, &top_secret_fin_usk, None)
             .is_ok());
 
+        //
+        // Rotating the disabled attribute should only change the msk
+        policy.rotate(&Attribute::new("Department", "FIN"))?;
+        cover_crypt.update_master_keys(&policy, &mut msk, &mut mpk)?;
+        let new_partitions_msk: Vec<Partition> = msk.subkeys.clone().into_keys().collect();
+        let new_partitions_mpk: Vec<Partition> = mpk.subkeys.clone().into_keys().collect();
+        // 5 new partitions added to the msk
+        assert_eq!(new_partitions_msk.len() - 10, new_partitions_mpk.len());
+        assert_eq!(new_partitions_msk.len(), partitions_msk.len() + 5);
+
         Ok(())
     }