1package merkletree
  2
  3import (
  4	"encoding/hex"
  5	"fmt"
  6	"math/big"
  7	"testing"
  8
  9	"github.com/iden3/go-iden3-core/db"
 10	"github.com/stretchr/testify/assert"
 11	"github.com/stretchr/testify/require"
 12)
 13
 14var debug = false
 15
 16type Fatalable interface {
 17	Fatal(args ...interface{})
 18}
 19
 20func newTestingMerkle(f Fatalable, numLevels int) *MerkleTree {
 21	mt, err := NewMerkleTree(db.NewMemoryStorage(), numLevels)
 22	if err != nil {
 23		f.Fatal(err)
 24		return nil
 25	}
 26	return mt
 27}
 28
 29func TestNewTree(t *testing.T) {
 30	mt, err := NewMerkleTree(db.NewMemoryStorage(), 10)
 31	assert.Nil(t, err)
 32	assert.Equal(t, "0", mt.Root().String())
 33
 34	// test vectors generated using https://github.com/iden3/circomlib smt.js
 35	err = mt.Add(big.NewInt(1), big.NewInt(2))
 36	assert.Nil(t, err)
 37	assert.Equal(t, "4932297968297298434239270129193057052722409868268166443802652458940273154854", mt.Root().BigInt().String())
 38
 39	err = mt.Add(big.NewInt(33), big.NewInt(44))
 40	assert.Nil(t, err)
 41	assert.Equal(t, "13563340744765267202993741297198970774200042973817962221376874695587906013050", mt.Root().BigInt().String())
 42
 43	err = mt.Add(big.NewInt(1234), big.NewInt(9876))
 44	assert.Nil(t, err)
 45	assert.Equal(t, "16970503620176669663662021947486532860010370357132361783766545149750777353066", mt.Root().BigInt().String())
 46
 47	proof, err := mt.GenerateProof(big.NewInt(33), nil)
 48	assert.Nil(t, err)
 49
 50	assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(33), big.NewInt(44)))
 51	assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(33), big.NewInt(45)))
 52}
 53
 54func TestAddDifferentOrder(t *testing.T) {
 55	mt1 := newTestingMerkle(t, 140)
 56	defer mt1.db.Close()
 57	for i := 0; i < 16; i++ {
 58		k := big.NewInt(int64(i))
 59		v := big.NewInt(0)
 60		if err := mt1.Add(k, v); err != nil {
 61			t.Fatal(err)
 62		}
 63	}
 64
 65	mt2 := newTestingMerkle(t, 140)
 66	defer mt2.db.Close()
 67	for i := 16 - 1; i >= 0; i-- {
 68		k := big.NewInt(int64(i))
 69		v := big.NewInt(0)
 70		if err := mt2.Add(k, v); err != nil {
 71			t.Fatal(err)
 72		}
 73	}
 74
 75	assert.Equal(t, mt1.Root().Hex(), mt2.Root().Hex())
 76	assert.Equal(t, "ee3b91f9df12d26c1430b9d6693d4d5062b95b40f3f0a0a34ae560d677b76709", mt1.Root().Hex())
 77}
 78
 79func TestAddRepeatedIndex(t *testing.T) {
 80	mt := newTestingMerkle(t, 140)
 81	defer mt.db.Close()
 82	k := big.NewInt(int64(3))
 83	v := big.NewInt(int64(12))
 84	if err := mt.Add(k, v); err != nil {
 85		t.Fatal(err)
 86	}
 87	err := mt.Add(k, v)
 88	assert.NotNil(t, err)
 89	assert.Equal(t, err, ErrEntryIndexAlreadyExists)
 90}
 91
 92func TestGenerateAndVerifyProof128(t *testing.T) {
 93	mt, err := NewMerkleTree(db.NewMemoryStorage(), 140)
 94	require.Nil(t, err)
 95	defer mt.db.Close()
 96
 97	for i := 0; i < 128; i++ {
 98		k := big.NewInt(int64(i))
 99		v := big.NewInt(0)
100		if err := mt.Add(k, v); err != nil {
101			t.Fatal(err)
102		}
103	}
104	proof, err := mt.GenerateProof(big.NewInt(42), nil)
105	assert.Nil(t, err)
106	assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(42), big.NewInt(0)))
107}
108
109func TestTreeLimit(t *testing.T) {
110	mt, err := NewMerkleTree(db.NewMemoryStorage(), 5)
111	require.Nil(t, err)
112	defer mt.db.Close()
113
114	for i := 0; i < 16; i++ {
115		err = mt.Add(big.NewInt(int64(i)), big.NewInt(int64(i)))
116		assert.Nil(t, err)
117	}
118
119	// here the tree is full, should not allow to add more data as reaches the maximum number of levels
120	err = mt.Add(big.NewInt(int64(16)), big.NewInt(int64(16)))
121	assert.NotNil(t, err)
122	assert.Equal(t, ErrReachedMaxLevel, err)
123}
124
125func TestSiblingsFromProof(t *testing.T) {
126	mt, err := NewMerkleTree(db.NewMemoryStorage(), 140)
127	require.Nil(t, err)
128	defer mt.db.Close()
129
130	for i := 0; i < 64; i++ {
131		k := big.NewInt(int64(i))
132		v := big.NewInt(0)
133		if err := mt.Add(k, v); err != nil {
134			t.Fatal(err)
135		}
136	}
137
138	proof, err := mt.GenerateProof(big.NewInt(4), nil)
139	if err != nil {
140		t.Fatal(err)
141	}
142
143	siblings := SiblingsFromProof(proof)
144	assert.Equal(t, 6, len(siblings))
145	assert.Equal(t, "26bc69cfd3c982eba7b45cd2e6a2c75f218c546089f115777df47ab06f1fdb23", siblings[0].Hex())
146	assert.Equal(t, "ba7fbb5841bc8fa65193c124fbda4843aaa8d64d4132c79e7176cbed4de65621", siblings[1].Hex())
147	assert.Equal(t, "1339329b2d15e467ec3734ec06f6f4ae5a7c8c0b6ba98c26558b5a4db3e9a804", siblings[2].Hex())
148	assert.Equal(t, "b2261e6f47d9feb6fac2480928cc0fe03204d946cbc0a7b4de250d3e1384f40f", siblings[3].Hex())
149	assert.Equal(t, "b6c905f21c9928efa19a2c4e55d88d5fd0af493032f5b84620eb872e48ff5d01", siblings[4].Hex())
150	assert.Equal(t, "69873a951f49bbaff71039d672ffe52d73605aed6b40c1ce7ab068ad86a44d1e", siblings[5].Hex())
151}
152
153func TestVerifyProofCases(t *testing.T) {
154	mt := newTestingMerkle(t, 140)
155	defer mt.DB().Close()
156
157	for i := 0; i < 8; i++ {
158		if err := mt.Add(big.NewInt(int64(i)), big.NewInt(0)); err != nil {
159			t.Fatal(err)
160		}
161	}
162
163	// Existence proof
164
165	proof, err := mt.GenerateProof(big.NewInt(4), nil)
166	if err != nil {
167		t.Fatal(err)
168	}
169	assert.Equal(t, proof.Existence, true)
170	assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(4), big.NewInt(0)))
171	assert.Equal(t, "000300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b16621630b9104d995843c7cffa685009a1b28dcd371022a3b27b3a4d6987f7c8b39b0f2fffc165330710754ca0fc24451bdd5d5f82a05f42f1427fbdf17879c0b84be60f", hex.EncodeToString(proof.Bytes()))
172
173	for i := 8; i < 32; i++ {
174		proof, err = mt.GenerateProof(big.NewInt(int64(i)), nil)
175		assert.Nil(t, err)
176		if debug {
177			fmt.Println(i, proof)
178		}
179	}
180	// Non-existence proof, empty aux
181	proof, err = mt.GenerateProof(big.NewInt(12), nil)
182	if err != nil {
183		t.Fatal(err)
184	}
185	assert.Equal(t, proof.Existence, false)
186	// assert.True(t, proof.nodeAux == nil)
187	assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(12), big.NewInt(0)))
188	assert.Equal(t, "030300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b16621630b9104d995843c7cffa685009a1b28dcd371022a3b27b3a4d6987f7c8b39b0f2fffc165330710754ca0fc24451bdd5d5f82a05f42f1427fbdf17879c0b84be60f04000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", hex.EncodeToString(proof.Bytes()))
189
190	// Non-existence proof, diff. node aux
191	proof, err = mt.GenerateProof(big.NewInt(10), nil)
192	if err != nil {
193		t.Fatal(err)
194	}
195	assert.Equal(t, proof.Existence, false)
196	assert.True(t, proof.NodeAux != nil)
197	assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(10), big.NewInt(0)))
198	assert.Equal(t, "030300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b1662163097fcf8f911b271df196e0a75667b8a4f3024ef39f87201ed2b7cda349ba202296b7aeba35dc19ab0d4f65e175536c9952a90b6de18c3205611c3cd4fb408f01602000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", hex.EncodeToString(proof.Bytes()))
199}
200
201func TestVerifyProofFalse(t *testing.T) {
202	mt := newTestingMerkle(t, 140)
203	defer mt.DB().Close()
204
205	for i := 0; i < 8; i++ {
206		if err := mt.Add(big.NewInt(int64(i)), big.NewInt(0)); err != nil {
207			t.Fatal(err)
208		}
209	}
210
211	// Invalid existence proof (node used for verification doesn't
212	// correspond to node in the proof)
213	proof, err := mt.GenerateProof(big.NewInt(int64(4)), nil)
214	if err != nil {
215		t.Fatal(err)
216	}
217	assert.Equal(t, proof.Existence, true)
218	assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(int64(5)), big.NewInt(int64(5))))
219
220	// Invalid non-existence proof (Non-existence proof, diff. node aux)
221	proof, err = mt.GenerateProof(big.NewInt(int64(4)), nil)
222	if err != nil {
223		t.Fatal(err)
224	}
225	assert.Equal(t, proof.Existence, true)
226	// Now we change the proof from existence to non-existence, and add e's
227	// data as auxiliary node.
228	proof.Existence = false
229	proof.NodeAux = &NodeAux{Key: NewHashFromBigInt(big.NewInt(int64(4))), Value: NewHashFromBigInt(big.NewInt(4))}
230	assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(int64(4)), big.NewInt(0)))
231}