ERC4626 Vault Operations

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.13;

import {Assertion} from "credible-std/Assertion.sol";
import {PhEvm} from "credible-std/PhEvm.sol";

contract ERC4626DepositWithdrawAssertion is Assertion {
    // Used to store the last known good share value
    uint256 private lastKnownShareValue;

    // Tolerance for precision/rounding errors (0.01%)
    uint256 private constant PRECISION_TOLERANCE = 1e14; // 0.01% of 1e18

    function triggers() external view override {
        // Register triggers for deposit operations
        registerCallTrigger(this.assertionDepositAssets.selector, IERC4626.deposit.selector);
        registerCallTrigger(this.assertionDepositShares.selector, IERC4626.deposit.selector);

        // Register triggers for withdraw operations
        registerCallTrigger(this.assertionWithdrawAssets.selector, IERC4626.withdraw.selector);
        registerCallTrigger(this.assertionWithdrawShares.selector, IERC4626.withdraw.selector);

        // Register trigger for share value monotonicity
        // This can be triggered by various operations
        registerCallTrigger(this.assertionShareValueMonotonicity.selector, IERC4626.deposit.selector);
        registerCallTrigger(this.assertionShareValueMonotonicity.selector, IERC4626.withdraw.selector);
    }

    // Assert that deposits correctly update total assets
    function assertionDepositAssets() external {
        // Get the assertion adopter address
        IERC4626 adopter = IERC4626(ph.getAssertionAdopter());

        // Get all deposit calls to the ERC4626 vault
        PhEvm.CallInputs[] memory callInputs = ph.getCallInputs(address(adopter), adopter.deposit.selector);

        // First, do a simple pre/post state check for the overall transaction
        ph.forkPreTx();
        uint256 preTotalAssets = adopter.totalAssets();

        // Get post-state values
        ph.forkPostTx();
        uint256 postTotalAssets = adopter.totalAssets();

        // Calculate total assets deposited across all calls
        uint256 totalAssetsDeposited = 0;
        for (uint256 i = 0; i < callInputs.length; i++) {
            (uint256 assets,) = abi.decode(callInputs[i].input, (uint256, address));
            totalAssetsDeposited += assets;
        }

        // Verify total assets increased by exactly the deposited amount
        require(
            postTotalAssets == preTotalAssets + totalAssetsDeposited,
            "Deposit assets assertion failed: incorrect total assets"
        );
    }

    // Assert that deposits maintain correct share accounting
    function assertionDepositShares() external {
        // Get the assertion adopter address
        IERC4626 adopter = IERC4626(ph.getAssertionAdopter());

        // Get all deposit calls to the ERC4626 vault
        PhEvm.CallInputs[] memory callInputs = ph.getCallInputs(address(adopter), adopter.deposit.selector);

        // Get pre-state values for total assets and total supply
        ph.forkPreTx();
        uint256 preTotalAssets = adopter.totalAssets();
        uint256 preTotalSupply = adopter.totalSupply();

        // Get post-state values for total assets and total supply
        ph.forkPostTx();
        uint256 postTotalAssets = adopter.totalAssets();
        uint256 postTotalSupply = adopter.totalSupply();

        // Calculate total assets deposited and total shares minted across all calls
        uint256 totalAssetsDeposited = 0;
        uint256 totalSharesMinted = 0;

        for (uint256 i = 0; i < callInputs.length; i++) {
            (uint256 assets, address receiver) = abi.decode(callInputs[i].input, (uint256, address));
            totalAssetsDeposited += assets;

            // Get pre and post share balances for this receiver
            ph.forkPreTx();
            uint256 preShareBalance = adopter.balanceOf(receiver);
            ph.forkPostTx();
            uint256 postShareBalance = adopter.balanceOf(receiver);

            totalSharesMinted += (postShareBalance - preShareBalance);
        }

        // Verify that total assets increased by exactly the deposited amount
        require(
            postTotalAssets == preTotalAssets + totalAssetsDeposited,
            "Deposit shares assertion failed: total assets not updated correctly"
        );

        // Verify that total supply increased by exactly the minted shares
        require(
            postTotalSupply == preTotalSupply + totalSharesMinted,
            "Deposit shares assertion failed: total supply not updated correctly"
        );

        // Verify that the share-to-asset ratio remains consistent
        // This is the key check that will catch the vulnerability
        uint256 preAssetsPerShare = preTotalSupply > 0 ? (preTotalAssets * 1e18) / preTotalSupply : 0;
        uint256 postAssetsPerShare = postTotalSupply > 0 ? (postTotalAssets * 1e18) / postTotalSupply : 0;

        // Allow for minimal precision/rounding errors
        require(
            postAssetsPerShare >= preAssetsPerShare || preAssetsPerShare - postAssetsPerShare <= PRECISION_TOLERANCE,
            "Deposit shares assertion failed: share-to-asset ratio decreased unexpectedly"
        );
    }

    // Assert that withdrawals correctly update total assets
    function assertionWithdrawAssets() external {
        // Get the assertion adopter address
        IERC4626 adopter = IERC4626(ph.getAssertionAdopter());

        // Get all withdraw calls to the ERC4626 vault
        PhEvm.CallInputs[] memory callInputs = ph.getCallInputs(address(adopter), adopter.withdraw.selector);

        // First, do a simple pre/post state check for the overall transaction
        ph.forkPreTx();
        uint256 preTotalAssets = adopter.totalAssets();

        // Get post-state values
        ph.forkPostTx();
        uint256 postTotalAssets = adopter.totalAssets();

        // Calculate total assets withdrawn across all calls
        uint256 totalAssetsWithdrawn = 0;
        for (uint256 i = 0; i < callInputs.length; i++) {
            (uint256 assets,,) = abi.decode(callInputs[i].input, (uint256, address, address));
            totalAssetsWithdrawn += assets;
        }

        // Verify total assets decreased by exactly the withdrawn amount
        require(
            postTotalAssets == preTotalAssets - totalAssetsWithdrawn,
            "Withdraw assets assertion failed: incorrect total assets"
        );
    }

    // Assert that withdrawals maintain correct share accounting
    function assertionWithdrawShares() external {
        // Get the assertion adopter address
        IERC4626 adopter = IERC4626(ph.getAssertionAdopter());

        // Get all withdraw calls to the ERC4626 vault
        PhEvm.CallInputs[] memory callInputs = ph.getCallInputs(address(adopter), adopter.withdraw.selector);

        // Get pre-state values for total assets and total supply
        ph.forkPreTx();
        uint256 preTotalAssets = adopter.totalAssets();
        uint256 preTotalSupply = adopter.totalSupply();

        // Get post-state values for total assets and total supply
        ph.forkPostTx();
        uint256 postTotalAssets = adopter.totalAssets();
        uint256 postTotalSupply = adopter.totalSupply();

        // Calculate total assets withdrawn and total shares burned across all calls
        uint256 totalAssetsWithdrawn = 0;
        uint256 totalSharesBurned = 0;

        for (uint256 i = 0; i < callInputs.length; i++) {
            (uint256 assets,, address owner) = abi.decode(callInputs[i].input, (uint256, address, address));
            totalAssetsWithdrawn += assets;

            // Get pre and post share balances for this owner
            ph.forkPreTx();
            uint256 preShareBalance = adopter.balanceOf(owner);
            ph.forkPostTx();
            uint256 postShareBalance = adopter.balanceOf(owner);

            totalSharesBurned += (preShareBalance - postShareBalance);
        }

        // Verify that total assets decreased by exactly the withdrawn amount
        require(
            postTotalAssets == preTotalAssets - totalAssetsWithdrawn,
            "Withdraw shares assertion failed: total assets not updated correctly"
        );

        // Verify that total supply decreased by exactly the burned shares
        require(
            postTotalSupply == preTotalSupply - totalSharesBurned,
            "Withdraw shares assertion failed: total supply not updated correctly"
        );

        // Verify that the share-to-asset ratio remains consistent
        // This is the key check that will catch the vulnerability
        uint256 preAssetsPerShare = preTotalSupply > 0 ? (preTotalAssets * 1e18) / preTotalSupply : 0;
        uint256 postAssetsPerShare = postTotalSupply > 0 ? (postTotalAssets * 1e18) / postTotalSupply : 0;

        // Allow for minimal precision/rounding errors
        require(
            postAssetsPerShare >= preAssetsPerShare || preAssetsPerShare - postAssetsPerShare <= PRECISION_TOLERANCE,
            "Withdraw shares assertion failed: share-to-asset ratio decreased unexpectedly"
        );
    }

    // Assert that share value never decreases unexpectedly
    function assertionShareValueMonotonicity() external {
        // Get the assertion adopter address
        IERC4626 adopter = IERC4626(ph.getAssertionAdopter());

        // Create a snapshot of the state before any transactions
        ph.forkPreTx();
        uint256 assetsPerSharePre = _calculateAssetsPerShare(adopter);

        // Get state after transaction
        ph.forkPostTx();
        uint256 assetsPerSharePost = _calculateAssetsPerShare(adopter);

        // Allow for minimal precision/rounding errors with a small tolerance
        require(
            assetsPerSharePost >= assetsPerSharePre || assetsPerSharePre - assetsPerSharePost <= PRECISION_TOLERANCE,
            "Share value decreased unexpectedly"
        );

        // If share value increased, update the last known good value
        if (assetsPerSharePost > assetsPerSharePre) {
            lastKnownShareValue = assetsPerSharePost;
        }

        // Note: In practice, you would want to add detection for legitimate cases where
        // share value might decrease, such as:
        //
        // 1. Fee collection events - Some vaults charge periodic fees by reducing share value
        //    This could be detected by:
        //    - Checking specific fee collection function calls
        //    - Looking for fee events emitted by the vault
        //    - Using a timestamp-based approach to see if a fee collection period has occurred
        //    - Adding a fee collection flag or method to the vault interface
        //
        // 2. Investment loss events - When underlying investments lose value
        //    This could be detected by:
        //    - Monitoring for loss events emitted by the vault
        //    - Tracking investment strategy function calls that might result in losses
        //    - Detecting significant changes in underlying asset values
        //    - Adding a loss event flag or method to the vault interface
        //
        // These detection mechanisms should be customized based on the specific
        // implementation of the vault you're working with.
    }

    // Helper function to calculate assets per share with 1e18 precision
    function _calculateAssetsPerShare(IERC4626 vault) internal view returns (uint256) {
        uint256 totalSupply = vault.totalSupply();
        if (totalSupply == 0) {
            return lastKnownShareValue; // Return last known value if supply is zero
        }

        uint256 totalAssets = vault.totalAssets();
        return (totalAssets * 1e18) / totalSupply;
    }
}

