From f9ee4ea7af0b3cda68f821f0fee948a3d0f3d952 Mon Sep 17 00:00:00 2001 From: Tomasz Bekas Date: Sun, 27 Aug 2023 17:41:20 +0200 Subject: [PATCH] Coders for merkle trees --- codex/merkletree.nim | 4 ++ codex/merkletree/coders.nim | 46 ++++++++++++++++++++++ codex/merkletree/merkletree.nim | 48 +++++++++++++++++++++-- tests/codex/merkletree/testcoders.nim | 43 ++++++++++++++++++++ tests/codex/merkletree/testmerkletree.nim | 2 +- tests/codex/testmerkletree.nim | 1 + 6 files changed, 139 insertions(+), 5 deletions(-) create mode 100644 codex/merkletree.nim create mode 100644 codex/merkletree/coders.nim create mode 100644 tests/codex/merkletree/testcoders.nim diff --git a/codex/merkletree.nim b/codex/merkletree.nim new file mode 100644 index 00000000..366af992 --- /dev/null +++ b/codex/merkletree.nim @@ -0,0 +1,4 @@ +import ./merkletree/merkletree +import ./merkletree/coders + +export merkletree, coders diff --git a/codex/merkletree/coders.nim b/codex/merkletree/coders.nim new file mode 100644 index 00000000..b90575f8 --- /dev/null +++ b/codex/merkletree/coders.nim @@ -0,0 +1,46 @@ +## Nim-Codex +## Copyright (c) 2023 Status Research & Development GmbH +## Licensed under either of +## * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE)) +## * MIT license ([LICENSE-MIT](LICENSE-MIT)) +## at your option. +## This file may not be copied, modified, or distributed except according to +## those terms. + +import pkg/libp2p +import pkg/questionable +import pkg/questionable/results + +import ./merkletree +import ../units +import ../errors + +const MaxMerkleTreeSize = 100.MiBs.uint + +proc encode*(self: MerkleTree): seq[byte] = + var pb = initProtoBuffer(maxSize = MaxMerkleTreeSize) + pb.write(1, self.mcodec.uint64) + pb.write(2, self.digestSize.uint64) + pb.write(3, self.leavesCount.uint64) + pb.write(4, self.nodesBuffer) + pb.finish + pb.buffer + +proc decode*(_: type MerkleTree, data: seq[byte]): ?!MerkleTree = + var pb = initProtoBuffer(data, maxSize = MaxMerkleTreeSize) + var mcodecCode: uint64 + var digestSize: uint64 + var leavesCount: uint64 + discard ? pb.getField(1, mcodecCode).mapFailure + discard ? pb.getField(2, digestSize).mapFailure + discard ? pb.getField(3, leavesCount).mapFailure + + let mcodec = MultiCodec.codec(cast[int](mcodecCode)) + if mcodec == InvalidMultiCodec: + return failure("Invalid MultiCodec code " & $cast[int](mcodec)) + + var nodesBuffer = newSeq[byte]() + discard ? pb.getField(4, nodesBuffer).mapFailure + + let tree = ? MerkleTree.init(mcodec, digestSize, leavesCount, nodesBuffer) + success(tree) diff --git a/codex/merkletree/merkletree.nim b/codex/merkletree/merkletree.nim index 7c3442df..3ee76377 100644 --- a/codex/merkletree/merkletree.nim +++ b/codex/merkletree/merkletree.nim @@ -15,6 +15,7 @@ import std/sugar import pkg/questionable/results import pkg/nimcrypto/sha2 import pkg/libp2p/[multicodec, multihash, vbuffer] +import pkg/stew/base58 import ../errors @@ -23,7 +24,7 @@ type mcodec: MultiCodec digestSize: Natural leavesCount: Natural - nodesBuffer: seq[byte] + nodesBuffer*: seq[byte] MerkleProof* = object mcodec: MultiCodec digestSize: Natural @@ -111,8 +112,8 @@ proc build*(self: MerkleTreeBuilder): ?!MerkleTree = copyMem(addr tree.nodesBuffer[0], unsafeAddr self.buffer[0], leavesCount * digestSize) # calculate intermediate nodes - var zero = newSeq[byte](self.digestSize) - var one = newSeq[byte](self.digestSize) + var zero = newSeq[byte](digestSize) + var one = newSeq[byte](digestSize) one[^1] = 0x01 var concatBuf = newSeq[byte](2 * digestSize) @@ -154,6 +155,12 @@ proc len*(self: (MerkleTree | MerkleProof)): Natural = proc nodes*(self: (MerkleTree | MerkleProof)): seq[MultiHash] = toSeq(0.. self.nodeBufferToMultiHash(i)) +proc mcodec*(self: (MerkleTree | MerkleProof)): MultiCodec = + self.mcodec + +proc digestSize*(self: (MerkleTree | MerkleProof)): Natural = + self.digestSize + proc root*(self: MerkleTree): MultiHash = let rootIndex = self.len - 1 self.nodeBufferToMultiHash(rootIndex) @@ -161,6 +168,9 @@ proc root*(self: MerkleTree): MultiHash = proc leaves*(self: MerkleTree): seq[MultiHash] = toSeq(0.. self.nodeBufferToMultiHash(i)) +proc leavesCount*(self: MerkleTree): Natural = + self.leavesCount + proc height*(self: MerkleTree): Natural = computeTreeHeight(self.leavesCount) @@ -213,6 +223,33 @@ proc `$`*(self: MerkleTree): string = "\nleavesCount: " & $self.leavesCount & "\nnodes: " & $self.nodes +proc `==`*(a, b: MerkleTree): bool = + (a.mcodec == b.mcodec) and + (a.digestSize == b.digestSize) and + (a.leavesCount == b.leavesCount) and + (a.nodesBuffer == b.nodesBuffer) + +func init*( + T: type MerkleTree, + mcodec: MultiCodec, + digestSize: Natural, + leavesCount: Natural, + nodesBuffer: seq[byte] +): ?!MerkleTree = + let levels = computeLevels(leavesCount) + let totalNodes = levels[^1].offset + 1 + if totalNodes * digestSize == nodesBuffer.len: + success( + MerkleTree( + mcodec: mcodec, + digestSize: digestSize, + leavesCount: leavesCount, + nodesBuffer: nodesBuffer + ) + ) + else: + failure("Expected nodesBuffer len to be " & $(totalNodes * digestSize) & " but was " & $nodesBuffer.len) + ########################################################### # MerkleProof ########################################################### @@ -226,7 +263,10 @@ proc `$`*(self: MerkleProof): string = "\nnodes: " & $self.nodes func `==`*(a, b: MerkleProof): bool = - (a.index == b.index) and (a.mcodec == b.mcodec) and (a.digestSize == b.digestSize) == (a.nodesBuffer == b.nodesBuffer) + (a.index == b.index) and + (a.mcodec == b.mcodec) and + (a.digestSize == b.digestSize) and + (a.nodesBuffer == b.nodesBuffer) proc init*( T: type MerkleProof, diff --git a/tests/codex/merkletree/testcoders.nim b/tests/codex/merkletree/testcoders.nim new file mode 100644 index 00000000..51b4d8fd --- /dev/null +++ b/tests/codex/merkletree/testcoders.nim @@ -0,0 +1,43 @@ +import std/unittest + +import pkg/questionable/results +import pkg/stew/byteutils + +import pkg/codex/merkletree +import ../helpers + +checksuite "merkletree - coders": + const data = + [ + "0123456789012345678901234567890123456789".toBytes, + "1234567890123456789012345678901234567890".toBytes, + "2345678901234567890123456789012345678901".toBytes, + "3456789012345678901234567890123456789012".toBytes, + "4567890123456789012345678901234567890123".toBytes, + "5678901234567890123456789012345678901234".toBytes, + "6789012345678901234567890123456789012345".toBytes, + "7890123456789012345678901234567890123456".toBytes, + "8901234567890123456789012345678901234567".toBytes, + "9012345678901234567890123456789012345678".toBytes, + ] + + test "encoding and decoding a tree yields the same tree": + var builder = MerkleTreeBuilder.init(multiCodec("sha2-256")).tryGet() + builder.addDataBlock(data[0]).tryGet() + builder.addDataBlock(data[1]).tryGet() + builder.addDataBlock(data[2]).tryGet() + builder.addDataBlock(data[3]).tryGet() + builder.addDataBlock(data[4]).tryGet() + builder.addDataBlock(data[5]).tryGet() + builder.addDataBlock(data[6]).tryGet() + builder.addDataBlock(data[7]).tryGet() + builder.addDataBlock(data[8]).tryGet() + builder.addDataBlock(data[9]).tryGet() + + let tree = builder.build().tryGet() + let encodedBytes = tree.encode() + echo "encode success, size " & $encodedBytes.len + let decodedTree = MerkleTree.decode(encodedBytes).tryGet() + + check: + tree == decodedTree diff --git a/tests/codex/merkletree/testmerkletree.nim b/tests/codex/merkletree/testmerkletree.nim index ed0ed9fd..a7c23b73 100644 --- a/tests/codex/merkletree/testmerkletree.nim +++ b/tests/codex/merkletree/testmerkletree.nim @@ -6,7 +6,7 @@ import pkg/questionable/results import pkg/stew/byteutils import pkg/nimcrypto/sha2 -import codex/merkletree/merkletree +import pkg/codex/merkletree import ../helpers checksuite "merkletree": diff --git a/tests/codex/testmerkletree.nim b/tests/codex/testmerkletree.nim index 3b2af66e..a713e144 100644 --- a/tests/codex/testmerkletree.nim +++ b/tests/codex/testmerkletree.nim @@ -1,3 +1,4 @@ import ./merkletree/testmerkletree +import ./merkletree/testcoders {.warning[UnusedImport]: off.}