arbo/tree_test.go
tree_test.goRaw
1package arbo
2
3import (
4 "encoding/hex"
5 "math"
6 "math/big"
7 "runtime"
8 "testing"
9 "time"
10
11 qt "github.com/frankban/quicktest"
12 "go.vocdoni.io/dvote/db"
13 "go.vocdoni.io/dvote/db/badgerdb"
14)
15
16func checkRootBIString(c *qt.C, tree *Tree, expected string) {
17 root, err := tree.Root()
18 c.Assert(err, qt.IsNil)
19 rootBI := BytesToBigInt(root)
20 c.Check(rootBI.String(), qt.Equals, expected)
21}
22
23func TestDBTx(t *testing.T) {
24 c := qt.New(t)
25
26 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
27 c.Assert(err, qt.IsNil)
28 wTx := database.WriteTx()
29
30 _, err = wTx.Get([]byte("a"))
31 c.Assert(err, qt.Equals, db.ErrKeyNotFound)
32
33 err = wTx.Set([]byte("a"), []byte("b"))
34 c.Assert(err, qt.IsNil)
35
36 v, err := wTx.Get([]byte("a"))
37 c.Assert(err, qt.IsNil)
38 c.Assert(v, qt.DeepEquals, []byte("b"))
39}
40
41func TestAddTestVectors(t *testing.T) {
42 c := qt.New(t)
43
44 // Poseidon test vectors generated using https://github.com/iden3/circomlib smt.js
45 testVectorsPoseidon := []string{
46 "0000000000000000000000000000000000000000000000000000000000000000",
47 "13578938674299138072471463694055224830892726234048532520316387704878000008795",
48 "5412393676474193513566895793055462193090331607895808993925969873307089394741",
49 "14204494359367183802864593755198662203838502594566452929175967972147978322084",
50 }
51 testAdd(c, HashFunctionPoseidon, testVectorsPoseidon)
52
53 testVectorsSha256 := []string{
54 "0000000000000000000000000000000000000000000000000000000000000000",
55 "46910109172468462938850740851377282682950237270676610513794735904325820156367",
56 "59481735341404520835410489183267411392292882901306595567679529387376287440550",
57 "20573794434149960984975763118181266662429997821552560184909083010514790081771",
58 }
59 testAdd(c, HashFunctionSha256, testVectorsSha256)
60}
61
62func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
63 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
64 c.Assert(err, qt.IsNil)
65 tree, err := NewTree(database, 256, hashFunc)
66 c.Assert(err, qt.IsNil)
67 defer tree.db.Close() //nolint:errcheck
68
69 root, err := tree.Root()
70 c.Assert(err, qt.IsNil)
71 c.Check(hex.EncodeToString(root), qt.Equals, testVectors[0])
72
73 bLen := 32
74 err = tree.Add(
75 BigIntToBytes(bLen, big.NewInt(1)),
76 BigIntToBytes(bLen, big.NewInt(2)))
77 c.Assert(err, qt.IsNil)
78 checkRootBIString(c, tree, testVectors[1])
79
80 err = tree.Add(
81 BigIntToBytes(bLen, big.NewInt(33)),
82 BigIntToBytes(bLen, big.NewInt(44)))
83 c.Assert(err, qt.IsNil)
84 checkRootBIString(c, tree, testVectors[2])
85
86 err = tree.Add(
87 BigIntToBytes(bLen, big.NewInt(1234)),
88 BigIntToBytes(bLen, big.NewInt(9876)))
89 c.Assert(err, qt.IsNil)
90 checkRootBIString(c, tree, testVectors[3])
91}
92
93func TestAddBatch(t *testing.T) {
94 c := qt.New(t)
95 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
96 c.Assert(err, qt.IsNil)
97 tree, err := NewTree(database, 256, HashFunctionPoseidon)
98 c.Assert(err, qt.IsNil)
99 defer tree.db.Close() //nolint:errcheck
100
101 bLen := 32
102 for i := 0; i < 1000; i++ {
103 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
104 v := BigIntToBytes(bLen, big.NewInt(0))
105 if err := tree.Add(k, v); err != nil {
106 t.Fatal(err)
107 }
108 }
109
110 checkRootBIString(c, tree,
111 "296519252211642170490407814696803112091039265640052570497930797516015811235")
112
113 database, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()})
114 c.Assert(err, qt.IsNil)
115 tree2, err := NewTree(database, 256, HashFunctionPoseidon)
116 c.Assert(err, qt.IsNil)
117 defer tree2.db.Close() //nolint:errcheck
118
119 var keys, values [][]byte
120 for i := 0; i < 1000; i++ {
121 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
122 v := BigIntToBytes(bLen, big.NewInt(0))
123 keys = append(keys, k)
124 values = append(values, v)
125 }
126 indexes, err := tree2.AddBatch(keys, values)
127 c.Assert(err, qt.IsNil)
128 c.Check(len(indexes), qt.Equals, 0)
129
130 checkRootBIString(c, tree2,
131 "296519252211642170490407814696803112091039265640052570497930797516015811235")
132}
133
134func TestAddDifferentOrder(t *testing.T) {
135 c := qt.New(t)
136 database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
137 c.Assert(err, qt.IsNil)
138 tree1, err := NewTree(database1, 256, HashFunctionPoseidon)
139 c.Assert(err, qt.IsNil)
140 defer tree1.db.Close() //nolint:errcheck
141
142 bLen := 32
143 for i := 0; i < 16; i++ {
144 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
145 v := BigIntToBytes(bLen, big.NewInt(0))
146 if err := tree1.Add(k, v); err != nil {
147 t.Fatal(err)
148 }
149 }
150
151 database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
152 c.Assert(err, qt.IsNil)
153 tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
154 c.Assert(err, qt.IsNil)
155 defer tree2.db.Close() //nolint:errcheck
156
157 for i := 16 - 1; i >= 0; i-- {
158 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
159 v := BigIntToBytes(bLen, big.NewInt(0))
160 if err := tree2.Add(k, v); err != nil {
161 t.Fatal(err)
162 }
163 }
164
165 root1, err := tree1.Root()
166 c.Assert(err, qt.IsNil)
167 root2, err := tree2.Root()
168 c.Assert(err, qt.IsNil)
169 c.Check(hex.EncodeToString(root2), qt.Equals, hex.EncodeToString(root1))
170 c.Check(hex.EncodeToString(root1), qt.Equals,
171 "3b89100bec24da9275c87bc188740389e1d5accfc7d88ba5688d7fa96a00d82f")
172}
173
174func TestAddRepeatedIndex(t *testing.T) {
175 c := qt.New(t)
176 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
177 c.Assert(err, qt.IsNil)
178 tree, err := NewTree(database, 256, HashFunctionPoseidon)
179 c.Assert(err, qt.IsNil)
180 defer tree.db.Close() //nolint:errcheck
181
182 bLen := 32
183 k := BigIntToBytes(bLen, big.NewInt(int64(3)))
184 v := BigIntToBytes(bLen, big.NewInt(int64(12)))
185
186 err = tree.Add(k, v)
187 c.Assert(err, qt.IsNil)
188 err = tree.Add(k, v) // repeating same key-value
189 c.Check(err, qt.Equals, ErrKeyAlreadyExists)
190}
191
192func TestUpdate(t *testing.T) {
193 c := qt.New(t)
194 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
195 c.Assert(err, qt.IsNil)
196 tree, err := NewTree(database, 256, HashFunctionPoseidon)
197 c.Assert(err, qt.IsNil)
198 defer tree.db.Close() //nolint:errcheck
199
200 bLen := 32
201 k := BigIntToBytes(bLen, big.NewInt(int64(20)))
202 v := BigIntToBytes(bLen, big.NewInt(int64(12)))
203 if err := tree.Add(k, v); err != nil {
204 t.Fatal(err)
205 }
206
207 v = BigIntToBytes(bLen, big.NewInt(int64(11)))
208 err = tree.Update(k, v)
209 c.Assert(err, qt.IsNil)
210
211 gettedKey, gettedValue, err := tree.Get(k)
212 c.Assert(err, qt.IsNil)
213 c.Check(gettedKey, qt.DeepEquals, k)
214 c.Check(gettedValue, qt.DeepEquals, v)
215
216 // add more leafs to the tree to do another test
217 for i := 0; i < 16; i++ {
218 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
219 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
220 if err := tree.Add(k, v); err != nil {
221 t.Fatal(err)
222 }
223 }
224
225 k = BigIntToBytes(bLen, big.NewInt(int64(3)))
226 v = BigIntToBytes(bLen, big.NewInt(int64(11)))
227 // check that before the Update, value for 3 is !=11
228 gettedKey, gettedValue, err = tree.Get(k)
229 c.Assert(err, qt.IsNil)
230 c.Check(gettedKey, qt.DeepEquals, k)
231 c.Check(gettedValue, qt.Not(qt.DeepEquals), v)
232 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(6)))
233
234 err = tree.Update(k, v)
235 c.Assert(err, qt.IsNil)
236
237 // check that after Update, the value for 3 is ==11
238 gettedKey, gettedValue, err = tree.Get(k)
239 c.Assert(err, qt.IsNil)
240 c.Check(gettedKey, qt.DeepEquals, k)
241 c.Check(gettedValue, qt.DeepEquals, v)
242 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(11)))
243}
244
245func TestAux(t *testing.T) { // TODO split in proper tests
246 c := qt.New(t)
247 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
248 c.Assert(err, qt.IsNil)
249 tree, err := NewTree(database, 256, HashFunctionPoseidon)
250 c.Assert(err, qt.IsNil)
251 defer tree.db.Close() //nolint:errcheck
252
253 bLen := 32
254 k := BigIntToBytes(bLen, big.NewInt(int64(1)))
255 v := BigIntToBytes(bLen, big.NewInt(int64(0)))
256 err = tree.Add(k, v)
257 c.Assert(err, qt.IsNil)
258 k = BigIntToBytes(bLen, big.NewInt(int64(256)))
259 err = tree.Add(k, v)
260 c.Assert(err, qt.IsNil)
261
262 k = BigIntToBytes(bLen, big.NewInt(int64(257)))
263 err = tree.Add(k, v)
264 c.Assert(err, qt.IsNil)
265
266 k = BigIntToBytes(bLen, big.NewInt(int64(515)))
267 err = tree.Add(k, v)
268 c.Assert(err, qt.IsNil)
269 k = BigIntToBytes(bLen, big.NewInt(int64(770)))
270 err = tree.Add(k, v)
271 c.Assert(err, qt.IsNil)
272
273 k = BigIntToBytes(bLen, big.NewInt(int64(388)))
274 err = tree.Add(k, v)
275 c.Assert(err, qt.IsNil)
276 k = BigIntToBytes(bLen, big.NewInt(int64(900)))
277 err = tree.Add(k, v)
278 c.Assert(err, qt.IsNil)
279 //
280 // err = tree.PrintGraphviz(nil)
281 // c.Assert(err, qt.IsNil)
282}
283
284func TestGet(t *testing.T) {
285 c := qt.New(t)
286 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
287 c.Assert(err, qt.IsNil)
288 tree, err := NewTree(database, 256, HashFunctionPoseidon)
289 c.Assert(err, qt.IsNil)
290 defer tree.db.Close() //nolint:errcheck
291
292 bLen := 32
293 for i := 0; i < 10; i++ {
294 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
295 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
296 if err := tree.Add(k, v); err != nil {
297 t.Fatal(err)
298 }
299 }
300
301 k := BigIntToBytes(bLen, big.NewInt(int64(7)))
302 gettedKey, gettedValue, err := tree.Get(k)
303 c.Assert(err, qt.IsNil)
304 c.Check(gettedKey, qt.DeepEquals, k)
305 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(int64(7*2))))
306}
307
308func TestPackAndUnpackSiblings(t *testing.T) {
309 c := qt.New(t)
310
311 siblingsHex := []string{
312 "0000000000000000000000000000000000000000000000000000000000000000",
313 "0100000000000000000000000000000000000000000000000000000000000000",
314 "0200000000000000000000000000000000000000000000000000000000000000",
315 "0000000000000000000000000000000000000000000000000000000000000000",
316 "0000000000000000000000000000000000000000000000000000000000000000",
317 "0000000000000000000000000000000000000000000000000000000000000000",
318 "0300000000000000000000000000000000000000000000000000000000000000",
319 "0400000000000000000000000000000000000000000000000000000000000000",
320 "0000000000000000000000000000000000000000000000000000000000000000",
321 "0000000000000000000000000000000000000000000000000000000000000000",
322 "0500000000000000000000000000000000000000000000000000000000000000",
323 }
324 siblings := make([][]byte, len(siblingsHex))
325 var err error
326 for i := 0; i < len(siblingsHex); i++ {
327 siblings[i], err = hex.DecodeString(siblingsHex[i])
328 c.Assert(err, qt.IsNil)
329 }
330
331 packed, err := PackSiblings(HashFunctionPoseidon, siblings)
332 c.Assert(err, qt.IsNil)
333 c.Assert(hex.EncodeToString(packed), qt.Equals, "a6000200c604"+
334 "0100000000000000000000000000000000000000000000000000000000000000"+
335 "0200000000000000000000000000000000000000000000000000000000000000"+
336 "0300000000000000000000000000000000000000000000000000000000000000"+
337 "0400000000000000000000000000000000000000000000000000000000000000"+
338 "0500000000000000000000000000000000000000000000000000000000000000")
339
340 unpacked, err := UnpackSiblings(HashFunctionPoseidon, packed)
341 c.Assert(err, qt.IsNil)
342 c.Assert(unpacked, qt.DeepEquals, siblings)
343
344 // another test with other values
345 siblingsHex = []string{
346 "1ce165cb1124ed3a0a94b4e212aaf7e8079f49b2fbef916bc290c593fda9092a",
347 "0000000000000000000000000000000000000000000000000000000000000000",
348 "0000000000000000000000000000000000000000000000000000000000000000",
349 "0000000000000000000000000000000000000000000000000000000000000000",
350 "33018202c57d898b84338b16d1a4960e133c6a4d656cfec1bd62a9ea00611729",
351 "bdbee2bd246ba0259a37be9a8740b550eed01c566aff0dca9a07306bcf731d13",
352 "0000000000000000000000000000000000000000000000000000000000000000",
353 "0000000000000000000000000000000000000000000000000000000000000000",
354 "d43b04d7c2d0bba83b4291fea9ba0fec7830d17af54cbe9967fe90b8244d4e0d",
355 "0000000000000000000000000000000000000000000000000000000000000000",
356 "7def274dbb3a72dca44f01a8d9f2f21a5be84c171eecef8e2e4112e7277e262a",
357 }
358 siblings = make([][]byte, len(siblingsHex))
359 for i := 0; i < len(siblingsHex); i++ {
360 siblings[i], err = hex.DecodeString(siblingsHex[i])
361 c.Assert(err, qt.IsNil)
362 }
363
364 packed, err = PackSiblings(HashFunctionPoseidon, siblings)
365 c.Assert(err, qt.IsNil)
366 c.Assert(hex.EncodeToString(packed), qt.Equals, "a60002003105"+
367 "1ce165cb1124ed3a0a94b4e212aaf7e8079f49b2fbef916bc290c593fda9092a"+
368 "33018202c57d898b84338b16d1a4960e133c6a4d656cfec1bd62a9ea00611729"+
369 "bdbee2bd246ba0259a37be9a8740b550eed01c566aff0dca9a07306bcf731d13"+
370 "d43b04d7c2d0bba83b4291fea9ba0fec7830d17af54cbe9967fe90b8244d4e0d"+
371 "7def274dbb3a72dca44f01a8d9f2f21a5be84c171eecef8e2e4112e7277e262a")
372
373 unpacked, err = UnpackSiblings(HashFunctionPoseidon, packed)
374 c.Assert(err, qt.IsNil)
375 c.Assert(unpacked, qt.DeepEquals, siblings)
376}
377
378func TestGenProofAndVerify(t *testing.T) {
379 c := qt.New(t)
380 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
381 c.Assert(err, qt.IsNil)
382 tree, err := NewTree(database, 256, HashFunctionPoseidon)
383 c.Assert(err, qt.IsNil)
384 defer tree.db.Close() //nolint:errcheck
385
386 bLen := 32
387 for i := 0; i < 10; i++ {
388 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
389 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
390 if err := tree.Add(k, v); err != nil {
391 t.Fatal(err)
392 }
393 }
394
395 k := BigIntToBytes(bLen, big.NewInt(int64(7)))
396 v := BigIntToBytes(bLen, big.NewInt(int64(14)))
397 kAux, proofV, siblings, existence, err := tree.GenProof(k)
398 c.Assert(err, qt.IsNil)
399 c.Assert(proofV, qt.DeepEquals, v)
400 c.Assert(k, qt.DeepEquals, kAux)
401 c.Assert(existence, qt.IsTrue)
402
403 root, err := tree.Root()
404 c.Assert(err, qt.IsNil)
405 verif, err := CheckProof(tree.hashFunction, k, v, root, siblings)
406 c.Assert(err, qt.IsNil)
407 c.Check(verif, qt.IsTrue)
408}
409
410func TestDumpAndImportDump(t *testing.T) {
411 c := qt.New(t)
412 database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
413 c.Assert(err, qt.IsNil)
414 tree1, err := NewTree(database1, 256, HashFunctionPoseidon)
415 c.Assert(err, qt.IsNil)
416 defer tree1.db.Close() //nolint:errcheck
417
418 bLen := 32
419 for i := 0; i < 16; i++ {
420 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
421 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
422 if err := tree1.Add(k, v); err != nil {
423 t.Fatal(err)
424 }
425 }
426
427 e, err := tree1.Dump(nil)
428 c.Assert(err, qt.IsNil)
429
430 database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
431 c.Assert(err, qt.IsNil)
432 tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
433 c.Assert(err, qt.IsNil)
434 defer tree2.db.Close() //nolint:errcheck
435 err = tree2.ImportDump(e)
436 c.Assert(err, qt.IsNil)
437
438 root1, err := tree1.Root()
439 c.Assert(err, qt.IsNil)
440 root2, err := tree2.Root()
441 c.Assert(err, qt.IsNil)
442 c.Check(root2, qt.DeepEquals, root1)
443 c.Check(hex.EncodeToString(root2), qt.Equals,
444 "0d93aaa3362b2f999f15e15728f123087c2eee716f01c01f56e23aae07f09f08")
445}
446
447func TestRWMutex(t *testing.T) {
448 c := qt.New(t)
449 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
450 c.Assert(err, qt.IsNil)
451 tree, err := NewTree(database, 256, HashFunctionPoseidon)
452 c.Assert(err, qt.IsNil)
453 defer tree.db.Close() //nolint:errcheck
454
455 bLen := 32
456 var keys, values [][]byte
457 for i := 0; i < 1000; i++ {
458 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
459 v := BigIntToBytes(bLen, big.NewInt(0))
460 keys = append(keys, k)
461 values = append(values, v)
462 }
463 go func() {
464 _, err = tree.AddBatch(keys, values)
465 if err != nil {
466 panic(err)
467 }
468 }()
469
470 time.Sleep(500 * time.Millisecond)
471 k := BigIntToBytes(bLen, big.NewInt(int64(99999)))
472 v := BigIntToBytes(bLen, big.NewInt(int64(99999)))
473 if err := tree.Add(k, v); err != nil {
474 t.Fatal(err)
475 }
476}
477
478// TODO UPDATE
479// func TestSetGetNLeafs(t *testing.T) {
480// c := qt.New(t)
481// database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
482// c.Assert(err, qt.IsNil)
483// tree, err := NewTree(database, 100, HashFunctionPoseidon)
484// c.Assert(err, qt.IsNil)
485//
486// // 0
487// tree.dbBatch = tree.db.NewBatch()
488//
489// err = tree.setNLeafs(0)
490// c.Assert(err, qt.IsNil)
491//
492// err = tree.dbBatch.Write()
493// c.Assert(err, qt.IsNil)
494//
495// n, err := tree.GetNLeafs()
496// c.Assert(err, qt.IsNil)
497// c.Assert(n, qt.Equals, 0)
498//
499// // 1024
500// tree.dbBatch = tree.db.NewBatch()
501//
502// err = tree.setNLeafs(1024)
503// c.Assert(err, qt.IsNil)
504//
505// err = tree.dbBatch.Write()
506// c.Assert(err, qt.IsNil)
507//
508// n, err = tree.GetNLeafs()
509// c.Assert(err, qt.IsNil)
510// c.Assert(n, qt.Equals, 1024)
511//
512// // 2**64 -1
513// tree.dbBatch = tree.db.NewBatch()
514//
515// maxUint := ^uint(0)
516// maxInt := int(maxUint >> 1)
517//
518// err = tree.setNLeafs(maxInt)
519// c.Assert(err, qt.IsNil)
520//
521// err = tree.dbBatch.Write()
522// c.Assert(err, qt.IsNil)
523//
524// n, err = tree.GetNLeafs()
525// c.Assert(err, qt.IsNil)
526// c.Assert(n, qt.Equals, maxInt)
527// }
528
529func TestAddBatchFullyUsed(t *testing.T) {
530 c := qt.New(t)
531
532 database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
533 c.Assert(err, qt.IsNil)
534 tree1, err := NewTree(database1, 4, HashFunctionPoseidon)
535 c.Assert(err, qt.IsNil)
536
537 database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
538 c.Assert(err, qt.IsNil)
539 tree2, err := NewTree(database2, 4, HashFunctionPoseidon)
540 c.Assert(err, qt.IsNil)
541
542 var keys, values [][]byte
543 for i := 0; i < 16; i++ {
544 k := BigIntToBytes(1, big.NewInt(int64(i)))
545 v := k
546
547 keys = append(keys, k)
548 values = append(values, v)
549
550 // add one by one expecting no error
551 err := tree1.Add(k, v)
552 c.Assert(err, qt.IsNil)
553 }
554
555 invalids, err := tree2.AddBatch(keys, values)
556 c.Assert(err, qt.IsNil)
557 c.Assert(0, qt.Equals, len(invalids))
558
559 root1, err := tree1.Root()
560 c.Assert(err, qt.IsNil)
561 root2, err := tree2.Root()
562 c.Assert(err, qt.IsNil)
563 c.Assert(root1, qt.DeepEquals, root2)
564
565 // get all key-values and check that are equal between both trees
566 for i := 0; i < 16; i++ {
567 auxK1, auxV1, err := tree1.Get(BigIntToBytes(1, big.NewInt(int64(i))))
568 c.Assert(err, qt.IsNil)
569
570 auxK2, auxV2, err := tree2.Get(BigIntToBytes(1, big.NewInt(int64(i))))
571 c.Assert(err, qt.IsNil)
572
573 c.Assert(auxK1, qt.DeepEquals, auxK2)
574 c.Assert(auxV1, qt.DeepEquals, auxV2)
575 }
576
577 // try adding one more key to both trees (through Add & AddBatch) and
578 // expect not being added due the tree is already full
579 k := BigIntToBytes(1, big.NewInt(int64(16)))
580 v := k
581 err = tree1.Add(k, v)
582 c.Assert(err, qt.Equals, ErrMaxVirtualLevel)
583
584 invalids, err = tree2.AddBatch([][]byte{k}, [][]byte{v})
585 c.Assert(err, qt.IsNil)
586 c.Assert(1, qt.Equals, len(invalids))
587}
588
589func TestSetRoot(t *testing.T) {
590 c := qt.New(t)
591 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
592 c.Assert(err, qt.IsNil)
593 tree, err := NewTree(database, 256, HashFunctionPoseidon)
594 c.Assert(err, qt.IsNil)
595
596 expectedRoot := "13742386369878513332697380582061714160370929283209286127733983161245560237407"
597
598 // fill the tree
599 bLen := 32
600 var keys, values [][]byte
601 for i := 0; i < 1000; i++ {
602 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
603 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
604 keys = append(keys, k)
605 values = append(values, v)
606 }
607 indexes, err := tree.AddBatch(keys, values)
608 c.Assert(err, qt.IsNil)
609 c.Check(len(indexes), qt.Equals, 0)
610 checkRootBIString(c, tree,
611 expectedRoot)
612
613 // add one more k-v
614 k := BigIntToBytes(bLen, big.NewInt(1000))
615 v := BigIntToBytes(bLen, big.NewInt(1000))
616 err = tree.Add(k, v)
617 c.Assert(err, qt.IsNil)
618 checkRootBIString(c, tree,
619 "10747149055773881257049574592162159501044114324358186833013814735296193179713")
620
621 // do a SetRoot, and expect the same root than the original tree
622 pastRootBI, ok := new(big.Int).SetString(expectedRoot, 10)
623 c.Assert(ok, qt.IsTrue)
624 pastRoot := BigIntToBytes(32, pastRootBI)
625
626 err = tree.SetRoot(pastRoot)
627 c.Assert(err, qt.IsNil)
628 checkRootBIString(c, tree, expectedRoot)
629
630 // check that the tree can be updated
631 err = tree.Add([]byte("test"), []byte("test"))
632 c.Assert(err, qt.IsNil)
633 err = tree.Update([]byte("test"), []byte("test"))
634 c.Assert(err, qt.IsNil)
635
636 // check that the k-v '1000' does not exist in the new tree
637 _, _, err = tree.Get(k)
638 c.Assert(err, qt.Equals, ErrKeyNotFound)
639
640 // check that can be set an empty root
641 err = tree.SetRoot(tree.emptyHash)
642 c.Assert(err, qt.IsNil)
643}
644
645func TestSnapshot(t *testing.T) {
646 c := qt.New(t)
647 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
648 c.Assert(err, qt.IsNil)
649 tree, err := NewTree(database, 256, HashFunctionPoseidon)
650 c.Assert(err, qt.IsNil)
651
652 // fill the tree
653 bLen := 32
654 var keys, values [][]byte
655 for i := 0; i < 1000; i++ {
656 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
657 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
658 keys = append(keys, k)
659 values = append(values, v)
660 }
661 indexes, err := tree.AddBatch(keys, values)
662 c.Assert(err, qt.IsNil)
663 c.Check(len(indexes), qt.Equals, 0)
664 checkRootBIString(c, tree,
665 "13742386369878513332697380582061714160370929283209286127733983161245560237407")
666
667 // do a snapshot, and expect the same root than the original tree
668 snapshotTree, err := tree.Snapshot(nil)
669 c.Assert(err, qt.IsNil)
670 checkRootBIString(c, snapshotTree,
671 "13742386369878513332697380582061714160370929283209286127733983161245560237407")
672
673 // check that the snapshotTree can not be updated
674 _, err = snapshotTree.AddBatch(keys, values)
675 c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
676 err = snapshotTree.Add([]byte("test"), []byte("test"))
677 c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
678 err = snapshotTree.Update([]byte("test"), []byte("test"))
679 c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
680 err = snapshotTree.ImportDump(nil)
681 c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
682
683 // update the original tree by adding a new key-value, and check that
684 // snapshotTree still has the old root, but the original tree has a new
685 // root
686 err = tree.Add([]byte("test"), []byte("test"))
687 c.Assert(err, qt.IsNil)
688 checkRootBIString(c, snapshotTree,
689 "13742386369878513332697380582061714160370929283209286127733983161245560237407")
690 checkRootBIString(c, tree,
691 "1025190963769001718196479367844646783678188389989148142691917685159698888868")
692}
693
694func TestGetFromSnapshotExpectArboErrKeyNotFound(t *testing.T) {
695 c := qt.New(t)
696
697 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
698 c.Assert(err, qt.IsNil)
699 tree, err := NewTree(database, 256, HashFunctionPoseidon)
700 c.Assert(err, qt.IsNil)
701 defer tree.db.Close() //nolint:errcheck
702
703 bLen := 32
704 k := BigIntToBytes(bLen, big.NewInt(int64(3)))
705
706 root, err := tree.Root()
707 c.Assert(err, qt.IsNil)
708 tree, err = tree.Snapshot(root)
709 c.Assert(err, qt.IsNil)
710
711 _, _, err = tree.Get(k)
712 c.Assert(err, qt.Equals, ErrKeyNotFound) // and not equal to db.ErrKeyNotFound
713}
714
715func TestKeyLen(t *testing.T) {
716 c := qt.New(t)
717 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
718 c.Assert(err, qt.IsNil)
719 // maxLevels is 100, keyPath length = ceil(maxLevels/8) = 13
720 maxLevels := 100
721 tree, err := NewTree(database, maxLevels, HashFunctionBlake2b)
722 c.Assert(err, qt.IsNil)
723
724 // expect no errors when adding a key of only 4 bytes (when the
725 // required length of keyPath for 100 levels would be 13 bytes)
726 bLen := 4
727 k := BigIntToBytes(bLen, big.NewInt(1))
728 v := BigIntToBytes(bLen, big.NewInt(1))
729
730 err = tree.Add(k, v)
731 c.Assert(err, qt.IsNil)
732
733 err = tree.Update(k, v)
734 c.Assert(err, qt.IsNil)
735
736 _, _, _, _, err = tree.GenProof(k)
737 c.Assert(err, qt.IsNil)
738
739 _, _, err = tree.Get(k)
740 c.Assert(err, qt.IsNil)
741
742 k = BigIntToBytes(bLen, big.NewInt(2))
743 v = BigIntToBytes(bLen, big.NewInt(2))
744 invalids, err := tree.AddBatch([][]byte{k}, [][]byte{v})
745 c.Assert(err, qt.IsNil)
746 c.Assert(len(invalids), qt.Equals, 0)
747
748 // expect errors when adding a key bigger than maximum capacity of the
749 // tree: ceil(maxLevels/8)
750 maxLevels = 32
751 database, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()})
752 c.Assert(err, qt.IsNil)
753 tree, err = NewTree(database, maxLevels, HashFunctionBlake2b)
754 c.Assert(err, qt.IsNil)
755
756 maxKeyLen := int(math.Ceil(float64(maxLevels) / float64(8))) //nolint:gomnd
757 k = BigIntToBytes(maxKeyLen+1, big.NewInt(1))
758 v = BigIntToBytes(maxKeyLen+1, big.NewInt(1))
759
760 expectedErrMsg := "len(k) can not be bigger than ceil(maxLevels/8)," +
761 " where len(k): 5, maxLevels: 32, max key len=ceil(maxLevels/8): 4." +
762 " Might need a bigger tree depth (maxLevels>=40) in order to input" +
763 " keys of length 5"
764
765 err = tree.Add(k, v)
766 c.Assert(err.Error(), qt.Equals, expectedErrMsg)
767
768 err = tree.Update(k, v)
769 c.Assert(err.Error(), qt.Equals, expectedErrMsg)
770
771 _, _, _, _, err = tree.GenProof(k)
772 c.Assert(err.Error(), qt.Equals, expectedErrMsg)
773
774 _, _, err = tree.Get(k)
775 c.Assert(err.Error(), qt.Equals, expectedErrMsg)
776
777 // check AddBatch with few key-values
778 invalids, err = tree.AddBatch([][]byte{k}, [][]byte{v})
779 c.Assert(err, qt.IsNil)
780 c.Assert(len(invalids), qt.Equals, 1)
781
782 // check AddBatch with many key-values
783 nCPU := flp2(runtime.NumCPU())
784 nKVs := nCPU + 1
785 var ks, vs [][]byte
786 for i := 0; i < nKVs; i++ {
787 ks = append(ks, BigIntToBytes(maxKeyLen+1, big.NewInt(1)))
788 vs = append(vs, BigIntToBytes(maxKeyLen+1, big.NewInt(1)))
789 }
790 invalids, err = tree.AddBatch(ks, vs)
791 c.Assert(err, qt.IsNil)
792 c.Assert(len(invalids), qt.Equals, nKVs)
793
794 // check that with maxKeyLen it can be added
795 k = BigIntToBytes(maxKeyLen, big.NewInt(1))
796 err = tree.Add(k, v)
797 c.Assert(err, qt.IsNil)
798
799 // check CheckProof check with key longer
800 kAux, vAux, packedSiblings, existence, err := tree.GenProof(k)
801 c.Assert(err, qt.IsNil)
802 c.Assert(existence, qt.IsTrue)
803
804 root, err := tree.Root()
805 c.Assert(err, qt.IsNil)
806 verif, err := CheckProof(tree.HashFunction(), kAux, vAux, root, packedSiblings)
807 c.Assert(err, qt.IsNil)
808 c.Assert(verif, qt.IsTrue)
809
810 // use a similar key but with one zero, expect that CheckProof fails on
811 // the verification
812 kAux = append(kAux, 0)
813 verif, err = CheckProof(tree.HashFunction(), kAux, vAux, root, packedSiblings)
814 c.Assert(err, qt.IsNil)
815 c.Assert(verif, qt.IsFalse)
816}
817
818func TestKeyLenBiggerThan32(t *testing.T) {
819 c := qt.New(t)
820 maxLevels := 264
821 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
822 c.Assert(err, qt.IsNil)
823 tree, err := NewTree(database, maxLevels, HashFunctionBlake2b)
824 c.Assert(err, qt.IsNil)
825
826 bLen := 33
827 err = tree.Add(
828 randomBytes(bLen),
829 randomBytes(bLen))
830 c.Assert(err, qt.IsNil)
831
832 // 2nd key that we add, will find a node with len(key)==32 (due the
833 // hash output size, expect that next Add does not give any error, as
834 // internally it will use a keyPath of size corresponent to the
835 // maxLevels size of the tree
836 err = tree.Add(
837 randomBytes(bLen),
838 randomBytes(bLen))
839 c.Assert(err, qt.IsNil)
840}
841
842func BenchmarkAdd(b *testing.B) {
843 bLen := 32 // for both Poseidon & Sha256
844 // prepare inputs
845 var ks, vs [][]byte
846 for i := 0; i < 1000; i++ {
847 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
848 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
849 ks = append(ks, k)
850 vs = append(vs, v)
851 }
852
853 b.Run("Poseidon", func(b *testing.B) {
854 benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
855 })
856 b.Run("Sha256", func(b *testing.B) {
857 benchmarkAdd(b, HashFunctionSha256, ks, vs)
858 })
859}
860
861func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
862 c := qt.New(b)
863 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
864 c.Assert(err, qt.IsNil)
865 tree, err := NewTree(database, 140, hashFunc)
866 c.Assert(err, qt.IsNil)
867 defer tree.db.Close() //nolint:errcheck
868
869 for i := 0; i < len(ks); i++ {
870 if err := tree.Add(ks[i], vs[i]); err != nil {
871 b.Fatal(err)
872 }
873 }
874}