1package statedb
  2
  3import (
  4	"errors"
  5	"fmt"
  6	"math/big"
  7	"os"
  8	"os/exec"
  9	"strconv"
 10
 11	"github.com/hermeznetwork/hermez-node/common"
 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)
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	}
178
179	// execute Checkpoint
180	err = s.db.Pebble().Checkpoint(checkpointPath)
181	if err != nil {
182		return err
183	}
184
185	return nil
186}
187
188// DeleteCheckpoint removes if exist the checkpoint of the given batchNum
189func (s *StateDB) DeleteCheckpoint(batchNum common.BatchNum) error {
190	checkpointPath := s.path + PathBatchNum + strconv.Itoa(int(batchNum))
191
192	if _, err := os.Stat(checkpointPath); os.IsNotExist(err) {
193		return fmt.Errorf("Checkpoint with batchNum %d does not exist in DB", batchNum)
194	}
195
196	return os.RemoveAll(checkpointPath)
197}
198
199// Reset resets the StateDB to the checkpoint at the given batchNum. Reset
200// does not delete the checkpoints between old current and the new current,
201// those checkpoints will remain in the storage, and eventually will be
202// deleted when MakeCheckpoint overwrites them.
203func (s *StateDB) Reset(batchNum common.BatchNum) error {
204	checkpointPath := s.path + PathBatchNum + strconv.Itoa(int(batchNum))
205	currentPath := s.path + PathCurrent
206
207	// remove 'current'
208	err := os.RemoveAll(currentPath)
209	if err != nil {
210		return err
211	}
212	if batchNum == 0 {
213		// if batchNum == 0, open the new fresh 'current'
214		sto, err := pebble.NewPebbleStorage(currentPath, false)
215		if err != nil {
216			return err
217		}
218		s.db = sto
219		s.idx = 255
220		s.currentBatch = batchNum
221		return nil
222	}
223
224	// copy 'BatchNumX' to 'current'
225	cmd := exec.Command("cp", "-r", checkpointPath, currentPath) //nolint:gosec
226	err = cmd.Run()
227	if err != nil {
228		return err
229	}
230
231	// open the new 'current'
232	sto, err := pebble.NewPebbleStorage(currentPath, false)
233	if err != nil {
234		return err
235	}
236	s.db = sto
237
238	// get currentBatch num
239	s.currentBatch, err = s.GetCurrentBatch()
240	if err != nil {
241		return err
242	}
243	// idx is obtained from the statedb reset
244	s.idx, err = s.getIdx()
245	if err != nil {
246		return err
247	}
248
249	if s.mt != nil {
250		// open the MT for the current s.db
251		mt, err := merkletree.NewMerkleTree(s.db.WithPrefix(PrefixKeyMT), s.mt.MaxLevels())
252		if err != nil {
253			return err
254		}
255		s.mt = mt
256	}
257
258	return nil
259}
260
261// GetAccount returns the account for the given Idx
262func (s *StateDB) GetAccount(idx common.Idx) (*common.Account, error) {
263	return getAccountInTreeDB(s.db, idx)
264}
265
266// GetAccounts returns all the accounts in the db.  Use for debugging pruposes
267// only.
268func (s *StateDB) GetAccounts() ([]common.Account, error) {
269	idxDB := s.db.WithPrefix(PrefixKeyIdx)
270	idxs := []common.Idx{}
271	// NOTE: Current implementation of Iterate in the pebble interface is
272	// not efficient, as it iterates over all keys.  Improve it following
273	// this example: https://github.com/cockroachdb/pebble/pull/923/files
274	if err := idxDB.Iterate(func(k []byte, v []byte) (bool, error) {
275		idx, err := common.IdxFromBytes(k)
276		if err != nil {
277			return false, err
278		}
279		idxs = append(idxs, idx)
280		return true, nil
281	}); err != nil {
282		return nil, err
283	}
284	accs := []common.Account{}
285	for i := range idxs {
286		acc, err := s.GetAccount(idxs[i])
287		if err != nil {
288			return nil, err
289		}
290		accs = append(accs, *acc)
291	}
292	return accs, nil
293}
294
295// getAccountInTreeDB is abstracted from StateDB to be used from StateDB and
296// from ExitTree.  GetAccount returns the account for the given Idx
297func getAccountInTreeDB(sto db.Storage, idx common.Idx) (*common.Account, error) {
298	idxBytes, err := idx.Bytes()
299	if err != nil {
300		return nil, err
301	}
302	vBytes, err := sto.Get(append(PrefixKeyIdx, idxBytes[:]...))
303	if err != nil {
304		return nil, err
305	}
306	accBytes, err := sto.Get(append(PrefixKeyAccHash, vBytes...))
307	if err != nil {
308		return nil, err
309	}
310	var b [32 * common.NLeafElems]byte
311	copy(b[:], accBytes)
312	account, err := common.AccountFromBytes(b)
313	if err != nil {
314		return nil, err
315	}
316	account.Idx = idx
317	return account, nil
318}
319
320// CreateAccount creates a new Account in the StateDB for the given Idx.  If
321// StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
322// MerkleTree, returning a CircomProcessorProof.
323func (s *StateDB) CreateAccount(idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
324	cpp, err := createAccountInTreeDB(s.db, s.mt, idx, account)
325	if err != nil {
326		return cpp, err
327	}
328	// store idx by EthAddr & BJJ
329	err = s.setIdxByEthAddrBJJ(idx, account.EthAddr, account.PublicKey)
330	return cpp, err
331}
332
333// createAccountInTreeDB is abstracted from StateDB to be used from StateDB and
334// from ExitTree.  Creates a new Account in the StateDB for the given Idx.  If
335// StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
336// MerkleTree, returning a CircomProcessorProof.
337func createAccountInTreeDB(sto db.Storage, mt *merkletree.MerkleTree, idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
338	// store at the DB the key: v, and value: leaf.Bytes()
339	v, err := account.HashValue()
340	if err != nil {
341		return nil, err
342	}
343	accountBytes, err := account.Bytes()
344	if err != nil {
345		return nil, err
346	}
347
348	// store the Leaf value
349	tx, err := sto.NewTx()
350	if err != nil {
351		return nil, err
352	}
353
354	idxBytes, err := idx.Bytes()
355	if err != nil {
356		return nil, err
357	}
358	_, err = tx.Get(append(PrefixKeyIdx, idxBytes[:]...))
359	if err != db.ErrNotFound {
360		return nil, ErrAccountAlreadyExists
361	}
362
363	err = tx.Put(append(PrefixKeyAccHash, v.Bytes()...), accountBytes[:])
364	if err != nil {
365		return nil, err
366	}
367	err = tx.Put(append(PrefixKeyIdx, idxBytes[:]...), v.Bytes())
368	if err != nil {
369		return nil, err
370	}
371
372	if err := tx.Commit(); err != nil {
373		return nil, err
374	}
375
376	if mt != nil {
377		return mt.AddAndGetCircomProof(idx.BigInt(), v)
378	}
379
380	return nil, nil
381}
382
383// UpdateAccount updates the Account in the StateDB for the given Idx.  If
384// StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
385// MerkleTree, returning a CircomProcessorProof.
386func (s *StateDB) UpdateAccount(idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
387	return updateAccountInTreeDB(s.db, s.mt, idx, account)
388}
389
390// updateAccountInTreeDB is abstracted from StateDB to be used from StateDB and
391// from ExitTree.  Updates the Account in the StateDB for the given Idx.  If
392// StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
393// MerkleTree, returning a CircomProcessorProof.
394func updateAccountInTreeDB(sto db.Storage, mt *merkletree.MerkleTree, idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
395	// store at the DB the key: v, and value: account.Bytes()
396	v, err := account.HashValue()
397	if err != nil {
398		return nil, err
399	}
400	accountBytes, err := account.Bytes()
401	if err != nil {
402		return nil, err
403	}
404
405	tx, err := sto.NewTx()
406	if err != nil {
407		return nil, err
408	}
409	err = tx.Put(append(PrefixKeyAccHash, v.Bytes()...), accountBytes[:])
410	if err != nil {
411		return nil, err
412	}
413	idxBytes, err := idx.Bytes()
414	if err != nil {
415		return nil, err
416	}
417	err = tx.Put(append(PrefixKeyIdx, idxBytes[:]...), v.Bytes())
418	if err != nil {
419		return nil, err
420	}
421
422	if err := tx.Commit(); err != nil {
423		return nil, err
424	}
425
426	if mt != nil {
427		return mt.Update(idx.BigInt(), v)
428	}
429	return nil, nil
430}
431
432// MTGetProof returns the CircomVerifierProof for a given Idx
433func (s *StateDB) MTGetProof(idx common.Idx) (*merkletree.CircomVerifierProof, error) {
434	if s.mt == nil {
435		return nil, ErrStateDBWithoutMT
436	}
437	return s.mt.GenerateCircomVerifierProof(idx.BigInt(), s.mt.Root())
438}
439
440// MTGetRoot returns the current root of the underlying Merkle Tree
441func (s *StateDB) MTGetRoot() *big.Int {
442	return s.mt.Root().BigInt()
443}
444
445// LocalStateDB represents the local StateDB which allows to make copies from
446// the synchronizer StateDB, and is used by the tx-selector and the
447// batch-builder. LocalStateDB is an in-memory storage.
448type LocalStateDB struct {
449	*StateDB
450	synchronizerStateDB *StateDB
451}
452
453// NewLocalStateDB returns a new LocalStateDB connected to the given
454// synchronizerDB
455func NewLocalStateDB(path string, synchronizerDB *StateDB, typ TypeStateDB, nLevels int) (*LocalStateDB, error) {
456	s, err := NewStateDB(path, typ, nLevels)
457	if err != nil {
458		return nil, err
459	}
460	return &LocalStateDB{
461		s,
462		synchronizerDB,
463	}, nil
464}
465
466// Reset performs a reset in the LocaStateDB. If fromSynchronizer is true, it
467// gets the state from LocalStateDB.synchronizerStateDB for the given batchNum. If fromSynchronizer is false, get the state from LocalStateDB checkpoints.
468func (l *LocalStateDB) Reset(batchNum common.BatchNum, fromSynchronizer bool) error {
469	if batchNum == 0 {
470		l.idx = 0
471		return nil
472	}
473
474	synchronizerCheckpointPath := l.synchronizerStateDB.path + PathBatchNum + strconv.Itoa(int(batchNum))
475	checkpointPath := l.path + PathBatchNum + strconv.Itoa(int(batchNum))
476	currentPath := l.path + PathCurrent
477
478	if fromSynchronizer {
479		// use checkpoint from SynchronizerStateDB
480		if _, err := os.Stat(synchronizerCheckpointPath); os.IsNotExist(err) {
481			// if synchronizerStateDB does not have checkpoint at batchNum, return err
482			return fmt.Errorf("Checkpoint not exist in Synchronizer")
483		}
484
485		// remove 'current'
486		err := os.RemoveAll(currentPath)
487		if err != nil {
488			return err
489		}
490		// copy synchronizer'BatchNumX' to 'current'
491		cmd := exec.Command("cp", "-r", synchronizerCheckpointPath, currentPath) //nolint:gosec
492		err = cmd.Run()
493		if err != nil {
494			return err
495		}
496		// copy synchronizer-'BatchNumX' to 'BatchNumX'
497		cmd = exec.Command("cp", "-r", synchronizerCheckpointPath, checkpointPath) //nolint:gosec
498		err = cmd.Run()
499		if err != nil {
500			return err
501		}
502
503		// open the new 'current'
504		sto, err := pebble.NewPebbleStorage(currentPath, false)
505		if err != nil {
506			return err
507		}
508		l.db = sto
509
510		// get currentBatch num
511		l.currentBatch, err = l.GetCurrentBatch()
512		if err != nil {
513			return err
514		}
515		// open the MT for the current s.db
516		mt, err := merkletree.NewMerkleTree(l.db.WithPrefix(PrefixKeyMT), l.mt.MaxLevels())
517		if err != nil {
518			return err
519		}
520		l.mt = mt
521
522		return nil
523	}
524	// use checkpoint from LocalStateDB
525	return l.StateDB.Reset(batchNum)
526}