arnaucube
arnaucube/go-merkletree-iden3Public Self-hosted repository
go-merkletree-iden3/merkletree_test.go
merkletree_test.goRaw
  1package merkletree
  2
  3import (
  4	"bytes"
  5	"encoding/hex"
  6	"fmt"
  7	"math/big"
  8	"testing"
  9
 10	"github.com/iden3/go-merkletree/db/memory"
 11	"github.com/stretchr/testify/assert"
 12	"github.com/stretchr/testify/require"
 13)
 14
 15var debug = false
 16
 17type Fatalable interface {
 18	Fatal(args ...interface{})
 19}
 20
 21func newTestingMerkle(f Fatalable, numLevels int) *MerkleTree {
 22	mt, err := NewMerkleTree(memory.NewMemoryStorage(), numLevels)
 23	if err != nil {
 24		f.Fatal(err)
 25		return nil
 26	}
 27	return mt
 28}
 29
 30func TestHashParsers(t *testing.T) {
 31	h0 := NewHashFromBigInt(big.NewInt(0))
 32	assert.Equal(t, "0", h0.String())
 33	h1 := NewHashFromBigInt(big.NewInt(1))
 34	assert.Equal(t, "1", h1.String())
 35	h10 := NewHashFromBigInt(big.NewInt(10))
 36	assert.Equal(t, "10", h10.String())
 37
 38	h7l := NewHashFromBigInt(big.NewInt(1234567))
 39	assert.Equal(t, "1234567", h7l.String())
 40	h8l := NewHashFromBigInt(big.NewInt(12345678))
 41	assert.Equal(t, "12345678...", h8l.String())
 42
 43	b, ok := new(big.Int).SetString("4932297968297298434239270129193057052722409868268166443802652458940273154854", 10)
 44	assert.True(t, ok)
 45	h := NewHashFromBigInt(b)
 46	assert.Equal(t, "4932297968297298434239270129193057052722409868268166443802652458940273154854", h.BigInt().String())
 47	assert.Equal(t, "49322979...", h.String())
 48	assert.Equal(t, "0ae794eb9c3d8bbb9002e993fc2ed301dcbd2af5508ed072c375e861f1aa5b26", h.Hex())
 49}
 50
 51func TestNewTree(t *testing.T) {
 52	mt, err := NewMerkleTree(memory.NewMemoryStorage(), 10)
 53	assert.Nil(t, err)
 54	assert.Equal(t, "0", mt.Root().String())
 55
 56	// test vectors generated using https://github.com/iden3/circomlib smt.js
 57	err = mt.Add(big.NewInt(1), big.NewInt(2))
 58	assert.Nil(t, err)
 59	assert.Equal(t, "4932297968297298434239270129193057052722409868268166443802652458940273154854", mt.Root().BigInt().String())
 60
 61	err = mt.Add(big.NewInt(33), big.NewInt(44))
 62	assert.Nil(t, err)
 63	assert.Equal(t, "13563340744765267202993741297198970774200042973817962221376874695587906013050", mt.Root().BigInt().String())
 64
 65	err = mt.Add(big.NewInt(1234), big.NewInt(9876))
 66	assert.Nil(t, err)
 67	assert.Equal(t, "16970503620176669663662021947486532860010370357132361783766545149750777353066", mt.Root().BigInt().String())
 68
 69	proof, err := mt.GenerateProof(big.NewInt(33), nil)
 70	assert.Nil(t, err)
 71
 72	assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(33), big.NewInt(44)))
 73	assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(33), big.NewInt(45)))
 74}
 75
 76func TestAddDifferentOrder(t *testing.T) {
 77	mt1 := newTestingMerkle(t, 140)
 78	defer mt1.db.Close()
 79	for i := 0; i < 16; i++ {
 80		k := big.NewInt(int64(i))
 81		v := big.NewInt(0)
 82		if err := mt1.Add(k, v); err != nil {
 83			t.Fatal(err)
 84		}
 85	}
 86
 87	mt2 := newTestingMerkle(t, 140)
 88	defer mt2.db.Close()
 89	for i := 16 - 1; i >= 0; i-- {
 90		k := big.NewInt(int64(i))
 91		v := big.NewInt(0)
 92		if err := mt2.Add(k, v); err != nil {
 93			t.Fatal(err)
 94		}
 95	}
 96
 97	assert.Equal(t, mt1.Root().Hex(), mt2.Root().Hex())
 98	assert.Equal(t, "0967b777d660e54aa3a0f0f3405bb962504d3d69d6b930146cd212dff9913bee", mt1.Root().Hex())
 99}
