From 382a6fc9a45abc6462b6e8b8983cc0c0ea54c7e3 Mon Sep 17 00:00:00 2001 From: optke3 <108488464+optke3@users.noreply.github.com> Date: Mon, 26 Jun 2023 13:05:06 -0400 Subject: [PATCH] [sui 10/x] - merkle tree (#903) * merkle tree impl * - take leftmost 20 bytes in hash - don't assign output of cursor::take_rest to _, instead just drop it * push PREFIXes (MERKLE_LEAF_PREFIX, MERKLE_NODE_PREFIX) to front instead of back * delete testXOR * test construct merkle tree depth exceeded error * invalid merkle proof test cases * comments * rename failure tests * simplification for initializing a vector * fix leafHash bug, add tests for hashLeaf and hashNode --- .../sui/contracts/sources/merkle_tree.move | 381 ++++++++++++++++++ 1 file changed, 381 insertions(+) create mode 100644 target_chains/sui/contracts/sources/merkle_tree.move diff --git a/target_chains/sui/contracts/sources/merkle_tree.move b/target_chains/sui/contracts/sources/merkle_tree.move new file mode 100644 index 00000000..53529af0 --- /dev/null +++ b/target_chains/sui/contracts/sources/merkle_tree.move @@ -0,0 +1,381 @@ +// Implementation of a Merkle tree in Move. Supports constructing a new tree +// with a given depth, as well as proving that a leaf node belongs to the tree. +module pyth::merkle_tree { + use std::vector::{Self}; + use sui::hash::{keccak256}; + use wormhole::bytes::{Self}; + use wormhole::bytes20::{Self, Bytes20, data}; + use wormhole::cursor::{Self, Cursor}; + use pyth::deserialize::{Self}; + + const MERKLE_LEAF_PREFIX: u8 = 0; + const MERKLE_NODE_PREFIX: u8 = 1; + const MERKLE_EMPTY_LEAF_PREFIX: u8 = 2; + + const E_DEPTH_NOT_LARGE_ENOUGH_FOR_MESSAGES: u64 = 1212121; + + // take keccak256 of input data, then return 20 leftmost bytes of result + fun hash(bytes: &vector): Bytes20 { + let hashed_bytes = keccak256(bytes); + let cursor = cursor::new(hashed_bytes); + let bytes20 = bytes::take_bytes(&mut cursor, 20); + cursor::take_rest(cursor); + bytes20::from_bytes(bytes20) + } + + fun emptyLeafHash(): Bytes20 { + let v = vector[MERKLE_EMPTY_LEAF_PREFIX]; + hash(&v) + } + + fun leafHash(data: &vector): Bytes20 { + let v = vector[MERKLE_LEAF_PREFIX]; + vector::append(&mut v, *data); + hash(&v) + } + + fun nodeHash( + childA: Bytes20, + childB: Bytes20 + ): Bytes20 { + if (greaterThan(childA, childB)) { + (childA, childB) = (childB, childA); + }; + // append data_B to data_A + let data_A = bytes20::data(&childA); + let data_B = bytes20::data(&childB); + vector::append(&mut data_A, data_B); + + // create a vector containing MERKLE_NODE_PREFIX and append data_A to back + let v = vector::empty(); + vector::push_back(&mut v, MERKLE_NODE_PREFIX); + vector::append(&mut v, data_A); + hash(&v) + } + + // greaterThan returns whether a is strictly greater than b + // note that data(&a) and data(&b) are both vectors of length 20 + fun greaterThan(a: Bytes20, b: Bytes20): bool{ + // aa and bb both have length 20 + let aa = data(&a); + let bb = data(&b); + let i = 0; + while (i < 20){ + if (*vector::borrow(&aa, i) > *vector::borrow(&bb, i)){ + return true + } else if (*vector::borrow(&bb, i) > *vector::borrow(&aa, i)){ + return false + }; + i = i + 1; + }; + false + } + + // The Sui Move stdlb insert function shifts v[i] and subsequent elements to the right. + // We don't want this behavior, so we define our own setElement function that instead replaces the ith element. + // Reference: https://github.com/MystenLabs/sui/blob/main/crates/sui-framework/packages/move-stdlib/sources/vector.move + fun setElement(a: &mut vector, value: T, index: u64){ + vector::push_back(a, value); // push value to end + vector::swap_remove(a, index); // swap value to correct position and pop last value + } + + // isProofValid returns whether a merkle proof is valid + fun isProofValid( + encodedProof: &mut Cursor, + root: Bytes20, + leafData: vector, + ): bool { + + let currentDigest: Bytes20 = leafHash(&leafData); + let proofSize: u8 = deserialize::deserialize_u8(encodedProof); + let i: u8 = 0; + while (i < proofSize){ + let siblingDigest: Bytes20 = bytes20::new( + deserialize::deserialize_vector(encodedProof, 20) + ); + + currentDigest = nodeHash( + currentDigest, + siblingDigest + ); + i = i + 1; + }; + bytes20::data(¤tDigest) == bytes20::data(&root) + } + + // constructProofs constructs a merkle tree and returns the root of the tree as + // a Bytes20 as well as the vector of encoded proofs + fun constructProofs( + messages: &vector>, + depth: u8 + ) : (Bytes20, vector) { + + if ( 1 << depth < vector::length(messages)) { + abort E_DEPTH_NOT_LARGE_ENOUGH_FOR_MESSAGES + }; + + // empty tree + // The tree is structured as follows: + // 1 + // 2 3 + // 4 5 6 7 + // ... + // In this structure the parent of node x is x//2 and the children + // of node x are x*2 and x*2 + 1. Also, the sibling of the node x + // is x^1. The root is at index 1 and index 0 is not used. + let tree = vector::empty(); + + // empty leaf hash + let cachedEmptyLeafHash: Bytes20 = emptyLeafHash(); + + // Instantiate tree to be a full binary tree with the appropriate depth. + // Add an entry at the end for swapping + let i: u64 = 0; + while (i < (1 << (depth+1)) + 1){ + vector::push_back(&mut tree, cachedEmptyLeafHash); + i = i + 1; + }; + + // Fill in bottom row with leaf hashes + let j: u64 = 0; + while (j < vector::length(messages)){ + setElement(&mut tree, leafHash(vector::borrow(messages, j)), (1 << depth) + j); + j = j + 1; + }; + + // Filling the node hashes from bottom to top + let k: u8 = depth; + while (k>0){ + let level: u8 = k-1; + let levelNumNodes = 1 << level; + let i: u64 = 0; + while (i < levelNumNodes ){ + let id = (1 << level) + i; + let nodeHash = nodeHash(*vector::borrow(&tree, id * 2), *vector::borrow(&tree, id * 2 + 1)); + setElement(&mut tree, nodeHash, id); + i = i + 1; + }; + k = k - 1; + }; + + let root = *vector::borrow(&tree, 1); + + // construct proofs and create encoded proofs vector + let proofs = vector::empty(); + let i: u64 = 0; + while (i < vector::length(messages)){ + let cur_proof = vector::empty(); + vector::push_back(&mut cur_proof, depth); + let idx = (1 << depth) + i; + while (idx > 1) { + vector::append(&mut cur_proof, bytes20::data(vector::borrow(&tree, idx ^ 1))); + + // Jump to parent + idx = idx / 2; + }; + vector::append(&mut proofs, cur_proof); + i = i + 1; + }; + + (root, proofs) + } + + #[test] + fun testGreaterThan(){ + // test 1 + let x = bytes20::new(x"0000000000000000000000000000000000001000"); + let y = bytes20::new(x"0000000000000000000000000000000000000001"); + let res = greaterThan(x, y); + assert!(res==true, 0); + res = greaterThan(y, x); + assert!(res==false, 0); + + // test 2 + x = bytes20::new(x"1100000000000000000000000000000000001000"); + y = bytes20::new(x"1100000000000000000000000000000000000001"); + res = greaterThan(x, y); + assert!(res==true, 0); + + // equality case + x = bytes20::new(x"1100000000000000000000000000000000001001"); + y = bytes20::new(x"1100000000000000000000000000000000001001"); + res = greaterThan(x, y); + assert!(res==false, 0); + } + + #[test] + fun test_hash_leaf() { + let data: vector = x"00640000000000000000000000000000000000000000000000000000000000000000000000000000640000000000000064000000640000000000000064000000000000006400000000000000640000000000000064"; + let hash = leafHash(&data); + let expected = bytes20::new(x"afc6a8ac466430f35895055f8a4c951785dad5ce"); + assert!(hash == expected, 1); + } + + #[test] + fun test_hash_node() { + let h1 = bytes20::new(x"05c51b04b820c0f704e3fdd2e4fc1e70aff26dff"); + let h2 = bytes20::new(x"1e108841c8d21c7a5c4860c8c3499c918ea9e0ac"); + let hash = nodeHash(h1, h2); + let expected = bytes20::new(x"2d0e4fde68184c7ce8af426a0865bd41ef84dfa4"); + assert!(hash == expected, 1); + } + + #[test] + fun testMerkleTreeDepth1(){ + let messages = vector::empty>(); + vector::push_back(&mut messages, x"1234"); + + let (root, proofs) = constructProofs(&messages, 1); + + let proofsCursor = cursor::new(proofs); + let valid = isProofValid(&mut proofsCursor, root, x"1234"); + assert!(valid==true, 0); + + // destroy cursor + cursor::take_rest(proofsCursor); + } + + #[test] + fun testMerkleTreeDepth2(){ + let messages = vector::empty>(); + vector::push_back(&mut messages, x"1234"); + vector::push_back(&mut messages, x"4321"); + vector::push_back(&mut messages, x"11"); + vector::push_back(&mut messages, x"22"); + + let (root, proofs) = constructProofs(&messages, 2); + + let proofsCursor = cursor::new(proofs); + assert!(isProofValid(&mut proofsCursor, root, x"1234")==true, 0); + assert!(isProofValid(&mut proofsCursor, root, x"4321")==true, 0); + assert!(isProofValid(&mut proofsCursor, root, x"11")==true, 0); + assert!(isProofValid(&mut proofsCursor, root, x"22")==true, 0); + + // destroy cursor + cursor::take_rest(proofsCursor); + } + + #[test] + fun testMerkleTreeDepth3(){ + let messages = vector::empty>(); + vector::push_back(&mut messages, x"00"); + vector::push_back(&mut messages, x"4321"); + vector::push_back(&mut messages, x"444444"); + vector::push_back(&mut messages, x"22222222"); + vector::push_back(&mut messages, x"22"); + vector::push_back(&mut messages, x"11"); + vector::push_back(&mut messages, x"100000"); + vector::push_back(&mut messages, x"eeeeee"); + + let (root, proofs) = constructProofs(&messages, 3); + + let proofsCursor = cursor::new(proofs); + assert!(isProofValid(&mut proofsCursor, root, x"00")==true, 0); + assert!(isProofValid(&mut proofsCursor, root, x"4321")==true, 0); + assert!(isProofValid(&mut proofsCursor, root, x"444444")==true, 0); + assert!(isProofValid(&mut proofsCursor, root, x"22222222")==true, 0); + assert!(isProofValid(&mut proofsCursor, root, x"22")==true, 0); + assert!(isProofValid(&mut proofsCursor, root, x"11")==true, 0); + assert!(isProofValid(&mut proofsCursor, root, x"100000")==true, 0); + assert!(isProofValid(&mut proofsCursor, root, x"eeeeee")==true, 0); + + // destroy cursor + cursor::take_rest(proofsCursor); + } + + #[test] + fun testMerkleTreeDepth1InvalidProofs(){ + let messages = vector::empty>(); + vector::push_back(&mut messages, x"1234"); + + let (root, proofs) = constructProofs(&messages, 1); + + let proofsCursor = cursor::new(proofs); + + // use wrong leaf data (it is not included in the tree) + let valid = isProofValid(&mut proofsCursor, root, x"432222"); + assert!(valid==false, 0); + + // destroy cursor + cursor::take_rest(proofsCursor); + } + + #[test] + fun testMerkleTreeDepth2InvalidProofs(){ + let messages = vector::empty>(); + vector::push_back(&mut messages, x"1234"); + vector::push_back(&mut messages, x"4321"); + vector::push_back(&mut messages, x"11"); + vector::push_back(&mut messages, x"22"); + + let (root, proofs) = constructProofs(&messages, 2); + + let proofsCursor = cursor::new(proofs); + // proof fails because we used the proof of x"1234" to try to prove that x"4321" is in the tree + assert!(isProofValid(&mut proofsCursor, root, x"4321")==false, 0); + // proof succeeds + assert!(isProofValid(&mut proofsCursor, root, x"4321")==true, 0); + // proof fails because we used the proof of x"11" to try to prove that x"22" is in the tree + assert!(isProofValid(&mut proofsCursor, root, x"22")==false, 0); + // proof succeeds + assert!(isProofValid(&mut proofsCursor, root, x"22")==true, 0); + + // destroy cursor + cursor::take_rest(proofsCursor); + } + + #[test] + fun testMerkleTreeDepth3InvalidProofs(){ + let messages = vector::empty>(); + vector::push_back(&mut messages, x"00"); + vector::push_back(&mut messages, x"4321"); + vector::push_back(&mut messages, x"444444"); + vector::push_back(&mut messages, x"22222222"); + vector::push_back(&mut messages, x"22"); + vector::push_back(&mut messages, x"11"); + vector::push_back(&mut messages, x"100000"); + vector::push_back(&mut messages, x"eeeeee"); + + let (root, proofs) = constructProofs(&messages, 3); + + let proofsCursor = cursor::new(proofs); + + // test various proof failure cases (because of mismatch between proof and leaf data) + assert!(isProofValid(&mut proofsCursor, root, x"00")==true, 0); + assert!(isProofValid(&mut proofsCursor, root, x"22")==false, 0); + assert!(isProofValid(&mut proofsCursor, root, x"22222222")==false, 0); + assert!(isProofValid(&mut proofsCursor, root, x"22222222")==true, 0); + assert!(isProofValid(&mut proofsCursor, root, x"22")==true, 0); + assert!(isProofValid(&mut proofsCursor, root, x"eeeeee")==false, 0); + assert!(isProofValid(&mut proofsCursor, root, x"4321")==false, 0); + assert!(isProofValid(&mut proofsCursor, root, x"eeeeee")==true, 0); + + // destroy cursor + cursor::take_rest(proofsCursor); + } + + #[test] + #[expected_failure(abort_code = pyth::merkle_tree::E_DEPTH_NOT_LARGE_ENOUGH_FOR_MESSAGES)] + fun testMerkleTreeDepthExceeded1(){ + let messages = vector::empty>(); + vector::push_back(&mut messages, x"00"); + vector::push_back(&mut messages, x"4321"); + vector::push_back(&mut messages, x"444444"); + + constructProofs(&messages, 1); //depth 1 + } + + #[test] + #[expected_failure(abort_code = pyth::merkle_tree::E_DEPTH_NOT_LARGE_ENOUGH_FOR_MESSAGES)] + fun testMerkleTreeDepthExceeded2(){ + let messages = vector::empty>(); + vector::push_back(&mut messages, x"00"); + vector::push_back(&mut messages, x"4321"); + vector::push_back(&mut messages, x"444444"); + vector::push_back(&mut messages, x"22222222"); + vector::push_back(&mut messages, x"22"); + + constructProofs(&messages, 2); // depth 2 + } + +}