refactor: use runtime env in ibc-wasm client module

cosmos · Dec 20, 2024 · 43ac57c · 43ac57c
1 parent bdcb8fd
commit 43ac57c
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 29 deletions.
diff --git a/modules/light-clients/08-wasm/keeper/keeper.go b/modules/light-clients/08-wasm/keeper/keeper.go
@@ -8,7 +8,7 @@ import (
 	wasmvm "github.com/CosmWasm/wasmvm/v2"
 
 	"cosmossdk.io/collections"
-	"cosmossdk.io/core/store"
+	"cosmossdk.io/core/appmodule"
 	errorsmod "cosmossdk.io/errors"
 	"cosmossdk.io/log"
 
@@ -22,16 +22,15 @@ import (
 
 // Keeper defines the 08-wasm keeper
 type Keeper struct {
-	// implements gRPC QueryServer interface
+	appmodule.Environment
 	types.QueryServer
 
 	cdc          codec.BinaryCodec
 	clientKeeper types.ClientKeeper
 
 	vm types.WasmEngine
 
-	checksums    collections.KeySet[[]byte]
-	storeService store.KVStoreService
+	checksums collections.KeySet[[]byte]
 
 	queryPlugins QueryPlugins
 
@@ -87,10 +86,12 @@ func (k Keeper) newQueryHandler(ctx sdk.Context, callerID string) *queryHandler
 // - Size bounds are checked. Contract length must not be 0 or exceed a specific size (maxWasmSize).
 // - The contract must not have already been stored in store.
 func (k Keeper) storeWasmCode(ctx context.Context, code []byte, storeFn func(code wasmvm.WasmCode, gasLimit uint64) (wasmvm.Checksum, uint64, error)) ([]byte, error) {
-	sdkCtx := sdk.UnwrapSDKContext(ctx) // TODO: https://github.com/cosmos/ibc-go/issues/5917
 	var err error
 	if types.IsGzip(code) {
-		sdkCtx.GasMeter().ConsumeGas(types.VMGasRegister.UncompressCosts(len(code)), "Uncompress gzip bytecode")
+		if err := k.GasService.GasMeter(ctx).Consume(types.VMGasRegister.UncompressCosts(len(code)), "Uncompress gzip bytecode"); err != nil {
+			return nil, errorsmod.Wrap(err, "failed to consume gas for decompression")
+		}
+
 		code, err = types.Uncompress(code, types.MaxWasmSize)
 		if err != nil {
 			return nil, errorsmod.Wrap(err, "failed to store contract")
@@ -113,6 +114,7 @@ func (k Keeper) storeWasmCode(ctx context.Context, code []byte, storeFn func(cod
 	}
 
 	// create the code in the vm
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
 	gasLeft := types.VMGasRegister.RuntimeGasForContract(sdkCtx)
 	vmChecksum, gasUsed, err := storeFn(code, gasLeft)
 	types.VMGasRegister.ConsumeRuntimeGas(sdkCtx, gasUsed)
@@ -186,7 +188,7 @@ func (k Keeper) migrateContractCode(ctx sdk.Context, clientID string, newChecksu
 }
 
 // GetWasmClientState returns the 08-wasm client state for the given client identifier.
-func (k Keeper) GetWasmClientState(ctx sdk.Context, clientID string) (*types.ClientState, error) {
+func (k Keeper) GetWasmClientState(ctx context.Context, clientID string) (*types.ClientState, error) {
 	clientState, found := k.clientKeeper.GetClientState(ctx, clientID)
 	if !found {
 		return nil, errorsmod.Wrapf(clienttypes.ErrClientTypeNotFound, "clientID %s", clientID)
@@ -233,7 +235,7 @@ func (k Keeper) HasChecksum(ctx context.Context, checksum types.Checksum) bool {
 }
 
 // InitializePinnedCodes updates wasmvm to pin to cache all contracts marked as pinned
-func (k Keeper) InitializePinnedCodes(ctx sdk.Context) error {
+func (k Keeper) InitializePinnedCodes(ctx context.Context) error {
 	checksums, err := k.GetAllChecksums(ctx)
 	if err != nil {
 		return err

diff --git a/modules/light-clients/08-wasm/keeper/keeper_test.go b/modules/light-clients/08-wasm/keeper/keeper_test.go
@@ -153,7 +153,7 @@ func (suite *KeeperTestSuite) TestNewKeeper() {
 			func() {
 				keeper.NewKeeperWithVM(
 					GetSimApp(suite.chainA).AppCodec(),
-					runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)),
+					runtime.NewEnvironment(runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)), log.NewNopLogger()),
 					GetSimApp(suite.chainA).IBCKeeper.ClientKeeper,
 					GetSimApp(suite.chainA).WasmClientKeeper.GetAuthority(),
 					GetSimApp(suite.chainA).WasmClientKeeper.GetVM(),
@@ -168,7 +168,7 @@ func (suite *KeeperTestSuite) TestNewKeeper() {
 			func() {
 				keeper.NewKeeperWithVM(
 					GetSimApp(suite.chainA).AppCodec(),
-					runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)),
+					runtime.NewEnvironment(runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)), log.NewNopLogger()),
 					GetSimApp(suite.chainA).IBCKeeper.ClientKeeper,
 					"", // authority
 					GetSimApp(suite.chainA).WasmClientKeeper.GetVM(),
@@ -183,7 +183,7 @@ func (suite *KeeperTestSuite) TestNewKeeper() {
 			func() {
 				keeper.NewKeeperWithVM(
 					GetSimApp(suite.chainA).AppCodec(),
-					runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)),
+					runtime.NewEnvironment(runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)), log.NewNopLogger()),
 					nil, // client keeper,
 					GetSimApp(suite.chainA).WasmClientKeeper.GetAuthority(),
 					GetSimApp(suite.chainA).WasmClientKeeper.GetVM(),
@@ -198,7 +198,7 @@ func (suite *KeeperTestSuite) TestNewKeeper() {
 			func() {
 				keeper.NewKeeperWithVM(
 					GetSimApp(suite.chainA).AppCodec(),
-					runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)),
+					runtime.NewEnvironment(runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)), log.NewNopLogger()),
 					GetSimApp(suite.chainA).IBCKeeper.ClientKeeper,
 					GetSimApp(suite.chainA).WasmClientKeeper.GetAuthority(),
 					nil,
@@ -213,7 +213,7 @@ func (suite *KeeperTestSuite) TestNewKeeper() {
 			func() {
 				keeper.NewKeeperWithVM(
 					GetSimApp(suite.chainA).AppCodec(),
-					nil,
+					runtime.NewEnvironment(nil, log.NewNopLogger()),
 					GetSimApp(suite.chainA).IBCKeeper.ClientKeeper,
 					GetSimApp(suite.chainA).WasmClientKeeper.GetAuthority(),
 					GetSimApp(suite.chainA).WasmClientKeeper.GetVM(),
@@ -228,7 +228,7 @@ func (suite *KeeperTestSuite) TestNewKeeper() {
 			func() {
 				keeper.NewKeeperWithVM(
 					GetSimApp(suite.chainA).AppCodec(),
-					runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)),
+					runtime.NewEnvironment(runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)), log.NewNopLogger()),
 					GetSimApp(suite.chainA).IBCKeeper.ClientKeeper,
 					GetSimApp(suite.chainA).WasmClientKeeper.GetAuthority(),
 					GetSimApp(suite.chainA).WasmClientKeeper.GetVM(),

diff --git a/modules/light-clients/08-wasm/keeper/keeper_vm.go b/modules/light-clients/08-wasm/keeper/keeper_vm.go
@@ -10,7 +10,7 @@ import (
 	wasmvm "github.com/CosmWasm/wasmvm/v2"
 
 	"cosmossdk.io/collections"
-	"cosmossdk.io/core/store"
+	"cosmossdk.io/core/appmodule"
 
 	"github.com/cosmos/cosmos-sdk/codec"
 
@@ -22,7 +22,7 @@ import (
 // and the same Wasm VM instance should be shared with it.
 func NewKeeperWithVM(
 	cdc codec.BinaryCodec,
-	storeService store.KVStoreService,
+	env appmodule.Environment,
 	clientKeeper types.ClientKeeper,
 	authority string,
 	vm types.WasmEngine,
@@ -41,21 +41,21 @@ func NewKeeperWithVM(
 		panic(errors.New("wasm VM must not be nil"))
 	}
 
-	if storeService == nil {
+	if env.KVStoreService == nil {
 		panic(errors.New("store service must not be nil"))
 	}
 
 	if strings.TrimSpace(authority) == "" {
 		panic(errors.New("authority must be non-empty"))
 	}
 
-	sb := collections.NewSchemaBuilder(storeService)
+	sb := collections.NewSchemaBuilder(env.KVStoreService)
 
 	keeper := &Keeper{
+		Environment:  env,
 		cdc:          cdc,
 		vm:           vm,
 		checksums:    collections.NewKeySet(sb, types.ChecksumsKey, "checksums", collections.BytesKey),
-		storeService: storeService,
 		clientKeeper: clientKeeper,
 		authority:    authority,
 	}
@@ -81,7 +81,7 @@ func NewKeeperWithVM(
 // and a Wasm VM needs to be instantiated using the provided parameters.
 func NewKeeperWithConfig(
 	cdc codec.BinaryCodec,
-	storeService store.KVStoreService,
+	env appmodule.Environment,
 	clientKeeper types.ClientKeeper,
 	authority string,
 	wasmConfig types.WasmConfig,
@@ -93,5 +93,5 @@ func NewKeeperWithConfig(
 		panic(fmt.Errorf("failed to instantiate new Wasm VM instance: %v", err))
 	}
 
-	return NewKeeperWithVM(cdc, storeService, clientKeeper, authority, vm, queryRouter, opts...)
+	return NewKeeperWithVM(cdc, env, clientKeeper, authority, vm, queryRouter, opts...)
 }
diff --git a/modules/light-clients/08-wasm/keeper/migrations.go b/modules/light-clients/08-wasm/keeper/migrations.go
@@ -47,7 +47,7 @@ func (m Migrator) MigrateChecksums(ctx sdk.Context) error {
 
 // getStoredChecksums returns the checksums stored under the KeyChecksums key.
 func (m Migrator) getStoredChecksums(ctx sdk.Context) ([][]byte, error) {
-	store := m.keeper.storeService.OpenKVStore(ctx)
+	store := m.keeper.KVStoreService.OpenKVStore(ctx)
 
 	bz, err := store.Get([]byte(types.KeyChecksums))
 	if err != nil {
@@ -65,7 +65,7 @@ func (m Migrator) getStoredChecksums(ctx sdk.Context) ([][]byte, error) {
 
 // deleteChecksums deletes the checksums stored under the KeyChecksums key.
 func (m Migrator) deleteChecksums(ctx sdk.Context) error {
-	store := m.keeper.storeService.OpenKVStore(ctx)
+	store := m.keeper.KVStoreService.OpenKVStore(ctx)
 	err := store.Delete([]byte(types.KeyChecksums))
 
 	return err

diff --git a/modules/light-clients/08-wasm/keeper/options_test.go b/modules/light-clients/08-wasm/keeper/options_test.go
@@ -4,6 +4,8 @@ import (
 	"encoding/json"
 	"errors"
 
+	"cosmossdk.io/log"
+
 	wasmvmtypes "github.com/CosmWasm/wasmvm/v2/types"
 
 	"github.com/cosmos/cosmos-sdk/runtime"
@@ -37,7 +39,7 @@ func (suite *KeeperTestSuite) TestNewKeeperWithOptions() {
 			func() {
 				k = keeper.NewKeeperWithVM(
 					GetSimApp(suite.chainA).AppCodec(),
-					runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)),
+					runtime.NewEnvironment(runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)), log.NewNopLogger()),
 					GetSimApp(suite.chainA).IBCKeeper.ClientKeeper,
 					GetSimApp(suite.chainA).WasmClientKeeper.GetAuthority(),
 					GetSimApp(suite.chainA).WasmClientKeeper.GetVM(),
@@ -62,7 +64,7 @@ func (suite *KeeperTestSuite) TestNewKeeperWithOptions() {
 				})
 				k = keeper.NewKeeperWithVM(
 					GetSimApp(suite.chainA).AppCodec(),
-					runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)),
+					runtime.NewEnvironment(runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)), log.NewNopLogger()),
 					GetSimApp(suite.chainA).IBCKeeper.ClientKeeper,
 					GetSimApp(suite.chainA).WasmClientKeeper.GetAuthority(),
 					GetSimApp(suite.chainA).WasmClientKeeper.GetVM(),
@@ -88,7 +90,7 @@ func (suite *KeeperTestSuite) TestNewKeeperWithOptions() {
 				})
 				k = keeper.NewKeeperWithVM(
 					GetSimApp(suite.chainA).AppCodec(),
-					runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)),
+					runtime.NewEnvironment(runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)), log.NewNopLogger()),
 					GetSimApp(suite.chainA).IBCKeeper.ClientKeeper,
 					GetSimApp(suite.chainA).WasmClientKeeper.GetAuthority(),
 					GetSimApp(suite.chainA).WasmClientKeeper.GetVM(),
@@ -115,7 +117,7 @@ func (suite *KeeperTestSuite) TestNewKeeperWithOptions() {
 				})
 				k = keeper.NewKeeperWithVM(
 					GetSimApp(suite.chainA).AppCodec(),
-					runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)),
+					runtime.NewEnvironment(runtime.NewKVStoreService(GetSimApp(suite.chainA).GetKey(types.StoreKey)), log.NewNopLogger()),
 					GetSimApp(suite.chainA).IBCKeeper.ClientKeeper,
 					GetSimApp(suite.chainA).WasmClientKeeper.GetAuthority(),
 					GetSimApp(suite.chainA).WasmClientKeeper.GetVM(),

diff --git a/modules/light-clients/08-wasm/testing/simapp/app.go b/modules/light-clients/08-wasm/testing/simapp/app.go
@@ -547,12 +547,16 @@ func NewSimApp(
 	if mockVM != nil {
 		// NOTE: mockVM is used for testing purposes only!
 		app.WasmClientKeeper = wasmkeeper.NewKeeperWithVM(
-			appCodec, runtime.NewKVStoreService(keys[wasmtypes.StoreKey]), app.IBCKeeper.ClientKeeper,
+			appCodec,
+			runtime.NewEnvironment(runtime.NewKVStoreService(keys[wasmtypes.StoreKey]), logger.With(log.ModuleKey, "x/ibc-wasm")),
+			app.IBCKeeper.ClientKeeper,
 			authtypes.NewModuleAddress(govtypes.ModuleName).String(), mockVM, app.GRPCQueryRouter(),
 		)
 	} else {
 		app.WasmClientKeeper = wasmkeeper.NewKeeperWithConfig(
-			appCodec, runtime.NewKVStoreService(keys[wasmtypes.StoreKey]), app.IBCKeeper.ClientKeeper,
+			appCodec,
+			runtime.NewEnvironment(runtime.NewKVStoreService(keys[wasmtypes.StoreKey]), logger.With(log.ModuleKey, "x/ibc-wasm")),
+			app.IBCKeeper.ClientKeeper,
 			authtypes.NewModuleAddress(govtypes.ModuleName).String(), wasmConfig, app.GRPCQueryRouter(),
 		)
 	}