arnaucube
arnaucube/arboPublic Self-hosted repository
Merge pull request #21 from vocdoni/feature/keyPath-key-len-check

Update keyPath to ceil(maxLvl/8), add checks that len(key)<=maxKeyLen
arnau 17317030+arnaucube@users.noreply.github.com
Tue, 05 Oct 2021 18:08:55 +0200
commit

63f7e8866e1ca24f6cdec8ebdde212fdaf99faf1

parent

9eb7c8e265c430f89a0625fb84066f407c18c0d9

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,17 +312,34 @@ }

return nil } -func (t *Tree) add(wTx db.WriteTx, root []byte, fromLvl int, k, v []byte) ([]byte, error) { - keyPath := make([]byte, t.hashFunction.Len()) - // if len(k) > t.hashFunction.Len() { // WIP - // return nil, fmt.Errorf("len(k) > hashFunction.Len()") - // } +// 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 }

@@ -386,12 +405,10 @@ if bytes.Equal(newKey, oldLeafKey) {

return nil, nil, nil, ErrKeyAlreadyExists } - oldLeafKeyFull := make([]byte, t.hashFunction.Len()) - // if len(oldLeafKey) > t.hashFunction.Len() { // WIP - // return nil, nil, nil, - // fmt.Errorf("len(oldLeafKey) > hashFunction.Len()") - // } - copy(oldLeafKeyFull[:], oldLeafKey) + oldLeafKeyFull, err := keyPathFromKey(t.maxLevels, oldLeafKey) + if err != nil { + return nil, nil, nil, err + } // if currKey is already used, go down until paths diverge oldPath := getPath(t.maxLevels, oldLeafKeyFull)

@@ -593,13 +610,10 @@ if !t.editable() {

return ErrSnapshotNotEditable } - var err error - - keyPath := make([]byte, t.hashFunction.Len()) - // if len(k) > t.hashFunction.Len() { // WIP - // return fmt.Errorf("len(k) > hashFunction.Len()") - // } - copy(keyPath[:], k) + keyPath, err := keyPathFromKey(t.maxLevels, k) + if err != nil { + return err + } path := getPath(t.maxLevels, keyPath) root, err := t.RootWithTx(wTx)

@@ -655,18 +669,17 @@

// 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, t.hashFunction.Len()) - // if len(k) > t.hashFunction.Len() { // WIP - // return nil, nil, nil, false, fmt.Errorf("len(k) > hashFunction.Len()") - // } - 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) if err != nil { return nil, nil, nil, false, err } - path := getPath(t.maxLevels, keyPath) // go down to the leaf var siblings [][]byte _, value, siblings, err := t.down(rTx, k, root, siblings, path, 0, true)

@@ -793,18 +806,17 @@ // 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, t.hashFunction.Len()) - // if len(k) > t.hashFunction.Len() { // WIP - // return nil, nil, fmt.Errorf("len(k) > hashFunction.Len()") - // } - 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) if err != nil { return nil, nil, err } - path := getPath(t.maxLevels, keyPath) // go down to the leaf var siblings [][]byte _, value, _, err := t.down(rTx, k, root, siblings, path, 0, true)

@@ -827,7 +839,7 @@ if err != nil {

return false, err } - keyPath := make([]byte, hashFunc.Len()) + keyPath := make([]byte, int(math.Ceil(float64(len(siblings))/float64(8)))) //nolint:gomnd copy(keyPath[:], k) key, _, err := newLeafValue(hashFunc, k, v)
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()

@@ -638,6 +640,133 @@ c.Assert(err, qt.IsNil)

_, _, err = tree.Get(k) c.Assert(err, qt.Equals, ErrKeyNotFound) // and not equal to db.ErrKeyNotFound +} + +func TestKeyLen(t *testing.T) { + c := qt.New(t) + 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, HashFunctionBlake2b) + c.Assert(err, qt.IsNil) + + // expect no errors when adding a key of only 4 bytes (when the + // required length of keyPath for 100 levels would be 13 bytes) + bLen := 4 + k := BigIntToBytes(bLen, big.NewInt(1)) + v := BigIntToBytes(bLen, big.NewInt(1)) + + err = tree.Add(k, v) + c.Assert(err, qt.IsNil) + + err = tree.Update(k, v) + c.Assert(err, qt.IsNil) + + _, _, _, _, err = tree.GenProof(k) + c.Assert(err, qt.IsNil) + + _, _, err = tree.Get(k) + c.Assert(err, qt.IsNil) + + k = BigIntToBytes(bLen, big.NewInt(2)) + 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 TestKeyLenBiggerThan32(t *testing.T) { + c := qt.New(t) + maxLevels := 264 + 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) + + bLen := 33 + err = tree.Add( + randomBytes(bLen), + randomBytes(bLen)) + c.Assert(err, qt.IsNil) + + // 2nd key that we add, will find a node with len(key)==32 (due the + // hash output size, expect that next Add does not give any error, as + // internally it will use a keyPath of size corresponent to the + // maxLevels size of the tree + err = tree.Add( + randomBytes(bLen), + randomBytes(bLen)) + c.Assert(err, qt.IsNil) } func BenchmarkAdd(b *testing.B) {
M utils.go → utils.go

@@ -1,6 +1,8 @@

package arbo -import "math/big" +import ( + "math/big" +) // SwapEndianness swaps the order of the bytes in the byte slice. func SwapEndianness(b []byte) []byte {
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, p.hashFunction.Len()) - 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)