From 759b0d250b0347e7e9c15f4127d641c58ec410d5 Mon Sep 17 00:00:00 2001 From: Lior Bondarevski Date: Mon, 3 Apr 2023 20:05:32 +0300 Subject: [PATCH] Create a presistent marker --- app/ante.go | 8 ++++--- app/app.go | 7 +++--- .../test-compute-contract/src/contract.rs | 22 +++++++++++++++++ .../test-compute-contract/src/msg.rs | 1 + x/compute/alias.go | 1 + x/compute/internal/keeper/ante.go | 24 +++++++++++++++++++ x/compute/internal/keeper/keeper.go | 4 ++++ x/compute/internal/keeper/msg_dispatcher.go | 13 +++++++++- .../keeper/secret_contracts_query_test.go | 19 +++++++++++++++ x/compute/internal/keeper/test_common.go | 6 +++++ x/compute/internal/types/keys.go | 1 + 11 files changed, 99 insertions(+), 7 deletions(-) diff --git a/app/ante.go b/app/ante.go index 416439d37..a4e9c27c1 100644 --- a/app/ante.go +++ b/app/ante.go @@ -14,9 +14,10 @@ import ( type HandlerOptions struct { ante.HandlerOptions - IBCKeeper *keeper.Keeper - WasmConfig *compute.WasmConfig - TXCounterStoreKey sdk.StoreKey + IBCKeeper *keeper.Keeper + WasmConfig *compute.WasmConfig + TXCounterStoreKey sdk.StoreKey + LastMsgMarkerStoreKey sdk.StoreKey } func NewAnteHandler(options HandlerOptions) (sdk.AnteHandler, error) { @@ -39,6 +40,7 @@ func NewAnteHandler(options HandlerOptions) (sdk.AnteHandler, error) { anteDecorators := []sdk.AnteDecorator{ compute.NewCountTXDecorator(options.TXCounterStoreKey), + compute.NewLastMsgDecorator(options.LastMsgMarkerStoreKey), ante.NewSetUpContextDecorator(), // outermost AnteDecorator. SetUpContext must be called first ante.NewRejectExtensionOptionsDecorator(), ante.NewMempoolFeeDecorator(), diff --git a/app/app.go b/app/app.go index 2c4253835..8a587e9fb 100644 --- a/app/app.go +++ b/app/app.go @@ -375,9 +375,10 @@ func NewSecretNetworkApp( SignModeHandler: encodingConfig.TxConfig.SignModeHandler(), SigGasConsumer: ante.DefaultSigVerificationGasConsumer, }, - IBCKeeper: app.AppKeepers.IbcKeeper, - WasmConfig: computeConfig, - TXCounterStoreKey: app.AppKeepers.GetKey(compute.StoreKey), + IBCKeeper: app.AppKeepers.IbcKeeper, + WasmConfig: computeConfig, + TXCounterStoreKey: app.AppKeepers.GetKey(compute.StoreKey), + LastMsgMarkerStoreKey: app.AppKeepers.GetKey(compute.StoreKey), }) if err != nil { panic(fmt.Errorf("failed to create AnteHandler: %s", err)) diff --git a/cosmwasm/contracts/v1/compute-tests/test-compute-contract/src/contract.rs b/cosmwasm/contracts/v1/compute-tests/test-compute-contract/src/contract.rs index 5d081339e..e9ae2fd83 100644 --- a/cosmwasm/contracts/v1/compute-tests/test-compute-contract/src/contract.rs +++ b/cosmwasm/contracts/v1/compute-tests/test-compute-contract/src/contract.rs @@ -342,6 +342,28 @@ pub fn wasm_msg(ty: String) -> StdResult { #[entry_point] pub fn execute(deps: DepsMut, env: Env, info: MessageInfo, msg: ExecuteMsg) -> StdResult { match msg { + ExecuteMsg::Special {} => { + let s = SubMsg { + id: 0, + msg: CosmosMsg::Wasm(WasmMsg::Execute { + code_hash: env.contract.code_hash, + contract_addr: env.contract.address.into_string(), + msg: Binary::from("{\"increment\":{\"addition\":1}}".as_bytes().to_vec()), + funds: vec![], + }) + .into(), + reply_on: ReplyOn::Never, + gas_limit: None, + }; + + Ok(Response::new().add_submessages(vec![ + s.clone(), + s.clone(), + s.clone(), + s.clone(), + s.clone(), + ])) + } ExecuteMsg::WasmMsg { ty } => wasm_msg(ty), ExecuteMsg::Increment { addition } => increment(env, deps, addition), ExecuteMsg::SendFundsWithErrorWithReply {} => Ok(Response::new() diff --git a/cosmwasm/contracts/v1/compute-tests/test-compute-contract/src/msg.rs b/cosmwasm/contracts/v1/compute-tests/test-compute-contract/src/msg.rs index eb1c4b601..39e3b99bf 100644 --- a/cosmwasm/contracts/v1/compute-tests/test-compute-contract/src/msg.rs +++ b/cosmwasm/contracts/v1/compute-tests/test-compute-contract/src/msg.rs @@ -114,6 +114,7 @@ pub enum InstantiateMsg { #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] #[serde(rename_all = "snake_case")] pub enum ExecuteMsg { + Special {}, WasmMsg { ty: String, }, diff --git a/x/compute/alias.go b/x/compute/alias.go index e949f45e6..3342747c4 100644 --- a/x/compute/alias.go +++ b/x/compute/alias.go @@ -83,6 +83,7 @@ var ( NewWasmSnapshotter = keeper.NewWasmSnapshotter ContractFromPortID = keeper.ContractFromPortID NewCountTXDecorator = keeper.NewCountTXDecorator + NewLastMsgDecorator = keeper.NewLastMsgDecorator // variable aliases ModuleCdc = types.ModuleCdc diff --git a/x/compute/internal/keeper/ante.go b/x/compute/internal/keeper/ante.go index 40e684ab0..742185277 100644 --- a/x/compute/internal/keeper/ante.go +++ b/x/compute/internal/keeper/ante.go @@ -2,6 +2,7 @@ package keeper import ( "encoding/binary" + "fmt" sdk "github.com/cosmos/cosmos-sdk/types" "github.com/scrtlabs/SecretNetwork/x/compute/internal/types" @@ -12,11 +13,20 @@ type CountTXDecorator struct { storeKey sdk.StoreKey } +type LastMsgDecorator struct { + storeKey sdk.StoreKey +} + // NewCountTXDecorator constructor func NewCountTXDecorator(storeKey sdk.StoreKey) *CountTXDecorator { return &CountTXDecorator{storeKey: storeKey} } +// LastMsgDecorator constructor +func NewLastMsgDecorator(storeKey sdk.StoreKey) *LastMsgDecorator { + return &LastMsgDecorator{storeKey: storeKey} +} + // AnteHandle handler stores a tx counter with current height encoded in the store to let the app handle // global rollback behavior instead of keeping state in the handler itself. // The ante handler passes the counter value via sdk.Context upstream. See `types.TXCounter(ctx)` to read the value. @@ -43,6 +53,20 @@ func (a CountTXDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate bool, return next(types.WithTXCounter(ctx, txCounter), tx, simulate) } +// AnteHandle handler resets the market of the last message in the block +func (a LastMsgDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate bool, next sdk.AnteHandler) (sdk.Context, error) { + if simulate { + return next(ctx, tx, simulate) + } + + store := ctx.KVStore(a.storeKey) + // Reset last msg marker + fmt.Println("LIORRR Setting marker to off") + store.Set(types.LastMsgPrefix, nil) + + return next(ctx, tx, simulate) +} + func encodeHeightCounter(height int64, counter uint32) []byte { b := make([]byte, 4) binary.BigEndian.PutUint32(b, counter) diff --git a/x/compute/internal/keeper/keeper.go b/x/compute/internal/keeper/keeper.go index e5711f4c5..12afb64cb 100644 --- a/x/compute/internal/keeper/keeper.go +++ b/x/compute/internal/keeper/keeper.go @@ -1075,3 +1075,7 @@ func (k Keeper) reply(ctx sdk.Context, contractAddress sdk.AccAddress, reply v1w return nil, sdkerrors.Wrap(types.ErrReplyFailed, fmt.Sprintf("cannot detect response type: %+v", res)) } } + +func (k Keeper) GetStoreKey() sdk.StoreKey { + return k.storeKey +} diff --git a/x/compute/internal/keeper/msg_dispatcher.go b/x/compute/internal/keeper/msg_dispatcher.go index 9af8845e1..c9a65e38a 100644 --- a/x/compute/internal/keeper/msg_dispatcher.go +++ b/x/compute/internal/keeper/msg_dispatcher.go @@ -27,6 +27,7 @@ type Messenger interface { // Replyer is a subset of keeper that can handle replies to submessages type Replyer interface { reply(ctx sdk.Context, contractAddress sdk.AccAddress, reply v1wasmTypes.Reply, ogTx []byte, ogSigInfo wasmTypes.VerificationInfo) ([]byte, error) + GetStoreKey() sdk.StoreKey } // MessageDispatcher coordinates message sending and submessage reply/ state commits @@ -185,7 +186,17 @@ func redactError(err error) (bool, error) { // that dispatched them, both on success as well as failure func (d MessageDispatcher) DispatchSubmessages(ctx sdk.Context, contractAddr sdk.AccAddress, ibcPort string, msgs []v1wasmTypes.SubMsg, ogTx []byte, ogSigInfo wasmTypes.VerificationInfo, ogCosmosMessageVersion wasmTypes.CosmosMsgVersion) ([]byte, error) { var rsp []byte - for _, msg := range msgs { + for idx, msg := range msgs { + store := ctx.KVStore(d.keeper.GetStoreKey()) + if idx == 2 { + store.Set(types.LastMsgPrefix, []byte{1}) + fmt.Printf("LIORRR Setting marker to on: %+v\n", store.Get(types.LastMsgPrefix)) + } + + if store.Get(types.LastMsgPrefix) != nil { + break + } + // Check replyOn validity switch msg.ReplyOn { case v1wasmTypes.ReplySuccess, v1wasmTypes.ReplyError, v1wasmTypes.ReplyAlways, v1wasmTypes.ReplyNever: diff --git a/x/compute/internal/keeper/secret_contracts_query_test.go b/x/compute/internal/keeper/secret_contracts_query_test.go index c7b45c6b3..6b609dd61 100644 --- a/x/compute/internal/keeper/secret_contracts_query_test.go +++ b/x/compute/internal/keeper/secret_contracts_query_test.go @@ -598,6 +598,25 @@ func TestV1EndpointsSanity(t *testing.T) { require.Equal(t, uint32(23), resp.Get.Count) } +func TestSpecial(t *testing.T) { + ctx, keeper, codeID, _, walletA, privKeyA, _, _ := setupTest(t, TestContractPaths[v1Contract], sdk.NewCoins()) + + _, _, contractAddress, _, _ := initHelper(t, keeper, ctx, codeID, walletA, privKeyA, `{"counter":{"counter":10, "expires":100}}`, true, true, defaultGasForTests) + + _, _, _, _, _, err := execHelper(t, keeper, ctx, contractAddress, walletA, privKeyA, `{"special":{}}`, true, true, math.MaxUint64, 0) + + require.Empty(t, err) + + queryRes, qErr := queryHelper(t, keeper, ctx, contractAddress, `{"get":{}}`, true, true, math.MaxUint64) + require.Empty(t, qErr) + + // assert result is 32 byte sha256 hash (if hashed), or contractAddr if not + var resp v1QueryResponse + e := json.Unmarshal([]byte(queryRes), &resp) + require.NoError(t, e) + fmt.Printf("LIORRR count: %+v\n", resp.Get.Count) +} + func TestV1QueryWorksWithEnv(t *testing.T) { ctx, keeper, codeID, _, walletA, privKeyA, _, _ := setupTest(t, TestContractPaths[v1Contract], sdk.NewCoins()) diff --git a/x/compute/internal/keeper/test_common.go b/x/compute/internal/keeper/test_common.go index a8e0b2f6a..afbd71c14 100644 --- a/x/compute/internal/keeper/test_common.go +++ b/x/compute/internal/keeper/test_common.go @@ -535,6 +535,12 @@ func TestHandler(k Keeper) sdk.Handler { return func(ctx sdk.Context, msg sdk.Msg) (*sdk.Result, error) { ctx = ctx.WithEventManager(sdk.NewEventManager()) + store := ctx.KVStore(k.GetStoreKey()) + if bz := store.Get(wasmtypes.LastMsgPrefix); bz != nil { + fmt.Printf("Message of type: %+v was filtered out\n", msg) + return &sdk.Result{}, nil + } + switch msg := msg.(type) { case *wasmtypes.MsgInstantiateContract: return handleInstantiate(ctx, k, msg) diff --git a/x/compute/internal/types/keys.go b/x/compute/internal/types/keys.go index 511634ef3..125a4a0fb 100644 --- a/x/compute/internal/types/keys.go +++ b/x/compute/internal/types/keys.go @@ -31,6 +31,7 @@ var ( ContractEnclaveIdPrefix = []byte{0x06} ContractLabelPrefix = []byte{0x07} TXCounterPrefix = []byte{0x08} + LastMsgPrefix = []byte{0x09} // RandomPrefix = []byte{0xFF} KeyLastCodeID = append(SequenceKeyPrefix, []byte("lastCodeId")...) KeyLastInstanceID = append(SequenceKeyPrefix, []byte("lastContractId")...)