Prevent second preimage attack (#16)

* feat: be consistent with solidity merkle tree library (prevent second preimage attack) * fix: apply code review changes Co-authored-by: pgebal <[email protected]>
yosriady · Jan 30, 2020 · 5dfaf8a · 5dfaf8a
1 parent e73c4b1
commit 5dfaf8a
Show file tree

Hide file tree

Showing 4 changed files with 109 additions and 147 deletions.
diff --git a/lib/merkle_tree.ex b/lib/merkle_tree.ex
@@ -7,23 +7,60 @@ defmodule MerkleTree do
 
       ## Usage Example
 
-      iex> MerkleTree.new ['a', 'b', 'c', 'd']
-      %MerkleTree{blocks: ['a', 'b', 'c', 'd'], hash_function: &MerkleTree.Crypto.sha256/1,
-            root: %MerkleTree.Node{children: [%MerkleTree.Node{children: [%MerkleTree.Node{children: [], height: 0,
-                 value: "ca978112ca1bbdcafac231b39a23dc4da786eff8147c4e72b9807785afee48bb"},
-                %MerkleTree.Node{children: [], height: 0, value: "3e23e8160039594a33894f6564e1b1348bbd7a0088d42c4acb73eeaed59c009d"}], height: 1,
-               value: "62af5c3cb8da3e4f25061e829ebeea5c7513c54949115b1acc225930a90154da"},
-              %MerkleTree.Node{children: [%MerkleTree.Node{children: [], height: 0,
-                 value: "2e7d2c03a9507ae265ecf5b5356885a53393a2029d241394997265a1a25aefc6"},
-                %MerkleTree.Node{children: [], height: 0, value: "18ac3e7343f016890c510e93f935261169d9e3f565436429830faf0934f4f8e4"}], height: 1,
-               value: "d3a0f1c792ccf7f1708d5422696263e35755a86917ea76ef9242bd4a8cf4891a"}], height: 2,
-             value: "58c89d709329eb37285837b042ab6ff72c7c8f74de0446b091b6a0131c102cfd"}}
+    iex> MerkleTree.new ["a", "b", "c", "d"]
+    %MerkleTree{
+      blocks: ["a", "b", "c", "d"],
+      hash_function: &MerkleTree.Crypto.sha256/1,
+      root: %MerkleTree.Node{
+        children: [
+          %MerkleTree.Node{
+            children: [
+              %MerkleTree.Node{
+                children: [],
+                height: 0,
+                value: "022a6979e6dab7aa5ae4c3e5e45f7e977112a7e63593820dbec1ec738a24f93c"
+              },
+              %MerkleTree.Node{
+                children: [],
+                height: 0,
+                value: "57eb35615d47f34ec714cacdf5fd74608a5e8e102724e80b24b287c0c27b6a31"
+              }
+            ],
+            height: 1,
+            value: "4c64254e6636add7f281ff49278beceb26378bd0021d1809974994e6e233ec35"
+          },
+          %MerkleTree.Node{
+            children: [
+              %MerkleTree.Node{
+                children: [],
+                height: 0,
+                value: "597fcb31282d34654c200d3418fca5705c648ebf326ec73d8ddef11841f876d8"
+              },
+              %MerkleTree.Node{
+                children: [],
+                height: 0,
+                value: "d070dc5b8da9aea7dc0f5ad4c29d89965200059c9a0ceca3abd5da2492dcb71d"
+              }
+            ],
+            height: 1,
+            value: "40e2511a6323177e537acb2e90886e0da1f84656fd6334b89f60d742a3967f09"
+          }
+        ],
+        height: 2,
+        value: "9dc1674ae1ee61c90ba50b6261e8f9a47f7ea07d92612158edfe3c2a37c6d74c"
+      }
+    }
   """
 
   defstruct [:blocks, :root, :hash_function]
 
   # Number of children per node. Configurable.
   @number_of_children 2
+  # values prepended to a leaf and node to differentiate between them when calculating values
+  # of parent in Merkle Tree, added to prevent a second preimage attack
+  # where a proof for node can be validated as a proof for leaf
+  @leaf_salt <<0>>
+  @node_salt <<1>>
 
   @type blocks :: [String.t(), ...]
   @type hash_function :: (String.t() -> String.t())
@@ -38,7 +75,6 @@ defmodule MerkleTree do
     Creates a new merkle tree, given a blocks and hash function or opts.
      available options:
       :hash_function - used hash in mercle tree default :sha256 from :cryto
-      :hash_leaves - flag says whether the leaves should be hashed, default true
       :height - allows to construct tree of provided height,
           empty leaves data will be taken from `:default_data_block` parameter
       :default_data_block - this data will be used to supply empty
@@ -70,10 +106,10 @@ defmodule MerkleTree do
   """
   @spec fast_root(blocks, Keyword.t()) :: MerkleTree.Node.hash()
   def fast_root(blocks, opts \\ []) do
-    {hash_function, height, hash_leaves?, default_data_block} = get_from_options(opts, blocks)
+    {hash_function, height, default_data_block} = get_from_options(opts, blocks)
 
-    default_leaf_value = if hash_leaves?, do: hash_function.(default_data_block), else: default_data_block
-    leaf_values = if hash_leaves?, do: Enum.map(blocks, hash_function), else: blocks
+    default_leaf_value = hash_function.(@leaf_salt <> default_data_block)
+    leaf_values = Enum.map(blocks, fn block -> hash_function.(@leaf_salt <> block) end)
 
     _fast_root(leaf_values, hash_function, 0, default_leaf_value, height)
   end
@@ -83,29 +119,28 @@ defmodule MerkleTree do
 
   defp _fast_root([root], _, final_height, _, final_height), do: root
 
-  defp _fast_root(nodes, hash_function, height, default_leaf, final_height) do
+  defp _fast_root(nodes, hash_function, height, default_node, final_height) do
     count = step = @number_of_children
-    leftover = List.duplicate(default_leaf, count - 1)
+    leftover = List.duplicate(default_node, count - 1)
     children_partitions = Enum.chunk_every(nodes, count, step, leftover)
     new_height = height + 1
 
     parents =
       Enum.map(children_partitions, fn partition ->
-        concatenated_values = partition |> Enum.join()
+        concatenated_values = [@node_salt | partition] |> Enum.join()
         hash_function.(concatenated_values)
       end)
 
-    new_default_leaf_value = hash_function.(default_leaf <> default_leaf)
+    new_default_node = hash_function.(@node_salt <> default_node <> default_node)
 
-    _fast_root(parents, hash_function, new_height, new_default_leaf_value, final_height)
+    _fast_root(parents, hash_function, new_height, new_default_node, final_height)
   end
 
   # takes care of the defaults etc
   defp get_from_options(opts, blocks) do
     {
       Keyword.get(opts, :hash_function, &MerkleTree.Crypto.sha256/1),
       Keyword.get(opts, :height, guess_height(Enum.count(blocks))),
-      Keyword.get(opts, :hash_leaves, true),
       Keyword.get(opts, :default_data_block, "")
     }
   end
@@ -122,10 +157,10 @@ defmodule MerkleTree do
     do: build(blocks, hash_function: hash_function)
 
   def build(blocks, opts) do
-    {hash_function, height, hash_leaves?, default_data_block} = get_from_options(opts, blocks)
+    {hash_function, height, default_data_block} = get_from_options(opts, blocks)
 
-    default_leaf_value = if hash_leaves?, do: hash_function.(default_data_block), else: default_data_block
-    leaf_values = if hash_leaves?, do: Enum.map(blocks, hash_function), else: blocks
+    default_leaf_value = hash_function.(@leaf_salt <> default_data_block)
+    leaf_values = Enum.map(blocks, fn block -> hash_function.(@leaf_salt <> block) end)
     default_leaf = %MerkleTree.Node{value: default_leaf_value, children: [], height: 0}
     leaves = Enum.map(leaf_values, &%MerkleTree.Node{value: &1, children: [], height: 0})
 
@@ -147,6 +182,7 @@ defmodule MerkleTree do
     parents =
       Enum.map(children_partitions, fn partition ->
         concatenated_values = partition |> Enum.map(& &1.value) |> Enum.join()
+        concatenated_values = @node_salt <> concatenated_values
 
         %MerkleTree.Node{
           value: hash_function.(concatenated_values),
@@ -155,7 +191,8 @@ defmodule MerkleTree do
         }
       end)
 
-    new_default_leaf_value = hash_function.(default_leaf.value <> default_leaf.value)
+    new_default_leaf_value =
+        hash_function.(@node_salt <>  default_leaf.value <> default_leaf.value)
 
     new_default_leaf = %MerkleTree.Node{
       value: new_default_leaf_value,
@@ -183,10 +220,10 @@ defmodule MerkleTree do
   end
 
   defp fill_blocks(blocks, _, _) when blocks != [] do
-    amout_elements = Enum.count(blocks)
-    required_leaves_count = :math.pow(2, _ceil(:math.log2(amout_elements)))
+    blocks_count = Enum.count(blocks)
+    required_leaves_count = :math.pow(2, _ceil(:math.log2(blocks_count)))
 
-    if required_leaves_count != amout_elements,
+    if required_leaves_count != blocks_count,
       do: raise(MerkleTree.ArgumentError),
       else: blocks
   end

diff --git a/lib/merkle_tree/proof.ex b/lib/merkle_tree/proof.ex
@@ -5,31 +5,28 @@ defmodule MerkleTree.Proof do
       ## Usage Example
       iex> proof = MerkleTree.new(~w/a b c d/) |>
       ...> MerkleTree.Proof.prove(1)
-      %MerkleTree.Proof{hash_function: &MerkleTree.Crypto.sha256/1,
-       hashes: ["d3a0f1c792ccf7f1708d5422696263e35755a86917ea76ef9242bd4a8cf4891a",
-        "ca978112ca1bbdcafac231b39a23dc4da786eff8147c4e72b9807785afee48bb"]}
-      iex> MerkleTree.Proof.proven?({"b", 1}, "58c89d709329eb37285837b042ab6ff72c7c8f74de0446b091b6a0131c102cfd", proof)
+      ["40e2511a6323177e537acb2e90886e0da1f84656fd6334b89f60d742a3967f09",
+        "022a6979e6dab7aa5ae4c3e5e45f7e977112a7e63593820dbec1ec738a24f93c"]
+      iex> MerkleTree.Proof.proven?({"b", 1}, "9dc1674ae1ee61c90ba50b6261e8f9a47f7ea07d92612158edfe3c2a37c6d74c", &MerkleTree.Crypto.sha256/1, proof)
       true
 
   """
+  @leaf_salt <<0>>
+  @node_salt <<1>>
+
   defstruct [:hashes, :hash_function]
 
-  @type t :: %MerkleTree.Proof{
-          hashes: [String.t(), ...],
-          # TODO: remove when deprecated MerkleTree.Proof.proven?/3 support ends
-          hash_function: MerkleTree.hash_function() | nil
-        }
+  @type proof_t() :: list(String.t())
 
   @doc """
   Generates proof for a block at a specific index
   """
-  @spec prove(MerkleTree.t() | MerkleTree.Node.t(), non_neg_integer) :: t
-  def prove(%MerkleTree{root: root} = tree, index),
-    # TODO: remove the struct update with hash function, when deprecated MerkleTree.Proof.proven?/3 support ends
-    do: %MerkleTree.Proof{prove(root, index) | hash_function: tree.hash_function}
+  @spec prove(MerkleTree.t() | MerkleTree.Node.t(), non_neg_integer) :: proof_t()
+  def prove(%MerkleTree{root: root}, index),
+    do: prove(root, index)
 
   def prove(%MerkleTree.Node{height: height} = root, index),
-    do: %MerkleTree.Proof{hashes: _prove(root, binarize(index, height))}
+    do: _prove(root, binarize(index, height))
 
   defp _prove(_, ""), do: []
 
@@ -51,30 +48,22 @@ defmodule MerkleTree.Proof do
   @doc """
   Verifies proof for a block at a specific index
   """
-  @spec proven?({String.t(), non_neg_integer}, String.t(), MerkleTree.hash_function(), t) :: boolean
-  def proven?({block, index}, root_hash, hash_function, %MerkleTree.Proof{hashes: proof}) do
-    height = length(proof)
-    root_hash == _hash_proof(block, binarize(index, height), proof, hash_function)
-  end
-
-  @doc false
-  @deprecated "Use proven?/4 instead"
-  # TODO: remove when deprecated MerkleTree.Proof.proven?/3 support ends
-  def proven?({block, index}, root_hash, %MerkleTree.Proof{hashes: proof, hash_function: hash_function}) do
+  @spec proven?({String.t(), non_neg_integer}, String.t(), MerkleTree.hash_function(), proof_t()) :: boolean
+  def proven?({block, index}, root_hash, hash_function, proof) do
     height = length(proof)
     root_hash == _hash_proof(block, binarize(index, height), proof, hash_function)
   end
 
   defp _hash_proof(block, "", [], hash_function) do
-    hash_function.(block)
+    hash_function.(@leaf_salt <> block)
   end
 
   defp _hash_proof(block, index_binary, [proof_head | proof_tail], hash_function) do
     {path_head, path_tail} = path_from_binary(index_binary)
 
     case path_head do
-      1 -> hash_function.(proof_head <> _hash_proof(block, path_tail, proof_tail, hash_function))
-      0 -> hash_function.(_hash_proof(block, path_tail, proof_tail, hash_function) <> proof_head)
+      1 -> hash_function.(@node_salt <> proof_head <> _hash_proof(block, path_tail, proof_tail, hash_function))
+      0 -> hash_function.(@node_salt <> _hash_proof(block, path_tail, proof_tail, hash_function) <> proof_head)
     end
   end
 

diff --git a/test/merkle_tree/proof_test.exs b/test/merkle_tree/proof_test.exs
@@ -4,6 +4,8 @@ defmodule MerkleTree.ProofTest do
 
   import MerkleTree.Proof
 
+  @node_salt <<1>>
+
   test "correct proofs" do
     blocks = ~w/a b c d e f g h/
     tree = MerkleTree.new(blocks)
@@ -20,105 +22,43 @@ defmodule MerkleTree.ProofTest do
            |> Enum.all?()
   end
 
-  # TODO: remove when deprecated MerkleTree.Proof.proven?/3 support ends
-  test "correct proofs with deprecated proven?/3" do
-    blocks = ~w/a b c d e f g h/
-    tree = MerkleTree.new(blocks)
-
-    proofs =
-      blocks
-      |> Enum.with_index()
-      |> Enum.map(fn {_, idx} -> prove(tree, idx) end)
-
-    assert blocks
-           |> Enum.with_index()
-           |> Enum.zip(proofs)
-           |> Enum.map(fn {x, proof} -> proven?(x, tree.root.value, proof) end)
-           |> Enum.all?()
-  end
-
   test "incorrect proof" do
     blocks = ~w/a b c d e f g h/
     tree = MerkleTree.new(blocks)
-    %MerkleTree.Proof{hashes: hashes} = proof = prove(tree, 5)
+    proof = prove(tree, 5)
 
     # test sanity
     assert proven?({"f", 5}, tree.root.value, tree.hash_function, proof)
 
     # bad index
     assert not proven?({"f", 6}, tree.root.value, tree.hash_function, proof)
 
-    incomplete_proof = %MerkleTree.Proof{
-      hashes: tl(hashes),
-      hash_function: tree.hash_function
-    }
-
+    incomplete_proof = tl(proof)
     assert not proven?({"f", 5}, tree.root.value, tree.hash_function, incomplete_proof)
 
     # different hash function
     assert not proven?({"f", 5}, tree.root.value, &MerkleTree.Crypto.hash(&1, :sha224), proof)
 
-    corrupted_proof = %MerkleTree.Proof{
-      hashes: [tree.hash_function.("z")] ++ tl(hashes),
-      hash_function: tree.hash_function
-    }
-
+    corrupted_proof = [tree.hash_function.("z")] ++ tl(proof)
     assert not proven?({"f", 5}, tree.root.value, tree.hash_function, corrupted_proof)
 
     # corrupted root hash
     assert not proven?({"f", 5}, tree.hash_function.("z"), tree.hash_function, proof)
 
     # fake proof rejected with proven?/4
-    fake_proof = %MerkleTree.Proof{
-      hashes: Enum.map(hashes, fn _ -> "" end),
-      hash_function: fn _ -> tree.root.value end
-    }
-
+    fake_proof = Enum.map(proof, fn _ -> "" end)
     assert not proven?({"z", 5}, tree.root.value, tree.hash_function, fake_proof)
   end
 
-  # TODO: remove when deprecated MerkleTree.Proof.proven?/3 support ends
-  test "incorrect proof with deprecated proven?/3" do
-    blocks = ~w/a b c d e f g h/
-    tree = MerkleTree.new(blocks)
-    %MerkleTree.Proof{hashes: hashes} = proof = prove(tree, 5)
-
-    # test sanity
-    assert proven?({"f", 5}, tree.root.value, proof)
-
-    # bad index
-    assert not proven?({"f", 6}, tree.root.value, proof)
+  test "can not prove that node is a leaf" do
+      blocks = ~w/a b c d e f g h/
+      tree = MerkleTree.new(blocks)
 
-    incomplete_proof = %MerkleTree.Proof{
-      hashes: tl(hashes),
-      hash_function: tree.hash_function
-    }
-
-    assert not proven?({"f", 5}, tree.root.value, incomplete_proof)
-
-    different_hash = %MerkleTree.Proof{
-      hashes: hashes,
-      hash_function: &MerkleTree.Crypto.hash(&1, :sha224)
-    }
-
-    assert not proven?({"f", 5}, tree.root.value, different_hash)
-
-    corrupted_proof = %MerkleTree.Proof{
-      hashes: [tree.hash_function.("z")] ++ tl(hashes),
-      hash_function: tree.hash_function
-    }
-
-    assert not proven?({"f", 5}, tree.root.value, corrupted_proof)
-
-    # corrupted root hash
-    assert not proven?({"f", 5}, tree.hash_function.("z"), proof)
+      proof = prove(tree, 0)
+      node_proof = tl(proof)
 
-    # fake proof accepted with deprecated proven?/3
-    fake_proof = %MerkleTree.Proof{
-      hashes: Enum.map(hashes, fn _ -> "" end),
-      hash_function: fn _ -> tree.root.value end
-    }
+      node = tree.hash_function.(@node_salt <> "a" <> "b")
 
-    assert proven?({"z", 5}, tree.root.value, fake_proof)
+      assert not proven?({node, 0}, tree.root.value, tree.hash_function, node_proof)
   end
 end