[template-graph] Parse Partial MPK From The WTG

Author: csegarragonz

Cargo.lock

@@ -1,7 +1,7 @@
 mod curve;
 mod hashing;
 pub mod policy;
-mod scheme;
+pub mod scheme;
 
 use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
 use base64::engine::{Engine as _, general_purpose};

accless/libs/abe4/src/lib.rs

@@ -17,3 +17,5 @@ serde_yaml = { workspace = true }
 
 [dev-dependencies]
 ark-std = { workspace = true }
+base64 = { workspace = true }
+ark-serialize = { workspace = true }

template-graph/Cargo.toml

@@ -1,9 +1,54 @@
-use accless_abe4::{UserAttribute, decrypt, encrypt, iota::Iota, keygen, setup, tau::Tau};
+use accless_abe4::{
+    UserAttribute, decrypt, encrypt,
+    iota::Iota,
+    keygen,
+    scheme::types::{MPK, PartialMPK},
+    setup,
+    tau::Tau,
+};
+use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
+use base64::engine::{Engine as _, general_purpose};
 use std::collections::HashSet;
 use template_graph::{TemplateGraph, policy_compiler};
 
+///
+/// # Description
+///
+/// Helper method to construct a full MPK from a template graph.
+///
+fn get_full_mpk_from_template_graph(template_graph: &TemplateGraph) -> MPK {
+    let mut mpk = MPK::new();
+
+    // User authority
+    let user_mpk_bytes = general_purpose::STANDARD
+        .decode(&template_graph.authorities.user.mpk_abe)
+        .unwrap();
+    let user_partial_mpk = PartialMPK::deserialize_compressed(&user_mpk_bytes[..]).unwrap();
+    mpk.add_partial_key(user_partial_mpk);
+
+    // Attestation services
+    for as_service in &template_graph.authorities.attestation_services {
+        let as_mpk_bytes = general_purpose::STANDARD
+            .decode(&as_service.mpk_abe)
+            .unwrap();
+        let as_partial_mpk = PartialMPK::deserialize_compressed(&as_mpk_bytes[..]).unwrap();
+        mpk.add_partial_key(as_partial_mpk);
+    }
+
+    // APS
+    if let Some(aps_vec) = &template_graph.authorities.aps {
+        for aps in aps_vec {
+            let aps_mpk_bytes = general_purpose::STANDARD.decode(&aps.mpk_abe).unwrap();
+            let aps_partial_mpk = PartialMPK::deserialize_compressed(&aps_mpk_bytes[..]).unwrap();
+            mpk.add_partial_key(aps_partial_mpk);
+        }
+    }
+
+    mpk
+}
+
 #[test]
-fn test_encrypt_decrypt_workflow() {
+fn test_encrypt_decrypt_workflow_with_real_mpk() {
     let yaml_content = r#"
 version: 1
 workflow:
@@ -12,13 +57,13 @@ workflow:
 authorities:
   user:
     id: user_42
-    mpk_abe: base64:mpk_abe_user
+    mpk_abe: ""
   attestation-services:
     - id: maa
-      mpk_abe: base64:mpk_abe_maa
+      mpk_abe: ""
   aps:
     - id: finra
-      mpk_abe: base64:mpk_abe_finra
+      mpk_abe: ""
 
 nodes:
 - name: fetch_public
@@ -44,12 +89,11 @@ output:
   dir: ./tests/out-ciphertexts
     "#;
 
-    let template_graph = TemplateGraph::from_yaml(yaml_content).unwrap();
-    let policies = policy_compiler::compile_policies(&template_graph);
+    let mut template_graph = TemplateGraph::from_yaml(yaml_content).unwrap();
 
     let mut rng = ark_std::test_rng();
 
-    // Setup ABE
+    // Setup ABE and get partial MPKs
     let mut auths: HashSet<String> = template_graph
         .authorities
         .attestation_services
@@ -65,10 +109,38 @@ output:
     let auths_str: Vec<&str> = auths.iter().map(|s| s.as_str()).collect();
     let (msk, mpk) = setup(&mut rng, &auths_str);
 
+    // Update template graph with real MPKs
+    for (auth, partial_mpk) in &mpk.partial_keys {
+        let mut mpk_bytes = Vec::new();
+        partial_mpk.serialize_compressed(&mut mpk_bytes).unwrap();
+        let mpk_b64 = general_purpose::STANDARD.encode(&mpk_bytes);
+
+        if auth == &template_graph.authorities.user.id {
+            template_graph.authorities.user.mpk_abe = mpk_b64;
+        } else if let Some(as_service) = template_graph
+            .authorities
+            .attestation_services
+            .iter_mut()
+            .find(|a| &a.id == auth)
+        {
+            as_service.mpk_abe = mpk_b64;
+        } else if let Some(aps) = template_graph
+            .authorities
+            .aps
+            .as_mut()
+            .and_then(|aps_vec| aps_vec.iter_mut().find(|a| &a.id == auth))
+        {
+            aps.mpk_abe = mpk_b64;
+        }
+    }
+
+    let policies = policy_compiler::compile_policies(&template_graph);
+    let full_mpk = get_full_mpk_from_template_graph(&template_graph);
+
     // Test encryption and decryption for each node
     for node in &template_graph.nodes {
         if let Some(policy) = policies.get(&node.name) {
-            let (k_enc, ct) = encrypt(&mut rng, &mpk, &policy, &Tau::new(&policy));
+            let (k_enc, ct) = encrypt(&mut rng, &full_mpk, &policy, &Tau::new(&policy));
 
             // Simulate a user with the required attributes
             let mut user_attrs = Vec::new();