fixed unsafe force withdraw to support native and cw20 tokens

DA0-DA0 · Nov 1, 2024 · 321de22 · 321de22
1 parent 78226e5
commit 321de22
Show file tree

Hide file tree

Showing 5 changed files with 59 additions and 36 deletions.
diff --git a/contracts/distribution/dao-rewards-distributor/schema/dao-rewards-distributor.json b/contracts/distribution/dao-rewards-distributor/schema/dao-rewards-distributor.json
@@ -235,11 +235,25 @@
           "unsafe_force_withdraw": {
             "type": "object",
             "required": [
-              "amount"
+              "amount",
+              "denom"
             ],
             "properties": {
               "amount": {
-                "$ref": "#/definitions/Coin"
+                "description": "amount to withdraw",
+                "allOf": [
+                  {
+                    "$ref": "#/definitions/Uint128"
+                  }
+                ]
+              },
+              "denom": {
+                "description": "denom to withdraw",
+                "allOf": [
+                  {
+                    "$ref": "#/definitions/UncheckedDenom"
+                  }
+                ]
               }
             },
             "additionalProperties": false
@@ -321,21 +335,6 @@
         "description": "Binary is a wrapper around Vec<u8> to add base64 de/serialization with serde. It also adds some helper methods to help encode inline.\n\nThis is only needed as serde-json-{core,wasm} has a horrible encoding for Vec<u8>. See also <https://github.com/CosmWasm/cosmwasm/blob/main/docs/MESSAGE_TYPES.md>.",
         "type": "string"
       },
