Getting and setting Merkle tree last stable index (client self-healin…

…g support)

libzkbob-rs-node/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "libzkbob-rs-node"
-version = "0.2.0"
+version = "0.2.1"
 authors = ["Dmitry Vdovin <voidxnull@gmail.com>"]
 repository = "https://github.com/zkBob/libzkbob-rs/"
 license = "MIT OR Apache-2.0"
@@ -11,8 +11,8 @@ exclude = ["index.node"]
 crate-type = ["cdylib"]
 
 [dependencies]
-libzkbob-rs = { version = "0.9.0", features = ["native"] }
-#libzkbob-rs =  { path = "../libzkbob-rs", features = ["native"] }
+#libzkbob-rs = { version = "0.9.0", features = ["native"] }
+libzkbob-rs =  { path = "../libzkbob-rs", features = ["native"] }
 neon = { version = "0.10.0", default-features = false, features = ["channel-api", "napi-6", "promise-api"] }
 # FIXME: Using a random fork for now
 neon-serde = { git = "https://github.com/NZXTCorp/neon-serde.git", branch = "refactor/update-neon-0.10" }

libzkbob-rs-node/index.d.ts

@@ -24,7 +24,11 @@ declare class MerkleTree {
         new_hashes_left_index: number,
         new_hashes_right_index: number,
     ): any;
