Add snapshotter extension (#4723)

* add snapshotter

* defind AbsoluteCodePositionLen and testing

* fix IterateCode

* init TestIterateCode

* nits

* add snapshotter manager to simapp

* Fix minor issues, lint.

* Add export_test.go file for exporting functions/methods for use in tests.

* spec -> tc.

* use ElementsMatch in TestIterateCode

* linting package import

* Refactor to use 08-wasm simapp.

* Fix sdk 0.50 issues.

* chore: adapt snapshotter code to obtain registered codeHashes and query wasm code from vm

* chore: adding godocs to all exported methods of WasmSnapshotter

* chore: adding reference link to cosmwasm's x/wasm snapshotter implementation

* chore: do not discard codeHash checksum and pin code to in-memory cache upon restore

* make lint-fix

* Use mockVM with snapshotter.

* Move default setup for callbacks into constructor.

* chore: rm finalize func and cleanup error returns

* chore: update casing on func name and add additional assertions to snapshotter test

---------

Co-authored-by: Carlos Rodriguez <carlos@interchain.io>
Co-authored-by: DimitrisJim <d.f.hilliard@gmail.com>
Co-authored-by: Damian Nolan <damiannolan@gmail.com>
Co-authored-by: Colin Axnér <25233464+colin-axner@users.noreply.github.com>

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/keeper_test.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/cosmos/ibc-go/modules/light-clients/08-wasm/internal/ibcwasm"
 	"github.com/cosmos/ibc-go/modules/light-clients/08-wasm/keeper"
+	wasmtesting "github.com/cosmos/ibc-go/modules/light-clients/08-wasm/testing"
 	"github.com/cosmos/ibc-go/modules/light-clients/08-wasm/testing/simapp"
 	"github.com/cosmos/ibc-go/modules/light-clients/08-wasm/types"
 	ibctesting "github.com/cosmos/ibc-go/v8/testing"
@@ -26,6 +27,8 @@ type KeeperTestSuite struct {
 
 	coordinator *ibctesting.Coordinator
 
+	// mockVM is a mock wasm VM that implements the WasmEngine interface
+	mockVM *wasmtesting.MockWasmEngine
 	chainA *ibctesting.TestChain
 }
 
@@ -58,6 +61,12 @@ func (suite *KeeperTestSuite) SetupTest() {
 	types.RegisterQueryServer(queryHelper, GetSimApp(suite.chainA).WasmClientKeeper)
 }
 
+func (suite *KeeperTestSuite) SetupSnapshotterWithMockVM() *simapp.SimApp {
+	suite.mockVM = wasmtesting.NewMockWasmEngine()
+
+	return simapp.SetupWithSnapshotter(suite.T(), suite.mockVM)
+}
+
 func TestKeeperTestSuite(t *testing.T) {
 	testifysuite.Run(t, new(KeeperTestSuite))
 }

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

@@ -0,0 +1,146 @@
+package keeper
+
+import (
+	"encoding/hex"
+	"io"
+
+	errorsmod "cosmossdk.io/errors"
+	snapshot "cosmossdk.io/store/snapshots/types"
+	storetypes "cosmossdk.io/store/types"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	tmproto "github.com/cometbft/cometbft/proto/tendermint/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 {
+	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 == ws.SnapshotFormat() {
+		return ws.processAllItems(height, payloadReader, restoreV1)
+	}
+
+	return errorsmod.Wrapf(snapshot.ErrUnknownFormat, "expected %d, got %d", ws.SnapshotFormat(), format)
+}
+
+func restoreV1(ctx sdk.Context, k *Keeper, compressedCode []byte) error {
+	if !types.IsGzip(compressedCode) {
+		return errorsmod.Wrap(types.ErrInvalidData, "expected wasm code is not gzip format")
+	}
+
+	wasmCode, err := types.Uncompress(compressedCode, types.MaxWasmByteSize())
+	if err != nil {
+		return errorsmod.Wrap(err, "failed to uncompress wasm code")
+	}
+
+	codeHash, err := k.wasmVM.StoreCode(wasmCode)
+	if err != nil {
+		return errorsmod.Wrap(err, "failed to store wasm code")
+	}
+
+	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 (ws *WasmSnapshotter) processAllItems(
+	height uint64,
+	payloadReader snapshot.ExtensionPayloadReader,
+	cb func(sdk.Context, *Keeper, []byte) 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, "failure processing snapshot item")
+		}
+	}
+
+	return nil
+}

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

