From 9b520acf5a0e1036edee32159c2acebd39f280a1 Mon Sep 17 00:00:00 2001 From: Hakeem Kazeem Date: Wed, 24 Jun 2026 10:20:44 +0100 Subject: [PATCH] add overflow checks to update_reward --- contract/contracts/src/storage.rs | 28 +++++++++++++++++++++------- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/contract/contracts/src/storage.rs b/contract/contracts/src/storage.rs index f2fc19cb..b61cf5ad 100644 --- a/contract/contracts/src/storage.rs +++ b/contract/contracts/src/storage.rs @@ -230,19 +230,33 @@ pub fn update_reward(env: &Env, user: Option<&Address>) { let config = read_config(env); let reward_token = token::Client::new(env, &config.reward_token); let staking_token = token::Client::new(env, &config.staking_token); - let total_supply = staking_token.total_supply(); + + // TODO: I'm not sure why total_supply is used, the method does not exist. + // let total_supply = staking_token.total_supply(); + // if total_supply > 0 { + // let reward_per_token = (config.reward_rate * PRECISION) / total_supply; + + let total_staked = staking_token.balance(&env.current_contract_address()); + if total_staked > 0 { + + let numerator = config.reward_rate.checked_mul(PRECISION).expect("reward_rate * PRECISION overflow"); + let reward_per_token = numerator.checked_div(total_staked).expect("division by zero in reward_per_token"); - if total_supply > 0 { - let reward_per_token = (config.reward_rate * PRECISION) / total_supply; let mut reward_per_token_stored = read_reward_per_token_stored(env); - reward_per_token_stored += reward_per_token; + reward_per_token_stored = reward_per_token_stored.checked_add(reward_per_token).expect("reward_per_token_stored overflow"); write_reward_per_token_stored(env, reward_per_token_stored); if let Some(user_addr) = user { if let Some(mut user_info) = read_user_info(env, user_addr) { - let rewards = (user_info.shares * reward_per_token_stored) / PRECISION - - user_info.reward_per_token_paid; - user_info.rewards += rewards; + + let user_info_reward_per_token = user_info.shares.checked_mul(reward_per_token_stored).expect("shares * reward_per_token_stored overflow"); + + let user_info_reward_per_token_div = user_info_reward_per_token.checked_div(PRECISION).expect("PRECISION is zero"); + + let reward_delta = user_info_reward_per_token_div.checked_sub(user_info.reward_per_token_paid).expect("reward delta underflow — stored < paid"); + + user_info.rewards = user_info.rewards.checked_add(reward_delta).expect("user rewards overflow"); + user_info.reward_per_token_paid = reward_per_token_stored; write_user_info(env, user_addr, &user_info); }