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
fedealconada committed Jan 29, 2024
1 parent ec434d5 commit 0aaf9e1
Show file tree
Hide file tree
Showing 7 changed files with 98 additions and 41 deletions.
4 changes: 3 additions & 1 deletion src/paymasters/SponsorPaymaster.sol
View file
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
View file
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
View file
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
27 changes: 20 additions & 7 deletions test/KintoWalletFactory.t.sol
View file
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
View file
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
View file
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
35 changes: 35 additions & 0 deletions test/wallet/validateSignature/ValidateSignature.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -810,4 +810,39 @@ contract ValidateSignatureTest is SharedSetup {
)
);
}

// special case 2: requiredSigners == 3, owners.length == 3 and only owners[0] is KYCd
function testValidateSignature_SpecialCase3() public {
assertEq(_kintoID.isKYC(_user), false);
assertEq(_kintoID.isKYC(_user2), false);

// reset signers & change policy
address[] memory owners = new address[](3);
owners[0] = _owner;
owners[1] = _user;
owners[2] = _user2;
resetSigners(owners, _kintoWallet.ALL_SIGNERS());

// create user op with owners 1 as signer
privateKeys = new uint256[](3);
privateKeys[0] = _ownerPk;
privateKeys[1] = _userPk;
privateKeys[2] = _user2Pk;

UserOperation memory userOp = _createUserOperation(
address(_kintoWallet),
address(counter),
_kintoWallet.getNonce(),
privateKeys,
abi.encodeWithSignature("increment()"),
address(_paymaster)
);

assertEq(
SIG_VALIDATION_FAILED,
KintoWalletHarness(payable(address(_kintoWallet))).exposed_validateSignature(
userOp, _entryPoint.getUserOpHash(userOp)
)
);
}
}

0 comments on commit 0aaf9e1

Please sign in to comment.