Encrypt plaintext replies only if the original message was encrypted

cosmwasm/enclaves/shared/contract-engine/src/io.rs

@@ -412,7 +412,9 @@ pub fn encrypt_output(
     let encryption_key = calc_encryption_key(&secret_msg.nonce, &secret_msg.user_public_key);
     trace!(
         "Output before encryption: {:?} {:?} {:?}",
-        String::from_utf8_lossy(&output), secret_msg.nonce, secret_msg.user_public_key
+        String::from_utf8_lossy(&output),
+        secret_msg.nonce,
+        secret_msg.user_public_key
     );
 
     let mut output: RawWasmOutput = serde_json::from_slice(&output).map_err(|err| {
@@ -449,6 +451,7 @@ pub fn encrypt_output(
                     let reply = Reply {
                         id: msg_id.unwrap(),
                         result: SubMsgResult::Err(encrypted_err),
+                        was_msg_encrypted: true,
                     };
                     let reply_as_vec = serde_json::to_vec(&reply).map_err(|err| {
                         warn!(
@@ -537,6 +540,7 @@ pub fn encrypt_output(
                             events,
                             data: ok.data.clone(),
                         }),
+                        was_msg_encrypted: true,
                     };
 
                     let reply_as_vec = serde_json::to_vec(&reply).map_err(|err| {
@@ -580,6 +584,8 @@ pub fn encrypt_output(
                     // We don't encrypt it here to remain with the same type (u64)
                     sub_msg.id = 0;
                 }
+
+                sub_msg.was_msg_encrypted = true;
             }
 
             // v1: The attributes that will be emitted as part of a "wasm" event.
@@ -649,6 +655,7 @@ pub fn encrypt_output(
                             events,
                             data: ok.data.clone(),
                         }),
+                        was_msg_encrypted: true,
                     };
 
                     let reply_as_vec = serde_json::to_vec(&reply).map_err(|err| {
@@ -689,6 +696,8 @@ pub fn encrypt_output(
                     // We don't encrypt it here to remain with the same type (u64)
                     sub_msg.id = 0;
                 }
+
+                sub_msg.was_msg_encrypted = true;
             }
 
             // v1: The attributes that will be emitted as part of a "wasm" event.

cosmwasm/enclaves/shared/contract-engine/src/message.rs

@@ -199,7 +199,7 @@ pub fn parse_message(
                 return Ok(ParsedMessage {
                     should_validate_sig_info: false,
                     was_msg_encrypted: false,
-                    should_encrypt_output: true, // When replies are plaintext it doesn't mean we can't encrypt them
+                    should_encrypt_output: reply.was_msg_encrypted,
                     secret_msg: reply_secret_msg,
                     decrypted_msg: serialized_reply,
                     contract_hash_for_validation: None,

cosmwasm/enclaves/shared/cosmwasm-types/v1.0/src/results/submessages.rs

@@ -21,6 +21,10 @@ pub enum ReplyOn {
     Never,
 }
 
+fn bool_false() -> bool {
+    false
+}
+
 /// A submessage that will guarantee a `reply` call on success or error, depending on
 /// the `reply_on` setting. If you do not need to process the result, use regular messages instead.
 ///
@@ -38,6 +42,11 @@ where
     pub msg: CosmosMsg<T>,
     pub gas_limit: Option<u64>,
     pub reply_on: ReplyOn,
+    // An indication that will be passed to the reply that will indicate wether the original message,
+    // which is the one who create the submessages, was encrypted or not.
+    // Plaintext replies will be encrypted only if the original message was.
+    #[serde(default = "bool_false")]
+    pub was_msg_encrypted: bool,
 }
 
 /// The information we get back from a successful sub message execution,
@@ -65,6 +74,7 @@ pub struct Reply {
     /// Use this to identify which submessage triggered the `reply`.
     pub id: Binary,
     pub result: SubMsgResult,
+    pub was_msg_encrypted: bool,
 }
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
 pub struct DecryptedReply {

go-cosmwasm/types/v1/subcall.go

@@ -66,15 +66,17 @@ type SubMsgResult struct {
 // SubMsg wraps a CosmosMsg with some metadata for handling replies (ID) and optionally
 // limiting the gas usage (GasLimit)
 type SubMsg struct {
-	ID       uint64    `json:"id"`
-	Msg      CosmosMsg `json:"msg"`
-	GasLimit *uint64   `json:"gas_limit,omitempty"`
-	ReplyOn  replyOn   `json:"reply_on"`
+	ID              uint64    `json:"id"`
+	Msg             CosmosMsg `json:"msg"`
+	GasLimit        *uint64   `json:"gas_limit,omitempty"`
+	ReplyOn         replyOn   `json:"reply_on"`
+	WasMsgEncrypted bool      `json:"was_msg_encrypted"`
 }
 
 type Reply struct {
-	ID     []byte       `json:"id"`
-	Result SubMsgResult `json:"result"`
+	ID              []byte       `json:"id"`
+	Result          SubMsgResult `json:"result"`
+	WasMsgEncrypted bool         `json:"was_msg_encrypted"`
 }
 
 // SubcallResult is the raw response we return from the sdk -> reply after executing a SubMsg.

x/compute/internal/keeper/ibc_test.go

@@ -3,8 +3,10 @@ package keeper
 import (
 	"encoding/binary"
 	"encoding/hex"
+	"encoding/json"
 	"fmt"
 	"math"
+	"os"
 	"testing"
 
 	crypto "github.com/cosmos/cosmos-sdk/crypto/types"
@@ -719,6 +721,14 @@ func TestIBCPacketReceive(t *testing.T) {
 				hasAttributes: false,
 				hasEvents:     false,
 			},
+			{
+				description:   "SubmessageWithReplyThatCallsToSubmessage",
+				sequence:      6,
+				output:        "35",
+				isSuccess:     true,
+				hasAttributes: false,
+				hasEvents:     false,
+			},
 		} {
 			t.Run(fmt.Sprintf("%s-Encryption:%t", test.description, isEncrypted), func(t *testing.T) {
 				ibcPacket := createIBCPacket(createIBCEndpoint(PortIDForContract(contractAddress), "channel.1"),
@@ -781,6 +791,68 @@ func TestIBCPacketReceive(t *testing.T) {
 	}
 }
 
+type ContractInfo struct {
+	Address string `json:"address"`
+	Hash    string `json:"hash"`
+}
+
+func TestIBCPacketReceiveCallsV010Contract(t *testing.T) {
+	ctx, keeper, codeID, _, walletA, privKeyA, _, _ := setupTest(t, "./testdata/ibc/contract.wasm", sdk.NewCoins())
+
+	wasmCode, err := os.ReadFile("./testdata/test-contract/contract.wasm")
+	require.NoError(t, err)
+
+	v010CodeID, err := keeper.Create(ctx, walletA, wasmCode, "", "")
+	require.NoError(t, err)
+
+	codeInfo, err := keeper.GetCodeInfo(ctx, v010CodeID)
+	require.NoError(t, err)
+	v010CodeHash := hex.EncodeToString(codeInfo.CodeHash)
+
+	_, _, contractAddress, _, err := initHelper(t, keeper, ctx, codeID, walletA, privKeyA, `{"init":{}}`, true, true, defaultGasForTests)
+	require.Empty(t, err)
+	_, _, v010ContractAddress, _, err := initHelper(t, keeper, ctx, v010CodeID, walletA, privKeyA, `{"counter":{"counter":10}}`, true, false, defaultGasForTests)
+	require.Empty(t, err)
+	contractInfo := ContractInfo{
+		Address: v010ContractAddress.String(),
+		Hash:    v010CodeHash,
+	}
+
+	contractInfoBz, err := json.Marshal(contractInfo)
+	require.NoError(t, err)
+
+	ibcPacket := createIBCPacket(createIBCEndpoint(PortIDForContract(contractAddress), "channel.1"),
+		createIBCEndpoint(PortIDForContract(contractAddress), "channel.0"),
+		7,
+		createIBCTimeout(math.MaxUint64),
+		contractInfoBz,
+	)
+
+	expected_v010_result := uint32(15)
+
+	for _, isEncrypted := range []bool{true, true} {
+		t.Run(fmt.Sprintf("Encryption:%t", isEncrypted), func(t *testing.T) {
+			ctx, _, _, data, err := ibcPacketReceiveHelper(t, keeper, ctx, contractAddress, privKeyA, isEncrypted, defaultGasForTests, ibcPacket)
+			require.Empty(t, err)
+			require.Equal(t, "\"out\"", string(data))
+
+			queryRes, err := queryHelper(t, keeper, ctx, contractAddress, `{"q":{}}`, true, true, math.MaxUint64)
+
+			require.Empty(t, err)
+			require.Equal(t, "20", queryRes)
+
+			queryRes, qErr := queryHelper(t, keeper, ctx, v010ContractAddress, `{"get":{}}`, true, false, math.MaxUint64)
+			require.Empty(t, qErr)
+
+			var resp v1QueryResponse
+			e := json.Unmarshal([]byte(queryRes), &resp)
+			require.NoError(t, e)
+			require.Equal(t, expected_v010_result, resp.Get.Count)
+			expected_v010_result += 5
+		})
+	}
+}
+
 func TestIBCPacketAck(t *testing.T) {
 	ctx, keeper, codeID, _, walletA, privKeyA, _, _ := setupTest(t, "./testdata/ibc/contract.wasm", sdk.NewCoins())
 

x/compute/internal/keeper/msg_dispatcher.go

@@ -274,8 +274,9 @@ func (d MessageDispatcher) DispatchSubmessages(ctx sdk.Context, contractAddr sdk
 		msg_id := []byte(fmt.Sprint(msg.ID))
 		// now handle the reply, we use the parent context, and abort on error
 		reply := v1wasmTypes.Reply{
-			ID:     msg_id,
-			Result: result,
+			ID:              msg_id,
+			Result:          result,
+			WasMsgEncrypted: msg.WasMsgEncrypted,
 		}
 
 		// we can ignore any result returned as there is nothing to do with the data

x/compute/internal/keeper/test_common.go

@@ -364,7 +364,7 @@ func CreateTestInput(t *testing.T, isCheckTx bool, supportedFeatures string, enc
 		authtypes.StoreKey, banktypes.StoreKey, stakingtypes.StoreKey,
 		minttypes.StoreKey, distrtypes.StoreKey, slashingtypes.StoreKey,
 		govtypes.StoreKey, paramstypes.StoreKey, ibchost.StoreKey, upgradetypes.StoreKey,
-		evidencetypes.StoreKey, ibctransfertypes.StoreKey, capabilitytypes.StoreKey,
+		evidencetypes.StoreKey, ibctransfertypes.StoreKey, capabilitytypes.StoreKey, "compute",
 		feegrant.StoreKey, authzkeeper.StoreKey, icahosttypes.StoreKey,
 	)
 

x/compute/internal/keeper/testdata/ibc/src/contract.rs

@@ -7,6 +7,8 @@ use cosmwasm_std::{
     IbcReceiveResponse, MessageInfo, Reply, ReplyOn, Response, StdError, StdResult, SubMsg,
     SubMsgResult, WasmMsg,
 };
+use serde::{Deserialize, Serialize};
+use serde_json_wasm as serde_json;
 
 pub const IBC_APP_VERSION: &str = "ibc-v1";
 
@@ -39,12 +41,28 @@ pub fn query(deps: Deps, _env: Env, _msg: QueryMsg) -> StdResult<Binary> {
 }
 
 #[entry_point]
-pub fn reply(deps: DepsMut, _env: Env, reply: Reply) -> StdResult<Response> {
+pub fn reply(deps: DepsMut, env: Env, reply: Reply) -> StdResult<Response> {
     match (reply.id, reply.result) {
         (1, SubMsgResult::Err(_)) => Err(StdError::generic_err("Failed to inc")),
         (1, SubMsgResult::Ok(_)) => {
             increment(deps, 6)?;
-            Ok(Response::new().set_data(to_binary(&"out2".to_string())?))
+            Ok(Response::new().set_data(to_binary(&"out".to_string())?))
+        }
+        (2, SubMsgResult::Err(_)) => Err(StdError::generic_err("Failed to inc")),
+        (2, SubMsgResult::Ok(_)) => {
+            increment(deps, 6)?;
+            Ok(Response::new().add_submessage(SubMsg {
+                id: 1,
+                msg: CosmosMsg::Wasm(WasmMsg::Execute {
+                    code_hash: env.contract.code_hash,
+                    contract_addr: env.contract.address.into_string(),
+                    msg: Binary::from("{\"increment\":{\"addition\":5}}".as_bytes().to_vec()),
+                    funds: vec![],
+                })
+                .into(),
+                reply_on: ReplyOn::Always,
+                gas_limit: None,
+            }))
         }
         _ => Err(StdError::generic_err("invalid reply id or result")),
     }
@@ -115,7 +133,24 @@ pub fn get_resp_based_on_num(env: Env, num: u64) -> StdResult<IbcBasicResponse>
     }
 }
 
-pub fn get_recv_resp_based_on_num(env: Env, num: u64) -> StdResult<IbcReceiveResponse> {
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
+pub struct ContractInfo {
+    pub address: String,
+    pub hash: String,
+}
+
+pub fn is_reply_on(num: u64) -> bool {
+    match num {
+        2 | 6 => true,
+        _ => false,
+    }
+}
+
+pub fn get_recv_resp_based_on_num(
+    env: Env,
+    num: u64,
+    data: Binary,
+) -> StdResult<IbcReceiveResponse> {
     match num {
         0 => Ok(IbcReceiveResponse::default()),
         1 => Ok(IbcReceiveResponse::new().add_submessage(SubMsg {
@@ -146,6 +181,37 @@ pub fn get_recv_resp_based_on_num(env: Env, num: u64) -> StdResult<IbcReceiveRes
         4 => Ok(IbcReceiveResponse::new()
             .add_event(Event::new("cyber1".to_string()).add_attribute("attr1", "🤯"))),
         5 => Err(StdError::generic_err("Intentional")),
+        6 => Ok(IbcReceiveResponse::new().add_submessage(SubMsg {
+            id: 2,
+            msg: CosmosMsg::Wasm(WasmMsg::Execute {
+                code_hash: env.contract.code_hash,
+                contract_addr: env.contract.address.into_string(),
+                msg: Binary::from("{\"increment\":{\"addition\":5}}".as_bytes().to_vec()),
+                funds: vec![],
+            })
+            .into(),
+            reply_on: ReplyOn::Always,
+            gas_limit: None,
+        })),
+        7 => {
+            let contract_info: ContractInfo = serde_json::from_slice(data.as_slice()).unwrap();
+            Ok(IbcReceiveResponse::new().add_submessage(SubMsg {
+                id: 1,
+                msg: CosmosMsg::Wasm(WasmMsg::Execute {
+                    code_hash: contract_info.hash,
+                    contract_addr: contract_info.address,
+                    msg: Binary::from(
+                        "{\"increment_from_v1\":{\"addition\":5}}"
+                            .as_bytes()
+                            .to_vec(),
+                    ),
+                    funds: vec![],
+                })
+                .into(),
+                reply_on: ReplyOn::Always,
+                gas_limit: None,
+            }))
+        }
         _ => Err(StdError::generic_err("Unsupported channel connect type")),
     }
 }
@@ -207,10 +273,12 @@ pub fn ibc_packet_receive(
     msg: IbcPacketReceiveMsg,
 ) -> StdResult<IbcReceiveResponse> {
     count(deps.storage).save(&(msg.packet.sequence + 7))?;
-    let mut resp = get_recv_resp_based_on_num(env, msg.packet.sequence);
+    let mut resp = get_recv_resp_based_on_num(env, msg.packet.sequence, msg.packet.data);
     match &mut resp {
         Ok(r) => {
-            r.acknowledgement = to_binary(&"out".to_string())?;
+            if !is_reply_on(msg.packet.sequence) {
+                r.acknowledgement = to_binary(&"out".to_string())?;
+            }
         }
         Err(_) => {}
     }