Skip to content

Commit

Permalink
Merge pull request #52 from G7DAO/fix/paymaster-security-of-trx
Browse files Browse the repository at this point in the history 
feat: adding tests and fixing high risk of paymasters
  • Loading branch information
@ogarciarevett
ogarciarevett authored Dec 10, 2023
2 parents 99ca943 + 99d961d commit 44fb576
Show file tree
Hide file tree
Showing 6 changed files with 110 additions and 17 deletions.
21 changes: 19 additions & 2 deletions contracts/paymasters/ERC20Paymaster.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ contract ERC20Paymaster is IPaymaster, Pausable, AccessControl {
bytes32 public constant MANAGER_ROLE = keccak256("MANAGER_ROLE");
bytes32 public constant PAUSER_ROLE = keccak256("MANAGER_ROLE");

mapping(address => bool) public allowedRecipients;

address public allowedERC20Token;
bytes32 public USDCPriceId;
bytes32 public ETHPriceId;
Expand All @@ -37,7 +39,6 @@ contract ERC20Paymaster is IPaymaster, Pausable, AccessControl {

modifier onlyBootloader() {
require(msg.sender == BOOTLOADER_FORMAL_ADDRESS, "Only bootloader can call this method");
// Continue execution if called from the bootloader.
_;
}

Expand Down Expand Up @@ -103,6 +104,12 @@ contract ERC20Paymaster is IPaymaster, Pausable, AccessControl {
bytes32,
Transaction calldata _transaction
) external payable onlyBootloader whenNotPaused returns (bytes4 magic, bytes memory context) {
// check the target and function of the transaction: _transaction

address recipient = address(uint160(_transaction.to));

require(allowedRecipients[recipient], "Invalid recipient");

// By default we consider the transaction as accepted.
magic = PAYMASTER_VALIDATION_SUCCESS_MAGIC;
require(_transaction.paymasterInput.length >= 4, "The standard paymaster input must be at least 4 bytes long");
Expand Down Expand Up @@ -198,14 +205,24 @@ contract ERC20Paymaster is IPaymaster, Pausable, AccessControl {
_unpause();
}

function addRecipient(address _recipient) public onlyRole(DEFAULT_ADMIN_ROLE) {
require(_recipient != address(0), "NonAddressZero");
allowedRecipients[_recipient] = true;
}

function removeRecipient(address _recipient) public onlyRole(DEFAULT_ADMIN_ROLE) {
require(_recipient != address(0), "NonAddressZero");
allowedRecipients[_recipient] = false;
}

function withdrawETH(address payable _to) external onlyRole(MANAGER_ROLE) {
// send paymaster funds to the owner
uint256 balance = address(this).balance;
(bool success, ) = _to.call{ value: balance }("");
require(success, "Failed to withdraw funds from paymaster.");
}

function withDrawERC20(address _to, uint256 _amount) external onlyRole(MANAGER_ROLE) {
function withdrawERC20(address _to, uint256 _amount) external onlyRole(MANAGER_ROLE) {
// send paymaster funds to the owner
IERC20 token = IERC20(allowedERC20Token);
uint256 balance = token.balanceOf(address(this));
Expand Down
16 changes: 16 additions & 0 deletions contracts/paymasters/GeneralPaymaster.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@ contract GasLessOpenMintPaymasterETH is IPaymaster, AccessControl, Pausable {
bytes32 public constant MANAGER_ROLE = keccak256("MANAGER_ROLE");
bytes32 public constant PAUSER_ROLE = keccak256("MANAGER_ROLE");

mapping(address => bool) public allowedRecipients;

event PaymasterPayment(address indexed from, uint256 amount);

constructor() {
Expand All @@ -39,6 +41,10 @@ contract GasLessOpenMintPaymasterETH is IPaymaster, AccessControl, Pausable {
bytes32,
Transaction calldata _transaction
) external payable onlyBootloader whenNotPaused returns (bytes4 magic, bytes memory context) {

address recipient = address(uint160(_transaction.to));
require(allowedRecipients[recipient], "Invalid recipient");

// By default we consider the transaction as accepted.
magic = PAYMASTER_VALIDATION_SUCCESS_MAGIC;
require(_transaction.paymasterInput.length >= 4, "The standard paymaster input must be at least 4 bytes long");
Expand Down Expand Up @@ -84,5 +90,15 @@ contract GasLessOpenMintPaymasterETH is IPaymaster, AccessControl, Pausable {
require(success, "Failed to withdraw funds from paymaster.");
}

function addRecipient(address _recipient) public onlyRole(DEFAULT_ADMIN_ROLE) {
require(_recipient != address(0), "NonAddressZero");
allowedRecipients[_recipient] = true;
}

function removeRecipient(address _recipient) public onlyRole(DEFAULT_ADMIN_ROLE) {
require(_recipient != address(0), "NonAddressZero");
allowedRecipients[_recipient] = false;
}

receive() external payable {}
}
5 changes: 3 additions & 2 deletions helpers/zkUtils.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,14 @@ import { Contract, Wallet, Provider } from 'zksync2-js';
import { Deployer } from '@matterlabs/hardhat-zksync-deploy';
import * as hre from 'hardhat';
import * as ethers from 'ethers';
import { log } from '@helpers/logger';

async function deployContract(deployer: Deployer, contract: string, params: any[]): Promise<Contract> {
const artifact = await deployer.loadArtifact(contract);

const deploymentFee = await deployer.estimateDeployFee(artifact, params);
const parsedFee = ethers.formatEther(deploymentFee.toString());
console.log(`${contract} deployment is estimated to cost ${parsedFee} ETH`);
log(`${contract} deployment is estimated to cost ${parsedFee} ETH`);

return await deployer.deploy(artifact, params);
}
Expand All @@ -22,7 +23,7 @@ async function fundAccount(wallet: Wallet, address: string, amount: string) {
})
).wait();

console.log(`Account ${address} funded with ${amount}`);
log(`Address ${address} funded with ${amount}`);
}

function setupDeployer(url: string, privateKey: string): [Provider, Wallet, Deployer] {
Expand Down
11 changes: 11 additions & 0 deletions metadata.example.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
{
"image": "https://summon.mypinata.cloud/ipfs/QmUdLtY99nFvF1yyKyvyNXpcUqqha323yGMn18oezBdBDt",
"name": "LIBERTAS OMNIBUS nft test open capsule",
"description": "OPEN CAPSULE",
"attributes": [
{
"trait_type": "Revealed",
"value": "YES"
}
]
}
8 changes: 7 additions & 1 deletion scripts/libraries/selectors.ts
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
// get function selectors from ABI

import { Contract } from 'ethers';
import { Contract, dataSlice, keccak256, toUtf8Bytes } from 'ethers';

export enum FacetCutAction {
Add,
Expand All @@ -20,6 +20,12 @@ export function getSelectors(contract: Contract): string[] {
return selectors;
}

export function getFunctionSignature(functionABI: string): string {
const functionHash = keccak256(toUtf8Bytes(functionABI));

return dataSlice(functionHash, 0, 4);
}

export function getSelectorsFacet(contract: Contract) {
const signatures = Object.keys(contract.interface.functions);
// log function names and signatures
Expand Down
66 changes: 54 additions & 12 deletions test/zkSync/erc20Paymaster.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import hardhatConfig from '../../hardhat.config';
import zkSyncConfig from '../../zkSync.config';
import { deployContract, fundAccount, setupDeployer } from '../../helpers/zkUtils';
import { Deployer } from '@matterlabs/hardhat-zksync-deploy';
import { getFunctionSignature } from '../../scripts/libraries/selectors';
dotenv.config();

// load wallet private key from env file
Expand Down Expand Up @@ -39,8 +40,6 @@ describe('ERC20Paymaster', function () {
innerInput: new Uint8Array(),
});

log(`Minting 1 erc721 for empty wallet via paymaster...`);

const mintTrx = await mockERC721Soulbound.connect(user).mint(user.address, {
customData: {
gasPerPubdata: utils.DEFAULT_GAS_PER_PUBDATA_LIMIT,
Expand All @@ -49,7 +48,6 @@ describe('ERC20Paymaster', function () {
});

await mintTrx.wait();
log('paymaster params:', paymasterParams);
}

beforeEach(async function () {
Expand All @@ -59,13 +57,28 @@ describe('ERC20Paymaster', function () {
[provider, minterAccount, deployer] = setupDeployer(deployUrl, PRIVATE_KEY);

const emptyWallet = Wallet.createRandom();
log(`Empty wallet's address: ${emptyWallet.address}`);

playerAccount = new Wallet(emptyWallet.privateKey, provider);

token = await deployContract(deployer, 'MockUSDC', ['MyUSDC', 'mUSDC', 18]);
await token.waitForDeployment();

mockERC721Soulbound = await deployContract(deployer, 'Mock721Soulbound', []);
paymaster = await deployContract(deployer, 'ERC20Paymaster', [await token.getAddress(), USDC_PRICE_ID, ETH_PRICE_ID, PYTH_ORACLE_ADDRESS]);
await mockERC721Soulbound.waitForDeployment();

paymaster = await deployContract(deployer, 'ERC20Paymaster', [
await token.getAddress(),
USDC_PRICE_ID,
ETH_PRICE_ID,
PYTH_ORACLE_ADDRESS,
]);

await paymaster.waitForDeployment();

log(`Empty wallet's address: ${emptyWallet.address}`);
log('Token deployed', await token.getAddress());
log('mockERC721Soulbound deployed', await mockERC721Soulbound.getAddress());
log('paymaster deployed', await paymaster.getAddress());

// fund paymaster
await fundAccount(minterAccount, await paymaster.getAddress(), '3');
Expand All @@ -82,16 +95,15 @@ describe('ERC20Paymaster', function () {

const initialPaymasterBalance = await provider.getBalance(paymasterAddress);

const addRecipientTx = await paymaster.addRecipient(await mockERC721Soulbound.getAddress());
await addRecipientTx.wait();

await executeMintTransaction(playerAccount);

const finalETHBalance = await provider.getBalance(playerAccount);
const finalUserTokenBalance = await token.balanceOf(playerAccount.address);
const finalPaymasterBalance = await provider.getBalance(paymasterAddress);

console.log('finalETHBalance ETH balance:', finalETHBalance.toString());
console.log('finalUserTokenBalance ETH balance:', finalUserTokenBalance.toString());
console.log('finalPaymasterBalance ETH balance:', finalPaymasterBalance.toString());

expect(await mockERC721Soulbound.balanceOf(playerAccount.address)).to.equal(1);
expect(initialPaymasterBalance > finalPaymasterBalance).to.be.true;
expect(userInitialETHBalance).to.eql(finalETHBalance);
Expand All @@ -101,19 +113,49 @@ describe('ERC20Paymaster', function () {
it('should allow owner to withdraw all funds', async function () {
const paymasterAddress = await paymaster.getAddress();

const tx = await paymaster.connect(minterAccount).withdraw(playerAccount.address);
const tx = await paymaster.connect(minterAccount).withdrawETH(playerAccount.address);
await tx.wait();

const finalContractBalance = await provider.getBalance(paymasterAddress);

expect(finalContractBalance).to.eql(BigInt(0));
});

it('should allow owner to withdraw all erc20 funds', async function () {
const paymasterAddress = await paymaster.getAddress();

const success = await token.mint(playerAccount.address, BigInt(3));
await success.wait();

const managerInitialTokenBalance = await token.balanceOf(minterAccount.address);
const addRecipientTx = await paymaster.addRecipient(await mockERC721Soulbound.getAddress());
await addRecipientTx.wait();

// +1n ERC20 to the paymaster
await executeMintTransaction(playerAccount);
// +1n ERC20 to the paymaster
await executeMintTransaction(playerAccount);

const paymasterERC20AfterTrx = await token.balanceOf(paymasterAddress);

log('initialPaymasterERC20Balance', paymasterERC20AfterTrx.toString());

const tx = await paymaster.connect(minterAccount).withdrawERC20(minterAccount.address, 1n);
await tx.wait();
const paymasterAfterWithdrawERC20Balance = await token.balanceOf(paymasterAddress);

const finalManagerTokenBalance = await token.balanceOf(minterAccount.address);
expect(paymasterERC20AfterTrx).to.greaterThan(0n);
expect(finalManagerTokenBalance).to.eql(1n);
expect(managerInitialTokenBalance).to.eql(0n);
expect(paymasterAfterWithdrawERC20Balance).to.eql(1n);
});

it('should prevent non-owners from withdrawing funds', async function () {
try {
await paymaster.connect(playerAccount).withdraw(minterAccount.address);
await paymaster.connect(playerAccount).withdrawETH(minterAccount.address);
} catch (e) {
expect(e.message).to.include('Ownable: caller is not the owner');
expect(e.message).to.include('AccessControl: a..');
}
});
});

0 comments on commit 44fb576

Please sign in to comment.