[audit] 5. Units of duration are not distinguished (#871)

* fix cargo errors

* clarify scalar values a bit more and update comments

* added extensions that make math more readable and maintainable

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

@@ -2,7 +2,7 @@
 use cosmwasm_std::entry_point;
 use cosmwasm_std::{
     ensure, from_json, to_json_binary, Binary, Deps, DepsMut, Env, MessageInfo, Order, Response,
-    StdResult, Uint128, Uint256,
+    StdError, StdResult, Uint128, Uint256,
 };
 use cw2::{get_contract_version, set_contract_version};
 use cw20::{Cw20ReceiveMsg, Denom};
@@ -542,7 +542,8 @@ fn query_pending_rewards(
     for (id, distribution) in distributions {
         // first we get the active epoch earned puvp value
         let active_total_earned_puvp =
-            get_active_total_earned_puvp(deps, &env.block, &distribution)?;
+            get_active_total_earned_puvp(deps, &env.block, &distribution)
+                .map_err(|e| StdError::generic_err(e.to_string()))?;
 
         // then we add that to the historical rewards earned puvp
         let total_earned_puvp =

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

@@ -36,16 +36,6 @@ pub fn get_voting_power_at_block(
     Ok(resp.power)
 }
 
-/// returns underlying scalar value for a given duration.
-/// if the duration is in blocks, returns the block height.
-/// if the duration is in time, returns the time in seconds.
-pub fn get_duration_scalar(duration: &Duration) -> u64 {
-    match duration {
-        Duration::Height(h) => *h,
-        Duration::Time(t) => *t,
-    }
-}
-
 /// Returns the appropriate CosmosMsg for transferring the reward token.
 pub fn get_transfer_msg(recipient: Addr, amount: Uint128, denom: Denom) -> StdResult<CosmosMsg> {
     match denom {
@@ -73,32 +63,6 @@ pub(crate) fn scale_factor() -> Uint256 {
     Uint256::from(10u8).pow(39)
 }
 
-/// Calculate the duration from start to end. If the end is at or before the
-/// start, return 0. The first argument is end, and the second is start.
-pub fn get_exp_diff(end: &Expiration, start: &Expiration) -> StdResult<u64> {
-    match (end, start) {
-        (Expiration::AtHeight(end), Expiration::AtHeight(start)) => {
-            if end > start {
-                Ok(end - start)
-            } else {
-                Ok(0)
-            }
-        }
-        (Expiration::AtTime(end), Expiration::AtTime(start)) => {
-            if end > start {
-                Ok(end.seconds() - start.seconds())
-            } else {
-                Ok(0)
-            }
-        }
-        (Expiration::Never {}, Expiration::Never {}) => Ok(0),
-        _ => Err(StdError::generic_err(format!(
-            "incompatible expirations: got end {:?}, start {:?}",
-            end, start
-        ))),
-    }
-}
-
 pub fn validate_voting_power_contract(
     deps: &DepsMut,
     vp_contract: String,
@@ -110,3 +74,74 @@ pub fn validate_voting_power_contract(
     )?;
     Ok(vp_contract)
 }
+
+pub trait ExpirationExt {
+    /// Compute the duration since the start, flooring at 0 if the current
+    /// expiration is before the start. If either is never, or if they have
+    /// different units, returns an error as those cannot be compared.
+    fn duration_since(&self, start: &Self) -> StdResult<Duration>;
+}
+
+impl ExpirationExt for Expiration {
+    fn duration_since(&self, start: &Self) -> StdResult<Duration> {
+        match (self, start) {
+            (Expiration::AtHeight(end), Expiration::AtHeight(start)) => {
+                if end > start {
+                    Ok(Duration::Height(end - start))
+                } else {
+                    Ok(Duration::Height(0))
+                }
+            }
+            (Expiration::AtTime(end), Expiration::AtTime(start)) => {
+                if end > start {
+                    Ok(Duration::Time(end.seconds() - start.seconds()))
+                } else {
+                    Ok(Duration::Time(0))
+                }
+            }
+            (Expiration::Never {}, _) | (_, Expiration::Never {}) => {
+                Err(StdError::generic_err(format!(
+                "can't compute diff between expirations with never: got end {:?} and start {:?}",
+                self, start
+            )))
+            }
+            _ => Err(StdError::generic_err(format!(
+                "incompatible expirations: got end {:?} and start {:?}",
+                self, start
+            ))),
+        }
+    }
+}
+
+pub trait DurationExt {
+    /// Returns true if the duration is 0 blocks or 0 seconds.
+    fn is_zero(&self) -> bool;
+
+    /// Perform checked integer division between two durations, erroring if the
+    /// units do not match or denominator is 0.
+    fn checked_div(&self, denominator: &Self) -> Result<Uint128, ContractError>;
+}
+
+impl DurationExt for Duration {
+    fn is_zero(&self) -> bool {
+        match self {
+            Duration::Height(h) => *h == 0,
+            Duration::Time(t) => *t == 0,
+        }
+    }
+
+    fn checked_div(&self, denominator: &Self) -> Result<Uint128, ContractError> {
+        match (self, denominator) {
+            (Duration::Height(numerator), Duration::Height(denominator)) => {
+                Ok(Uint128::from(*numerator).checked_div(Uint128::from(*denominator))?)
+            }
+            (Duration::Time(numerator), Duration::Time(denominator)) => {
+                Ok(Uint128::from(*numerator).checked_div(Uint128::from(*denominator))?)
+            }
+            _ => Err(ContractError::Std(StdError::generic_err(format!(
+                "incompatible durations: got numerator {:?} and denominator {:?}",
+                self, denominator
+            )))),
+        }
+    }
+}

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

@@ -2,8 +2,8 @@ use cosmwasm_std::{Addr, BlockInfo, Deps, DepsMut, Env, StdResult, Uint128, Uint
 
 use crate::{
     helpers::{
-        get_duration_scalar, get_exp_diff, get_prev_block_total_vp, get_voting_power_at_block,
-        scale_factor,
+        get_prev_block_total_vp, get_voting_power_at_block, scale_factor, DurationExt,
+        ExpirationExt,
     },
     state::{DistributionState, EmissionRate, UserRewardState, DISTRIBUTIONS, USER_REWARDS},
     ContractError,
@@ -82,7 +82,7 @@ pub fn get_active_total_earned_puvp(
     deps: Deps,
     block: &BlockInfo,
     distribution: &DistributionState,
-) -> StdResult<Uint256> {
+) -> Result<Uint256, ContractError> {
     match distribution.active_epoch.emission_rate {
         EmissionRate::Paused {} => Ok(Uint256::zero()),
         // this is updated manually during funding, so just return it here.
@@ -98,11 +98,8 @@ pub fn get_active_total_earned_puvp(
             // get the duration from the last time rewards were updated to the
             // last time rewards were distributed. this will be 0 if the rewards
             // were updated at or after the last time rewards were distributed.
-            let new_reward_distribution_duration: Uint128 = get_exp_diff(
-                &last_time_rewards_distributed,
-                &distribution.active_epoch.last_updated_total_earned_puvp,
-            )?
-            .into();
+            let new_reward_distribution_duration = last_time_rewards_distributed
+                .duration_since(&distribution.active_epoch.last_updated_total_earned_puvp)?;
 
             // no need to query total voting power and do math if distribution
             // is already up to date.
@@ -118,8 +115,8 @@ pub fn get_active_total_earned_puvp(
             } else {
                 // count intervals of the rewards emission that have passed
                 // since the last update which need to be distributed
-                let complete_distribution_periods = new_reward_distribution_duration
-                    .checked_div(get_duration_scalar(&duration).into())?;
+                let complete_distribution_periods =
+                    new_reward_distribution_duration.checked_div(&duration)?;
 
                 // It is impossible for this to overflow as total rewards can
                 // never exceed max value of Uint128 as total tokens in

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_duration_scalar, get_exp_diff, get_prev_block_total_vp, scale_factor},
+    helpers::{get_prev_block_total_vp, scale_factor, DurationExt, ExpirationExt},
     rewards::get_active_total_earned_puvp,
     ContractError,
 };
@@ -70,12 +70,12 @@ impl EmissionRate {
             EmissionRate::Linear {
                 amount, duration, ..
             } => {
-                if *amount == Uint128::zero() {
+                if amount.is_zero() {
                     return Err(ContractError::InvalidEmissionRateFieldZero {
                         field: "amount".to_string(),
                     });
                 }
-                if get_duration_scalar(duration) == 0 {
+                if duration.is_zero() {
                     return Err(ContractError::InvalidEmissionRateFieldZero {
                         field: "duration".to_string(),
                     });
@@ -229,24 +229,23 @@ impl DistributionState {
 
     /// get the total rewards to be distributed based on the active epoch's
     /// emission rate
-    pub fn get_total_rewards(&self) -> StdResult<Uint128> {
+    pub fn get_total_rewards(&self) -> Result<Uint128, ContractError> {
         match self.active_epoch.emission_rate {
             EmissionRate::Paused {} => Ok(Uint128::zero()),
             EmissionRate::Immediate {} => Ok(self.funded_amount),
             EmissionRate::Linear {
                 amount, duration, ..
             } => {
-                let epoch_duration =
-                    get_exp_diff(&self.active_epoch.ends_at, &self.active_epoch.started_at)?;
+                let epoch_duration = self
+                    .active_epoch
+                    .ends_at
+                    .duration_since(&self.active_epoch.started_at)?;
 
-                let emission_rate_duration_scalar = match duration {
-                    Duration::Height(h) => h,
-                    Duration::Time(t) => t,
-                };
+                // count total intervals of the rewards emission that will pass
+                // based on the start and end times.
+                let complete_distribution_periods = epoch_duration.checked_div(&duration)?;
 
-                amount
-                    .checked_multiply_ratio(epoch_duration, emission_rate_duration_scalar)
-                    .map_err(|e| StdError::generic_err(e.to_string()))
+                Ok(amount.checked_mul(complete_distribution_periods)?)
             }
         }
     }