-    rollback(index: number): void;
+    rollback(index: number): number;
+    wipe(): void;
+    getLeftSiblings(index: number): string[];
+    getLastStableIndex(): number;
+    setLastStableIndex(index: number): void;
 }
 
 declare class TxStorage {

libzkbob-rs-node/index.js

@@ -62,9 +62,21 @@ class MerkleTree {
         return zp.merkleRollback(this.inner, index)
     }
 
+    wipe() {
+        zp.merkleWipe(this.inner)
+    }
+
     getLeftSiblings(index) {
         return zp.merkleGetLeftSiblings(this.inner, index)
     }
+
+    getLastStableIndex() {
+        return zp.merkleGetLastStableIndex(this.inner)
+    }
+
+    setLastStableIndex(index) {
+        zp.merkleSetLastStableIndex(this.inner, index)
+    };
 }
 
 class TxStorage {

libzkbob-rs-node/package.json

@@ -1,6 +1,6 @@
 {
   "name": "libzkbob-rs-node",
-  "version": "0.2.0",
+  "version": "0.2.1",
   "description": "Neon version of libzkbob-rs",
   "main": "index.js",
   "types": "index.d.ts",

libzkbob-rs-node/src/lib.rs

@@ -63,7 +63,10 @@ fn main(mut cx: ModuleContext) -> NeonResult<()> {
     cx.export_function("merkleGetAllNodes", merkle::merkle_get_all_nodes)?;
     cx.export_function("merkleGetVirtualNode", merkle::merkle_get_virtual_node)?;
     cx.export_function("merkleRollback", merkle::merkle_rollback)?;
+    cx.export_function("merkleWipe", merkle::merkle_wipe)?;
     cx.export_function("merkleGetLeftSiblings", merkle::merkle_get_left_siblings)?;
+    cx.export_function("merkleGetLastStableIndex", merkle::merkle_get_last_stable_index)?;
+    cx.export_function("merkleSetLastStableIndex", merkle::merkle_set_last_stable_index)?;
 
     cx.export_function("txStorageNew", storage::tx_storage_new)?;
     cx.export_function("txStorageAdd", storage::tx_storage_add)?;

libzkbob-rs-node/src/merkle.rs

@@ -148,9 +148,9 @@ pub fn merkle_get_node(mut cx: FunctionContext) -> JsResult<JsValue> {
 pub fn merkle_get_next_index(mut cx: FunctionContext) -> JsResult<JsValue> {
     let tree = cx.argument::<BoxedMerkleTree>(0)?;
 
-    let root = tree.read().unwrap().inner.next_index();
+    let next_index = tree.read().unwrap().inner.next_index();
 
-    let result = neon_serde::to_value(&mut cx, &root).unwrap();
+    let result = neon_serde::to_value(&mut cx, &next_index).unwrap();
 
     Ok(result)
 }
@@ -209,15 +209,25 @@ pub fn merkle_get_virtual_node(mut cx: FunctionContext) -> JsResult<JsValue> {
     Ok(result)
 }
 
-pub fn merkle_rollback(mut cx: FunctionContext) -> JsResult<JsUndefined> {
+pub fn merkle_rollback(mut cx: FunctionContext) -> JsResult<JsValue> {
     let tree = cx.argument::<BoxedMerkleTree>(0)?;
     let rollback_index = {
         let num = cx.argument::<JsNumber>(1)?;
         num.value(&mut cx) as u64
     };
 
     let mut a = tree.write().unwrap();
-    a.inner.rollback(rollback_index);
+    let new_next_index = a.inner.rollback(rollback_index);
+
+    let result = neon_serde::to_value(&mut cx, &new_next_index).unwrap();
+
+    Ok(result)
+}
+
+pub fn merkle_wipe(mut cx: FunctionContext) -> JsResult<JsUndefined> {
+    let tree = cx.argument::<BoxedMerkleTree>(0)?;
+
+    tree.write().unwrap().inner.wipe();
 
     Ok(cx.undefined())
 }
@@ -271,4 +281,26 @@ pub fn merkle_get_left_siblings(mut cx: FunctionContext) -> JsResult<JsValue> {
         Ok(val) => Ok(neon_serde::to_value(&mut cx, &val).unwrap()),
         Err(e) => cx.throw_error(e),
     }
+}
+
+pub fn merkle_get_last_stable_index(mut cx: FunctionContext) -> JsResult<JsValue> {
+    let tree = cx.argument::<BoxedMerkleTree>(0)?;
+
+    let stable_index = tree.read().unwrap().inner.get_last_stable_index();
+
+    let result = neon_serde::to_value(&mut cx, &stable_index).unwrap();
+
+    Ok(result)
+}
+
+pub fn merkle_set_last_stable_index(mut cx: FunctionContext) -> JsResult<JsUndefined> {
+    let tree = cx.argument::<BoxedMerkleTree>(0)?;
+    let stable_index = {    
+        let num = cx.argument::<JsNumber>(1)?;
+        num.value(&mut cx) as u64
+    };
+
+    tree.write().unwrap().inner.set_last_stable_index(Some(stable_index));
+
+    Ok(cx.undefined())
 }

libzkbob-rs-node/test.js

@@ -2,19 +2,47 @@ const zp = require('./index.js')
 
 let tree = new zp.MerkleTree('./testdb');
 
+console.log(`Root at index ${tree.getNextIndex()}: ${tree.getRoot()}`);
+
 console.log('Building the tree...');
-for (let i = 0; i < 512; ++i) {
-    tree.addHash(i, Buffer.alloc(32));
+for (let i = 0; i < 512; ++i) {;
+    tree.addHash(i, Buffer.alloc(32).fill(i % 10));
 }
 
 console.log('Calculating proof...')
 let proof = tree.getProof(50);
 console.log('Proof', proof);
 
-console.log('Getting siblings...')
+let fixed_index = tree.getNextIndex();
+console.log(`Root at index ${tree.getNextIndex()}: ${tree.getRoot()}`);
+tree.setLastStableIndex(fixed_index);
+
+console.log('Appending hashes...');
+for (let i = 0; i < 256; ++i) {
+    tree.addHash(i + fixed_index, Buffer.alloc(32).fill(i % 20));
+}
+
+const siblingsIndex = 640;
+console.log(`Getting siblings at index ${siblingsIndex} ...`);
 try {
-    let siblings = tree.getLeftSiblings(384);
+    let siblings = tree.getLeftSiblings(siblingsIndex);
     console.log(siblings);
 } catch(err) {
     console.error(`Cannot get siblings: ${err}`);
-}
+}
+
+console.log(`Root at index ${tree.getNextIndex()}: ${tree.getRoot()}`);
+
+console.log(`Rollback tree to the index ${fixed_index}...`);
+tree.rollback(tree.getLastStableIndex());
+
+console.log(`Root at index ${tree.getNextIndex()}: ${tree.getRoot()}`);
+
+console.log(`Wipe Merkle tree...`);
+tree.wipe();
+
+console.log(`Root at index ${tree.getNextIndex()}: ${tree.getRoot()}`);
+
+
+
+

libzkbob-rs-wasm/src/client/mod.rs

@@ -533,4 +533,14 @@ impl UserAccount {
     pub fn wipe_state(&self) {
         self.inner.borrow_mut().state.wipe();
     }
+
+    #[wasm_bindgen(js_name = "treeGetStableIndex")]
+    pub fn tree_get_stable_index(&self) -> u64 {
+        self.inner.borrow_mut().state.tree.get_last_stable_index()
+    }
+
+    #[wasm_bindgen(js_name = "treeSetStableIndex")]
+    pub fn tree_set_stable_index(&self, stable_index: u64) {
+        self.inner.borrow_mut().state.tree.set_last_stable_index(Some(stable_index));
+    }
 }

libzkbob-rs/src/merkle.rs

@@ -927,6 +927,10 @@ impl<D: KeyValueDB, P: PoolParams> MerkleTree<D, P> {
         let new_root = self.get_node_full_batched(constants::HEIGHT as u32, 0, &mut update_batch, new_next_index);
         self.db.write(update_batch).unwrap();
 
+        if self.next_index() < self.get_last_stable_index() {
+            self.set_last_stable_index(None);
+        }
+
         new_root
     }
 
@@ -1022,6 +1026,7 @@ impl<D: KeyValueDB, P: PoolParams> MerkleTree<D, P> {
         self.next_index = 0;
         self.write_index_to_database(FIRST_INDEX_KEY, None);
         self.write_index_to_database(NEXT_INDEX_KEY, Some(self.next_index));
+        self.set_last_stable_index(None);
     }
 
     pub fn get_all_nodes(&self) -> Vec<Node<P::Fr>> {
@@ -1052,6 +1057,17 @@ impl<D: KeyValueDB, P: PoolParams> MerkleTree<D, P> {
         self.next_index
     }
 
+    pub fn get_last_stable_index(&self) -> u64 {
+        match self.get_named_index_opt("stable_index") {
+            Some(val) => val,
+            _ => 0,
+        }
+    }
+
+    pub fn set_last_stable_index(&mut self, value: Option<u64>) {
+        self.set_named_index("stable_index", value);
+    }
+
     fn update_first_index(&mut self, first_index: Option<u64>) -> bool {
         if first_index != self.first_index {
             if first_index.is_some() &&
@@ -1225,6 +1241,15 @@ impl<D: KeyValueDB, P: PoolParams> MerkleTree<D, P> {
         }
     }
 
+    fn set_named_index(&mut self, key: &str, value: Option<u64>) {
+        let mut batch = self.db.transaction();
+        match value {
+            Some(val) => batch.put(DbCols::NamedIndex as u32, key.as_bytes(), &val.to_be_bytes()),
+            None => batch.delete(DbCols::NamedIndex as u32, key.as_bytes())
+        }
+        self.db.write(batch).unwrap();
+    }
+
     fn set_named_index_batched(&mut self, batch: &mut DBTransaction, key: &str, value: u64) {
         batch.put(DbCols::NamedIndex as u32, key.as_bytes(), &value.to_be_bytes());
     }
@@ -2163,4 +2188,18 @@ mod tests {
         assert_eq!(tree.next_index(), 0);
         assert_eq!(tree.get_leaves().len(), 0);
     }
+
+
+    #[test]
+    fn test_stable_index() {
+        let mut tree = MerkleTree::new(create(3), POOL_PARAMS.clone());
+
+        assert_eq!(tree.get_last_stable_index(), 0);
+
+        tree.set_last_stable_index(Some(1024));
+        assert_eq!(tree.get_last_stable_index(), 1024);
+
+        tree.set_last_stable_index(None);
+        assert_eq!(tree.get_last_stable_index(), 0);
+    }
 }