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