Merge pull request #318 from AdExNetwork/hotfix-316-merkle-tree-missing-0-leaves-check

Hotfix #316 MerkleTree missing 0 leaves check

Author: samparsky

Cargo.lock

@@ -230,12 +230,12 @@ mod test {
             .expect("The timestamp should be able to be converted to u64");
         BigEndian::write_uint(&mut timestamp_buf[26..], n, 6);
 
-        let merkle_tree = MerkleTree::new(&[timestamp_buf]);
+        let merkle_tree = MerkleTree::new(&[timestamp_buf]).expect("Should instantiate");
 
         let channel_id = "061d5e2a67d0a9a10f1c732bca12a676d83f79663a396f7d87b3e30b9b411088";
 
         let state_root = get_signable_state_root(
-            &hex::decode(&channel_id).expect("fialed"),
+            &hex::decode(&channel_id).expect("failed"),
             &merkle_tree.root(),
         )
         .expect("Should get state_root");

adapter/src/lib.rs

@@ -20,6 +20,7 @@ slog = { version = "^2.5.2" , features = ["max_level_trace"] }
 slog-term = "^2.4.2"
 slog-async = "^2.3.0"
 # Domain
+thiserror = "^1.0"
 chrono = { version = "0.4", features = ["serde"] }
 time = "0.1.42"
 hex = "0.3.2"

primitives/Cargo.toml

@@ -1,14 +1,19 @@
-use merkletree::hash::Algorithm;
-use merkletree::merkle;
-use merkletree::merkle::VecStore;
-use merkletree::proof::Proof;
+use merkletree::{hash::Algorithm, merkle, merkle::VecStore, proof::Proof};
+use std::fmt;
 use std::hash::Hasher;
 use std::iter::FromIterator;
+use thiserror::Error;
 use tiny_keccak::Keccak;
 
 #[derive(Clone)]
 struct KeccakAlgorithm(Keccak);
 
+impl fmt::Debug for KeccakAlgorithm {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "Keccak256 Algorithm")
+    }
+}
+
 impl KeccakAlgorithm {
     pub fn new() -> KeccakAlgorithm {
         KeccakAlgorithm(Keccak::new_keccak256())
@@ -72,39 +77,49 @@ impl Algorithm<MerkleItem> for KeccakAlgorithm {
 type ExternalMerkleTree =
     merkletree::merkle::MerkleTree<MerkleItem, KeccakAlgorithm, VecStore<MerkleItem>>;
 
-#[derive(Clone)]
+#[derive(Debug, Clone)]
 enum Tree {
     SingleItem(MerkleItem),
     MerkleTree(ExternalMerkleTree),
 }
 
+#[derive(Debug, Error, Eq, PartialEq)]
+pub enum Error {
+    #[error("No leaves were provided")]
+    ZeroLeaves,
+}
+
+#[derive(Debug)]
 pub struct MerkleTree {
     tree: Tree,
     root: MerkleItem,
 }
 
 impl MerkleTree {
-    pub fn new(data: &[MerkleItem]) -> MerkleTree {
+    pub fn new(data: &[MerkleItem]) -> Result<MerkleTree, Error> {
         let mut leaves: Vec<MerkleItem> = data.to_owned();
-
-        let tree: Tree = if leaves.len() == 1 {
-            Tree::SingleItem(leaves[0].to_owned())
-        } else {
-            // sort the merkle tree leaves
-            leaves.sort();
-            // remove duplicates
-            leaves.dedup_by(|a, b| a == b);
-
-            let merkletree = merkle::MerkleTree::from_iter(leaves);
-            Tree::MerkleTree(merkletree)
+        // sort the MerkleTree leaves
+        leaves.sort();
+        // remove duplicates **before** we check the leaves length
+        leaves.dedup_by(|a, b| a == b);
+
+        let tree = match leaves.len() {
+            0 => return Err(Error::ZeroLeaves),
+            // should never `panic!`, we have a single leaf after all
+            1 => Tree::SingleItem(leaves.remove(0)),
+            _ => {
+                let merkletree = merkle::MerkleTree::from_iter(leaves);
+
+                Tree::MerkleTree(merkletree)
+            }
         };
 
         let root: MerkleItem = match &tree {
             Tree::SingleItem(root) => root.to_owned(),
             Tree::MerkleTree(merkletree) => merkletree.root(),
         };
 
-        MerkleTree { tree, root }
+        Ok(MerkleTree { tree, root })
     }
 
     pub fn root(&self) -> MerkleItem {
@@ -134,6 +149,12 @@ mod test {
     use super::*;
     use hex::FromHex;
 
+    #[test]
+    fn it_returns_error_on_zero_leaves() {
+        let error = MerkleTree::new(&[]).expect_err("ZeroLeaves error expected");
+        assert_eq!(Error::ZeroLeaves, error);
+    }
+
     #[test]
     fn it_generates_correct_merkle_tree_that_correlates_with_js_impl() {
         let h1 = <[u8; 32]>::from_hex(
@@ -145,7 +166,7 @@ mod test {
         )
         .unwrap();
 
-        let top = MerkleTree::new(&[h1, h2]);
+        let top = MerkleTree::new(&[h1, h2]).expect("Should create MerkleTree");
 
         let root = hex::encode(&top.root());
 
@@ -172,7 +193,7 @@ mod test {
         .unwrap();
 
         // duplicate leaves
-        let top = MerkleTree::new(&[h1, h2, h2]);
+        let top = MerkleTree::new(&[h1, h2, h2]).expect("Should create MerkleTree");
 
         let root = hex::encode(&top.root());
 
@@ -204,7 +225,7 @@ mod test {
         .unwrap();
 
         // odd leaves
-        let top = MerkleTree::new(&[h1, h2, h3]);
+        let top = MerkleTree::new(&[h1, h2, h3]).expect("Should create MerkleTree");
 
         let root = hex::encode(&top.root());
 

primitives/src/merkle_tree.rs

@@ -46,7 +46,7 @@ pub async fn postgres_connection() -> Result<DbPool, bb8_postgres::tokio_postgre
         .user(POSTGRES_USER.as_str())
         .password(POSTGRES_PASSWORD.as_str())
         .host(POSTGRES_HOST.as_str())
-        .port(POSTGRES_PORT.clone());
+        .port(*POSTGRES_PORT);
     if let Some(db) = POSTGRES_DB.clone() {
         config.dbname(&db);
     }
@@ -62,7 +62,7 @@ pub async fn setup_migrations(environment: &str) {
         .database_user(POSTGRES_USER.as_str())
         .database_password(POSTGRES_PASSWORD.as_str())
         .database_host(POSTGRES_HOST.as_str())
-        .database_port(POSTGRES_PORT.clone())
+        .database_port(*POSTGRES_PORT)
         .database_name(&POSTGRES_DB.as_ref().unwrap_or(&POSTGRES_USER))
         .build()
         .expect("Should build migration settings");

sentry/src/db.rs

@@ -58,7 +58,7 @@ async fn store(db: &DbPool, channel_id: &ChannelId, logger: &Logger, recorder: R
                 aggregate: new_aggr(&channel_id),
             };
             channel_recorder.insert(channel_id.to_owned(), record);
-        };
+        }
     }
 }
 

sentry/src/event_aggregator.rs

@@ -2,8 +2,6 @@ use crate::analytics_recorder::get_payout;
 use primitives::sentry::{AggregateEvents, Event, EventAggregate};
 use primitives::{BigNum, Channel, ValidatorId};
 
-// @TODO: Remove attribute once we use this function!
-#[allow(dead_code)]
 pub(crate) fn reduce(channel: &Channel, initial_aggr: &mut EventAggregate, ev: &Event) {
     match ev {
         Event::Impression { publisher, .. } => {

sentry/src/event_reducer.rs

@@ -77,7 +77,7 @@ fn get_request_ip(req: &Request<Body>) -> Option<String> {
         .get("true-client-ip")
         .or_else(|| req.headers().get("x-forwarded-for"))
         .and_then(|hv| hv.to_str().map(ToString::to_string).ok())
-        .map(|token| token.split(',').nth(0).map(ToString::to_string))
+        .map(|token| token.split(',').next().map(ToString::to_string))
         .flatten()
 }
 

sentry/src/middleware/auth.rs

@@ -18,7 +18,7 @@ async fn send_heartbeat<A: Adapter + 'static>(iface: &SentryApi<A>) -> Result<()
         .expect("The timestamp should be able to be converted to u64");
     BigEndian::write_uint(&mut timestamp_buf[26..], milliseconds, 6);
 
-    let merkle_tree = MerkleTree::new(&[timestamp_buf]);
+    let merkle_tree = MerkleTree::new(&[timestamp_buf])?;
 
     let state_root_raw = get_signable_state_root(&iface.channel.id, &merkle_tree.root())?;
     let state_root = hex::encode(state_root_raw);

validator_worker/src/heartbeat.rs

@@ -33,7 +33,7 @@ pub(crate) fn get_state_root_hash<A: Adapter + 'static>(
         .map(|(acc, amount)| get_balance_leaf(acc, amount))
         .collect::<Result<_, _>>()?;
 
-    let tree = MerkleTree::new(&elems);
+    let tree = MerkleTree::new(&elems)?;
     // keccak256(channelId, balanceRoot
     get_signable_state_root(&iface.channel.id, &tree.root())
 }