smom-dbis-138/test/flash/DBISEngineXVirtualBatchVault.t.sol

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

import {Test} from "forge-std/Test.sol";
import {IERC3156FlashBorrower} from "@openzeppelin/contracts/interfaces/IERC3156FlashBorrower.sol";
import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import {DBISEngineXFlashProofBorrower} from "../../contracts/flash/DBISEngineXFlashProofBorrower.sol";
import {DBISEngineXVirtualBatchVault} from "../../contracts/flash/DBISEngineXVirtualBatchVault.sol";
import {MockMintableToken} from "../dbis/MockMintableToken.sol";

contract EngineXFlashBorrower is IERC3156FlashBorrower {
    bytes32 private constant _RETURN_VALUE = keccak256("ERC3156FlashBorrower.onFlashLoan");
    bool public repay = true;

    function setRepay(bool repay_) external {
        repay = repay_;
    }

    function onFlashLoan(address, address token, uint256 amount, uint256 fee, bytes calldata)
        external
        override
        returns (bytes32)
    {
        if (repay) {
            IERC20(token).transfer(msg.sender, amount + fee);
        }
        return _RETURN_VALUE;
    }
}

contract DBISEngineXVirtualBatchVaultTest is Test {
    MockMintableToken internal cwusdc;
    MockMintableToken internal usdc;
    MockMintableToken internal xaut;
    DBISEngineXVirtualBatchVault internal vault;

    address internal constant USER = address(0xBEEF);
    address internal constant SURPLUS_RECEIVER = address(0xCAFE);
    address internal constant OUTPUT_RECIPIENT = address(0xD00D);
    address internal constant ROUNDING_RECEIVER = address(0xA11CE);

    uint256 internal constant LIVE_POOL_RESERVE = 85_763_529;
    uint256 internal constant LENDER_USDC = 5_000_000;
    uint256 internal constant XAUT_USD_PRICE6 = 3_226_640_000;
    uint256 internal constant LTV_BPS = 8_000;
    uint256 internal constant MAX_ROUND_TRIP_LOSS_BPS = 100;
    bytes32 internal constant ISO_HASH = bytes32(uint256(0x1001));
    bytes32 internal constant AUDIT_HASH = bytes32(uint256(0x1002));
    bytes32 internal constant PEG_HASH = bytes32(uint256(0x1003));

    event VirtualProofAuditEvidence(
        bytes32 indexed proofId,
        address indexed operator,
        address indexed outputRecipient,
        uint256 exactOutputAmount,
        uint256 outputRoundingAmount,
        address roundingReceiver,
        bytes32 iso20022DocumentHash,
        bytes32 auditEnvelopeHash,
        bytes32 pegProofHash
    );

    function setUp() public {
        cwusdc = new MockMintableToken("Wrapped cWUSDC", "cWUSDC", 6, address(this));
        usdc = new MockMintableToken("USD Coin", "USDC", 6, address(this));
        xaut = new MockMintableToken("Tether Gold", "XAUt", 6, address(this));

        vault = new DBISEngineXVirtualBatchVault(
            address(cwusdc),
            address(usdc),
            address(xaut),
            address(this),
            SURPLUS_RECEIVER,
            XAUT_USD_PRICE6,
            LTV_BPS,
            MAX_ROUND_TRIP_LOSS_BPS
        );

        cwusdc.mint(address(this), LIVE_POOL_RESERVE);
        usdc.mint(address(this), LIVE_POOL_RESERVE + LENDER_USDC);
        cwusdc.approve(address(vault), type(uint256).max);
        usdc.approve(address(vault), type(uint256).max);
        vault.seedPool(LIVE_POOL_RESERVE, LIVE_POOL_RESERVE);
        vault.fundLender(LENDER_USDC);

        cwusdc.mint(USER, 10_000_000_000_000_000);
        xaut.mint(USER, 10_000_000);
        vm.startPrank(USER);
        cwusdc.approve(address(vault), type(uint256).max);
        xaut.approve(address(vault), type(uint256).max);
        vm.stopPrank();
    }

    function testPreviewMatchesLiveMaintainedLoopMath() public view {
        (
            uint256 collateralXaut,
            uint256 cwusdcInPerLoop,
            uint256 cwusdcOutPerLoop,
            uint256 cwusdcLossPerLoop,
            uint256 totalCwusdcInAmount,
            uint256 totalCwusdcOutAmount,
            uint256 totalNeutralizedCwusdcAmount
        ) = vault.previewVirtualProof(LENDER_USDC, 3);

        assertEq(collateralXaut, 1_937, "collateral should match live 5 USDC floor");
        assertEq(cwusdcInPerLoop, 5_325_523, "cwusdc in should match live proof");
        assertEq(cwusdcOutPerLoop, 5_295_471, "cwusdc out should match live proof");
        assertEq(cwusdcLossPerLoop, 30_052, "neutralized loss should match live proof");
        assertEq(totalCwusdcInAmount, 15_976_569, "three-loop cWUSDC in should match live proof");
        assertEq(totalCwusdcOutAmount, 15_886_413, "three-loop cWUSDC out should match live proof");
        assertEq(totalNeutralizedCwusdcAmount, 90_156, "three-loop neutralization should match live proof");
    }

    function testVirtualProofNetSettlesAndKeepsPoolMaintained() public {
        uint256 userCwusdcBefore = cwusdc.balanceOf(USER);
        uint256 userXautBefore = xaut.balanceOf(USER);

        vm.prank(USER);
        vault.runVirtualProof(bytes32("proof-1"), LENDER_USDC, 3);

        assertEq(vault.poolCwusdcReserve(), LIVE_POOL_RESERVE, "pool cWUSDC reserve should not drift");
        assertEq(vault.poolUsdcReserve(), LIVE_POOL_RESERVE, "pool USDC reserve should not drift");
        assertEq(vault.currentSurplusCwusdc(), 0, "pool surplus should stay neutralized");
        assertEq(vault.lenderUsdcAvailable(), LENDER_USDC, "lender bucket should be reusable");

        assertEq(cwusdc.balanceOf(USER), userCwusdcBefore - 90_156, "user only loses neutralized amount");
        assertEq(cwusdc.balanceOf(SURPLUS_RECEIVER), 90_156, "receiver gets neutralized cWUSDC");
        assertEq(xaut.balanceOf(USER), userXautBefore, "XAUt collateral should be returned");

        assertEq(vault.totalVirtualLoops(), 3, "virtual loop counter");
        assertEq(vault.totalVirtualDebtUsdc(), 15_000_000, "virtual debt counter");
        assertEq(vault.totalVirtualCwusdcIn(), 15_976_569, "virtual cWUSDC in counter");
        assertEq(vault.totalNeutralizedCwusdc(), 90_156, "neutralized counter");
        assertTrue(vault.usedProofIds(bytes32("proof-1")), "proof id should be consumed");
    }

    function testVirtualProofCanSendOutputToRecipient() public {
        uint256 userCwusdcBefore = cwusdc.balanceOf(USER);
        uint256 userXautBefore = xaut.balanceOf(USER);

        vm.prank(USER);
        vault.runVirtualProofTo(bytes32("proof-to"), LENDER_USDC, 3, OUTPUT_RECIPIENT);

        assertEq(vault.poolCwusdcReserve(), LIVE_POOL_RESERVE, "pool cWUSDC reserve should not drift");
        assertEq(vault.poolUsdcReserve(), LIVE_POOL_RESERVE, "pool USDC reserve should not drift");
        assertEq(cwusdc.balanceOf(USER), userCwusdcBefore - 15_976_569, "sender supplies gross cWUSDC");
        assertEq(cwusdc.balanceOf(OUTPUT_RECIPIENT), 15_886_413, "recipient gets cWUSDC output");
        assertEq(cwusdc.balanceOf(SURPLUS_RECEIVER), 90_156, "receiver gets neutralized cWUSDC");
        assertEq(xaut.balanceOf(USER), userXautBefore, "XAUt collateral should be returned");
    }

    function testVirtualProofExactOutToAnchorsAuditProofsAndSendsRounding() public {
        uint256 userCwusdcBefore = cwusdc.balanceOf(USER);
        uint256 userXautBefore = xaut.balanceOf(USER);
        uint256 exactOutput = 15_886_000;
        bytes32 proofId = bytes32("proof-audit");

        vm.expectEmit(true, true, true, true, address(vault));
        emit VirtualProofAuditEvidence(
            proofId, USER, OUTPUT_RECIPIENT, exactOutput, 413, ROUNDING_RECEIVER, ISO_HASH, AUDIT_HASH, PEG_HASH
        );
        vm.prank(USER);
        vault.runVirtualProofExactOutTo(
            proofId, LENDER_USDC, 3, OUTPUT_RECIPIENT, exactOutput, ROUNDING_RECEIVER, ISO_HASH, AUDIT_HASH, PEG_HASH
        );

        assertEq(vault.poolCwusdcReserve(), LIVE_POOL_RESERVE, "pool cWUSDC reserve should not drift");
        assertEq(vault.poolUsdcReserve(), LIVE_POOL_RESERVE, "pool USDC reserve should not drift");
        assertEq(cwusdc.balanceOf(USER), userCwusdcBefore - 15_976_569, "sender supplies gross cWUSDC");
        assertEq(cwusdc.balanceOf(OUTPUT_RECIPIENT), exactOutput, "recipient gets exact cWUSDC");
        assertEq(cwusdc.balanceOf(ROUNDING_RECEIVER), 413, "rounding receiver gets output dust");
        assertEq(cwusdc.balanceOf(SURPLUS_RECEIVER), 90_156, "receiver gets neutralized cWUSDC");
        assertEq(xaut.balanceOf(USER), userXautBefore, "XAUt collateral should be returned");
    }

    function testVirtualLoopCountScalesTotalsButNotSequentialCollateral() public view {
        uint256 loops = 938_874_924;
        (
            uint256 collateralXaut,
            uint256 cwusdcInPerLoop,,
            uint256 cwusdcLossPerLoop,
            uint256 totalCwusdcInAmount,,
            uint256 totalNeutralizedCwusdcAmount
        ) = vault.previewVirtualProof(LENDER_USDC, loops);

        assertEq(collateralXaut, 1_937, "sequential virtual batch reuses one loop of XAUt collateral");
        assertEq(cwusdcInPerLoop, 5_325_523, "per-loop input stays constant");
        assertEq(cwusdcLossPerLoop, 30_052, "per-loop neutralization stays constant");
        assertEq(totalCwusdcInAmount, 5_000_000_001_885_252, "virtual batch covers 5B cWUSDC gross");
        assertEq(totalNeutralizedCwusdcAmount, 28_215_069_216_048, "5B batch neutralized amount");
    }

    function testPreviewRevertsWhenPoolIsNotMaintained() public {
        DBISEngineXVirtualBatchVault unevenVault = new DBISEngineXVirtualBatchVault(
            address(cwusdc),
            address(usdc),
            address(xaut),
            address(this),
            SURPLUS_RECEIVER,
            XAUT_USD_PRICE6,
            LTV_BPS,
            MAX_ROUND_TRIP_LOSS_BPS
        );
        cwusdc.mint(address(this), LIVE_POOL_RESERVE + 1);
        usdc.mint(address(this), LIVE_POOL_RESERVE + LENDER_USDC);
        cwusdc.approve(address(unevenVault), type(uint256).max);
        usdc.approve(address(unevenVault), type(uint256).max);
        unevenVault.seedPool(LIVE_POOL_RESERVE + 1, LIVE_POOL_RESERVE);
        unevenVault.fundLender(LENDER_USDC);

        vm.expectRevert(bytes("pool not maintained"));
        unevenVault.previewVirtualProof(LENDER_USDC, 1);
    }

    function testRunVirtualProofRejectsDuplicateProofIds() public {
        vm.prank(USER);
        vault.runVirtualProof(bytes32("proof-1"), LENDER_USDC, 1);

        vm.expectRevert(bytes("proof used"));
        vm.prank(USER);
        vault.runVirtualProof(bytes32("proof-1"), LENDER_USDC, 1);
    }

    function testRunVirtualProofToRejectsZeroRecipient() public {
        vm.expectRevert(bytes("zero output"));
        vm.prank(USER);
        vault.runVirtualProofTo(bytes32("proof-zero"), LENDER_USDC, 1, address(0));
    }

    function testRunVirtualProofExactOutToRejectsMissingAuditHashes() public {
        vm.expectRevert(bytes("zero iso hash"));
        vm.prank(USER);
        vault.runVirtualProofExactOutTo(
            bytes32("proof-no-iso"),
            LENDER_USDC,
            1,
            OUTPUT_RECIPIENT,
            1,
            ROUNDING_RECEIVER,
            bytes32(0),
            AUDIT_HASH,
            PEG_HASH
        );
    }

    function testRunVirtualProofExactOutToRejectsOutputAbovePreview() public {
        vm.expectRevert(bytes("exact output too high"));
        vm.prank(USER);
        vault.runVirtualProofExactOutTo(
            bytes32("proof-too-high"),
            LENDER_USDC,
            1,
            OUTPUT_RECIPIENT,
            5_295_472,
            ROUNDING_RECEIVER,
            ISO_HASH,
            AUDIT_HASH,
            PEG_HASH
        );
    }

    function testWithdrawPoolLiquidityUpdatesAccountingAndPreservesMaintainedPool() public {
        uint256 withdrawAmount = 10_000_000;
        uint256 ownerCwusdcBefore = cwusdc.balanceOf(address(this));
        uint256 ownerUsdcBefore = usdc.balanceOf(address(this));

        vault.withdrawPoolLiquidity(address(this), withdrawAmount, withdrawAmount);

        assertEq(vault.poolCwusdcReserve(), LIVE_POOL_RESERVE - withdrawAmount, "pool cWUSDC accounting");
        assertEq(vault.poolUsdcReserve(), LIVE_POOL_RESERVE - withdrawAmount, "pool USDC accounting");
        assertEq(vault.lenderUsdcAvailable(), LENDER_USDC, "lender accounting unchanged");
        assertEq(cwusdc.balanceOf(address(this)), ownerCwusdcBefore + withdrawAmount, "owner cWUSDC received");
        assertEq(usdc.balanceOf(address(this)), ownerUsdcBefore + withdrawAmount, "owner USDC received");
    }

    function testWithdrawPoolLiquidityRejectsBreakingMaintainedPool() public {
        vm.expectRevert(bytes("would break maintained pool"));
        vault.withdrawPoolLiquidity(address(this), 1, 0);
    }

    function testWithdrawLenderUsdcUpdatesAccounting() public {
        uint256 withdrawAmount = 1_000_000;
        uint256 ownerUsdcBefore = usdc.balanceOf(address(this));

        vault.withdrawLenderUsdc(address(this), withdrawAmount);

        assertEq(vault.lenderUsdcAvailable(), LENDER_USDC - withdrawAmount, "lender accounting");
        assertEq(vault.poolUsdcReserve(), LIVE_POOL_RESERVE, "pool accounting unchanged");
        assertEq(usdc.balanceOf(address(this)), ownerUsdcBefore + withdrawAmount, "owner USDC received");
    }

    function testGenericWithdrawCannotTouchAccountedBalances() public {
        vm.expectRevert(bytes("accounting undercollateralized"));
        vault.withdraw(address(usdc), address(this), 1);
    }

    function testGenericWithdrawCanRescueUnaccountedTokens() public {
        uint256 dust = 123;
        usdc.mint(address(vault), dust);
        uint256 ownerUsdcBefore = usdc.balanceOf(address(this));

        vault.withdraw(address(usdc), address(this), dust);

        assertEq(vault.poolUsdcReserve(), LIVE_POOL_RESERVE, "pool accounting unchanged");
        assertEq(vault.lenderUsdcAvailable(), LENDER_USDC, "lender accounting unchanged");
        assertEq(usdc.balanceOf(address(this)), ownerUsdcBefore + dust, "owner receives unaccounted dust");
    }

    function testFlashLoanUsesLenderBucketAndCollectsFee() public {
        EngineXFlashBorrower borrower = new EngineXFlashBorrower();
        uint256 amount = 1_000_000;
        uint256 fee = vault.flashFee(address(usdc), amount);
        usdc.mint(address(borrower), fee);

        vm.prank(USER);
        vault.flashLoan(IERC3156FlashBorrower(address(borrower)), address(usdc), amount, "");

        assertEq(vault.lenderUsdcAvailable(), LENDER_USDC + fee, "fee stays in lender bucket");
        assertEq(vault.totalFlashFeesCollectedUsdc(), fee, "fee accounting");
        assertEq(usdc.balanceOf(address(vault)), LIVE_POOL_RESERVE + LENDER_USDC + fee, "USDC backing");
    }

    function testFlashLoanCanPullRepaymentByAllowance() public {
        EngineXFlashBorrower borrower = new EngineXFlashBorrower();
        uint256 amount = 1_000_000;
        uint256 fee = vault.flashFee(address(usdc), amount);
        usdc.mint(address(borrower), fee);

        borrower.setRepay(false);
        vm.prank(address(borrower));
        usdc.approve(address(vault), type(uint256).max);

        vm.prank(USER);
        vault.flashLoan(IERC3156FlashBorrower(address(borrower)), address(usdc), amount, "");

        assertEq(vault.lenderUsdcAvailable(), LENDER_USDC + fee, "fee stays in lender bucket");
        assertEq(vault.totalFlashFeesCollectedUsdc(), fee, "fee accounting");
    }

    function testFlashLoanRejectsBorrowingPoolUsdc() public {
        EngineXFlashBorrower borrower = new EngineXFlashBorrower();

        vm.expectRevert(bytes("insufficient lender usdc"));
        vault.flashLoan(IERC3156FlashBorrower(address(borrower)), address(usdc), LENDER_USDC + 1, "");
    }

    function testFlashLoanRejectsUnsupportedToken() public {
        EngineXFlashBorrower borrower = new EngineXFlashBorrower();

        vm.expectRevert(bytes("unsupported flash token"));
        vault.flashLoan(IERC3156FlashBorrower(address(borrower)), address(cwusdc), 1, "");
    }

    function testFlashLoanCanBeCapped() public {
        EngineXFlashBorrower borrower = new EngineXFlashBorrower();
        vault.setMaxFlashLoanAmount(999_999);

        vm.expectRevert(bytes("flash amount too high"));
        vault.flashLoan(IERC3156FlashBorrower(address(borrower)), address(usdc), 1_000_000, "");
    }

    function testFlashLoanAllowlistRejectsUnapprovedBorrower() public {
        EngineXFlashBorrower borrower = new EngineXFlashBorrower();
        vault.setFlashBorrowerAllowlistEnabled(true);

        vm.expectRevert(bytes("flash borrower not approved"));
        vault.flashLoan(IERC3156FlashBorrower(address(borrower)), address(usdc), 1, "");
    }

    function testFlashLoanAllowlistAllowsApprovedBorrower() public {
        EngineXFlashBorrower borrower = new EngineXFlashBorrower();
        uint256 amount = 1_000_000;
        uint256 fee = vault.flashFee(address(usdc), amount);
        usdc.mint(address(borrower), fee);

        vault.setFlashBorrowerAllowlistEnabled(true);
        vault.setFlashBorrowerApproved(address(borrower), true);

        vault.flashLoan(IERC3156FlashBorrower(address(borrower)), address(usdc), amount, "");

        assertEq(vault.lenderUsdcAvailable(), LENDER_USDC + fee, "fee stays in lender bucket");
    }

    function testPauseBlocksProofsAndFlashLoans() public {
        EngineXFlashBorrower borrower = new EngineXFlashBorrower();
        vault.pause();

        vm.expectRevert(bytes("paused"));
        vm.prank(USER);
        vault.runVirtualProof(bytes32("proof-paused"), LENDER_USDC, 1);

        vm.expectRevert(bytes("paused"));
        vault.flashLoan(IERC3156FlashBorrower(address(borrower)), address(usdc), 1, "");

        assertEq(vault.maxFlashLoan(address(usdc)), 0, "paused max flash");
    }

    function testEngineXFlashProofBorrowerRunsProofFlash() public {
        DBISEngineXFlashProofBorrower borrower =
            new DBISEngineXFlashProofBorrower(address(vault), address(usdc), address(this));
        uint256 amount = 1_000_000;
        uint256 fee = vault.flashFee(address(usdc), amount);
        bytes32 proofId = bytes32("flash-proof");
        usdc.mint(address(borrower), fee);

        borrower.runFlashProof(amount, proofId, ISO_HASH, AUDIT_HASH, PEG_HASH);

        assertTrue(borrower.usedProofIds(proofId), "proof consumed");
        assertEq(vault.lenderUsdcAvailable(), LENDER_USDC + fee, "fee stays in lender bucket");
    }
}