Check internal node keys during sliced tree recovery

shared/src/main/scala/scorex/crypto/authds/avltree/batch/serialization/BatchAVLProverManifest.scala

@@ -20,9 +20,14 @@ case class BatchAVLProverManifest[D <: Digest](root: ProverNodes[D], rootHeight:
   /**
     * Verify that manifest corresponds to expected digest and height provided by a trusted party
     * (for blockchain protocols, it can be digest and height included by a miner)
+    *
+    * Node labels do not commit to the keys of internal nodes, so matching the digest is not enough on its
+    * own. A manifest ends at nodes whose children are only labels, so the keys below it are not available
+    * yet and only the search-tree ordering of its own keys can be checked here; the complete check runs on
+    * the assembled tree in `BatchAVLProverSerializer.combine`.
     */
   def verify(expectedDigest: D, expectedHeight: Int): Boolean = {
-    id.sameElements(expectedDigest) && expectedHeight == rootHeight
+    id.sameElements(expectedDigest) && expectedHeight == rootHeight && NodeKeyChecks.keysAreOrdered(root)
   }
 
   /**

shared/src/main/scala/scorex/crypto/authds/avltree/batch/serialization/BatchAVLProverSerializer.scala

@@ -167,6 +167,13 @@ class BatchAVLProverSerializer[D <: Digest, HF <: CryptographicHash[D]]
     require(actualHeight == manifest.rootHeight,
       s"manifest height ${manifest.rootHeight} does not match actual tree height $actualHeight")
 
+    // Node labels commit to the balance and the children labels, not to the keys of internal nodes, so a
+    // manifest and its subtrees can carry keys their digests do not authenticate. Now that they are joined
+    // the tree is materialized, and every internal key can be checked against the one its leaves imply.
+    // Recovering a tree with a forged internal key fails here rather than producing a prover which then
+    // routes lookups down the wrong branch.
+    NodeKeyChecks.checkKeys(manifest.root).get
+
     new BatchAVLProver[D, HF](keyLength, valueLengthOpt, Some(manifest.root -> manifest.rootHeight)) {
       override val logger = serializer.logger
     }

shared/src/main/scala/scorex/crypto/authds/avltree/batch/serialization/BatchAVLProverSubtree.scala

@@ -17,9 +17,13 @@ case class BatchAVLProverSubtree[D <: Digest](subtreeTop: ProverNodes[D]) {
 
   /**
     * Verify that manifest corresponds to expected digest (e.g. got from a manifest)
+    *
+    * Node labels do not commit to the keys of internal nodes, so matching the digest is not enough on its
+    * own: the keys are checked against the ones the leaves imply. A subtree is fully materialized, so that
+    * check is complete here.
     */
   def verify(expectedDigest: D): Boolean = {
-    subtreeTop.label.sameElements(expectedDigest)
+    subtreeTop.label.sameElements(expectedDigest) && NodeKeyChecks.keysAreConsistent(subtreeTop)
   }
 
   /**

shared/src/main/scala/scorex/crypto/authds/avltree/batch/serialization/NodeKeyChecks.scala

@@ -0,0 +1,93 @@
+package scorex.crypto.authds.avltree.batch.serialization
+
+import scorex.crypto.authds.ADKey
+import scorex.crypto.authds.avltree.batch.{InternalProverNode, ProverLeaf, ProverNodes}
+import scorex.crypto.hash.Digest
+import scorex.util.encode.Base16
+import scorex.utils.ByteArray
+
+import scala.annotation.tailrec
+import scala.collection.mutable
+import scala.util.Try
+
+/**
+  * Keys of internal nodes are not committed to by node labels: `InternalNode.computeLabel` hashes the
+  * balance and the two children labels, and nothing else. A manifest or a subtree therefore carries keys
+  * which its own digest does not authenticate, so a verifier checking labels alone takes them on trust.
+  *
+  * They do not have to be trusted, because they are fully determined by the leaves: the key of an internal
+  * node is the smallest key of its right subtree. That is the invariant `BatchAVLProver.checkTree` enforces
+  * ("min of right subtree doesn't match"), and it lets a received subtree be checked against itself.
+  */
+private[serialization] object NodeKeyChecks {
+
+  /**
+    * Smallest key of the subtree rooted at `node`, checking on the way up that every internal node carries
+    * the key its right subtree implies.
+    *
+    * Requires the subtree to be materialized: a [[ProxyInternalNode]] whose children are still labels
+    * cannot be descended into, and is rejected rather than skipped.
+    */
+  private def checkedMinKey[D <: Digest](node: ProverNodes[D]): ADKey = node match {
+    case l: ProverLeaf[D] =>
+      l.key
+    case n: ProxyInternalNode[D] if n.isEmpty =>
+      throw new IllegalArgumentException(
+        s"Node ${Base16.encode(n.label)} still has unresolved children, its keys cannot be checked")
+    case n: InternalProverNode[D] =>
+      val minLeft = checkedMinKey(n.left)
+      val minRight = checkedMinKey(n.right)
+      if (ByteArray.compare(n.key, minRight) != 0) {
+        throw new IllegalArgumentException(
+          s"Key ${Base16.encode(n.key)} of internal node ${Base16.encode(n.label)} does not match " +
+            s"the smallest key of its right subtree, ${Base16.encode(minRight)}")
+      }
+      minLeft
+  }
+
+  /**
+    * Whether every internal node of a materialized subtree carries the key its leaves imply.
+    *
+    * Complete: a subtree passing this check cannot contain a forged internal key.
+    */
+  def keysAreConsistent[D <: Digest](node: ProverNodes[D]): Boolean = checkKeys(node).isSuccess
+
+  /** As [[keysAreConsistent]], but keeps the reason so callers can report it */
+  def checkKeys[D <: Digest](node: ProverNodes[D]): Try[Unit] = Try(checkedMinKey(node)).map(_ => ())
+
+  /**
+    * Whether the internal node keys of a manifest are strictly increasing in in-order traversal.
+    *
+    * A manifest ends at [[ProxyInternalNode]]s whose children are only labels, so the keys below them are
+    * not available and [[keysAreConsistent]] cannot be used. What can still be checked is the search-tree
+    * ordering: with every internal key being the smallest key of its right subtree, an in-order walk of the
+    * internal nodes yields strictly increasing keys, the terminal proxy nodes included. This is the check
+    * available "assuming bottom keys of subtrees are correct"; the complete one runs once the subtrees are
+    * attached, in `BatchAVLProverSerializer.combine`.
+    */
+  def keysAreOrdered[D <: Digest](root: ProverNodes[D]): Boolean = {
+    val keys = mutable.Buffer.empty[ADKey]
+
+    def inOrder(node: ProverNodes[D]): Unit = node match {
+      case _: ProverLeaf[D] =>
+      case n: ProxyInternalNode[D] if n.isEmpty =>
+        keys += n.key
+      case n: InternalProverNode[D] =>
+        inOrder(n.left)
+        keys += n.key
+        inOrder(n.right)
+    }
+
+    inOrder(root)
+
+    @tailrec
+    def strictlyIncreasing(i: Int): Boolean = {
+      if (i >= keys.length) true
+      else if (ByteArray.compare(keys(i - 1), keys(i)) >= 0) false
+      else strictlyIncreasing(i + 1)
+    }
+
+    strictlyIncreasing(1)
+  }
+
+}

shared/src/test/scala/scorex/crypto/authds/avltree/batch/serialization/AVLBatchSerializationSpecification.scala

@@ -91,7 +91,10 @@ class AVLBatchSerializationSpecification extends AnyPropSpec with ScalaCheckDriv
         ADKey @@ keyBytes(2 * i + 1, KL)
       )
       current = new InternalProverNode[D](
-        ADKey @@ keyBytes(2 * i + 1, KL),
+        // The key of an internal node is the smallest key of its right subtree - the invariant
+        // BatchAVLProver.checkTree enforces. Here the right subtree is the chain built at i + 1,
+        // whose leftmost leaf carries key 2 * (i + 1).
+        ADKey @@ keyBytes(2 * i + 2, KL),
         left,
         current,
         Balance @@ 0.toByte
@@ -488,6 +491,74 @@ class AVLBatchSerializationSpecification extends AnyPropSpec with ScalaCheckDriv
     }
   }
 
