Merge pull request #139 from NIC619/Update_deposit_contract

Add receipt tree merkle branch getter

contracts/validator_registration.v.py

@@ -6,10 +6,10 @@
 TWO_TO_POWER_OF_TREE_DEPTH: constant(uint256) = 4294967296  # 2**32
 SECONDS_PER_DAY: constant(uint256) = 86400
 
-Eth1Deposit: event({previous_receipt_root: bytes32, data: bytes[2064], deposit_count: uint256})
-ChainStart: event({receipt_root: bytes32, time: bytes[8]})
+Deposit: event({previous_deposit_root: bytes32, data: bytes[2064], merkle_tree_index: bytes[8]})
+ChainStart: event({deposit_root: bytes32, time: bytes[8]})
 
-receipt_tree: map(uint256, bytes32)
+deposit_tree: map(uint256, bytes32)
 deposit_count: uint256
 full_deposit_count: uint256
 
@@ -23,24 +23,35 @@ def deposit(deposit_input: bytes[2048]):
     msg_gwei_bytes8: bytes[8] = slice(concat("", convert(msg.value / GWEI_PER_ETH, bytes32)), start=24, len=8)
     timestamp_bytes8: bytes[8] = slice(concat("", convert(block.timestamp, bytes32)), start=24, len=8)
     deposit_data: bytes[2064] = concat(msg_gwei_bytes8, timestamp_bytes8, deposit_input)
+    merkle_tree_index: bytes[8] = slice(concat("", convert(index, bytes32)), start=24, len=8)
 
-    log.Eth1Deposit(self.receipt_tree[1], deposit_data, self.deposit_count)
+    log.Deposit(self.deposit_tree[1], deposit_data, merkle_tree_index)
 
     # add deposit to merkle tree
-    self.receipt_tree[index] = sha3(deposit_data)
-    for i in range(32):  # DEPOSIT_CONTRACT_TREE_DEPTH (range of constant var not yet supported)
+    self.deposit_tree[index] = sha3(deposit_data)
+    for i in range(DEPOSIT_CONTRACT_TREE_DEPTH):  # DEPOSIT_CONTRACT_TREE_DEPTH (range of constant var not yet supported)
         index /= 2
-        self.receipt_tree[index] = sha3(concat(self.receipt_tree[index * 2], self.receipt_tree[index * 2 + 1]))
+        self.deposit_tree[index] = sha3(concat(self.deposit_tree[index * 2], self.deposit_tree[index * 2 + 1]))
 
     self.deposit_count += 1
     if msg.value == as_wei_value(MAX_DEPOSIT, "ether"):
         self.full_deposit_count += 1
         if self.full_deposit_count == CHAIN_START_FULL_DEPOSIT_THRESHOLD:
             timestamp_day_boundary: uint256 = as_unitless_number(block.timestamp) - as_unitless_number(block.timestamp) % SECONDS_PER_DAY + SECONDS_PER_DAY
             timestamp_day_boundary_bytes8: bytes[8] = slice(concat("", convert(timestamp_day_boundary, bytes32)), start=24, len=8)
-            log.ChainStart(self.receipt_tree[1], timestamp_day_boundary_bytes8)
+            log.ChainStart(self.deposit_tree[1], timestamp_day_boundary_bytes8)
 
 @public
 @constant
-def get_receipt_root() -> bytes32:
-    return self.receipt_tree[1]
+def get_deposit_root() -> bytes32:
+    return self.deposit_tree[1]
+
+@public
+@constant
+def get_branch(leaf: uint256) -> bytes32[32]: # size is DEPOSIT_CONTRACT_TREE_DEPTH (symbolic const not supported)
+    branch: bytes32[32] # size is DEPOSIT_CONTRACT_TREE_DEPTH
+    index: uint256 = leaf + TWO_TO_POWER_OF_TREE_DEPTH
+    for i in range(DEPOSIT_CONTRACT_TREE_DEPTH):
+        branch[i] = self.deposit_tree[bitwise_xor(index, 1)]
+        index /= 2
+    return branch

tests/contracts/conftest.py

@@ -16,6 +16,7 @@
 MIN_DEPOSIT = 1  # ETH
 MAX_DEPOSIT = 32  # ETH
 DEPOSIT_CONTRACT_TREE_DEPTH = 32
+TWO_TO_POWER_OF_TREE_DEPTH = 4294967296  # 2**32
 
 
 def get_dirs(path):

tests/contracts/test_deposit.py

@@ -4,14 +4,17 @@
 
 import eth_utils
 
