add WarpMsgs logic to job-account

terra-money · Nov 13, 2023 · dd5f2eb · dd5f2eb
1 parent 2421ebc
commit dd5f2eb
Show file tree

Hide file tree

Showing 5 changed files with 67 additions and 7 deletions.
diff --git a/contracts/warp-job-account/src/contract.rs b/contracts/warp-job-account/src/contract.rs
@@ -54,9 +54,10 @@ pub fn execute(
             .add_attribute("action", "generic")),
         ExecuteMsg::WithdrawAssets(data) => {
             nonpayable(&info).unwrap();
-            execute::withdraw::withdraw_assets(deps, env, data, config)
+            execute::withdraw::withdraw_assets(deps.as_ref(), env, data, config)
         }
         ExecuteMsg::IbcTransfer(data) => execute::ibc::ibc_transfer(env, data),
+        ExecuteMsg::WarpMsgs(data) => execute::msgs::execute_warp_msgs(deps, env, data, config),
     }
 }
 

diff --git a/contracts/warp-job-account/src/execute/mod.rs b/contracts/warp-job-account/src/execute/mod.rs
@@ -1,2 +1,3 @@
 pub(crate) mod ibc;
+pub(crate) mod msgs;
 pub(crate) mod withdraw;
diff --git a/contracts/warp-job-account/src/execute/msgs.rs b/contracts/warp-job-account/src/execute/msgs.rs
@@ -0,0 +1,44 @@
+use crate::ContractError;
+use cosmwasm_std::{Env, Response};
+use job_account::{Config, WarpMsg, WarpMsgs};
+
+use cosmwasm_std::{CosmosMsg, DepsMut};
+
+use super::ibc::ibc_transfer;
+use super::withdraw::withdraw_assets;
+
+pub fn execute_warp_msgs(
+    deps: DepsMut,
+    env: Env,
+    data: WarpMsgs,
+    config: Config,
+) -> Result<Response, ContractError> {
+    let msgs = data
+        .msgs
+        .into_iter()
+        .flat_map(|msg| -> Vec<CosmosMsg> {
+            match msg {
+                WarpMsg::Generic(msg) => vec![msg],
+                WarpMsg::IbcTransfer(msg) => ibc_transfer(env.clone(), msg)
+                    .map(extract_messages)
+                    .unwrap(),
+                WarpMsg::WithdrawAssets(msg) => {
+                    withdraw_assets(deps.as_ref(), env.clone(), msg, config.clone())
+                        .map(extract_messages)
+                        .unwrap()
+                }
+            }
+        })
+        .collect::<Vec<CosmosMsg>>();
+
+    Ok(Response::new()
+        .add_messages(msgs)
+        .add_attribute("action", "warp_msgs"))
+}
+
+fn extract_messages(resp: Response) -> Vec<CosmosMsg> {
+    resp.messages
+        .into_iter()
+        .map(|cosmos_msg| cosmos_msg.msg)
+        .collect()
+}
diff --git a/contracts/warp-job-account/src/execute/withdraw.rs b/contracts/warp-job-account/src/execute/withdraw.rs
@@ -1,5 +1,5 @@
 use cosmwasm_std::{
-    to_binary, Addr, BankMsg, CosmosMsg, Deps, DepsMut, Env, Response, StdResult, Uint128, WasmMsg,
+    to_binary, Addr, BankMsg, CosmosMsg, Deps, Env, Response, StdResult, Uint128, WasmMsg,
 };
 use cw20::{BalanceResponse, Cw20ExecuteMsg};
 use cw721::{Cw721QueryMsg, OwnerOfResponse};
@@ -9,7 +9,7 @@ use controller::account::{AssetInfo, Cw721ExecuteMsg};
 use job_account::{Config, WithdrawAssetsMsg};
 
 pub fn withdraw_assets(
-    deps: DepsMut,
+    deps: Deps,
     env: Env,
     data: WithdrawAssetsMsg,
     config: Config,
@@ -20,7 +20,7 @@ pub fn withdraw_assets(
         match asset_info {
             AssetInfo::Native(denom) => {
                 let withdraw_native_msg =
-                    withdraw_asset_native(deps.as_ref(), env.clone(), &config.owner, denom)?;
+                    withdraw_asset_native(deps, env.clone(), &config.owner, denom)?;
 
                 match withdraw_native_msg {
                     None => {}
@@ -29,16 +29,15 @@ pub fn withdraw_assets(
             }
             AssetInfo::Cw20(addr) => {
                 let withdraw_cw20_msg =
-                    withdraw_asset_cw20(deps.as_ref(), env.clone(), &config.owner, addr)?;
+                    withdraw_asset_cw20(deps, env.clone(), &config.owner, addr)?;
 
                 match withdraw_cw20_msg {
                     None => {}
                     Some(msg) => withdraw_msgs.push(msg),
                 }
             }
             AssetInfo::Cw721(addr, token_id) => {
-                let withdraw_cw721_msg =
-                    withdraw_asset_cw721(deps.as_ref(), &config.owner, addr, token_id)?;
+                let withdraw_cw721_msg = withdraw_asset_cw721(deps, &config.owner, addr, token_id)?;
                 match withdraw_cw721_msg {
                     None => {}
                     Some(msg) => withdraw_msgs.push(msg),

diff --git a/packages/job-account/src/lib.rs b/packages/job-account/src/lib.rs
@@ -28,6 +28,9 @@ pub struct InstantiateMsg {
 #[cw_serde]
 #[allow(clippy::large_enum_variant)]
 pub enum ExecuteMsg {
+    WarpMsgs(WarpMsgs),
+
+    // legacy flow
     Generic(GenericMsg),
     WithdrawAssets(WithdrawAssetsMsg),
     IbcTransfer(IbcTransferMsg),
@@ -38,6 +41,18 @@ pub struct GenericMsg {
     pub msgs: Vec<CosmosMsg>,
 }
 
+#[cw_serde]
+pub struct WarpMsgs {
+    pub msgs: Vec<WarpMsg>,
+}
+
+#[cw_serde]
+pub enum WarpMsg {
+    Generic(CosmosMsg),
+    IbcTransfer(IbcTransferMsg),
+    WithdrawAssets(WithdrawAssetsMsg),
+}
+
 #[derive(Serialize, Deserialize, Clone, PartialEq, JsonSchema, prost::Message)]
 pub struct Coin {
     #[prost(string, tag = "1")]