Skip to content

Commit

Permalink
feat: add setter for user op max cost + polishes
Browse files Browse the repository at this point in the history
  • Loading branch information
fedealconada committed Jan 29, 2024
1 parent ec434d5 commit ded38b9
Show file tree
Hide file tree
Showing 7 changed files with 80 additions and 67 deletions.
4 changes: 3 additions & 1 deletion src/paymasters/SponsorPaymaster.sol
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ contract SponsorPaymaster is Initializable, BasePaymaster, UUPSUpgradeable, Reen
uint256 public constant RATE_LIMIT_PERIOD = 1 minutes;
uint256 public constant RATE_LIMIT_THRESHOLD_TOTAL = 50;

uint256 public userOpMaxCost = 0.03 ether;
uint256 public userOpMaxCost;
IKintoAppRegistry public override appRegistry;

mapping(address => uint256) public balances;
Expand Down Expand Up @@ -64,6 +64,8 @@ contract SponsorPaymaster is Initializable, BasePaymaster, UUPSUpgradeable, Reen
function initialize(address _owner) external virtual initializer {
__UUPSUpgradeable_init();
_transferOwnership(_owner);

userOpMaxCost = 0.03 ether;
// unlocks owner
unlockBlock[_owner] = block.number;
}
Expand Down
28 changes: 15 additions & 13 deletions src/wallet/KintoWallet.sol
Original file line number Diff line number Diff line change
Expand Up @@ -387,26 +387,28 @@ contract KintoWallet is Initializable, BaseAccount, TokenCallbackHandler, IKinto
}

function _resetSigners(address[] calldata newSigners, uint8 _policy) internal {
require(newSigners.length > 0 && newSigners.length <= MAX_SIGNERS, "KW-rs: invalid array");
require(newSigners.length > 0 && newSigners.length <= MAX_SIGNERS, "KW-rs: signers exceed max limit");
require(newSigners[0] != address(0) && kintoID.isKYC(newSigners[0]), "KW-rs: KYC Required");
require(
newSigners.length == 1 || (newSigners.length == 2 && newSigners[0] != newSigners[1])
|| (
newSigners.length == 3 && (newSigners[0] != newSigners[1]) && (newSigners[1] != newSigners[2])
&& newSigners[0] != newSigners[2]
),
"duplicate owners"
);

// ensure no duplicate signers
for (uint256 i = 0; i < newSigners.length; i++) {
for (uint256 j = i + 1; j < newSigners.length; j++) {
require(newSigners[i] != newSigners[j], "KW-rs: duplicate signers");
}
}

// ensure all signers are valid
for (uint256 i = 0; i < newSigners.length; i++) {
require(newSigners[i] != address(0), "KW-rs: invalid signer address");
}

emit SignersChanged(owners, newSigners);
owners = newSigners;
emit SignersChanged(newSigners, owners);
// Change policy if needed.

// change policy if needed
require(_policy == 1 || newSigners.length > 1, "KW-rs: invalid policy for single signer");
if (_policy != signerPolicy) {
setSignerPolicy(_policy);
} else {
require(_policy == 1 || newSigners.length > 1, "KW-rs: invalid policy");
}
}

Expand Down
24 changes: 14 additions & 10 deletions src/wallet/KintoWalletFactory.sol
Original file line number Diff line number Diff line change
Expand Up @@ -264,18 +264,18 @@ contract KintoWalletFactory is Initializable, UUPSUpgradeable, OwnableUpgradeabl
(newImplementation);
}

