add ability to fund latest distribution

contracts/distribution/dao-rewards-distributor/src/contract.rs

@@ -1,8 +1,8 @@
 #[cfg(not(feature = "library"))]
 use cosmwasm_std::entry_point;
 use cosmwasm_std::{
-    ensure, from_json, to_json_binary, Binary, 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};
@@ -68,7 +68,7 @@ pub fn execute(
         ExecuteMsg::MemberChangedHook(msg) => execute_membership_changed(deps, env, info, msg),
         ExecuteMsg::UpdateOwnership(action) => execute_update_owner(deps, info, env, action),
         ExecuteMsg::Receive(msg) => execute_receive_cw20(deps, env, info, msg),
-        ExecuteMsg::Create(create_msg) => execute_create(deps, env, info, create_msg),
+        ExecuteMsg::Create(create_msg) => execute_create_native(deps, env, info, create_msg),
         ExecuteMsg::Update {
             id,
             emission_rate,
@@ -86,6 +86,7 @@ pub fn execute(
             withdraw_destination,
         ),
         ExecuteMsg::Fund(FundMsg { id }) => execute_fund_native(deps, env, info, id),
+        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),
     }
@@ -119,22 +120,68 @@ fn execute_receive_cw20(
                 }
             };
 
+            execute_fund(deps, env, distribution, wrapper.amount)
+        }
+        ReceiveCw20Msg::FundLatest {} => {
+            let id = COUNT.load(deps.storage)?;
+            let distribution = DISTRIBUTIONS
+                .load(deps.storage, id)
+                .map_err(|_| ContractError::DistributionNotFound { id })?;
+
+            match &distribution.denom {
+                Denom::Native(_) => return Err(ContractError::InvalidFunds {}),
+                Denom::Cw20(addr) => {
+                    // ensure funding is coming from the cw20 we are currently
+                    // distributing
+                    if addr != info.sender {
+                        return Err(ContractError::InvalidCw20 {});
+                    }
+                }
+            };
+
             execute_fund(deps, env, distribution, wrapper.amount)
         }
     }
 }
 
