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