interface IERC4626 {
    // Deposit/mint
    function deposit(uint256 assets, address receiver) external returns (uint256);
    function previewDeposit(uint256 assets) external view returns (uint256);

    // Withdraw/redeem
    function withdraw(uint256 assets, address receiver, address owner) external returns (uint256);
    function previewWithdraw(uint256 assets) external view returns (uint256);

    // View functions
    function totalAssets() external view returns (uint256);
    function totalSupply() external view returns (uint256);
    function balanceOf(address account) external view returns (uint256);
    function asset() external view returns (address);
}

interface IERC20 {
    function balanceOf(address account) external view returns (uint256);
}

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.13;

/**
 * @title ERC4626OperationsAssertion
 * @notice This assertion contract validates the correctness of ERC4626 vault operations by checking:
 *
 * 1. Batch Operations Consistency:
 *    - Validates that all ERC4626 operations (deposit, mint, withdraw, redeem) maintain correct
 *      accounting of total assets and total supply
 *    - Ensures that the net changes in assets and shares match the expected changes
 *    - Handles multiple operations in a single transaction
 *
 * 2. Deposit Operation Validation:
 *    - Verifies that deposit operations correctly increase the vault's asset balance
 *    - Ensures depositors receive the correct number of shares based on previewDeposit
 *    - Validates that the vault's total assets increase by exactly the deposited amount
 *
 * 3. Base Invariants:
 *    - Ensures the vault always has at least as many assets as shares
 *    - Validates this invariant after any storage changes
 *
 * The contract uses the Credible Layer's fork mechanism to compare pre and post-state
 * changes, ensuring that all operations maintain the vault's accounting integrity.
 */