100
101func TestAddRepeatedIndex(t *testing.T) {
102	mt := newTestingMerkle(t, 140)
103	defer mt.db.Close()
104	k := big.NewInt(int64(3))
105	v := big.NewInt(int64(12))
106	if err := mt.Add(k, v); err != nil {
107		t.Fatal(err)
108	}
109	err := mt.Add(k, v)
110	assert.NotNil(t, err)
111	assert.Equal(t, err, ErrEntryIndexAlreadyExists)
112}
113
114func TestGet(t *testing.T) {
115	mt := newTestingMerkle(t, 140)
116	defer mt.db.Close()
117
118	for i := 0; i < 16; i++ {
119		k := big.NewInt(int64(i))
120		v := big.NewInt(int64(i * 2))
121		if err := mt.Add(k, v); err != nil {
122			t.Fatal(err)
123		}
124	}
125	v, err := mt.Get(big.NewInt(10))
126	assert.Nil(t, err)
127	assert.Equal(t, big.NewInt(20), v)
128
129	v, err = mt.Get(big.NewInt(15))
130	assert.Nil(t, err)
131	assert.Equal(t, big.NewInt(30), v)
132
133	v, err = mt.Get(big.NewInt(16))
134	assert.NotNil(t, err)
135	assert.Equal(t, ErrKeyNotFound, err)
136	assert.Nil(t, v)
137}
138
139func TestGenerateAndVerifyProof128(t *testing.T) {
140	mt, err := NewMerkleTree(memory.NewMemoryStorage(), 140)
141	require.Nil(t, err)
142	defer mt.db.Close()
143
144	for i := 0; i < 128; i++ {
145		k := big.NewInt(int64(i))
146		v := big.NewInt(0)
147		if err := mt.Add(k, v); err != nil {
148			t.Fatal(err)
149		}
150	}
151	proof, err := mt.GenerateProof(big.NewInt(42), nil)
152	assert.Nil(t, err)
153	assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(42), big.NewInt(0)))
154}
155
156func TestTreeLimit(t *testing.T) {
157	mt, err := NewMerkleTree(memory.NewMemoryStorage(), 5)
158	require.Nil(t, err)
159	defer mt.db.Close()
160
161	for i := 0; i < 16; i++ {
162		err = mt.Add(big.NewInt(int64(i)), big.NewInt(int64(i)))
163		assert.Nil(t, err)
164	}
165
166	// here the tree is full, should not allow to add more data as reaches the maximum number of levels
167	err = mt.Add(big.NewInt(int64(16)), big.NewInt(int64(16)))
168	assert.NotNil(t, err)
169	assert.Equal(t, ErrReachedMaxLevel, err)
170}
171
172func TestSiblingsFromProof(t *testing.T) {
173	mt, err := NewMerkleTree(memory.NewMemoryStorage(), 140)
174	require.Nil(t, err)
175	defer mt.db.Close()
176
177	for i := 0; i < 64; i++ {
178		k := big.NewInt(int64(i))
179		v := big.NewInt(0)
180		if err := mt.Add(k, v); err != nil {
181			t.Fatal(err)
182		}
183	}
184
185	proof, err := mt.GenerateProof(big.NewInt(4), nil)
186	if err != nil {
187		t.Fatal(err)
188	}
189
190	siblings := SiblingsFromProof(proof)
191	assert.Equal(t, 6, len(siblings))
192	assert.Equal(t, "23db1f6fb07af47d7715f18960548c215fc7a2e6d25cb4a7eb82c9d3cf69bc26", siblings[0].Hex())
193	assert.Equal(t, "2156e64dedcb76719ec732414dd6a8aa4348dafb24c19351a68fbc4158bb7fba", siblings[1].Hex())
194	assert.Equal(t, "04a8e9b34d5a8b55268ca96b0b8c7c5aaef4f606ec3437ec67e4152d9b323913", siblings[2].Hex())
195	assert.Equal(t, "0ff484133e0d25deb4a7c0cb46d90432e00fcc280948c2fab6fed9476f1e26b2", siblings[3].Hex())
196	assert.Equal(t, "015dff482e87eb2046b8f5323049afd05f8dd8554e2c9aa1ef28991cf205c9b6", siblings[4].Hex())
197	assert.Equal(t, "1e4da486ad68b07acec1406bed5a60732de5ff72d63910f7afbb491f953a8769", siblings[5].Hex())
198}
199
200func TestVerifyProofCases(t *testing.T) {
201	mt := newTestingMerkle(t, 140)
202	defer mt.DB().Close()
203
204	for i := 0; i < 8; i++ {
205		if err := mt.Add(big.NewInt(int64(i)), big.NewInt(0)); err != nil {
206			t.Fatal(err)
207		}
208	}
209
210	// Existence proof
211
212	proof, err := mt.GenerateProof(big.NewInt(4), nil)
213	if err != nil {
214		t.Fatal(err)
215	}
216	assert.Equal(t, proof.Existence, true)
217	assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(4), big.NewInt(0)))
218	assert.Equal(t, "000300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b16621630b9104d995843c7cffa685009a1b28dcd371022a3b27b3a4d6987f7c8b39b0f2fffc165330710754ca0fc24451bdd5d5f82a05f42f1427fbdf17879c0b84be60f", hex.EncodeToString(proof.Bytes()))
219
220	for i := 8; i < 32; i++ {
221		proof, err = mt.GenerateProof(big.NewInt(int64(i)), nil)
222		assert.Nil(t, err)
223		if debug {
224			fmt.Println(i, proof)
225		}
226	}
227	// Non-existence proof, empty aux
228	proof, err = mt.GenerateProof(big.NewInt(12), nil)
229	if err != nil {
230		t.Fatal(err)
231	}
232	assert.Equal(t, proof.Existence, false)
233	// assert.True(t, proof.nodeAux == nil)
234	assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(12), big.NewInt(0)))
235	assert.Equal(t, "030300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b16621630b9104d995843c7cffa685009a1b28dcd371022a3b27b3a4d6987f7c8b39b0f2fffc165330710754ca0fc24451bdd5d5f82a05f42f1427fbdf17879c0b84be60f04000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", hex.EncodeToString(proof.Bytes()))
236
237	// Non-existence proof, diff. node aux
238	proof, err = mt.GenerateProof(big.NewInt(10), nil)
239	if err != nil {
240		t.Fatal(err)
241	}
242	assert.Equal(t, proof.Existence, false)
243	assert.True(t, proof.NodeAux != nil)
244	assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(10), big.NewInt(0)))
245	assert.Equal(t, "030300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b1662163097fcf8f911b271df196e0a75667b8a4f3024ef39f87201ed2b7cda349ba202296b7aeba35dc19ab0d4f65e175536c9952a90b6de18c3205611c3cd4fb408f01602000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", hex.EncodeToString(proof.Bytes()))
246}
247
248func TestVerifyProofFalse(t *testing.T) {
249	mt := newTestingMerkle(t, 140)
250	defer mt.DB().Close()
251
252	for i := 0; i < 8; i++ {
253		if err := mt.Add(big.NewInt(int64(i)), big.NewInt(0)); err != nil {
254			t.Fatal(err)
255		}
256	}
257
258	// Invalid existence proof (node used for verification doesn't
259	// correspond to node in the proof)
260	proof, err := mt.GenerateProof(big.NewInt(int64(4)), nil)
261	if err != nil {
262		t.Fatal(err)
263	}
264	assert.Equal(t, proof.Existence, true)
265	assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(int64(5)), big.NewInt(int64(5))))
266
267	// Invalid non-existence proof (Non-existence proof, diff. node aux)
268	proof, err = mt.GenerateProof(big.NewInt(int64(4)), nil)
269	if err != nil {
270		t.Fatal(err)
271	}
272	assert.Equal(t, proof.Existence, true)
273	// Now we change the proof from existence to non-existence, and add e's
274	// data as auxiliary node.
275	proof.Existence = false
276	proof.NodeAux = &NodeAux{Key: NewHashFromBigInt(big.NewInt(int64(4))), Value: NewHashFromBigInt(big.NewInt(4))}
277	assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(int64(4)), big.NewInt(0)))
278}
279
280func TestGraphViz(t *testing.T) {
281	mt, err := NewMerkleTree(memory.NewMemoryStorage(), 10)
282	assert.Nil(t, err)
283
284	mt.Add(big.NewInt(1), big.NewInt(0))
285	mt.Add(big.NewInt(2), big.NewInt(0))
286	mt.Add(big.NewInt(3), big.NewInt(0))
287	mt.Add(big.NewInt(4), big.NewInt(0))
288	mt.Add(big.NewInt(5), big.NewInt(0))
289	mt.Add(big.NewInt(100), big.NewInt(0))
290
291	// mt.PrintGraphViz(nil)
292
293	expected := `digraph hierarchy {
294node [fontname=Monospace,fontsize=10,shape=box]
295"60195538..." -> {"19759736..." "18893277..."}
296"19759736..." -> {"16152312..." "43945008..."}
297"16152312..." -> {"empty0" "13952255..."}
298"empty0" [style=dashed,label=0];
299"13952255..." -> {"61769925..." "empty1"}
300"empty1" [style=dashed,label=0];
301"61769925..." -> {"92723289..." "empty2"}
302"empty2" [style=dashed,label=0];
303"92723289..." -> {"21082735..." "82784818..."}
304"21082735..." [style=filled];
305"82784818..." [style=filled];
306"43945008..." [style=filled];
307"18893277..." -> {"19855703..." "17718670..."}
308"19855703..." -> {"11499909..." "15828714..."}
309"11499909..." [style=filled];
310"15828714..." [style=filled];
311"17718670..." [style=filled];
312}
313`
314	w := bytes.NewBufferString("")
315	mt.GraphViz(w, nil)
316	assert.Equal(t, []byte(expected), w.Bytes())
317}
318
319func TestDelete(t *testing.T) {
320	mt, err := NewMerkleTree(memory.NewMemoryStorage(), 10)
321	assert.Nil(t, err)
322	assert.Equal(t, "0", mt.Root().String())
323
324	// test vectors generated using https://github.com/iden3/circomlib smt.js
325	err = mt.Add(big.NewInt(1), big.NewInt(2))
326	assert.Nil(t, err)
327	assert.Equal(t, "4932297968297298434239270129193057052722409868268166443802652458940273154854", mt.Root().BigInt().String())
328
329	err = mt.Add(big.NewInt(33), big.NewInt(44))
330	assert.Nil(t, err)
331	assert.Equal(t, "13563340744765267202993741297198970774200042973817962221376874695587906013050", mt.Root().BigInt().String())
332
333	err = mt.Add(big.NewInt(1234), big.NewInt(9876))
334	assert.Nil(t, err)
335	assert.Equal(t, "16970503620176669663662021947486532860010370357132361783766545149750777353066", mt.Root().BigInt().String())
336
337	// mt.PrintGraphViz(nil)
338
339	err = mt.Delete(big.NewInt(33))
340	// mt.PrintGraphViz(nil)
341	assert.Nil(t, err)
342	assert.Equal(t, "12820263606494630162816839760750120928463716794691735985748071431547370997091", mt.Root().BigInt().String())
343
344	err = mt.Delete(big.NewInt(1234))
345	assert.Nil(t, err)
346	err = mt.Delete(big.NewInt(1))
347	assert.Nil(t, err)
348	assert.Equal(t, "0", mt.Root().String())
349}
350
351func TestDelete2(t *testing.T) {
352	mt := newTestingMerkle(t, 140)
353	defer mt.db.Close()
354	for i := 0; i < 8; i++ {
355		k := big.NewInt(int64(i))
356		v := big.NewInt(0)
357		if err := mt.Add(k, v); err != nil {
358			t.Fatal(err)
359		}
360	}
361
362	expectedRoot := mt.Root()
363
364	k := big.NewInt(8)
365	v := big.NewInt(0)
366	err := mt.Add(k, v)
367	require.Nil(t, err)
368
369	err = mt.Delete(big.NewInt(8))
370	assert.Nil(t, err)
371	assert.Equal(t, expectedRoot, mt.Root())
372
373	mt2 := newTestingMerkle(t, 140)
374	defer mt2.db.Close()
375	for i := 0; i < 8; i++ {
376		k := big.NewInt(int64(i))
377		v := big.NewInt(0)
378		if err := mt2.Add(k, v); err != nil {
379			t.Fatal(err)
380		}
381	}
382	assert.Equal(t, mt2.Root(), mt.Root())
383}
384
385func TestDelete3(t *testing.T) {
386	mt := newTestingMerkle(t, 140)
387	defer mt.db.Close()
388
389	err := mt.Add(big.NewInt(1), big.NewInt(1))
390	assert.Nil(t, err)
391
392	err = mt.Add(big.NewInt(2), big.NewInt(2))
393	assert.Nil(t, err)
394
395	assert.Equal(t, "2427629547967522489273866134471574861207714751136138191708011221765688788661", mt.Root().BigInt().String())
396	err = mt.Delete(big.NewInt(1))
397	assert.Nil(t, err)
398	assert.Equal(t, "10822920717809411688334493481050035035708810950159417482558569847174767667301", mt.Root().BigInt().String())
399
400	mt2 := newTestingMerkle(t, 140)
401	defer mt2.db.Close()
402	err = mt2.Add(big.NewInt(2), big.NewInt(2))
403	assert.Nil(t, err)
404	assert.Equal(t, mt2.Root(), mt.Root())
405}
406
407func TestDelete4(t *testing.T) {
408	mt := newTestingMerkle(t, 140)
409	defer mt.db.Close()
410
411	err := mt.Add(big.NewInt(1), big.NewInt(1))
412	assert.Nil(t, err)
413
414	err = mt.Add(big.NewInt(2), big.NewInt(2))
415	assert.Nil(t, err)
416
417	err = mt.Add(big.NewInt(3), big.NewInt(3))
418	assert.Nil(t, err)
419
420	assert.Equal(t, "16614298246517994771186095530428786749320098419259206061045083278756632941513", mt.Root().BigInt().String())
421	err = mt.Delete(big.NewInt(1))
422	assert.Nil(t, err)
423	assert.Equal(t, "6117330520107511783353383870014397665359816230889739699667943862706617498952", mt.Root().BigInt().String())
424
425	mt2 := newTestingMerkle(t, 140)
426	defer mt2.db.Close()
427	err = mt2.Add(big.NewInt(2), big.NewInt(2))
428	assert.Nil(t, err)
429	err = mt2.Add(big.NewInt(3), big.NewInt(3))
430	assert.Nil(t, err)
431	assert.Equal(t, mt2.Root(), mt.Root())
432}
433
434func TestDelete5(t *testing.T) {
435	mt, err := NewMerkleTree(memory.NewMemoryStorage(), 10)
436	assert.Nil(t, err)
437
438	err = mt.Add(big.NewInt(1), big.NewInt(2))
439	assert.Nil(t, err)
440	err = mt.Add(big.NewInt(33), big.NewInt(44))
441	assert.Nil(t, err)
442	assert.Equal(t, "13563340744765267202993741297198970774200042973817962221376874695587906013050", mt.Root().BigInt().String())
443
444	err = mt.Delete(big.NewInt(1))
445	assert.Nil(t, err)
446	assert.Equal(t, "12075524681474630909546786277734445038384732558409197537058769521806571391765", mt.Root().BigInt().String())
447
448	mt2 := newTestingMerkle(t, 140)
449	defer mt2.db.Close()
450	err = mt2.Add(big.NewInt(33), big.NewInt(44))
451	assert.Nil(t, err)
452	assert.Equal(t, mt2.Root(), mt.Root())
453}
454
455func TestDeleteNonExistingKeys(t *testing.T) {
456	mt, err := NewMerkleTree(memory.NewMemoryStorage(), 10)
457	assert.Nil(t, err)
458
459	err = mt.Add(big.NewInt(1), big.NewInt(2))
460	assert.Nil(t, err)
461	err = mt.Add(big.NewInt(33), big.NewInt(44))
462	assert.Nil(t, err)
463
464	err = mt.Delete(big.NewInt(33))
465	assert.Nil(t, err)
466	err = mt.Delete(big.NewInt(33))
467	assert.Equal(t, ErrKeyNotFound, err)
468
469	err = mt.Delete(big.NewInt(1))
470	assert.Nil(t, err)
471
472	assert.Equal(t, "0", mt.Root().String())
473
474	err = mt.Delete(big.NewInt(33))
475	assert.Equal(t, ErrKeyNotFound, err)
476}
477
478func TestDumpLeafsImportLeafs(t *testing.T) {
479	mt, err := NewMerkleTree(memory.NewMemoryStorage(), 140)
480	require.Nil(t, err)
481	defer mt.db.Close()
482
483	for i := 0; i < 10; i++ {
484		k := big.NewInt(int64(i))
485		v := big.NewInt(0)
486		err = mt.Add(k, v)
487		require.Nil(t, err)
488	}
489
490	d, err := mt.DumpLeafs(nil)
491	assert.Nil(t, err)
492
493	mt2, err := NewMerkleTree(memory.NewMemoryStorage(), 140)
494	require.Nil(t, err)
495	defer mt2.db.Close()
496	err = mt2.ImportDumpedLeafs(d)
497	assert.Nil(t, err)
498
499	assert.Equal(t, mt.Root(), mt2.Root())
500}