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 TestUpdate(t *testing.T) {
140 mt := newTestingMerkle(t, 140)
141 defer mt.db.Close()
142
143 for i := 0; i < 16; i++ {
144 k := big.NewInt(int64(i))
145 v := big.NewInt(int64(i * 2))
146 if err := mt.Add(k, v); err != nil {
147 t.Fatal(err)
148 }
149 }
150 v, err := mt.Get(big.NewInt(10))
151 assert.Nil(t, err)
152 assert.Equal(t, big.NewInt(20), v)
153
154 err = mt.Update(big.NewInt(10), big.NewInt(1024))
155 assert.Nil(t, err)
156 v, err = mt.Get(big.NewInt(10))
157 assert.Nil(t, err)
158 assert.Equal(t, big.NewInt(1024), v)
159
160 err = mt.Update(big.NewInt(1000), big.NewInt(1024))
161 assert.Equal(t, ErrKeyNotFound, err)
162}
163
164func TestUpdate2(t *testing.T) {
165 mt1 := newTestingMerkle(t, 140)
166 defer mt1.db.Close()
167 mt2 := newTestingMerkle(t, 140)
168 defer mt2.db.Close()
169
170 err := mt1.Add(big.NewInt(1), big.NewInt(119))
171 assert.Nil(t, err)
172 err = mt1.Add(big.NewInt(2), big.NewInt(229))
173 assert.Nil(t, err)
174 err = mt1.Add(big.NewInt(9876), big.NewInt(6789))
175 assert.Nil(t, err)
176
177 err = mt2.Add(big.NewInt(1), big.NewInt(11))
178 assert.Nil(t, err)
179 err = mt2.Add(big.NewInt(2), big.NewInt(22))
180 assert.Nil(t, err)
181 err = mt2.Add(big.NewInt(9876), big.NewInt(10))
182 assert.Nil(t, err)
183
184 err = mt1.Update(big.NewInt(1), big.NewInt(11))
185 assert.Nil(t, err)
186 err = mt1.Update(big.NewInt(2), big.NewInt(22))
187 assert.Nil(t, err)
188 err = mt2.Update(big.NewInt(9876), big.NewInt(6789))
189 assert.Nil(t, err)
190
191 assert.Equal(t, mt1.Root(), mt2.Root())
192}
193
194func TestGenerateAndVerifyProof128(t *testing.T) {
195 mt, err := NewMerkleTree(memory.NewMemoryStorage(), 140)
196 require.Nil(t, err)
197 defer mt.db.Close()
198
199 for i := 0; i < 128; i++ {
200 k := big.NewInt(int64(i))
201 v := big.NewInt(0)
202 if err := mt.Add(k, v); err != nil {
203 t.Fatal(err)
204 }
205 }
206 proof, err := mt.GenerateProof(big.NewInt(42), nil)
207 assert.Nil(t, err)
208 assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(42), big.NewInt(0)))
209}
210
211func TestTreeLimit(t *testing.T) {
212 mt, err := NewMerkleTree(memory.NewMemoryStorage(), 5)
213 require.Nil(t, err)
214 defer mt.db.Close()
215
216 for i := 0; i < 16; i++ {
217 err = mt.Add(big.NewInt(int64(i)), big.NewInt(int64(i)))
218 assert.Nil(t, err)
219 }
220
221 // here the tree is full, should not allow to add more data as reaches the maximum number of levels
222 err = mt.Add(big.NewInt(int64(16)), big.NewInt(int64(16)))
223 assert.NotNil(t, err)
224 assert.Equal(t, ErrReachedMaxLevel, err)
225}
226
227func TestSiblingsFromProof(t *testing.T) {
228 mt, err := NewMerkleTree(memory.NewMemoryStorage(), 140)
229 require.Nil(t, err)
230 defer mt.db.Close()
231
232 for i := 0; i < 64; i++ {
233 k := big.NewInt(int64(i))
234 v := big.NewInt(0)
235 if err := mt.Add(k, v); err != nil {
236 t.Fatal(err)
237 }
238 }
239
240 proof, err := mt.GenerateProof(big.NewInt(4), nil)
241 if err != nil {
242 t.Fatal(err)
243 }
244
245 siblings := SiblingsFromProof(proof)
246 assert.Equal(t, 6, len(siblings))
247 assert.Equal(t, "23db1f6fb07af47d7715f18960548c215fc7a2e6d25cb4a7eb82c9d3cf69bc26", siblings[0].Hex())
248 assert.Equal(t, "2156e64dedcb76719ec732414dd6a8aa4348dafb24c19351a68fbc4158bb7fba", siblings[1].Hex())
249 assert.Equal(t, "04a8e9b34d5a8b55268ca96b0b8c7c5aaef4f606ec3437ec67e4152d9b323913", siblings[2].Hex())
250 assert.Equal(t, "0ff484133e0d25deb4a7c0cb46d90432e00fcc280948c2fab6fed9476f1e26b2", siblings[3].Hex())
251 assert.Equal(t, "015dff482e87eb2046b8f5323049afd05f8dd8554e2c9aa1ef28991cf205c9b6", siblings[4].Hex())
252 assert.Equal(t, "1e4da486ad68b07acec1406bed5a60732de5ff72d63910f7afbb491f953a8769", siblings[5].Hex())
253}
254
255func TestVerifyProofCases(t *testing.T) {
256 mt := newTestingMerkle(t, 140)
257 defer mt.DB().Close()
258
259 for i := 0; i < 8; i++ {
260 if err := mt.Add(big.NewInt(int64(i)), big.NewInt(0)); err != nil {
261 t.Fatal(err)
262 }
263 }
264
265 // Existence proof
266
267 proof, err := mt.GenerateProof(big.NewInt(4), nil)
268 if err != nil {
269 t.Fatal(err)
270 }
271 assert.Equal(t, proof.Existence, true)
272 assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(4), big.NewInt(0)))
273 assert.Equal(t, "000300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b16621630b9104d995843c7cffa685009a1b28dcd371022a3b27b3a4d6987f7c8b39b0f2fffc165330710754ca0fc24451bdd5d5f82a05f42f1427fbdf17879c0b84be60f", hex.EncodeToString(proof.Bytes()))
274
275 for i := 8; i < 32; i++ {
276 proof, err = mt.GenerateProof(big.NewInt(int64(i)), nil)
277 assert.Nil(t, err)
278 if debug {
279 fmt.Println(i, proof)
280 }
281 }
282 // Non-existence proof, empty aux
283 proof, err = mt.GenerateProof(big.NewInt(12), nil)
284 if err != nil {
285 t.Fatal(err)
286 }
287 assert.Equal(t, proof.Existence, false)
288 // assert.True(t, proof.nodeAux == nil)
289 assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(12), big.NewInt(0)))
290 assert.Equal(t, "030300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b16621630b9104d995843c7cffa685009a1b28dcd371022a3b27b3a4d6987f7c8b39b0f2fffc165330710754ca0fc24451bdd5d5f82a05f42f1427fbdf17879c0b84be60f04000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", hex.EncodeToString(proof.Bytes()))
291
292 // Non-existence proof, diff. node aux
293 proof, err = mt.GenerateProof(big.NewInt(10), nil)
294 if err != nil {
295 t.Fatal(err)
296 }
297 assert.Equal(t, proof.Existence, false)
298 assert.True(t, proof.NodeAux != nil)
299 assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(10), big.NewInt(0)))
300 assert.Equal(t, "030300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b1662163097fcf8f911b271df196e0a75667b8a4f3024ef39f87201ed2b7cda349ba202296b7aeba35dc19ab0d4f65e175536c9952a90b6de18c3205611c3cd4fb408f01602000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", hex.EncodeToString(proof.Bytes()))
301}
302
303func TestVerifyProofFalse(t *testing.T) {
304 mt := newTestingMerkle(t, 140)
305 defer mt.DB().Close()
306
307 for i := 0; i < 8; i++ {
308 if err := mt.Add(big.NewInt(int64(i)), big.NewInt(0)); err != nil {
309 t.Fatal(err)
310 }
311 }
312
313 // Invalid existence proof (node used for verification doesn't
314 // correspond to node in the proof)
315 proof, err := mt.GenerateProof(big.NewInt(int64(4)), nil)
316 if err != nil {
317 t.Fatal(err)
318 }
319 assert.Equal(t, proof.Existence, true)
320 assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(int64(5)), big.NewInt(int64(5))))
321
322 // Invalid non-existence proof (Non-existence proof, diff. node aux)
323 proof, err = mt.GenerateProof(big.NewInt(int64(4)), nil)
324 if err != nil {
325 t.Fatal(err)
326 }
327 assert.Equal(t, proof.Existence, true)
328 // Now we change the proof from existence to non-existence, and add e's
329 // data as auxiliary node.
330 proof.Existence = false
331 proof.NodeAux = &NodeAux{Key: NewHashFromBigInt(big.NewInt(int64(4))), Value: NewHashFromBigInt(big.NewInt(4))}
332 assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(int64(4)), big.NewInt(0)))
333}
334
335func TestGraphViz(t *testing.T) {
336 mt, err := NewMerkleTree(memory.NewMemoryStorage(), 10)
337 assert.Nil(t, err)
338
339 mt.Add(big.NewInt(1), big.NewInt(0))
340 mt.Add(big.NewInt(2), big.NewInt(0))
341 mt.Add(big.NewInt(3), big.NewInt(0))
342 mt.Add(big.NewInt(4), big.NewInt(0))
343 mt.Add(big.NewInt(5), big.NewInt(0))
344 mt.Add(big.NewInt(100), big.NewInt(0))
345
346 // mt.PrintGraphViz(nil)
347
348 expected := `digraph hierarchy {
349node [fontname=Monospace,fontsize=10,shape=box]
350"60195538..." -> {"19759736..." "18893277..."}
351"19759736..." -> {"16152312..." "43945008..."}
352"16152312..." -> {"empty0" "13952255..."}
353"empty0" [style=dashed,label=0];
354"13952255..." -> {"61769925..." "empty1"}
355"empty1" [style=dashed,label=0];
356"61769925..." -> {"92723289..." "empty2"}
357"empty2" [style=dashed,label=0];
358"92723289..." -> {"21082735..." "82784818..."}
359"21082735..." [style=filled];
360"82784818..." [style=filled];
361"43945008..." [style=filled];
362"18893277..." -> {"19855703..." "17718670..."}
363"19855703..." -> {"11499909..." "15828714..."}
364"11499909..." [style=filled];
365"15828714..." [style=filled];
366"17718670..." [style=filled];
367}
368`
369 w := bytes.NewBufferString("")
370 mt.GraphViz(w, nil)
371 assert.Equal(t, []byte(expected), w.Bytes())
372}
373
374func TestDelete(t *testing.T) {
375 mt, err := NewMerkleTree(memory.NewMemoryStorage(), 10)
376 assert.Nil(t, err)
377 assert.Equal(t, "0", mt.Root().String())
378
379 // test vectors generated using https://github.com/iden3/circomlib smt.js
380 err = mt.Add(big.NewInt(1), big.NewInt(2))
381 assert.Nil(t, err)
382 assert.Equal(t, "4932297968297298434239270129193057052722409868268166443802652458940273154854", mt.Root().BigInt().String())
383
384 err = mt.Add(big.NewInt(33), big.NewInt(44))
385 assert.Nil(t, err)
386 assert.Equal(t, "13563340744765267202993741297198970774200042973817962221376874695587906013050", mt.Root().BigInt().String())
387
388 err = mt.Add(big.NewInt(1234), big.NewInt(9876))
389 assert.Nil(t, err)
390 assert.Equal(t, "16970503620176669663662021947486532860010370357132361783766545149750777353066", mt.Root().BigInt().String())
391
392 // mt.PrintGraphViz(nil)
393
394 err = mt.Delete(big.NewInt(33))
395 // mt.PrintGraphViz(nil)
396 assert.Nil(t, err)
397 assert.Equal(t, "12820263606494630162816839760750120928463716794691735985748071431547370997091", mt.Root().BigInt().String())
398
399 err = mt.Delete(big.NewInt(1234))
400 assert.Nil(t, err)
401 err = mt.Delete(big.NewInt(1))
402 assert.Nil(t, err)
403 assert.Equal(t, "0", mt.Root().String())
404
405}
406
407func TestDelete2(t *testing.T) {
408 mt := newTestingMerkle(t, 140)
409 defer mt.db.Close()
410 for i := 0; i < 8; i++ {
411 k := big.NewInt(int64(i))
412 v := big.NewInt(0)
413 if err := mt.Add(k, v); err != nil {
414 t.Fatal(err)
415 }
416 }
417
418 expectedRoot := mt.Root()
419
420 k := big.NewInt(8)
421 v := big.NewInt(0)
422 err := mt.Add(k, v)
423 require.Nil(t, err)
424
425 err = mt.Delete(big.NewInt(8))
426 assert.Nil(t, err)
427 assert.Equal(t, expectedRoot, mt.Root())
428
429 mt2 := newTestingMerkle(t, 140)
430 defer mt2.db.Close()
431 for i := 0; i < 8; i++ {
432 k := big.NewInt(int64(i))
433 v := big.NewInt(0)
434 if err := mt2.Add(k, v); err != nil {
435 t.Fatal(err)
436 }
437 }
438 assert.Equal(t, mt2.Root(), mt.Root())
439}
440
441func TestDelete3(t *testing.T) {
442 mt := newTestingMerkle(t, 140)
443 defer mt.db.Close()
444
445 err := mt.Add(big.NewInt(1), big.NewInt(1))
446 assert.Nil(t, err)
447
448 err = mt.Add(big.NewInt(2), big.NewInt(2))
449 assert.Nil(t, err)
450
451 assert.Equal(t, "2427629547967522489273866134471574861207714751136138191708011221765688788661", mt.Root().BigInt().String())
452 err = mt.Delete(big.NewInt(1))
453 assert.Nil(t, err)
454 assert.Equal(t, "10822920717809411688334493481050035035708810950159417482558569847174767667301", mt.Root().BigInt().String())
455
456 mt2 := newTestingMerkle(t, 140)
457 defer mt2.db.Close()
458 err = mt2.Add(big.NewInt(2), big.NewInt(2))
459 assert.Nil(t, err)
460 assert.Equal(t, mt2.Root(), mt.Root())
461}
462
463func TestDelete4(t *testing.T) {
464 mt := newTestingMerkle(t, 140)
465 defer mt.db.Close()
466
467 err := mt.Add(big.NewInt(1), big.NewInt(1))
468 assert.Nil(t, err)
469
470 err = mt.Add(big.NewInt(2), big.NewInt(2))
471 assert.Nil(t, err)
472
473 err = mt.Add(big.NewInt(3), big.NewInt(3))
474 assert.Nil(t, err)
475
476 assert.Equal(t, "16614298246517994771186095530428786749320098419259206061045083278756632941513", mt.Root().BigInt().String())
477 err = mt.Delete(big.NewInt(1))
478 assert.Nil(t, err)
479 assert.Equal(t, "6117330520107511783353383870014397665359816230889739699667943862706617498952", mt.Root().BigInt().String())
480
481 mt2 := newTestingMerkle(t, 140)
482 defer mt2.db.Close()
483 err = mt2.Add(big.NewInt(2), big.NewInt(2))
484 assert.Nil(t, err)
485 err = mt2.Add(big.NewInt(3), big.NewInt(3))
486 assert.Nil(t, err)
487 assert.Equal(t, mt2.Root(), mt.Root())
488}
489
490func TestDelete5(t *testing.T) {
491 mt, err := NewMerkleTree(memory.NewMemoryStorage(), 10)
492 assert.Nil(t, err)
493
494 err = mt.Add(big.NewInt(1), big.NewInt(2))
495 assert.Nil(t, err)
496 err = mt.Add(big.NewInt(33), big.NewInt(44))
497 assert.Nil(t, err)
498 assert.Equal(t, "13563340744765267202993741297198970774200042973817962221376874695587906013050", mt.Root().BigInt().String())
499
500 err = mt.Delete(big.NewInt(1))
501 assert.Nil(t, err)
502 assert.Equal(t, "12075524681474630909546786277734445038384732558409197537058769521806571391765", mt.Root().BigInt().String())
503
504 mt2 := newTestingMerkle(t, 140)
505 defer mt2.db.Close()
506 err = mt2.Add(big.NewInt(33), big.NewInt(44))
507 assert.Nil(t, err)
508 assert.Equal(t, mt2.Root(), mt.Root())
509}
510
511func TestDeleteNonExistingKeys(t *testing.T) {
512 mt, err := NewMerkleTree(memory.NewMemoryStorage(), 10)
513 assert.Nil(t, err)
514
515 err = mt.Add(big.NewInt(1), big.NewInt(2))
516 assert.Nil(t, err)
517 err = mt.Add(big.NewInt(33), big.NewInt(44))
518 assert.Nil(t, err)
519
520 err = mt.Delete(big.NewInt(33))
521 assert.Nil(t, err)
522 err = mt.Delete(big.NewInt(33))
523 assert.Equal(t, ErrKeyNotFound, err)
524
525 err = mt.Delete(big.NewInt(1))
526 assert.Nil(t, err)
527
528 assert.Equal(t, "0", mt.Root().String())
529
530 err = mt.Delete(big.NewInt(33))
531 assert.Equal(t, ErrKeyNotFound, err)
532}
533
534func TestDumpLeafsImportLeafs(t *testing.T) {
535 mt, err := NewMerkleTree(memory.NewMemoryStorage(), 140)
536 require.Nil(t, err)
537 defer mt.db.Close()
538
539 for i := 0; i < 10; i++ {
540 k := big.NewInt(int64(i))
541 v := big.NewInt(0)
542 err = mt.Add(k, v)
543 require.Nil(t, err)
544 }
545
546 d, err := mt.DumpLeafs(nil)
547 assert.Nil(t, err)
548
549 mt2, err := NewMerkleTree(memory.NewMemoryStorage(), 140)
550 require.Nil(t, err)
551 defer mt2.db.Close()
552 err = mt2.ImportDumpedLeafs(d)
553 assert.Nil(t, err)
554
555 assert.Equal(t, mt.Root(), mt2.Root())
556}