Add checks that len(key)<=maxKeyLen Add checks that the key is not bigger than maximum key length for the tree maxLevels size, where maximum key len = ceil(maxLevels/8). This is because if the key bits length is bigger than the maxLevels of the tree, two different keys that their difference is at the end, will collision in the same leaf of the tree (at the max depth).
arnaucube root@arnaucube.com
Mon, 04 Oct 2021 11:27:21 +0200
8 files changed,
263 insertions(+),
146 deletions(-)
M
addbatch_test.go
→
addbatch_test.go
@@ -39,12 +39,12 @@
func testInit(c *qt.C, n int) (*Tree, *Tree) { database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree1, err := NewTree(database1, 100, HashFunctionPoseidon) + tree1, err := NewTree(database1, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree2, err := NewTree(database2, 100, HashFunctionPoseidon) + tree2, err := NewTree(database2, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) bLen := HashFunctionPoseidon.Len()@@ -70,11 +70,11 @@ nLeafs := 1024
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 100, HashFunctionPoseidon) + tree, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree.db.Close() //nolint:errcheck - bLen := tree.HashFunction().Len() + bLen := 32 var keys, values [][]byte for i := 0; i < nLeafs; i++ { k := BigIntToBytes(bLen, big.NewInt(int64(i)))@@ -93,7 +93,7 @@ time1 := time.Since(start)
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree2, err := NewTree(database2, 100, HashFunctionPoseidon) + tree2, err := NewTree(database2, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree2.db.Close() //nolint:errcheck tree2.dbgInit()@@ -120,11 +120,11 @@ nLeafs := 1027
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 100, HashFunctionPoseidon) + tree, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree.db.Close() //nolint:errcheck - bLen := tree.HashFunction().Len() + bLen := 32 for i := 0; i < nLeafs; i++ { k := BigIntToBytes(bLen, big.NewInt(int64(i))) v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))@@ -135,7 +135,7 @@ }
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree2, err := NewTree(database2, 100, HashFunctionPoseidon) + tree2, err := NewTree(database2, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree2.db.Close() //nolint:errcheck@@ -167,13 +167,13 @@ func TestAddBatchTestVector1(t *testing.T) {
c := qt.New(t) database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree1, err := NewTree(database1, 100, HashFunctionBlake2b) + tree1, err := NewTree(database1, 256, HashFunctionBlake2b) c.Assert(err, qt.IsNil) defer tree1.db.Close() //nolint:errcheck database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree2, err := NewTree(database2, 100, HashFunctionBlake2b) + tree2, err := NewTree(database2, 256, HashFunctionBlake2b) c.Assert(err, qt.IsNil) defer tree2.db.Close() //nolint:errcheck@@ -207,13 +207,13 @@
// 2nd test vectors database1, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree1, err = NewTree(database1, 100, HashFunctionBlake2b) + tree1, err = NewTree(database1, 256, HashFunctionBlake2b) c.Assert(err, qt.IsNil) defer tree1.db.Close() //nolint:errcheck database2, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree2, err = NewTree(database2, 100, HashFunctionBlake2b) + tree2, err = NewTree(database2, 256, HashFunctionBlake2b) c.Assert(err, qt.IsNil) defer tree2.db.Close() //nolint:errcheck@@ -255,13 +255,13 @@ c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree1, err := NewTree(database, 100, HashFunctionPoseidon) + tree1, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree1.db.Close() //nolint:errcheck database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree2, err := NewTree(database2, 100, HashFunctionPoseidon) + tree2, err := NewTree(database2, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree2.db.Close() //nolint:errcheck@@ -300,13 +300,13 @@ c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree1, err := NewTree(database, 100, HashFunctionPoseidon) + tree1, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree1.db.Close() //nolint:errcheck database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree2, err := NewTree(database2, 100, HashFunctionPoseidon) + tree2, err := NewTree(database2, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree2.db.Close() //nolint:errcheck@@ -349,13 +349,13 @@ nLeafs := 8
database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree1, err := NewTree(database1, 100, HashFunctionBlake2b) + tree1, err := NewTree(database1, 256, HashFunctionBlake2b) c.Assert(err, qt.IsNil) defer tree1.db.Close() //nolint:errcheck database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree2, err := NewTree(database2, 100, HashFunctionBlake2b) + tree2, err := NewTree(database2, 256, HashFunctionBlake2b) c.Assert(err, qt.IsNil) defer tree2.db.Close() //nolint:errcheck@@ -699,7 +699,7 @@ time1 := time.Since(start)
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree2, err := NewTree(database2, 100, HashFunctionPoseidon) + tree2, err := NewTree(database2, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree2.db.Close() //nolint:errcheck tree2.dbgInit()@@ -776,7 +776,7 @@ c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 140, HashFunctionBlake2b) + tree, err := NewTree(database, 256, HashFunctionBlake2b) c.Assert(err, qt.IsNil) defer tree.db.Close() //nolint:errcheck@@ -796,7 +796,7 @@ c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 140, HashFunctionBlake2b) + tree, err := NewTree(database, 256, HashFunctionBlake2b) c.Assert(err, qt.IsNil) defer tree.db.Close() //nolint:errcheck@@ -829,7 +829,7 @@
// 1 database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree1, err := NewTree(database1, 100, HashFunctionBlake2b) + tree1, err := NewTree(database1, 256, HashFunctionBlake2b) c.Assert(err, qt.IsNil) defer tree1.db.Close() //nolint:errcheck@@ -843,7 +843,7 @@
// 2 database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree2, err := NewTree(database2, 100, HashFunctionBlake2b) + tree2, err := NewTree(database2, 256, HashFunctionBlake2b) c.Assert(err, qt.IsNil) defer tree2.db.Close() //nolint:errcheck@@ -856,7 +856,7 @@
// 3 database3, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree3, err := NewTree(database3, 100, HashFunctionBlake2b) + tree3, err := NewTree(database3, 256, HashFunctionBlake2b) c.Assert(err, qt.IsNil) defer tree3.db.Close() //nolint:errcheck@@ -891,7 +891,7 @@ nLeafs := 1024
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 100, HashFunctionPoseidon) + tree, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree.db.Close() //nolint:errcheck@@ -927,11 +927,11 @@ nLeafs := 1024
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 100, HashFunctionPoseidon) + tree, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree.db.Close() //nolint:errcheck - bLen := tree.HashFunction().Len() + bLen := 32 var keys, values [][]byte for i := 0; i < nLeafs; i++ { k := BigIntToBytes(bLen, big.NewInt(int64(i)))@@ -948,7 +948,7 @@ }
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree2, err := NewTree(database2, 100, HashFunctionPoseidon) + tree2, err := NewTree(database2, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree2.db.Close() //nolint:errcheck tree2.dbgInit()@@ -962,7 +962,7 @@
// use tree3 to add nil value array database3, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree3, err := NewTree(database3, 100, HashFunctionPoseidon) + tree3, err := NewTree(database3, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree3.db.Close() //nolint:errcheck
M
circomproofs_test.go
→
circomproofs_test.go
@@ -17,14 +17,13 @@ tree, err := NewTree(database, 4, HashFunctionPoseidon)
c.Assert(err, qt.IsNil) defer tree.db.Close() //nolint:errcheck - bLen := tree.HashFunction().Len() - testVector := [][]int64{ {1, 11}, {2, 22}, {3, 33}, {4, 44}, } + bLen := 1 for i := 0; i < len(testVector); i++ { k := BigIntToBytes(bLen, big.NewInt(testVector[i][0])) v := BigIntToBytes(bLen, big.NewInt(testVector[i][1]))
M
testvectors/circom/go-data-generator/generator_test.go
→
testvectors/circom/go-data-generator/generator_test.go
@@ -18,14 +18,13 @@ c.Assert(err, qt.IsNil)
tree, err := arbo.NewTree(database, 4, arbo.HashFunctionPoseidon) c.Assert(err, qt.IsNil) - bLen := tree.HashFunction().Len() - testVector := [][]int64{ {1, 11}, {2, 22}, {3, 33}, {4, 44}, } + bLen := 1 for i := 0; i < len(testVector); i++ { k := arbo.BigIntToBytes(bLen, big.NewInt(testVector[i][0])) v := arbo.BigIntToBytes(bLen, big.NewInt(testVector[i][1]))
M
tree.go
→
tree.go
@@ -229,8 +229,10 @@ }
} // store root (from the vt) to db - if err := wTx.Set(dbKeyRoot, vt.root.h); err != nil { - return nil, err + if vt.root != nil { + if err := wTx.Set(dbKeyRoot, vt.root.h); err != nil { + return nil, err + } } // update nLeafs@@ -310,14 +312,34 @@ }
return nil } -func (t *Tree) add(wTx db.WriteTx, root []byte, fromLvl int, k, v []byte) ([]byte, error) { - keyPath := make([]byte, int(math.Ceil(float64(t.maxLevels)/float64(8)))) //nolint:gomnd +// keyPathFromKey returns the keyPath and checks that the key is not bigger +// than maximum key length for the tree maxLevels size. +// This is because if the key bits length is bigger than the maxLevels of the +// tree, two different keys that their difference is at the end, will collision +// in the same leaf of the tree (at the max depth). +func keyPathFromKey(maxLevels int, k []byte) ([]byte, error) { + maxKeyLen := int(math.Ceil(float64(maxLevels) / float64(8))) //nolint:gomnd + if len(k) > maxKeyLen { + return nil, fmt.Errorf("len(k) can not be bigger than ceil(maxLevels/8), where"+ + " len(k): %d, maxLevels: %d, max key len=ceil(maxLevels/8): %d. Might need"+ + " a bigger tree depth (maxLevels>=%d) in order to input keys of length %d", + len(k), maxLevels, maxKeyLen, len(k)*8, len(k)) //nolint:gomnd + } + keyPath := make([]byte, maxKeyLen) //nolint:gomnd copy(keyPath[:], k) + return keyPath, nil +} + +func (t *Tree) add(wTx db.WriteTx, root []byte, fromLvl int, k, v []byte) ([]byte, error) { + keyPath, err := keyPathFromKey(t.maxLevels, k) + if err != nil { + return nil, err + } path := getPath(t.maxLevels, keyPath) // go down to the leaf var siblings [][]byte - _, _, siblings, err := t.down(wTx, k, root, siblings, path, fromLvl, false) + _, _, siblings, err = t.down(wTx, k, root, siblings, path, fromLvl, false) if err != nil { return nil, err }@@ -590,8 +612,10 @@ if !t.editable() {
return ErrSnapshotNotEditable } - keyPath := make([]byte, int(math.Ceil(float64(t.maxLevels)/float64(8)))) //nolint:gomnd - copy(keyPath[:], k) + keyPath, err := keyPathFromKey(t.maxLevels, k) + if err != nil { + return err + } path := getPath(t.maxLevels, keyPath) root, err := t.RootWithTx(wTx)@@ -647,8 +671,10 @@
// GenProofWithTx does the same than the GenProof method, but allowing to pass // the db.ReadTx that is used. func (t *Tree) GenProofWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, []byte, bool, error) { - keyPath := make([]byte, int(math.Ceil(float64(t.maxLevels)/float64(8)))) //nolint:gomnd - copy(keyPath[:], k) + keyPath, err := keyPathFromKey(t.maxLevels, k) + if err != nil { + return nil, nil, nil, false, err + } path := getPath(t.maxLevels, keyPath) root, err := t.RootWithTx(rTx)@@ -782,8 +808,10 @@ // db.ReadTx that is used. If the key is not found, will return the error
// ErrKeyNotFound, and in the leafK & leafV parameters will be placed the data // found in the tree in the leaf that was on the path going to the input key. func (t *Tree) GetWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, error) { - keyPath := make([]byte, int(math.Ceil(float64(t.maxLevels)/float64(8)))) //nolint:gomnd - copy(keyPath[:], k) + keyPath, err := keyPathFromKey(t.maxLevels, k) + if err != nil { + return nil, nil, err + } path := getPath(t.maxLevels, keyPath) root, err := t.RootWithTx(rTx)
M
tree_test.go
→
tree_test.go
@@ -2,7 +2,9 @@ package arbo
import ( "encoding/hex" + "math" "math/big" + "runtime" "testing" "time"@@ -60,7 +62,7 @@
func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) { database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 10, hashFunc) + tree, err := NewTree(database, 256, hashFunc) c.Assert(err, qt.IsNil) defer tree.db.Close() //nolint:errcheck@@ -68,7 +70,7 @@ root, err := tree.Root()
c.Assert(err, qt.IsNil) c.Check(hex.EncodeToString(root), qt.Equals, testVectors[0]) - bLen := hashFunc.Len() + bLen := 32 err = tree.Add( BigIntToBytes(bLen, big.NewInt(1)), BigIntToBytes(bLen, big.NewInt(2)))@@ -92,11 +94,11 @@ func TestAddBatch(t *testing.T) {
c := qt.New(t) database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 100, HashFunctionPoseidon) + tree, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree.db.Close() //nolint:errcheck - bLen := tree.HashFunction().Len() + bLen := 32 for i := 0; i < 1000; i++ { k := BigIntToBytes(bLen, big.NewInt(int64(i))) v := BigIntToBytes(bLen, big.NewInt(0))@@ -110,7 +112,7 @@ "296519252211642170490407814696803112091039265640052570497930797516015811235")
database, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree2, err := NewTree(database, 100, HashFunctionPoseidon) + tree2, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree2.db.Close() //nolint:errcheck@@ -133,11 +135,11 @@ func TestAddDifferentOrder(t *testing.T) {
c := qt.New(t) database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree1, err := NewTree(database1, 100, HashFunctionPoseidon) + tree1, err := NewTree(database1, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree1.db.Close() //nolint:errcheck - bLen := tree1.HashFunction().Len() + bLen := 32 for i := 0; i < 16; i++ { k := BigIntToBytes(bLen, big.NewInt(int64(i))) v := BigIntToBytes(bLen, big.NewInt(0))@@ -148,7 +150,7 @@ }
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree2, err := NewTree(database2, 100, HashFunctionPoseidon) + tree2, err := NewTree(database2, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree2.db.Close() //nolint:errcheck@@ -173,11 +175,11 @@ func TestAddRepeatedIndex(t *testing.T) {
c := qt.New(t) database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 100, HashFunctionPoseidon) + tree, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree.db.Close() //nolint:errcheck - bLen := tree.HashFunction().Len() + bLen := 32 k := BigIntToBytes(bLen, big.NewInt(int64(3))) v := BigIntToBytes(bLen, big.NewInt(int64(12)))@@ -191,11 +193,11 @@ func TestUpdate(t *testing.T) {
c := qt.New(t) database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 100, HashFunctionPoseidon) + tree, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree.db.Close() //nolint:errcheck - bLen := tree.HashFunction().Len() + bLen := 32 k := BigIntToBytes(bLen, big.NewInt(int64(20))) v := BigIntToBytes(bLen, big.NewInt(int64(12))) if err := tree.Add(k, v); err != nil {@@ -244,11 +246,11 @@ func TestAux(t *testing.T) { // TODO split in proper tests
c := qt.New(t) database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 100, HashFunctionPoseidon) + tree, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree.db.Close() //nolint:errcheck - bLen := tree.HashFunction().Len() + bLen := 32 k := BigIntToBytes(bLen, big.NewInt(int64(1))) v := BigIntToBytes(bLen, big.NewInt(int64(0))) err = tree.Add(k, v)@@ -283,11 +285,11 @@ func TestGet(t *testing.T) {
c := qt.New(t) database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 100, HashFunctionPoseidon) + tree, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree.db.Close() //nolint:errcheck - bLen := tree.HashFunction().Len() + bLen := 32 for i := 0; i < 10; i++ { k := BigIntToBytes(bLen, big.NewInt(int64(i))) v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))@@ -307,11 +309,11 @@ func TestGenProofAndVerify(t *testing.T) {
c := qt.New(t) database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 100, HashFunctionPoseidon) + tree, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree.db.Close() //nolint:errcheck - bLen := tree.HashFunction().Len() - 1 + bLen := 32 for i := 0; i < 10; i++ { k := BigIntToBytes(bLen, big.NewInt(int64(i))) v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))@@ -339,11 +341,11 @@ func TestDumpAndImportDump(t *testing.T) {
c := qt.New(t) database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree1, err := NewTree(database1, 100, HashFunctionPoseidon) + tree1, err := NewTree(database1, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree1.db.Close() //nolint:errcheck - bLen := tree1.HashFunction().Len() + bLen := 32 for i := 0; i < 16; i++ { k := BigIntToBytes(bLen, big.NewInt(int64(i))) v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))@@ -357,7 +359,7 @@ c.Assert(err, qt.IsNil)
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree2, err := NewTree(database2, 100, HashFunctionPoseidon) + tree2, err := NewTree(database2, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree2.db.Close() //nolint:errcheck err = tree2.ImportDump(e)@@ -376,11 +378,11 @@ func TestRWMutex(t *testing.T) {
c := qt.New(t) database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 100, HashFunctionPoseidon) + tree, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree.db.Close() //nolint:errcheck - bLen := tree.HashFunction().Len() + bLen := 32 var keys, values [][]byte for i := 0; i < 1000; i++ { k := BigIntToBytes(bLen, big.NewInt(int64(i)))@@ -469,7 +471,7 @@ c.Assert(err, qt.IsNil)
var keys, values [][]byte for i := 0; i < 16; i++ { - k := BigIntToBytes(32, big.NewInt(int64(i))) + k := BigIntToBytes(1, big.NewInt(int64(i))) v := k keys = append(keys, k)@@ -492,10 +494,10 @@ c.Assert(root1, qt.DeepEquals, root2)
// get all key-values and check that are equal between both trees for i := 0; i < 16; i++ { - auxK1, auxV1, err := tree1.Get(BigIntToBytes(32, big.NewInt(int64(i)))) + auxK1, auxV1, err := tree1.Get(BigIntToBytes(1, big.NewInt(int64(i)))) c.Assert(err, qt.IsNil) - auxK2, auxV2, err := tree2.Get(BigIntToBytes(32, big.NewInt(int64(i)))) + auxK2, auxV2, err := tree2.Get(BigIntToBytes(1, big.NewInt(int64(i)))) c.Assert(err, qt.IsNil) c.Assert(auxK1, qt.DeepEquals, auxK2)@@ -504,7 +506,7 @@ }
// try adding one more key to both trees (through Add & AddBatch) and // expect not being added due the tree is already full - k := BigIntToBytes(32, big.NewInt(int64(16))) + k := BigIntToBytes(1, big.NewInt(int64(16))) v := k err = tree1.Add(k, v) c.Assert(err, qt.Equals, ErrMaxVirtualLevel)@@ -518,13 +520,13 @@ func TestSetRoot(t *testing.T) {
c := qt.New(t) database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 100, HashFunctionPoseidon) + tree, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) expectedRoot := "13742386369878513332697380582061714160370929283209286127733983161245560237407" // fill the tree - bLen := tree.HashFunction().Len() + bLen := 32 var keys, values [][]byte for i := 0; i < 1000; i++ { k := BigIntToBytes(bLen, big.NewInt(int64(i)))@@ -574,11 +576,11 @@ func TestSnapshot(t *testing.T) {
c := qt.New(t) database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 100, HashFunctionPoseidon) + tree, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) // fill the tree - bLen := tree.HashFunction().Len() + bLen := 32 var keys, values [][]byte for i := 0; i < 1000; i++ { k := BigIntToBytes(bLen, big.NewInt(int64(i)))@@ -624,11 +626,11 @@ c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) c.Assert(err, qt.IsNil) - tree, err := NewTree(database, 100, HashFunctionPoseidon) + tree, err := NewTree(database, 256, HashFunctionPoseidon) c.Assert(err, qt.IsNil) defer tree.db.Close() //nolint:errcheck - bLen := tree.HashFunction().Len() + bLen := 32 k := BigIntToBytes(bLen, big.NewInt(int64(3))) root, err := tree.Root()@@ -646,7 +648,7 @@ database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil) // maxLevels is 100, keyPath length = ceil(maxLevels/8) = 13 maxLevels := 100 - tree, err := NewTree(database, maxLevels, HashFunctionPoseidon) + tree, err := NewTree(database, maxLevels, HashFunctionBlake2b) c.Assert(err, qt.IsNil) // expect no errors when adding a key of only 4 bytes (when the@@ -672,6 +674,75 @@ v = BigIntToBytes(bLen, big.NewInt(2))
invalids, err := tree.AddBatch([][]byte{k}, [][]byte{v}) c.Assert(err, qt.IsNil) c.Assert(len(invalids), qt.Equals, 0) + + // expect errors when adding a key bigger than maximum capacity of the + // tree: ceil(maxLevels/8) + maxLevels = 32 + database, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()}) + c.Assert(err, qt.IsNil) + tree, err = NewTree(database, maxLevels, HashFunctionBlake2b) + c.Assert(err, qt.IsNil) + + maxKeyLen := int(math.Ceil(float64(maxLevels) / float64(8))) //nolint:gomnd + k = BigIntToBytes(maxKeyLen+1, big.NewInt(1)) + v = BigIntToBytes(maxKeyLen+1, big.NewInt(1)) + + expectedErrMsg := "len(k) can not be bigger than ceil(maxLevels/8)," + + " where len(k): 5, maxLevels: 32, max key len=ceil(maxLevels/8): 4." + + " Might need a bigger tree depth (maxLevels>=40) in order to input" + + " keys of length 5" + + err = tree.Add(k, v) + c.Assert(err.Error(), qt.Equals, expectedErrMsg) + + err = tree.Update(k, v) + c.Assert(err.Error(), qt.Equals, expectedErrMsg) + + _, _, _, _, err = tree.GenProof(k) + c.Assert(err.Error(), qt.Equals, expectedErrMsg) + + _, _, err = tree.Get(k) + c.Assert(err.Error(), qt.Equals, expectedErrMsg) + + // check AddBatch with few key-values + invalids, err = tree.AddBatch([][]byte{k}, [][]byte{v}) + c.Assert(err, qt.IsNil) + c.Assert(len(invalids), qt.Equals, 1) + + // check AddBatch with many key-values + nCPU := flp2(runtime.NumCPU()) + nKVs := nCPU + 1 + var ks, vs [][]byte + for i := 0; i < nKVs; i++ { + ks = append(ks, BigIntToBytes(maxKeyLen+1, big.NewInt(1))) + vs = append(vs, BigIntToBytes(maxKeyLen+1, big.NewInt(1))) + } + invalids, err = tree.AddBatch(ks, vs) + c.Assert(err, qt.IsNil) + c.Assert(len(invalids), qt.Equals, nKVs) + + // check that with maxKeyLen it can be added + k = BigIntToBytes(maxKeyLen, big.NewInt(1)) + err = tree.Add(k, v) + c.Assert(err, qt.IsNil) + + // check CheckProof check with key longer + kAux, vAux, packedSiblings, existence, err := tree.GenProof(k) + c.Assert(err, qt.IsNil) + c.Assert(existence, qt.IsTrue) + + root, err := tree.Root() + c.Assert(err, qt.IsNil) + verif, err := CheckProof(tree.HashFunction(), kAux, vAux, root, packedSiblings) + c.Assert(err, qt.IsNil) + c.Assert(verif, qt.IsTrue) + + // use a similar key but with one zero, expect that CheckProof fails on + // the verification + kAux = append(kAux, 0) + verif, err = CheckProof(tree.HashFunction(), kAux, vAux, root, packedSiblings) + c.Assert(err, qt.IsNil) + c.Assert(verif, qt.IsFalse) } func BenchmarkAdd(b *testing.B) {
M
vt.go
→
vt.go
@@ -37,22 +37,32 @@ k []byte
v []byte } -func (p *params) keysValuesToKvs(ks, vs [][]byte) ([]kv, error) { +func (p *params) keysValuesToKvs(ks, vs [][]byte) ([]kv, []int, error) { if len(ks) != len(vs) { - return nil, fmt.Errorf("len(keys)!=len(values) (%d!=%d)", + return nil, nil, fmt.Errorf("len(keys)!=len(values) (%d!=%d)", len(ks), len(vs)) } - kvs := make([]kv, len(ks)) + var invalids []int + var kvs []kv for i := 0; i < len(ks); i++ { - keyPath := make([]byte, int(math.Ceil(float64(p.maxLevels)/float64(8)))) //nolint:gomnd - copy(keyPath[:], ks[i]) - kvs[i].pos = i - kvs[i].keyPath = keyPath - kvs[i].k = ks[i] - kvs[i].v = vs[i] + keyPath, err := keyPathFromKey(p.maxLevels, ks[i]) + if err != nil { + // TODO in a future iteration, invalids will contain + // the reason of the error of why each index is + // invalid. + invalids = append(invalids, i) + continue + } + + var kvsI kv + kvsI.pos = i + kvsI.keyPath = keyPath + kvsI.k = ks[i] + kvsI.v = vs[i] + kvs = append(kvs, kvsI) } - return kvs, nil + return kvs, invalids, nil } // vt stands for virtual tree. It's a tree that does not have any computed hash@@ -94,9 +104,9 @@ }
l := int(math.Log2(float64(nCPU))) - kvs, err := t.params.keysValuesToKvs(ks, vs) + kvs, invalids, err := t.params.keysValuesToKvs(ks, vs) if err != nil { - return nil, err + return invalids, err } buckets := splitInBuckets(kvs, nCPU)@@ -186,7 +196,6 @@ }(i)
} wg.Wait() - var invalids []int for i := 0; i < len(invalidsInBucket); i++ { invalids = append(invalids, invalidsInBucket[i]...) }@@ -284,7 +293,10 @@ }
// add adds a key&value as a leaf in the VirtualTree func (t *vt) add(fromLvl int, k, v []byte) error { - leaf := newLeafNode(t.params, k, v) + leaf, err := newLeafNode(t.params, k, v) + if err != nil { + return err + } if t.root == nil { t.root = leaf return nil@@ -366,16 +378,18 @@
return pairs, nil } -func newLeafNode(p *params, k, v []byte) *node { - keyPath := make([]byte, p.hashFunction.Len()) - copy(keyPath[:], k) +func newLeafNode(p *params, k, v []byte) (*node, error) { + keyPath, err := keyPathFromKey(p.maxLevels, k) + if err != nil { + return nil, err + } path := getPath(p.maxLevels, keyPath) n := &node{ k: k, v: v, path: path, } - return n + return n, nil } type virtualNodeType int
M
vt_test.go
→
vt_test.go
@@ -2,6 +2,7 @@ package arbo
import ( "encoding/hex" + "math" "math/big" "testing"@@ -9,28 +10,62 @@ qt "github.com/frankban/quicktest"
"go.vocdoni.io/dvote/db/badgerdb" ) +// testVirtualTree adds the given key-values and tests the vt root against the +// Tree +func testVirtualTree(c *qt.C, maxLevels int, keys, values [][]byte) { + c.Assert(len(keys), qt.Equals, len(values)) + + // normal tree, to have an expected root value + database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) + c.Assert(err, qt.IsNil) + tree, err := NewTree(database, maxLevels, HashFunctionSha256) + c.Assert(err, qt.IsNil) + for i := 0; i < len(keys); i++ { + err := tree.Add(keys[i], values[i]) + c.Assert(err, qt.IsNil) + } + + // virtual tree + vTree := newVT(maxLevels, HashFunctionSha256) + + c.Assert(vTree.root, qt.IsNil) + + for i := 0; i < len(keys); i++ { + err := vTree.add(0, keys[i], values[i]) + c.Assert(err, qt.IsNil) + } + + // compute hashes, and check Root + _, err = vTree.computeHashes() + c.Assert(err, qt.IsNil) + root, err := tree.Root() + c.Assert(err, qt.IsNil) + c.Assert(vTree.root.h, qt.DeepEquals, root) +} + func TestVirtualTreeTestVectors(t *testing.T) { c := qt.New(t) - bLen := 32 + maxLevels := 32 + keyLen := int(math.Ceil(float64(maxLevels) / float64(8))) //nolint:gomnd keys := [][]byte{ - BigIntToBytes(bLen, big.NewInt(1)), - BigIntToBytes(bLen, big.NewInt(33)), - BigIntToBytes(bLen, big.NewInt(1234)), - BigIntToBytes(bLen, big.NewInt(123456789)), + BigIntToBytes(keyLen, big.NewInt(1)), + BigIntToBytes(keyLen, big.NewInt(33)), + BigIntToBytes(keyLen, big.NewInt(1234)), + BigIntToBytes(keyLen, big.NewInt(123456789)), } values := [][]byte{ - BigIntToBytes(bLen, big.NewInt(2)), - BigIntToBytes(bLen, big.NewInt(44)), - BigIntToBytes(bLen, big.NewInt(9876)), - BigIntToBytes(bLen, big.NewInt(987654321)), + BigIntToBytes(keyLen, big.NewInt(2)), + BigIntToBytes(keyLen, big.NewInt(44)), + BigIntToBytes(keyLen, big.NewInt(9876)), + BigIntToBytes(keyLen, big.NewInt(987654321)), } // check the root for different batches of leafs - testVirtualTree(c, 10, keys[:1], values[:1]) - testVirtualTree(c, 10, keys[:2], values[:2]) - testVirtualTree(c, 10, keys[:3], values[:3]) - testVirtualTree(c, 10, keys[:4], values[:4]) + testVirtualTree(c, maxLevels, keys[:1], values[:1]) + testVirtualTree(c, maxLevels, keys[:2], values[:2]) + testVirtualTree(c, maxLevels, keys[:3], values[:3]) + testVirtualTree(c, maxLevels, keys[:4], values[:4]) // test with hardcoded values testvectorKeys := []string{@@ -53,8 +88,8 @@ values = append(values, []byte{0})
} // check the root for different batches of leafs - testVirtualTree(c, 10, keys[:1], values[:1]) - testVirtualTree(c, 10, keys, values) + testVirtualTree(c, 256, keys[:1], values[:1]) + testVirtualTree(c, 256, keys, values) } func TestVirtualTreeRandomKeys(t *testing.T) {@@ -69,45 +104,14 @@ keys[i] = randomBytes(32)
values[i] = randomBytes(32) } - testVirtualTree(c, 100, keys, values) -} - -func testVirtualTree(c *qt.C, maxLevels int, keys, values [][]byte) { - c.Assert(len(keys), qt.Equals, len(values)) - - // normal tree, to have an expected root value - database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()}) - c.Assert(err, qt.IsNil) - tree, err := NewTree(database, maxLevels, HashFunctionSha256) - c.Assert(err, qt.IsNil) - for i := 0; i < len(keys); i++ { - err := tree.Add(keys[i], values[i]) - c.Assert(err, qt.IsNil) - } - - // virtual tree - vTree := newVT(maxLevels, HashFunctionSha256) - - c.Assert(vTree.root, qt.IsNil) - - for i := 0; i < len(keys); i++ { - err := vTree.add(0, keys[i], values[i]) - c.Assert(err, qt.IsNil) - } - - // compute hashes, and check Root - _, err = vTree.computeHashes() - c.Assert(err, qt.IsNil) - root, err := tree.Root() - c.Assert(err, qt.IsNil) - c.Assert(vTree.root.h, qt.DeepEquals, root) + testVirtualTree(c, 256, keys, values) } func TestVirtualTreeAddBatch(t *testing.T) { c := qt.New(t) nLeafs := 2000 - maxLevels := 100 + maxLevels := 256 keys := make([][]byte, nLeafs) values := make([][]byte, nLeafs)@@ -151,7 +155,7 @@ vTree2 := newVT(7, HashFunctionPoseidon) // used for addBatch
var keys, values [][]byte for i := 0; i < 128; i++ { - k := BigIntToBytes(32, big.NewInt(int64(i))) + k := BigIntToBytes(1, big.NewInt(int64(i))) v := k keys = append(keys, k)