Merge pull request #109 from ccellado/fix_overflow

Stack overflow fix for BatchMerkeProof
input-output-hk · Feb 13, 2024 · 1d2db26 · 1d2db26
2 parents d925f93 + 9eb6b85
commit 1d2db26
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/shared/src/main/scala/scorex/crypto/authds/merkle/BatchMerkleProof.scala b/shared/src/main/scala/scorex/crypto/authds/merkle/BatchMerkleProof.scala
@@ -1,6 +1,7 @@
 package scorex.crypto.authds.merkle
 
 import scorex.crypto.authds.Side
+import scorex.crypto.authds.merkle.MerkleProof.LeftSide
 import scorex.crypto.authds.merkle.MerkleTree.InternalNodePrefix
 import scorex.crypto.hash.{CryptographicHash, Digest}
 import scorex.util.ScorexEncoding
@@ -68,10 +69,10 @@ case class BatchMerkleProof[D <: Digest](indices: Seq[(Int, Digest)], proofs: Se
  } else {
 
  // hash the corresponding value inside E with the first hash inside M, taking note of the side
- if (m_new.head._2 == MerkleProof.LeftSide) {
- e_new = e_new :+ hf.prefixedHash(MerkleTree.InternalNodePrefix, m_new.head._1 ++ e.apply(i)._2)
+ if (m_new.head._2 == LeftSide) {
+ e_new = e_new :+ hf.prefixedHash(InternalNodePrefix, m_new.head._1 ++ e.apply(i)._2)
  } else {
- e_new = e_new :+ hf.prefixedHash(MerkleTree.InternalNodePrefix, e.apply(i)._2 ++ m_new.head._1)
+ e_new = e_new :+ hf.prefixedHash(InternalNodePrefix, e.apply(i)._2 ++ m_new.head._1)
  }
 
  // remove the used value from m
@@ -84,7 +85,7 @@ case class BatchMerkleProof[D <: Digest](indices: Seq[(Int, Digest)], proofs: Se
  a_new = b.distinct.map(_._1 / 2)
 
  // Repeat until the root of the tree is reached (M has no more elements)
- if (m_new.nonEmpty || e_new.size > 1) {
+ if ((m_new.nonEmpty || e_new.size > 1) && a_new.nonEmpty) {
  e_new = loop(a_new, a_new zip e_new, m_new)
  }
  e_new

diff --git a/shared/src/test/scala/scorex/crypto/authds/merkle/MerkleTreeSpecification.scala b/shared/src/test/scala/scorex/crypto/authds/merkle/MerkleTreeSpecification.scala
@@ -5,7 +5,7 @@ import org.scalatest.matchers.should.Matchers
 import org.scalatestplus.scalacheck.ScalaCheckDrivenPropertyChecks
 import scorex.crypto.TestingCommons
 import scorex.crypto.authds.LeafData
-import scorex.crypto.hash.Blake2b256
+import scorex.crypto.hash.{Blake2b256, Digest32}
 
 import scala.util.Random
 
@@ -87,6 +87,11 @@ class MerkleTreeSpecification extends AnyPropSpec with ScalaCheckDrivenPropertyC
  tree.proofByIndices(Seq.empty[Int]) shouldBe None
  }
 
+ property("Proof for empty node caused stack overflow") {
+ val batch = BatchMerkleProof(Seq(), Seq((Digest32 @@ Array.fill[Byte](32)(0),MerkleProof.LeftSide)))
+ batch.valid(Digest32 @@ Array.fill[Byte](32)(0))
+ }
+
  property("Tree creation from 0 elements") {
  val tree = MerkleTree(Seq.empty)(hf)
  tree.rootHash shouldEqual Array.fill(hf.DigestSize)(0: Byte)