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