feat: merkle root ffi

.gitignore

@@ -1,4 +1,6 @@
 node_modules
+target
+
 .hypothesis
 .vscode
 *cache*

merkly/accelerator/Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "accelerator"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+tiny-keccak = { version = "2.0.2", features = ["keccak"] }
+
+
+[lib]
+crate-type = ["cdylib"]
+
+[profile.release]
+strip = true

merkly/accelerator/build.sh

@@ -0,0 +1,2 @@
+#!/bin/sh
+cargo b -r --out-dir . -Z unstable-options

merkly/accelerator/mtreers.py

@@ -0,0 +1,80 @@
+from ctypes import CDLL, POINTER, c_size_t, c_ubyte, c_bool
+from typing import List
+
+from merkly.node import Node
+
+
+class MTreers:
+    def __init__(self) -> None:
+        self.lib = CDLL("./merkly/accelerator/libaccelerator.dylib")
+
+        self.lib.make_root.argtypes = [POINTER(POINTER(c_ubyte)), c_size_t]
+        self.lib.make_root.restype = POINTER(c_ubyte)
+
+        self.lib.free_root.argtypes = [POINTER(c_ubyte)]
+        self.lib.free_root.restype = None
+
+        # self.lib.make_proof.argtypes = [
+        #     POINTER(POINTER(c_ubyte)),
+        #     c_size_t,
+        #     POINTER(c_ubyte),
+        # ]
+        # self.lib.make_proof.restype = POINTER(POINTER(c_ubyte))
+
+        # self.lib.verify.argtypes = [
+        #     POINTER(POINTER(c_ubyte)),
+        #     c_size_t,
+        #     POINTER(c_ubyte),
+        # ]
+        # self.lib.verify.restype = c_bool
+
+    def make_root(self, leaves: List[bytes]) -> bytes:
+        len_leaves = len(leaves)
+        leaves_pointers = (POINTER(c_ubyte) * len_leaves)()
+
+        for i, leaf in enumerate(leaves):
+            array_type = c_ubyte * 32
+            leaves_pointers[i] = array_type(*leaf)
+
+        root_ptr = self.lib.make_root(leaves_pointers, len_leaves)
+        root = bytes(root_ptr[:32])
+        self.lib.free_root(root_ptr)
+        return root
+
+    def make_proof(self, leaves: List[bytes], leaf: bytes) -> bytes:
+        leaf_pointer = (c_ubyte * 32)(*leaf)
+        len_leaves = len(leaves)
+        leaves_pointers = (POINTER(c_ubyte) * len_leaves)()
+
+        for i, leaf in enumerate(leaves):
+            array_type = c_ubyte * 32
+            leaves_pointers[i] = array_type(*leaf)
+
+        result_pointers = self.lib.make_proof(leaves_pointers, len_leaves, leaf_pointer)
+
+        result = []
+        for i in range(len_leaves):
+            result = list(bytes(result_pointers[i][:33]))
+            flag = result[32]
+            data = result[:32]
+            result.append(Node(data=data, side=flag))
+        return result
+
+    def verify(self, proof: List[Node], leaf: bytes) -> bytes:
+        proof_bytes = []
+
+        for node in proof:
+            flag = node.side.value
+            data = node.data
+            proof_bytes.append([*data, flag])
+
+        len_proof = len(proof_bytes)
+        proof_pointers = (POINTER(c_ubyte) * len_proof)()
+
+        for i, node in enumerate(proof_bytes):
+            array_type = c_ubyte * 33
+            proof_pointers[i] = array_type(*node)
+
+        leaf_pointer = (c_ubyte * 32)(*leaf)
+
+        return self.lib.verify(proof_pointers, len_proof, leaf_pointer)

merkly/accelerator/src/lib.rs

