Incorrect higher mint price for the first token minted in each generation starting from the second generation

[howlbot-integration]

Lines of code

https://github.com/code-423n4/2024-07-traitforge/blob/279b2887e3d38bc219a05d332cbcb0655b2dc644/contracts/TraitForgeNft/TraitForgeNft.sol#L190
https://github.com/code-423n4/2024-07-traitforge/blob/279b2887e3d38bc219a05d332cbcb0655b2dc644/contracts/TraitForgeNft/TraitForgeNft.sol#L193
https://github.com/code-423n4/2024-07-traitforge/blob/279b2887e3d38bc219a05d332cbcb0655b2dc644/contracts/TraitForgeNft/TraitForgeNft.sol#L282

Vulnerability details

Summary

Users will need to pay a higher price to mint the first token for each generation starting from generaton 2, instead of the 0.005 ether starting price. This is because the TraitForgeNft.currentGeneration variable isn't incremented on the 10_000 th mint of each generation. It is incremented on the 10_001 th mint (basically, the first mint of the next generation). So on the first mint of the next generation, the TraitForgeNft contract still assumes to be in the last generation and calculates a wrong mint price which is higher than 0.005 ether. Only then is the TraitForgeNft.currentGeneration variable incremented. This is applicable to both TraitForgeNft::mintToken() and TraitForgeNft::mintWithBudget() functions.

Vulnerability Details

Let's see two code segments, one from TraitForgeNft::mintToken() and the other from TraitForgeNft::_mintInternal() function,

    function mintToken(bytes32[] calldata proof)
        public
        payable
        whenNotPaused
        nonReentrant
        onlyWhitelisted(proof, keccak256(abi.encodePacked(msg.sender)))
    {
@>      uint256 mintPrice = calculateMintPrice();
        require(msg.value >= mintPrice, "Insufficient ETH send for minting.");

@>      _mintInternal(msg.sender, mintPrice);

        // more code here
    }

The price for minting is calculated first.

    function _mintInternal(address to, uint256 mintPrice) internal {
        if (generationMintCounts[currentGeneration] >= maxTokensPerGen) {
@>          _incrementGeneration();
        }

        _tokenIds++;
        uint256 newItemId = _tokenIds;
        _mint(to, newItemId);
        uint256 entropyValue = entropyGenerator.getNextEntropy();

        tokenCreationTimestamps[newItemId] = block.timestamp;
        tokenEntropy[newItemId] = entropyValue;
        tokenGenerations[newItemId] = currentGeneration;
@>      generationMintCounts[currentGeneration]++;
        initialOwners[newItemId] = to;

        // more code here
    }

Links:

https://github.com/code-423n4/2024-07-traitforge/blob/279b2887e3d38bc219a05d332cbcb0655b2dc644/contracts/TraitForgeNft/TraitForgeNft.sol#L190

https://github.com/code-423n4/2024-07-traitforge/blob/279b2887e3d38bc219a05d332cbcb0655b2dc644/contracts/TraitForgeNft/TraitForgeNft.sol#L193

https://github.com/code-423n4/2024-07-traitforge/blob/279b2887e3d38bc219a05d332cbcb0655b2dc644/contracts/TraitForgeNft/TraitForgeNft.sol#L282.

The current generation is incremented after price calculation for the mint. Thus, while moving from one generation to the next, the contract still thinks it's in the last generation during the price calculation phase, and gives an incorrect high price when it should be 0.005 ether instead. Then the TraitForgeNft.currentGeneration value is incremented and the contract functions as expected from the second mint of the next generation.

This issue is applicable to TraitForgeNft::mintWithBudget() function as well. However, it currently has a bug which doesn't allow minting beyond the first generation.

    function mintWithBudget(bytes32[] calldata proof)
        public
        payable
        whenNotPaused
        nonReentrant
        onlyWhitelisted(proof, keccak256(abi.encodePacked(msg.sender)))
    {
        uint256 mintPrice = calculateMintPrice();
        
        // more code here

@>      while (budgetLeft >= mintPrice && _tokenIds < maxTokensPerGen) {
            _mintInternal(msg.sender, mintPrice);
            amountMinted++;
            budgetLeft -= mintPrice;
            mintPrice = calculateMintPrice();
        }
        
        // more code here
    }

Link:

https://github.com/code-423n4/2024-07-traitforge/blob/279b2887e3d38bc219a05d332cbcb0655b2dc644/contracts/TraitForgeNft/TraitForgeNft.sol#L215

