WIP: CosmWasm v1 IBC Receive Hooks

.vscode/settings.json

@@ -3,7 +3,9 @@
     "./cosmwasm/Cargo.toml",
     "./cosmwasm/enclaves/Cargo.toml",
     "./x/compute/internal/keeper/testdata/v1-sanity-contract/Cargo.toml",
-    "./x/compute/internal/keeper/testdata/test-contract/Cargo.toml"
+    "./x/compute/internal/keeper/testdata/test-contract/Cargo.toml",
+    "./cosmwasm/enclaves/shared/cosmwasm-v1-types/Cargo.toml",
+    "./cosmwasm/enclaves/shared/cosmwasm-v010-types/Cargo.toml"
   ],
   "rust-analyzer.diagnostics.experimental.enable": true,
   "rust-analyzer.rustfmt.rangeFormatting.enable": true,

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

@@ -1,4 +1,7 @@
 use log::*;
+use serde::de::DeserializeOwned;
+use serde::Serialize;
+use std::fmt::Debug;
 
 use enclave_ffi_types::{Ctx, EnclaveError};
 
@@ -11,6 +14,9 @@ use enclave_cosmos_types::types::{ContractCode, HandleType, SigInfo};
 use enclave_cosmwasm_v010_types as cosmwasm_v010_types;
 use enclave_cosmwasm_v010_types::encoding::Binary;
 use enclave_cosmwasm_v1_types::addresses::Addr;
+use enclave_cosmwasm_v1_types::ibc::{
+    IbcPacketAckMsg, IbcPacketReceiveMsg, IbcPacketTimeoutMsg, IbcPacketTrait,
+};
 use enclave_cosmwasm_v1_types::results::{
     DecryptedReply, Event, Reply, SubMsgResponse, SubMsgResult,
 };
@@ -24,7 +30,7 @@ use super::contract_validation::{
     verify_params, ContractKey,
 };
 use super::gas::WasmCosts;
-use super::io::encrypt_output;
+use super::io::{encrypt_output, finalize_raw_output, RawWasmOutput};
 use super::module_cache::create_module_instance;
 use super::types::{IoNonce, SecretMessage};
 use super::wasm::{ContractInstance, ContractOperation, Engine};
@@ -178,6 +184,11 @@ pub struct ParsedMessage {
     pub contract_hash_for_validation: Option<Vec<u8>>,
 }
 
