feat(nexus)!: allow wasm contracts to query for tx hash and index

app/keepers.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/CosmWasm/wasmd/x/wasm"
 	wasmkeeper "github.com/CosmWasm/wasmd/x/wasm/keeper"
+	wasmtypes "github.com/CosmWasm/wasmd/x/wasm/types"
 	bam "github.com/cosmos/cosmos-sdk/baseapp"
 	"github.com/cosmos/cosmos-sdk/codec"
 	"github.com/cosmos/cosmos-sdk/server/types"
@@ -163,20 +164,28 @@ func initStakingKeeper(appCodec codec.Codec, keys map[string]*sdk.KVStoreKey, ke
 
 func initWasmKeeper(encodingConfig axelarParams.EncodingConfig, keys map[string]*sdk.KVStoreKey, keepers *KeeperCache, bApp *bam.BaseApp, appOpts types.AppOptions, wasmOpts []wasm.Option, wasmDir string) *wasm.Keeper {
 	wasmConfig := mustReadWasmConfig(appOpts)
+	nexusK := GetKeeper[nexusKeeper.Keeper](keepers)
 
 	// The last arguments can contain custom message handlers, and custom query handlers,
 	// if we want to allow any custom callbacks
-	wasmOpts = append(wasmOpts, wasmkeeper.WithMessageHandlerDecorator(
-		func(old wasmkeeper.Messenger) wasmkeeper.Messenger {
-			encoders := wasm.DefaultEncoders(encodingConfig.Codec, GetKeeper[ibctransferkeeper.Keeper](keepers))
-			encoders.Custom = nexusKeeper.EncodeRoutingMessage
-
-			return WithAnteHandlers(
-				encoders,
-				initMessageAnteDecorators(encodingConfig, keepers),
-				// for security reasons we disallow some msg types that can be used for arbitrary calls
-				wasmkeeper.NewMessageHandlerChain(NewMsgTypeBlacklistMessenger(), old, nexusKeeper.NewMessenger(GetKeeper[nexusKeeper.Keeper](keepers))))
-		}))
+	wasmOpts = append(
+		wasmOpts,
+		wasmkeeper.WithMessageHandlerDecorator(
+			func(old wasmkeeper.Messenger) wasmkeeper.Messenger {
+				encoders := wasm.DefaultEncoders(encodingConfig.Codec, GetKeeper[ibctransferkeeper.Keeper](keepers))
+				encoders.Custom = nexusKeeper.EncodeRoutingMessage
+
+				return WithAnteHandlers(
+					encoders,
+					initMessageAnteDecorators(encodingConfig, keepers),
+					// for security reasons we disallow some msg types that can be used for arbitrary calls
+					wasmkeeper.NewMessageHandlerChain(NewMsgTypeBlacklistMessenger(), old, nexusKeeper.NewMessenger(nexusK)))
+			}),
+		wasmkeeper.WithWasmEngineDecorator(func(old wasmtypes.WasmerEngine) wasmtypes.WasmerEngine {
+			return nexusKeeper.NewWasmerEngine(old, nexusK)
+		}),
+		wasmkeeper.WithQueryPlugins(NewQueryPlugins(nexusK)),
+	)
 
 	scopedWasmK := GetKeeper[capabilitykeeper.Keeper](keepers).ScopeToModule(wasm.ModuleName)
 	ibcKeeper := GetKeeper[ibckeeper.Keeper](keepers)

app/wasm.go

@@ -15,6 +15,9 @@ import (
 	"golang.org/x/exp/maps"
 
 	"github.com/axelarnetwork/axelar-core/x/ante"
+	nexus "github.com/axelarnetwork/axelar-core/x/nexus/exported"
+	nexusKeeper "github.com/axelarnetwork/axelar-core/x/nexus/keeper"
+	nexustypes "github.com/axelarnetwork/axelar-core/x/nexus/types"
 	"github.com/axelarnetwork/utils/funcs"
 )
 
@@ -146,3 +149,28 @@ func (m WasmAppModuleBasicOverride) DefaultGenesis(cdc codec.JSONCodec) json.Raw
 		},
 	})
 }
