Preserve duplicate leaves in Merkle proofs

shared/src/main/scala/scorex/crypto/authds/merkle/MerkleTree.scala

@@ -10,13 +10,34 @@ import scala.collection.mutable
 case class MerkleTree[D <: Digest](topNode: Node[D],
                                    elementsHashIndex: Map[mutable.WrappedArray.ofByte, Int]) {
 
+  private lazy val leafHashes: Seq[D] = {
+    @tailrec
+    def loop(nodes: List[Node[D]], hashes: List[D]): List[D] = nodes match {
+      case Nil => hashes.reverse
+      case (node: InternalNode[D]) :: rest =>
+        loop(node.left :: node.right :: rest, hashes)
+      case (node: Leaf[D]) :: rest =>
+        loop(rest, node.hash :: hashes)
+      case _ :: rest =>
+        loop(rest, hashes)
+    }
+
+    loop(topNode :: Nil, Nil)
+  }
+
   lazy val rootHash: D = topNode.hash
-  lazy val length: Int = elementsHashIndex.size
+  // Equal leaf hashes share an element lookup entry but remain separate tree positions.
+  lazy val length: Int = leafHashes.length
 
   def proofByElement(element: Leaf[D]): Option[MerkleProof[D]] = proofByElementHash(element.hash)
 
   def proofByElementHash(hash: D): Option[MerkleProof[D]] = {
-    elementsHashIndex.get(new mutable.WrappedArray.ofByte(hash)).flatMap(i => proofByIndex(i))
+    elementsHashIndex
+      .get(new mutable.WrappedArray.ofByte(hash))
+      .flatMap(proofByIndex)
+      .filter { proof =>
+        proof.hf.prefixedHash(MerkleTree.LeafPrefix, proof.leafData).sameElements(hash)
+      }
   }
 
   def proofByIndex(index: Int): Option[MerkleProof[D]] = if (index >= 0 && index < length) {
@@ -96,7 +117,7 @@ case class MerkleTree[D <: Digest](topNode: Node[D],
     }
 
     if (indices.nonEmpty && indices.forall(index => index >= 0 && index < length)) {
-      val hashes: Seq[Digest] = elementsHashIndex.toSeq.sortBy(_._2).map(_._1.toArray.asInstanceOf[Digest])
+      val hashes: Seq[Digest] = leafHashes
       val normalized_indices = indices.distinct.sorted
       val multiproof = loop(normalized_indices, hashes)
       Some(BatchMerkleProof(normalized_indices zip (normalized_indices map hashes.apply), multiproof))

shared/src/test/scala/scorex/crypto/authds/merkle/MerkleTreeSpecification.scala

@@ -7,13 +7,22 @@ import scorex.crypto.TestingCommons
 import scorex.crypto.authds.LeafData
 import scorex.crypto.hash.{Blake2b256, Digest32}
 
+import scala.collection.mutable
 import scala.util.Random
 
 class MerkleTreeSpecification extends AnyPropSpec with ScalaCheckDrivenPropertyChecks with Matchers with TestingCommons {
   implicit val hf: Blake2b256.type = Blake2b256
 
   private val LeafSize = 32
 
+  private def payloadsWithDuplicateLeaves: Seq[Seq[LeafData]] =
+    Seq(
+      Seq(1, 1),
+      Seq(1, 2, 1),
+      Seq(1, 2, 1, 3),
+      Seq(1, 2, 3, 4, 1)
+    ).map(_.map(value => LeafData @@ Array.fill[Byte](LeafSize)(value.toByte)))
+
   property("Proof generation by element") {
     forAll(smallInt) { (N: Int) =>
       whenever(N > 0) {
@@ -29,6 +38,18 @@ class MerkleTreeSpecification extends AnyPropSpec with ScalaCheckDrivenPropertyC
     }
   }
 
+  property("Proof by element hash rejects an inconsistent element index") {
+    val first = LeafData @@ Array.fill[Byte](LeafSize)(1: Byte)
+    val second = LeafData @@ Array.fill[Byte](LeafSize)(2: Byte)
+    val firstLeaf = Leaf(first)
+    val tree = MerkleTree(
+      InternalNode(firstLeaf, Leaf(second)),
+      Map(new mutable.WrappedArray.ofByte(firstLeaf.hash) -> 1)
+    )
+
+    tree.proofByElementHash(firstLeaf.hash) shouldBe None
+  }
+
   property("Proof generation by index") {
     forAll(smallInt) { (N: Int) =>
       whenever(N > 0) {
@@ -48,6 +69,40 @@ class MerkleTreeSpecification extends AnyPropSpec with ScalaCheckDrivenPropertyC
     }
   }
 
+  property("Proof by index remains bound to the requested position with duplicate leaves") {
+    val first = LeafData @@ Array.fill[Byte](LeafSize)(1: Byte)
+    val second = LeafData @@ Array.fill[Byte](LeafSize)(2: Byte)
+    val tree = MerkleTree(Seq(first, second, first))
+    val proof = tree.proofByIndex(0).get
+
+    proof.valid(tree.rootHash) shouldBe true
+    proof.leafData.sameElements(first) shouldBe true
+  }
+
+  property("Proof generation with duplicate leaves") {
+    payloadsWithDuplicateLeaves.foreach { data =>
+      val tree = MerkleTree(data)
+
+      withClue(s"leaf count ${data.length}: ") {
+        tree.length shouldBe data.length
+        data.indices.foreach { index =>
+          val proof = tree.proofByIndex(index).get
+          proof.leafData.sameElements(data(index)) shouldBe true
+          proof.valid(tree.rootHash) shouldBe true
+        }
+
+        Seq(data.head, data(1)).foreach { leafData =>
+          val leaf = Leaf(leafData)
+          Seq(tree.proofByElement(leaf), tree.proofByElementHash(leaf.hash)).foreach { proofOpt =>
+            val proof = proofOpt.get
+            proof.leafData.sameElements(leafData) shouldBe true
+            proof.valid(tree.rootHash) shouldBe true
+          }
+        }
+      }
+    }
+  }
+
   property("Batch proof generation by indices") {
     val r = new Random()
     forAll(smallInt) { (N: Int) =>
@@ -69,6 +124,24 @@ class MerkleTreeSpecification extends AnyPropSpec with ScalaCheckDrivenPropertyC
     tree.proofByIndices(Seq(2,2,2,3,6,6,8,9,9)).get.valid(tree.rootHash) shouldBe true
   }
 
+  property("Batch proof generation with duplicate leaves") {
+    payloadsWithDuplicateLeaves.foreach { data =>
+      val tree = MerkleTree(data)
+      val duplicateIndex = data.lastIndexWhere(_.sameElements(data.head))
+      val requestedIndices = Seq(duplicateIndex, 0, duplicateIndex)
+      val expectedIndices = requestedIndices.distinct.sorted
+      val proof = tree.proofByIndices(requestedIndices).get
+
+      withClue(s"leaf count ${data.length}: ") {
+        proof.indices.map(_._1) shouldEqual expectedIndices
+        proof.indices.zip(expectedIndices).foreach { case ((_, hash), index) =>
+          hash.sameElements(Leaf(data(index)).hash) shouldBe true
+        }
+        proof.valid(tree.rootHash) shouldBe true
+      }
+    }
+  }
+
   property("Batch proof generation by negative indices") {
     val d = (0 until 5).map(_ => LeafData @@ scorex.utils.Random.randomBytes(LeafSize))
     val tree = MerkleTree(d)