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.

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