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