+from eth_hash.auto import keccak as hash
+
 from tests.contracts.conftest import (
     MAX_DEPOSIT,
     MIN_DEPOSIT,
     DEPOSIT_CONTRACT_TREE_DEPTH,
+    TWO_TO_POWER_OF_TREE_DEPTH,
 )
 
 
-def compute_merkle_root(w3, leaf_nodes):
+def compute_merkle_root(leaf_nodes):
     assert len(leaf_nodes) >= 1
     empty_node = b'\x00' * 32
     child_nodes = leaf_nodes[:]
@@ -20,11 +23,21 @@ def compute_merkle_root(w3, leaf_nodes):
         if len(child_nodes) % 2 == 1:
             child_nodes.append(empty_node)
         for j in range(0, len(child_nodes), 2):
-            parent_nodes.append(w3.sha3(child_nodes[j] + child_nodes[j+1]))
+            parent_nodes.append(hash(child_nodes[j] + child_nodes[j+1]))
         child_nodes = parent_nodes
     return child_nodes[0]
 
 
+def verify_merkle_branch(leaf, branch, depth, index, root):
+    value = leaf
+    for i in range(depth):
+        if index // (2**i) % 2:
+            value = hash(branch[i] + value)
+        else:
+            value = hash(value + branch[i])
+    return value == root
+
+
 @pytest.mark.parametrize(
     'success,amount_deposit',
     [
@@ -46,7 +59,7 @@ def test_deposit_amount(registration_contract, w3, success, amount_deposit, asse
 
 
 def test_deposit_log(registration_contract, a0, w3):
-    log_filter = registration_contract.events.Eth1Deposit.createFilter(
+    log_filter = registration_contract.events.Deposit.createFilter(
         fromBlock='latest',
     )
 
@@ -64,14 +77,14 @@ def test_deposit_log(registration_contract, a0, w3):
         amount_bytes8 = deposit_amount[i].to_bytes(8, 'big')
         timestamp_bytes8 = int(w3.eth.getBlock(w3.eth.blockNumber)['timestamp']).to_bytes(8, 'big')
         if i == 0:
-            assert log['previous_receipt_root'] == b'\x00' * 32
+            assert log['previous_deposit_root'] == b'\x00' * 32
         else:
-            assert log['previous_receipt_root'] != b'\x00' * 32
+            assert log['previous_deposit_root'] != b'\x00' * 32
         assert log['data'] == amount_bytes8 + timestamp_bytes8 + deposit_input
-        assert log['deposit_count'] == i
+        assert log['merkle_tree_index'] == (i + TWO_TO_POWER_OF_TREE_DEPTH).to_bytes(8, 'big')
 
 
-def test_reciept_tree(registration_contract, w3, assert_tx_failed):
+def test_receipt_tree(registration_contract, w3, assert_tx_failed):
     deposit_amount = [randint(MIN_DEPOSIT, MAX_DEPOSIT) * eth_utils.denoms.gwei for _ in range(10)]
 
     leaf_nodes = []
@@ -87,8 +100,17 @@ def test_reciept_tree(registration_contract, w3, assert_tx_failed):
         amount_bytes8 = deposit_amount[i].to_bytes(8, 'big')
         data = amount_bytes8 + timestamp_bytes8 + deposit_input
         leaf_nodes.append(w3.sha3(data))
-        root = compute_merkle_root(w3, leaf_nodes)
-        assert registration_contract.functions.get_receipt_root().call() == root
+        root = compute_merkle_root(leaf_nodes)
+        assert registration_contract.functions.get_deposit_root().call() == root
+        index = randint(0, i)
+        branch = registration_contract.functions.get_branch(index).call()
+        assert verify_merkle_branch(
+            leaf_nodes[index],
+            branch,
+            DEPOSIT_CONTRACT_TREE_DEPTH,
+            index,
+            root
+        )
 
 
 def test_chain_start(modified_registration_contract, w3, assert_tx_failed):
@@ -131,7 +153,7 @@ def test_chain_start(modified_registration_contract, w3, assert_tx_failed):
     timestamp = int(w3.eth.getBlock(w3.eth.blockNumber)['timestamp'])
     timestamp_day_boundary = timestamp + (86400 - timestamp % 86400)
     log = logs[0]['args']
-    assert log['receipt_root'] == modified_registration_contract.functions.get_receipt_root().call()
+    assert log['deposit_root'] == modified_registration_contract.functions.get_deposit_root().call()
     assert int.from_bytes(log['time'], byteorder='big') == timestamp_day_boundary
 
     # Make 1 deposit with value MAX_DEPOSIT and check that ChainStart event is not triggered