+  private val forgedKey: ADKey = ADKey @@ Blake2b256("forged").take(KL)
+
+  /** Same children and balance, different key - and therefore the same label */
+  private def withForgedKey(n: InternalProverNode[D], key: ADKey = forgedKey): InternalProverNode[D] =
+    new InternalProverNode[D](key, n.left, n.right, n.balance)
+
+  private def firstInternalSubtree(sliced: serializer.SlicedTree): BatchAVLProverSubtree[D] =
+    sliced._2.find(_.subtreeTop.isInstanceOf[InternalProverNode[D]]).get
+
+  property("a label does not commit to the key of an internal node") {
+    val tree = generateProver()
+    val subtree = firstInternalSubtree(slice(tree))
+    val top = subtree.subtreeTop.asInstanceOf[InternalProverNode[D]]
+
+    val forged = withForgedKey(top)
+
+    // this is the hole the checks below close: the key changed, the digest did not
+    forged.key.sameElements(top.key) shouldBe false
+    forged.label.sameElements(top.label) shouldBe true
+  }
+
+  property("subtree with a forged internal key does not verify") {
+    val tree = generateProver()
+    val subtree = firstInternalSubtree(slice(tree))
+    val top = subtree.subtreeTop.asInstanceOf[InternalProverNode[D]]
+
+    val forgedSubtree = BatchAVLProverSubtree[D](withForgedKey(top))
+
+    // it still passes the digest check it used to be accepted on ...
+    forgedSubtree.id.sameElements(subtree.id) shouldBe true
+    // ... and is rejected anyway
+    forgedSubtree.verify(subtree.id) shouldBe false
+
+    subtree.verify(subtree.id) shouldBe true
+  }
+
+  property("manifest with an internal key out of order does not verify") {
+    val tree = generateProver()
+    val manifest = slice(tree)._1
+    val root = manifest.root.asInstanceOf[InternalProverNode[D]]
+
+    // smaller than every key in the tree, so the in-order walk stops increasing at the root
+    val smallest = ADKey @@ Array.fill(KL)(0.toByte)
+    val forgedManifest = BatchAVLProverManifest[D](withForgedKey(root, smallest), manifest.rootHeight)
+
+    forgedManifest.id.sameElements(manifest.id) shouldBe true
+    forgedManifest.verify(manifest.id, manifest.rootHeight) shouldBe false
+
+    manifest.verify(manifest.id, manifest.rootHeight) shouldBe true
+  }
+
+  property("combining a tree with a forged internal key fails") {
+    val tree = generateProver()
+
+    // combine mutates the manifest it is given, so each case gets its own slice of the same tree
+    serializer.combine(slice(tree), tree.keyLength, tree.valueLengthOpt).isSuccess shouldBe true
+
+    val sliced = slice(tree)
+    val subtree = firstInternalSubtree(sliced)
+    val top = subtree.subtreeTop.asInstanceOf[InternalProverNode[D]]
+
+    val forgedSubtree = BatchAVLProverSubtree[D](withForgedKey(top))
+    val forgedSliced = (sliced._1, sliced._2.map(st => if (st.id.sameElements(subtree.id)) forgedSubtree else st))
+
+    // the label is unchanged, so the forged subtree is still found and attached where the real one belongs
+    serializer.combine(forgedSliced, tree.keyLength, tree.valueLengthOpt).isFailure shouldBe true
+  }
+
   def leftTree(n: ProverNodes[D]): Seq[ProverNodes[D]] = n match {
     case n: ProxyInternalNode[D] if n.isEmpty =>
       Seq(n)