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	return getAccountInTreeDB(s.db, idx)
252}
253
254// getAccountInTreeDB is abstracted from StateDB to be used from StateDB and
255// from ExitTree.  GetAccount returns the account for the given Idx
256func getAccountInTreeDB(sto db.Storage, idx common.Idx) (*common.Account, error) {
257	idxBytes, err := idx.Bytes()
258	if err != nil {
259		return nil, err
260	}
261	vBytes, err := sto.Get(idxBytes[:])
262	if err != nil {
263		return nil, err
264	}
265	accBytes, err := sto.Get(vBytes)
266	if err != nil {
267		return nil, err
268	}
269	var b [32 * common.NLeafElems]byte
270	copy(b[:], accBytes)
271	return common.AccountFromBytes(b)
272}
273
274// CreateAccount creates a new Account in the StateDB for the given Idx.  If
275// StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
276// MerkleTree, returning a CircomProcessorProof.
277func (s *StateDB) CreateAccount(idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
278	cpp, err := createAccountInTreeDB(s.db, s.mt, idx, account)
279	if err != nil {
280		return cpp, err
281	}
282	// store idx by EthAddr & BJJ
283	err = s.setIdxByEthAddrBJJ(idx, account.EthAddr, account.PublicKey)
284	return cpp, err
285}
286
287// createAccountInTreeDB is abstracted from StateDB to be used from StateDB and
288// from ExitTree.  Creates a new Account in the StateDB for the given Idx.  If
289// StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
290// MerkleTree, returning a CircomProcessorProof.
291func createAccountInTreeDB(sto db.Storage, mt *merkletree.MerkleTree, idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
292	// store at the DB the key: v, and value: leaf.Bytes()
293	v, err := account.HashValue()
294	if err != nil {
295		return nil, err
296	}
297	accountBytes, err := account.Bytes()
298	if err != nil {
299		return nil, err
300	}
301
302	// store the Leaf value
303	tx, err := sto.NewTx()
304	if err != nil {
305		return nil, err
306	}
307
308	idxBytes, err := idx.Bytes()
309	if err != nil {
310		return nil, err
311	}
312	_, err = tx.Get(idxBytes[:])
313	if err != db.ErrNotFound {
314		return nil, ErrAccountAlreadyExists
315	}
316
317	err = tx.Put(v.Bytes(), accountBytes[:])
318	if err != nil {
319		return nil, err
320	}
321	err = tx.Put(idxBytes[:], v.Bytes())
322	if err != nil {
323		return nil, err
324	}
325
326	if err := tx.Commit(); err != nil {
327		return nil, err
328	}
329
330	if mt != nil {
331		return mt.AddAndGetCircomProof(idx.BigInt(), v)
332	}
333
334	return nil, nil
335}
336
337// UpdateAccount updates the Account in the StateDB for the given Idx.  If
338// StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
339// MerkleTree, returning a CircomProcessorProof.
340func (s *StateDB) UpdateAccount(idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
341	return updateAccountInTreeDB(s.db, s.mt, idx, account)
342}
343
344// updateAccountInTreeDB is abstracted from StateDB to be used from StateDB and
345// from ExitTree.  Updates the Account in the StateDB for the given Idx.  If
346// StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
347// MerkleTree, returning a CircomProcessorProof.
348func updateAccountInTreeDB(sto db.Storage, mt *merkletree.MerkleTree, idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
349	// store at the DB the key: v, and value: account.Bytes()
350	v, err := account.HashValue()
351	if err != nil {
352		return nil, err
353	}
354	accountBytes, err := account.Bytes()
355	if err != nil {
356		return nil, err
357	}
358
359	tx, err := sto.NewTx()
360	if err != nil {
361		return nil, err
362	}
363	err = tx.Put(v.Bytes(), accountBytes[:])
364	if err != nil {
365		return nil, err
366	}
367	idxBytes, err := idx.Bytes()
368	if err != nil {
369		return nil, err
370	}
371	err = tx.Put(idxBytes[:], v.Bytes())
372	if err != nil {
373		return nil, err
374	}
375
376	if err := tx.Commit(); err != nil {
377		return nil, err
378	}
379
380	if mt != nil {
381		return mt.Update(idx.BigInt(), v)
382	}
383	return nil, nil
384}
385
386// MTGetProof returns the CircomVerifierProof for a given Idx
387func (s *StateDB) MTGetProof(idx common.Idx) (*merkletree.CircomVerifierProof, error) {
388	if s.mt == nil {
389		return nil, ErrStateDBWithoutMT
390	}
391	return s.mt.GenerateCircomVerifierProof(idx.BigInt(), s.mt.Root())
392}
393
394// LocalStateDB represents the local StateDB which allows to make copies from
395// the synchronizer StateDB, and is used by the tx-selector and the
396// batch-builder. LocalStateDB is an in-memory storage.
397type LocalStateDB struct {
398	*StateDB
399	synchronizerStateDB *StateDB
400}
401
402// NewLocalStateDB returns a new LocalStateDB connected to the given
403// synchronizerDB
404func NewLocalStateDB(path string, synchronizerDB *StateDB, typ TypeStateDB, nLevels int) (*LocalStateDB, error) {
405	s, err := NewStateDB(path, typ, nLevels)
406	if err != nil {
407		return nil, err
408	}
409	return &LocalStateDB{
410		s,
411		synchronizerDB,
412	}, nil
413}
414
415// Reset performs a reset in the LocaStateDB. If fromSynchronizer is true, it
416// gets the state from LocalStateDB.synchronizerStateDB for the given batchNum. If fromSynchronizer is false, get the state from LocalStateDB checkpoints.
417func (l *LocalStateDB) Reset(batchNum common.BatchNum, fromSynchronizer bool) error {
418	if batchNum == 0 {
419		l.idx = 0
420		return nil
421	}
422
423	synchronizerCheckpointPath := l.synchronizerStateDB.path + PathBatchNum + strconv.Itoa(int(batchNum))
424	checkpointPath := l.path + PathBatchNum + strconv.Itoa(int(batchNum))
425	currentPath := l.path + PathCurrent
426
427	if fromSynchronizer {
428		// use checkpoint from SynchronizerStateDB
429		if _, err := os.Stat(synchronizerCheckpointPath); os.IsNotExist(err) {
430			// if synchronizerStateDB does not have checkpoint at batchNum, return err
431			return fmt.Errorf("Checkpoint not exist in Synchronizer")
432		}
433
434		// remove 'current'
435		err := os.RemoveAll(currentPath)
436		if err != nil {
437			return err
438		}
439		// copy synchronizer'BatchNumX' to 'current'
440		cmd := exec.Command("cp", "-r", synchronizerCheckpointPath, currentPath) //nolint:gosec
441		err = cmd.Run()
442		if err != nil {
443			return err
444		}
445		// copy synchronizer-'BatchNumX' to 'BatchNumX'
446		cmd = exec.Command("cp", "-r", synchronizerCheckpointPath, checkpointPath) //nolint:gosec
447		err = cmd.Run()
448		if err != nil {
449			return err
450		}
451
452		// open the new 'current'
453		sto, err := pebble.NewPebbleStorage(currentPath, false)
454		if err != nil {
455			return err
456		}
457		l.db = sto
458
459		// get currentBatch num
460		l.currentBatch, err = l.GetCurrentBatch()
461		if err != nil {
462			return err
463		}
464		// open the MT for the current s.db
465		mt, err := merkletree.NewMerkleTree(l.db, l.mt.MaxLevels())
466		if err != nil {
467			return err
468		}
469		l.mt = mt
470
471		return nil
472	}
473	// use checkpoint from LocalStateDB
474	return l.StateDB.Reset(batchNum)
475}