function _preventCreationBytecode(bytes calldata _bytes) internal view {
// Prevent direct deployment of KintoWallet contracts
function _preventWalletDeployment(bytes calldata _bytecode) internal view {
bytes memory walletInitCode = type(SafeBeaconProxy).creationCode;
if (_bytes.length > walletInitCode.length + 32) {
// Make sure this beacon cannot be called with creation code
if (_bytecode.length > walletInitCode.length + 32) {
uint256 offset = 12 + walletInitCode.length;
address beaconAddress = address(0);
bytes memory slice = _bytes[offset:offset + 20];

// extract beacon address directly from calldata
address beaconAddress;
bytes memory slice = _bytecode[offset:offset + 20];
assembly {
beaconAddress := mload(add(slice, 20))
}
// Compare the extracted beacon address with the disallowed address

require(beaconAddress != address(beacon), "Direct KintoWallet deployment not allowed");
}
}
Expand All @@ -284,10 +284,14 @@ contract KintoWalletFactory is Initializable, UUPSUpgradeable, OwnableUpgradeabl
internal
returns (address)
{
require(amount == msg.value, "amount mismatch");
_preventCreationBytecode(bytecode);
require(amount == msg.value, "Amount mismatch");
require(bytecode.length > 0, "Bytecode is empty");
_preventWalletDeployment(bytecode);

// deploy the contract using `CREATE2`
address created = Create2.deploy(amount, salt, bytecode);
// Assign ownership to the deployer if needed

// assign ownership to newOwner if the contract is Ownable
try OwnableUpgradeable(created).owner() returns (address owner) {
if (owner == address(this)) {
OwnableUpgradeable(created).transferOwnership(newOwner);
Expand Down
43 changes: 17 additions & 26 deletions test/KintoID.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -180,12 +180,12 @@ contract KintoIDTest is KYCSignature, AATestScaffolding, UserOp {

/* ============ Burn tests ============ */

function testBurnKYC_RevertWhen_UsingBurn() public {
function testBurn_RevertWhen_UsingBurn() public {
vm.expectRevert("Use burnKYC instead");
_kintoID.burn(1);
}

function test_RevertWhen_BurnIsCalled() public {
function testBurn_RevertWhen_BurnIsCalled() public {
approveKYC(_kycProvider, _user, _userPk);
uint256 tokenIdx = _kintoID.tokenOfOwnerByIndex(_user, 0);
vm.prank(_user);
Expand All @@ -194,49 +194,40 @@ contract KintoIDTest is KYCSignature, AATestScaffolding, UserOp {
}

function testBurnKYC() public {
approveKYC(_kycProvider, _user, _userPk);

IKintoID.SignatureData memory sigdata = _auxCreateSignature(_kintoID, _user, _userPk, block.timestamp + 1000);
uint16[] memory traits = new uint16[](1);
traits[0] = 1;
vm.startPrank(_kycProvider);
_kintoID.mintIndividualKyc(sigdata, traits);
assertEq(_kintoID.isKYC(_user), true);
sigdata = _auxCreateSignature(_kintoID, _user, _userPk, block.timestamp + 1000);
vm.prank(_kycProvider);
_kintoID.burnKYC(sigdata);
assertEq(_kintoID.balanceOf(_user), 0);
}

function testOnlyProviderCanBurnKYC() public {
function testBurnKYC_WhenCallerIsNotProvider() public {
approveKYC(_kycProvider, _user, _userPk);

IKintoID.SignatureData memory sigdata = _auxCreateSignature(_kintoID, _user, _userPk, block.timestamp + 1000);
uint16[] memory traits = new uint16[](1);
traits[0] = 1;
vm.startPrank(_kycProvider);
_kintoID.mintIndividualKyc(sigdata, traits);
assertEq(_kintoID.isKYC(_user), true);
sigdata = _auxCreateSignature(_kintoID, _user, _userPk, block.timestamp + 1000);
vm.stopPrank();
vm.startPrank(_user);
vm.expectRevert("Invalid Provider");
vm.startPrank(_user);
_kintoID.burnKYC(sigdata);
}

function testBurnFailsWithoutMinting() public {
function testBurnKYC_WhenUserIsNotKYCd() public {
IKintoID.SignatureData memory sigdata = _auxCreateSignature(_kintoID, _user, _userPk, block.timestamp + 1000);
vm.startPrank(_kycProvider);
vm.expectRevert("Nothing to burn");
vm.prank(_kycProvider);
_kintoID.burnKYC(sigdata);
}

function testBurningTwiceFails() public {
function testBurnKYC_WhenBurningTwice() public {
approveKYC(_kycProvider, _user, _userPk);

IKintoID.SignatureData memory sigdata = _auxCreateSignature(_kintoID, _user, _userPk, block.timestamp + 1000);
uint16[] memory traits = new uint16[](1);
traits[0] = 1;
vm.startPrank(_kycProvider);
_kintoID.mintIndividualKyc(sigdata, traits);
assertEq(_kintoID.isKYC(_user), true);
sigdata = _auxCreateSignature(_kintoID, _user, _userPk, block.timestamp + 1000);
vm.prank(_kycProvider);
_kintoID.burnKYC(sigdata);

sigdata = _auxCreateSignature(_kintoID, _user, _userPk, block.timestamp + 1000);
vm.expectRevert("Nothing to burn");
vm.prank(_kycProvider);
_kintoID.burnKYC(sigdata);
}

Expand Down
27 changes: 20 additions & 7 deletions test/KintoWalletFactory.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -165,7 +165,26 @@ contract KintoWalletFactoryTest is SharedSetup {
assertEq(Counter(created).count(), 1);
}

function testDeployContract_RevertWhen_CreateWalletThroughDeploy() public {
function testDeployContract_RevertWhen_SenderNotKYCd() public {
vm.prank(_user2);
vm.expectRevert("KYC required");
_walletFactory.deployContract(_owner, 0, abi.encodePacked(type(Counter).creationCode), bytes32(0));
}

function testDeployContract_RevertWhen_AmountMismatch() public {
vm.deal(_owner, 1 ether);
vm.prank(_owner);
vm.expectRevert("Amount mismatch");
_walletFactory.deployContract(_owner, 1 ether + 1, abi.encodePacked(type(Counter).creationCode), bytes32(0));
}

function testDeployContract_RevertWhen_ZeroBytecode() public {
vm.expectRevert("Bytecode is empty");
vm.prank(_owner);
_walletFactory.deployContract(_owner, 0, bytes(""), bytes32(0));
}

function testDeployContract_RevertWhen_CreateWallet() public {
bytes memory initialize = abi.encodeWithSelector(IKintoWallet.initialize.selector, _owner, _owner);
bytes memory bytecode = abi.encodePacked(
type(SafeBeaconProxy).creationCode, abi.encode(address(_walletFactory.beacon()), initialize)
Expand All @@ -175,12 +194,6 @@ contract KintoWalletFactoryTest is SharedSetup {
_walletFactory.deployContract(_owner, 0, bytecode, bytes32(0));
}

function testDeployContract_RevertWhen_SenderNotKYCd() public {
vm.prank(_user2);
vm.expectRevert("KYC required");
_walletFactory.deployContract(_owner, 0, abi.encodePacked(type(Counter).creationCode), bytes32(0));
}

function testFundWallet() public {
vm.prank(_owner);
_walletFactory.fundWallet{value: 1e18}(payable(address(_kintoWallet)));
Expand Down
1 change: 1 addition & 0 deletions test/SponsorPaymaster.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ contract SponsorPaymasterTest is SharedSetup {
function testUp() public override {
super.testUp();
assertEq(_paymaster.COST_OF_POST(), 200_000);
assertEq(_paymaster.userOpMaxCost(), 0.03 ether);
}

/* ============ Events ============ */
Expand Down
20 changes: 10 additions & 10 deletions test/wallet/policy/Policy.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ import "../../SharedSetup.t.sol";
contract ResetSignerTest is SharedSetup {
/* ============ Signers & Policy tests ============ */

function testAddingOneSigner() public {
function testResetSigners_WhenAddingOneSigner() public {
vm.startPrank(_owner);
address[] memory owners = new address[](2);
owners[0] = _owner;
Expand All @@ -29,7 +29,7 @@ contract ResetSignerTest is SharedSetup {
vm.stopPrank();
}

function test_RevertWhen_DuplicateSigner() public {
function testResetSigners_RevertWhen_DuplicateSigner() public {
address[] memory owners = new address[](2);
owners[0] = _owner;
owners[1] = _owner;
Expand All @@ -52,10 +52,10 @@ contract ResetSignerTest is SharedSetup {
);
vm.recordLogs();
_entryPoint.handleOps(userOps, payable(_owner));
assertRevertReasonEq("duplicate owners");
assertRevertReasonEq("KW-rs: duplicate signers");
}

function test_RevertWhen_WithEmptyArray() public {
function testResetSigners_RevertWhen_EmptyArray() public {
address[] memory owners = new address[](0);

UserOperation memory userOp = _createUserOperation(
Expand All @@ -79,7 +79,7 @@ contract ResetSignerTest is SharedSetup {
// fixme: assertRevertReasonEq(stdError.indexOOBError)F;
}

function test_RevertWhen_WithManyOwners() public {
function testResetSigners_RevertWhen_ManyOwners() public {
address[] memory owners = new address[](4);
owners[0] = _owner;
owners[1] = _user;
Expand All @@ -104,10 +104,10 @@ contract ResetSignerTest is SharedSetup {
);
vm.recordLogs();
_entryPoint.handleOps(userOps, payable(_owner));
assertRevertReasonEq("KW-rs: invalid array");
assertRevertReasonEq("KW-rs: signers exceed max limit");
}

function test_RevertWhen_WithoutKYCSigner() public {
function testResetSigners_RevertWhen_WithoutKYCSigner() public {
address[] memory owners = new address[](1);
owners[0] = _user;

Expand All @@ -125,7 +125,7 @@ contract ResetSignerTest is SharedSetup {
_entryPoint.handleOps(userOps, payable(_owner));
}

function testChangingPolicyWithTwoSigners() public {
function testResetSigners_WhenChangingPolicy_WhenTwoSigners() public {
address[] memory owners = new address[](2);
owners[0] = _owner;
owners[1] = _user;
Expand All @@ -148,7 +148,7 @@ contract ResetSignerTest is SharedSetup {
assertEq(_kintoWallet.signerPolicy(), _kintoWallet.ALL_SIGNERS());
}

function testChangingPolicyWithThreeSigners() public {
function testResetSigners_WhenChangingPolicy_WhenThreeSigners() public {
vm.startPrank(_owner);
address[] memory owners = new address[](3);
owners[0] = _owner;
Expand All @@ -173,7 +173,7 @@ contract ResetSignerTest is SharedSetup {
}

// todo: separate into 2 different tests
function test_RevertWhen_ChangingPolicyWithoutRightSigners() public {
function testResetSigners_RevertWhen_ChangingPolicy_WhenNotRightSigners() public {
address[] memory owners = new address[](2);
owners[0] = _owner;
owners[1] = _user;
Expand Down

0 comments on commit ded38b9

Please sign in to comment.