testing/mozharness/test/test_mozilla_merkle.py

   1 import codecs
   2 import hashlib
   3 import random
   4 import unittest
   5
   6 from mozharness.mozilla.merkle import InclusionProof, MerkleTree
   7
   8 decode_hex = codecs.getdecoder("hex_codec")
   9 encode_hex = codecs.getencoder("hex_codec")
  10
  11 # Pre-computed tree on 7 inputs
  12 #
  13 #         ______F_____
  14 #        /            \
  15 #     __D__           _E_
  16 #    /     \         /   \
  17 #   A       B       C     |
  18 #  / \     / \     / \    |
  19 # 0   1   2   3   4   5   6
  20 hash_fn = hashlib.sha256
  21
  22 data = [
  23     decode_hex("fbc459361fc111024c6d1fd83d23a9ff")[0],
  24     decode_hex("ae3a44925afec860451cd8658b3cadde")[0],
  25     decode_hex("418903fe6ef29fc8cab93d778a7b018b")[0],
  26     decode_hex("3d1c53c00b2e137af8c4c23a06388c6b")[0],
  27     decode_hex("e656ebd8e2758bc72599e5896be357be")[0],
  28     decode_hex("81aae91cf90be172eedd1c75c349bf9e")[0],
  29     decode_hex("00c262edf8b0bc345aca769e8733e25e")[0],
  30 ]
  31
  32 leaves = [
  33     decode_hex("5cb551f87797381a24a5359a986e2cef25b1f2113b387197fe48e8babc9ad5c7")[0],
  34     decode_hex("9899dc0be00306bda2a8e69cec32525ca6244f132479bcf840d8c1bc8bdfbff2")[0],
  35     decode_hex("fdd27d0393e32637b474efb9b3efad29568c3ec9b091fdda40fd57ec9196f06d")[0],
  36     decode_hex("c87292a6c8528c2a0679b6c1eefb47e4dbac7840d23645d5b7cb47cf1a8d365f")[0],
  37     decode_hex("2ff3bdac9bec3580b82da8a357746f15919414d9cbe517e2dd96910c9814c30c")[0],
  38     decode_hex("883e318240eccc0e2effafebdb0fd4fd26d0996da1b01439566cb9babef8725f")[0],
  39     decode_hex("bb13dfb7b202a95f241ea1715c8549dc048d9936ec747028002f7c795de72fcf")[0],
  40 ]
  41
  42 nodeA = decode_hex("06447a7baa079cb0b4b6119d0f575bec508915403fdc30923eba982b63759805")[
  43     0
  44 ]
  45 nodeB = decode_hex("3db98027c655ead4fe897bef3a4b361839a337941a9e624b475580c9d4e882ee")[
  46     0
  47 ]
  48 nodeC = decode_hex("17524f8b0169b2745c67846925d55449ae80a8022ef8189dcf4cbb0ec7fcc470")[
  49     0
  50 ]
  51 nodeD = decode_hex("380d0dc6fd7d4f37859a12dbfc7171b3cce29ab0688c6cffd2b15f3e0b21af49")[
  52     0
  53 ]
  54 nodeE = decode_hex("3a9c2886a5179a6e1948876034f99d52a8f393f47a09887adee6d1b4a5c5fbd6")[
  55     0
  56 ]
  57 nodeF = decode_hex("d1a0d3947db4ae8305f2ac32985957e02659b2ea3c10da52a48d2526e9af3bbc")[
  58     0
  59 ]
  60
  61 proofs = [
  62     [leaves[1], nodeB, nodeE],
  63     [leaves[0], nodeB, nodeE],
  64     [leaves[3], nodeA, nodeE],
  65     [leaves[2], nodeA, nodeE],
  66     [leaves[5], leaves[6], nodeD],
  67     [leaves[4], leaves[6], nodeD],
  68     [nodeC, nodeD],
  69 ]
  70
  71 known_proof5 = decode_hex(
  72     "020000"
  73     + "0000000000000007"
  74     + "0000000000000005"
  75     + "0063"
  76     + "20"
  77     + encode_hex(leaves[4])[0].decode()
  78     + "20"
  79     + encode_hex(leaves[6])[0].decode()
  80     + "20"
  81     + encode_hex(nodeD)[0].decode()
  82 )[0]
  83
  84
  85 class TestMerkleTree(unittest.TestCase):
  86     def testPreComputed(self):
  87         tree = MerkleTree(hash_fn, data)
  88         head = tree.head()
  89         self.assertEqual(head, nodeF)
  90
  91         for i in range(len(data)):
  92             proof = tree.inclusion_proof(i)
  93
  94             self.assertTrue(proof.verify(hash_fn, data[i], i, len(data), head))
  95             self.assertEqual(proof.leaf_index, i)
  96             self.assertEqual(proof.tree_size, tree.n)
  97             self.assertEqual(proof.path_elements, proofs[i])
  98
  99     def testInclusionProofEncodeDecode(self):
 100         tree = MerkleTree(hash_fn, data)
 101
 102         # Inclusion proof encode/decode round trip test
 103         proof5 = tree.inclusion_proof(5)
 104         serialized5 = proof5.to_rfc6962_bis()
 105         deserialized5 = InclusionProof.from_rfc6962_bis(serialized5)
 106         reserialized5 = deserialized5.to_rfc6962_bis()
 107         self.assertEqual(serialized5, reserialized5)
 108
 109         # Inclusion proof encode known answer test
 110         serialized5 = proof5.to_rfc6962_bis()
 111         self.assertEqual(serialized5, known_proof5)
 112
 113         # Inclusion proof decode known answer test
 114         known_deserialized5 = InclusionProof.from_rfc6962_bis(known_proof5)
 115         self.assertEqual(proof5.leaf_index, known_deserialized5.leaf_index)
 116         self.assertEqual(proof5.tree_size, known_deserialized5.tree_size)
 117         self.assertEqual(proof5.path_elements, known_deserialized5.path_elements)
 118
 119     def testLargeTree(self):
 120         TEST_SIZE = 5000
 121         ELEM_SIZE_BYTES = 16
 122         data = [
 123             bytearray(random.getrandbits(8) for _ in range(ELEM_SIZE_BYTES))
 124             for _ in range(TEST_SIZE)
 125         ]
 126         tree = MerkleTree(hash_fn, data)
 127         head = tree.head()
 128
 129         for i in range(len(data)):
 130             proof = tree.inclusion_proof(i)
 131
 132             self.assertTrue(proof.verify(hash_fn, data[i], i, len(data), head))
 133             self.assertEqual(proof.leaf_index, i)
 134             self.assertEqual(proof.tree_size, tree.n)