If the number of token ids minted goes beyond TraitForgeNft.maxTokensPerGen, the function will not mint further, however, this is a finding on its own. If this was to be resolved, the function would also be susceptible to the same problem while minting the first token in each generation starting from generation 2.

Impact

Users will need to pay an incorrect/high amount for minting the first token for each generation starting from generation 2 instead of the 0.005 ether starting price. This essentially causes a fund loss for users (paying higher instead of 0.005 ether). Also, the amount to pay for the first token of each generation continues to become larger and larger as the generations progress (due to the TraitForgeNft.priceIncrement variable increasing by TraitForgeNft.priceIncrementByGen each generation). However, the fund loss is only limited to 9 mints upto the 10th generation.

Proof of Concept

To write the PoC, some changes had to be made to the TraitForgeNft contract to stay within the block gas limit and avoid a revert due to another issue highlighted below.

The following changes were made to write the PoC for this finding.

Changed the TraitForgeNft.maxTokensPerGen variable from 1e4 to 1e3, to avoid OutOfGas error.

Commented out the following line from TraitForgeNft::_incrementGeneration() function because it causes a revert, and is a completely different finding on its own. Here, the TraitForgeNft contract tries to re-initialize the alpha indices in the EntropyGenerator contract -- the slots which hold the God entropy value of 999999. However, it will always revert because TraitForgeNft isn't the owner of EntropyGenerator contract and EntropyGenerator::initializeAlphaIndices() function has an onlyOwner modifier.

    function _incrementGeneration() private {
        require(generationMintCounts[currentGeneration] >= maxTokensPerGen, "Generation limit not yet reached");
        currentGeneration++;
        generationMintCounts[currentGeneration] = 0;
        priceIncrement = priceIncrement + priceIncrementByGen;
@>      entropyGenerator.initializeAlphaIndices();
        emit GenerationIncremented(currentGeneration);
    }

We will be using the following setup to write PoCs in Foundry,

// SPDX-License-Identifier: MIT
pragma solidity 0.8.24;

import {console} from "lib/forge-std/src/console.sol";
import {Test} from "lib/forge-std/src/Test.sol";

import {Trait} from "../contracts/Trait/Trait.sol";
import {TraitForgeNft} from "../contracts/TraitForgeNft/TraitForgeNft.sol";
import {EntropyGenerator} from "../contracts/EntropyGenerator/EntropyGenerator.sol";
import {EntityTrading} from "../contracts/EntityTrading/EntityTrading.sol";
import {EntityForging} from "../contracts/EntityForging/EntityForging.sol";
import {DevFund} from "../contracts/DevFund/DevFund.sol";
import {Airdrop} from "../contracts/Airdrop/Airdrop.sol";
import {DAOFund} from "../contracts/DAOFund/DAOFund.sol";
import {NukeFund} from "../contracts/NukeFund/NukeFund.sol";

contract POCs is Test {
    address public owner;
    address public user1;

    address public dev1;
    address public dev2;
    address public dev3;

    Trait public traitToken;
    TraitForgeNft public traitForgeNft;
    EntropyGenerator public entropyGenerator;
    EntityTrading public entityTrading;
    EntityForging public entityForging;
    DevFund public devFund;
    Airdrop public airdrop;
    DAOFund public daoFund;
    NukeFund public nukeFund;

    function setUp() public {
        owner = makeAddr("owner");
        user1 = makeAddr("user1");
        dev1 = makeAddr("dev1");
        dev2 = makeAddr("dev2");
        dev3 = makeAddr("dev3");

        string memory name = "Trait";
        string memory symbol = "TRT";
        uint8 decimals = 18;
        uint256 tokenTotalSupply = 10_000e18;
        address uniswapV3Router = makeAddr("uniswap v3 router");

        vm.startPrank(owner);
        traitToken = new Trait(name, symbol, decimals, tokenTotalSupply);
        traitForgeNft = new TraitForgeNft();
        entropyGenerator = new EntropyGenerator(address(traitForgeNft));
        entityTrading = new EntityTrading(address(traitForgeNft));
        entityForging = new EntityForging(address(traitForgeNft));
        devFund = new DevFund();
        airdrop = new Airdrop();
        daoFund = new DAOFund(address(traitToken), uniswapV3Router);
        nukeFund = new NukeFund(address(traitForgeNft), address(airdrop), payable(devFund), payable(daoFund));

        traitForgeNft.setEntropyGenerator(address(entropyGenerator));
        traitForgeNft.setEntityForgingContract(address(entityForging));
        traitForgeNft.setNukeFundContract(payable(nukeFund));
        traitForgeNft.setAirdropContract(address(airdrop));
        entropyGenerator.setAllowedCaller(address(traitForgeNft));
        entityTrading.setNukeFundAddress(payable(nukeFund));
        entityForging.setNukeFundAddress(payable(nukeFund));
        airdrop.setTraitToken(address(traitToken));
        airdrop.transferOwnership(address(traitForgeNft));

        uint256 weights = 100;
        devFund.addDev(dev1, weights);
        devFund.addDev(dev2, weights);
        devFund.addDev(dev3, weights);

        entropyGenerator.writeEntropyBatch1();
        entropyGenerator.writeEntropyBatch2();
        entropyGenerator.writeEntropyBatch3();
        vm.stopPrank();
    }

    // PoCs go here
}

