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