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