From cb57c17eb6212624704443cc918bcb173642a061 Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@z.cash>
Date: Fri, 3 Apr 2020 23:10:38 +1300
Subject: [PATCH] test: Verify ZIP 221 logic against reference implementation

---
 qa/pull-tester/rpc-tests.sh              |   1 +
 qa/rpc-tests/feature_zip221.py           | 133 +++++++++++++++++
 qa/rpc-tests/test_framework/flyclient.py | 177 +++++++++++++++++++++++
 qa/rpc-tests/test_framework/mininode.py  |  15 ++
 4 files changed, 326 insertions(+)
 create mode 100755 qa/rpc-tests/feature_zip221.py
 create mode 100644 qa/rpc-tests/test_framework/flyclient.py

diff --git a/qa/pull-tester/rpc-tests.sh b/qa/pull-tester/rpc-tests.sh
index 6fb3c0f3c..66115b51d 100755
--- a/qa/pull-tester/rpc-tests.sh
+++ b/qa/pull-tester/rpc-tests.sh
@@ -83,6 +83,7 @@ testScripts=(
     'turnstile.py'
     'mining_shielded_coinbase.py'
     'framework.py'
+    'feature_zip221.py'
 );
 testScriptsExt=(
     'getblocktemplate_longpoll.py'
diff --git a/qa/rpc-tests/feature_zip221.py b/qa/rpc-tests/feature_zip221.py
new file mode 100755
index 000000000..f6a9d65e8
--- /dev/null
+++ b/qa/rpc-tests/feature_zip221.py
@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+# Copyright (c) 2020 The Zcash developers
+# Distributed under the MIT software license, see the accompanying
+# file COPYING or https://www.opensource.org/licenses/mit-license.php .
+
+
+from test_framework.flyclient import (ZcashMMRNode, append, delete, make_root_commitment)
+from test_framework.mininode import (HEARTWOOD_BRANCH_ID, CBlockHeader)
+from test_framework.test_framework import BitcoinTestFramework
+from test_framework.util import (
+    assert_equal,
+    bytes_to_hex_str,
+    hex_str_to_bytes,
+    initialize_chain_clean,
+    start_nodes,
+)
+
+from io import BytesIO
+
+NULL_FIELD = "0000000000000000000000000000000000000000000000000000000000000000"
+CHAIN_HISTORY_ROOT_VERSION = 2010200
+
+# Verify block header field 'hashLightClientRoot' is set correctly for Heartwood blocks.
+class Zip221Test(BitcoinTestFramework):
+
+    def setup_chain(self):
+        print("Initializing test directory "+self.options.tmpdir)
+        initialize_chain_clean(self.options.tmpdir, 4)
+
+    def setup_nodes(self):
+        return start_nodes(4, self.options.tmpdir, extra_args=[[
+            '-nuparams=2bb40e60:1', # Blossom
+            '-nuparams=f5b9230b:10', # Heartwood
+            '-nurejectoldversions=false',
+        ]] * 4)
+
+    def node_for_block(self, height):
+        block_header = CBlockHeader()
+        block_header.deserialize(BytesIO(hex_str_to_bytes(
+            self.nodes[0].getblock(str(height), 0))))
+        sapling_root = hex_str_to_bytes(
+            self.nodes[0].getblock(str(height))["finalsaplingroot"])[::-1]
+        return ZcashMMRNode.from_block(
+            block_header, height, sapling_root, 0, HEARTWOOD_BRANCH_ID)
+
+    def run_test(self):
+        self.nodes[0].generate(10)
+        self.sync_all()
+
+        # Verify all blocks up to and including Heartwood activation set
+        # hashChainHistoryRoot to null.
+        print("Verifying blocks up to and including Heartwood activation")
+        blockcount = self.nodes[0].getblockcount()
+        for height in range(0, blockcount + 1):
+            blk = self.nodes[0].getblock(str(height))
+            assert_equal(blk["chainhistoryroot"], NULL_FIELD)
+
+
+        # Create the initial history tree, containing a single node.
+        root = self.node_for_block(10)
+
+        # Generate the first block that contains a non-null chain history root.
+        print("Verifying first non-null chain history root")
+        self.nodes[0].generate(1)
+        self.sync_all()
+
+        # Verify that hashChainHistoryRoot is set correctly.
+        assert_equal(
+            self.nodes[0].getblock('11')["chainhistoryroot"],
+            bytes_to_hex_str(make_root_commitment(root)[::-1]))
+
+        # Generate 9 more blocks on node 0, and verify their chain history roots.
+        print("Mining 9 blocks on node 0")
+        self.nodes[0].generate(9)
+        self.sync_all()
+
+        print("Verifying node 0's chain history")
+        for height in range(12, 21):
+            leaf = self.node_for_block(height - 1)
+            root = append(root, leaf)
+            assert_equal(
+                self.nodes[0].getblock(str(height))["chainhistoryroot"],
+                bytes_to_hex_str(make_root_commitment(root)[::-1]))
+
+        # The rest of the test only applies to Heartwood-aware node versions.
+        # Earlier versions won't serialize chain history roots in the block
+        # index, and splitting the network below requires restarting the nodes.
+        if self.nodes[0].getnetworkinfo()["version"] < CHAIN_HISTORY_ROOT_VERSION:
+            print("Node's block index is not Heartwood-aware, skipping reorg test")
+            return
+
+        # Split the network so we can test the effect of a reorg.
+        print("Splitting the network")
+        self.split_network()
+
+        # Generate 10 more blocks on node 0, and verify their chain history roots.
+        print("Mining 10 more blocks on node 0")
+        self.nodes[0].generate(10)
+        self.sync_all()
+
+        print("Verifying node 0's chain history")
+        for height in range(21, 31):
+            leaf = self.node_for_block(height - 1)
+            root = append(root, leaf)
+            assert_equal(
+                self.nodes[0].getblock(str(height))["chainhistoryroot"],
+                bytes_to_hex_str(make_root_commitment(root)[::-1]))
+
+        # Generate 11 blocks on node 2.
+        print("Mining alternate chain on node 2")
+        self.nodes[2].generate(11)
+        self.sync_all()
+
+        # Reconnect the nodes; node 0 will re-org to node 2's chain.
+        print("Re-joining the network so that node 0 reorgs")
+        self.join_network()
+
+        # Verify that node 0's chain history was correctly updated.
+        print("Deleting orphaned blocks from the expected chain history")
+        for _ in range(10):
+            root = delete(root)
+
+        print("Verifying that node 0 is now on node 1's chain history")
+        for height in range(21, 32):
+            leaf = self.node_for_block(height - 1)
+            root = append(root, leaf)
+            assert_equal(
+                self.nodes[2].getblock(str(height))["chainhistoryroot"],
+                bytes_to_hex_str(make_root_commitment(root)[::-1]))
+
+
+if __name__ == '__main__':
+    Zip221Test().main()
diff --git a/qa/rpc-tests/test_framework/flyclient.py b/qa/rpc-tests/test_framework/flyclient.py
new file mode 100644
index 000000000..39cfa4627
--- /dev/null
+++ b/qa/rpc-tests/test_framework/flyclient.py
@@ -0,0 +1,177 @@
+from pyblake2 import blake2b
+import struct
+from typing import (List, Optional)
+
+from .mininode import (CBlockHeader, block_work_from_compact, ser_compactsize, ser_uint256)
+
+def H(msg: bytes, consensusBranchId: int) -> bytes:
+    digest = blake2b(
+        digest_size=32,
+        person=b'ZcashHistory' + struct.pack("<I", consensusBranchId))
+    digest.update(msg)
+    return digest.digest()
+
+class ZcashMMRNode():
+    # leaf nodes have no children
+    left_child: Optional['ZcashMMRNode']
+    right_child: Optional['ZcashMMRNode']
+
+    # commitments
+    hashSubtreeCommitment: bytes
+    nEarliestTimestamp: int
+    nLatestTimestamp: int
+    nEarliestTargetBits: int
+    nLatestTargetBits: int
+    hashEarliestSaplingRoot: bytes # left child's sapling root
+    hashLatestSaplingRoot: bytes # right child's sapling root
+    nSubTreeTotalWork: int # total difficulty accumulated within each subtree
+    nEarliestHeight: int
+    nLatestHeight: int
+    nSaplingTxCount: int # number of Sapling transactions in block
+
+    consensusBranchId: bytes
+
+    @classmethod
+    def from_block(Z, block: CBlockHeader, height, sapling_root, sapling_tx_count, consensusBranchId) -> 'ZcashMMRNode':
+        '''Create a leaf node from a block'''
+        node = Z()
+        node.left_child = None
+        node.right_child = None
+        node.hashSubtreeCommitment = ser_uint256(block.rehash())
+        node.nEarliestTimestamp = block.nTime
+        node.nLatestTimestamp = block.nTime
+        node.nEarliestTargetBits = block.nBits
+        node.nLatestTargetBits = block.nBits
+        node.hashEarliestSaplingRoot = sapling_root
+        node.hashLatestSaplingRoot = sapling_root
+        node.nSubTreeTotalWork = block_work_from_compact(block.nBits)
+        node.nEarliestHeight = height
+        node.nLatestHeight = height
+        node.nSaplingTxCount = sapling_tx_count
+        node.consensusBranchId = consensusBranchId
+        return node
+
+    def serialize(self) -> bytes:
+        '''serializes a node'''
+        buf = b''
+        buf += self.hashSubtreeCommitment
+        buf += struct.pack("<I", self.nEarliestTimestamp)
+        buf += struct.pack("<I", self.nLatestTimestamp)
+        buf += struct.pack("<I", self.nEarliestTargetBits)
+        buf += struct.pack("<I", self.nLatestTargetBits)
+        buf += self.hashEarliestSaplingRoot
+        buf += self.hashLatestSaplingRoot
+        buf += ser_uint256(self.nSubTreeTotalWork)
+        buf += ser_compactsize(self.nEarliestHeight)
+        buf += ser_compactsize(self.nLatestHeight)
+        buf += ser_compactsize(self.nSaplingTxCount)
+        return buf
+
+def make_parent(
+        left_child: ZcashMMRNode,
+        right_child: ZcashMMRNode) -> ZcashMMRNode:
+    parent = ZcashMMRNode()
+    parent.left_child = left_child
+    parent.right_child = right_child
+    parent.hashSubtreeCommitment = H(
+        left_child.serialize() + right_child.serialize(),
+        left_child.consensusBranchId,
+    )
+    parent.nEarliestTimestamp = left_child.nEarliestTimestamp
+    parent.nLatestTimestamp = right_child.nLatestTimestamp
+    parent.nEarliestTargetBits = left_child.nEarliestTargetBits
+    parent.nLatestTargetBits = right_child.nLatestTargetBits
+    parent.hashEarliestSaplingRoot = left_child.hashEarliestSaplingRoot
+    parent.hashLatestSaplingRoot = right_child.hashLatestSaplingRoot
+    parent.nSubTreeTotalWork = left_child.nSubTreeTotalWork + right_child.nSubTreeTotalWork
+    parent.nEarliestHeight = left_child.nEarliestHeight
+    parent.nLatestHeight = right_child.nLatestHeight
+    parent.nSaplingTxCount = left_child.nSaplingTxCount + right_child.nSaplingTxCount
+    parent.consensusBranchId = left_child.consensusBranchId
+    return parent
+
+def make_root_commitment(root: ZcashMMRNode) -> bytes:
+    '''Makes the root commitment for a blockheader'''
+    return H(root.serialize(), root.consensusBranchId)
+
+def get_peaks(node: ZcashMMRNode) -> List[ZcashMMRNode]:
+    peaks: List[ZcashMMRNode] = []
+
+    # Get number of leaves.
+    leaves = node.nLatestHeight - (node.nEarliestHeight - 1)
+    assert(leaves > 0)
+
+    # Check if the number of leaves is a power of two.
+    if (leaves & (leaves - 1)) == 0:
+        # Tree is full, hence a single peak. This also covers the
+        # case of a single isolated leaf.
+        peaks.append(node)
+    else:
+        # If the number of leaves is not a power of two, then this
+        # node must be internal, and cannot be a peak.
+        peaks.extend(get_peaks(node.left_child))
+        peaks.extend(get_peaks(node.right_child))
+
+    return peaks
+
+
+def bag_peaks(peaks: List[ZcashMMRNode]) -> ZcashMMRNode:
+    '''
+    "Bag" a list of peaks, and return the final root
+    '''
+    root = peaks[0]
+    for i in range(1, len(peaks)):
+        root = make_parent(root, peaks[i])
+    return root
+
+
+def append(root: ZcashMMRNode, leaf: ZcashMMRNode) -> ZcashMMRNode:
+    '''Append a leaf to an existing tree, return the new tree root'''
+    # recursively find a list of peaks in the current tree
+    peaks: List[ZcashMMRNode] = get_peaks(root)
+    merged: List[ZcashMMRNode] = []
+
+    # Merge peaks from right to left.
+    # This will produce a list of peaks in reverse order
+    current = leaf
+    for peak in peaks[::-1]:
+        current_leaves = current.nLatestHeight - (current.nEarliestHeight - 1)
+        peak_leaves = peak.nLatestHeight - (peak.nEarliestHeight - 1)
+
+        if current_leaves == peak_leaves:
+            current = make_parent(peak, current)
+        else:
+            merged.append(current)
+            current = peak
+    merged.append(current)
+
+    # finally, bag the merged peaks
+    return bag_peaks(merged[::-1])
+
+def delete(root: ZcashMMRNode) -> ZcashMMRNode:
+    '''
+    Delete the rightmost leaf node from an existing MMR
+    Return the new tree root
+    '''
+
+    n_leaves = root.nLatestHeight - (root.nEarliestHeight - 1)
+    # if there were an odd number of leaves,
+    # simply replace root with left_child
+    if n_leaves & 1:
+        return root.left_child
+
+    # otherwise, we need to re-bag the peaks.
+    else:
+        # first peak
+        peaks = [root.left_child]
+
+        # we do this traversing the right (unbalanced) side of the tree
+        # we keep the left side (balanced subtree or leaf) of each subtree
+        # until we reach a leaf
+        subtree_root = root.right_child
+        while subtree_root.left_child:
+            peaks.append(subtree_root.left_child)
+            subtree_root = subtree_root.right_child
+
+    new_root = bag_peaks(peaks)
+    return new_root
diff --git a/qa/rpc-tests/test_framework/mininode.py b/qa/rpc-tests/test_framework/mininode.py
index 1e86e8bc6..359553def 100755
--- a/qa/rpc-tests/test_framework/mininode.py
+++ b/qa/rpc-tests/test_framework/mininode.py
@@ -57,6 +57,7 @@ SPROUT_BRANCH_ID = 0x00000000
 OVERWINTER_BRANCH_ID = 0x5BA81B19
 SAPLING_BRANCH_ID = 0x76B809BB
 BLOSSOM_BRANCH_ID = 0x2BB40E60
+HEARTWOOD_BRANCH_ID = 0xF5B9230B
 
 MAX_INV_SZ = 50000
 
@@ -85,6 +86,15 @@ def hash256(s):
     return sha256(sha256(s))
 
 
+def ser_compactsize(n):
+    if n < 253:
+        return struct.pack("B", n)
+    elif n < 0x10000:
+        return struct.pack("<BH", 253, n)
+    elif n < 0x100000000:
+        return struct.pack("<BI", 254, n)
+    return struct.pack("<BQ", 255, n)
+
 def deser_string(f):
     nit = struct.unpack("<B", f.read(1))[0]
     if nit == 253:
@@ -134,6 +144,11 @@ def uint256_from_compact(c):
     return v
 
 
+def block_work_from_compact(c):
+    target = uint256_from_compact(c)
+    return 2**256 // (target + 1)
+
+
 def deser_vector(f, c):
     nit = struct.unpack("<B", f.read(1))[0]
     if nit == 253: