feat(app): rollback many blocks

client/app/rollback.go

@@ -1,76 +1,24 @@
 package app
 
 import (
-	"context"
-	"fmt"
-	"math/big"
-
-	"cosmossdk.io/store"
-	storemetrics "cosmossdk.io/store/metrics"
-
 	cmtcmd "github.com/cometbft/cometbft/cmd/cometbft/commands"
-	dbm "github.com/cosmos/cosmos-db"
-	"github.com/ethereum/go-ethereum/common"
+	cmtcfg "github.com/cometbft/cometbft/config"
 
+	"github.com/piplabs/story/client/config"
 	"github.com/piplabs/story/lib/errors"
-	"github.com/piplabs/story/lib/log"
 )
 
-func Rollback(ctx context.Context, cfg Config, removeBlock bool) error {
-	engineCl, err := newEngineClient(ctx, cfg)
-	if err != nil {
-		return err
-	}
-
-	latestHeigth, err := engineCl.BlockNumber(ctx)
-	if err != nil {
-		return err
-	}
-
-	latestBlock, err := engineCl.BlockByNumber(ctx, big.NewInt(int64(latestHeigth)))
-	if err != nil {
-		return err
-	} else if latestBlock.BeaconRoot() == nil {
-		return errors.New("cannot rollback EVM with nil beacon root", "height", latestHeigth)
-	}
-
-	db, err := dbm.NewDB("application", cfg.BackendType(), cfg.DataDir())
-	if err != nil {
-		return errors.Wrap(err, "create db")
-	}
-
-	// Rollback CometBFT state
-	height, hash, err := cmtcmd.RollbackState(&cfg.Comet, removeBlock)
-	if err != nil {
-		return errors.Wrap(err, "rollback comet state")
+func RollbackCometAndAppState(a *App, cometCfg cmtcfg.Config, rollbackCfg config.RollbackConfig) (lastHeight int64, lastHash []byte, err error) {
+	for range rollbackCfg.RollbackHeights {
+		lastHeight, lastHash, err = cmtcmd.RollbackState(&cometCfg, rollbackCfg.RemoveBlock)
+		if err != nil {
+			return lastHeight, lastHash, errors.Wrap(err, "failed to rollback CometBFT state")
+		}
 	}
 
-	// Rollback the multistore
-	cms := store.NewCommitMultiStore(db, newSDKLogger(ctx), storemetrics.NewNoOpMetrics())
-	if err := cms.RollbackToVersion(height); err != nil {
-		return errors.Wrap(err, "rollback to height")
+	if err = a.CommitMultiStore().RollbackToVersion(lastHeight); err != nil {
+		return lastHeight, lastHash, errors.Wrap(err, "failed to rollback to version")
 	}
 
-	log.Info(ctx, "Rolled back consensus state", "height", height, "hash", fmt.Sprintf("%X", hash))
-
-	// Rollback EVM if latest EVM block built on-top of new rolled-back consensus head.
-	if *latestBlock.BeaconRoot() != common.BytesToHash(hash) {
-		return errors.New("cannot rollback EVM, latest EVM block not built on new rolled-back state",
-			"evm_height", latestHeigth,
-			"evm_beacon_root", *latestBlock.BeaconRoot(),
-		)
-	}
-
-	if err := engineCl.SetHead(ctx, latestHeigth-1); err != nil {
-		return errors.Wrap(err, "set head")
-	}
-
-	rolledBackBlock, err := engineCl.BlockByNumber(ctx, big.NewInt(int64(latestHeigth-1)))
-	if err != nil {
-		return err
-	}
-
-	log.Info(ctx, "Rolled back execution state", "height", rolledBackBlock.Number(), "hash", rolledBackBlock.Hash())
-
-	return nil
+	return lastHeight, lastHash, nil
 }

client/cmd/cmd.go

@@ -3,16 +3,13 @@ package cmd
 
 import (
 	"context"
-	"fmt"
 
-	cmtcmd "github.com/cometbft/cometbft/cmd/cometbft/commands"
 	"github.com/spf13/cobra"
 
 	"github.com/piplabs/story/client/app"
 	storycfg "github.com/piplabs/story/client/config"
 	"github.com/piplabs/story/lib/buildinfo"
 	libcmd "github.com/piplabs/story/lib/cmd"
-	"github.com/piplabs/story/lib/errors"
 	"github.com/piplabs/story/lib/log"
 )
 
@@ -65,59 +62,3 @@ func newRunCmd(name string, runFunc func(context.Context, app.Config) error) *co
 
 	return cmd
 }
-
-// newRollbackCmd returns a new cobra command that rolls back one block of the story consensus client.
-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 errors.Wrap(err, "failed to rollback CometBFT state")
-			}
-
-			if err = app.CommitMultiStore().RollbackToVersion(height); err != nil {
-				return errors.Wrap(err, "failed to rollback to version")
-			}
-
-			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
-}

