Update withdrawing permission check

Replace Staking::Withdraw with Distribution::*
CosmWasm · Apr 21, 2021 · 5a745c8 · 5a745c8
1 parent 9daea73
commit 5a745c8
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 4 deletions.
diff --git a/contracts/cw1-subkeys/src/contract.rs b/contracts/cw1-subkeys/src/contract.rs
@@ -5,8 +5,8 @@ use std::ops::{AddAssign, Sub};
 #[cfg(not(feature = "library"))]
 use cosmwasm_std::entry_point;
 use cosmwasm_std::{
-    attr, to_binary, BankMsg, Binary, Coin, CosmosMsg, Deps, DepsMut, Empty, Env, MessageInfo,
-    Order, Response, StakingMsg, StdError, StdResult,
+    attr, to_binary, BankMsg, Binary, Coin, CosmosMsg, Deps, DepsMut, DistributionMsg, Empty, Env,
+    MessageInfo, Order, Response, StakingMsg, StdError, StdResult,
 };
 use cw0::Expiration;
 use cw1::CanExecuteResponse;
@@ -100,6 +100,11 @@ where
                     let perm = perm.ok_or(ContractError::NotAllowed {})?;
                     check_staking_permissions(staking_msg, perm)?;
                 }
+                CosmosMsg::Distribution(distribution_msg) => {
+                    let perm = PERMISSIONS.may_load(deps.storage, &info.sender)?;
+                    let perm = perm.ok_or(ContractError::NotAllowed {})?;
+                    check_distribution_permissions(distribution_msg, perm)?;
+                }
                 CosmosMsg::Bank(BankMsg::Send {
                     to_address: _,
                     amount,
@@ -147,12 +152,27 @@ pub fn check_staking_permissions(
                 return Err(ContractError::ReDelegatePerm {});
             }
         }
-        StakingMsg::Withdraw { .. } => {
+        s => panic!("Unsupported staking message: {:?}", s),
+    }
+    Ok(true)
+}
+
+pub fn check_distribution_permissions(
+    distribution_msg: &DistributionMsg,
+    permissions: Permissions,
+) -> Result<bool, ContractError> {
+    match distribution_msg {
+        DistributionMsg::SetWithdrawAddress { .. } => {
+            if !permissions.withdraw {
+                return Err(ContractError::WithdrawAddrPerm {});
+            }
+        }
+        DistributionMsg::WithdrawDelegatorReward { .. } => {
             if !permissions.withdraw {
                 return Err(ContractError::WithdrawPerm {});
             }
         }
-        s => panic!("Unsupported staking message: {:?}", s),
+        s => panic!("Unsupported distribution message: {:?}", s),
     }
     Ok(true)
 }

diff --git a/contracts/cw1-subkeys/src/error.rs b/contracts/cw1-subkeys/src/error.rs
@@ -32,6 +32,9 @@ pub enum ContractError {
 
     #[error("Withdraw is not allowed")]
     WithdrawPerm {},
+
+    #[error("Set withdraw address is not allowed")]
+    WithdrawAddrPerm {},
 }
 
 impl From<cw1_whitelist::ContractError> for ContractError {