Now run the following test case,

    function testIncorrectPriceCalculatedForFirstMintOfEachGenerationStartingfromGeneration2() public {
        // Provide some ether to user1
        uint256 maxEthToDeal = 1000 ether;
        vm.deal(user1, maxEthToDeal);

        uint256 startPriceForEachGen = traitForgeNft.startPrice();
        // Changed this value to 1_000 Nfts in the TraitForgeNft contract to stay within the block gas limit
        uint256 nftsToMint = traitForgeNft.maxTokensPerGen();
        // To avoid the whitelist, we will warp by 1 day
        uint256 warpBy = 1 days + 1;
        vm.warp(block.timestamp + warpBy);
        bytes32[] memory proof = new bytes32[](0);
        uint256 count = 0;

        vm.startPrank(user1);
        // Mint 1k Nfts
        // This allows us to move from the first generation to the second generation
        while (count < nftsToMint) {
            // `TraitForgeNft::mintToken()` and `TraitForgeNft::mintWithBudget()` both use the
            // `TraitForgeNft::calculateMintPrice()` function to calculate the price for the next mint
            uint256 ethToSupply = traitForgeNft.calculateMintPrice();
            traitForgeNft.mintToken{value: ethToSupply}(proof);
            ++count;
        }
        vm.stopPrank();

        uint256 expectedPriceForNextMint = startPriceForEachGen;
        uint256 userEthBalanceBefore = user1.balance;
        uint256 userEthBalanceAfter;
        // This is the mint price for the first mint in the second generation
        // It is expected to be 0.005 ether
        uint256 actualEthPriceToSend = traitForgeNft.calculateMintPrice();

        vm.startPrank(user1);
        traitForgeNft.mintToken{value: actualEthPriceToSend}(proof);
        vm.stopPrank();

        userEthBalanceAfter = user1.balance;
        // This is going to be 5000000000000000 (5e15/0.005 ether)
        console.log("Expected mint price for the first token in generation 2: ", expectedPriceForNextMint);
        // This is going to be 29500000000000000 (29.5e15/ 0.0295 ether)
        console.log("Actual mint price for the first token in generation 2: ", actualEthPriceToSend);
        // However, remember that this is for 1000 tokens per generation and the second generation first token mint
        // For 10k tokens per generation and higher generations (with higher price increment), the value will be larger
        // Calculating some:
        // Mint price for first token of second generation: 250000000000000000 (0.25 ether)
        // Mint price for first token of third generation: 300000000000000000 (0.3 ether)
        // And so on

        assert(actualEthPriceToSend > expectedPriceForNextMint);
        assert((userEthBalanceBefore - userEthBalanceAfter) > expectedPriceForNextMint);
    }

The test passes with the following logs,

Ran 1 test for test/POCs.sol:POCs
[PASS] testWrongPriceCalculatedForFirstMintOfEachGenStartingfromGen2() (gas: 280943740)
Logs:
  Expected mint price for the first token in generation 2:  5000000000000000
  Actual mint price for the first token in generation 2:  29500000000000000

Suite result: ok. 1 passed; 0 failed; 0 skipped; finished in 94.33ms (90.61ms CPU time)

Ran 1 test suite in 102.50ms (94.33ms CPU time): 1 tests passed, 0 failed, 0 skipped (1 total tests)

Tools Used

Manual review. Foundry for writing the PoC.

Recommended Mitigation

Increment the generation before calculating the price.

Make the following changes in TraitForgeNft::mintToken() function,

    function mintToken(bytes32[] calldata proof)
        public
        payable
        whenNotPaused
        nonReentrant
        onlyWhitelisted(proof, keccak256(abi.encodePacked(msg.sender)))
    {
+       if (generationMintCounts[currentGeneration] >= maxTokensPerGen) {
+           _incrementGeneration();
+       }
        uint256 mintPrice = calculateMintPrice();
        require(msg.value >= mintPrice, "Insufficient ETH send for minting.");

        // more code here
    }

Remember, the finding is also valid for TraitForgeNft::mintWithBudget() function, so changes need to be made there as well.

    function mintWithBudget(bytes32[] calldata proof)
        public
        payable
        whenNotPaused
        nonReentrant
        onlyWhitelisted(proof, keccak256(abi.encodePacked(msg.sender)))
    {
-       uint256 mintPrice = calculateMintPrice();
+       uint256 mintPrice;
        uint256 amountMinted = 0;
        uint256 budgetLeft = msg.value;

        while (budgetLeft >= mintPrice && _tokenIds < maxTokensPerGen) {
+           if (generationMintCounts[currentGeneration] >= maxTokensPerGen) {
+              _incrementGeneration();
+           }
+           mintPrice = calculateMintPrice();
            _mintInternal(msg.sender, mintPrice);
            amountMinted++;
            budgetLeft -= mintPrice;
-           mintPrice = calculateMintPrice();
        }
        if (budgetLeft > 0) {
            (bool refundSuccess,) = msg.sender.call{value: budgetLeft}("");
            require(refundSuccess, "Refund failed.");
        }
    }

Let's unit test the TraitForgeNft::mintToken() function (in Foundry),

    function testMitigation() public {
        // Provide some ether to user1
        uint256 maxEthToDeal = 1000 ether;
        vm.deal(user1, maxEthToDeal);

        uint256 startPriceForEachGen = traitForgeNft.startPrice();
        // Changed this value to 1_000 Nfts in the TraitForgeNft contract to stay within the block gas limit
        uint256 nftsToMint = traitForgeNft.maxTokensPerGen();
        // To avoid the whitelist, we will warp by 1 day
        uint256 warpBy = 1 days + 1;
        vm.warp(block.timestamp + warpBy);
        bytes32[] memory proof = new bytes32[](0);
        uint256 count = 0;

        vm.startPrank(user1);
        // Mint 1k Nfts
        // This allows us to move from the first generation to the second generation
        while (count < nftsToMint) {
            // `TraitForgeNft::mintToken()` and `TraitForgeNft::mintWithBudget()` both use the
            // `TraitForgeNft::calculateMintPrice()` function to calculate the price for the next mint
            uint256 ethToSupply = traitForgeNft.calculateMintPrice();
            traitForgeNft.mintToken{value: ethToSupply}(proof);
            ++count;
        }
        vm.stopPrank();

        uint256 expectedPriceForNextMint = startPriceForEachGen;
        uint256 userEthBalanceBefore = user1.balance;
        uint256 userEthBalanceAfter;

        vm.startPrank(user1);
        traitForgeNft.mintToken{value: expectedPriceForNextMint}(proof);
        vm.stopPrank();

        userEthBalanceAfter = user1.balance;
        // Both the expected and actual mint price will be 0.005 ether
        console.log("Expected mint price for the first token in generation 2: ", expectedPriceForNextMint);
        console.log(
            "Actual mint price for the first token in generation 2: ", (userEthBalanceBefore - userEthBalanceAfter)
        );

        assertEq((userEthBalanceBefore - userEthBalanceAfter), expectedPriceForNextMint);
    }

The unit test passes with the following logs,

Ran 1 test for test/POCs.sol:POCs
[PASS] testMitigation() (gas: 281335207)
Logs:
  Expected mint price for the first token in generation 2:  5000000000000000
  Actual mint price for the first token in generation 2:  5000000000000000

Suite result: ok. 1 passed; 0 failed; 0 skipped; finished in 131.10ms (122.20ms CPU time)

Ran 1 test suite in 141.59ms (131.10ms CPU time): 1 tests passed, 0 failed, 0 skipped (1 total tests)

Assessed type

Other

[c4-judge]

koolexcrypto marked the issue as duplicate of #564