feat: Implement fraudproof generation mode in cosmos-sdk

baseapp/baseapp.go

@@ -769,6 +769,53 @@ func makeABCIData(msgResponses []*codectypes.Any) ([]byte, error) {
 	return proto.Marshal(&sdk.TxMsgData{MsgResponses: msgResponses})
 }
 
+// enableFraudProofGenerationMode rolls back an app's state to a previous
+// state and enables tracing for the list of store keys
+// It returns the tracing-enabled app along with the trace buffers used
+func (app *BaseApp) enableFraudProofGenerationMode(storeKeys []types.StoreKey, routerOpts map[string]AppOptionFunc) (*BaseApp, map[string]*bytes.Buffer, error) {
+	cms := app.cms.(*multi.Store)
+	lastVersion := cms.LastCommitID().Version
+	previousCMS, err := cms.GetVersion(lastVersion)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	// Add options for tracing
+	storeTraceBuf := &bytes.Buffer{}
+
+	storeKeyToSubstoreTraceBuf := make(map[string]*bytes.Buffer)
+
+	// Initialize params from previousCMS
+	storeToLoadFrom := make(map[string]types.KVStore)
+	storeKeyNames := make([]string, 0, len(storeKeys))
+	for _, storeKey := range storeKeys {
+		storeKeyName := storeKey.Name()
+		storeKeyNames = append(storeKeyNames, storeKeyName)
+		storeToLoadFrom[storeKeyName] = previousCMS.GetKVStore(storeKey)
+		storeKeyToSubstoreTraceBuf[storeKeyName] = &bytes.Buffer{}
+	}
+
+	// BaseApp, B1
+	options := []AppOption{
+		SetSubstoreTracer(storeTraceBuf),
+	}
+
+	for _, storeKey := range storeKeys {
+		options = append(options, SetTracerFor(storeKey.Name(), storeKeyToSubstoreTraceBuf[storeKey.Name()]))
+		options = append(options, AppOptionFunc(routerOpts[storeKey.Name()]))
+	}
+	newBlockHeight := app.LastBlockHeight() + 1
+	newApp, err := SetupBaseAppFromParams(app.name+"WithTracing", app.logger, dbm.NewMemDB(), app.txDecoder, storeKeyNames, newBlockHeight, storeToLoadFrom, options...)
+
+	// Need to reset all the buffers to remove anything logged while setting up baseapp
+	storeTraceBuf.Reset()
+	for _, storeKey := range storeKeys {
+		storeKeyToSubstoreTraceBuf[storeKey.Name()].Reset()
+	}
+	return newApp, storeKeyToSubstoreTraceBuf, err
+}
+
+// Generate a fraudproof for an app with the given trace buffers
 func (app *BaseApp) generateFraudProof(storeKeyToSubstoreTraceBuf map[types.StoreKey]*bytes.Buffer) (FraudProof, error) {
 	fraudProof := FraudProof{}
 	fraudProof.stateWitness = make(map[string]StateWitness)
@@ -803,3 +850,36 @@ func (app *BaseApp) generateFraudProof(storeKeyToSubstoreTraceBuf map[types.Stor
 
 	return fraudProof, nil
 }
+
+// set up a new baseapp from given params
+func SetupBaseAppFromParams(appName string, logger log.Logger, db dbm.Connection, txDecoder sdk.TxDecoder, storeKeyNames []string, blockHeight int64, storeToLoadFrom map[string]types.KVStore, options ...AppOption) (*BaseApp, error) {
+	storeKeys := make([]types.StoreKey, 0, len(storeKeyNames))
+
+	for _, storeKeyName := range storeKeyNames {
+		storeKey := sdk.NewKVStoreKey(storeKeyName)
+		storeKeys = append(storeKeys, storeKey)
+		subStore := storeToLoadFrom[storeKeyName]
+		it := subStore.Iterator(nil, nil)
+		for ; it.Valid(); it.Next() {
+			key, val := it.Key(), it.Value()
+			options = append(options, SetSubstoreKVPair(storeKey, key, val))
+		}
+	}
+	options = append(options, SetSubstores(storeKeys...))
+
+	// This initial height is used in `BeginBlock` in `validateHeight`
+	options = append(options, SetInitialHeight(blockHeight))
+
+	// make list of options to pass by parsing fraudproof
+	app := NewBaseApp(appName, logger, db, txDecoder, options...)
+
+	// stores are mounted
+	err := app.Init()
+
+	return app, err
+}
+
+// set up a new baseapp from a fraudproof
+func SetupBaseAppFromFraudProof(appName string, logger log.Logger, db dbm.Connection, txDecoder sdk.TxDecoder, fraudProof FraudProof, options ...AppOption) (*BaseApp, error) {
+	return SetupBaseAppFromParams(appName, logger, db, txDecoder, fraudProof.getModules(), fraudProof.blockHeight, fraudProof.extractStore(), options...)
+}

baseapp/baseapp_test.go

@@ -17,6 +17,8 @@ import (
 	pruningtypes "github.com/cosmos/cosmos-sdk/pruning/types"
 	"github.com/cosmos/cosmos-sdk/snapshots"
 	snapshottypes "github.com/cosmos/cosmos-sdk/snapshots/types"
+	"github.com/cosmos/cosmos-sdk/store/tracekv"
+	"github.com/cosmos/cosmos-sdk/store/types"
 	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	stypes "github.com/cosmos/cosmos-sdk/store/v2alpha1"
 	"github.com/cosmos/cosmos-sdk/store/v2alpha1/multi"
@@ -98,51 +100,6 @@ func setupBaseApp(t *testing.T, options ...AppOption) *BaseApp {
 	return app
 }
 
-// baseapp loaded from a fraudproof
-func setupBaseAppFromFraudProof(t *testing.T, fraudProof FraudProof, options ...AppOption) *BaseApp {
-	storeKeys := make([]stypes.StoreKey, 0, len(fraudProof.stateWitness))
-	routerOpts := make([]func(bapp *BaseApp), 0, len(fraudProof.stateWitness))
-	for storeKeyName := range fraudProof.stateWitness {
-		storeKey := sdk.NewKVStoreKey(storeKeyName)
-		storeKeys = append(storeKeys, storeKey)
-
-		routerOpt := func(bapp *BaseApp) {
-			bapp.Router().AddRoute(sdk.NewRoute(routeMsgKeyValue, func(ctx sdk.Context, msg sdk.Msg) (*sdk.Result, error) {
-				kv := msg.(*msgKeyValue)
-				bapp.cms.GetKVStore(storeKey).Set(kv.Key, kv.Value)
-				return &sdk.Result{}, nil
-			}))
-		}
-		routerOpts = append(routerOpts, routerOpt)
-
-		stateWitness := fraudProof.stateWitness[storeKeyName]
-		witnessData := stateWitness.WitnessData
-		for _, witness := range witnessData {
-			options = append(options, SetSubstoreKVPair(storeKey, witness.Key, witness.Value))
-		}
-	}
-	options = append(options, SetSubstores(storeKeys...))
-
-	// RouterOpts should only be called after call to `SetSubstores`
-	for _, routerOpt := range routerOpts {
-		options = append(options, AppOptionFunc(routerOpt))
-	}
-
-	// This initial height is used in `BeginBlock` in `validateHeight`
-	options = append(options, SetInitialHeight(fraudProof.blockHeight))
-
-	// make list of options to pass by parsing fraudproof
-	app := newBaseApp(t.Name(), options...)
-	require.Equal(t, t.Name(), app.Name())
-
-	app.SetParamStore(mock.NewParamStore(dbm.NewMemDB()))
-
-	// stores are mounted
-	err := app.Init()
-	require.Nil(t, err)
-	return app
-}
-
 // simple one store baseapp with data and snapshots. Each tx is 1 MB in size (uncompressed).
 func setupBaseAppWithSnapshots(t *testing.T, config *setupConfig) (*BaseApp, error) {
 	codec := codec.NewLegacyAmino()
@@ -2253,6 +2210,86 @@ func checkSMTStoreEqual(appB1 *BaseApp, appB2 *BaseApp, storeKeyName string) boo
 	return bytes.Equal(storeHashB1, storeHashB2)
 }
 
+func TestFraudProofGenerationMode(t *testing.T) {
+	/*
+		Tests fraudproof generation mode. Steps:
+		1. Initialize a baseapp, B1, with some state, S0
+		2. Make some state transition to state S1 by doing some transactions, commit those transactions, and save
+		resulting SMT root hash
+		3. Make another set of state transitions to state S2 but do not commit
+		4. Get a new tracing enabled app by a call to enableFraudProofGenerationMode
+		5. Make sure the SMT root hashes of the new app is the same as the saved one
+		6. Execute another set of transactions and check if the traced keys are the same as the ones in the transactions
+	*/
+	routerOpt := func(bapp *BaseApp) {
+		bapp.Router().AddRoute(sdk.NewRoute(routeMsgKeyValue, func(ctx sdk.Context, msg sdk.Msg) (*sdk.Result, error) {
+			kv := msg.(*msgKeyValue)
+			bapp.cms.GetKVStore(capKey2).Set(kv.Key, kv.Value)
+			return &sdk.Result{}, nil
+		}))
+	}
+
+	// BaseApp, B1 with no Tracing
+	appB1 := setupBaseApp(t,
+		AppOptionFunc(routerOpt),
+	)
+
+	// B1 <- S0
+	appB1.InitChain(abci.RequestInitChain{})
+
+	numTransactions := 1
+	// B1 <- S1
+	executeBlockWithArbitraryTxs(t, appB1, numTransactions, 1)
+	appB1.Commit()
+
+	// storeHash of B1 at S1
+	cmsB1 := appB1.cms.(*multi.Store)
+	storeHashB1AtS1 := cmsB1.GetSubstoreSMT(capKey2.Name()).Root()
+
+	// B1 <- S2
+	executeBlockWithArbitraryTxs(t, appB1, numTransactions, 2)
+
+	// Do not commit here in order to preserve saved previous versions
+
+	// the only store key we'd like to enable tracing for
+	storeKeys := []types.StoreKey{capKey2}
+	routerOpts := make(map[string]AppOptionFunc)
+	routerOpts[capKey2.Name()] = func(bapp *BaseApp) {
+		bapp.Router().AddRoute(sdk.NewRoute(routeMsgKeyValue, func(ctx sdk.Context, msg sdk.Msg) (*sdk.Result, error) {
+			kv := msg.(*msgKeyValue)
+			cms := bapp.cms.(*multi.Store)
+			// There is only storeKey in the test for now
+			storeKey := cms.GetStoreKeys()[len(storeKeys)-1]
+			bapp.cms.GetKVStore(storeKey).Set(kv.Key, kv.Value)
+			return &sdk.Result{}, nil
+		}))
+	}
+	appB2, storeKeyToSubstoreTraceBuf, err := appB1.enableFraudProofGenerationMode(storeKeys, routerOpts)
+
+	require.Nil(t, err)
+	cmsB2 := appB2.cms.(*multi.Store)
+	for _, storeKey := range storeKeys {
+		storeHashB2 := cmsB2.GetSubstoreSMT(storeKey.Name()).Root()
+		require.Equal(t, storeHashB1AtS1, storeHashB2)
+	}
+
+	txs1 := executeBlockWithArbitraryTxs(t, appB2, numTransactions, 2)
+
+	for _, storeKey := range cmsB2.GetStoreKeys() {
+		subStoreBuf := storeKeyToSubstoreTraceBuf[storeKey.Name()]
+		tracedKeys := appB2.cms.GetKVStore(storeKey).(*tracekv.Store).GetAllKeysUsedInTrace(*subStoreBuf).Values()
+		txKeys := make([]string, 0)
+		for _, tx := range txs1 {
+			for _, msg := range tx.GetMsgs() {
+				msgKV := msg.(msgKeyValue)
+				txKeys = append(txKeys, (string(msgKV.Key)))
+			}
+		}
+		require.Equal(t, tracedKeys, txKeys)
+	}
+
+}
+
 func TestGenerateAndLoadFraudProof(t *testing.T) {
 	/*
 		Tests switch between a baseapp and fraudproof and covers parts of the fraudproof cycle. Steps:
@@ -2288,13 +2325,13 @@ func TestGenerateAndLoadFraudProof(t *testing.T) {
 	appB1 := setupBaseApp(t,
 		AppOptionFunc(routerOpt),
 		SetSubstoreTracer(storeTraceBuf),
-		SetTracerFor(capKey2, subStoreTraceBuf),
+		SetTracerFor(capKey2.Name(), subStoreTraceBuf),
 	)
 
 	// B1 <- S0
 	appB1.InitChain(abci.RequestInitChain{})
 
-	numTransactions := 5
+	numTransactions := 1
 	// B1 <- S1
 	executeBlockWithArbitraryTxs(t, appB1, numTransactions, 1)
 	commitB1 := appB1.Commit()
@@ -2315,8 +2352,11 @@ func TestGenerateAndLoadFraudProof(t *testing.T) {
 	require.Nil(t, err)
 	require.True(t, fraudProofVerified)
 
-	// Now we take contents of the fraud proof which was recorded with S1 and try to populate a fresh baseapp B2 with it
-	// B2 <- S1
-	appB2 := setupBaseAppFromFraudProof(t, fraudProof)
+	// Now we take contents of the fraud proof which was recorded with S2 and try to populate a fresh baseapp B2 with it
+	// B2 <- S2
+	codec := codec.NewLegacyAmino()
+	registerTestCodec(codec)
+	appB2, err := SetupBaseAppFromFraudProof(t.Name(), defaultLogger(), dbm.NewMemDB(), testTxDecoder(codec), fraudProof, AppOptionFunc(routerOpt))
+	require.Nil(t, err)
 	require.True(t, checkSMTStoreEqual(appB1, appB2, capKey2.Name()))
 }

baseapp/fraudproof.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"fmt"
 
+	"github.com/cosmos/cosmos-sdk/store/mem"
 	"github.com/cosmos/cosmos-sdk/store/types"
 	"github.com/cosmos/cosmos-sdk/store/v2alpha1/smt"
 	tmcrypto "github.com/tendermint/tendermint/proto/tendermint/crypto"
@@ -34,6 +35,27 @@ type WitnessData struct {
 	proof tmcrypto.ProofOp
 }
 
+func (fraudProof *FraudProof) getModules() []string {
+	keys := make([]string, 0, len(fraudProof.stateWitness))
+	for k := range fraudProof.stateWitness {
+		keys = append(keys, k)
+	}
+	return keys
+}
+
+func (fraudProof *FraudProof) extractStore() map[string]types.KVStore {
+	store := make(map[string]types.KVStore)
+	for storeKey, stateWitness := range fraudProof.stateWitness {
+		subStore := mem.NewStore()
+		for _, witnessData := range stateWitness.WitnessData {
+			key, val := witnessData.Key, witnessData.Value
+			subStore.Set(key, val)
+		}
+		store[storeKey] = subStore
+	}
+	return store
+}
+
 func (fraudProof *FraudProof) verifyFraudProof(headerAppHash []byte) (bool, error) {
 	for storeKey, stateWitness := range fraudProof.stateWitness {
 		proofOp := stateWitness.proof

baseapp/options.go

@@ -32,7 +32,7 @@ func SetSubstoreTracer(w io.Writer) StoreOption {
 
 // SetTracerFor provides a BaseApp option function that sets the
 // tracer for a substore with given skey inside a multistore.
-func SetTracerFor(skey storetypes.StoreKey, w io.Writer) StoreOption {
+func SetTracerFor(skey string, w io.Writer) StoreOption {
 	return func(cfg *multi.StoreParams, _ uint64) error {
 		cfg.SetTracerFor(skey, w)
 		return nil

store/v2alpha1/multi/params.go

@@ -16,7 +16,7 @@ func DefaultStoreParams() StoreParams {
 		SchemaBuilder:             newSchemaBuilder(),
 		storeKeys:                 storeKeys{},
 		traceListenMixin:          newTraceListenMixin(),
-		substoreTraceListenMixins: make(map[types.StoreKey]*traceListenMixin),
+		substoreTraceListenMixins: make(map[string]*traceListenMixin),
 	}
 }
 
@@ -43,7 +43,7 @@ func (par *StoreParams) RegisterSubstore(skey types.StoreKey, typ types.StoreTyp
 	return nil
 }
 
-func (par *StoreParams) SetTracerFor(skey types.StoreKey, w io.Writer) {
+func (par *StoreParams) SetTracerFor(skey string, w io.Writer) {
 	tlm := newTraceListenMixin()
 	tlm.SetTracer(w)
 	tlm.SetTracingContext(par.TraceContext)

store/v2alpha1/multi/store.go

@@ -77,7 +77,7 @@ type StoreParams struct {
 	Upgrades *types.StoreUpgrades
 	// Contains The trace context and listeners that can also be set from store methods.
 	*traceListenMixin
-	substoreTraceListenMixins map[types.StoreKey]*traceListenMixin
+	substoreTraceListenMixins map[string]*traceListenMixin
 }
 
 // StoreSchema defineds a mapping of substore keys to store types
@@ -114,7 +114,7 @@ type Store struct {
 
 	PersistentCache           types.MultiStorePersistentCache
 	substoreCache             map[string]*substore
-	substoreTraceListenMixins map[types.StoreKey]*traceListenMixin
+	substoreTraceListenMixins map[string]*traceListenMixin
 }
 
 type substore struct {
@@ -531,6 +531,14 @@ func (s *Store) SetSubstoreKVPair(skey types.StoreKey, kv, val []byte) {
 	sub.Set(kv, val)
 }
 
+func (s *Store) GetStoreKeys() []types.StoreKey {
+	storeKeys := make([]types.StoreKey, 0, len(s.schema))
+	for sk := range s.schema {
+		storeKeys = append(storeKeys, sk)
+	}
+	return storeKeys
+}
+
 // Resets a substore's state after commit (because root stateTxn has been discarded)
 func (s *substore) refresh(rootHash []byte) {
 	pfx := prefixSubstore(s.name)
@@ -1047,11 +1055,11 @@ func (tlm *traceListenMixin) getTracingContext() types.TraceContext {
 
 func (s *Store) wrapTraceListen(store types.KVStore, skey types.StoreKey) types.KVStore {
 	if s.TracingEnabled() {
-		subStoreTlm := s.substoreTraceListenMixins[skey]
+		subStoreTlm := s.substoreTraceListenMixins[skey.Name()]
 		store = tracekv.NewStore(store, subStoreTlm.TraceWriter, s.getTracingContext())
 	}
 	if s.ListeningEnabled(skey) {
-		subStoreTlm := s.substoreTraceListenMixins[skey]
+		subStoreTlm := s.substoreTraceListenMixins[skey.Name()]
 		store = listenkv.NewStore(store, skey, subStoreTlm.listeners[skey])
 	}
 	return store