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