@@ -0,0 +1,105 @@
+use std::slice;
+use tiny_keccak::{Hasher, Keccak};
+
+pub fn hash_it(data: &[u8], buffer: &mut [u8; 32]) {
+    let mut k256 = Keccak::v256();
+
+    k256.update(data);
+    k256.finalize(buffer);
+}
+
+fn hash_function(left: &[u8], right: &[u8], buffer: &mut [u8; 32]) {
+    let concat = [left, right].concat();
+
+    hash_it(&concat, buffer)
+}
+
+/// # Safety
+///
+/// FFI to Python.
+#[no_mangle]
+pub unsafe extern "C" fn make_root(leaves_ptr: *const *const u8, len_leaves: usize) -> *mut u8 {
+    let mut leaves = unsafe { slice::from_raw_parts(leaves_ptr, len_leaves) }
+        .iter()
+        .map(|leaf_ptr| unsafe { slice::from_raw_parts(*leaf_ptr, 32).to_vec() })
+        .collect::<Vec<Vec<u8>>>();
+    let mut node = [0u8; 32];
+
+    while leaves.len() > 1 {
+        let mut next_level = Vec::new();
+
+        for leaf_pair in leaves.chunks(2) {
+            match leaf_pair {
+                [left, right] => hash_function(left, right, &mut node),
+                [left] => node.copy_from_slice(left),
+                _ => unreachable!(),
+            };
+            next_level.push(node.to_vec());
+        }
+
+        leaves = next_level;
+    }
+
+    let root = node;
+    let boxed_root = root.to_vec().into_boxed_slice();
+    Box::into_raw(boxed_root) as *mut u8
+}
+
+/// # Safety
+///
+#[no_mangle]
+pub unsafe extern "C" fn free_root(ptr: *mut u8) {
+    unsafe {
+        let _ = Box::from_raw(slice::from_raw_parts_mut(ptr, 32));
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn setup_leaves() -> Vec<[u8; 32]> {
+        vec![
+            // a
+            [
+                58, 194, 37, 22, 141, 245, 66, 18, 162, 92, 28, 1, 253, 53, 190, 191, 234, 64, 143,
+                218, 194, 227, 29, 221, 111, 128, 164, 187, 249, 165, 241, 203,
+            ],
+            // b
+            [
+                181, 85, 61, 227, 21, 224, 237, 245, 4, 217, 21, 10, 248, 45, 175, 165, 196, 102,
+                127, 166, 24, 237, 10, 111, 25, 198, 155, 65, 22, 108, 85, 16,
+            ],
+            // c
+            [
+                11, 66, 182, 57, 60, 31, 83, 6, 15, 227, 221, 191, 205, 122, 173, 204, 168, 148,
+                70, 90, 90, 67, 143, 105, 200, 125, 121, 11, 34, 153, 185, 178,
+            ],
+            // d
+            [
+                241, 145, 142, 133, 98, 35, 110, 177, 122, 220, 133, 2, 51, 47, 76, 156, 130, 188,
+                20, 225, 155, 252, 10, 161, 10, 182, 116, 255, 117, 179, 210, 243,
+            ],
+        ]
+    }
+
+    #[test]
+    fn test_make_root() {
+        let root = [
+            104, 32, 63, 144, 233, 208, 125, 197, 133, 146, 89, 215, 83, 110, 135, 166, 186, 157,
+            52, 95, 37, 82, 181, 185, 222, 41, 153, 221, 206, 156, 225, 191,
+        ];
+        let leaves = setup_leaves();
+        let leaves_ptrs: Vec<*const u8> = leaves.iter().map(|leaf| leaf.as_ptr()).collect();
+
+        let root_ptr = unsafe { make_root(leaves_ptrs.as_ptr(), leaves.len()) };
+        let result_slice = unsafe { slice::from_raw_parts(root_ptr, 32) };
+        let result = result_slice.to_vec();
+
+        assert_eq!(result, root);
+        unsafe { free_root(root_ptr) };
+
+        let result_dangling_ptr = unsafe { slice::from_raw_parts(root_ptr, 32) };
+        assert_eq!(result_dangling_ptr, [0u8; 32]);
+    }
+}

merkly/mtree.py

@@ -56,7 +56,7 @@ def proof(self, raw_leaf: str) -> List[Node]:
             self.leaves, [], self.hash_function(raw_leaf.encode(), bytes())
         )
 
-    def verify(self, proof: List[bytes], raw_leaf: str) -> bool:
+    def verify(self, proof: List[Node], raw_leaf: str) -> bool:
         full_proof = [self.hash_function(raw_leaf.encode(), bytes())]
         full_proof.extend(proof)
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "merkly"
-version = "1.1.0"
+version = "1.2.0-beta"
 description = "🌳 The simple and easy implementation of Merkle Tree"
 authors = ["Lucas Oliveira <[email protected]>"]
 repository = "https://github.com/olivmath/merkly.git"
@@ -26,21 +26,22 @@ classifiers = [
     "Programming Language :: Python :: 3.10",
     "Topic :: Software Development :: Libraries",
     "Topic :: Software Development :: Libraries :: Python Modules",
-    "Topic :: Security :: Cryptography"
+    "Topic :: Security :: Cryptography",
 ]
 
 [tool.poetry.dependencies]
-    pycryptodome = "^3.19.0"
-    pydantic = "^1.10.2"
-    python = "^3.8"
+pycryptodome = "^3.19.0"
+pydantic = "^1.10.2"
+python = "^3.8"
+keccaky = "^0.3.1"
 
 [tool.poetry.dev-dependencies]
-    conventional-pre-commit = "^3.0.0"
-    pre-commit = "^3.0.3"
-    coverage = "^7.2.7"
-    pyclean = "^2.2.0"
-    pytest = "^7.2.1"
-    black = "^23.1.0"
+conventional-pre-commit = "^3.0.0"
+pre-commit = "^3.0.3"
+coverage = "^7.2.7"
+pyclean = "^2.2.0"
+pytest = "^7.2.1"
+black = "^23.1.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]
@@ -52,6 +53,7 @@ lint = "scripts.poetry:lint"
 
 [tool.pytest.ini_options]
 markers = [
+    "ffi: mark test as requiring Rust FFI compilation (deselect with '-m \"not ffi\"')",
     "merkletreejs: marks tests as merkletreejs (deselect with '-m \"not merkletreejs\"')",
     "benchmark: marks tests as benchmark (deselect with '-m \"not benchmark\"')",
-]
+]

test/accelerator/test_accelerator.py

@@ -0,0 +1,19 @@
+from merkly.accelerator.mtreers import MTreers
+from merkly.mtree import MerkleTree
+from keccaky import hash_it_bytes
+from pytest import mark
+
+
+@mark.ffi
+def test_merkle_root_ffi(compile_rust_ffi):
+    leaves_short = ["a", "b", "c", "d"]
+    leaves_bytes = list(map(lambda x: hash_it_bytes(x.encode()), leaves_short))
+
+    tree = MerkleTree(leaves_short)
+    treers = MTreers()
+
+
+    root = tree.root
+    result = treers.make_root(leaves_bytes)
+
+    assert list(root) == list(result)

test/conftest.py

@@ -0,0 +1,10 @@
+import subprocess
+
+from pytest import fixture
+
+
+@fixture(scope="session", autouse=True)
+def compile_rust_ffi(request):
+    if any("ffi" in item.keywords for item in request.session.items):
+        print("Compiling Rust FFI")
+        subprocess.run(["./build.sh"], cwd="./merkly/accelerator", check=True)