You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

392 lines
11 KiB

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