diff --git a/build.zig b/build.zig index 448a6a1..f223733 100644 --- a/build.zig +++ b/build.zig @@ -24,11 +24,17 @@ pub fn build(b: *Builder) void { .target = target, }); const run_main_tests = b.addRunArtifact(main_tests); + const tests_tests = b.addTest(.{ .root_source_file = .{ .cwd_relative = "src/tests.zig" }, .optimize = optimize, .target = target, }); + const tweak_hash = b.dependency("tweak_hash", .{ + .target = target, + .optimize = optimize, + }).module("tweak_hash"); + tests_tests.root_module.addImport("tweak_hash", tweak_hash); tests_tests.root_module.addImport("ssz.zig", mod); const run_tests_tests = b.addRunArtifact(tests_tests); diff --git a/build.zig.zon b/build.zig.zon index 41d7a45..3acb9c0 100644 --- a/build.zig.zon +++ b/build.zig.zon @@ -2,5 +2,11 @@ .name = .ssz, .fingerprint = 0x1d34bd0ceb1dfc2d, .version = "0.0.4", + .dependencies = .{ + .tweak_hash = .{ + .url = "git+https://github.com/bhaskar1001101/tweak-hash#63611fd65d9db28a7ea342ff2fd61b1e201727fd", + .hash = "tweak_hash-0.0.0-b5xZ7JkRAABSp1CPFJpayRnzIt2tnLjzjTPfveqLiUXu", + }, + }, .paths = .{""}, } diff --git a/src/lib.zig b/src/lib.zig index b5c3aaf..f45a1fe 100644 --- a/src/lib.zig +++ b/src/lib.zig @@ -818,3 +818,88 @@ fn bytesToBits(comptime N: usize, src: [N]u8) [N * 8]bool { } return bitvector; } + +pub const MerkleTree = [][]u8; +pub const MerklePath = [][]u8; + +// returns a full merkle tree with a tweakable hash function from provided leaf hashes +pub fn merkleizeWithTweak(comptime TweakHash: type, hash: TweakHash, allocator: std.mem.Allocator, parameter: []u8, chunks: [][]u8) !MerkleTree { + const num_leaves = chunks.len; + std.debug.assert(std.math.isPowerOfTwo(num_leaves)); + + const node_count = (2 * num_leaves) - 1; + var nodes = try allocator.alloc([]u8, node_count); + + for (chunks, 0..) |c, i| { + const leaf_pos = node_count - num_leaves + i; + nodes[leaf_pos] = try allocator.dupe(u8, c); + } + + var level: u8 = 1; + var level_size: usize = num_leaves / 2; + var level_offset: usize = node_count - num_leaves - level_size; + + while (level_size > 0) { + for (0..level_size) |i| { + const left_child = nodes[level_offset + level_size + i * 2]; + const right_child = nodes[level_offset + level_size + i * 2 + 1]; + + var combined = [_][]u8{ left_child, right_child }; + + nodes[level_offset + i] = try allocator.alloc(u8, hash.hash_size); + const tweak = hash.tree_tweak(level, @as(u32, @intCast(i))); + hash.hash(parameter, tweak, &combined, nodes[level_offset + i]); + } + + level += 1; + level_size /= 2; + level_offset -= level_size; + } + + return nodes; +} + +// returns a merkle path from a leaf to the root +pub fn buildPath(allocator: std.mem.Allocator, tree: MerkleTree, leaf_index: usize) !MerklePath { + const height = std.math.log2_int(usize, tree.len); + + var path = try allocator.alloc([]u8, height); + var current_index = leaf_index; + const num_leaves = @as(u32, 1) << @intCast(height); + const total_nodes = (2 * num_leaves) - 1; + var node_index = total_nodes - num_leaves + current_index; + for (0..height) |level| { + const is_left = current_index % 2 == 0; + const sibling_offset: isize = if (is_left) 1 else -1; + + const sibling_node_index: usize = @intCast(@as(isize, @intCast(node_index)) + sibling_offset); + path[level] = try allocator.dupe(u8, tree[sibling_node_index]); + + current_index /= 2; + node_index = (node_index - 1) / 2; + } + + return path; +} + +// verify a merkle path +pub fn verifyPath(comptime TweakHash: type, parameter: []u8, hash: TweakHash, leaf_index: usize, root: []u8, leaf_hash: []u8, path: MerklePath) bool { + var current_index = leaf_index; + for (0..path.len) |level| { + const is_left = current_index % 2 == 0; + const sibling = path[level]; + + const combined = if (is_left) [_][]u8{ leaf_hash, sibling } else [_][]u8{ sibling, leaf_hash }; + const tweak = hash.tree_tweak(@as(u8, @intCast(level + 1)), @as(u32, @intCast(current_index / 2))); + + hash.hash(parameter, tweak, &combined, leaf_hash); + + current_index /= 2; + } + + return std.mem.eql(u8, leaf_hash, root); +} + +pub fn treeRoot(tree: MerkleTree) []u8 { + return tree[0]; +} diff --git a/src/tests.zig b/src/tests.zig index 9a67e31..9bbd5ac 100644 --- a/src/tests.zig +++ b/src/tests.zig @@ -857,3 +857,84 @@ test "structs with nested fixed/variable size u8 array" { try expect(var_signed_block.message.body.slot == deserialized_var_block.message.body.slot); try expect(std.mem.eql(u8, var_signed_block.message.body.data[0..], deserialized_var_block.message.body.data[0..])); } + +test "MerkleTree build, path, and verify" { + const testing = std.testing; + const allocator = testing.allocator; + + const num_leaves: usize = 1024; + // Each leaf will consist of `leaf_len` chunks of `hash.hash_size` bytes. + const leaf_len = 3; + + const parameter = try allocator.alloc(u8, 16); + defer allocator.free(parameter); + std.crypto.random.bytes(parameter); + + const ShaTweakHash = @import("tweak_hash").ShaTweakHash; + const hash = ShaTweakHash.init(16, 24); + + // Generate random leaves. Each leaf is a slice of leaf_len chunks. + var leaves = try allocator.alloc([][]u8, num_leaves); + defer { + for (leaves) |leaf_chunks| { + for (leaf_chunks) |chunk| { + allocator.free(chunk); + } + allocator.free(leaf_chunks); + } + allocator.free(leaves); + } + + // Fill every leaf with `leaf_len` random chunks. + for (0..num_leaves) |i| { + leaves[i] = try allocator.alloc([]u8, leaf_len); + for (0..leaf_len) |j| { + leaves[i][j] = try allocator.alloc(u8, hash.hash_size); + std.crypto.random.bytes(leaves[i][j]); + } + } + + // Hash the leaves with the level 0 tweak. + var leaf_hashes = try allocator.alloc([]u8, num_leaves); + defer { + for (leaf_hashes) |h| allocator.free(h); + allocator.free(leaf_hashes); + } + for (leaves, 0..) |leaf_chunks, i| { + leaf_hashes[i] = try allocator.alloc(u8, hash.hash_size); + const tweak = hash.tree_tweak(0, @as(u32, @intCast(i))); + hash.hash(parameter, tweak, leaf_chunks, leaf_hashes[i]); + } + + const tree = try libssz.merkleizeWithTweak(ShaTweakHash, hash, allocator, parameter, leaf_hashes); + defer { + for (tree) |node| { + allocator.free(node); + } + allocator.free(tree); + } + + const root = libssz.treeRoot(tree); + + // build path for all leaves and verify them + for (0..num_leaves) |idx| { + const path = try libssz.buildPath(allocator, tree, idx); + defer { + for (path) |node| { + allocator.free(node); + } + allocator.free(path); + } + + const leaf_tweak = hash.tree_tweak(0, @as(u32, @intCast(idx))); + const leaf_hash_expected = try allocator.alloc(u8, hash.hash_size); + defer allocator.free(leaf_hash_expected); + hash.hash(parameter, leaf_tweak, leaves[idx], leaf_hash_expected); + + const leaf_hash_to_verify = try allocator.dupe(u8, leaf_hash_expected); + defer allocator.free(leaf_hash_to_verify); + + const ok = libssz.verifyPath(ShaTweakHash, parameter, hash, idx, root, leaf_hash_to_verify, path); + try testing.expect(ok); + } +}