feat(tree): UpdateProof fixes

crates/common/src/hashchain.rs

@@ -67,77 +67,77 @@ impl Hashchain {
         }
     }
 
-    pub fn validate(&self) -> Result<()> {
+    pub fn verify_last_entry(&self) -> Result<()> {
+        if self.entries.is_empty() {
+            return Ok(());
+        }
+
         let mut valid_keys: HashSet<PublicKey> = HashSet::new();
 
-        for (index, entry) in self.entries.iter().enumerate() {
-            self.validate_entry(index, entry, &valid_keys)?;
-            self.update_valid_keys(&entry.operation, &mut valid_keys);
+        for (index, entry) in self.entries.iter().enumerate().take(self.entries.len() - 1) {
+            match &entry.operation {
+                Operation::CreateAccount(args) => {
+                    if index != 0 {
+                        bail!("CreateAccount operation must be the first entry");
+                    }
+                    valid_keys.insert(args.value.clone());
+                }
+                Operation::AddKey(args) => {
+                    valid_keys.insert(args.value.clone());
+                }
+                Operation::RevokeKey(args) => {
+                    valid_keys.remove(&args.value);
+                }
+            }
         }
 
-        Ok(())
-    }
+        let last_entry = self.entries.last().unwrap();
+        let last_index = self.entries.len() - 1;
+        if last_index > 0 {
+            let prev_entry = &self.entries[last_index - 1];
+            if last_entry.previous_hash != prev_entry.hash {
+                bail!("Previous hash mismatch for the last entry");
+            }
+        }
 
-    fn validate_entry(
-        &self,
-        index: usize,
-        entry: &HashchainEntry,
-        valid_keys: &HashSet<PublicKey>,
-    ) -> Result<()> {
-        match &entry.operation {
+        match &last_entry.operation {
             Operation::CreateAccount(args) => {
-                if index != 0 {
+                if last_index != 0 {
                     bail!("CreateAccount operation must be the first entry");
                 }
                 self.verify_signature(
                     &args.value,
-                    &bincode::serialize(&entry.operation.without_signature())?,
+                    &bincode::serialize(&last_entry.operation.without_signature())?,
                     &args.signature,
                 )?;
             }
             Operation::AddKey(args) | Operation::RevokeKey(args) => {
                 let signing_key = self.get_key_at_index(args.signature.key_idx as usize)?;
-                if !valid_keys.contains(&signing_key) {
-                    bail!("Operation signed with revoked or non-existent key");
+                if !valid_keys.contains(signing_key) {
+                    bail!("Last operation signed with revoked or non-existent key");
                 }
                 self.verify_signature(
-                    &signing_key,
-                    &bincode::serialize(&entry.operation.without_signature())?,
+                    signing_key,
+                    &bincode::serialize(&last_entry.operation.without_signature())?,
                     &args.signature.signature,
                 )?;
             }
         }
+
         Ok(())
     }
 
-    fn update_valid_keys(&self, operation: &Operation, valid_keys: &mut HashSet<PublicKey>) {
-        match operation {
-            Operation::CreateAccount(args) => {
-                valid_keys.insert(args.value.clone());
-            }
-            Operation::AddKey(args) => {
-                valid_keys.insert(args.value.clone());
-            }
-            Operation::RevokeKey(args) => {
-                valid_keys.remove(&args.value);
-            }
-        }
+    pub(crate) fn insert_unsafe(&self, new_entry: HashchainEntry) -> Hashchain {
+        let mut new = self.clone();
+        new.entries.push(new_entry);
+        new
     }
 
-    pub fn get_key_at_index(&self, idx: usize) -> Result<PublicKey> {
-        let hc_entry: Option<&HashchainEntry> = self.entries.get(idx);
-
-        if let Some(entry) = hc_entry {
-            match entry.operation.get_public_key() {
-                Some(key) => {
-                    return Ok(key);
-                }
-                None => {
-                    bail!("Key at index {idx} does not exist");
-                }
-            }
-        }
-        Err(anyhow!("No hashchain entry found at idx {idx}"))
+    pub fn get_key_at_index(&self, idx: usize) -> Result<&PublicKey> {
+        self.entries
+            .get(idx)
+            .and_then(|entry| entry.operation.get_public_key())
+            .ok_or_else(|| anyhow!("No valid public key found at index {}", idx))
     }
 
     pub fn get_valid_keys(&self) -> HashSet<PublicKey> {
@@ -231,7 +231,7 @@ impl Hashchain {
                 }
 
                 let message = bincode::serialize(&operation.without_signature())?;
-                self.verify_signature(&signing_key, &message, &args.signature.signature)
+                self.verify_signature(signing_key, &message, &args.signature.signature)
             }
             Operation::CreateAccount(args) => {
                 let message = bincode::serialize(&operation.without_signature())?;

crates/common/src/operation.rs

@@ -77,11 +77,11 @@ impl Operation {
         }
     }
 
-    pub fn get_public_key(&self) -> Option<PublicKey> {
+    pub fn get_public_key(&self) -> Option<&PublicKey> {
         match self {
-            Operation::AddKey(args) => Some(args.value.clone()),
-            Operation::RevokeKey(args) => Some(args.value.clone()),
-            Operation::CreateAccount(args) => Some(args.value.clone()),
+            Operation::AddKey(args) => Some(&args.value),
+            Operation::RevokeKey(args) => Some(&args.value),
+            Operation::CreateAccount(args) => Some(&args.value),
         }
     }
 

crates/common/src/test_utils.rs

@@ -65,7 +65,7 @@ impl TestTreeState {
 
         let proof = self
             .tree
-            .update(account.key_hash, account.hashchain)
+            .update(account.key_hash, account.hashchain.last().unwrap().clone())
             .expect("Update should succeed");
         Ok(proof)
     }
@@ -121,12 +121,14 @@ pub fn create_random_insert(state: &mut TestTreeState, rng: &mut StdRng) -> Inse
         let random_string: String = (0..10)
             .map(|_| rng.sample(rand::distributions::Alphanumeric) as char)
             .collect();
-        let hc = Hashchain::new(random_string);
+        let sk = create_mock_signing_key();
+        let hc = create_mock_hashchain(random_string.as_str(), &sk); //Hashchain::new(random_string);
         let key = hc.get_keyhash();
 
         if !state.inserted_keys.contains(&key) {
             let proof = state.tree.insert(key, hc).expect("Insert should succeed");
             state.inserted_keys.insert(key);
+            state.signing_keys.insert(random_string, sk);
             return proof;
         }
     }
@@ -170,7 +172,10 @@ pub fn create_random_update(state: &mut TestTreeState, rng: &mut StdRng) -> Upda
     hc.perform_operation(operation)
         .expect("Adding to hashchain should succeed");
 
-    state.tree.update(key, hc).expect("Update should succeed")
+    state
+        .tree
+        .update(key, hc.last().unwrap().clone())
+        .expect("Update should succeed")
 }
 
 pub fn create_mock_signature(signing_key: &SigningKey, message: &[u8]) -> SignatureBundle {

crates/common/src/tree.rs

@@ -11,7 +11,7 @@ use serde::{ser::SerializeTupleStruct, Deserialize, Serialize};
 use std::sync::Arc;
 
 use crate::{
-    hashchain::Hashchain,
+    hashchain::{Hashchain, HashchainEntry},
     operation::{CreateAccountArgs, KeyOperationArgs, Operation, ServiceChallengeInput},
 };
 
@@ -205,7 +205,7 @@ impl InsertProof {
             value,
         )?;
 
-        self.value.validate()?;
+        self.value.verify_last_entry()?;
 
         Ok(())
     }
@@ -217,29 +217,43 @@ pub struct UpdateProof {
     pub new_root: RootHash,
 
     pub key: KeyHash,
-    pub new_value: Hashchain,
+    pub old_value: Hashchain,
+    pub new_entry: HashchainEntry,
 
-    pub proof: UpdateMerkleProof<Hasher>,
+    /// Inclusion proof of [`old_value`]
+    pub inclusion_proof: SparseMerkleProof<Hasher>,
+    /// Update proof for [`key`] to be updated with [`new_entry`]
+    pub update_proof: UpdateMerkleProof<Hasher>,
 }
 
 impl UpdateProof {
     pub fn verify(&self) -> Result<()> {
-        let new_value = bincode::serialize(&self.new_value)?;
-
-        self.proof.clone().verify_update(
+        // Verify existence of old value.
+        // Otherwise, any arbitrary hashchain could be set
+        let old_value = bincode::serialize(&self.old_value)?;
+        self.inclusion_proof
+            .verify_existence(self.old_root, self.key, old_value)?;
+
+        // Append the new entry and verify it's validity
+        let new_hashchain = self.old_value.insert_unsafe(self.new_entry.clone());
+        new_hashchain.verify_last_entry()?;
+
+        // Ensure the update proof corresponds to the new hashchain value
+        let new_value = bincode::serialize(&new_hashchain)?;
+        self.update_proof.clone().verify_update(
             self.old_root,
             self.new_root,
             vec![(self.key, Some(new_value))],
         )?;
 
-        self.new_value.validate()
+        Ok(())
     }
 }
 
 pub trait SnarkableTree {
     fn process_operation(&mut self, operation: &Operation) -> Result<Proof>;
     fn insert(&mut self, key: KeyHash, value: Hashchain) -> Result<InsertProof>;
-    fn update(&mut self, key: KeyHash, value: Hashchain) -> Result<UpdateProof>;
+    fn update(&mut self, key: KeyHash, value: HashchainEntry) -> Result<UpdateProof>;
     fn get(&self, key: KeyHash) -> Result<Result<Hashchain, NonMembershipProof>>;
 }
 
@@ -323,10 +337,10 @@ where
                     .get(key_hash)?
                     .map_err(|_| anyhow!("Failed to get hashchain for ID {}", id))?;
 
-                current_chain.perform_operation(operation.clone())?;
+                let new_entry = current_chain.perform_operation(operation.clone())?;
 
                 debug!("updating hashchain for user id {}", id.clone());
-                let proof = self.update(key_hash, current_chain.clone())?;
+                let proof = self.update(key_hash, new_entry.clone())?;
 
                 Ok(Proof::Update(proof))
             }
@@ -401,17 +415,21 @@ where
         })
     }
 
-    fn update(&mut self, key: KeyHash, value: Hashchain) -> Result<UpdateProof> {
-        let serialized_value = Self::serialize_value(&value)?;
-
+    fn update(&mut self, key: KeyHash, new_entry: HashchainEntry) -> Result<UpdateProof> {
         let old_root = self.get_current_root()?;
-        let (old_value, _) = self.jmt.get_with_proof(key, self.epoch)?;
+        let (old_value, inclusion_proof) = self.jmt.get_with_proof(key, self.epoch)?;
 
         if old_value.is_none() {
             bail!("Key does not exist");
         }
 
-        let (new_root, proof, tree_update_batch) = self.jmt.put_value_set_with_proof(
+        let old_value: Hashchain = bincode::deserialize(old_value.unwrap().as_slice())?;
+        let new_hashchain = old_value.insert_unsafe(new_entry.clone());
+        new_hashchain.verify_last_entry()?;
+
+        let serialized_value = Self::serialize_value(&new_hashchain)?;
+
+        let (new_root, update_proof, tree_update_batch) = self.jmt.put_value_set_with_proof(
             vec![(key, Some(serialized_value.clone()))],
             self.epoch + 1,
         )?;
@@ -421,9 +439,11 @@ where
         Ok(UpdateProof {
             old_root,
             new_root,
+            inclusion_proof,
+            old_value,
             key,
-            new_value: value,
-            proof,
+            new_entry,
+            update_proof,
         })
     }
 

crates/nova/src/update.rs

@@ -80,7 +80,7 @@ where
             |lc| lc + pre_insertion_root.get_variable(),
         );
 
-        let update_proof = &self.update_proof.proof.proofs()[0];
+        let update_proof = &self.update_proof.update_proof.proofs()[0];
 
         let leaf = &update_proof
             .leaf()