1package statedb
  2
  3import (
  4	"errors"
  5	"fmt"
  6	"math/big"
  7	"os"
  8	"strconv"
  9
 10	"github.com/hermeznetwork/hermez-node/common"
 11	"github.com/hermeznetwork/hermez-node/log"
 12	"github.com/hermeznetwork/tracerr"
 13	"github.com/iden3/go-merkletree"
 14	"github.com/iden3/go-merkletree/db"
 15	"github.com/iden3/go-merkletree/db/pebble"
 16)
 17
 18// TODO(Edu): Document here how StateDB is kept consistent
 19
 20var (
 21	// ErrStateDBWithoutMT is used when a method that requires a MerkleTree
 22	// is called in a StateDB that does not have a MerkleTree defined
 23	ErrStateDBWithoutMT = errors.New("Can not call method to use MerkleTree in a StateDB without MerkleTree")
 24
 25	// ErrAccountAlreadyExists is used when CreateAccount is called and the
 26	// Account already exists
 27	ErrAccountAlreadyExists = errors.New("Can not CreateAccount because Account already exists")
 28
 29	// ErrToIdxNotFound is used when trying to get the ToIdx from ToEthAddr
 30	// or ToEthAddr&ToBJJ
 31	ErrToIdxNotFound = errors.New("ToIdx can not be found")
 32	// ErrGetIdxNoCase is used when trying to get the Idx from EthAddr &
 33	// BJJ with not compatible combination
 34	ErrGetIdxNoCase = errors.New("Can not get Idx due unexpected combination of ethereum Address & BabyJubJub PublicKey")
 35
 36	// KeyCurrentBatch is used as key in the db to store the current BatchNum
 37	KeyCurrentBatch = []byte("k:currentbatch")
 38
 39	// PrefixKeyIdx is the key prefix for idx in the db
 40	PrefixKeyIdx = []byte("i:")
 41	// PrefixKeyAccHash is the key prefix for account hash in the db
 42	PrefixKeyAccHash = []byte("h:")
 43	// PrefixKeyMT is the key prefix for merkle tree in the db
 44	PrefixKeyMT = []byte("m:")
 45	// PrefixKeyAddr is the key prefix for address in the db
 46	PrefixKeyAddr = []byte("a:")
 47	// PrefixKeyAddrBJJ is the key prefix for address-babyjubjub in the db
 48	PrefixKeyAddrBJJ = []byte("ab:")
 49)
 50
 51const (
 52	// PathBatchNum defines the subpath of the Batch Checkpoint in the
 53	// subpath of the StateDB
 54	PathBatchNum = "/BatchNum"
 55	// PathCurrent defines the subpath of the current Batch in the subpath
 56	// of the StateDB
 57	PathCurrent = "/current"
 58	// TypeSynchronizer defines a StateDB used by the Synchronizer, that
 59	// generates the ExitTree when processing the txs
 60	TypeSynchronizer = "synchronizer"
 61	// TypeTxSelector defines a StateDB used by the TxSelector, without
 62	// computing ExitTree neither the ZKInputs
 63	TypeTxSelector = "txselector"
 64	// TypeBatchBuilder defines a StateDB used by the BatchBuilder, that
 65	// generates the ExitTree and the ZKInput when processing the txs
 66	TypeBatchBuilder = "batchbuilder"
 67)
 68
 69// TypeStateDB determines the type of StateDB
 70type TypeStateDB string
 71
 72// StateDB represents the StateDB object
 73type StateDB struct {
 74	path         string
 75	currentBatch common.BatchNum
 76	db           *pebble.PebbleStorage
 77	mt           *merkletree.MerkleTree
 78	typ          TypeStateDB
 79	chainID      uint16
 80	// idx holds the current Idx that the BatchBuilder is using
 81	idx common.Idx
 82	zki *common.ZKInputs
 83	// i is the current transaction index in the ZKInputs generation (zki)
 84	i int
 85	// AccumulatedFees contains the accumulated fees for each token (Coord
 86	// Idx) in the processed batch
 87	AccumulatedFees map[common.Idx]*big.Int
 88}
 89
 90// NewStateDB creates a new StateDB, allowing to use an in-memory or in-disk
 91// storage
 92func NewStateDB(path string, typ TypeStateDB, nLevels int, chainID uint16) (*StateDB, error) {
 93	var sto *pebble.PebbleStorage
 94	var err error
 95	sto, err = pebble.NewPebbleStorage(path+PathCurrent, false)
 96	if err != nil {
 97		return nil, tracerr.Wrap(err)
 98	}
 99
100	var mt *merkletree.MerkleTree = nil
101	if typ == TypeSynchronizer || typ == TypeBatchBuilder {
102		mt, err = merkletree.NewMerkleTree(sto.WithPrefix(PrefixKeyMT), nLevels)
103		if err != nil {
104			return nil, tracerr.Wrap(err)
105		}
106	}
107	if typ == TypeTxSelector && nLevels != 0 {
108		return nil, tracerr.Wrap(fmt.Errorf("invalid StateDB parameters: StateDB type==TypeStateDB can not have nLevels!=0"))
109	}
110
111	sdb := &StateDB{
112		path:    path,
113		db:      sto,
114		mt:      mt,
115		typ:     typ,
116		chainID: chainID,
117	}
118
119	// load currentBatch
120	sdb.currentBatch, err = sdb.GetCurrentBatch()
121	if err != nil {
122		return nil, tracerr.Wrap(err)
123	}
124
125	// make reset (get checkpoint) at currentBatch
126	err = sdb.reset(sdb.currentBatch, false)
127	if err != nil {
128		return nil, tracerr.Wrap(err)
129	}
130
131	return sdb, nil
132}
133
134// DB returns the *pebble.PebbleStorage from the StateDB
135func (s *StateDB) DB() *pebble.PebbleStorage {
136	return s.db
137}
138
139// GetCurrentBatch returns the current BatchNum stored in the StateDB
140func (s *StateDB) GetCurrentBatch() (common.BatchNum, error) {
141	cbBytes, err := s.db.Get(KeyCurrentBatch)
142	if tracerr.Unwrap(err) == db.ErrNotFound {
143		return 0, nil
144	}
145	if err != nil {
146		return 0, tracerr.Wrap(err)
147	}
148	return common.BatchNumFromBytes(cbBytes)
149}
150
151// setCurrentBatch stores the current BatchNum in the StateDB
152func (s *StateDB) setCurrentBatch() error {
153	tx, err := s.db.NewTx()
154	if err != nil {
155		return tracerr.Wrap(err)
156	}
157	err = tx.Put(KeyCurrentBatch, s.currentBatch.Bytes())
158	if err != nil {
159		return tracerr.Wrap(err)
160	}
161	if err := tx.Commit(); err != nil {
162		return tracerr.Wrap(err)
163	}
164	return nil
165}
166
167// MakeCheckpoint does a checkpoint at the given batchNum in the defined path. Internally this advances & stores the current BatchNum, and then stores a Checkpoint of the current state of the StateDB.
168func (s *StateDB) MakeCheckpoint() error {
169	// advance currentBatch
170	s.currentBatch++
171	log.Debugw("Making StateDB checkpoint", "batch", s.currentBatch, "type", s.typ)
172
173	checkpointPath := s.path + PathBatchNum + strconv.Itoa(int(s.currentBatch))
174
175	err := s.setCurrentBatch()
176	if err != nil {
177		return tracerr.Wrap(err)
178	}
179
180	// if checkpoint BatchNum already exist in disk, delete it
181	if _, err := os.Stat(checkpointPath); !os.IsNotExist(err) {
182		err := os.RemoveAll(checkpointPath)
183		if err != nil {
184			return tracerr.Wrap(err)
185		}
186	} else if err != nil && !os.IsNotExist(err) {
187		return tracerr.Wrap(err)
188	}
189
190	// execute Checkpoint
191	err = s.db.Pebble().Checkpoint(checkpointPath)
192	if err != nil {
193		return tracerr.Wrap(err)
194	}
195
196	return nil
197}
198
199// DeleteCheckpoint removes if exist the checkpoint of the given batchNum
200func (s *StateDB) DeleteCheckpoint(batchNum common.BatchNum) error {
201	checkpointPath := s.path + PathBatchNum + strconv.Itoa(int(batchNum))
202
203	if _, err := os.Stat(checkpointPath); os.IsNotExist(err) {
204		return tracerr.Wrap(fmt.Errorf("Checkpoint with batchNum %d does not exist in DB", batchNum))
205	}
206
207	return os.RemoveAll(checkpointPath)
208}
209
210func pebbleMakeCheckpoint(source, dest string) error {
211	// Remove dest folder (if it exists) before doing the checkpoint
212	if _, err := os.Stat(dest); !os.IsNotExist(err) {
213		err := os.RemoveAll(dest)
214		if err != nil {
215			return tracerr.Wrap(err)
216		}
217	} else if err != nil && !os.IsNotExist(err) {
218		return tracerr.Wrap(err)
219	}
220
221	sto, err := pebble.NewPebbleStorage(source, false)
222	if err != nil {
223		return tracerr.Wrap(err)
224	}
225	defer func() {
226		errClose := sto.Pebble().Close()
227		if errClose != nil {
228			log.Errorw("Pebble.Close", "err", errClose)
229		}
230	}()
231
232	// execute Checkpoint
233	err = sto.Pebble().Checkpoint(dest)
234	if err != nil {
235		return tracerr.Wrap(err)
236	}
237
238	return nil
239}
240
241// Reset resets the StateDB to the checkpoint at the given batchNum. Reset
242// does not delete the checkpoints between old current and the new current,
243// those checkpoints will remain in the storage, and eventually will be
244// deleted when MakeCheckpoint overwrites them.
245func (s *StateDB) Reset(batchNum common.BatchNum) error {
246	return s.reset(batchNum, true)
247}
248
249// reset resets the StateDB to the checkpoint at the given batchNum. Reset
250// does not delete the checkpoints between old current and the new current,
251// those checkpoints will remain in the storage, and eventually will be
252// deleted when MakeCheckpoint overwrites them.  `closeCurrent` will close the
253// currently opened db before doing the reset.
254func (s *StateDB) reset(batchNum common.BatchNum, closeCurrent bool) error {
255	currentPath := s.path + PathCurrent
256
257	if closeCurrent {
258		if err := s.db.Pebble().Close(); err != nil {
259			return tracerr.Wrap(err)
260		}
261	}
262	// remove 'current'
263	err := os.RemoveAll(currentPath)
264	if err != nil {
265		return tracerr.Wrap(err)
266	}
267	if batchNum == 0 {
268		// if batchNum == 0, open the new fresh 'current'
269		sto, err := pebble.NewPebbleStorage(currentPath, false)
270		if err != nil {
271			return tracerr.Wrap(err)
272		}
273		s.db = sto
274		s.idx = 255
275		s.currentBatch = batchNum
276
277		if s.mt != nil {
278			// open the MT for the current s.db
279			mt, err := merkletree.NewMerkleTree(s.db.WithPrefix(PrefixKeyMT), s.mt.MaxLevels())
280			if err != nil {
281				return tracerr.Wrap(err)
282			}
283			s.mt = mt
284		}
285		return nil
286	}
287
288	checkpointPath := s.path + PathBatchNum + strconv.Itoa(int(batchNum))
289	// copy 'BatchNumX' to 'current'
290	err = pebbleMakeCheckpoint(checkpointPath, currentPath)
291	if err != nil {
292		return tracerr.Wrap(err)
293	}
294
295	// open the new 'current'
296	sto, err := pebble.NewPebbleStorage(currentPath, false)
297	if err != nil {
298		return tracerr.Wrap(err)
299	}
300	s.db = sto
301
302	// get currentBatch num
303	s.currentBatch, err = s.GetCurrentBatch()
304	if err != nil {
305		return tracerr.Wrap(err)
306	}
307	// idx is obtained from the statedb reset
308	s.idx, err = s.GetIdx()
309	if err != nil {
310		return tracerr.Wrap(err)
311	}
312
313	if s.mt != nil {
314		// open the MT for the current s.db
315		mt, err := merkletree.NewMerkleTree(s.db.WithPrefix(PrefixKeyMT), s.mt.MaxLevels())
316		if err != nil {
317			return tracerr.Wrap(err)
318		}
319		s.mt = mt
320	}
321
322	return nil
323}
324
325// GetAccount returns the account for the given Idx
326func (s *StateDB) GetAccount(idx common.Idx) (*common.Account, error) {
327	return getAccountInTreeDB(s.db, idx)
328}
329
330// GetAccounts returns all the accounts in the db.  Use for debugging pruposes
331// only.
332func (s *StateDB) GetAccounts() ([]common.Account, error) {
333	idxDB := s.db.WithPrefix(PrefixKeyIdx)
334	idxs := []common.Idx{}
335	// NOTE: Current implementation of Iterate in the pebble interface is
336	// not efficient, as it iterates over all keys.  Improve it following
337	// this example: https://github.com/cockroachdb/pebble/pull/923/files
338	if err := idxDB.Iterate(func(k []byte, v []byte) (bool, error) {
339		idx, err := common.IdxFromBytes(k)
340		if err != nil {
341			return false, tracerr.Wrap(err)
342		}
343		idxs = append(idxs, idx)
344		return true, nil
345	}); err != nil {
346		return nil, tracerr.Wrap(err)
347	}
348	accs := []common.Account{}
349	for i := range idxs {
350		acc, err := s.GetAccount(idxs[i])
351		if err != nil {
352			return nil, tracerr.Wrap(err)
353		}
354		accs = append(accs, *acc)
355	}
356	return accs, nil
357}
358
359// getAccountInTreeDB is abstracted from StateDB to be used from StateDB and
360// from ExitTree.  GetAccount returns the account for the given Idx
361func getAccountInTreeDB(sto db.Storage, idx common.Idx) (*common.Account, error) {
362	idxBytes, err := idx.Bytes()
363	if err != nil {
364		return nil, tracerr.Wrap(err)
365	}
366	vBytes, err := sto.Get(append(PrefixKeyIdx, idxBytes[:]...))
367	if err != nil {
368		return nil, tracerr.Wrap(err)
369	}
370	accBytes, err := sto.Get(append(PrefixKeyAccHash, vBytes...))
371	if err != nil {
372		return nil, tracerr.Wrap(err)
373	}
374	var b [32 * common.NLeafElems]byte
375	copy(b[:], accBytes)
376	account, err := common.AccountFromBytes(b)
377	if err != nil {
378		return nil, tracerr.Wrap(err)
379	}
380	account.Idx = idx
381	return account, nil
382}
383
384// CreateAccount creates a new Account in the StateDB for the given Idx.  If
385// StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
386// MerkleTree, returning a CircomProcessorProof.
387func (s *StateDB) CreateAccount(idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
388	cpp, err := createAccountInTreeDB(s.db, s.mt, idx, account)
389	if err != nil {
390		return cpp, tracerr.Wrap(err)
391	}
392	// store idx by EthAddr & BJJ
393	err = s.setIdxByEthAddrBJJ(idx, account.EthAddr, account.PublicKey, account.TokenID)
394	return cpp, tracerr.Wrap(err)
395}
396
397// createAccountInTreeDB is abstracted from StateDB to be used from StateDB and
398// from ExitTree.  Creates a new Account in the StateDB for the given Idx.  If
399// StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
400// MerkleTree, returning a CircomProcessorProof.
401func createAccountInTreeDB(sto db.Storage, mt *merkletree.MerkleTree, idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
402	// store at the DB the key: v, and value: leaf.Bytes()
403	v, err := account.HashValue()
404	if err != nil {
405		return nil, tracerr.Wrap(err)
406	}
407	accountBytes, err := account.Bytes()
408	if err != nil {
409		return nil, tracerr.Wrap(err)
410	}
411
412	// store the Leaf value
413	tx, err := sto.NewTx()
414	if err != nil {
415		return nil, tracerr.Wrap(err)
416	}
417
418	idxBytes, err := idx.Bytes()
419	if err != nil {
420		return nil, tracerr.Wrap(err)
421	}
422	_, err = tx.Get(append(PrefixKeyIdx, idxBytes[:]...))
423	if tracerr.Unwrap(err) != db.ErrNotFound {
424		return nil, tracerr.Wrap(ErrAccountAlreadyExists)
425	}
426
427	err = tx.Put(append(PrefixKeyAccHash, v.Bytes()...), accountBytes[:])
428	if err != nil {
429		return nil, tracerr.Wrap(err)
430	}
431	err = tx.Put(append(PrefixKeyIdx, idxBytes[:]...), v.Bytes())
432	if err != nil {
433		return nil, tracerr.Wrap(err)
434	}
435
436	if err := tx.Commit(); err != nil {
437		return nil, tracerr.Wrap(err)
438	}
439
440	if mt != nil {
441		return mt.AddAndGetCircomProof(idx.BigInt(), v)
442	}
443
444	return nil, nil
445}
446
447// UpdateAccount updates the Account in the StateDB for the given Idx.  If
448// StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
449// MerkleTree, returning a CircomProcessorProof.
450func (s *StateDB) UpdateAccount(idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
451	return updateAccountInTreeDB(s.db, s.mt, idx, account)
452}
453
454// updateAccountInTreeDB is abstracted from StateDB to be used from StateDB and
455// from ExitTree.  Updates the Account in the StateDB for the given Idx.  If
456// StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
457// MerkleTree, returning a CircomProcessorProof.
458func updateAccountInTreeDB(sto db.Storage, mt *merkletree.MerkleTree, idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
459	// store at the DB the key: v, and value: account.Bytes()
460	v, err := account.HashValue()
461	if err != nil {
462		return nil, tracerr.Wrap(err)
463	}
464	accountBytes, err := account.Bytes()
465	if err != nil {
466		return nil, tracerr.Wrap(err)
467	}
468
469	tx, err := sto.NewTx()
470	if err != nil {
471		return nil, tracerr.Wrap(err)
472	}
473	err = tx.Put(append(PrefixKeyAccHash, v.Bytes()...), accountBytes[:])
474	if err != nil {
475		return nil, tracerr.Wrap(err)
476	}
477	idxBytes, err := idx.Bytes()
478	if err != nil {
479		return nil, tracerr.Wrap(err)
480	}
481	err = tx.Put(append(PrefixKeyIdx, idxBytes[:]...), v.Bytes())
482	if err != nil {
483		return nil, tracerr.Wrap(err)
484	}
485
486	if err := tx.Commit(); err != nil {
487		return nil, tracerr.Wrap(err)
488	}
489
490	if mt != nil {
491		proof, err := mt.Update(idx.BigInt(), v)
492		return proof, tracerr.Wrap(err)
493	}
494	return nil, nil
495}
496
497// MTGetProof returns the CircomVerifierProof for a given Idx
498func (s *StateDB) MTGetProof(idx common.Idx) (*merkletree.CircomVerifierProof, error) {
499	if s.mt == nil {
500		return nil, tracerr.Wrap(ErrStateDBWithoutMT)
501	}
502	return s.mt.GenerateCircomVerifierProof(idx.BigInt(), s.mt.Root())
503}
504
505// MTGetRoot returns the current root of the underlying Merkle Tree
506func (s *StateDB) MTGetRoot() *big.Int {
507	return s.mt.Root().BigInt()
508}
509
510// MerkleTree returns the underlying StateDB merkle tree.  It can be nil.
511func (s *StateDB) MerkleTree() *merkletree.MerkleTree {
512	return s.mt
513}
514
515// LocalStateDB represents the local StateDB which allows to make copies from
516// the synchronizer StateDB, and is used by the tx-selector and the
517// batch-builder. LocalStateDB is an in-memory storage.
518type LocalStateDB struct {
519	*StateDB
520	synchronizerStateDB *StateDB
521}
522
523// NewLocalStateDB returns a new LocalStateDB connected to the given
524// synchronizerDB
525func NewLocalStateDB(path string, synchronizerDB *StateDB, typ TypeStateDB, nLevels int) (*LocalStateDB, error) {
526	s, err := NewStateDB(path, typ, nLevels, synchronizerDB.chainID)
527	if err != nil {
528		return nil, tracerr.Wrap(err)
529	}
530	return &LocalStateDB{
531		s,
532		synchronizerDB,
533	}, nil
534}
535
536// Reset performs a reset in the LocaStateDB. If fromSynchronizer is true, it
537// gets the state from LocalStateDB.synchronizerStateDB for the given batchNum. If fromSynchronizer is false, get the state from LocalStateDB checkpoints.
538func (l *LocalStateDB) Reset(batchNum common.BatchNum, fromSynchronizer bool) error {
539	if batchNum == 0 {
540		l.idx = 0
541		return nil
542	}
543
544	synchronizerCheckpointPath := l.synchronizerStateDB.path + PathBatchNum + strconv.Itoa(int(batchNum))
545	checkpointPath := l.path + PathBatchNum + strconv.Itoa(int(batchNum))
546	currentPath := l.path + PathCurrent
547
548	if fromSynchronizer {
549		// use checkpoint from SynchronizerStateDB
550		if _, err := os.Stat(synchronizerCheckpointPath); os.IsNotExist(err) {
551			// if synchronizerStateDB does not have checkpoint at batchNum, return err
552			return tracerr.Wrap(fmt.Errorf("Checkpoint \"%v\" not exist in Synchronizer",
553				synchronizerCheckpointPath))
554		}
555
556		if err := l.db.Pebble().Close(); err != nil {
557			return tracerr.Wrap(err)
558		}
559		// remove 'current'
560		err := os.RemoveAll(currentPath)
561		if err != nil {
562			return tracerr.Wrap(err)
563		}
564		// copy synchronizer'BatchNumX' to 'current'
565		err = pebbleMakeCheckpoint(synchronizerCheckpointPath, currentPath)
566		if err != nil {
567			return tracerr.Wrap(err)
568		}
569		// copy synchronizer'BatchNumX' to 'BatchNumX'
570		err = pebbleMakeCheckpoint(synchronizerCheckpointPath, checkpointPath)
571		if err != nil {
572			return tracerr.Wrap(err)
573		}
574
575		// open the new 'current'
576		sto, err := pebble.NewPebbleStorage(currentPath, false)
577		if err != nil {
578			return tracerr.Wrap(err)
579		}
580		l.db = sto
581
582		// get currentBatch num
583		l.currentBatch, err = l.GetCurrentBatch()
584		if err != nil {
585			return tracerr.Wrap(err)
586		}
587		// open the MT for the current s.db
588		if l.mt != nil {
589			mt, err := merkletree.NewMerkleTree(l.db.WithPrefix(PrefixKeyMT), l.mt.MaxLevels())
590			if err != nil {
591				return tracerr.Wrap(err)
592			}
593			l.mt = mt
594		}
595
596		return nil
597	}
598	// use checkpoint from LocalStateDB
599	return l.StateDB.reset(batchNum, true)
600}