Skip to content

feat: add overflow safety for all math operations #313

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Mkalbani wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: add overflow safety for all math operations #313

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 11 additions & 7 deletions contracts/liquid_staking/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,9 @@ impl LiquidStaking {
.instance()
.get(&DataKey::RewardPerTokenStored)
.unwrap_or(0);
rpt += amount * PRECISION / total_staked;
rpt = rpt.checked_add(
amount.checked_mul(PRECISION).expect("rpt overflow") / total_staked
).expect("rpt overflow");
env.storage()
.instance()
.set(&DataKey::RewardPerTokenStored, &rpt);
Expand Down Expand Up @@ -116,15 +118,15 @@ impl LiquidStaking {
);

env.storage().persistent().set(&DataKey::StakeAmount(token_id), &amount);
let lock_time = env.ledger().timestamp() + lock_duration;
let lock_time = env.ledger().timestamp().checked_add(lock_duration).expect("lock time overflow");
env.storage().persistent().set(&DataKey::StakeLockTime(token_id), &lock_time);

let rpt: i128 = env.storage().instance().get(&DataKey::RewardPerTokenStored).unwrap_or(0);
env.storage().persistent().set(&DataKey::NftRewardPerTokenPaid(token_id), &rpt);
env.storage().persistent().set(&DataKey::NftRewards(token_id), &0_i128);

let total: i128 = env.storage().instance().get(&DataKey::TotalStaked).unwrap_or(0);
env.storage().instance().set(&DataKey::TotalStaked, &(total + amount));
env.storage().instance().set(&DataKey::TotalStaked, &total.checked_add(amount).expect("total staked overflow"));

env.events().publish((symbol_short!("staked"), user), (token_id, amount, lock_time));

Expand Down Expand Up @@ -161,7 +163,7 @@ impl LiquidStaking {
}

let total: i128 = env.storage().instance().get(&DataKey::TotalStaked).unwrap_or(0);
env.storage().instance().set(&DataKey::TotalStaked, &(total - amount));
env.storage().instance().set(&DataKey::TotalStaked, &total.checked_sub(amount).expect("total staked underflow"));

let stake_token: Address = env.storage().instance().get(&DataKey::StakeToken).unwrap();
token::Client::new(&env, &stake_token).transfer(
Expand Down Expand Up @@ -222,7 +224,9 @@ impl LiquidStaking {
let amount: i128 = env.storage().persistent().get(&DataKey::StakeAmount(token_id)).unwrap_or(0);
let accrued: i128 = env.storage().persistent().get(&DataKey::NftRewards(token_id)).unwrap_or(0);

let pending = accrued + amount * (rpt - nft_rpt) / PRECISION;
let pending = accrued.checked_add(
amount.checked_mul(rpt - nft_rpt).expect("rewards overflow") / PRECISION
).expect("rewards overflow");
let lock_time: u64 = env.storage().persistent().get(&DataKey::StakeLockTime(token_id)).unwrap_or(0);

StakeInfo {
Expand All @@ -237,11 +241,11 @@ impl LiquidStaking {
let rpt: i128 = env.storage().instance().get(&DataKey::RewardPerTokenStored).unwrap_or(0);
let nft_rpt: i128 = env.storage().persistent().get(&DataKey::NftRewardPerTokenPaid(token_id)).unwrap_or(0);
let amount: i128 = env.storage().persistent().get(&DataKey::StakeAmount(token_id)).unwrap_or(0);
let earned = amount * (rpt - nft_rpt) / PRECISION;
let earned = amount.checked_mul(rpt - nft_rpt).expect("rewards overflow") / PRECISION;

if earned > 0 {
let prev: i128 = env.storage().persistent().get(&DataKey::NftRewards(token_id)).unwrap_or(0);
env.storage().persistent().set(&DataKey::NftRewards(token_id), &(prev + earned));
env.storage().persistent().set(&DataKey::NftRewards(token_id), &prev.checked_add(earned).expect("rewards overflow"));
}

env.storage().persistent().set(&DataKey::NftRewardPerTokenPaid(token_id), &rpt);
Expand Down
7 changes: 4 additions & 3 deletions contracts/nft_metadata/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -223,7 +223,7 @@ impl NftMetadataContract {
.instance()
.get(&DataKey::TokenCounter)
.unwrap_or(0);
let token_id = counter + 1;
let token_id = counter.checked_add(1).expect("token id overflow");

let metadata = NftMetadata {
token_id,
Expand Down Expand Up @@ -285,7 +285,7 @@ impl NftMetadataContract {
.instance()
.get(&DataKey::TokenCounter)
.unwrap_or(0);
let token_id = counter + 1;
let token_id = counter.checked_add(1).expect("token id overflow");

let mut final_metadata = metadata;
final_metadata.token_id = token_id;
Expand Down Expand Up @@ -741,7 +741,8 @@ impl NftMetadataContract {
.get(&DataKey::NftMetadata(token_id))
.expect("token not found");

let royalty_amount = (sale_price * metadata.royalty_percentage as i128) / 10000;
let royalty_amount = sale_price
.checked_mul(metadata.royalty_percentage as i128).expect("royalty overflow") / 10000;

(metadata.royalty_recipient, royalty_amount)
}
Expand Down
18 changes: 10 additions & 8 deletions contracts/staking/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,9 @@ impl MultiTokenStaking {
.get(&DataKey::RewardPerTokenStored(reward_token.clone()))
.unwrap_or(0);

rpt += amount * PRECISION / total_staked;
rpt = rpt.checked_add(
amount.checked_mul(PRECISION).expect("rpt overflow") / total_staked
).expect("rpt overflow");
env.storage()
.instance()
.set(&DataKey::RewardPerTokenStored(reward_token), &rpt);
Expand Down Expand Up @@ -147,7 +149,7 @@ impl MultiTokenStaking {
let prev: i128 = Self::_stake_of(&env, &user);
env.storage()
.persistent()
.set(&DataKey::Stake(user.clone()), &(prev + amount));
.set(&DataKey::Stake(user.clone()), &prev.checked_add(amount).expect("stake overflow"));

let total: i128 = env
.storage()
Expand All @@ -156,7 +158,7 @@ impl MultiTokenStaking {
.unwrap_or(0);
env.storage()
.instance()
.set(&DataKey::TotalStaked, &(total + amount));
.set(&DataKey::TotalStaked, &total.checked_add(amount).expect("total staked overflow"));

env.events().publish((symbol_short!("staked"), user), amount);
}
Expand All @@ -172,7 +174,7 @@ impl MultiTokenStaking {

env.storage()
.persistent()
.set(&DataKey::Stake(user.clone()), &(prev - amount));
.set(&DataKey::Stake(user.clone()), &prev.checked_sub(amount).expect("stake underflow"));

let total: i128 = env
.storage()
Expand All @@ -181,7 +183,7 @@ impl MultiTokenStaking {
.unwrap_or(0);
env.storage()
.instance()
.set(&DataKey::TotalStaked, &(total - amount));
.set(&DataKey::TotalStaked, &total.checked_sub(amount).expect("total staked underflow"));

let stake_token: Address = env.storage().instance().get(&DataKey::StakeToken).unwrap();
token::Client::new(&env, &stake_token).transfer(
Expand Down Expand Up @@ -276,7 +278,7 @@ impl MultiTokenStaking {
.get(&DataKey::Rewards(user, reward_token))
.unwrap_or(0);

accrued + stake * (rpt - user_rpt) / PRECISION
accrued + stake.checked_mul(rpt - user_rpt).expect("rewards overflow") / PRECISION
}

pub fn total_staked(env: Env) -> i128 {
Expand Down Expand Up @@ -322,7 +324,7 @@ impl MultiTokenStaking {
.unwrap_or(0);

let stake = Self::_stake_of(env, user);
let earned = stake * (rpt - user_rpt) / PRECISION;
let earned = stake.checked_mul(rpt - user_rpt).expect("rewards overflow") / PRECISION;

if earned > 0 {
let prev: i128 = env
Expand All @@ -332,7 +334,7 @@ impl MultiTokenStaking {
.unwrap_or(0);
env.storage()
.persistent()
.set(&DataKey::Rewards(user.clone(), reward_token.clone()), &(prev + earned));
.set(&DataKey::Rewards(user.clone(), reward_token.clone()), &prev.checked_add(earned).expect("rewards overflow"));
}

// Snapshot current global rate for this user
Expand Down
18 changes: 10 additions & 8 deletions contracts/yield/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,9 @@ impl YieldDistribution {
.get(&DataKey::RewardPerTokenStored)
.unwrap_or(0);
// Δ reward_per_token = amount * PRECISION / total_staked
rpt += amount * PRECISION / total_staked;
rpt = rpt.checked_add(
amount.checked_mul(PRECISION).expect("rpt overflow") / total_staked
).expect("rpt overflow");
env.storage()
.instance()
.set(&DataKey::RewardPerTokenStored, &rpt);
Expand Down Expand Up @@ -154,7 +156,7 @@ impl YieldDistribution {
let prev: i128 = Self::_stake_of(&env, &user);
env.storage()
.persistent()
.set(&DataKey::Stake(user.clone()), &(prev + amount));
.set(&DataKey::Stake(user.clone()), &prev.checked_add(amount).expect("stake overflow"));

let total: i128 = env
.storage()
Expand All @@ -163,7 +165,7 @@ impl YieldDistribution {
.unwrap_or(0);
env.storage()
.instance()
.set(&DataKey::TotalStaked, &(total + amount));
.set(&DataKey::TotalStaked, &total.checked_add(amount).expect("total staked overflow"));

env.events()
.publish((symbol_short!("staked"), user), amount);
Expand All @@ -189,7 +191,7 @@ impl YieldDistribution {

env.storage()
.persistent()
.set(&DataKey::Stake(user.clone()), &(prev - amount));
.set(&DataKey::Stake(user.clone()), &prev.checked_sub(amount).expect("stake underflow"));

let total: i128 = env
.storage()
Expand All @@ -198,7 +200,7 @@ impl YieldDistribution {
.unwrap_or(0);
env.storage()
.instance()
.set(&DataKey::TotalStaked, &(total - amount));
.set(&DataKey::TotalStaked, &total.checked_sub(amount).expect("total staked underflow"));

let stake_token: Address = env.storage().instance().get(&DataKey::StakeToken).unwrap();
token::Client::new(&env, &stake_token).transfer(
Expand Down Expand Up @@ -273,7 +275,7 @@ impl YieldDistribution {
.get(&DataKey::Rewards(user))
.unwrap_or(0);

accrued + stake * (rpt - user_rpt) / PRECISION
accrued + stake.checked_mul(rpt - user_rpt).expect("rewards overflow") / PRECISION
}

pub fn total_staked(env: Env) -> i128 {
Expand Down Expand Up @@ -306,7 +308,7 @@ impl YieldDistribution {
.unwrap_or(0);

let stake = Self::_stake_of(env, user);
let earned = stake * (rpt - user_rpt) / PRECISION;
let earned = stake.checked_mul(rpt - user_rpt).expect("rewards overflow") / PRECISION;

if earned > 0 {
let prev: i128 = env
Expand All @@ -316,7 +318,7 @@ impl YieldDistribution {
.unwrap_or(0);
env.storage()
.persistent()
.set(&DataKey::Rewards(user.clone()), &(prev + earned));
.set(&DataKey::Rewards(user.clone()), &prev.checked_add(earned).expect("rewards overflow"));
}

// Snapshot current global rate for this user
Expand Down
40 changes: 20 additions & 20 deletions src/amm/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,11 +71,11 @@ impl AMM {
// Calculate shares to mint
let shares = if total_shares == 0 {
// Initial liquidity = geometric mean
sqrt(amount_a * amount_b)
sqrt(amount_a.checked_mul(amount_b).expect("deposit overflow"))
} else {
// Proportional liquidity: min(amount_a/reserve_a, amount_b/reserve_b) * total_shares
let shares_a = (amount_a * total_shares) / reserve_a;
let shares_b = (amount_b * total_shares) / reserve_b;
let shares_a = amount_a.checked_mul(total_shares).expect("deposit overflow") / reserve_a;
let shares_b = amount_b.checked_mul(total_shares).expect("deposit overflow") / reserve_b;
if shares_a < shares_b {
shares_a
} else {
Expand Down Expand Up @@ -106,13 +106,13 @@ impl AMM {
// Update state
env.storage()
.instance()
.set(&DataKey::ReserveA, &(reserve_a + amount_a));
.set(&DataKey::ReserveA, &reserve_a.checked_add(amount_a).expect("reserve overflow"));
env.storage()
.instance()
.set(&DataKey::ReserveB, &(reserve_b + amount_b));
.set(&DataKey::ReserveB, &reserve_b.checked_add(amount_b).expect("reserve overflow"));
env.storage()
.instance()
.set(&DataKey::TotalShares, &(total_shares + shares));
.set(&DataKey::TotalShares, &total_shares.checked_add(shares).expect("shares overflow"));

let old_shares: i128 = env
.storage()
Expand All @@ -121,7 +121,7 @@ impl AMM {
.unwrap_or(0);
env.storage()
.persistent()
.set(&DataKey::Shares(from.clone()), &(old_shares + shares));
.set(&DataKey::Shares(from.clone()), &old_shares.checked_add(shares).expect("shares overflow"));

env.events().publish(
(symbol_short!("deposit"), from),
Expand Down Expand Up @@ -171,9 +171,9 @@ impl AMM {
);

// Constant product formula with 0.3% fee: dy = (reserve_out * dx * 997) / (reserve_in * 1000 + dx * 997)
let amount_in_with_fee = amount_in * 997;
let numerator = amount_in_with_fee * reserve_out;
let denominator = (reserve_in * 1000) + amount_in_with_fee;
let amount_in_with_fee = amount_in.checked_mul(997).expect("swap overflow");
let numerator = amount_in_with_fee.checked_mul(reserve_out).expect("swap overflow");
let denominator = reserve_in.checked_mul(1000).expect("swap overflow").checked_add(amount_in_with_fee).expect("swap overflow");
let amount_out = numerator / denominator;

if amount_out < min_amount_out {
Expand All @@ -182,11 +182,11 @@ impl AMM {

// Update state
if token_in == token_a {
reserve_a += amount_in;
reserve_b -= amount_out;
reserve_a = reserve_a.checked_add(amount_in).expect("reserve overflow");
reserve_b = reserve_b.checked_sub(amount_out).expect("reserve underflow");
} else {
reserve_b += amount_in;
reserve_a -= amount_out;
reserve_b = reserve_b.checked_add(amount_in).expect("reserve overflow");
reserve_a = reserve_a.checked_sub(amount_out).expect("reserve underflow");
}

env.storage().instance().set(&DataKey::ReserveA, &reserve_a);
Expand Down Expand Up @@ -233,22 +233,22 @@ impl AMM {
panic!("insufficient shares");
}

let amount_a = (shares * reserve_a) / total_shares;
let amount_b = (shares * reserve_b) / total_shares;
let amount_a = shares.checked_mul(reserve_a).expect("withdraw overflow") / total_shares;
let amount_b = shares.checked_mul(reserve_b).expect("withdraw overflow") / total_shares;

// Update state
env.storage()
.instance()
.set(&DataKey::ReserveA, &(reserve_a - amount_a));
.set(&DataKey::ReserveA, &reserve_a.checked_sub(amount_a).expect("reserve underflow"));
env.storage()
.instance()
.set(&DataKey::ReserveB, &(reserve_b - amount_b));
.set(&DataKey::ReserveB, &reserve_b.checked_sub(amount_b).expect("reserve underflow"));
env.storage()
.instance()
.set(&DataKey::TotalShares, &(total_shares - shares));
.set(&DataKey::TotalShares, &total_shares.checked_sub(shares).expect("shares underflow"));
env.storage()
.persistent()
.set(&DataKey::Shares(from.clone()), &(user_shares - shares));
.set(&DataKey::Shares(from.clone()), &user_shares.checked_sub(shares).expect("shares underflow"));

// Transfer tokens back to user
transfer(
Expand Down
2 changes: 1 addition & 1 deletion src/batch/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ impl BatchExecutor {
.unwrap_or(0);
env.storage()
.instance()
.set(&DataKey::Nonce(caller.clone()), &(current_nonce + 1));
.set(&DataKey::Nonce(caller.clone()), &current_nonce.checked_add(1).expect("nonce overflow"));

// Execute all calls atomically
let mut results = Vec::new(&env);
Expand Down
4 changes: 2 additions & 2 deletions src/circuit_breaker/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -272,7 +272,7 @@ impl CircuitBreaker {
.get(&DataKey::TimelockSeconds)
.unwrap_or(DEFAULT_TIMELOCK_SECONDS);

let unlocks_at = env.ledger().timestamp() + timelock;
let unlocks_at = env.ledger().timestamp().checked_add(timelock).expect("timelock overflow");

env.storage()
.instance()
Expand Down Expand Up @@ -465,7 +465,7 @@ impl CircuitBreaker {
.unwrap_or(0);
env.storage()
.instance()
.set(&DataKey::TripCount, &(count + 1));
.set(&DataKey::TripCount, &count.checked_add(1).expect("trip count overflow"));

env.events()
.publish((symbol_short!("tripped"), tier), (caller, source));
Expand Down
Loading