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.
205// MerkleTree is not affected.
206func (s *StateDB) CreateAccount(idx common.Idx, account *common.Account) error {
207	// store at the DB the key: v, and value: leaf.Bytes()
208	v, err := account.HashValue()
209	if err != nil {
210		return err
211	}
212	accountBytes, err := account.Bytes()
213	if err != nil {
214		return err
215	}
216
217	// store the Leaf value
218	tx, err := s.db.NewTx()
219	if err != nil {
220		return err
221	}
222
223	_, err = tx.Get(idx.Bytes())
224	if err != db.ErrNotFound {
225		return ErrAccountAlreadyExists
226	}
227
228	tx.Put(v.Bytes(), accountBytes[:])
229	tx.Put(idx.Bytes(), v.Bytes())
230
231	return tx.Commit()
232}
233
234// UpdateAccount updates the Account in the StateDB for the given Idx.
235// MerkleTree is not affected.
236func (s *StateDB) UpdateAccount(idx common.Idx, account *common.Account) error {
237	// store at the DB the key: v, and value: leaf.Bytes()
238	v, err := account.HashValue()
239	if err != nil {
240		return err
241	}
242	accountBytes, err := account.Bytes()
243	if err != nil {
244		return err
245	}
246
247	tx, err := s.db.NewTx()
248	if err != nil {
249		return err
250	}
251	tx.Put(v.Bytes(), accountBytes[:])
252	tx.Put(idx.Bytes(), v.Bytes())
253
254	return tx.Commit()
255}
256
257// MTCreateAccount creates a new Account in the StateDB for the given Idx,
258// and updates the MerkleTree, returning a CircomProcessorProof
259func (s *StateDB) MTCreateAccount(idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
260	if s.mt == nil {
261		return nil, ErrStateDBWithoutMT
262	}
263	err := s.CreateAccount(idx, account)
264	if err != nil {
265		return nil, err
266	}
267
268	v, err := account.HashValue() // already computed in s.CreateAccount, next iteration reuse first computation
269	if err != nil {
270		return nil, err
271	}
272	// Add k & v into the MT
273	return s.mt.AddAndGetCircomProof(idx.BigInt(), v)
274}
275
276// MTUpdateAccount updates the Account in the StateDB for the given Idx, and
277// updates the MerkleTree, returning a CircomProcessorProof
278func (s *StateDB) MTUpdateAccount(idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
279	if s.mt == nil {
280		return nil, ErrStateDBWithoutMT
281	}
282	err := s.UpdateAccount(idx, account)
283	if err != nil {
284		return nil, err
285	}
286
287	v, err := account.HashValue() // already computed in s.CreateAccount, next iteration reuse first computation
288	if err != nil {
289		return nil, err
290	}
291	// Add k & v into the MT
292	return s.mt.Update(idx.BigInt(), v)
293}
294
295// MTGetProof returns the CircomVerifierProof for a given Idx
296func (s *StateDB) MTGetProof(idx common.Idx) (*merkletree.CircomVerifierProof, error) {
297	if s.mt == nil {
298		return nil, ErrStateDBWithoutMT
299	}
300	return s.mt.GenerateCircomVerifierProof(idx.BigInt(), s.mt.Root())
301}
302
303// LocalStateDB represents the local StateDB which allows to make copies from
304// the synchronizer StateDB, and is used by the tx-selector and the
305// batch-builder. LocalStateDB is an in-memory storage.
306type LocalStateDB struct {
307	*StateDB
308	synchronizerStateDB *StateDB
309}
310
311// NewLocalStateDB returns a new LocalStateDB connected to the given
312// synchronizerDB
313func NewLocalStateDB(path string, synchronizerDB *StateDB, withMT bool, nLevels int) (*LocalStateDB, error) {
314	s, err := NewStateDB(path, withMT, nLevels)
315	if err != nil {
316		return nil, err
317	}
318	return &LocalStateDB{
319		s,
320		synchronizerDB,
321	}, nil
322}
323
324// Reset performs a reset in the LocaStateDB. If fromSynchronizer is true, it
325// gets the state from LocalStateDB.synchronizerStateDB for the given batchNum. If fromSynchronizer is false, get the state from LocalStateDB checkpoints.
326func (l *LocalStateDB) Reset(batchNum uint64, fromSynchronizer bool) error {
327	if batchNum == 0 {
328		l.idx = 0
329		return nil
330	}
331
332	synchronizerCheckpointPath := l.synchronizerStateDB.path + PATHBATCHNUM + strconv.Itoa(int(batchNum))
333	checkpointPath := l.path + PATHBATCHNUM + strconv.Itoa(int(batchNum))
334	currentPath := l.path + PATHCURRENT
335
336	if fromSynchronizer {
337		// use checkpoint from SynchronizerStateDB
338		if _, err := os.Stat(synchronizerCheckpointPath); os.IsNotExist(err) {
339			// if synchronizerStateDB does not have checkpoint at batchNum, return err
340			return fmt.Errorf("Checkpoint not exist in Synchronizer")
341		}
342
343		// remove 'current'
344		err := os.RemoveAll(currentPath)
345		if err != nil {
346			return err
347		}
348		// copy synchronizer'BatchNumX' to 'current'
349		cmd := exec.Command("cp", "-r", synchronizerCheckpointPath, currentPath)
350		err = cmd.Run()
351		if err != nil {
352			return err
353		}
354		// copy synchronizer-'BatchNumX' to 'BatchNumX'
355		cmd = exec.Command("cp", "-r", synchronizerCheckpointPath, checkpointPath)
356		err = cmd.Run()
357		if err != nil {
358			return err
359		}
360
361		// open the new 'current'
362		sto, err := pebble.NewPebbleStorage(currentPath, false)
363		if err != nil {
364			return err
365		}
366		l.db = sto
367
368		// get currentBatch num
369		l.currentBatch, err = l.GetCurrentBatch()
370		if err != nil {
371			return err
372		}
373		return nil
374	}
375	// use checkpoint from LocalStateDB
376	return l.StateDB.Reset(batchNum)
377}