From 532de89fa7c364acd66407baef7637446685d5cc Mon Sep 17 00:00:00 2001 From: arnaucube Date: Wed, 1 Jul 2020 15:35:07 +0200 Subject: [PATCH] Add more tests, update README.md --- README.md | 32 +++++++++ go.mod | 2 +- merkletree.go | 5 ++ merkletree_test.go | 168 +++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 206 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 44837b1..a28475c 100644 --- a/README.md +++ b/README.md @@ -3,3 +3,35 @@ MerkleTree compatible with version from [circomlib](https://github.com/iden3/circomlib). Adaptation of the merkletree from https://github.com/iden3/go-iden3-core/tree/v0.0.8 + +## Usage +More detailed examples can be found at the [tests](https://github.com/iden3/go-merkletree/blob/master/merkletree_test.go), and in the [documentation](https://godoc.org/github.com/iden3/go-merkletree). + +```go +import ( + "fmt" + "math/big" + "testing" + + "github.com/iden3/go-iden3-core/db" + "github.com/stretchr/testify/assert" +) + +[...] + +func TestExampleMerkleTree(t *testing.T) { + mt, err := NewMerkleTree(db.NewMemoryStorage(), 10) + assert.Nil(t, err) + + key := big.NewInt(1) + value := big.NewInt(2) + err = mt.Add(key, value) + assert.Nil(t, err) + fmt.Println(mt.Root().String()) + + proof, err := mt.GenerateProof(key, nil) + assert.Nil(t, err) + + assert.True(t, VerifyProof(mt.Root(), proof, key, value)) +} +``` diff --git a/go.mod b/go.mod index dfb1805..d3ce793 100644 --- a/go.mod +++ b/go.mod @@ -1,4 +1,4 @@ -module go-merkletree +module github.com/iden3/go-merkletree go 1.14 diff --git a/merkletree.go b/merkletree.go index 659bec4..0f3f1b1 100644 --- a/merkletree.go +++ b/merkletree.go @@ -106,6 +106,11 @@ func NewMerkleTree(storage db.Storage, maxLevels int) (*MerkleTree, error) { return &mt, nil } +// DB returns the MerkleTree.DB() +func (mt *MerkleTree) DB() db.Storage { + return mt.db +} + // Root returns the MerkleRoot func (mt *MerkleTree) Root() *Hash { return mt.rootKey diff --git a/merkletree_test.go b/merkletree_test.go index 27224c0..ad922c0 100644 --- a/merkletree_test.go +++ b/merkletree_test.go @@ -1,6 +1,8 @@ package merkletree import ( + "encoding/hex" + "fmt" "math/big" "testing" @@ -9,6 +11,21 @@ import ( "github.com/stretchr/testify/require" ) +var debug = false + +type Fatalable interface { + Fatal(args ...interface{}) +} + +func newTestingMerkle(f Fatalable, numLevels int) *MerkleTree { + mt, err := NewMerkleTree(db.NewMemoryStorage(), numLevels) + if err != nil { + f.Fatal(err) + return nil + } + return mt +} + func TestNewTree(t *testing.T) { mt, err := NewMerkleTree(db.NewMemoryStorage(), 10) assert.Nil(t, err) @@ -34,6 +51,77 @@ func TestNewTree(t *testing.T) { assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(33), big.NewInt(45))) } +func TestAddDifferentOrder(t *testing.T) { + mt1 := newTestingMerkle(t, 140) + defer mt1.db.Close() + for i := 0; i < 16; i++ { + k := big.NewInt(int64(i)) + v := big.NewInt(0) + if err := mt1.Add(k, v); err != nil { + t.Fatal(err) + } + } + + mt2 := newTestingMerkle(t, 140) + defer mt2.db.Close() + for i := 16 - 1; i >= 0; i-- { + k := big.NewInt(int64(i)) + v := big.NewInt(0) + if err := mt2.Add(k, v); err != nil { + t.Fatal(err) + } + } + + assert.Equal(t, mt1.Root().Hex(), mt2.Root().Hex()) + assert.Equal(t, "ee3b91f9df12d26c1430b9d6693d4d5062b95b40f3f0a0a34ae560d677b76709", mt1.Root().Hex()) +} + +func TestAddRepeatedIndex(t *testing.T) { + mt := newTestingMerkle(t, 140) + defer mt.db.Close() + k := big.NewInt(int64(3)) + v := big.NewInt(int64(12)) + if err := mt.Add(k, v); err != nil { + t.Fatal(err) + } + err := mt.Add(k, v) + assert.NotNil(t, err) + assert.Equal(t, err, ErrEntryIndexAlreadyExists) +} + +func TestGenerateAndVerifyProof128(t *testing.T) { + mt, err := NewMerkleTree(db.NewMemoryStorage(), 140) + require.Nil(t, err) + defer mt.db.Close() + + for i := 0; i < 128; i++ { + k := big.NewInt(int64(i)) + v := big.NewInt(0) + if err := mt.Add(k, v); err != nil { + t.Fatal(err) + } + } + proof, err := mt.GenerateProof(big.NewInt(42), nil) + assert.Nil(t, err) + assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(42), big.NewInt(0))) +} + +func TestTreeLimit(t *testing.T) { + mt, err := NewMerkleTree(db.NewMemoryStorage(), 5) + require.Nil(t, err) + defer mt.db.Close() + + for i := 0; i < 16; i++ { + err = mt.Add(big.NewInt(int64(i)), big.NewInt(int64(i))) + assert.Nil(t, err) + } + + // here the tree is full, should not allow to add more data as reaches the maximum number of levels + err = mt.Add(big.NewInt(int64(16)), big.NewInt(int64(16))) + assert.NotNil(t, err) + assert.Equal(t, ErrReachedMaxLevel, err) +} + func TestSiblingsFromProof(t *testing.T) { mt, err := NewMerkleTree(db.NewMemoryStorage(), 140) require.Nil(t, err) @@ -61,3 +149,83 @@ func TestSiblingsFromProof(t *testing.T) { assert.Equal(t, "b6c905f21c9928efa19a2c4e55d88d5fd0af493032f5b84620eb872e48ff5d01", siblings[4].Hex()) assert.Equal(t, "69873a951f49bbaff71039d672ffe52d73605aed6b40c1ce7ab068ad86a44d1e", siblings[5].Hex()) } + +func TestVerifyProofCases(t *testing.T) { + mt := newTestingMerkle(t, 140) + defer mt.DB().Close() + + for i := 0; i < 8; i++ { + if err := mt.Add(big.NewInt(int64(i)), big.NewInt(0)); err != nil { + t.Fatal(err) + } + } + + // Existence proof + + proof, err := mt.GenerateProof(big.NewInt(4), nil) + if err != nil { + t.Fatal(err) + } + assert.Equal(t, proof.Existence, true) + assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(4), big.NewInt(0))) + assert.Equal(t, "000300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b16621630b9104d995843c7cffa685009a1b28dcd371022a3b27b3a4d6987f7c8b39b0f2fffc165330710754ca0fc24451bdd5d5f82a05f42f1427fbdf17879c0b84be60f", hex.EncodeToString(proof.Bytes())) + + for i := 8; i < 32; i++ { + proof, err = mt.GenerateProof(big.NewInt(int64(i)), nil) + assert.Nil(t, err) + if debug { + fmt.Println(i, proof) + } + } + // Non-existence proof, empty aux + proof, err = mt.GenerateProof(big.NewInt(12), nil) + if err != nil { + t.Fatal(err) + } + assert.Equal(t, proof.Existence, false) + // assert.True(t, proof.nodeAux == nil) + assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(12), big.NewInt(0))) + assert.Equal(t, "030300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b16621630b9104d995843c7cffa685009a1b28dcd371022a3b27b3a4d6987f7c8b39b0f2fffc165330710754ca0fc24451bdd5d5f82a05f42f1427fbdf17879c0b84be60f04000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", hex.EncodeToString(proof.Bytes())) + + // Non-existence proof, diff. node aux + proof, err = mt.GenerateProof(big.NewInt(10), nil) + if err != nil { + t.Fatal(err) + } + assert.Equal(t, proof.Existence, false) + assert.True(t, proof.NodeAux != nil) + assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(10), big.NewInt(0))) + assert.Equal(t, "030300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b1662163097fcf8f911b271df196e0a75667b8a4f3024ef39f87201ed2b7cda349ba202296b7aeba35dc19ab0d4f65e175536c9952a90b6de18c3205611c3cd4fb408f01602000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", hex.EncodeToString(proof.Bytes())) +} + +func TestVerifyProofFalse(t *testing.T) { + mt := newTestingMerkle(t, 140) + defer mt.DB().Close() + + for i := 0; i < 8; i++ { + if err := mt.Add(big.NewInt(int64(i)), big.NewInt(0)); err != nil { + t.Fatal(err) + } + } + + // Invalid existence proof (node used for verification doesn't + // correspond to node in the proof) + proof, err := mt.GenerateProof(big.NewInt(int64(4)), nil) + if err != nil { + t.Fatal(err) + } + assert.Equal(t, proof.Existence, true) + assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(int64(5)), big.NewInt(int64(5)))) + + // Invalid non-existence proof (Non-existence proof, diff. node aux) + proof, err = mt.GenerateProof(big.NewInt(int64(4)), nil) + if err != nil { + t.Fatal(err) + } + assert.Equal(t, proof.Existence, true) + // Now we change the proof from existence to non-existence, and add e's + // data as auxiliary node. + proof.Existence = false + proof.NodeAux = &NodeAux{Key: NewHashFromBigInt(big.NewInt(int64(4))), Value: NewHashFromBigInt(big.NewInt(4))} + assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(int64(4)), big.NewInt(0))) +}