@@ -0,0 +1,119 @@
+package keeper_test
+
+import (
+	"encoding/hex"
+	"time"
+
+	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"
+	wasmtesting "github.com/cosmos/ibc-go/modules/light-clients/08-wasm/testing"
+	"github.com/cosmos/ibc-go/modules/light-clients/08-wasm/testing/simapp"
+	"github.com/cosmos/ibc-go/modules/light-clients/08-wasm/types"
+)
+
+func (suite *KeeperTestSuite) TestSnapshotter() {
+	gzippedContract, err := types.GzipIt([]byte("gzipped-contract"))
+	suite.Require().NoError(err)
+
+	testCases := []struct {
+		name      string
+		contracts [][]byte
+	}{
+		{
+			name:      "single contract",
+			contracts: [][]byte{wasmtesting.Code},
+		},
+		{
+			name:      "multiple contracts",
+			contracts: [][]byte{wasmtesting.Code, gzippedContract},
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+
+		suite.Run(tc.name, func() {
+			t := suite.T()
+			wasmClientApp := suite.SetupSnapshotterWithMockVM()
+
+			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 _, contract := range tc.contracts {
+				signer := authtypes.NewModuleAddress(govtypes.ModuleName).String()
+				msg := types.NewMsgStoreCode(signer, contract)
+
+				res, err := wasmClientApp.WasmClientKeeper.StoreCode(ctx, msg)
+				suite.Require().NoError(err)
+
+				codeHashes = append(codeHashes, res.Checksum)
+				srcChecksumCodes = append(srcChecksumCodes, res.Checksum...)
+
+				suite.Require().NoError(err)
+			}
+
+			// create snapshot
+			res, err := wasmClientApp.Commit()
+			suite.Require().NoError(err)
+			suite.Require().NotNil(res)
+
+			snapshotHeight := uint64(wasmClientApp.LastBlockHeight())
+			snapshot, err := wasmClientApp.SnapshotManager().Create(snapshotHeight)
+			suite.Require().NoError(err)
+			suite.Require().NotNil(snapshot)
+
+			// setup dest app with snapshot imported
+			destWasmClientApp := simapp.SetupWithEmptyStore(t, suite.mockVM)
+			destCtx := destWasmClientApp.NewUncachedContext(false, tmproto.Header{
+				ChainID: "bar",
+				Height:  destWasmClientApp.LastBlockHeight() + 1,
+				Time:    time.Now(),
+			})
+
+			resp, err := destWasmClientApp.WasmClientKeeper.CodeHashes(destCtx, &types.QueryCodeHashesRequest{})
+			suite.Require().NoError(err)
+			suite.Require().Empty(resp.CodeHashes)
+
+			suite.Require().NoError(destWasmClientApp.SnapshotManager().Restore(*snapshot))
+
+			for i := uint32(0); i < snapshot.Chunks; i++ {
+				chunkBz, err := wasmClientApp.SnapshotManager().LoadChunk(snapshot.Height, snapshot.Format, i)
+				suite.Require().NoError(err)
+
+				end, err := destWasmClientApp.SnapshotManager().RestoreChunk(chunkBz)
+				suite.Require().NoError(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)})
+				suite.Require().NoError(err)
+
+				allDestAppCodeHashInWasmVMStore = append(allDestAppCodeHashInWasmVMStore, keeper.GenerateWasmCodeHash(resp.Data)...)
+			}
+
+			suite.Require().Equal(srcChecksumCodes, allDestAppCodeHashInWasmVMStore)
+		})
+	}
+}

modules/light-clients/08-wasm/testing/mock_engine.go

@@ -1,7 +1,10 @@
 package testing
 
 import (
+	"crypto/sha256"
+	"encoding/binary"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"reflect"
 
@@ -34,8 +37,9 @@ type (
 
 func NewMockWasmEngine() *MockWasmEngine {
 	m := &MockWasmEngine{
-		queryCallbacks: map[string]queryFn{},
-		sudoCallbacks:  map[string]sudoFn{},
+		queryCallbacks:  map[string]queryFn{},
+		sudoCallbacks:   map[string]sudoFn{},
+		storedContracts: map[uint32][]byte{},
 	}
 
 	for _, msgType := range queryTypes {
@@ -52,6 +56,27 @@ func NewMockWasmEngine() *MockWasmEngine {
 		}
 	}
 
+	// Set up default behavior for Store/Pin/Get
+	m.StoreCodeFn = func(code wasmvm.WasmCode) (wasmvm.Checksum, error) {
+		hash := sha256.Sum256(code)
+		checkSum := wasmvm.Checksum(hash[:])
+
+		m.storedContracts[binary.LittleEndian.Uint32(checkSum)] = code
+		return checkSum, nil
+	}
+
+	m.PinFn = func(codeID wasmvm.Checksum) error {
+		return nil
+	}
+
+	m.GetCodeFn = func(codeID wasmvm.Checksum) (wasmvm.WasmCode, error) {
+		code, ok := m.storedContracts[binary.LittleEndian.Uint32(codeID)]
+		if !ok {
+			return nil, errors.New("code not found")
+		}
+		return code, nil
+	}
+
 	return m
 }
 
@@ -85,6 +110,9 @@ type MockWasmEngine struct {
 	// queryCallbacks contains a mapping of queryMsg field type name to callback function.
 	queryCallbacks map[string]queryFn
 	sudoCallbacks  map[string]sudoFn
+
+	// contracts contains a mapping of code hash to code.
+	storedContracts map[uint32][]byte
 }
 
 // StoreCode implements the WasmEngine interface.

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.