From 3d1b7886ea484c98981baedf7aba6c87742ef252 Mon Sep 17 00:00:00 2001 From: arnaucube Date: Wed, 8 Jul 2020 15:39:03 +0200 Subject: [PATCH] Add DumpLeafs & ImportDumpedLeafs for MerkleTree backup functionality --- merkletree.go | 69 ++++++++++++++++++++++++++++++++++++++++++---- merkletree_test.go | 24 ++++++++++++++++ node.go | 2 +- utils.go | 4 --- 4 files changed, 89 insertions(+), 10 deletions(-) diff --git a/merkletree.go b/merkletree.go index b89fa29..1711832 100644 --- a/merkletree.go +++ b/merkletree.go @@ -76,6 +76,22 @@ func (h *Hash) BigInt() *big.Int { return new(big.Int).SetBytes(common.SwapEndianness(h[:])) } +// Bytes returns the [32]byte representation of the *Hash +func (h *Hash) Bytes() [32]byte { + bi := new(big.Int).SetBytes(common.SwapEndianness(h[:])).Bytes() + b := [32]byte{} + copy(b[:], bi[:]) + return b +} + +// NewBigIntFromBytes returns a *big.Int from a byte array, swapping the endianness in the process. This is the intended method to get a *big.Int from a byte array that previously has ben generated by the Hash.Bytes() method. +func NewBigIntFromBytes(b []byte) (*big.Int, error) { + if len(b) != 32 { + return nil, fmt.Errorf("Expected 32 bytes, found %d bytes", len(b)) + } + return new(big.Int).SetBytes(common.SwapEndianness(b[:32])), nil +} + // NewHashFromBigInt returns a *Hash representation of the given *big.Int func NewHashFromBigInt(b *big.Int) *Hash { r := &Hash{} @@ -279,8 +295,16 @@ func (mt *MerkleTree) addNode(tx db.Tx, n *Node) (*Hash, error) { return k, nil } -// Delete removes the specified Key from the MerkleTree, and updates the pad from the delted key to the Root with the new values. -// This method removes the key from the MerkleTree, but does not remove the old nodes from the key-value database; this means that if the tree is accessed by an old Root where the key was not deleted yet, the key will still exist. If is desired to remove the key-values from the database that are not under the current Root, an option could be to dump all the claims and import them in a new MerkleTree in a new database, but this will loose all the Root history of the MerkleTree +// Delete removes the specified Key from the MerkleTree and updates the path +// from the deleted key to the Root with the new values. This method removes +// the key from the MerkleTree, but does not remove the old nodes from the +// key-value database; this means that if the tree is accessed by an old Root +// where the key was not deleted yet, the key will still exist. If is desired +// to remove the key-values from the database that are not under the current +// Root, an option could be to dump all the leafs (using mt.DumpLeafs) and +// import them in a new MerkleTree in a new database (using +// mt.ImportDumpedLeafs), but this will loose all the Root history of the +// MerkleTree func (mt *MerkleTree) Delete(k *big.Int) error { // verify that the MerkleTree is writable if !mt.writable { @@ -630,7 +654,7 @@ func VerifyProof(rootKey *Hash, proof *Proof, k, v *big.Int) bool { } // RootFromProof calculates the root that would correspond to a tree whose -// siblings are the ones in the proof with the claim hashing to hIndex and +// siblings are the ones in the proof with the leaf hashing to hIndex and // hValue. func RootFromProof(proof *Proof, k, v *big.Int) (*Hash, error) { kHash := NewHashFromBigInt(k) @@ -708,8 +732,8 @@ func (mt *MerkleTree) walk(key *Hash, f func(*Node)) error { // Walk iterates over all the branches of a MerkleTree with the given rootKey // if rootKey is nil, it will get the current RootKey of the current state of the MerkleTree. // For each node, it calls the f function given in the parameters. -// See some examples of the Walk function usage in the merkletree_test.go -// test functions: TestMTWalk, TestMTWalkGraphViz, TestMTWalkDumpClaims +// See some examples of the Walk function usage in the merkletree.go and +// merkletree_test.go func (mt *MerkleTree) Walk(rootKey *Hash, f func(*Node)) error { if rootKey == nil { rootKey = mt.Root() @@ -773,3 +797,38 @@ func (mt *MerkleTree) PrintGraphViz(rootKey *Hash) error { fmt.Println(w) return nil } + +// DumpLeafs returns all the Leafs that exist under the given Root. If no Root +// is given (nil), it uses the current Root of the MerkleTree. +func (mt *MerkleTree) DumpLeafs(rootKey *Hash) ([]byte, error) { + var b []byte + err := mt.Walk(rootKey, func(n *Node) { + if n.Type == NodeTypeLeaf { + l := n.Entry[0].Bytes() + r := n.Entry[1].Bytes() + b = append(b, append(l[:], r[:]...)...) + } + }) + return b, err +} + +// ImportDumpedLeafs parses and adds to the MerkleTree the dumped list of leafs +// from the DumpLeafs function. +func (mt *MerkleTree) ImportDumpedLeafs(b []byte) error { + for i := 0; i < len(b); i += 64 { + lr := b[i : i+64] + lB, err := NewBigIntFromBytes(lr[:32]) + if err != nil { + return err + } + rB, err := NewBigIntFromBytes(lr[32:]) + if err != nil { + return err + } + err = mt.Add(lB, rB) + if err != nil { + return err + } + } + return nil +} diff --git a/merkletree_test.go b/merkletree_test.go index 83e18c2..18baf5d 100644 --- a/merkletree_test.go +++ b/merkletree_test.go @@ -449,3 +449,27 @@ func TestDeleteNonExistingKeys(t *testing.T) { err = mt.Delete(big.NewInt(33)) assert.Equal(t, ErrKeyNotFound, err) } + +func TestDumpLeafsImportLeafs(t *testing.T) { + mt, err := NewMerkleTree(db.NewMemoryStorage(), 140) + require.Nil(t, err) + defer mt.db.Close() + + for i := 0; i < 10; i++ { + k := big.NewInt(int64(i)) + v := big.NewInt(0) + err = mt.Add(k, v) + require.Nil(t, err) + } + + d, err := mt.DumpLeafs(nil) + assert.Nil(t, err) + + mt2, err := NewMerkleTree(db.NewMemoryStorage(), 140) + require.Nil(t, err) + defer mt2.db.Close() + err = mt2.ImportDumpedLeafs(d) + assert.Nil(t, err) + + assert.Equal(t, mt.Root(), mt2.Root()) +} diff --git a/node.go b/node.go index 57ea09a..4ed6749 100644 --- a/node.go +++ b/node.go @@ -11,7 +11,7 @@ type NodeType byte const ( // NodeTypeMiddle indicates the type of middle Node that has children. NodeTypeMiddle NodeType = 0 - // NodeTypeLeaf indicates the type of a leaf Node that contains a claim. + // NodeTypeLeaf indicates the type of a leaf Node that contains a key & value. NodeTypeLeaf NodeType = 1 // NodeTypeEmpty indicates the type of an empty Node. NodeTypeEmpty NodeType = 2 diff --git a/utils.go b/utils.go index 17b3b60..c4dee94 100644 --- a/utils.go +++ b/utils.go @@ -4,7 +4,6 @@ import ( "fmt" "math/big" - cryptoConstants "github.com/iden3/go-iden3-crypto/constants" "github.com/iden3/go-iden3-crypto/poseidon" ) @@ -24,9 +23,6 @@ func HashElems(elems ...*big.Int) (*Hash, error) { poseidonHash, err := poseidon.PoseidonHash(bi) if err != nil { - fmt.Println("ERR HashElems PoseidonHash") - fmt.Println("e", bi[0]) - fmt.Println("q", cryptoConstants.Q) return nil, err } return NewHashFromBigInt(poseidonHash), nil