import {Assertion} from "credible-std/Assertion.sol";
import {CoolVault} from "../../src/ass2-erc4626-operations.sol";
import {IERC4626} from "openzeppelin-contracts/interfaces/IERC4626.sol";
import {IERC20} from "openzeppelin-contracts/token/ERC20/IERC20.sol";
import {PhEvm} from "credible-std/PhEvm.sol";

contract ERC4626OperationsAssertion is Assertion {
    // The triggers function tells the Credible Layer which assertion functions to run
    function triggers() external view override {
        // Batch operations assertion - triggers on any of the four functions
        registerCallTrigger(this.assertionBatchOperationsConsistency.selector, CoolVault.deposit.selector);
        registerCallTrigger(this.assertionBatchOperationsConsistency.selector, CoolVault.mint.selector);
        registerCallTrigger(this.assertionBatchOperationsConsistency.selector, CoolVault.withdraw.selector);
        registerCallTrigger(this.assertionBatchOperationsConsistency.selector, CoolVault.redeem.selector);

        // Deposit-specific assertions
        registerCallTrigger(this.assertionDepositIncreasesBalance.selector, CoolVault.deposit.selector);
        registerCallTrigger(this.assertionDepositerSharesIncreases.selector, CoolVault.deposit.selector);

        // Base invariant assertion - triggers on storage changes
        registerStorageChangeTrigger(this.assertionVaultAlwaysAccumulatesAssets.selector, bytes32(uint256(2)));
    }

    /**
     * @dev Comprehensive assertion for batch operations: validates that all ERC4626 operations
     * in a single transaction maintain consistency across total supply and total assets
     * This function checks all four operations (deposit, mint, withdraw, redeem) that may occur
     * within the same transaction and ensures the final state is mathematically correct
     */
    function assertionBatchOperationsConsistency() external {
        // Get the assertion adopter address
        CoolVault adopter = CoolVault(ph.getAssertionAdopter());

        // Get call inputs for all four functions
        PhEvm.CallInputs[] memory depositInputs = ph.getCallInputs(address(adopter), adopter.deposit.selector);
        PhEvm.CallInputs[] memory mintInputs = ph.getCallInputs(address(adopter), adopter.mint.selector);
        PhEvm.CallInputs[] memory withdrawInputs = ph.getCallInputs(address(adopter), adopter.withdraw.selector);
        PhEvm.CallInputs[] memory redeemInputs = ph.getCallInputs(address(adopter), adopter.redeem.selector);

        // Calculate net changes from all operations
        uint256 totalAssetsAdded = 0;
        uint256 totalAssetsRemoved = 0;
        uint256 totalSharesAdded = 0;
        uint256 totalSharesRemoved = 0;

        // Process deposit operations (increase assets and supply)
        for (uint256 i = 0; i < depositInputs.length; i++) {
            (uint256 assets,) = abi.decode(depositInputs[i].input, (uint256, address));
            totalAssetsAdded += assets;
            totalSharesAdded += adopter.previewDeposit(assets);
        }

        // Process mint operations (increase assets and supply)
        for (uint256 i = 0; i < mintInputs.length; i++) {
            (uint256 shares,) = abi.decode(mintInputs[i].input, (uint256, address));
            totalSharesAdded += shares;
            totalAssetsAdded += adopter.previewMint(shares);
        }

        // Process withdraw operations (decrease assets and supply)
        for (uint256 i = 0; i < withdrawInputs.length; i++) {
            (uint256 assets,,) = abi.decode(withdrawInputs[i].input, (uint256, address, address));
            totalAssetsRemoved += assets;
            totalSharesRemoved += adopter.previewWithdraw(assets);
        }

        // Process redeem operations (decrease assets and supply)
        for (uint256 i = 0; i < redeemInputs.length; i++) {
            (uint256 shares,,) = abi.decode(redeemInputs[i].input, (uint256, address, address));
            totalSharesRemoved += shares;
            totalAssetsRemoved += adopter.previewRedeem(shares);
        }

        ph.forkPreTx();
        uint256 preVaultAssets = adopter.totalAssets();
        uint256 preVaultSupply = adopter.totalSupply();

        ph.forkPostTx();
        uint256 postVaultAssets = adopter.totalAssets();
        uint256 postVaultSupply = adopter.totalSupply();

        // Calculate expected changes
        uint256 expectedAssetsAdded = postVaultAssets > preVaultAssets ? postVaultAssets - preVaultAssets : 0;
        uint256 expectedAssetsRemoved = preVaultAssets > postVaultAssets ? preVaultAssets - postVaultAssets : 0;
        uint256 expectedSharesAdded = postVaultSupply > preVaultSupply ? postVaultSupply - preVaultSupply : 0;
        uint256 expectedSharesRemoved = preVaultSupply > postVaultSupply ? preVaultSupply - postVaultSupply : 0;

        // Ensure operations had some effect if there were calls
        require(totalAssetsAdded == expectedAssetsAdded, "Batch Operations: Assets added mismatch");
        require(totalAssetsRemoved == expectedAssetsRemoved, "Batch Operations: Assets removed mismatch");
        require(totalSharesAdded == expectedSharesAdded, "Batch Operations: Shares added mismatch");
        require(totalSharesRemoved == expectedSharesRemoved, "Batch Operations: Shares removed mismatch");
    }

    /**
     * @dev Assertion to verify that deposit operations correctly increase the vault's asset balance
     * This ensures that when users deposit assets, the vault's total assets increase by exactly the deposited amount
     */
    function assertionDepositIncreasesBalance() external {
        // Get the assertion adopter address
        CoolVault adopter = CoolVault(ph.getAssertionAdopter());

        // Get all deposit calls to the vault
        PhEvm.CallInputs[] memory inputs = ph.getCallInputs(address(adopter), adopter.deposit.selector);

        for (uint256 i = 0; i < inputs.length; i++) {
            (uint256 assets, address receiver) = abi.decode(inputs[i].input, (uint256, address));

            // Check pre-state
            ph.forkPreTx();
            uint256 vaultAssetPreBalance = adopter.totalAssets();
            uint256 userSharesPreBalance = adopter.balanceOf(receiver);
            uint256 expectedShares = adopter.previewDeposit(assets);

            // Check post-state
            ph.forkPostTx();
            uint256 vaultAssetPostBalance = adopter.totalAssets();
            uint256 userSharesPostBalance = adopter.balanceOf(receiver);

            // Verify vault assets increased by exactly the deposited amount
            require(
                vaultAssetPostBalance == vaultAssetPreBalance + assets,
                "Deposit assertion failed: Vault assets did not increase by the correct amount"
            );

            // Verify user received exactly the expected number of shares
            require(
                userSharesPostBalance == userSharesPreBalance + expectedShares,
                "Deposit assertion failed: User did not receive the correct number of shares"
            );
        }
    }

    /**
     * @dev Assertion to verify that deposit operations correctly increase the depositor's share balance
     * This ensures that when users deposit assets, they receive the correct number of shares
     */
    function assertionDepositerSharesIncreases() external {
        // Get the assertion adopter address
        CoolVault adopter = CoolVault(ph.getAssertionAdopter());

        PhEvm.CallInputs[] memory inputs = ph.getCallInputs(address(adopter), adopter.deposit.selector);

        for (uint256 i = 0; i < inputs.length; i++) {
            ph.forkPreTx();
            (uint256 assets,) = abi.decode(inputs[i].input, (uint256, address));
            uint256 previewPreAssets = adopter.previewDeposit(assets);
            address depositer = inputs[0].caller;
            uint256 preShares = adopter.balanceOf(depositer);

            ph.forkPostTx();

            uint256 postShares = adopter.balanceOf(depositer);

            require(
                postShares == preShares + previewPreAssets,
                "Depositer shares assertion failed: Share balance did not increase correctly"
            );
        }
    }

    /**
     * @dev Base invariant assertion to verify that the vault always has at least as many assets as shares
     * This is a fundamental invariant of ERC4626 vaults - they should never have more shares than assets
     */
    function assertionVaultAlwaysAccumulatesAssets() external {
        // Get the assertion adopter address
        CoolVault adopter = CoolVault(ph.getAssertionAdopter());

        ph.forkPostTx();

        uint256 vaultAssetPostBalance = adopter.totalAssets();
        uint256 vaultSharesPostBalance = adopter.balanceOf(address(adopter));

        require(
            vaultAssetPostBalance >= vaultSharesPostBalance, "Base invariant failed: Vault has more shares than assets"
        );
    }
}