+
+// QueryRequest is the custom queries wasm contracts can make for the core modules
+type QueryRequest struct {
+	Nexus *nexus.WasmQueryRequest
+}
+
+// NewQueryPlugins returns a new instance of the custom query plugins
+func NewQueryPlugins(txIDGenerator nexustypes.TxIDGenerator) *wasmkeeper.QueryPlugins {
+	nexusWasmQuerier := nexusKeeper.NewWasmQuerier(txIDGenerator)
+
+	return &wasmkeeper.QueryPlugins{
+		Custom: func(ctx sdk.Context, request json.RawMessage) ([]byte, error) {
+			req := QueryRequest{}
+			if err := json.Unmarshal(request, &req); err != nil {
+				return nil, wasmvmtypes.InvalidRequest{Err: "invalid Custom query request", Request: request}
+			}
+
+			if req.Nexus != nil {
+				return nexusWasmQuerier.Query(ctx, *req.Nexus)
+			}
+
+			return nil, wasmvmtypes.UnsupportedRequest{Kind: "unknown Custom query request"}
+		},
+	}
+}

app/wasm_test.go

@@ -2,6 +2,7 @@ package app_test
 
 import (
 	"encoding/json"
+	"fmt"
 	"testing"
 
 	"github.com/CosmWasm/wasmd/x/wasm"
@@ -20,6 +21,8 @@ import (
 	"github.com/axelarnetwork/axelar-core/cmd/axelard/cmd"
 	"github.com/axelarnetwork/axelar-core/testutils/fake"
 	"github.com/axelarnetwork/axelar-core/testutils/rand"
+	nexusmock "github.com/axelarnetwork/axelar-core/x/nexus/types/mock"
+	"github.com/axelarnetwork/utils/funcs"
 	. "github.com/axelarnetwork/utils/test"
 )
 
@@ -328,3 +331,62 @@ func TestMaxSizeOverrideForClient(t *testing.T) {
 
 	assert.NoError(t, msg.ValidateBasic())
 }
+
+func TestQueryPlugins(t *testing.T) {
+	var (
+		txIDGenerator *nexusmock.TxIDGeneratorMock
+		req           json.RawMessage
+		ctx           sdk.Context
+	)
+
+	Given("the tx id generator", func() {
+		ctx = sdk.NewContext(nil, tmproto.Header{}, false, log.TestingLogger())
+		txIDGenerator = &nexusmock.TxIDGeneratorMock{}
+	}).
+		Branch(
+			When("request is invalid", func() {
+				req = []byte("{\"invalid\"}")
+			}).
+				Then("it should return an error", func(t *testing.T) {
+					_, err := app.NewQueryPlugins(txIDGenerator).Custom(ctx, req)
+
+					assert.ErrorContains(t, err, "invalid Custom query request")
+				}),
+
+			When("request is unknown", func() {
+				req = []byte("{\"unknown\":{}}")
+			}).
+				Then("it should return an error", func(t *testing.T) {
+					_, err := app.NewQueryPlugins(txIDGenerator).Custom(ctx, req)
+
+					assert.ErrorContains(t, err, "unknown Custom query request")
+				}),
+
+			When("request is a nexus wasm query but unknown", func() {
+				req = []byte("{\"nexus\":{}}")
+			}).
+				Then("it should return an error", func(t *testing.T) {
+					_, err := app.NewQueryPlugins(txIDGenerator).Custom(ctx, req)
+
+					assert.ErrorContains(t, err, "unknown Nexus query request")
+				}),
+
+			When("request is a nexus wasm TxID query", func() {
+				req = []byte("{\"nexus\":{\"tx_id\":{}}}")
+			}).
+				Then("it should return an error", func(t *testing.T) {
+					txHash := [32]byte(rand.Bytes(32))
+					index := uint64(rand.PosI64())
+					txIDGenerator.CurrFunc = func(ctx sdk.Context) ([32]byte, uint64) {
+						return txHash, index
+					}
+
+					actual, err := app.NewQueryPlugins(txIDGenerator).Custom(ctx, req)
+
+					assert.NoError(t, err)
+					assert.Equal(t, fmt.Sprintf("{\"tx_hash\":%s,\"index\":%d}", funcs.Must(json.Marshal(txHash)), index), string(actual))
+				}),
+		).
+		Run(t)
+
+}

x/nexus/exported/types.go

@@ -414,3 +414,14 @@ func (bz *WasmBytes) UnmarshalJSON(data []byte) error {
 
 	return nil
 }
+
+// WasmQueryRequest is the request for wasm contracts to query
+type WasmQueryRequest struct {
+	TxID *struct{} `json:"tx_id,omitempty"`
+}
+
+// WasmQueryTxIDResponse is the response for the TxID query
+type WasmQueryTxIDResponse struct {
+	TxHash [32]byte `json:"tx_hash,omitempty"` // the hash of the current transaction
+	Index  uint64   `json:"index,omitempty"`   // the index of the current execution, which increments with each entry of any wasm execution
+}

x/nexus/keeper/general_message.go

@@ -1,7 +1,6 @@
 package keeper
 
 import (
-	"crypto/sha256"
 	"encoding/hex"
 	"fmt"
 
@@ -32,9 +31,7 @@ func getProcessingMessageKey(destinationChain exported.ChainName, id string) key
 // GenerateMessageID generates a unique general message ID, and returns the message ID, current transacation ID and a unique integer nonce
 // The message ID is just a concatenation of the transaction ID and the nonce
 func (k Keeper) GenerateMessageID(ctx sdk.Context) (string, []byte, uint64) {
-	counter := utils.NewCounter[uint64](messageNonceKey, k.getStore(ctx))
-	nonce := counter.Incr(ctx)
-	hash := sha256.Sum256(ctx.TxBytes())
+	hash, nonce := k.Next(ctx)
 
 	return fmt.Sprintf("0x%s-%d", hex.EncodeToString(hash[:]), nonce), hash[:], nonce
 }

x/nexus/keeper/tx_id_generator.go

@@ -0,0 +1,26 @@
+package keeper
+
+import (
+	"crypto/sha256"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/axelarnetwork/axelar-core/utils"
+	"github.com/axelarnetwork/axelar-core/x/nexus/types"
+)
+
+var _ types.TxIDGenerator = &Keeper{}
+
+func getTxHash(ctx sdk.Context) [32]byte {
+	return sha256.Sum256(ctx.TxBytes())
+}
+
+// Next returns the next transaction hash and index, and increments the index
+func (k Keeper) Next(ctx sdk.Context) ([32]byte, uint64) {
+	return getTxHash(ctx), utils.NewCounter[uint64](messageNonceKey, k.getStore(ctx)).Incr(ctx)
+}
+
+// Curr returns the current transaction hash and index
+func (k Keeper) Curr(ctx sdk.Context) ([32]byte, uint64) {
+	return getTxHash(ctx), utils.NewCounter[uint64](messageNonceKey, k.getStore(ctx)).Curr(ctx)
+}

x/nexus/keeper/wasm_querier.go

@@ -0,0 +1,36 @@
+package keeper
+
+import (
+	"encoding/json"
+
+	wasmvmtypes "github.com/CosmWasm/wasmvm/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/axelarnetwork/axelar-core/x/nexus/exported"
+	"github.com/axelarnetwork/axelar-core/x/nexus/types"
+	"github.com/axelarnetwork/utils/funcs"
+)
+
+// WasmQuerier is a querier for the wasm contracts
+type WasmQuerier struct {
+	txIDGenerator types.TxIDGenerator
+}
+
+// NewWasmQuerier creates a new WasmQuerier
+func NewWasmQuerier(txIDGenerator types.TxIDGenerator) *WasmQuerier {
+	return &WasmQuerier{txIDGenerator}
+}
+
+// Query handles the wasm queries for the nexus module
+func (q WasmQuerier) Query(ctx sdk.Context, req exported.WasmQueryRequest) ([]byte, error) {
+	if req.TxID != nil {
+		txHash, index := q.txIDGenerator.Curr(ctx)
+
+		return funcs.Must(json.Marshal(exported.WasmQueryTxIDResponse{
+			TxHash: txHash,
+			Index:  index,
+		})), nil
+	}
+
+	return nil, wasmvmtypes.UnsupportedRequest{Kind: "unknown Nexus query request"}
+}

x/nexus/keeper/wasmer_engine.go

@@ -0,0 +1,113 @@
+package keeper
+
+import (
+	wasmkeeper "github.com/CosmWasm/wasmd/x/wasm/keeper"
+	wasmtypes "github.com/CosmWasm/wasmd/x/wasm/types"
+	wasmvm "github.com/CosmWasm/wasmvm"
+	wasmvmtypes "github.com/CosmWasm/wasmvm/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/axelarnetwork/axelar-core/x/nexus/types"
+)
+
+// WasmerEngine is a wrapper around the WasmerEngine to add a transaction ID generator
+type WasmerEngine struct {
+	wasmtypes.WasmerEngine
+	txIDGenerator types.TxIDGenerator
+}
+
+// NewWasmerEngine wraps the given engine with a transaction ID generator
+func NewWasmerEngine(inner wasmtypes.WasmerEngine, txIDGenerator types.TxIDGenerator) wasmtypes.WasmerEngine {
+	return &WasmerEngine{WasmerEngine: inner, txIDGenerator: txIDGenerator}
+}
+
+func getCtx(querier wasmvm.Querier) sdk.Context {
+	return querier.(wasmkeeper.QueryHandler).Ctx
+}
+
+// Instantiate calls the inner engine and increments the transaction ID
+func (w *WasmerEngine) Instantiate(
+	checksum wasmvm.Checksum,
+	env wasmvmtypes.Env,
+	info wasmvmtypes.MessageInfo,
+	initMsg []byte,
+	store wasmvm.KVStore,
+	goapi wasmvm.GoAPI,
+	querier wasmvm.Querier,
+	gasMeter wasmvm.GasMeter,
+	gasLimit uint64,
+	deserCost wasmvmtypes.UFraction,
+) (*wasmvmtypes.Response, uint64, error) {
+	defer w.txIDGenerator.Next(getCtx(querier))
+
+	return w.WasmerEngine.Instantiate(checksum, env, info, initMsg, store, goapi, querier, gasMeter, gasLimit, deserCost)
+}
+
+// Execute calls the inner engine and increments the transaction ID
+func (w *WasmerEngine) Execute(
+	code wasmvm.Checksum,
+	env wasmvmtypes.Env,
+	info wasmvmtypes.MessageInfo,
+	executeMsg []byte,
+	store wasmvm.KVStore,
+	goapi wasmvm.GoAPI,
+	querier wasmvm.Querier,
+	gasMeter wasmvm.GasMeter,
+	gasLimit uint64,
+	deserCost wasmvmtypes.UFraction,
+) (*wasmvmtypes.Response, uint64, error) {
+	defer w.txIDGenerator.Next(getCtx(querier))
+
+	return w.WasmerEngine.Execute(code, env, info, executeMsg, store, goapi, querier, gasMeter, gasLimit, deserCost)
+}
+
+// Migrate calls the inner engine and increments the transaction ID
+func (w *WasmerEngine) Migrate(
+	checksum wasmvm.Checksum,
+	env wasmvmtypes.Env,
+	migrateMsg []byte,
+	store wasmvm.KVStore,
+	goapi wasmvm.GoAPI,
+	querier wasmvm.Querier,
+	gasMeter wasmvm.GasMeter,
+	gasLimit uint64,
+	deserCost wasmvmtypes.UFraction,
+) (*wasmvmtypes.Response, uint64, error) {
+	defer w.txIDGenerator.Next(getCtx(querier))
+
+	return w.WasmerEngine.Migrate(checksum, env, migrateMsg, store, goapi, querier, gasMeter, gasLimit, deserCost)
+}
+
+// Sudo calls the inner engine and increments the transaction ID
+func (w *WasmerEngine) Sudo(
+	checksum wasmvm.Checksum,
+	env wasmvmtypes.Env,
+	sudoMsg []byte,
+	store wasmvm.KVStore,
+	goapi wasmvm.GoAPI,
+	querier wasmvm.Querier,
+	gasMeter wasmvm.GasMeter,
+	gasLimit uint64,
+	deserCost wasmvmtypes.UFraction,
+) (*wasmvmtypes.Response, uint64, error) {
+	defer w.txIDGenerator.Next(getCtx(querier))
+
+	return w.WasmerEngine.Sudo(checksum, env, sudoMsg, store, goapi, querier, gasMeter, gasLimit, deserCost)
+}
+
+// Reply calls the inner engine and increments the transaction ID
+func (w *WasmerEngine) Reply(
+	checksum wasmvm.Checksum,
+	env wasmvmtypes.Env,
+	reply wasmvmtypes.Reply,
+	store wasmvm.KVStore,
+	goapi wasmvm.GoAPI,
+	querier wasmvm.Querier,
+	gasMeter wasmvm.GasMeter,
+	gasLimit uint64,
+	deserCost wasmvmtypes.UFraction,
+) (*wasmvmtypes.Response, uint64, error) {
+	defer w.txIDGenerator.Next(getCtx(querier))
+
+	return w.WasmerEngine.Reply(checksum, env, reply, store, goapi, querier, gasMeter, gasLimit, deserCost)
+}