From 19531cdc0dbb34e63f76d7be75afc0e320b183fd Mon Sep 17 00:00:00 2001
From: Mauro Lacy <mauro@babylonlabs.io>
Date: Wed, 18 Dec 2024 08:35:13 +0100
Subject: [PATCH] Fix: Rewards distribution amount

---
 contracts/btc-finality/src/contract.rs |  3 ++-
 contracts/btc-finality/src/finality.rs | 18 +++++++++++++-----
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/contracts/btc-finality/src/contract.rs b/contracts/btc-finality/src/contract.rs
index b15ff58..2e8e667 100644
--- a/contracts/btc-finality/src/contract.rs
+++ b/contracts/btc-finality/src/contract.rs
@@ -50,6 +50,7 @@ pub fn instantiate(
     let params = msg.params.unwrap_or_default();
     PARAMS.save(deps.storage, &params)?;
     // initialize storage, so no issue when reading for the first time
+    TOTAL_REWARDS.save(deps.storage, &Uint128::zero())?;
 
     set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;
     Ok(Response::new().add_attribute("action", "instantiate"))
@@ -261,7 +262,7 @@ fn handle_end_block(
     // On an epoch boundary, send rewards to Babylon through the babylon contract
     let params = PARAMS.load(deps.storage)?;
     if env.block.height > 0 && env.block.height % params.epoch_length == 0 {
-        let rewards = TOTAL_REWARDS.may_load(deps.storage)?.unwrap_or_default();
+        let rewards = TOTAL_REWARDS.load(deps.storage)?;
         if rewards.u128() > 0 {
             let wasm_msg = send_rewards_msg(deps, rewards.u128(), &cfg)?;
             res = res.add_message(wasm_msg);
diff --git a/contracts/btc-finality/src/finality.rs b/contracts/btc-finality/src/finality.rs
index 2cf6ecf..7214b52 100644
--- a/contracts/btc-finality/src/finality.rs
+++ b/contracts/btc-finality/src/finality.rs
@@ -632,14 +632,20 @@ pub fn distribute_rewards(deps: &mut DepsMut, env: &Env) -> Result<(), ContractE
     };
     // Get the voting power of the active FPS
     let total_voting_power = active_fps.iter().map(|fp| fp.power as u128).sum::<u128>();
-    // Get the rewards to distribute (bank balance of the finality contract)
+    // Get the rewards to distribute (bank balance of the finality contract minus already distributed rewards)
+    let distributed_rewards = TOTAL_REWARDS.load(deps.storage)?;
     let cfg = CONFIG.load(deps.storage)?;
     let rewards_amount = deps
         .querier
         .query_balance(env.contract.address.clone(), cfg.denom)?
-        .amount;
+        .amount
+        .saturating_sub(distributed_rewards);
+    // Short-circuit if there are no rewards to distribute
+    if rewards_amount.is_zero() {
+        return Ok(());
+    }
     // Compute the rewards for each active FP
-    let mut total_rewards = Uint128::zero();
+    let mut accumulated_rewards = Uint128::zero();
     for fp in active_fps {
         let reward = (Decimal::from_ratio(fp.power as u128, total_voting_power)
             * Decimal::from_ratio(rewards_amount, 1u128))
@@ -649,9 +655,11 @@ pub fn distribute_rewards(deps: &mut DepsMut, env: &Env) -> Result<(), ContractE
             Ok::<Uint128, ContractError>(r.unwrap_or_default() + reward)
         })?;
         // Compute the total rewards
-        total_rewards += reward;
+        accumulated_rewards += reward;
     }
     // Update the total rewards
-    TOTAL_REWARDS.save(deps.storage, &total_rewards)?;
+    TOTAL_REWARDS.update(deps.storage, |r| {
+        Ok::<Uint128, ContractError>(r + accumulated_rewards)
+    })?;
     Ok(())
 }