Add StateDB set & get Idx by EthAddr & BJJ methods

common/account.go

@@ -5,6 +5,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"math/big"
+	"strconv"
 
 	ethCommon "github.com/ethereum/go-ethereum/common"
 	"github.com/iden3/go-iden3-crypto/babyjub"
@@ -34,6 +35,11 @@ const (
 // Idx represents the account Index in the MerkleTree
 type Idx uint32
 
+// String returns a string representation of the Idx
+func (idx Idx) String() string {
+	return strconv.Itoa(int(idx))
+}
+
 // Bytes returns a byte array representing the Idx
 func (idx Idx) Bytes() []byte {
 	var b [4]byte

db/statedb/statedb.go

@@ -23,6 +23,10 @@ var ErrStateDBWithoutMT = errors.New("Can not call method to use MerkleTree in a
 // already exists
 var ErrAccountAlreadyExists = errors.New("Can not CreateAccount because Account already exists")
 
+// ErrToIdxNotFound is used when trying to get the ToIdx from ToEthAddr or
+// ToEthAddr&ToBJJ
+var ErrToIdxNotFound = errors.New("ToIdx can not be found")
+
 // KeyCurrentBatch is used as key in the db to store the current BatchNum
 var KeyCurrentBatch = []byte("currentbatch")
 
@@ -248,7 +252,13 @@ func getAccountInTreeDB(sto db.Storage, idx common.Idx) (*common.Account, error)
 // StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
 // MerkleTree, returning a CircomProcessorProof.
 func (s *StateDB) CreateAccount(idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
-	return createAccountInTreeDB(s.db, s.mt, idx, account)
+	cpp, err := createAccountInTreeDB(s.db, s.mt, idx, account)
+	if err != nil {
+		return cpp, err
+	}
+	// store idx by EthAddr & BJJ
+	err = s.setIdxByEthAddrBJJ(idx, account.EthAddr, account.PublicKey)
+	return cpp, err
 }
 
 // createAccountInTreeDB is abstracted from StateDB to be used from StateDB and

db/statedb/txprocessors.go

@@ -3,7 +3,6 @@ package statedb
 import (
 	"bytes"
 	"errors"
-	"fmt"
 	"math/big"
 
 	ethCommon "github.com/ethereum/go-ethereum/common"
@@ -20,7 +19,8 @@ var (
 	// keyidx is used as key in the db to store the current Idx
 	keyidx = []byte("idx")
 
-	ffAddr = ethCommon.HexToAddress("0xffffffffffffffffffffffffffffffffffffffff")
+	// ffAddr    = ethCommon.HexToAddress("0xffffffffffffffffffffffffffffffffffffffff")
+	emptyAddr = ethCommon.HexToAddress("0x0000000000000000000000000000000000000000")
 )
 
 func (s *StateDB) resetZKInputs() {
@@ -322,6 +322,7 @@ func (s *StateDB) processL1Tx(exitTree *merkletree.MerkleTree, tx *common.L1Tx)
 // related to the Exit (in case of): Idx, ExitAccount, boolean determining if
 // the Exit created a new Leaf in the ExitTree.
 func (s *StateDB) processL2Tx(exitTree *merkletree.MerkleTree, tx *common.PoolL2Tx) (*common.Idx, *common.Account, bool, error) {
+	var err error
 	// ZKInputs
 	if s.zki != nil {
 		// Txs
@@ -332,18 +333,22 @@ func (s *StateDB) processL2Tx(exitTree *merkletree.MerkleTree, tx *common.PoolL2
 
 		// fill AuxToIdx if needed
 		if tx.ToIdx == common.Idx(0) {
-			// Idx not set in the Tx, get it from DB through ToEthAddr or ToBJJ
 			var idx common.Idx
-			if !bytes.Equal(tx.ToEthAddr.Bytes(), ffAddr.Bytes()) {
-				idx = s.GetIdxByEthAddr(tx.ToEthAddr)
-				if idx == common.Idx(0) {
-					return nil, nil, false, fmt.Errorf("Idx can not be found for given tx.FromEthAddr")
+			if !bytes.Equal(tx.ToEthAddr.Bytes(), emptyAddr.Bytes()) && tx.ToBJJ == nil {
+				// case ToEthAddr!=0 && ToBJJ=0
+				idx, err = s.GetIdxByEthAddr(tx.ToEthAddr)
+				if err != nil {
+					return nil, nil, false, ErrToIdxNotFound
 				}
-			} else {
-				idx = s.GetIdxByBJJ(tx.ToBJJ)
-				if idx == common.Idx(0) {
-					return nil, nil, false, fmt.Errorf("Idx can not be found for given tx.FromBJJ")
+			} else if !bytes.Equal(tx.ToEthAddr.Bytes(), emptyAddr.Bytes()) && tx.ToBJJ != nil {
+				// case ToEthAddr!=0 && ToBJJ!=0
+				idx, err = s.GetIdxByEthAddrBJJ(tx.ToEthAddr, tx.ToBJJ)
+				if err != nil {
+					return nil, nil, false, ErrToIdxNotFound
 				}
+			} else {
+				// rest of cases (included case ToEthAddr==0) are not possible
+				return nil, nil, false, ErrToIdxNotFound
 			}
 			s.zki.AuxToIdx[s.i] = idx.BigInt()
 		}
@@ -367,7 +372,7 @@ func (s *StateDB) processL2Tx(exitTree *merkletree.MerkleTree, tx *common.PoolL2
 	case common.TxTypeTransfer:
 		// go to the MT account of sender and receiver, and update
 		// balance & nonce
-		err := s.applyTransfer(tx.Tx())
+		err = s.applyTransfer(tx.Tx())
 		if err != nil {
 			return nil, nil, false, err
 		}

db/statedb/utils.go

@@ -5,32 +5,98 @@ import (
 
 	ethCommon "github.com/ethereum/go-ethereum/common"
 	"github.com/hermeznetwork/hermez-node/common"
+	"github.com/hermeznetwork/hermez-node/log"
 	"github.com/iden3/go-iden3-crypto/babyjub"
 	"github.com/iden3/go-merkletree"
 )
 
-// GetIdxByEthAddr returns the smallest Idx in the StateDB for the given
-// Ethereum Address. Will return 0 in case that Idx is not found in the
-// StateDB.
-func (s *StateDB) GetIdxByEthAddr(addr ethCommon.Address) common.Idx {
-	// TODO
-	return common.Idx(0)
+func concatEthAddrBJJ(addr ethCommon.Address, pk *babyjub.PublicKey) []byte {
+	pkComp := pk.Compress()
+	var b []byte
+	b = append(b, addr.Bytes()...)
+	b = append(b[:], pkComp[:]...)
+	return b
+}
+
+// setIdxByEthAddrBJJ stores the given Idx in the StateDB as follows:
+// - key: Eth Address, value: idx
+// - key: EthAddr & BabyJubJub PublicKey Compressed, value: idx
+// If Idx already exist for the given EthAddr & BJJ, the remaining Idx will be
+// always the smallest one.
+func (s *StateDB) setIdxByEthAddrBJJ(idx common.Idx, addr ethCommon.Address, pk *babyjub.PublicKey) error {
+	oldIdx, err := s.GetIdxByEthAddrBJJ(addr, pk)
+	if err == nil {
+		// EthAddr & BJJ already have an Idx
+		// check which Idx is smaller
+		// if new idx is smaller, store the new one
+		// if new idx is bigger, don't store and return, as the used one will be the old
+		if idx >= oldIdx {
+			log.Debug("StateDB.setIdxByEthAddrBJJ: Idx not stored because there already exist a smaller Idx for the given EthAddr & BJJ")
+			return nil
+		}
+	}
+
+	// store idx for EthAddr & BJJ assuming that EthAddr & BJJ still don't
+	// have an Idx stored in the DB, and if so, the already stored Idx is
+	// bigger than the given one, so should be updated to the new one
+	// (smaller)
+	tx, err := s.db.NewTx()
+	if err != nil {
+		return err
+	}
+	k := concatEthAddrBJJ(addr, pk)
+	// store Addr&BJJ-idx
+	err = tx.Put(k, idx.Bytes())
+	if err != nil {
+		return err
+	}
+	// store Addr-idx
+	err = tx.Put(addr.Bytes(), idx.Bytes())
+	if err != nil {
+		return err
+	}
+	err = tx.Commit()
+	if err != nil {
+		return err
+	}
+	return nil
 }
 
-// GetIdxByBJJ returns the smallest Idx in the StateDB for the given BabyJubJub
-// PublicKey. Will return 0 in case that Idx is not found in the StateDB.
-func (s *StateDB) GetIdxByBJJ(pk *babyjub.PublicKey) common.Idx {
-	// TODO
-	return common.Idx(0)
+// GetIdxByEthAddr returns the smallest Idx in the StateDB for the given
+// Ethereum Address. Will return common.Idx(0) and error in case that Idx is
+// not found in the StateDB.
+func (s *StateDB) GetIdxByEthAddr(addr ethCommon.Address) (common.Idx, error) {
+	b, err := s.db.Get(addr.Bytes())
+	if err != nil {
+		return common.Idx(0), err
+	}
+	idx, err := common.IdxFromBytes(b)
+	if err != nil {
+		return common.Idx(0), err
+	}
+	return idx, nil
 }
 
 // GetIdxByEthAddrBJJ returns the smallest Idx in the StateDB for the given
 // Ethereum Address AND the given BabyJubJub PublicKey. If `addr` is the zero
 // address, it's ignored in the query.  If `pk` is nil, it's ignored in the
-// query.  Will return 0 in case that Idx is not found in the StateDB.
-func (s *StateDB) GetIdxByEthAddrBJJ(addr ethCommon.Address, pk *babyjub.PublicKey) common.Idx {
-	// TODO
-	return common.Idx(0)
+// query.  Will return common.Idx(0) and error in case that Idx is not found in
+// the StateDB.
+func (s *StateDB) GetIdxByEthAddrBJJ(addr ethCommon.Address, pk *babyjub.PublicKey) (common.Idx, error) {
+	if pk == nil {
+		return s.GetIdxByEthAddr(addr)
+	}
+
+	k := concatEthAddrBJJ(addr, pk)
+	b, err := s.db.Get(k)
+	if err != nil {
+		return common.Idx(0), err
+	}
+	idx, err := common.IdxFromBytes(b)
+	if err != nil {
+		return common.Idx(0), err
+	}
+	return idx, nil
 }
 
 func siblingsToZKInputFormat(s []*merkletree.Hash) []*big.Int {

db/statedb/utils_test.go

@@ -1,13 +1,84 @@
 package statedb
 
 import (
+	"io/ioutil"
 	"math/big"
 	"testing"
 
+	ethCommon "github.com/ethereum/go-ethereum/common"
+	"github.com/hermeznetwork/hermez-node/common"
 	"github.com/iden3/go-iden3-crypto/babyjub"
+	"github.com/iden3/go-merkletree/db"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
+func TestGetIdx(t *testing.T) {
+	dir, err := ioutil.TempDir("", "tmpdb")
+	require.Nil(t, err)
+
+	sdb, err := NewStateDB(dir, false, 0)
+	assert.Nil(t, err)
+
+	var sk babyjub.PrivateKey
+	copy(sk[:], []byte("1234")) // only for testing
+	pk := sk.Public()
+	var sk2 babyjub.PrivateKey
+	copy(sk2[:], []byte("12345")) // only for testing
+	pk2 := sk2.Public()
+	addr := ethCommon.HexToAddress("0x74E803744B7EEFc272E852f89a05D41515d431f2")
+	addr2 := ethCommon.HexToAddress("0x54A0706531cEa2ee8F09bAd22f604e377bb56948")
+	idx := common.Idx(1234)
+	idx2 := common.Idx(12345)
+	idx3 := common.Idx(1233)
+
+	// store the keys for idx by Addr & BJJ
+	err = sdb.setIdxByEthAddrBJJ(idx, addr, pk)
+	require.Nil(t, err)
+
+	idxR, err := sdb.GetIdxByEthAddrBJJ(addr, pk)
+	assert.Nil(t, err)
+	assert.Equal(t, idx, idxR)
+
+	// expect error when getting only by EthAddr, as value does not exist
+	// in the db for only EthAddr
+	_, err = sdb.GetIdxByEthAddr(addr)
+	assert.Nil(t, err)
+	_, err = sdb.GetIdxByEthAddr(addr2)
+	assert.NotNil(t, err)
+
+	// expect to fail
+	idxR, err = sdb.GetIdxByEthAddrBJJ(addr2, pk)
+	assert.NotNil(t, err)
+	assert.Equal(t, common.Idx(0), idxR)
+	idxR, err = sdb.GetIdxByEthAddrBJJ(addr, pk2)
+	assert.NotNil(t, err)
+	assert.Equal(t, common.Idx(0), idxR)
+
+	// try to store bigger idx, will not affect as already exist a smaller
+	// Idx for that Addr & BJJ
+	err = sdb.setIdxByEthAddrBJJ(idx2, addr, pk)
+	assert.Nil(t, err)
+
+	// store smaller idx
+	err = sdb.setIdxByEthAddrBJJ(idx3, addr, pk)
+	assert.Nil(t, err)
+
+	idxR, err = sdb.GetIdxByEthAddrBJJ(addr, pk)
+	assert.Nil(t, err)
+	assert.Equal(t, idx3, idxR)
+
+	// by EthAddr should work
+	idxR, err = sdb.GetIdxByEthAddr(addr)
+	assert.Nil(t, err)
+	assert.Equal(t, idx3, idxR)
+	// expect error when trying to get Idx by addr2 & pk2
+	idxR, err = sdb.GetIdxByEthAddrBJJ(addr2, pk2)
+	assert.NotNil(t, err)
+	assert.Equal(t, db.ErrNotFound, err)
+	assert.Equal(t, common.Idx(0), idxR)
+}
+
 func TestBJJCompressedTo256BigInt(t *testing.T) {
 	var pkComp babyjub.PublicKeyComp
 	r := BJJCompressedTo256BigInts(pkComp)

txselector/txselector.go

@@ -124,16 +124,16 @@ func (txsel *TxSelector) GetL1L2TxSelection(batchNum common.BatchNum, l1Txs []*c
 			// Account found in the DB, include the l2Tx in the selection
 			validTxs = append(validTxs, l2TxsRaw[i])
 		} else if l2TxsRaw[i].ToIdx == common.Idx(0) {
-			idx := txsel.localAccountsDB.GetIdxByEthAddrBJJ(l2TxsRaw[i].ToEthAddr, l2TxsRaw[i].ToBJJ)
-			if idx != common.Idx(0) {
-				// account for ToEthAddr/ToBJJ already exist,
+			_, err := txsel.localAccountsDB.GetIdxByEthAddrBJJ(l2TxsRaw[i].ToEthAddr, l2TxsRaw[i].ToBJJ)
+			if err == nil {
+				// account for ToEthAddr&ToBJJ already exist,
 				// there is no need to create a new one.
 				// tx valid, StateDB will use the ToIdx==0 to define the AuxToIdx
 				validTxs = append(validTxs, l2TxsRaw[i])
 				continue
 			}
 			// check if ToEthAddr is in AccountCreationAuths
-			_, err := txsel.l2db.GetAccountCreationAuth(l2TxsRaw[i].ToEthAddr) // TODO once l2db.GetAccountCreationAuth is ready, use the value returned as 'accAuth'
+			_, err = txsel.l2db.GetAccountCreationAuth(l2TxsRaw[i].ToEthAddr) // TODO once l2db.GetAccountCreationAuth is ready, use the value returned as 'accAuth'
 			if err != nil {
 				// not found, l2Tx will not be added in the selection
 				continue