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