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