Skip to content

Commit

Permalink
👷🏻 ✅ Add and test management funcs
Browse files Browse the repository at this point in the history
  • Loading branch information
@tommyrharper
tommyrharper committed Jul 11, 2024
1 parent 0700baa commit 9b04be9
Show file tree
Hide file tree
Showing 2 changed files with 209 additions and 18 deletions.
53 changes: 40 additions & 13 deletions src/MarginPaymaster.sol
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
// SPDX-License-Identifier: UNLICENSED
pragma solidity 0.8.20;

import {EntryPoint} from "@account-abstraction/contracts/core/EntryPoint.sol";
import {IPaymaster, UserOperation} from "lib/account-abstraction/contracts/interfaces/IPaymaster.sol";
import {IPerpsMarketProxy} from "src/interfaces/external/IPerpsMarketProxy.sol";
import {IV3SwapRouter} from "src/interfaces/external/IV3SwapRouter.sol";
Expand All @@ -26,7 +27,7 @@ contract MarginPaymaster is IPaymaster, Zap, Ownable {
IMMUTABLES
//////////////////////////////////////////////////////////////*/

address public immutable entryPoint;
EntryPoint public immutable entryPoint;
IEngine public immutable smartMarginV3;
IPerpsMarketProxy public immutable perpsMarketSNXV3;
IV3SwapRouter public immutable uniV3Router;
Expand Down Expand Up @@ -64,7 +65,7 @@ contract MarginPaymaster is IPaymaster, Zap, Ownable {
address _weth,
address _pool
) Zap(_usdc, _sUSDProxy, _spotMarketProxy, _sUSDCId) {
entryPoint = _entryPoint;
entryPoint = EntryPoint(payable(_entryPoint));
smartMarginV3 = IEngine(_smartMarginV3);
perpsMarketSNXV3 = IPerpsMarketProxy(_perpsMarketSNXV3);
uniV3Router = IV3SwapRouter(_uniV3Router);
Expand All @@ -81,7 +82,7 @@ contract MarginPaymaster is IPaymaster, Zap, Ownable {
//////////////////////////////////////////////////////////////*/

modifier onlyEntryPoint() {
if (msg.sender != entryPoint) revert InvalidEntryPoint();
if (msg.sender != address(entryPoint)) revert InvalidEntryPoint();
_;
}

Expand Down Expand Up @@ -116,13 +117,13 @@ contract MarginPaymaster is IPaymaster, Zap, Ownable {
POST OP
//////////////////////////////////////////////////////////////*/

// TODO: handle user not having enough funds
function postOp(
PostOpMode,
bytes calldata context,
uint256 actualGasCostInWei
) external onlyEntryPoint {
uint256 costOfGasInUSDC = getCostOfGasInUSDC(actualGasCostInWei); // TODO: account for gas costs of postOp func
uint256 USDCToSwapForWETH = costOfGasInUSDC;
address sender = abi.decode(context, (address));
(uint256 availableUSDCInWallet, , ) = getUSDCAvailableInWallet(sender);

Expand All @@ -145,20 +146,43 @@ contract MarginPaymaster is IPaymaster, Zap, Ownable {
// TODO: handle users who don't have an snx account or margin or enough margin
withdrawFromMargin(sender, sUSDToWithdrawFromMargin);
// zap sUSD into USDC
USDCToSwapForWETH =
_zapOut(sUSDToWithdrawFromMargin) +
availableUSDCInWallet;
_zapOut(sUSDToWithdrawFromMargin);
}
}

console.log("actualGasCostInWei", actualGasCostInWei); // 43350920000000 = 0.00004335092 ETH = 0.13 USD
/*//////////////////////////////////////////////////////////////
FUND MANAGEMENT
//////////////////////////////////////////////////////////////*/

// TODO: remove these steps
function swapUSDCToETH(uint256 amountOutMinimum) external onlyOwner {
uint256 amountIn = _USDC.balanceOf(address(this));
// swap USDC for WETH
uint256 amountOut = swapUSDCForWETH(USDCToSwapForWETH);
uint256 amountOut = swapUSDCForWETH(amountIn, amountOutMinimum);
// unwrap WETH to ETH
weth.withdraw(amountOut);
}

function depositToEntryPoint(uint256 amount) external onlyOwner {
entryPoint.depositTo{value: amount}(address(this));
}

function stake(uint256 amount, uint32 unstakeDelaySec) external onlyOwner {
entryPoint.addStake{value: amount}(unstakeDelaySec);
}

function unlockStake() external onlyOwner {
entryPoint.unlockStake();
}

function withdrawStake(address payable withdrawAddress) external onlyOwner {
entryPoint.withdrawStake(withdrawAddress);
}

// TODO: add renew deposit logic if it is running low
function withdrawTo(
address payable withdrawAddress,
uint256 withdrawAmount
) external onlyOwner {
entryPoint.withdrawTo(withdrawAddress, withdrawAmount);
}

/*//////////////////////////////////////////////////////////////
Expand Down Expand Up @@ -200,15 +224,18 @@ contract MarginPaymaster is IPaymaster, Zap, Ownable {
);
}

function swapUSDCForWETH(uint256 amountIn) internal returns (uint256) {
function swapUSDCForWETH(
uint256 amountIn,
uint256 amountOutMinimum
) internal returns (uint256) {
IV3SwapRouter.ExactInputSingleParams memory params = IV3SwapRouter
.ExactInputSingleParams({
tokenIn: address(_USDC),
tokenOut: address(weth),
fee: POOL_FEE,
recipient: address(this),
amountIn: amountIn,
amountOutMinimum: 0, // TODO: think, should this be actualGasCostInWei???
amountOutMinimum: amountOutMinimum,
sqrtPriceLimitX96: 0
});
return uniV3Router.exactInputSingle(params);
Expand Down
174 changes: 169 additions & 5 deletions test/MarginPaymaster.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import {Bootstrap} from "test/utils/Bootstrap.sol";
import {EntryPoint, UserOperation} from "lib/account-abstraction/contracts/core/EntryPoint.sol";
import {AccountFactory, MockAccount} from "src/MockAccount.sol";
import {MarginPaymaster, IPaymaster} from "src/MarginPaymaster.sol";
import {IStakeManager} from "lib/account-abstraction/contracts/interfaces/IStakeManager.sol";
import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
import {console} from "forge-std/console.sol";

Expand Down Expand Up @@ -37,7 +38,6 @@ contract MarginPaymasterTest is Bootstrap {
smartMarginV3Address,
usdcAddress
);
vm.deal(address(this), initialPaymasterBalance);
entryPoint.depositTo{value: initialPaymasterBalance}(
marginPaymasterAddress
);
Expand Down Expand Up @@ -88,7 +88,7 @@ contract MarginPaymasterTest is Bootstrap {
}

/*//////////////////////////////////////////////////////////////
TESTS
MANAGEMENT TESTS
//////////////////////////////////////////////////////////////*/

function testOwner() public {
Expand Down Expand Up @@ -129,6 +129,171 @@ contract MarginPaymasterTest is Bootstrap {
assertTrue(marginPaymaster.authorizers(authorizer));
}

function testSwapUSDCToETH() public {
uint256 initialETHBalance = marginPaymasterAddress.balance;
uint256 amountOutMinimum = 1e18; // 1 ETH minimum output

// Mint USDC to the contract
mintUSDC(address(marginPaymaster), 5000 * 1e6);

// Swap USDC to ETH
marginPaymaster.swapUSDCToETH(amountOutMinimum);

// Check if ETH balance has increased
uint256 finalETHBalance = marginPaymasterAddress.balance;
assertGt(finalETHBalance, initialETHBalance);
}

function testSwapUSDCToETH_onlyOwner() public {
uint256 amountOutMinimum = 1e18; // 1 ETH minimum output

vm.prank(address(0x123)); // some non-owner address
vm.expectRevert("Ownable: caller is not the owner");
marginPaymaster.swapUSDCToETH(amountOutMinimum);
}

function testDepositToEntryPoint() public {
uint256 depositAmount = 1e18; // 1 ETH
vm.deal(marginPaymasterAddress, depositAmount);

// Deposit ETH to EntryPoint
marginPaymaster.depositToEntryPoint(depositAmount);

// Check if the deposit was successful
uint256 entryPointBalance = entryPoint.balanceOf(
address(marginPaymaster)
);
assertEq(entryPointBalance, depositAmount + initialPaymasterBalance);
}

function testDepositToEntryPoint_onlyOwner() public {
uint256 depositAmount = 1e18; // 1 ETH

vm.prank(address(0x123)); // some non-owner address
vm.expectRevert("Ownable: caller is not the owner");
marginPaymaster.depositToEntryPoint(depositAmount);
}

function testStake() public {
uint256 stakeAmount = 1e18; // 1 ETH
uint32 unstakeDelaySec = 3600; // 1 hour

vm.deal(marginPaymasterAddress, stakeAmount);

// Stake ETH in EntryPoint
marginPaymaster.stake(stakeAmount, unstakeDelaySec);

// Check if the stake was successful
IStakeManager.DepositInfo memory depositInfo = entryPoint
.getDepositInfo(address(marginPaymaster));
assertEq(depositInfo.stake, stakeAmount);
assertEq(depositInfo.unstakeDelaySec, unstakeDelaySec);
}

function testStake_onlyOwner() public {
uint256 stakeAmount = 1e18; // 1 ETH
uint32 unstakeDelaySec = 3600; // 1 hour

vm.prank(address(0x123)); // some non-owner address
vm.expectRevert("Ownable: caller is not the owner");
marginPaymaster.stake(stakeAmount, unstakeDelaySec);
}

function testUnlockStake() public {
uint256 stakeAmount = 1e18; // 1 ETH
uint32 unstakeDelaySec = 3600; // 1 hour

vm.deal(marginPaymasterAddress, stakeAmount);

// Stake ETH in EntryPoint
marginPaymaster.stake(stakeAmount, unstakeDelaySec);

// Unlock the stake
marginPaymaster.unlockStake();

// Check if the stake is unlocked
IStakeManager.DepositInfo memory depositInfo = entryPoint
.getDepositInfo(address(marginPaymaster));
assertEq(depositInfo.stake, stakeAmount);
assertEq(depositInfo.unstakeDelaySec, unstakeDelaySec);
assertGt(depositInfo.withdrawTime, 0);
}

function testUnlockStake_onlyOwner() public {
vm.prank(address(0x123)); // some non-owner address
vm.expectRevert("Ownable: caller is not the owner");
marginPaymaster.unlockStake();
}

function testWithdrawStake() public {
uint256 stakeAmount = 1e18; // 1 ETH
uint32 unstakeDelaySec = 3600; // 1 hour
address payable withdrawAddress = payable(address(0x321));

vm.deal(marginPaymasterAddress, stakeAmount);

// Stake ETH in EntryPoint
marginPaymaster.stake(stakeAmount, unstakeDelaySec);

// Unlock the stake
marginPaymaster.unlockStake();

// Fast forward time to allow withdrawal
vm.warp(block.timestamp + unstakeDelaySec);

// Withdraw the stake
marginPaymaster.withdrawStake(withdrawAddress);

// Check if the stake was withdrawn
IStakeManager.DepositInfo memory depositInfo = entryPoint
.getDepositInfo(address(marginPaymaster));
assertEq(depositInfo.stake, 0);
assertEq(depositInfo.unstakeDelaySec, 0);
}

function testWithdrawStake_onlyOwner() public {
address payable withdrawAddress = payable(address(0x321));
vm.prank(withdrawAddress); // some non-owner address
vm.expectRevert("Ownable: caller is not the owner");
marginPaymaster.withdrawStake(withdrawAddress);
}

function testWithdrawTo() public {
uint256 depositAmount = 1e18; // 1 ETH
address payable withdrawAddress = payable(address(0x321));

vm.deal(marginPaymasterAddress, depositAmount);

// Deposit ETH to EntryPoint
marginPaymaster.depositToEntryPoint(depositAmount);

// // // Withdraw from EntryPoint
marginPaymaster.withdrawTo(withdrawAddress, depositAmount);

// // // Check if the withdrawal was successful
uint256 entryPointBalance = entryPoint.balanceOf(
address(marginPaymaster)
);
assertEq(entryPointBalance, initialPaymasterBalance);

// Check if the funds were transferred to the withdrawAddress
uint256 withdrawAddressBalance = withdrawAddress.balance;
assertEq(withdrawAddressBalance, depositAmount);
}

function testWithdrawTo_onlyOwner() public {
uint256 depositAmount = 1e18; // 1 ETH
address payable withdrawAddress = payable(address(0x321));
// Attempt to withdraw from EntryPoint as a non-owner
vm.prank(withdrawAddress); // some non-owner address
vm.expectRevert("Ownable: caller is not the owner");
marginPaymaster.withdrawTo(withdrawAddress, depositAmount);
}

/*//////////////////////////////////////////////////////////////
USER OP TESTS
//////////////////////////////////////////////////////////////*/

function testAccountDeployed() public {
ops.push(userOp);

Expand Down Expand Up @@ -163,7 +328,6 @@ contract MarginPaymasterTest is Bootstrap {
entryPoint.handleOps(ops, backEnd);
}


function testAccountSetup() public {
ops.push(userOp);

Expand All @@ -184,7 +348,7 @@ contract MarginPaymasterTest is Bootstrap {
);
assertEq(usdc.balanceOf(address(this)), 995 * 1e6);
assertEq(usdc.balanceOf(sender), 0);
assertEq(usdc.balanceOf(marginPaymasterAddress), 0);
assertGt(usdc.balanceOf(marginPaymasterAddress), 0);
uint256 colAmount = perpsMarketProxy.getCollateralAmount(
accountId,
sUSDId
Expand All @@ -205,7 +369,7 @@ contract MarginPaymasterTest is Bootstrap {

assertEq(usdc.balanceOf(address(this)), 995 * 1e6);
assertEq(usdc.balanceOf(sender), 0);
assertEq(usdc.balanceOf(marginPaymasterAddress), 0);
assertGt(usdc.balanceOf(marginPaymasterAddress), 0);
uint256 colAmount = perpsMarketProxy.getCollateralAmount(
account.accountId(),
sUSDId
Expand Down

0 comments on commit 9b04be9

Please sign in to comment.