+pub struct DecryptedSecretMessage {
+    pub secret_msg: SecretMessage,
+    pub decrypted_msg: Vec<u8>,
+}
+
 pub fn redact_custom_events(reply: &mut Reply) {
     reply.result = match &reply.result {
         SubMsgResult::Ok(r) => {
@@ -214,33 +225,171 @@ pub fn redact_custom_events(reply: &mut Reply) {
     };
 }
 
-// Parse the message that was passed to handle (Based on the assumption that it might be a reply or IBC as well)
-pub fn parse_message(
+pub fn get_secret_msg(message: &[u8]) -> SecretMessage {
+    match SecretMessage::from_slice(message) {
+        Ok(orig_secret_msg) => orig_secret_msg,
+        Err(_) => {
+            trace!(
+                "Msg is not SecretMessage (probably plaintext): {:?}",
+                base64::encode(&message)
+            );
+
+            return SecretMessage {
+                nonce: [0; 32],
+                user_public_key: [0; 32],
+                msg: message.into(),
+            };
+        }
+    }
+}
+
+pub fn try_get_decrypted_secret_msg(message: &[u8]) -> Option<DecryptedSecretMessage> {
+    let secret_msg = get_secret_msg(message);
+    match secret_msg.decrypt() {
+        Ok(decrypted_msg) => Some(DecryptedSecretMessage {
+            secret_msg,
+            decrypted_msg,
+        }),
+        Err(_) => None,
+    }
+}
+
+pub fn parse_ibc_packet<T>(
+    _t: T,
     message: &[u8],
-    sig_info: &SigInfo,
-    handle_type: &HandleType,
-) -> Result<ParsedMessage, EnclaveError> {
-    let orig_secret_msg = SecretMessage::from_slice(message)?;
+    orig_secret_msg: SecretMessage,
+    function_name: &str,
+) -> Result<ParsedMessage, EnclaveError>
+where
+    T: IbcPacketTrait<Data = Binary> + Serialize + DeserializeOwned + Debug,
+{
+    let mut parsed_encrypted_ibc_packet: T =
+        serde_json::from_slice(&orig_secret_msg.msg.as_slice().to_vec()).map_err(|err| {
+            warn!(
+            "{} msg got an error while trying to deserialize msg input bytes into json {:?}: {}",
+            function_name,
+            String::from_utf8_lossy(&orig_secret_msg.msg),
+            err
+        );
+            EnclaveError::FailedToDeserialize
+        })?;
+
+    let tmp_secret_data = SecretMessage {
+        nonce: orig_secret_msg.nonce,
+        user_public_key: orig_secret_msg.user_public_key,
+        msg: parsed_encrypted_ibc_packet.get_packet().as_slice().to_vec(),
+    };
+
+    match tmp_secret_data.decrypt() {
+        Ok(decrypted_msg) => {
+            // IBC packet was encrypted
 
-    return match handle_type {
-        HandleType::HANDLE_TYPE_EXECUTE => {
             trace!(
-                "handle input before decryption: {:?}",
+                "{} input before decryption: {:?}",
+                function_name,
                 base64::encode(&message)
             );
-            let decrypted_msg = orig_secret_msg.decrypt()?;
+
+            parsed_encrypted_ibc_packet.set_packet(decrypted_msg.as_slice().into());
+
+            match parsed_encrypted_ibc_packet.get_ack() {
+                Some(ack_data) => {
+                    let tmp_secret_ack_data = SecretMessage {
+                        nonce: orig_secret_msg.nonce,
+                        user_public_key: orig_secret_msg.user_public_key,
+                        msg: ack_data.as_slice().to_vec(),
+                    };
+
+                    parsed_encrypted_ibc_packet
+                        .set_ack(tmp_secret_ack_data.decrypt()?.as_slice().into());
+                }
+                None => {}
+            }
+
             Ok(ParsedMessage {
-                should_validate_sig_info: true,
+                should_validate_sig_info: false,
                 was_msg_encrypted: true,
                 secret_msg: orig_secret_msg,
+                decrypted_msg: serde_json::to_vec(&parsed_encrypted_ibc_packet).map_err(|err| {
+                    warn!(
+                        "got an error while trying to serialize {} msg into bytes {:?}: {}",
+                        function_name, parsed_encrypted_ibc_packet, err
+                    );
+                    EnclaveError::FailedToSerialize
+                })?,
+                contract_hash_for_validation: None,
+            })
+        }
+        Err(_) => {
+            // assume packet is not encrypted, continue in plaintext mode
+
+            trace!(
+                "{} input was plaintext: {:?}",
+                function_name,
+                base64::encode(&message)
+            );
+
+            let decrypted_msg = orig_secret_msg.msg.clone();
+
+            Ok(ParsedMessage {
+                should_validate_sig_info: false,
+                was_msg_encrypted: false,
+                secret_msg: orig_secret_msg,
                 decrypted_msg,
                 contract_hash_for_validation: None,
             })
         }
+    }
+}
 
+// Parse the message that was passed to handle (Based on the assumption that it might be a reply or IBC as well)
+pub fn parse_message(
+    message: &[u8],
+    sig_info: &SigInfo,
+    handle_type: &HandleType,
+) -> Result<ParsedMessage, EnclaveError> {
+    return match handle_type {
+        HandleType::HANDLE_TYPE_EXECUTE => match try_get_decrypted_secret_msg(message) {
+            Some(decrypted_secret_msg) => {
+                trace!(
+                    "execute input before decryption: {:?}",
+                    base64::encode(&message)
+                );
+
+                Ok(ParsedMessage {
+                    should_validate_sig_info: true,
+                    was_msg_encrypted: true,
+                    secret_msg: decrypted_secret_msg.secret_msg,
+                    decrypted_msg: decrypted_secret_msg.decrypted_msg,
+                    contract_hash_for_validation: None,
+                })
+            }
+            None => {
+                trace!(
+                    "execute input was plaintext: {:?}",
+                    base64::encode(&message)
+                );
+
+                let secret_msg = get_secret_msg(message);
+                let decrypted_msg = secret_msg.msg.clone();
+
+                Ok(ParsedMessage {
+                    should_validate_sig_info: true,
+                    was_msg_encrypted: false,
+                    secret_msg,
+                    decrypted_msg,
+                    contract_hash_for_validation: None,
+                })
+            }
+        },
         HandleType::HANDLE_TYPE_REPLY => {
+            let orig_secret_msg = SecretMessage::from_slice(message)?;
+
             if sig_info.sign_mode == SignMode::SIGN_MODE_UNSPECIFIED {
-                trace!("reply input is not encrypted");
+                trace!(
+                    "reply input is not encrypted: {:?}",
+                    base64::encode(&message)
+                );
                 let decrypted_msg = orig_secret_msg.msg.clone();
                 let mut reply: Reply = serde_json::from_slice(&decrypted_msg)
                     .map_err(|err| {
@@ -306,7 +455,6 @@ pub fn parse_message(
                 });
             }
 
-            // Here we are sure the reply is OK because only OK is encrypted
             trace!(
                 "reply input before decryption: {:?}",
                 base64::encode(&message)
@@ -507,6 +655,61 @@ pub fn parse_message(
                 }
             }
         }
+        HandleType::HANDLE_TYPE_IBC_CHANNEL_OPEN
+        | HandleType::HANDLE_TYPE_IBC_CHANNEL_CONNECT
+        | HandleType::HANDLE_TYPE_IBC_CHANNEL_CLOSE => {
+            trace!(
+                "parsing {} msg (Should always be plaintext): {:?}",
+                HandleType::to_export_name(&handle_type),
+                base64::encode(&message)
+            );
+
+            let scrt_msg = SecretMessage {
+                nonce: [0; 32],
+                user_public_key: [0; 32],
+                msg: message.into(),
+            };
+
+            let decrypted_msg = scrt_msg.msg.clone();
+
+            Ok(ParsedMessage {
+                should_validate_sig_info: false,
+                was_msg_encrypted: false,
+                secret_msg: scrt_msg,
+                decrypted_msg,
+                contract_hash_for_validation: None,
+            })
+        }
+        HandleType::HANDLE_TYPE_IBC_PACKET_RECEIVE => {
+            // LIORRR TODO: Maybe mark whether the message was encrypted or not.
+            let orig_secret_msg = get_secret_msg(message);
+            parse_ibc_packet(
+                IbcPacketReceiveMsg::default(),
+                message,
+                orig_secret_msg,
+                "ibc_packet_receive",
+            )
+        }
+        HandleType::HANDLE_TYPE_IBC_PACKET_ACK => {
+            // LIORRR TODO: Maybe mark whether the message was encrypted or not.
+            let orig_secret_msg = get_secret_msg(message);
+            parse_ibc_packet(
+                IbcPacketAckMsg::default(),
+                message,
+                orig_secret_msg,
+                "ibc_packet_receive",
+            )
+        }
+        HandleType::HANDLE_TYPE_IBC_PACKET_TIMEOUT => {
+            // LIORRR TODO: Maybe mark whether the message was encrypted or not.
+            let orig_secret_msg = get_secret_msg(message);
+            parse_ibc_packet(
+                IbcPacketTimeoutMsg::default(),
+                message,
+                orig_secret_msg,
+                "ibc_packet_timeout",
+            )
+        }
     };
 }
 
@@ -634,22 +837,41 @@ pub fn handle(
     let output = coalesce!(EnclaveError, {
         let vec_ptr = engine.handle(env_ptr, msg_info_ptr, msg_ptr, parsed_handle_type)?;
 
-        let output = engine.extract_vector(vec_ptr)?;
+        let mut output = engine.extract_vector(vec_ptr)?;
 
         debug!(
             "(2) nonce just before encrypt_output: nonce = {:?} pubkey = {:?}",
             secret_msg.nonce, secret_msg.user_public_key
         );
 
-        let output = encrypt_output(
-            output,
-            &secret_msg,
-            &canonical_contract_address,
-            &env_v010.contract_code_hash,
-            reply_params,
-            &canonical_sender_address,
-            false,
-        )?;
+        if was_msg_encrypted {
+            output = encrypt_output(
+                output,
+                &secret_msg,
+                &canonical_contract_address,
+                &env_v010.contract_code_hash,
+                reply_params,
+                &canonical_sender_address,
+                false,
+            )?;
+        } else {
+            let raw_output: RawWasmOutput = serde_json::from_slice(&output).map_err(|err| {
+                warn!("got an error while trying to deserialize output bytes into json");
+                trace!("output: {:?} error: {:?}", output, err);
+                EnclaveError::FailedToDeserialize
+            })?;
+
+            let finalized_output = finalize_raw_output(raw_output, false);
+
+            output = serde_json::to_vec(&finalized_output).map_err(|err| {
+                debug!(
+                    "got an error while trying to serialize output json into bytes {:?}: {}",
+                    finalized_output, err
+                );
+                EnclaveError::FailedToSerialize
+            })?;
+        }
+
         Ok(output)
     })
     .map_err(|err| {