[audit] 2. Incorrect reward distribution due to inconsistent voting power queries

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

@@ -10,13 +10,13 @@ use dao_interface::voting::{
 
 use crate::ContractError;
 
-pub fn get_prev_block_total_vp(
+pub fn get_total_voting_power_at_block(
     deps: Deps,
     block: &BlockInfo,
     contract_addr: &Addr,
 ) -> StdResult<Uint128> {
     let msg = VotingQueryMsg::TotalPowerAtHeight {
-        height: Some(block.height.checked_sub(1).unwrap_or_default()),
+        height: Some(block.height),
     };
     let resp: TotalPowerAtHeightResponse = deps.querier.query_wasm_smart(contract_addr, &msg)?;
     Ok(resp.power)

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

@@ -2,7 +2,7 @@ use cosmwasm_std::{Addr, BlockInfo, Deps, DepsMut, Env, StdResult, Uint128, Uint
 
 use crate::{
     helpers::{
-        get_prev_block_total_vp, get_voting_power_at_block, scale_factor, DurationExt,
+        get_total_voting_power_at_block, get_voting_power_at_block, scale_factor, DurationExt,
         ExpirationExt,
     },
     state::{DistributionState, EmissionRate, UserRewardState, DISTRIBUTIONS, USER_REWARDS},
@@ -107,10 +107,11 @@ pub fn get_active_total_earned_puvp(
                 return Ok(curr);
             }
 
-            let prev_total_power = get_prev_block_total_vp(deps, block, &distribution.vp_contract)?;
+            let total_power =
+                get_total_voting_power_at_block(deps, block, &distribution.vp_contract)?;
 
             // if no voting power is registered, no one should receive rewards.
-            if prev_total_power.is_zero() {
+            if total_power.is_zero() {
                 Ok(curr)
             } else {
                 // count intervals of the rewards emission that have passed
@@ -128,8 +129,7 @@ pub fn get_active_total_earned_puvp(
 
                 // the new rewards per unit voting power that have been
                 // distributed since the last update
-                let new_rewards_puvp =
-                    new_rewards_distributed.checked_div(prev_total_power.into())?;
+                let new_rewards_puvp = new_rewards_distributed.checked_div(total_power.into())?;
                 Ok(curr.checked_add(new_rewards_puvp)?)
             }
         }

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

@@ -9,7 +9,7 @@ use cw_utils::Duration;
 use std::{cmp::min, collections::HashMap};
 
 use crate::{
-    helpers::{get_prev_block_total_vp, scale_factor, DurationExt, ExpirationExt},
+    helpers::{get_total_voting_power_at_block, scale_factor, DurationExt, ExpirationExt},
     rewards::get_active_total_earned_puvp,
     ContractError,
 };
@@ -366,18 +366,18 @@ impl DistributionState {
 
         let curr = self.active_epoch.total_earned_puvp;
 
-        let prev_total_power = get_prev_block_total_vp(deps, block, &self.vp_contract)?;
+        let total_power = get_total_voting_power_at_block(deps, block, &self.vp_contract)?;
 
         // if no voting power is registered, error since rewards can't be
         // distributed.
-        if prev_total_power.is_zero() {
+        if total_power.is_zero() {
             Err(ContractError::NoVotingPowerNoRewards {})
         } else {
             // the new rewards per unit voting power based on the funded amount
             let new_rewards_puvp = Uint256::from(funded_amount_delta)
                 // this can never overflow since funded_amount is a Uint128
                 .checked_mul(scale_factor())?
-                .checked_div(prev_total_power.into())?;
+                .checked_div(total_power.into())?;
 
             self.active_epoch.total_earned_puvp = curr.checked_add(new_rewards_puvp)?;
 

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

@@ -2447,3 +2447,37 @@ fn test_fund_while_paused() {
     suite.assert_started_at(Expiration::AtHeight(200_000));
     suite.assert_ends_at(Expiration::AtHeight(1_000_000 + 100_000 + 1_000_000));
 }
+
+#[test]
+fn test_large_stake_before_claim() {
+    let mut suite = SuiteBuilder::base(super::suite::DaoType::Native)
+        .with_rewards_config(RewardsConfig {
+            amount: 3_000,
+            denom: UncheckedDenom::Native(DENOM.to_string()),
+            duration: Duration::Height(1),
+            destination: None,
+            continuous: true,
+        })
+        .build();
+
+    suite.assert_amount(3_000);
+    suite.assert_ends_at(Expiration::AtHeight(33_333));
+    suite.assert_duration(1);
+
+    // ADDR1 stake big amount of tokens
+    suite.skip_blocks(33_000);
+    suite.mint_native(coin(10_000, &suite.reward_denom), ADDR1);
+    suite.stake_native_tokens(ADDR1, 10_000);
+
+    // ADD1 claims rewards in the next block
+    suite.skip_blocks(1);
+    suite.claim_rewards(ADDR1, 1);
+
+    // skip to end
+    suite.skip_blocks(100_000_000);
+
+    // all users should be able to claim rewards
+    suite.claim_rewards(ADDR1, 1);
+    suite.claim_rewards(ADDR2, 1);
+    suite.claim_rewards(ADDR3, 1);
+}