Aztec TGE Contract Audit

Description. In the Aztec network, the locked AZTEC tokens in ATP can participate in staking. This is achieved by allowing the ATP contract to approve the locked token to the staker contract. It is crucial that the staking operation cannot be used to bypass the unlock schedule. In the contract, the ATP should be restricted to only use the staker contract that is whitelisted by the Registry. Unfortunately, due to an oversight, the ATP contract can use the upgradeStaker function to bypass the restriction and upgrade the staker contract to an arbitrary contract.

function upgradeStaker(StakerVersion _version, bytes memory _initdata)
    external
    override(IMATPCore)
    onlyBeneficiary
{
    address impl = REGISTRY.getMilestoneStakerImplementation(_version);
    UUPSUpgradeable(address(staker)).upgradeToAndCall(impl, _initdata);

    emit StakerUpgraded(_version);
}

In the upgradeStaker function above, the caller specifies the staker _version and _initdata. The staker implementation is then fetched from the global Registry. Then it will perform a UUPS upgrade for the staker and call the function specified in the _initdata. The _initdata is intended for initialization operations. However, there is no restriction on the _initdata, allowing the ATP to call arbitrary functions with arbitrary parameters during the upgrade. This includes calling the upgradeToAndCall function again.

To perform the attack, the ATP contract will encode the upgradeToAndCall function in the _initdata. Calling the upgradeStaker function then leads to two upgrades. The first upgrade has the implementation address checked. However, the second upgrade is embedded in upgradeToAndCall and is not checked. As a result, the ATP can upgrade the staker to an arbitrary implementation.

Below is a PoC of the attack:

// Add this test in test/btt/atps/milestone/upgradeStaker/upgradeStaker.t.sol

function test_arbitraryUpgrade() external {
    address initialStaker = registry.getMilestoneStakerImplementation(StakerVersion.wrap(0));
    FakeMilestoneStaker badStaker = new FakeMilestoneStaker(staking);
    // encode the UUPSUpgradeable.upgradeToAndCall function with badStaker in the _initdata
    atp.upgradeStaker(StakerVersion.wrap(0), abi.encodeCall(UUPSUpgradeable.upgradeToAndCall, (address(badStaker), "")));

    assertEq(atp.getStaker().getImplementation(), address(badStaker));
    assertNotEq(address(badStaker), initialStaker);
}

Impact. The ATP beneficiary can exploit this vulnerability to upgrade the staker to malicious implementation and then withdraw the locked token from the staker.

Recommendation. It is implicit that the ATP contract can call any function of the staker contract though the upgradeStaker function. To mitigate this, it is recommended to add a restriction to the _initdata in the function. For example, only allow calling the initialize function in the _initdata.

Description. The LockLib contains helper functions to handle the schedule lock. The createLock function creates a lock based on the provided parameters:

function createLock(LockParams memory _params, uint256 _allocation) internal pure returns (Lock memory) {
    LockLib.assertValid(_params);
    return Lock({
        startTime: _params.startTime,
        cliff: _params.startTime + _params.cliffDuration,
        endTime: _params.startTime + _params.lockDuration,
        allocation: _allocation
    });
}

The cliff time and end time are calculated as the sum of the start time and the duration. The sum could overflow the range of uint256 if the time and duration are large numbers.

Additionally, in the claim function of ATP contract, the total claimed amount is updated by adding the current claimed amount. This could cause overflow if the amount is very large and the beneficiary cyclically claims and transfers tokens (since additional transferred token can be claimed).

function claim() external override(IMATPCore) onlyBeneficiary returns (uint256) {
    uint256 amount = getClaimable();
    require(amount > 0, NoClaimable());

    claimed += amount;

    TOKEN.safeTransfer(msg.sender, amount);

    emit Claimed(amount);
    return amount;
}

Impact. This is unlikely to cause issues in practice, as such overflows are improbable in real-world scenarios.

Recommendation. It is recommended to use SafeMath for arithmetic operations to prevent potential overflow and underflow.

Description. Besides the token to be distributed, the ATP contract may also hold additional transferred tokens. For example, users may transfer their own AZTEC tokens to the ATP contract directly (by mistake). The contract supports that the beneficiary can claim these tokens only after the global lock ends. As described in the documentation “Any additional tokens transferred to the ATP after its creation cannot be retrieved until after the global lock ends.”. However, it turns out that, in the LATP contract, the beneficiary can still retrieve the additional transferred AZTEC tokens before the global lock ends.

In LATP, the claimable amount is calculated as follows:

function getClaimable() public view override(IATPCore) returns (uint256) {
    Lock memory globalLock = getGlobalLock();
    uint256 unlocked = globalLock.hasEnded(block.timestamp)
        ? type(uint256).max
        : (globalLock.unlockedAt(block.timestamp) - claimed);

    return Math.min(TOKEN.balanceOf(address(this)) - getRevokableAmount(), unlocked);
}

When the accumulated amount is lower than the global unlock amount, the claimable amount is calculated as TOKEN.balanceOf(address(this)) - getRevokableAmount(). This calculation includes additional transferred AZTEC tokens.

Impact. This behavior does not cause critical issues but is inconsistent with the documentation.

Recommendation. It is recommended to align the documentation with the contract behavior.