+fn execute_create_native(
+    deps: DepsMut,
+    env: Env,
+    info: MessageInfo,
+    msg: CreateMsg,
+) -> Result<Response, ContractError> {
+    let checked_denom = msg.denom.clone().into_checked(deps.as_ref())?;
+
+    // if native funds provided, ensure they are for this denom. if other native
+    // funds present, return error. if no funds, do nothing and leave registered
+    // denom with no funding, to be funded later.
+    let initial_funds = if !info.funds.is_empty() {
+        match &checked_denom {
+            Denom::Native(denom) => {
+                // ensures there is exactly 1 coin passed that matches the denom
+                Some(must_pay(&info, denom)?)
+            }
+            Denom::Cw20(_) => return Err(ContractError::NoFundsOnCw20Create {}),
+        }
+    } else {
+        None
+    };
+
+    execute_create(deps, env, info.sender, msg, initial_funds)
+}
+
 /// creates a new rewards distribution. only the owner can do this. if funds
 /// provided when creating a native token distribution, will start distributing
 /// rewards immediately.
 fn execute_create(
     deps: DepsMut,
     env: Env,
-    info: MessageInfo,
+    sender: Addr,
     msg: CreateMsg,
+    initial_funds: Option<Uint128>,
 ) -> Result<Response, ContractError> {
     // only the owner can create a new distribution
-    cw_ownable::assert_owner(deps.storage, &info.sender)?;
+    cw_ownable::assert_owner(deps.storage, &sender)?;
 
     // update count and use as the new distribution's ID
     let id = COUNT.update(deps.storage, |count| -> StdResult<u64> { Ok(count + 1) })?;
@@ -147,7 +194,7 @@ fn execute_create(
         // if withdraw destination is specified, we validate it
         Some(addr) => deps.api.addr_validate(&addr)?,
         // otherwise default to the owner
-        None => info.sender.clone(),
+        None => sender.clone(),
     };
 
     msg.emission_rate.validate()?;
@@ -185,20 +232,10 @@ fn execute_create(
         .add_attribute("id", id.to_string())
         .add_attribute("denom", distribution.get_denom_string());
 
-    // if native funds provided, ensure they are for this denom. if other native
-    // funds present, return error. if no funds, do nothing and leave registered
-    // denom with no funding, to be funded later.
-    if !info.funds.is_empty() {
-        match &distribution.denom {
-            Denom::Native(denom) => {
-                // ensures there is exactly 1 coin passed that matches the denom
-                let amount = must_pay(&info, denom)?;
-
-                execute_fund(deps, env, distribution, amount)?;
-
-                response = response.add_attribute("amount_funded", amount);
-            }
-            Denom::Cw20(_) => return Err(ContractError::NoFundsOnCw20Create {}),
+    if let Some(initial_funds) = initial_funds {
+        if !initial_funds.is_zero() {
+            execute_fund(deps, env, distribution, initial_funds)?;
+            response = response.add_attribute("amount_funded", initial_funds);
         }
     }
 
@@ -259,6 +296,15 @@ fn execute_update(
         .add_attribute("denom", distribution.get_denom_string()))
 }
 
+fn execute_fund_latest_native(
+    deps: DepsMut,
+    env: Env,
+    info: MessageInfo,
+) -> Result<Response, ContractError> {
+    let id = COUNT.load(deps.storage)?;
+    execute_fund_native(deps, env, info, id)
+}
+
 fn execute_fund_native(
     deps: DepsMut,
     env: Env,

contracts/distribution/dao-rewards-distributor/src/msg.rs

@@ -50,6 +50,8 @@ pub enum ExecuteMsg {
     Receive(Cw20ReceiveMsg),
     /// Used to fund this contract with native tokens.
     Fund(FundMsg),
+    /// Used to fund the latest distribution with native tokens.
+    FundLatest {},
     /// Claims rewards for the sender.
     Claim { id: u64 },
     /// withdraws the undistributed rewards for a distribution. members can
@@ -83,6 +85,8 @@ pub struct FundMsg {
 pub enum ReceiveCw20Msg {
     /// Used to fund this contract with cw20 tokens.
     Fund(FundMsg),
+    /// Used to fund the latest distribution with cw20 tokens.
+    FundLatest {},
 }
 
 #[cw_serde]

contracts/distribution/dao-rewards-distributor/src/testing/suite.rs

@@ -634,6 +634,19 @@ impl Suite {
             .unwrap();
     }
 
+    pub fn fund_latest_native(&mut self, coin: Coin) {
+        self.mint_native(coin.clone(), OWNER);
+        self.app
+            .borrow_mut()
+            .execute_contract(
+                Addr::unchecked(OWNER),
+                self.distribution_contract.clone(),
+                &ExecuteMsg::FundLatest {},
+                &[coin],
+            )
+            .unwrap();
+    }
+
     pub fn fund_cw20(&mut self, id: u64, coin: Cw20Coin) {
         let fund_sub_msg = to_json_binary(&ReceiveCw20Msg::Fund(FundMsg { id })).unwrap();
         self.app
@@ -650,6 +663,22 @@ impl Suite {
             .unwrap();
     }
 
+    pub fn fund_latest_cw20(&mut self, coin: Cw20Coin) {
+        let fund_sub_msg = to_json_binary(&ReceiveCw20Msg::FundLatest {}).unwrap();
+        self.app
+            .execute_contract(
+                Addr::unchecked(OWNER),
+                Addr::unchecked(coin.address),
+                &cw20::Cw20ExecuteMsg::Send {
+                    contract: self.distribution_contract.to_string(),
+                    amount: coin.amount,
+                    msg: fund_sub_msg,
+                },
+                &[],
+            )
+            .unwrap();
+    }
+
     pub fn skip_blocks(&mut self, blocks: u64) {
         self.app.borrow_mut().update_block(|b| {
             println!("skipping blocks {:?} -> {:?}", b.height, b.height + blocks);

contracts/distribution/dao-rewards-distributor/src/testing/tests.rs

@@ -2584,6 +2584,55 @@ fn test_large_stake_before_claim() {
     suite.claim_rewards(ADDR3, 1);
 }
 
+#[test]
+fn test_fund_latest_native() {
+    let mut suite = SuiteBuilder::base(super::suite::DaoType::Native).build();
+
+    suite.assert_amount(1_000);
+    suite.assert_ends_at(Expiration::AtHeight(1_000_000));
+    suite.assert_duration(10);
+
+    // double duration by 1_000_000 blocks
+    suite.fund_latest_native(coin(100_000_000, DENOM));
+
+    // skip all of the time
+    suite.skip_blocks(2_000_000);
+
+    suite.assert_pending_rewards(ADDR1, 1, 100_000_000);
+    suite.assert_pending_rewards(ADDR2, 1, 50_000_000);
+    suite.assert_pending_rewards(ADDR3, 1, 50_000_000);
+}
+
+#[test]
+fn test_fund_latest_cw20() {
+    let mut suite = SuiteBuilder::base(super::suite::DaoType::CW20)
+        .with_rewards_config(RewardsConfig {
+            amount: 1_000,
+            denom: UncheckedDenom::Cw20(DENOM.to_string()),
+            duration: Duration::Height(10),
+            destination: None,
+            continuous: true,
+        })
+        .build();
+
+    suite.assert_amount(1_000);
+    suite.assert_ends_at(Expiration::AtHeight(1_000_000));
+    suite.assert_duration(10);
+
+    // double duration by 1_000_000 blocks
+    suite.fund_latest_cw20(Cw20Coin {
+        address: suite.reward_denom.clone(),
+        amount: Uint128::new(100_000_000),
+    });
+
+    // skip all of the time
+    suite.skip_blocks(2_000_000);
+
+    suite.assert_pending_rewards(ADDR1, 1, 100_000_000);
+    suite.assert_pending_rewards(ADDR2, 1, 50_000_000);
+    suite.assert_pending_rewards(ADDR3, 1, 50_000_000);
+}
+
 #[test]
 fn test_migrate() {
     let mut deps = mock_dependencies();