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