feat: Test stateless execution for cosmos-sdk baseapp using SMT-based fraudproofs

baseapp/baseapp.go

@@ -1,6 +1,7 @@
 package baseapp
 
 import (
+	"bytes"
 	"fmt"
 	"strings"
 
@@ -13,6 +14,8 @@ import (
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
 	dbm "github.com/cosmos/cosmos-sdk/db"
 	"github.com/cosmos/cosmos-sdk/snapshots"
+	"github.com/cosmos/cosmos-sdk/store/tracekv"
+	"github.com/cosmos/cosmos-sdk/store/types"
 	stypes "github.com/cosmos/cosmos-sdk/store/v2alpha1"
 	"github.com/cosmos/cosmos-sdk/store/v2alpha1/multi"
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -241,7 +244,7 @@ func (app *BaseApp) loadStore() error {
 			return err
 		}
 	}
-	app.cms, err = multi.NewV1MultiStoreAsV2(app.db, config)
+	app.cms, err = multi.NewStore(app.db, config)
 	if err != nil {
 		return fmt.Errorf("failed to load store: %w", err)
 	}
@@ -293,6 +296,10 @@ func (app *BaseApp) setHaltHeight(haltHeight uint64) {
 	app.haltHeight = haltHeight
 }
 
+func (app *BaseApp) setInitialHeight(initialHeight int64) {
+	app.initialHeight = initialHeight
+}
+
 func (app *BaseApp) setHaltTime(haltTime uint64) {
 	app.haltTime = haltTime
 }
@@ -761,3 +768,30 @@ func (app *BaseApp) runMsgs(ctx sdk.Context, msgs []sdk.Msg, mode runTxMode) (*s
 func makeABCIData(msgResponses []*codectypes.Any) ([]byte, error) {
 	return proto.Marshal(&sdk.TxMsgData{MsgResponses: msgResponses})
 }
+
+func (app *BaseApp) generateFraudProof(storeKeyToSubstoreTraceBuf map[types.StoreKey]*bytes.Buffer) FraudProof {
+	var fraudProof FraudProof
+	fraudProof.stateWitness = make(map[string]StateWitness)
+	cms := app.cms.(*multi.Store)
+
+	for storeKey, subStoreTraceBuf := range storeKeyToSubstoreTraceBuf {
+		keys := cms.GetKVStore(storeKey).(*tracekv.Store).GetAllKeysUsedInTrace(*subStoreTraceBuf)
+
+		substoreSMT := cms.GetSubstoreSMT(storeKey.Name())
+		stateWitness := StateWitness{
+			WitnessData: make([]WitnessData, 0, keys.Len()),
+		}
+		for key := range keys {
+			value := substoreSMT.Get([]byte(key))
+			// Assumption: The keys exist in the SMT because they were traced
+			// TOOD: Investigate default value leaves
+			proof := substoreSMT.MustGetSMTProof([]byte(key))
+			bKey, bVal := []byte(key), []byte(value)
+			witnessData := WitnessData{bKey, bVal, *proof}
+			stateWitness.WitnessData = append(stateWitness.WitnessData, witnessData)
+		}
+		fraudProof.stateWitness[storeKey.Name()] = stateWitness
+	}
+
+	return fraudProof
+}

baseapp/baseapp_test.go

@@ -11,13 +11,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/gogo/protobuf/jsonpb"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-	abci "github.com/tendermint/tendermint/abci/types"
-	"github.com/tendermint/tendermint/libs/log"
-	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
-
 	"github.com/cosmos/cosmos-sdk/codec"
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
 	dbm "github.com/cosmos/cosmos-sdk/db"
@@ -33,11 +26,18 @@ import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	"github.com/cosmos/cosmos-sdk/x/auth/migrations/legacytx"
+	"github.com/gogo/protobuf/jsonpb"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	abci "github.com/tendermint/tendermint/abci/types"
+	"github.com/tendermint/tendermint/libs/log"
+	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
 )
 
 var (
-	capKey1 = sdk.NewKVStoreKey("key1")
-	capKey2 = sdk.NewKVStoreKey("key2")
+	capKey1    = sdk.NewKVStoreKey("key1")
+	capKey2    = sdk.NewKVStoreKey("key2")
+	randSource = int64(123456789)
 
 	// testTxPriority is the CheckTx priority that we set in the test
 	// antehandler.
@@ -98,6 +98,54 @@ 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 {
+			// TODO:
+			// Verify proof inside 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()
@@ -2153,3 +2201,122 @@ func TestBaseApp_Init(t *testing.T) {
 		require.Equal(t, tc.expectedSnapshot.KeepRecent, tc.bapp.snapshotManager.GetKeepRecent())
 	}
 }
+
+func executeBlockWithArbitraryTxs(t *testing.T, app *BaseApp, numTransactions int, blockHeight int64) []txTest {
+	codec := codec.NewLegacyAmino()
+	registerTestCodec(codec)
+	r := rand.New(rand.NewSource(randSource))
+	randSource += 1
+	keyCounter := r.Intn(100)
+	txs := make([]txTest, 0)
+
+	app.BeginBlock(abci.RequestBeginBlock{Header: tmproto.Header{Height: blockHeight}})
+	for txNum := 0; txNum < numTransactions; txNum++ {
+		tx := txTest{Msgs: []sdk.Msg{}}
+		for msgNum := 0; msgNum < 2; msgNum++ {
+			key := []byte(fmt.Sprintf("%v", keyCounter))
+			value := make([]byte, 10000)
+			_, err := r.Read(value)
+			require.NoError(t, err)
+			tx.Msgs = append(tx.Msgs, msgKeyValue{Key: key, Value: value})
+			keyCounter++
+		}
+		txBytes, err := codec.Marshal(tx)
+		require.NoError(t, err)
+		resp := app.DeliverTx(abci.RequestDeliverTx{Tx: txBytes})
+		require.True(t, resp.IsOK(), "%v", resp.String())
+		txs = append(txs, tx)
+	}
+	app.EndBlock(abci.RequestEndBlock{Height: blockHeight})
+	return txs
+}
+
+func executeBlock(t *testing.T, app *BaseApp, txs []txTest, blockHeight int64) {
+	codec := codec.NewLegacyAmino()
+	registerTestCodec(codec)
+
+	app.BeginBlock(abci.RequestBeginBlock{Header: tmproto.Header{Height: blockHeight}})
+	for _, tx := range txs {
+		txBytes, err := codec.Marshal(tx)
+		require.NoError(t, err)
+		resp := app.DeliverTx(abci.RequestDeliverTx{Tx: txBytes})
+		require.True(t, resp.IsOK(), "%v", resp.String())
+	}
+	app.EndBlock(abci.RequestEndBlock{Height: blockHeight})
+}
+
+func checkSubstoreSMTsEqual(appB1 *BaseApp, appB2 *BaseApp, storeKeyName string) bool {
+	cmsB1 := appB1.cms.(*multi.Store)
+	cmsB2 := appB2.cms.(*multi.Store)
+	smtB1 := cmsB1.GetSubstoreSMT(storeKeyName)
+	smtB2 := cmsB2.GetSubstoreSMT(storeKeyName)
+
+	return string(smtB1.Root()) == string(smtB2.Root())
+}
+
+func TestGenerateAndLoadFraudProof(t *testing.T) {
+	/*
+		Tests switch between a baseapp and fraudproof and covers parts of the fraudproof cycle. Steps:
+		1. Initialize a baseapp, B1, with some state, S0
+		2. Make some state transition to state S1 by doing some transactions and commit
+		3. Export that state S1 into a fraudProof data structure (minimal snapshot)
+		4. Load a fresh baseapp, B2, with the contents of fraud proof data structure from (3) so it can begin from state S1.
+		5. If the state of both B1 and B2 has to converge at a state S2, then we the test passes.
+
+		Note that the appHash from B1 and B2 will not be the same because the subset of storeKeys in both apps is different
+		so we compare subStores instead
+
+		Tests to write in future:
+
+		1. Block with bad txs: Txs that exceed gas limits, validateBasic fails, unregistered messages (see TestRunInvalidTransaction)
+		2. Block with invalid appHash at the end
+		3. Corrupted Fraud Proof: bad SMT format, insufficient key-value pairs inside SMT needed to verify fraud
+		4. Bad block, fraud proof needed, fraud proof works, chain halts (happy case)
+
+	*/
+
+	storeTraceBuf := &bytes.Buffer{}
+	subStoreTraceBuf := &bytes.Buffer{}
+
+	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
+
+	appB1 := setupBaseApp(t,
+		AppOptionFunc(routerOpt),
+		SetSubstoreTracer(storeTraceBuf),
+		SetTracerFor(capKey2, subStoreTraceBuf),
+	)
+
+	// B1 <- S0
+	appB1.InitChain(abci.RequestInitChain{})
+
+	numTransactions := 5
+	// B1 <- S1
+	executeBlockWithArbitraryTxs(t, appB1, numTransactions, 1)
+	appB1.Commit()
+
+	// Exports all data inside current multistore into a fraudProof (S0 -> S1) //
+	storeKeyToSubstoreTraceBuf := make(map[stypes.StoreKey]*bytes.Buffer)
+	storeKeyToSubstoreTraceBuf[capKey2] = subStoreTraceBuf
+
+	// Records S1 in fraudproof
+	fraudProof := appB1.generateFraudProof(storeKeyToSubstoreTraceBuf)
+	currentBlockHeight := appB1.LastBlockHeight() // Only changes on a Commit
+	fraudProof.blockHeight = currentBlockHeight + 1
+
+	// Light Client
+
+	// TODO: Insert fraudproof verification here
+
+	// Now we take contents of the fraud proof which was recorded with S2 and try to populate a fresh baseapp B2 with it
+	// B2 <- S1
+	appB2 := setupBaseAppFromFraudProof(t, fraudProof)
+	require.True(t, checkSubstoreSMTsEqual(appB1, appB2, capKey2.Name()))
+}

baseapp/fraudproof.go

@@ -0,0 +1,25 @@
+package baseapp
+
+import "github.com/lazyledger/smt"
+
+// Represents a single-round fraudProof
+type FraudProof struct {
+	// The block height to load state of
+	blockHeight int64
+
+	// A map from module name to state witness
+	stateWitness map[string]StateWitness
+}
+
+// State witness with a list of all witness data
+type StateWitness struct {
+	// List of witness data
+	WitnessData []WitnessData
+}
+
+// Witness data containing a key/value pair and a SMT proof for said key/value pair
+type WitnessData struct {
+	Key   []byte
+	Value []byte
+	proof smt.SparseMerkleProof
+}

baseapp/options.go

@@ -21,6 +21,33 @@ func SetPruning(opts pruningtypes.PruningOptions) StoreOption {
 	return func(config *multi.StoreParams, _ uint64) error { config.Pruning = opts; return nil }
 }
 
+// SetSubstoreTracer provides a BaseApp option function that sets the
+// tracer for a multistore.
+func SetSubstoreTracer(w io.Writer) StoreOption {
+	return func(cfg *multi.StoreParams, _ uint64) error {
+		cfg.SetTracer(w)
+		return nil
+	}
+}
+
+// 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 {
+	return func(cfg *multi.StoreParams, _ uint64) error {
+		cfg.SetTracerFor(skey, w)
+		return nil
+	}
+}
+
+// SetSubstoreKVPair sets a key, value pair for the given substore inside a multistore
+// Only works for v2alpha1/multi
+func SetSubstoreKVPair(skey storetypes.StoreKey, key, val []byte) AppOptionOrdered {
+	return AppOptionOrdered{
+		func(bapp *BaseApp) { bapp.cms.(*multi.Store).SetSubstoreKVPair(skey, key, val) },
+		OptionOrderAfterStore,
+	}
+}
+
 // SetMinGasPrices returns an option that sets the minimum gas prices on the app.
 func SetMinGasPrices(gasPricesStr string) AppOptionFunc {
 	gasPrices, err := sdk.ParseDecCoins(gasPricesStr)
@@ -36,6 +63,11 @@ func SetHaltHeight(blockHeight uint64) AppOptionFunc {
 	return func(bap *BaseApp) { bap.setHaltHeight(blockHeight) }
 }
 
+// SetInitialHeight returns a BaseApp option function that sets the initial block height.
+func SetInitialHeight(blockHeight int64) AppOptionFunc {
+	return func(bapp *BaseApp) { bapp.setInitialHeight(blockHeight) }
+}
+
 // SetHaltTime returns a BaseApp option function that sets the halt block time.
 func SetHaltTime(haltTime uint64) AppOptionFunc {
 	return func(bap *BaseApp) { bap.setHaltTime(haltTime) }

go.mod

@@ -65,7 +65,10 @@ require (
 	sigs.k8s.io/yaml v1.3.0
 )
 
-require github.com/cosmos/cosmos-sdk/store/tools/ics23 v0.0.0-00010101000000-000000000000
+require (
+	github.com/chrispappas/golang-generics-set v1.0.1
+	github.com/cosmos/cosmos-sdk/store/tools/ics23 v0.0.0-00010101000000-000000000000
+)
 
 require (
 	4d63.com/gochecknoglobals v0.1.0 // indirect

go.sum

@@ -229,6 +229,8 @@ github.com/chavacava/garif v0.0.0-20220316182200-5cad0b5181d4 h1:tFXjAxje9thrTF4
 github.com/chavacava/garif v0.0.0-20220316182200-5cad0b5181d4/go.mod h1:W8EnPSQ8Nv4fUjc/v1/8tHFqhuOJXnRub0dTfuAQktU=
 github.com/checkpoint-restore/go-criu/v5 v5.3.0/go.mod h1:E/eQpaFtUKGOOSEBZgmKAcn+zUUwWxqcaKZlF54wK8E=
 github.com/cheggaaa/pb v1.0.27/go.mod h1:pQciLPpbU0oxA0h+VJYYLxO+XeDQb5pZijXscXHm81s=
+github.com/chrispappas/golang-generics-set v1.0.1 h1:91l8cInAWTxCPwZ8UNg7qkkPsdFdkYS9hytsd8UJsIU=
+github.com/chrispappas/golang-generics-set v1.0.1/go.mod h1:cp8j73+rlDyFF9PrjUkrRvi8L4jSRIsRK6Q1nPPIoqo=
 github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI=
 github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI=
 github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU=