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