client/cmd/flags.go

@@ -141,8 +141,9 @@ func bindKeyConvertFlags(cmd *cobra.Command, cfg *keyConfig) {
 	cmd.Flags().StringVar(&cfg.PubKeyHexUncompressed, "pubkey-hex-uncompressed", "", "Uncompressed public key in hex format")
 }
 
-func bindRollbackFlags(cmd *cobra.Command, cfg *config.Config) {
-	cmd.Flags().BoolVar(&cfg.RemoveBlock, "hard", false, "remove last block as well as state")
+func bindRollbackFlags(cmd *cobra.Command, cfg *config.RollbackConfig) {
+	cmd.Flags().BoolVar(&cfg.RemoveBlock, "hard", false, "remove Cosmos & CometBFT blocks as well as states")
+	cmd.Flags().Uint64VarP(&cfg.RollbackHeights, "number", "n", 1, "number of blocks to rollback")
 }
 
 func bindValidatorUnjailFlags(cmd *cobra.Command, cfg *unjailConfig) {

client/cmd/rollback.go

@@ -0,0 +1,69 @@
+package cmd
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/spf13/cobra"
+
+	"github.com/piplabs/story/client/app"
+	cfg "github.com/piplabs/story/client/config"
+	libcmd "github.com/piplabs/story/lib/cmd"
+	"github.com/piplabs/story/lib/log"
+)
+
+// newRollbackCmd returns a new cobra command that rolls back one block of the story consensus client.
+func newRollbackCmd(appCreateFunc func(context.Context, app.Config) *app.App) *cobra.Command {
+	rollbackCfg := cfg.DefaultRollbackConfig()
+	storyCfg := cfg.DefaultConfig()
+	logCfg := log.DefaultConfig()
+
+	cmd := &cobra.Command{
+		Use:   "rollback",
+		Short: "rollback Cosmos SDK and CometBFT state by X 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 - X.
+The application also rolls back to height n - X. If --hard=true, the block
+itself will also be deleted and re-downloaded from the p2p network. Note that
+different blocks from n - X to n cannot be re-built/re-proposed since that would result in validator slashing.
+If --hard=false, No blocks are removed, so upon restarting CometBFT the transactions in blocks [n - X + 1, 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
+			}
+
+			appCfg := app.Config{
+				Config: storyCfg,
+				Comet:  cometCfg,
+			}
+			a := appCreateFunc(ctx, appCfg)
+			lastHeight, lastHash, err := app.RollbackCometAndAppState(a, cometCfg, rollbackCfg)
+			if err != nil {
+				return err
+			}
+
+			fmt.Printf("Rolled back state to height %d and hash %X", lastHeight, lastHash)
+
+			return nil
+		},
+	}
+
+	bindRunFlags(cmd, &storyCfg)
+	bindRollbackFlags(cmd, &rollbackCfg)
+	log.BindFlags(cmd.Flags(), &logCfg)
+
+	return cmd
+}

client/config/config.go

@@ -181,7 +181,6 @@ type Config struct {
 	ExternalAddress    string
 	Seeds              string
 	SeedMode           bool
-	RemoveBlock        bool // See cosmos-sdk/server/rollback.go
 }
 
 // ConfigFile returns the default path to the toml story config file.

client/config/rollback.go

@@ -0,0 +1,13 @@
+package config
+
+type RollbackConfig struct {
+	RemoveBlock     bool // See cosmos-sdk/server/rollback.go
+	RollbackHeights uint64
+}
+
+func DefaultRollbackConfig() RollbackConfig {
+	return RollbackConfig{
+		RemoveBlock:     false,
+		RollbackHeights: 1,
+	}
+}