feat: enhance PegInInvariant tests with non-LP refund liability and additional helper functions

Author: Hakob23

test/invariant/PegInInvariant.t.sol

@@ -10,6 +10,12 @@ import {Flyover} from "../../src/libraries/Flyover.sol";
 /// @notice Tests critical invariants for the PegInContract using a dedicated handler
 contract PegInInvariantTest is PegInTestBase {
     PegInHandler public handler;
+    address public nonLpRefundCreditor;
+    uint256 public baselineBalance;
+
+    bytes constant RAW_TX_MOCK = hex"112233";
+    bytes constant PMT_MOCK = hex"010203";
+    uint256 constant HEIGHT_MOCK = 10;
 
     function setUp() public {
         deployPegInContract();
@@ -30,6 +36,17 @@ contract PegInInvariantTest is PegInTestBase {
         );
         handler.addLP(extraPegInLp);
 
+        (address creditor, ) = _seedNonLpRefundLiabilityFixture(
+            fullLp,
+            fullLpKey,
+            1 ether,
+            HEIGHT_MOCK,
+            RAW_TX_MOCK,
+            PMT_MOCK
+        );
+        nonLpRefundCreditor = creditor;
+        baselineBalance = address(pegInContract).balance;
+
         targetContract(address(handler));
 
         bytes4[] memory selectors = new bytes4[](2);
@@ -42,7 +59,7 @@ contract PegInInvariantTest is PegInTestBase {
 
     // ============ Invariant Tests ============
 
-    /// @notice Contract balance must cover sum of all LP internal balances
+    /// @notice LP-only lower bound: contract balance must cover tracked LP internal balances
     function invariant_ContractSolvent() public view {
         uint256 totalLPBalances = handler.calculateTotalLPBalances();
         assertGe(
@@ -52,6 +69,20 @@ contract PegInInvariantTest is PegInTestBase {
         );
     }
 
+    /// @notice Contract balance must cover tracked LP liabilities plus seeded non-LP refund creditor
+    function invariant_ContractSolventIncludingNonLpRefundCreditor()
+        public
+        view
+    {
+        uint256 totalLPBalances = handler.calculateTotalLPBalances();
+        uint256 nonLpLiability = pegInContract.getBalance(nonLpRefundCreditor);
+        assertGe(
+            address(pegInContract).balance,
+            totalLPBalances + nonLpLiability,
+            "INVARIANT VIOLATED: PegIn contract insolvent with non-LP refund liability"
+        );
+    }
+
     /// @notice No individual LP balance should exceed total deposited through the handler
     function invariant_NoUnderflow() public view {
         uint256 deposited = handler.ghost_totalDeposited();
@@ -66,13 +97,18 @@ contract PegInInvariantTest is PegInTestBase {
         }
     }
 
-    /// @notice Contract balance must equal deposited minus withdrawn (tight equality)
+    /// @notice Contract balance must equal deposited minus withdrawn (for deposit/withdraw handler actions)
     function invariant_GhostAccounting() public view {
         uint256 deposited = handler.ghost_totalDeposited();
         uint256 withdrawn = handler.ghost_totalWithdrawn();
+        assertGe(
+            address(pegInContract).balance,
+            baselineBalance,
+            "INVARIANT VIOLATED: Contract balance unexpectedly below baseline"
+        );
 
         assertEq(
-            address(pegInContract).balance,
+            address(pegInContract).balance - baselineBalance,
             deposited - withdrawn,
             "INVARIANT VIOLATED: Contract balance != deposited - withdrawn"
         );
@@ -91,6 +127,11 @@ contract PegInInvariantTest is PegInTestBase {
             "Handler withdraw calls:",
             handler.getHandlerCalls("withdraw")
         );
+        console.log(
+            "Seeded non-LP liability:",
+            pegInContract.getBalance(nonLpRefundCreditor)
+        );
+        console.log("Baseline balance:", baselineBalance);
         console.log("Contract balance:", address(pegInContract).balance);
         console.log("-------------------------------\n");
     }

test/pegin/PegInTestBase.sol

@@ -7,6 +7,7 @@ import {CollateralManagementContract} from "../../src/CollateralManagement.sol";
 import {FlyoverDiscovery} from "../../src/FlyoverDiscovery.sol";
 import {PauseRegistry} from "../../src/PauseRegistry.sol";
 import {BridgeMock} from "../../src/test-contracts/BridgeMock.sol";
+import {WalletMock} from "../../src/test-contracts/WalletMock.sol";
 import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
 import {Quotes} from "../../src/libraries/Quotes.sol";
 import {Flyover} from "../../src/libraries/Flyover.sol";
@@ -184,4 +185,107 @@ abstract contract PegInTestBase is Test {
             Flyover.ProviderType.Both
         );
     }
+
+    /// @notice Shared helper for constructing a standard PegIn quote fixture
+    function _buildPegInQuoteFixture(
+        uint256 value,
+        address lp,
+        address payable refundAddress,
+        address destinationContract,
+        uint256 nonce
+    ) internal view returns (Quotes.PegInQuote memory) {
+        bytes memory testBtcAddress = new bytes(21);
+
+        return
+            Quotes.PegInQuote({
+                chainId: block.chainid,
+                callFee: 100000000000000,
+                penaltyFee: 10000000000000,
+                value: value,
+                gasFee: 100,
+                fedBtcAddress: bytes20(testBtcAddress),
+                lbcAddress: address(pegInContract),
+                liquidityProviderRskAddress: lp,
+                contractAddress: destinationContract,
+                rskRefundAddress: refundAddress,
+                nonce: int64(uint64(nonce)),
+                gasLimit: 21000,
+                agreementTimestamp: uint32(block.timestamp),
+                timeForDeposit: 3600,
+                callTime: 7200,
+                depositConfirmations: 10,
+                callOnRegister: false,
+                btcRefundAddress: testBtcAddress,
+                liquidityProviderBtcAddress: testBtcAddress,
+                data: new bytes(0)
+            });
+    }
+
+    /// @notice Shared helper for signing PegIn quotes
+    function _signPegInQuoteWithKey(
+        uint256 privateKey,
+        Quotes.PegInQuote memory quote
+    ) internal view returns (bytes memory) {
+        bytes32 eip712Hash = pegInContract.hashPegInQuoteEIP712(quote);
+        (uint8 v, bytes32 r, bytes32 s) = vm.sign(privateKey, eip712Hash);
+        return abi.encodePacked(r, s, v);
+    }
+
+    /// @notice Shared helper for creating BTC block headers with LE timestamp
+    function _btcHeaderFromTimestamp(
+        uint32 timestamp
+    ) internal pure returns (bytes memory) {
+        bytes memory header = new bytes(80);
+        header[68] = bytes1(uint8(timestamp));
+        header[69] = bytes1(uint8(timestamp >> 8));
+        header[70] = bytes1(uint8(timestamp >> 16));
+        header[71] = bytes1(uint8(timestamp >> 24));
+        return header;
+    }
+
+    /// @notice Seeds one non-LP internal balance by forcing refund transfer failure on registerPegIn
+    function _seedNonLpRefundLiabilityFixture(
+        address lp,
+        uint256 lpPrivateKey,
+        uint256 value,
+        uint256 height,
+        bytes memory rawTx,
+        bytes memory pmt
+    ) internal returns (address creditor, uint256 creditedAmount) {
+        WalletMock refundWallet = new WalletMock();
+        refundWallet.setRejectFunds(true);
+        creditor = address(refundWallet);
+
+        Quotes.PegInQuote memory quote = _buildPegInQuoteFixture(
+            value,
+            lp,
+            payable(creditor),
+            address(0xBEEF),
+            uint256(uint64(block.timestamp))
+        );
+        bytes32 quoteHash = pegInContract.hashPegInQuote(quote);
+        bytes memory signature = _signPegInQuoteWithKey(lpPrivateKey, quote);
+        uint256 peginAmount = quote.value + quote.callFee + quote.gasFee;
+
+        vm.deal(address(this), peginAmount);
+        bridgeMock.setPegin{value: peginAmount}(quoteHash);
+        bridgeMock.setHeader(
+            height,
+            _btcHeaderFromTimestamp(uint32(block.timestamp) + 300)
+        );
+        bridgeMock.setHeader(
+            height + uint256(quote.depositConfirmations) - 1,
+            _btcHeaderFromTimestamp(uint32(block.timestamp) + 600)
+        );
+
+        vm.prank(lp);
+        pegInContract.registerPegIn(quote, signature, rawTx, pmt, height);
+
+        creditedAmount = pegInContract.getBalance(creditor);
+        assertGt(
+            creditedAmount,
+            0,
+            "Expected non-LP refund creditor balance to be seeded"
+        );
+    }
 }