Add snapshotter extension

modules/light-clients/08-wasm/keeper/export_test.go

@@ -0,0 +1,9 @@
+package keeper
+
+/*
+	This file is to allow for unexported functions to be accessible to the testing package.
+*/
+
+func GenerateWasmCodeHash(code []byte) []byte {
+	return generateWasmCodeHash(code)
+}

modules/light-clients/08-wasm/keeper/snapshotter.go

@@ -0,0 +1,150 @@
+package keeper
+
+import (
+	"encoding/hex"
+	"io"
+
+	errorsmod "cosmossdk.io/errors"
+
+	snapshot "cosmossdk.io/store/snapshots/types"
+	storetypes "cosmossdk.io/store/types"
+	tmproto "github.com/cometbft/cometbft/proto/tendermint/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/cosmos/ibc-go/modules/light-clients/08-wasm/types"
+)
+
+var _ snapshot.ExtensionSnapshotter = &WasmSnapshotter{}
+
+// SnapshotFormat defines the default snapshot extension encoding format.
+// SnapshotFormat 1 is gzipped wasm byte code for each item payload. No protobuf envelope, no metadata.
+const SnapshotFormat = 1
+
+// WasmSnapshotter implements the snapshot.ExtensionSnapshotter interface and is used to
+// import and export state maintained within the wasmvm cache.
+// NOTE: The following ExtensionSnapshotter has been adapted from CosmWasm's x/wasm:
+// https://github.com/CosmWasm/wasmd/blob/v0.43.0/x/wasm/keeper/snapshotter.go
+type WasmSnapshotter struct {
+	cms    storetypes.MultiStore
+	keeper *Keeper
+}
+
+// NewWasmSnapshotter creates and returns a new snapshot.ExtensionSnapshotter implementation for the 08-wasm module.
+func NewWasmSnapshotter(cms storetypes.MultiStore, keeper *Keeper) snapshot.ExtensionSnapshotter {
+	return &WasmSnapshotter{
+		cms:    cms,
+		keeper: keeper,
+	}
+}
+
+// SnapshotName implements the snapshot.ExtensionSnapshotter interface.
+// A unique name should be provided such that the implementation can be identified by the manager.
+func (*WasmSnapshotter) SnapshotName() string {
+	return types.ModuleName
+}
+
+// SnapshotFormat implements the snapshot.ExtensionSnapshotter interface.
+// This is the default format used for encoding payloads when taking a snapshot.
+func (*WasmSnapshotter) SnapshotFormat() uint32 {
+	return SnapshotFormat
+}
+
+// SupportedFormats implements the snapshot.ExtensionSnapshotter interface.
+// This defines a list of supported formats the snapshotter extension can restore from.
+func (*WasmSnapshotter) SupportedFormats() []uint32 {
+	// If we support older formats, add them here and handle them in Restore
+	return []uint32{SnapshotFormat}
+}
+
+// SnapshotExtension implements the snapshot.ExntensionSnapshotter interface.
+// SnapshotExtension is used to write data payloads into the underlying protobuf stream from the 08-wasm module.
+func (ws *WasmSnapshotter) SnapshotExtension(height uint64, payloadWriter snapshot.ExtensionPayloadWriter) error {
+	cacheMS, err := ws.cms.CacheMultiStoreWithVersion(int64(height))
+	if err != nil {
+		return err
+	}
+
+	ctx := sdk.NewContext(cacheMS, tmproto.Header{}, false, nil)
+
+	codeHashes, err := types.GetAllCodeHashes(ctx)
+	if err != nil {
+		return err
+	}
+
+	for _, codeHash := range codeHashes {
+		wasmCode, err := ws.keeper.wasmVM.GetCode(codeHash)
+		if err != nil {
+			return err
+		}
+
+		compressedWasm, err := types.GzipIt(wasmCode)
+		if err != nil {
+			return err
+		}
+
+		if err = payloadWriter(compressedWasm); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// RestoreExtension implements the snapshot.ExtensionSnapshotter interface.
+// RestoreExtension is used to read data from an existing extension state snapshot into the 08-wasm module.
+// The payload reader returns io.EOF when it has reached the end of the extension state snapshot.
+func (ws *WasmSnapshotter) RestoreExtension(height uint64, format uint32, payloadReader snapshot.ExtensionPayloadReader) error {
+	if format == SnapshotFormat {
+		return ws.processAllItems(height, payloadReader, restoreV1, finalizeV1)
+	}
+	return snapshot.ErrUnknownFormat
+}
+
+func restoreV1(ctx sdk.Context, k *Keeper, compressedCode []byte) error {
+	if !types.IsGzip(compressedCode) {
+		return types.ErrInvalid.Wrap("not a gzip")
+	}
+
+	wasmCode, err := types.Uncompress(compressedCode, types.MaxWasmByteSize())
+	if err != nil {
+		return errorsmod.Wrap(errorsmod.Wrap(err, "failed to store contract"), err.Error())
+	}
+
+	codeHash, err := k.wasmVM.StoreCode(wasmCode)
+	if err != nil {
+		return errorsmod.Wrap(errorsmod.Wrap(err, "failed to store contract"), err.Error())
+	}
+
+	if err := k.wasmVM.Pin(codeHash); err != nil {
+		return errorsmod.Wrapf(err, "failed to pin code hash: %s to in-memory cache", hex.EncodeToString(codeHash))
+	}
+
+	return nil
+}
+
+func finalizeV1(ctx sdk.Context, k *Keeper) error {
+	return nil
+}
+
+func (ws *WasmSnapshotter) processAllItems(
+	height uint64,
+	payloadReader snapshot.ExtensionPayloadReader,
+	cb func(sdk.Context, *Keeper, []byte) error,
+	finalize func(sdk.Context, *Keeper) error,
+) error {
+	ctx := sdk.NewContext(ws.cms, tmproto.Header{Height: int64(height)}, false, nil)
+	for {
+		payload, err := payloadReader()
+		if err == io.EOF {
+			break
+		} else if err != nil {
+			return err
+		}
+
+		if err := cb(ctx, ws.keeper, payload); err != nil {
+			return errorsmod.Wrap(err, "processing snapshot item")
+		}
+	}
+
+	return finalize(ctx, ws.keeper)
+}

modules/light-clients/08-wasm/keeper/snapshotter_test.go

@@ -0,0 +1,103 @@
+package keeper_test
+
+import (
+	"encoding/hex"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	govtypes "github.com/cosmos/cosmos-sdk/x/gov/types"
+
+	tmproto "github.com/cometbft/cometbft/proto/tendermint/types"
+
+	"github.com/cosmos/ibc-go/modules/light-clients/08-wasm/keeper"
+	"github.com/cosmos/ibc-go/modules/light-clients/08-wasm/testing/simapp"
+	"github.com/cosmos/ibc-go/modules/light-clients/08-wasm/types"
+)
+
+func TestSnapshotter(t *testing.T) {
+	testCases := []struct {
+		name      string
+		wasmFiles []string
+	}{
+		{
+			name:      "single contract",
+			wasmFiles: []string{"../test_data/ics10_grandpa_cw.wasm.gz"},
+		},
+
+		{
+			name:      "multiple contract",
+			wasmFiles: []string{"../test_data/ics07_tendermint_cw.wasm.gz", "../test_data/ics10_grandpa_cw.wasm.gz"},
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+
+		t.Run(tc.name, func(t *testing.T) {
+			wasmClientApp := simapp.SetupWithSnapShotter(t)
+			ctx := wasmClientApp.NewUncachedContext(false, tmproto.Header{
+				ChainID: "foo",
+				Height:  wasmClientApp.LastBlockHeight() + 1,
+				Time:    time.Now(),
+			})
+
+			var srcChecksumCodes []byte
+			var codeHashes [][]byte
+			// store contract on chain
+			for _, contractDir := range tc.wasmFiles {
+				signer := authtypes.NewModuleAddress(govtypes.ModuleName).String()
+				code, _ := os.ReadFile(contractDir)
+				msg := types.NewMsgStoreCode(signer, code)
+
+				res, err := wasmClientApp.WasmClientKeeper.StoreCode(ctx, msg)
+				codeHashes = append(codeHashes, res.Checksum)
+				srcChecksumCodes = append(srcChecksumCodes, res.Checksum...)
+
+				require.NoError(t, err)
+			}
+
+			// create snapshot
+			wasmClientApp.Commit()
+			snapshotHeight := uint64(wasmClientApp.LastBlockHeight())
+			snapshot, err := wasmClientApp.SnapshotManager().Create(snapshotHeight)
+			require.NoError(t, err)
+			require.NotNil(t, snapshot)
+
+			// setup dest app with snapshot imported
+			destWasmClientApp := simapp.SetupWithEmptyStore(t)
+
+			require.NoError(t, destWasmClientApp.SnapshotManager().Restore(*snapshot))
+			for i := uint32(0); i < snapshot.Chunks; i++ {
+				chunkBz, err := wasmClientApp.SnapshotManager().LoadChunk(snapshot.Height, snapshot.Format, i)
+				require.NoError(t, err)
+				end, err := destWasmClientApp.SnapshotManager().RestoreChunk(chunkBz)
+				require.NoError(t, err)
+				if end {
+					break
+				}
+			}
+
+			var allDestAppCodeHashInWasmVMStore []byte
+			// check wasm contracts are imported
+			ctx = destWasmClientApp.NewUncachedContext(false, tmproto.Header{
+				ChainID: "foo",
+				Height:  destWasmClientApp.LastBlockHeight() + 1,
+				Time:    time.Now(),
+			})
+
+			for _, codeHash := range codeHashes {
+				resp, err := destWasmClientApp.WasmClientKeeper.Code(ctx, &types.QueryCodeRequest{CodeHash: hex.EncodeToString(codeHash)})
+				require.NoError(t, err)
+
+				allDestAppCodeHashInWasmVMStore = append(allDestAppCodeHashInWasmVMStore, keeper.GenerateWasmCodeHash(resp.Data)...)
+
+			}
+
+			require.Equal(t, srcChecksumCodes, allDestAppCodeHashInWasmVMStore)
+		})
+	}
+}

modules/light-clients/08-wasm/testing/simapp/app.go

@@ -767,6 +767,16 @@ func NewSimApp(
 	app.SetEndBlocker(app.EndBlocker)
 	app.setAnteHandler(txConfig)
 
+	// must be before Loading version
+	if manager := app.SnapshotManager(); manager != nil {
+		err := manager.RegisterExtensions(
+			wasmkeeper.NewWasmSnapshotter(app.CommitMultiStore(), &app.WasmClientKeeper),
+		)
+		if err != nil {
+			panic(fmt.Errorf("failed to register snapshot extension: %s", err))
+		}
+	}
+
 	// In v0.46, the SDK introduces _postHandlers_. PostHandlers are like
 	// antehandlers, but are run _after_ the `runMsgs` execution. They are also
 	// defined as a chain, and have the same signature as antehandlers.

modules/light-clients/08-wasm/testing/simapp/test_helpers.go

@@ -0,0 +1,109 @@
+package simapp
+
+import (
+	"encoding/json"
+	"path/filepath"
+	"testing"
+
+	"cosmossdk.io/log"
+	sdkmath "cosmossdk.io/math"
+	"cosmossdk.io/store/snapshots"
+	snapshottypes "cosmossdk.io/store/snapshots/types"
+	dbm "github.com/cosmos/cosmos-db"
+	"github.com/stretchr/testify/require"
+
+	bam "github.com/cosmos/cosmos-sdk/baseapp"
+	"github.com/cosmos/cosmos-sdk/client/flags"
+	"github.com/cosmos/cosmos-sdk/crypto/keys/secp256k1"
+	"github.com/cosmos/cosmos-sdk/server"
+	simtestutil "github.com/cosmos/cosmos-sdk/testutil/sims"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
+	"github.com/cosmos/ibc-go/v8/testing/mock"
+
+	abci "github.com/cometbft/cometbft/abci/types"
+	cmttypes "github.com/cometbft/cometbft/types"
+)
+
+func setup(tb testing.TB, chainID string, withGenesis bool, invCheckPeriod uint) (*SimApp, GenesisState) {
+	tb.Helper()
+	db := dbm.NewMemDB()
+	nodeHome := tb.TempDir()
+	snapshotDir := filepath.Join(nodeHome, "data", "snapshots")
+
+	snapshotDB, err := dbm.NewDB("metadata", dbm.GoLevelDBBackend, snapshotDir)
+	require.NoError(tb, err)
+	tb.Cleanup(func() { snapshotDB.Close() })
+	snapshotStore, err := snapshots.NewStore(snapshotDB, snapshotDir)
+	require.NoError(tb, err)
+
+	appOptions := make(simtestutil.AppOptionsMap, 0)
+	appOptions[flags.FlagHome] = nodeHome // ensure unique folder
+	appOptions[server.FlagInvCheckPeriod] = invCheckPeriod
+	app := NewSimApp(log.NewNopLogger(), db, nil, true, appOptions, nil, bam.SetChainID(chainID), bam.SetSnapshot(snapshotStore, snapshottypes.SnapshotOptions{KeepRecent: 2}))
+
+	if withGenesis {
+		return app, app.DefaultGenesis()
+	}
+
+	return app, GenesisState{}
+}
+
+// SetupWithEmptyStore set up a simapp instance with empty DB
+func SetupWithEmptyStore(tb testing.TB) *SimApp {
+	tb.Helper()
+
+	app, _ := setup(tb, "", false, 0)
+	return app
+}
+
+// SetupWithGenesisValSet initializes a new SimApp with a validator set and genesis accounts
+// that also act as delegators. For simplicity, each validator is bonded with a delegation
+// of one consensus engine unit in the default token of the simapp from first genesis
+// account. A Nop logger is set in SimApp.
+func SetupWithGenesisValSetSnapshotter(t *testing.T, valSet *cmttypes.ValidatorSet, genAccs []authtypes.GenesisAccount, balances ...banktypes.Balance) *SimApp {
+	t.Helper()
+
+	app, genesisState := setup(t, "", true, 5)
+	genesisState, err := simtestutil.GenesisStateWithValSet(app.AppCodec(), genesisState, valSet, genAccs, balances...)
+	require.NoError(t, err)
+
+	stateBytes, err := json.MarshalIndent(genesisState, "", " ")
+	require.NoError(t, err)
+
+	// init chain will set the validator set and initialize the genesis accounts
+	_, err = app.InitChain(&abci.RequestInitChain{
+		Validators:      []abci.ValidatorUpdate{},
+		ConsensusParams: simtestutil.DefaultConsensusParams,
+		AppStateBytes:   stateBytes,
+	})
+	require.NoError(t, err)
+
+	return app
+}
+
+// Setup initializes a new SimApp. A Nop logger is set in SimApp.
+func SetupWithSnapShotter(t *testing.T) *SimApp {
+	t.Helper()
+
+	privVal := mock.NewPV()
+	pubKey, err := privVal.GetPubKey()
+	require.NoError(t, err)
+
+	// create validator set with single validator
+	validator := cmttypes.NewValidator(pubKey, 1)
+	valSet := cmttypes.NewValidatorSet([]*cmttypes.Validator{validator})
+
+	// generate genesis account
+	senderPrivKey := secp256k1.GenPrivKey()
+	acc := authtypes.NewBaseAccount(senderPrivKey.PubKey().Address().Bytes(), senderPrivKey.PubKey(), 0, 0)
+	balance := banktypes.Balance{
+		Address: acc.GetAddress().String(),
+		Coins:   sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, sdkmath.NewInt(100000000000000))),
+	}
+
+	app := SetupWithGenesisValSetSnapshotter(t, valSet, []authtypes.GenesisAccount{acc}, balance)
+
+	return app
+}