From 5271497ab8962863f3668e659d58597513e9f9f8 Mon Sep 17 00:00:00 2001
From: zsystm <actor93kor@gmail.com>
Date: Mon, 30 Sep 2024 16:52:33 +0900
Subject: [PATCH] feat: rollback cmd

add cometBFT's one block rollback cmd
---
 client/app/start.go     | 43 ++++++++++++++++++++++++++++++
 client/cmd/cmd.go       | 58 +++++++++++++++++++++++++++++++++++++++++
 client/cmd/flags.go     |  9 ++++---
 client/config/config.go |  1 +
 4 files changed, 108 insertions(+), 3 deletions(-)

diff --git a/client/app/start.go b/client/app/start.go
index ed9c8f36..da096759 100644
--- a/client/app/start.go
+++ b/client/app/start.go
@@ -184,6 +184,49 @@ func Start(ctx context.Context, cfg Config) (func(context.Context) error, error)
 	}, nil
 }
 
+func CreateApp(ctx context.Context, cfg Config) *App {
+	privVal, err := loadPrivVal(cfg)
+	if err != nil {
+		panic(errors.Wrap(err, "load validator key"))
+	}
+
+	db, err := dbm.NewDB("application", cfg.BackendType(), cfg.DataDir())
+	if err != nil {
+		panic(errors.Wrap(err, "create db"))
+	}
+
+	baseAppOpts, err := makeBaseAppOpts(cfg)
+	if err != nil {
+		panic(errors.Wrap(err, "make base app opts"))
+	}
+
+	engineCl, err := newEngineClient(ctx, cfg)
+	if err != nil {
+		panic(err)
+	}
+
+	//nolint:contextcheck // False positive
+	app, err := newApp(
+		newSDKLogger(ctx),
+		db,
+		engineCl,
+		baseAppOpts...,
+	)
+	if err != nil {
+		panic(errors.Wrap(err, "create app"))
+	}
+	app.Keepers.EVMEngKeeper.SetBuildDelay(cfg.EVMBuildDelay)
+	app.Keepers.EVMEngKeeper.SetBuildOptimistic(cfg.EVMBuildOptimistic)
+
+	addr, err := k1util.PubKeyToAddress(privVal.Key.PrivKey.PubKey())
+	if err != nil {
+		panic(errors.Wrap(err, "convert validator pubkey to address"))
+	}
+	app.Keepers.EVMEngKeeper.SetValidatorAddress(addr)
+
+	return app
+}
+
 func newCometNode(ctx context.Context, cfg *cmtcfg.Config, app *App, privVal cmttypes.PrivValidator,
 ) (*node.Node, error) {
 	nodeKey, err := p2p.LoadOrGenNodeKey(cfg.NodeKeyFile())
diff --git a/client/cmd/cmd.go b/client/cmd/cmd.go
index 9d742e7e..2d344449 100644
--- a/client/cmd/cmd.go
+++ b/client/cmd/cmd.go
@@ -3,7 +3,9 @@ package cmd
 
 import (
 	"context"
+	"fmt"
 
+	cmtcmd "github.com/cometbft/cometbft/cmd/cometbft/commands"
 	"github.com/spf13/cobra"
 
 	"github.com/piplabs/story/client/app"
@@ -23,6 +25,7 @@ func New() *cobra.Command {
 		buildinfo.NewVersionCmd(),
 		newValidatorCmds(),
 		newStatusCmd(),
+		newRollbackCmd(app.CreateApp),
 	)
 }
 
@@ -60,3 +63,58 @@ func newRunCmd(name string, runFunc func(context.Context, app.Config) error) *co
 
 	return cmd
 }
+
+func newRollbackCmd(appCreateFunc func(context.Context, app.Config) *app.App) *cobra.Command {
+	storyCfg := storycfg.DefaultConfig()
+	logCfg := log.DefaultConfig()
+
+	cmd := &cobra.Command{
+		Use:   "rollback",
+		Short: "rollback Cosmos SDK and CometBFT state by one height",
+		Long: `
+A state rollback is performed to recover from an incorrect application state transition,
+when CometBFT has persisted an incorrect app hash and is thus unable to make
+progress. Rollback overwrites a state at height n with the state at height n - 1.
+The application also rolls back to height n - 1. No blocks are removed, so upon
+restarting CometBFT the transactions in block n will be re-executed against the
+application.
+`,
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			ctx, err := log.Init(cmd.Context(), logCfg)
+			if err != nil {
+				return err
+			}
+			if err := libcmd.LogFlags(ctx, cmd.Flags()); err != nil {
+				return err
+			}
+
+			cometCfg, err := parseCometConfig(ctx, storyCfg.HomeDir)
+			if err != nil {
+				return err
+			}
+
+			app := appCreateFunc(ctx, app.Config{
+				Config: storyCfg,
+				Comet:  cometCfg,
+			})
+			height, hash, err := cmtcmd.RollbackState(&cometCfg, storyCfg.RemoveBlock)
+			if err != nil {
+				return err
+			}
+
+			if err := app.CommitMultiStore().RollbackToVersion(height); err != nil {
+				return err
+			}
+
+			fmt.Printf("Rolled back state to height %d and hash %X", height, hash)
+			return nil
+		},
+	}
+
+	bindRunFlags(cmd, &storyCfg)
+	bindRollbackFlags(cmd, &storyCfg)
+	log.BindFlags(cmd.Flags(), &logCfg)
+
+	return cmd
+
+}
diff --git a/client/cmd/flags.go b/client/cmd/flags.go
index b787b9dc..e6b0f975 100644
--- a/client/cmd/flags.go
+++ b/client/cmd/flags.go
@@ -1,6 +1,8 @@
 package cmd
 
 import (
+	// Used for ABI embedding of the staking contract.
+	_ "embed"
 	"fmt"
 	"path/filepath"
 	"strings"
@@ -12,9 +14,6 @@ import (
 	libcmd "github.com/piplabs/story/lib/cmd"
 	"github.com/piplabs/story/lib/netconf"
 	"github.com/piplabs/story/lib/tracer"
-
-	// Used for ABI embedding of the staking contract.
-	_ "embed"
 )
 
 func bindRunFlags(cmd *cobra.Command, cfg *config.Config) {
@@ -120,6 +119,10 @@ func bindStatusFlags(flags *pflag.FlagSet, cfg *StatusConfig) {
 	libcmd.BindHomeFlag(flags, &cfg.HomeDir)
 }
 
+func bindRollbackFlags(cmd *cobra.Command, cfg *config.Config) {
+	cmd.Flags().BoolVar(&cfg.RemoveBlock, "hard", false, "remove last block as well as state")
+}
+
 // Flag Validation
 
 func validateFlags(flags map[string]string) error {
diff --git a/client/config/config.go b/client/config/config.go
index c9bba36b..81be0888 100644
--- a/client/config/config.go
+++ b/client/config/config.go
@@ -160,6 +160,7 @@ type Config struct {
 	ExternalAddress    string
 	Seeds              string
 	SeedMode           bool
+	RemoveBlock        bool // See rollback
 }
 
 // ConfigFile returns the default path to the toml story config file.