Overflow in stake cw20

This essay discusses an integer overflow bug in the stake-cw20 contract used by Junoswap and DAO DAO. In the ~most pathological worst case scenario~, the following would need to happen to trigger this bug:

LP providers on Junoswap stake eighteen quintillion four hundred forty-six quadrillion seven hundred forty-four trillion (2^64) LP tokens (there are fewer LP tokens than regular tokens, in an exponential fashion, so for a 1 cent token you can see how this would be actually impossible amount of money). With those many tokens staked, the RAW DAO would then have to provide an equal amount (2^64) of RAW ($258B) as liquidity incentives to the pool, and then the contract would lock in a way that was entirely recoverable in one DAO governance cycle.

So, I think it is unlikely that we will see it happen and if we ever do see it happen, we'll recover. Even so, here I describe the bug, describe how to recover from the bug, and provide a fix for the bug.

The bug

To calculate the staked balance change during staking, the stake-cw20 contract computes

        staked_total
            .checked_mul(amount)
            .map_err(StdError::overflow)?
            .checked_div(balance)
            .map_err(StdError::divide_by_zero)?

after checking that balance and staked_total are non-zero. This will overflow if staked_total * amount > 2^128, and indeed a valid cw20 could induce this overflow as the only guarantee the cw20 provides is that staked_total + amount <= total_supply <= 2^128.

During unstaking, to calculate the number of tokens to return, the contract computes:

	amount
        .checked_mul(balance)
        .map_err(StdError::overflow)?
        .checked_div(staked_total)
        .map_err(StdError::divide_by_zero)?;

This too will overflow if amount * balance > 2^128, which, for the same reasons provided earlier, is a valid state for a non-malicious cw20.

balance here is defined as balance := staked_total + rewards as this contract supports autocompoudning staking rewards. As rewards being provided can induce these overflows, for a sufficiently high supply any of staking, unstaking, or funding this staking contract could cause it to lock.

Fortunately, if this bug is encountered it is entirely recoverable and will not put funds at risk.

Recovery

There are two deployments of this contract at present:

1. Junoswap liquidity pools use stake-cw20 to lock up LP tokens and provide rewards for LPs.
2. Token based DAO DAO DAOs use stake-cw20 to manage staked governance tokens.

For Junoswap, enough tokens being staked to trigger this bug will prevent LP providers from unbonding their LPs or providing new ones. To recover from the bug RAW DAO can do a CosmWasm migration of the current staking contract to code ID CODE ID WHEN READY. This is safe and simple as no state migration is required.

For token based DAOs, enough tokens being staked to trigger this will not lock the DAO, but it will stop addresses from changing their voting power until the bug is fixed. To recover from the bug, do a CosmWasm migration of your current staking contract to code ID CODE ID WHEN READY. This is safe and simple as no state migration is required.

The fix

I've patched this issue in this commit. To do so I wrote and formally verified two functions to replace the faulty logic I have reproduced below.

/// Computes the amount to add to an address' staked balance when
/// staking.
///
/// # Arguments
///
/// * `staked_total` - The number of tokens that have been staked.
/// * `balance` - The number of tokens the contract has (staked_total + rewards).
/// * `sent` - The number of tokens the user has sent to be staked.
pub(crate) fn amount_to_stake(staked_total: Uint128, balance: Uint128, sent: Uint128) -> Uint128 {
    if staked_total.is_zero() || balance.is_zero() {
        sent
    } else {
        staked_total
            .full_mul(sent)
            .div(Uint256::from(balance))
            .try_into()
            .unwrap() // balance := staked_total + rewards
                      // => balance >= staked_total
                      // => staked_total / balance <= 1
                      // => staked_total * sent / balance <= sent
                      // => we can safely unwrap here as sent fits into a u128 by construction.
    }
}

/// Computes the number of tokens to return to an address when
/// claiming.
///
/// # Arguments
///
/// * `staked_total` - The number of tokens that have been staked.
/// * `balance` - The number of tokens the contract has (staked_total + rewards).
/// * `ask` - The number of tokens being claimed.
///
/// # Invariants
///
/// These must be checked by the caller. If checked, this function is
/// guaranteed not to panic.
///
/// 1. staked_total != 0.
/// 2. ask + balance <= 2^128
///
/// For information on the panic conditions for math, see:
/// <https://rust-lang.github.io/rfcs/0560-integer-overflow.html>
pub(crate) fn amount_to_claim(staked_total: Uint128, balance: Uint128, ask: Uint128) -> Uint128 {
    // we know that:
    //
    // 1. cw20's max supply is 2^128
    // 2. balance := staked_total + rewards
    //
    // for non-malicious inputs:
    //
    // 3. 1 => ask + balance <= 2^128
    // 4. ask <= staked_total
    // 5. staked_total != 0
    // 6. 4 => ask / staked_total <= 1
    // 7. 3 => balance <= 2^128
    // 8. 6 + 7 => ask / staked_total * balance <= 2^128
    //
    // which, as addition and division are commutative, proves that
    // ask * balance / staked_total will fit into a 128 bit integer.
    ask.full_mul(balance)
        .div(Uint256::from(staked_total))
        .try_into()
        .unwrap()
}