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