From 2fadc2e39c087cc0ddf2d9a16fced5903c0b80ee Mon Sep 17 00:00:00 2001
From: Sosthene <contact@sosthene.net>
Date: Fri, 29 Nov 2024 01:11:43 +0100
Subject: [PATCH] Update create_merkle_tree

---
 src/pcd.rs | 23 ++++++++++++++---------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/src/pcd.rs b/src/pcd.rs
index 4cc6bda..a2690d7 100644
--- a/src/pcd.rs
+++ b/src/pcd.rs
@@ -118,21 +118,26 @@ pub trait Pcd<'a>: Serialize + Deserialize<'a> {
         Ok(field2hash)
     }
 
-    /// We need to run `hash_fields` before
-    /// This will just take all the hash value and produces a merkle tree
-    fn create_merkle_tree(fields2hash: &Value) -> Result<MerkleTree<Sha256>> {
-        let leaves: Vec<[u8; 32]> = fields2hash
-            .as_object()
-            .unwrap()
+    /// We run this on the result of `hash_fields`
+    fn create_merkle_tree(&self) -> Result<MerkleTree<Sha256>> {
+        let map = self.to_value_object()?;
+        let leaves: Result<Vec<[u8; 32]>> = map
             .iter()
             .map(|(_, value)| {
                 let mut res = [0u8; 32];
-                res.copy_from_slice(&Vec::from_hex(value.as_str().unwrap()).unwrap());
-                res
+                if !value.is_string() {
+                    return Err(Error::msg("value is not a string"));
+                }
+                let vec = Vec::from_hex(value.as_str().unwrap())?;
+                if vec.len() != 32 {
+                    return Err(Error::msg("value must be 32B length"));
+                }
+                res.copy_from_slice(&vec);
+                Ok(res)
             })
             .collect();
 
-        let merkle_tree = MerkleTree::<Sha256>::from_leaves(&leaves);
+        let merkle_tree = MerkleTree::<Sha256>::from_leaves(&leaves?);
 
         Ok(merkle_tree)
     }