Encrypt the output of every reply (Should be changed)

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

@@ -186,6 +186,7 @@ pub fn handle(
     let ParsedMessage {
         should_validate_sig_info,
         was_msg_encrypted,
+        should_encrypt_output,
         secret_msg,
         decrypted_msg,
         contract_hash_for_validation,
@@ -262,7 +263,7 @@ pub fn handle(
             secret_msg.nonce, secret_msg.user_public_key
         );
 
-        if was_msg_encrypted {
+        if should_encrypt_output {
             output = encrypt_output(
                 output,
                 &secret_msg,

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

@@ -411,8 +411,8 @@ pub fn encrypt_output(
     // More info in: https://github.com/CosmWasm/cosmwasm/blob/v1.0.0/packages/std/src/results/submessages.rs#L192-L198
     let encryption_key = calc_encryption_key(&secret_msg.nonce, &secret_msg.user_public_key);
     trace!(
-        "Output before encryption: {:?}",
-        String::from_utf8_lossy(&output)
+        "Output before encryption: {:?} {:?} {:?}",
+        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| {

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

@@ -1,5 +1,4 @@
 use log::{trace, warn};
-use serde::Serialize;
 
 use cosmos_proto::tx::signing::SignMode;
 use cw_types_v010::encoding::Binary;
@@ -15,6 +14,7 @@ const HEX_ENCODED_HASH_SIZE: usize = 64;
 pub struct ParsedMessage {
     pub should_validate_sig_info: bool,
     pub was_msg_encrypted: bool,
+    pub should_encrypt_output: bool,
     pub secret_msg: SecretMessage,
     pub decrypted_msg: Vec<u8>,
     pub contract_hash_for_validation: Option<Vec<u8>>,
@@ -107,6 +107,7 @@ pub fn parse_message(
                 Ok(ParsedMessage {
                     should_validate_sig_info: true,
                     was_msg_encrypted: true,
+                    should_encrypt_output: true,
                     secret_msg: decrypted_secret_msg.secret_msg,
                     decrypted_msg: decrypted_secret_msg.decrypted_msg,
                     contract_hash_for_validation: None,
@@ -124,6 +125,7 @@ pub fn parse_message(
                 Ok(ParsedMessage {
                     should_validate_sig_info: false,
                     was_msg_encrypted: false,
+                    should_encrypt_output: false,
                     secret_msg,
                     decrypted_msg,
                     contract_hash_for_validation: None,
@@ -197,6 +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
                     secret_msg: reply_secret_msg,
                     decrypted_msg: serialized_reply,
                     contract_hash_for_validation: None,
@@ -301,6 +304,7 @@ pub fn parse_message(
                     Ok(ParsedMessage {
                         should_validate_sig_info: true,
                         was_msg_encrypted: true,
+                        should_encrypt_output: true,
                         secret_msg: reply_secret_msg,
                         decrypted_msg: decrypted_reply_as_vec,
                         contract_hash_for_validation: Some(
@@ -394,6 +398,7 @@ pub fn parse_message(
                     Ok(ParsedMessage {
                         should_validate_sig_info: true,
                         was_msg_encrypted: true,
+                        should_encrypt_output: true,
                         secret_msg: reply_secret_msg,
                         decrypted_msg: decrypted_reply_as_vec,
                         contract_hash_for_validation: Some(
@@ -425,6 +430,7 @@ pub fn parse_message(
             Ok(ParsedMessage {
                 should_validate_sig_info: false,
                 was_msg_encrypted: false,
+                should_encrypt_output: false,
                 secret_msg: scrt_msg,
                 decrypted_msg,
                 contract_hash_for_validation: None,
@@ -471,6 +477,7 @@ pub fn parse_message(
             Ok(ParsedMessage {
                 should_validate_sig_info: false,
                 was_msg_encrypted,
+                should_encrypt_output: was_msg_encrypted,
                 secret_msg: orig_secret_msg,
                 decrypted_msg: serde_json::to_vec(&parsed_encrypted_ibc_packet).map_err(|err| {
                     warn!(

x/compute/internal/keeper/ibc_test.go

@@ -905,3 +905,126 @@ func TestIBCPacketAck(t *testing.T) {
 		})
 	}
 }
+
+func TestIBCPacketTimeout(t *testing.T) {
+	ctx, keeper, codeID, _, walletA, privKeyA, _, _ := setupTest(t, "./testdata/ibc/contract.wasm", sdk.NewCoins())
+
+	_, _, contractAddress, _, err := initHelper(t, keeper, ctx, codeID, walletA, privKeyA, `{"init":{}}`, true, true, defaultGasForTests)
+	require.Empty(t, err)
+
+	for _, test := range []struct {
+		description   string
+		sequence      uint64
+		output        string
+		isSuccess     bool
+		hasAttributes bool
+		hasEvents     bool
+	}{
+		{
+			description:   "Default",
+			sequence:      0,
+			output:        "9",
+			isSuccess:     true,
+			hasAttributes: false,
+			hasEvents:     false,
+		},
+		{
+			description:   "SubmessageNoReply",
+			sequence:      1,
+			output:        "15",
+			isSuccess:     true,
+			hasAttributes: false,
+			hasEvents:     false,
+		},
+		{
+			description:   "SubmessageWithReply",
+			sequence:      2,
+			output:        "22",
+			isSuccess:     true,
+			hasAttributes: false,
+			hasEvents:     false,
+		},
+		{
+			description:   "Attributes",
+			sequence:      3,
+			output:        "12",
+			isSuccess:     true,
+			hasAttributes: true,
+			hasEvents:     false,
+		},
+		{
+			description:   "Events",
+			sequence:      4,
+			output:        "13",
+			isSuccess:     true,
+			hasAttributes: false,
+			hasEvents:     true,
+		},
+		{
+			description:   "Error",
+			sequence:      5,
+			output:        "",
+			isSuccess:     false,
+			hasAttributes: false,
+			hasEvents:     false,
+		},
+	} {
+		t.Run(test.description, func(t *testing.T) {
+			ibcPacket := createIBCPacket(createIBCEndpoint(PortIDForContract(contractAddress), "channel.1"),
+				createIBCEndpoint(PortIDForContract(contractAddress), "channel.0"),
+				test.sequence,
+				createIBCTimeout(math.MaxUint64),
+				[]byte{},
+			)
+
+			ctx, events, err := ibcPacketTimeoutHelper(t, keeper, ctx, contractAddress, privKeyA, defaultGasForTests, ibcPacket)
+
+			if !test.isSuccess {
+				require.Contains(t, fmt.Sprintf("%+v", err), "Intentional")
+			} else {
+				require.Empty(t, err)
+				if test.hasAttributes {
+					require.Equal(t,
+						[]ContractEvent{
+							{
+								{Key: "contract_address", Value: contractAddress.String()},
+								{Key: "attr1", Value: "😗"},
+							},
+						},
+						events,
+					)
+				}
+
+				if test.hasEvents {
+					hadCyber1 := false
+					evts := ctx.EventManager().Events()
+					for _, e := range evts {
+						if e.Type == "wasm-cyber1" {
+							require.False(t, hadCyber1)
+							attrs, err := parseAndDecryptAttributes(e.Attributes, []byte{}, false)
+							require.Empty(t, err)
+
+							require.Equal(t,
+								[]v010cosmwasm.LogAttribute{
+									{Key: "contract_address", Value: contractAddress.String()},
+									{Key: "attr1", Value: "🤯"},
+								},
+								attrs,
+							)
+
+							hadCyber1 = true
+						}
+					}
+
+					require.True(t, hadCyber1)
+				}
+
+				queryRes, err := queryHelper(t, keeper, ctx, contractAddress, `{"q":{}}`, true, true, math.MaxUint64)
+
+				require.Empty(t, err)
+
+				require.Equal(t, test.output, queryRes)
+			}
+		})
+	}
+}

x/compute/internal/keeper/secret_contracts_test.go

@@ -232,11 +232,8 @@ func tryDecryptWasmEvents(ctx sdk.Context, nonce []byte, shouldSkipAttributes ..
 
 				if !shouldSkip && newLog.Key != "contract_address" {
 					// key
-					fmt.Printf("LIORRR before %+v nonce %+v", newLog, nonce)
 					newAttr, err := decryptAttribute(newLog, nonce)
-					fmt.Printf("LIORRR after %+v", newAttr)
 					if err != nil {
-						fmt.Printf("LIORRR here %+v", err)
 						continue
 					}
 

x/compute/internal/keeper/test_common.go

@@ -369,6 +369,9 @@ func CreateTestInput(t *testing.T, isCheckTx bool, supportedFeatures string, enc
 	)
 
 	memKeys := sdk.NewMemoryStoreKeys(capabilitytypes.MemStoreKey)
+	for _, v := range memKeys {
+		ms.MountStoreWithDB(v, sdk.StoreTypeMemory, db)
+	}
 
 	upgradeKeeper := upgradekeeper.NewKeeper(
 		map[int64]bool{},

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

@@ -3,8 +3,9 @@ use crate::state::{count, count_read};
 use cosmwasm_std::{
     entry_point, to_binary, Binary, CosmosMsg, Deps, DepsMut, Env, Event, Ibc3ChannelOpenResponse,
     IbcBasicResponse, IbcChannelCloseMsg, IbcChannelConnectMsg, IbcChannelOpenMsg,
-    IbcChannelOpenResponse, IbcPacketAckMsg, IbcPacketReceiveMsg, IbcReceiveResponse, MessageInfo,
-    Reply, ReplyOn, Response, StdError, StdResult, SubMsg, SubMsgResult, WasmMsg,
+    IbcChannelOpenResponse, IbcPacketAckMsg, IbcPacketReceiveMsg, IbcPacketTimeoutMsg,
+    IbcReceiveResponse, MessageInfo, Reply, ReplyOn, Response, StdError, StdResult, SubMsg,
+    SubMsgResult, WasmMsg,
 };
 
 pub const IBC_APP_VERSION: &str = "ibc-v1";
@@ -43,7 +44,7 @@ pub fn reply(deps: DepsMut, _env: Env, reply: Reply) -> StdResult<Response> {
         (1, SubMsgResult::Err(_)) => Err(StdError::generic_err("Failed to inc")),
         (1, SubMsgResult::Ok(_)) => {
             increment(deps, 6)?;
-            Ok(Response::default())
+            Ok(Response::new().set_data(to_binary(&"out2".to_string())?))
         }
         _ => Err(StdError::generic_err("invalid reply id or result")),
     }
@@ -233,3 +234,13 @@ pub fn ibc_packet_ack(
     count(deps.storage).save(&(msg.original_packet.sequence + 8))?;
     get_resp_based_on_num(env, msg.original_packet.sequence)
 }
+
+// #[entry_point]
+// pub fn ibc_packet_timeout(
+//     deps: DepsMut,
+//     env: Env,
+//     msg: IbcPacketTimeoutMsg,
+// ) -> StdResult<IbcBasicResponse> {
+//     count(deps.storage).save(&(msg.packet.sequence + 9))?;
+//     get_resp_based_on_num(env, msg.packet.sequence)
+// }