-      "Coin": {
-        "type": "object",
-        "required": [
-          "amount",
-          "denom"
-        ],
-        "properties": {
-          "amount": {
-            "$ref": "#/definitions/Uint128"
-          },
-          "denom": {
-            "type": "string"
-          }
-        }
-      },
       "CreateMsg": {
         "type": "object",
         "required": [

diff --git a/contracts/distribution/dao-rewards-distributor/src/contract.rs b/contracts/distribution/dao-rewards-distributor/src/contract.rs
@@ -1,11 +1,11 @@
 #[cfg(not(feature = "library"))]
 use cosmwasm_std::entry_point;
 use cosmwasm_std::{
-    ensure, from_json, to_json_binary, Addr, BankMsg, Binary, Coin, CosmosMsg, Deps, DepsMut, Env,
-    MessageInfo, Order, Response, StdError, StdResult, Uint128, Uint256,
+    ensure, from_json, to_json_binary, Addr, Binary, Deps, DepsMut, Env, MessageInfo, Order,
+    Response, StdError, StdResult, Uint128, Uint256,
 };
 use cw2::{get_contract_version, set_contract_version};
-use cw20::{Cw20ReceiveMsg, Denom};
+use cw20::{Cw20ReceiveMsg, Denom, UncheckedDenom};
 use cw_storage_plus::Bound;
 use cw_utils::{must_pay, nonpayable, Duration, Expiration};
 use dao_interface::voting::InfoResponse;
@@ -91,8 +91,8 @@ pub fn execute(
         ExecuteMsg::FundLatest {} => execute_fund_latest_native(deps, env, info),
         ExecuteMsg::Claim { id } => execute_claim(deps, env, info, id),
         ExecuteMsg::Withdraw { id } => execute_withdraw(deps, info, env, id),
-        ExecuteMsg::UnsafeForceWithdraw { amount } => {
-            execute_unsafe_force_withdraw(deps, info, amount)
+        ExecuteMsg::UnsafeForceWithdraw { amount, denom } => {
+            execute_unsafe_force_withdraw(deps, info, amount, denom)
         }
     }
 }
@@ -600,22 +600,28 @@ fn execute_update_owner(
 fn execute_unsafe_force_withdraw(
     deps: DepsMut,
     info: MessageInfo,
-    amount: Coin,
+    amount: Uint128,
+    denom: UncheckedDenom,
 ) -> Result<Response, ContractError> {
     nonpayable(&info)?;
 
     // only the owner can initiate a force withdraw
     cw_ownable::assert_owner(deps.storage, &info.sender)?;
 
-    let send = CosmosMsg::Bank(BankMsg::Send {
-        to_address: info.sender.to_string(),
-        amount: vec![amount.clone()],
-    });
+    let checked_denom = denom.into_checked(deps.as_ref())?;
+
+    let denom_str = match &checked_denom {
+        Denom::Native(denom) => denom.to_string(),
+        Denom::Cw20(address) => address.to_string(),
+    };
+
+    let send = get_transfer_msg(info.sender, amount, checked_denom)?;
 
     Ok(Response::new()
         .add_message(send)
         .add_attribute("action", "unsafe_force_withdraw")
-        .add_attribute("amount", amount.to_string()))
+        .add_attribute("amount", amount.to_string())
+        .add_attribute("denom", denom_str))
 }
 
 #[cfg_attr(not(feature = "library"), entry_point)]

diff --git a/contracts/distribution/dao-rewards-distributor/src/msg.rs b/contracts/distribution/dao-rewards-distributor/src/msg.rs
@@ -1,5 +1,5 @@
 use cosmwasm_schema::{cw_serde, QueryResponses};
-use cosmwasm_std::{Coin, Uint128};
+use cosmwasm_std::Uint128;
 use cw20::{Cw20ReceiveMsg, Denom, UncheckedDenom};
 use cw4::MemberChangedHookMsg;
 use cw_ownable::cw_ownable_execute;
@@ -62,7 +62,12 @@ pub enum ExecuteMsg {
     Withdraw { id: u64 },
     /// forcibly withdraw funds from the contract. this is unsafe and should
     /// only be used to recover funds that are stuck in the contract.
-    UnsafeForceWithdraw { amount: Coin },
+    UnsafeForceWithdraw {
+        /// amount to withdraw
+        amount: Uint128,
+        /// denom to withdraw
+        denom: UncheckedDenom,
+    },
 }
 
 #[cw_serde]

diff --git a/contracts/distribution/dao-rewards-distributor/src/testing/suite.rs b/contracts/distribution/dao-rewards-distributor/src/testing/suite.rs
@@ -975,8 +975,11 @@ impl Suite {
             .unwrap();
     }
 
-    pub fn unsafe_force_withdraw(&mut self, amount: Coin) {
-        let msg = ExecuteMsg::UnsafeForceWithdraw { amount };
+    pub fn unsafe_force_withdraw(&mut self, amount: impl Into<Uint128>, denom: UncheckedDenom) {
+        let msg = ExecuteMsg::UnsafeForceWithdraw {
+            amount: amount.into(),
+            denom,
+        };
         self.base
             .app
             .execute_contract(
@@ -988,8 +991,15 @@ impl Suite {
             .unwrap();
     }
 
-    pub fn unsafe_force_withdraw_unauthorized(&mut self, amount: Coin) -> ContractError {
-        let msg = ExecuteMsg::UnsafeForceWithdraw { amount };
+    pub fn unsafe_force_withdraw_unauthorized(
+        &mut self,
+        amount: impl Into<Uint128>,
+        denom: UncheckedDenom,
+    ) -> ContractError {
+        let msg = ExecuteMsg::UnsafeForceWithdraw {
+            amount: amount.into(),
+            denom,
+        };
         self.base
             .app
             .execute_contract(

diff --git a/contracts/distribution/dao-rewards-distributor/src/testing/tests.rs b/contracts/distribution/dao-rewards-distributor/src/testing/tests.rs
@@ -2997,7 +2997,10 @@ fn test_unsafe_force_withdraw() {
         suite.get_balance_native(suite.distribution_contract.clone(), &suite.reward_denom);
 
     // non-owner cannot force withdraw
-    let err = suite.unsafe_force_withdraw_unauthorized(coin(100, &suite.reward_denom));
+    let err = suite.unsafe_force_withdraw_unauthorized(
+        100u128,
+        UncheckedDenom::Native(suite.reward_denom.clone()),
+    );
     assert_eq!(err, ContractError::Ownable(OwnershipError::NotOwner));
 
     let after_balance =
@@ -3009,7 +3012,7 @@ fn test_unsafe_force_withdraw() {
     assert_eq!(owner_balance, 0);
 
     // owner can force withdraw
-    suite.unsafe_force_withdraw(coin(100, &suite.reward_denom));
+    suite.unsafe_force_withdraw(100u128, UncheckedDenom::Native(suite.reward_denom.clone()));
 
     // owner has balance
     let owner_balance = suite.get_balance_native(OWNER, &suite.reward_denom);