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