chore: fix: add KYC checks on addDepositFor and setter for user op ma…

…x cost

script/deploy.sol

@@ -173,7 +173,8 @@ contract KintoInitialDeployScript is Create2Helper, ArtifactsReader {
             _sponsorPaymaster = SponsorPaymaster(address(_proxy));
             console.log("Paymaster proxy deployed at ", address(_sponsorPaymaster));
             // Initialize proxy
-            _sponsorPaymaster.initialize(address(msg.sender));
+            // TODO: deploy registry and pass it to the paymaster
+            _sponsorPaymaster.initialize(address(msg.sender), IKintoAppRegistry(address(0)), _kintoIDv1);
         }
 
         // KYC Viewer

src/interfaces/ISponsorPaymaster.sol

@@ -2,6 +2,7 @@
 pragma solidity ^0.8.18;
 
 import {IKintoAppRegistry} from "./IKintoAppRegistry.sol";
+import {IKintoID} from "./IKintoID.sol";
 
 interface ISponsorPaymaster {
     /* ============ Structs ============ */
@@ -14,7 +15,7 @@ interface ISponsorPaymaster {
 
     /* ============ State Change ============ */
 
-    function initialize(address owner) external;
+    function initialize(address owner, IKintoAppRegistry _appRegistry, IKintoID _kintoID) external;
 
     function setAppRegistry(address _appRegistry) external;
 

src/paymasters/SponsorPaymaster.sol

@@ -11,6 +11,7 @@ import "@aa/core/BasePaymaster.sol";
 import "../interfaces/ISponsorPaymaster.sol";
 import "../interfaces/IKintoAppRegistry.sol";
 import "../interfaces/IKintoWallet.sol";
+import "../interfaces/IKintoID.sol";
 
 /**
  * An ETH-based paymaster that accepts ETH deposits
@@ -47,6 +48,7 @@ contract SponsorPaymaster is Initializable, BasePaymaster, UUPSUpgradeable, Reen
     mapping(address => ISponsorPaymaster.RateLimitData) public globalRateLimit;
 
     IKintoAppRegistry public override appRegistry;
+    IKintoID public kintoID;
 
     // ========== Constructor & Upgrades ============
 
@@ -59,11 +61,17 @@ contract SponsorPaymaster is Initializable, BasePaymaster, UUPSUpgradeable, Reen
      * a new implementation of SimpleAccount must be deployed with the new EntryPoint address, then upgrading
      * the implementation by calling `upgradeTo()`
      */
-    function initialize(address _owner) external virtual initializer {
+    function initialize(address _owner, IKintoAppRegistry _appRegistry, IKintoID _kintoID)
+        external
+        virtual
+        initializer
+    {
         __UUPSUpgradeable_init();
         _transferOwnership(_owner);
-        // unlocks owner
-        unlockBlock[_owner] = block.number;
+
+        kintoID = _kintoID;
+        appRegistry = _appRegistry;
+        unlockBlock[_owner] = block.number; // unlocks owner
     }
 
     /**
@@ -76,16 +84,6 @@ contract SponsorPaymaster is Initializable, BasePaymaster, UUPSUpgradeable, Reen
         (newImplementation);
     }
 
-    /**
-     * @dev Set the app registry
-     * @param _appRegistry address of the app registry
-     */
-    function setAppRegistry(address _appRegistry) external override onlyOwner {
-        require(_appRegistry != address(0) && _appRegistry != address(appRegistry), "SP: appRegistry cannot be 0");
-        emit AppRegistrySet(_appRegistry, address(appRegistry));
-        appRegistry = IKintoAppRegistry(_appRegistry);
-    }
-
     // ========== Deposit Mgmt ============
 
     /**
@@ -98,6 +96,9 @@ contract SponsorPaymaster is Initializable, BasePaymaster, UUPSUpgradeable, Reen
      */
     function addDepositFor(address account) external payable override {
         require(msg.value > 0, "SP: requires a deposit");
+        require(kintoID.isKYC(msg.sender), "SP: sender KYC required");
+        if (account.code.length == 0 && !kintoID.isKYC(account)) revert("SP: account KYC required");
+
         // sender must have approval for the paymaster
         balances[account] += msg.value;
         if (msg.sender == account) {
@@ -173,9 +174,22 @@ contract SponsorPaymaster is Initializable, BasePaymaster, UUPSUpgradeable, Reen
     }
 
     /**
-     * Validate the request from the sender to fund it.
-     * The sender should have enough txs and gas left to be gasless.
-     * The contract developer funds the contract for its users and rate limits the app.
+     * @dev Set the app registry
+     * @param _newRegistry address of the app registry
+     */
+    function setAppRegistry(address _newRegistry) external override onlyOwner {
+        require(_newRegistry != address(0), "SP: new registry cannot be 0");
+        require(_newRegistry != address(appRegistry), "SP: new registry cannot be the same");
+        emit AppRegistrySet(address(appRegistry), _newRegistry);
+        appRegistry = IKintoAppRegistry(_newRegistry);
+    }
+
+    /* =============== AA overrides ============= */
+
+    /**
+     * @notice Validates the request from the sender to fund it.
+     * @dev sender should have enough txs and gas left to be gasless.
+     * @dev contract developer funds the contract for its users and rate limits the app.
      */
     function _validatePaymasterUserOp(UserOperation calldata userOp, bytes32 userOpHash, uint256 maxCost)
         internal

test/EngenCredits.t.sol

@@ -30,8 +30,8 @@ contract EngenCreditsTest is UserOp, AATestScaffolding {
         _owner.transfer(1e18);
         vm.stopPrank();
         deployAAScaffolding(_owner, 1, _kycProvider, _recoverer);
-        fundSponsorForApp(address(_engenCredits));
-        fundSponsorForApp(address(_kintoWallet));
+        fundSponsorForApp(_owner, address(_engenCredits));
+        fundSponsorForApp(_owner, address(_kintoWallet));
         vm.startPrank(_owner);
         _kintoAppRegistry.registerApp(
             "engen credits", address(_engenCredits), new address[](0), [uint256(0), uint256(0), uint256(0), uint256(0)]

test/KintoWalletFactory.t.sol

@@ -16,8 +16,8 @@ import "../src/wallet/KintoWallet.sol";
 import "./SharedSetup.t.sol";
 
 contract KintoWalletUpgrade is KintoWallet {
-    constructor(IEntryPoint _entryPoint, IKintoID _kintoID, IKintoAppRegistry _kintoAppRegistry)
-        KintoWallet(_entryPoint, _kintoID, _kintoAppRegistry)
+    constructor(IEntryPoint _entryPoint, IKintoID _kintoIDv1, IKintoAppRegistry _kintoAppRegistry)
+        KintoWallet(_entryPoint, _kintoIDv1, _kintoAppRegistry)
     {}
 
     function walletFunction() public pure returns (uint256) {
@@ -148,8 +148,8 @@ contract KintoWalletFactoryTest is SharedSetup {
     }
 
     function testWhitelistedSignerCanFundWallet() public {
+        fundSponsorForApp(_owner, address(_kintoWallet));
         vm.startPrank(_owner);
-        fundSponsorForApp(address(_kintoWallet));
         uint256 nonce = _kintoWallet.getNonce();
         address[] memory funders = new address[](1);
         funders[0] = _funder;

test/SharedSetup.t.sol

@@ -34,7 +34,7 @@ contract SharedSetup is UserOp, AATestScaffolding {
         deployAAScaffolding(_owner, 1, _kycProvider, _recoverer);
 
         // add paymaster to _kintoWallet
-        fundSponsorForApp(address(_kintoWallet));
+        fundSponsorForApp(_owner, address(_kintoWallet));
 
         // all tests will use 1 private key (_ownerPk) unless otherwise specified
         privateKeys = new uint256[](1);
@@ -46,7 +46,7 @@ contract SharedSetup is UserOp, AATestScaffolding {
 
         registerApp(_owner, "test", address(counter));
         whitelistApp(address(counter));
-        fundSponsorForApp(address(counter));
+        fundSponsorForApp(_owner, address(counter));
     }
 
     function testUp() public virtual {

test/SponsorPaymastExploit.t.sol

@@ -58,8 +58,8 @@ contract SponsorPaymasterExploitTest is MyOpCreator, AATestScaffolding {
         _owner.transfer(1e18);
         vm.stopPrank();
         deployAAScaffolding(_owner, 1, _kycProvider, _recoverer);
-        fundSponsorForApp(address(_engenCredits));
-        fundSponsorForApp(address(_kintoWallet));
+        fundSponsorForApp(_owner, address(_engenCredits));
+        fundSponsorForApp(_owner, address(_kintoWallet));
     }
 
     function testExploit() public {
@@ -68,7 +68,7 @@ contract SponsorPaymasterExploitTest is MyOpCreator, AATestScaffolding {
         Counter counter = new Counter();
         assertEq(counter.count(), 0);
         vm.stopPrank();
-        fundSponsorForApp(address(counter));
+        fundSponsorForApp(_owner, address(counter));
         vm.startPrank(_owner);
         // Let's send a transaction to the counter contract through our wallet
         uint256 nonce = _kintoWallet.getNonce();

test/SponsorPaymaster.t.sol

@@ -49,59 +49,108 @@ contract SponsorPaymasterTest is SharedSetup {
         assertEq(_newImplementation.newFunction(), 1);
     }
 
-    function testUpgrade_RevertWhen_CallerIsNotOwner() public {
+    function testUpgradeTo_RevertWhen_CallerIsNotOwner() public {
         SponsorPaymasterUpgrade _newImplementation = new SponsorPaymasterUpgrade(_entryPoint, _owner);
         vm.expectRevert("SP: not owner");
         _paymaster.upgradeTo(address(_newImplementation));
     }
 
     /* ============ Deposit & Stake ============ */
 
-    function testOwnerCanDepositStakeAndWithdraw() public {
-        vm.startPrank(_owner);
+    function testAddDepositFor_WhenAccountIsEOA_WhenAccountIsKYCd() public {
         uint256 balance = address(_owner).balance;
+        vm.prank(_owner);
         _paymaster.addDepositFor{value: 5e18}(address(_owner));
         assertEq(address(_owner).balance, balance - 5e18);
-        _paymaster.unlockTokenDeposit();
-        vm.roll(block.number + 1);
-        _paymaster.withdrawTokensTo(address(_owner), 5e18);
-        assertEq(address(_owner).balance, balance);
-        vm.stopPrank();
+        assertEq(_paymaster.balances(_owner), 5e18);
     }
 
-    function testUserCanDepositStakeAndWithdraw() public {
-        vm.startPrank(_user);
+    function testAddDepositFor_WhenAccountIsContract() public {
+        uint256 balance = address(_owner).balance;
+        vm.prank(_owner);
+        _paymaster.addDepositFor{value: 5e18}(address(_owner));
+        assertEq(address(_owner).balance, balance - 5e18);
+        assertEq(_paymaster.balances(_owner), 5e18);
+    }
+
+    function testAddDepositFor_RevertWhen_ZeroValue() public {
+        vm.expectRevert("SP: requires a deposit");
+        vm.prank(_owner);
+        _paymaster.addDepositFor{value: 0}(address(_owner));
+    }
+
+    function testAddDepositFor_RevertWhen_SenderIsNotKYCd() public {
+        assertEq(_kintoIDv1.isKYC(address(_user)), false);
+        vm.expectRevert("SP: sender KYC required");
+        vm.prank(_user);
+        _paymaster.addDepositFor{value: 5e18}(address(_user));
+    }
+
+    function testAddDepositFor_RevertWhen_AccountIsEOA_WhenAccountIsNotKYCd() public {
+        assertEq(_kintoIDv1.isKYC(address(_user)), false);
+        vm.expectRevert("SP: account KYC required");
+        vm.prank(_owner);
+        _paymaster.addDepositFor{value: 5e18}(address(_user));
+    }
+
+    function testWithdrawTokensTo() public {
         uint256 balance = address(_user).balance;
+        approveKYC(_kycProvider, _user, _userPk);
+
+        vm.startPrank(_user);
+
         _paymaster.addDepositFor{value: 5e18}(address(_user));
         assertEq(address(_user).balance, balance - 5e18);
+
         _paymaster.unlockTokenDeposit();
-        // advance block to allow withdraw
-        vm.roll(block.number + 1);
+        vm.roll(block.number + 1); // advance block to allow withdraw
         _paymaster.withdrawTokensTo(address(_user), 5e18);
+
         assertEq(address(_user).balance, balance);
+
         vm.stopPrank();
     }
 
-    function test_RevertWhen_UserCanDepositStakeAndWithdrawWithoutRoll() public {
-        // user deposits 5 eth
-        uint256 balance = address(this).balance;
-        _paymaster.addDepositFor{value: 5e18}(address(this));
-        assertEq(address(this).balance, balance - 5e18);
+    function testWithdrawTokensTo_OwnerCanDepositStakeAndWithdraw() public {
+        uint256 balance = address(_owner).balance;
+
+        vm.startPrank(_owner);
+
+        _paymaster.addDepositFor{value: 5e18}(address(_owner));
+        assertEq(address(_owner).balance, balance - 5e18);
+
+        _paymaster.unlockTokenDeposit();
+        vm.roll(block.number + 1);
+
+        _paymaster.withdrawTokensTo(address(_owner), 5e18);
+        assertEq(address(_owner).balance, balance);
+
+        vm.stopPrank();
+    }
+
+    function testWithdrawTokensTo_RevertWhen_DepositLocked() public {
+        uint256 balance = _user.balance;
+        approveKYC(_kycProvider, _user, _userPk);
+
+        vm.prank(_user);
+        _paymaster.addDepositFor{value: 5e18}(_user);
+        assertEq(_user.balance, balance - 5e18);
 
         // user unlocks token deposit
         _paymaster.unlockTokenDeposit();
 
         // user withdraws 5 eth
         vm.expectRevert("SP: must unlockTokenDeposit");
-        _paymaster.withdrawTokensTo(address(this), 5e18);
+        _paymaster.withdrawTokensTo(_user, 5e18);
 
-        assertEq(address(this).balance, balance - 5e18);
+        assertEq(_user.balance, balance - 5e18);
     }
 
-    function testOwnerCanWithdrawAllInEmergency() public {
+    function testWithrawTo_WhenCallerIsOwner() public {
         uint256 deposited = _paymaster.getDeposit();
+        approveKYC(_kycProvider, _user, _userPk);
 
-        vm.prank(_user);
+        vm.prank(_owner);
         _paymaster.addDepositFor{value: 5e18}(address(_user));
 
         vm.prank(_owner);
@@ -119,32 +168,25 @@ contract SponsorPaymasterTest is SharedSetup {
         assertEq(address(_owner).balance, balBefore + deposited);
     }
 
-    function test_RevertWhen_UserCanWithdrawAllInEmergency() public {
-        vm.prank(_owner);
-        _paymaster.addDepositFor{value: 5e18}(address(_owner));
-
-        // user deposits 5 eth and then tries to withdraw all
-        vm.startPrank(_user);
-        _paymaster.addDepositFor{value: 5e18}(address(_user));
+    function testWithrawTo_RevertWhen_CallerIsNotOwner() public {
         vm.expectRevert("Ownable: caller is not the owner");
         _paymaster.withdrawTo(payable(_user), address(_entryPoint).balance);
-        vm.stopPrank();
     }
 
-    function testDepositInfo_WhenCallerDepositsToSomeoneElse(address someone) public {
-        vm.assume(someone != _owner);
+    function testDepositInfo_WhenCallerDepositsToHimself() public {
         vm.prank(_owner);
-        _paymaster.addDepositFor{value: 5e18}(address(someone));
-
-        (uint256 amount, uint256 _unlockBlock) = _paymaster.depositInfo(address(someone));
+        _paymaster.addDepositFor{value: 5e18}(address(_owner));
+        (uint256 amount, uint256 _unlockBlock) = _paymaster.depositInfo(address(_owner));
         assertEq(amount, 5e18);
         assertEq(_unlockBlock, 0);
     }
 
-    function testDepositInfo_WhenCallerDepositsToHimself() public {
+    function testDepositInfo_WhenCallerDepositsToSomeoneElse() public {
+        approveKYC(_kycProvider, _user, _userPk);
         vm.prank(_owner);
-        _paymaster.addDepositFor{value: 5e18}(address(_owner));
-        (uint256 amount, uint256 _unlockBlock) = _paymaster.depositInfo(address(_owner));
+        _paymaster.addDepositFor{value: 5e18}(address(_user));
+
+        (uint256 amount, uint256 _unlockBlock) = _paymaster.depositInfo(address(_user));
         assertEq(amount, 5e18);
         assertEq(_unlockBlock, 0);
     }

test/helpers/AATestScaffolding.sol

@@ -84,8 +84,9 @@ abstract contract AATestScaffolding is KYCSignature {
         vm.deal(_owner, 1e20);
     }
 
-    function fundSponsorForApp(address _contract) internal {
+    function fundSponsorForApp(address _sender, address _contract) internal {
         // we add the deposit to the counter contract in the paymaster
+        vm.prank(_sender);
         _paymaster.addDepositFor{value: 1e19}(address(_contract));
     }
 
@@ -139,10 +140,7 @@ abstract contract AATestScaffolding is KYCSignature {
         _paymaster = SponsorPaymaster(address(_proxyPaymaster));
 
         // initialize proxy
-        _paymaster.initialize(_owner);
-
-        // Set the registry in the paymaster
-        _paymaster.setAppRegistry(address(_kintoAppRegistry));
+        _paymaster.initialize(_owner, _kintoAppRegistry, _kintoIDv1);
 
         vm.stopPrank();
     }

test/wallet/whitelist/Whitelist.t.sol

@@ -34,7 +34,7 @@ contract WhitelistTest is SharedSetup {
     //     assertEq(counter.count(), 0);
 
     //     // (2). fund paymaster for Counter contract
-    //     fundSponsorForApp(address(counter));
+    //     fundSponsorForApp(_owner, address(counter));
 
     //     // (3). Create whitelist app user op
     //     UserOperation[] memory userOps = new UserOperation[](1);