Repository: DA0-DA0/dao-contracts
Pull request #281 (2022-05-03 02:24:29)

Add withdraw method to stake-cw20-reward-distributor
This new execute message allows the owner to reclaim the reward funds. This would be used by a DAO that wants to change something significant about how rewards are being distributed such as lowering by a large amount or change the contract.


Patch:
diff --git a/contracts/stake-cw20-reward-distributor/schema/execute_msg.json b/contracts/stake-cw20-reward-distributor/schema/execute_msg.json
index 021e5d10..63597c3f 100644
--- a/contracts/stake-cw20-reward-distributor/schema/execute_msg.json
+++ b/contracts/stake-cw20-reward-distributor/schema/execute_msg.json
@@ -45,6 +45,18 @@
         }
       },
       "additionalProperties": false
+    },
+    {
+      "type": "object",
+      "required": [
+        "withdraw"
+      ],
+      "properties": {
+        "withdraw": {
+          "type": "object"
+        }
+      },
+      "additionalProperties": false
     }
   ],
   "definitions": {
diff --git a/contracts/stake-cw20-reward-distributor/src/contract.rs b/contracts/stake-cw20-reward-distributor/src/contract.rs
index 5f0b73db..b01a817b 100644
--- a/contracts/stake-cw20-reward-distributor/src/contract.rs
+++ b/contracts/stake-cw20-reward-distributor/src/contract.rs
@@ -62,6 +62,7 @@ pub fn execute(
             reward_token,
         } => execute_update_config(deps, info, owner, staking_addr, reward_rate, reward_token),
         ExecuteMsg::Distribute {} => execute_distribute(deps, env),
+        ExecuteMsg::Withdraw {} => execute_withdraw(deps, info, env),
     }
 }
 
@@ -153,6 +154,37 @@ pub fn execute_distribute(deps: DepsMut, env: Env) -> Result<Response, ContractE
     Ok(Response::default().add_message(send_msg))
 }
 
+pub fn execute_withdraw(
+    deps: DepsMut,
+    info: MessageInfo,
+    env: Env,
+) -> Result<Response, ContractError> {
+    let config = CONFIG.load(deps.storage)?;
+    if config.owner != info.sender {
+        return Err(ContractError::Unauthorized {});
+    }
+
+    let balance_info: cw20::BalanceResponse = deps.querier.query_wasm_smart(
+        config.reward_token.clone(),
+        &cw20::Cw20QueryMsg::Balance {
+            address: env.contract.address.to_string(),
+        },
+    )?;
+
+    let msg = to_binary(&cw20::Cw20ExecuteMsg::Transfer {
+        recipient: config.owner.into(),
+        amount: balance_info.balance,
+    })?;
+    let send_msg: CosmosMsg = WasmMsg::Execute {
+        contract_addr: config.reward_token.into(),
+        msg,
+        funds: vec![],
+    }
+    .into();
+
+    Ok(Response::new().add_message(send_msg))
+}
+
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
     match msg {
diff --git a/contracts/stake-cw20-reward-distributor/src/msg.rs b/contracts/stake-cw20-reward-distributor/src/msg.rs
index 4e8e7e88..8b757158 100644
--- a/contracts/stake-cw20-reward-distributor/src/msg.rs
+++ b/contracts/stake-cw20-reward-distributor/src/msg.rs
@@ -21,6 +21,7 @@ pub enum ExecuteMsg {
         reward_token: String,
     },
     Distribute {},
+    Withdraw {},
 }
 
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]


Test patch:
diff --git a/contracts/stake-cw20-reward-distributor/src/tests.rs b/contracts/stake-cw20-reward-distributor/src/tests.rs
index 738ff9cd..5ae7a07c 100644
--- a/contracts/stake-cw20-reward-distributor/src/tests.rs
+++ b/contracts/stake-cw20-reward-distributor/src/tests.rs
@@ -369,3 +369,72 @@ fn test_update_config_invalid_addrs() {
         .unwrap();
     assert_eq!(err, ContractError::InvalidStakingContract {});
 }
+
+#[test]
+fn test_withdraw() {
+    let mut app = App::default();
+
+    let cw20_addr = instantiate_cw20(
+        &mut app,
+        vec![cw20::Cw20Coin {
+            address: OWNER.to_string(),
+            amount: Uint128::from(1000u64),
+        }],
+    );
+    let staking_addr = instantiate_staking(&mut app, cw20_addr.clone());
+
+    let msg = InstantiateMsg {
+        owner: OWNER.to_string(),
+        staking_addr: staking_addr.to_string(),
+        reward_rate: Uint128::new(1),
+        reward_token: cw20_addr.to_string(),
+    };
+    let distributor_addr = instantiate_distributor(&mut app, msg);
+
+    let msg = cw20::Cw20ExecuteMsg::Transfer {
+        recipient: distributor_addr.to_string(),
+        amount: Uint128::from(1000u128),
+    };
+    app.execute_contract(Addr::unchecked(OWNER), cw20_addr.clone(), &msg, &[])
+        .unwrap();
+
+    app.update_block(|mut block| block.height += 10);
+    app.execute_contract(
+        Addr::unchecked(OWNER),
+        distributor_addr.clone(),
+        &ExecuteMsg::Distribute {},
+        &[],
+    )
+    .unwrap();
+
+    let staking_balance = get_balance_cw20(&app, cw20_addr.clone(), staking_addr);
+    assert_eq!(staking_balance, Uint128::new(10));
+
+    let distributor_info = get_info(&app, distributor_addr.clone());
+    assert_eq!(distributor_info.balance, Uint128::new(990));
+    assert_eq!(distributor_info.last_payment_block, app.block_info().height);
+
+    // Unauthorized user cannot withdraw funds
+    let err = app
+        .execute_contract(
+            Addr::unchecked(MANAGER),
+            distributor_addr.clone(),
+            &ExecuteMsg::Withdraw {},
+            &[],
+        )
+        .unwrap_err();
+
+    assert_eq!(ContractError::Unauthorized {}, err.downcast().unwrap());
+
+    // Withdraw funds
+    app.execute_contract(
+        Addr::unchecked(OWNER),
+        distributor_addr,
+        &ExecuteMsg::Withdraw {},
+        &[],
+    )
+    .unwrap();
+
+    let owner_balance = get_balance_cw20(&app, cw20_addr, Addr::unchecked(OWNER));
+    assert_eq!(owner_balance, Uint128::new(990));
+}