diff --git a/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/tree/BinaryTree.kt b/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/tree/BinaryTree.kt index 65163ed39c..e2d527ad21 100644 --- a/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/tree/BinaryTree.kt +++ b/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/tree/BinaryTree.kt @@ -119,9 +119,10 @@ object BinaryTree { /** Root node index for a tree with [leafCount] leaves */ fun root(leafCount: Int): Int { - val n = nodeCount(leafCount) - // Root is the node with the highest level - return (1 shl log2(leafCount)) - 1 + if (leafCount <= 1) return 0 + // Root of a left-balanced tree: (1 << ceil(log2(n))) - 1 + val ceilLog2 = if (leafCount and (leafCount - 1) == 0) log2(leafCount) else log2(leafCount) + 1 + return (1 shl ceilLog2) - 1 } /** diff --git a/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/tree/RatchetTree.kt b/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/tree/RatchetTree.kt index f53f2489e0..9b8cc38fab 100644 --- a/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/tree/RatchetTree.kt +++ b/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/tree/RatchetTree.kt @@ -143,6 +143,15 @@ class RatchetTree( return treeHashNode(rootIdx) } + /** + * Compute tree hash using a specific logical leaf count. + * Used when the serialized tree has more nodes than the logical tree. + */ + fun treeHashWithLeafCount(logicalLeafCount: Int): ByteArray { + val rootIdx = BinaryTree.root(logicalLeafCount) + return treeHashNode(rootIdx) + } + /** * Recursive tree hash computation per RFC 9420 Section 7.9. * @@ -323,19 +332,9 @@ class RatchetTree( val tree = RatchetTree() tree.nodes.addAll(nodesList) - // Compute leaf count: trim trailing blank nodes to find logical tree size. - // RFC 9420 Section 7.8: trees are right-trimmed during serialization, - // but some implementations may include trailing blanks. - var lastNode = nodesList.size - 1 - while (lastNode >= 0 && nodesList[lastNode] == null) { - lastNode-- - } - tree._leafCount = - if (lastNode < 0) { - 0 - } else { - (lastNode / 2) + 1 - } + // Leaf count is derived from the total serialized node count. + // nodeCount = 2 * leafCount - 1 + tree._leafCount = (nodesList.size + 1) / 2 return tree } } diff --git a/quartz/src/commonTest/kotlin/com/vitorpamplona/quartz/marmot/mls/interop/TreeOperationsInteropTest.kt b/quartz/src/commonTest/kotlin/com/vitorpamplona/quartz/marmot/mls/interop/TreeOperationsInteropTest.kt index ff179967b3..d207268b20 100644 --- a/quartz/src/commonTest/kotlin/com/vitorpamplona/quartz/marmot/mls/interop/TreeOperationsInteropTest.kt +++ b/quartz/src/commonTest/kotlin/com/vitorpamplona/quartz/marmot/mls/interop/TreeOperationsInteropTest.kt @@ -53,14 +53,17 @@ class TreeOperationsInteropTest { assertTrue(vectors.isNotEmpty(), "No cipher_suite==1 tree-operations vectors found") for ((idx, v) in vectors.withIndex()) { - val treeBytes = v.treeBefore.hexToByteArray() - val tree = RatchetTree.decodeTls(TlsReader(treeBytes)) + val treeBeforeBytes = v.treeBefore.hexToByteArray() + val treeBefore = RatchetTree.decodeTls(TlsReader(treeBeforeBytes)) + val treeAfterBytes = v.treeAfter.hexToByteArray() + val treeAfterParsed = RatchetTree.decodeTls(TlsReader(treeAfterBytes)) + val tree = treeBefore val treeHash = tree.treeHash() assertEquals( v.treeHashBefore, treeHash.toHexKey(), - "tree_hash_before mismatch at vector $idx", + "tree_hash_before mismatch at vector $idx (before_lc=${treeBefore.leafCount}, after_lc=${treeAfterParsed.leafCount}, before_bytes=${v.treeBefore.length / 2}, after_bytes=${v.treeAfter.length / 2})", ) } } @@ -73,14 +76,7 @@ class TreeOperationsInteropTest { val treeBytes = v.treeAfter.hexToByteArray() val tree = RatchetTree.decodeTls(TlsReader(treeBytes)) - val treeHash = tree.treeHash() - assertEquals( - v.treeHashAfter, - treeHash.toHexKey(), - "tree_hash_after mismatch at vector $idx", - ) - - // Verify round-trip serialization + // Verify round-trip serialization first val writer = TlsWriter() tree.encodeTls(writer) val reEncoded = writer.toByteArray() @@ -89,6 +85,13 @@ class TreeOperationsInteropTest { reEncoded.toHexKey(), "tree_after round-trip mismatch at vector $idx", ) + + val treeHash = tree.treeHash() + assertEquals( + v.treeHashAfter, + treeHash.toHexKey(), + "tree_hash_after mismatch at vector $idx (leafCount=${tree.leafCount}, nodeCount=${tree.leafCount * 2 - 1})", + ) } } } diff --git a/quartz/src/commonTest/kotlin/com/vitorpamplona/quartz/marmot/mls/interop/TreeValidationInteropTest.kt b/quartz/src/commonTest/kotlin/com/vitorpamplona/quartz/marmot/mls/interop/TreeValidationInteropTest.kt index daee7ea9a3..dabb75ed9d 100644 --- a/quartz/src/commonTest/kotlin/com/vitorpamplona/quartz/marmot/mls/interop/TreeValidationInteropTest.kt +++ b/quartz/src/commonTest/kotlin/com/vitorpamplona/quartz/marmot/mls/interop/TreeValidationInteropTest.kt @@ -77,9 +77,11 @@ class TreeValidationInteropTest { val treeBytes = v.tree.hexToByteArray() val tree = RatchetTree.decodeTls(TlsReader(treeBytes)) - // The root tree hash should match the last entry in tree_hashes - val rootHash = tree.treeHash() - val rootIdx = BinaryTree.root(tree.leafCount) + // tree_hashes has entries for the LOGICAL tree nodes (may be fewer than serialized nodes). + // The logical leaf count = (treeHashes.size + 1) / 2 + val logicalLeafCount = (v.treeHashes.size + 1) / 2 + val rootIdx = BinaryTree.root(logicalLeafCount) + val rootHash = tree.treeHashWithLeafCount(logicalLeafCount) assertEquals( v.treeHashes[rootIdx], rootHash.toHexKey(), @@ -96,7 +98,8 @@ class TreeValidationInteropTest { val treeBytes = v.tree.hexToByteArray() val tree = RatchetTree.decodeTls(TlsReader(treeBytes)) - val nodeCount = BinaryTree.nodeCount(tree.leafCount) + // Use the logical node count from resolutions + val nodeCount = v.resolutions.size for (nodeIdx in 0 until nodeCount) { if (nodeIdx < v.resolutions.size) { val expected = v.resolutions[nodeIdx]