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