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