chore: increase coverage

script/migrations/21-multiple_upgrade.sol

@@ -16,6 +16,12 @@ import "../../test/helpers/UserOp.sol";
 import "forge-std/Script.sol";
 import "forge-std/console.sol";
 
+contract KintoWalletV4 is KintoWallet {
+    constructor(IEntryPoint _entryPoint, IKintoID _kintoID, IKintoAppRegistry _appRegistry)
+        KintoWallet(_entryPoint, _kintoID, _appRegistry)
+    {}
+}
+
 contract KintoMigration21DeployScript is Create2Helper, ArtifactsReader, UserOp {
     using ECDSAUpgradeable for bytes32;
 

script/migrations/22-wallet_v5.sol

@@ -0,0 +1,27 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.18;
+
+import "../../src/wallet/KintoWallet.sol";
+import "./utils/MigrationHelper.sol";
+
+contract KintoMigration22DeployScript is MigrationHelper {
+    using ECDSAUpgradeable for bytes32;
+
+    // NOTE: this migration must be run from the ledger admin
+    function run() public {
+        run2();
+
+        // generate bytecode for KintoWalletV5
+        bytes memory bytecode = abi.encodePacked(
+            type(KintoWalletV5).creationCode,
+            abi.encode(
+                _getChainDeployment("EntryPoint"),
+                IKintoID(_getChainDeployment("KintoID")),
+                IKintoAppRegistry(_getChainDeployment("KintoAppRegistry"))
+            )
+        );
+
+        // upgrade KintoWallet to V5
+        deployAndUpgrade("KintoWallet", "V5", bytecode);
+    }
+}

script/migrations/utils/MigrationHelper.sol

@@ -0,0 +1,146 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.18;
+
+import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+import "@openzeppelin/contracts/access/Ownable.sol";
+
+import "../../../src/wallet/KintoWalletFactory.sol";
+
+import "../../../src/interfaces/ISponsorPaymaster.sol";
+import "../../../src/interfaces/IKintoWallet.sol";
+
+import "../../../test/helpers/Create2Helper.sol";
+import "../../../test/helpers/ArtifactsReader.sol";
+import "../../../test/helpers/UserOp.sol";
+
+import "forge-std/Script.sol";
+import "forge-std/console.sol";
+
+interface IInitialize {
+    function initialize(address anOwner, address _recoverer) external;
+}
+
+contract MigrationHelper is Create2Helper, ArtifactsReader, UserOp {
+    using ECDSAUpgradeable for bytes32;
+
+    uint256 deployerPrivateKey;
+    KintoWalletFactory _walletFactory;
+
+    function run2() public {
+        console.log("Running on chain: ", vm.toString(block.chainid));
+        console.log("Executing from address", msg.sender);
+
+        deployerPrivateKey = vm.envUint("PRIVATE_KEY");
+        _walletFactory = KintoWalletFactory(payable(_getChainDeployment("KintoWalletFactory")));
+    }
+
+    /// @dev deploys contracts from deployer and upgrades them using LEDGER_ADMIN
+    function deployAndUpgrade(string memory contractName, string memory version, bytes memory bytecode) public {
+        bool isWallet = keccak256(abi.encodePacked(contractName)) == keccak256(abi.encodePacked("KintoWallet"));
+        address proxy = _getChainDeployment(contractName);
+
+        if (!isWallet) require(proxy != address(0), "Need to execute main deploy script first");
+
+        // (1). deploy new implementation via wallet factory
+        vm.broadcast(deployerPrivateKey);
+        address _impl = _walletFactory.deployContract(vm.envAddress("LEDGER_ADMIN"), 0, bytecode, bytes32(0));
+
+        console.log(string.concat(contractName, version, "-impl: ", vm.toString(address(_impl))));
+
+        // (2). call upgradeTo to set new implementation
+        if (isWallet) {
+            vm.broadcast(); // requires LEDGER_ADMIN
+            // vm.prank(vm.envAddress("LEDGER_ADMIN"));
+            _walletFactory.upgradeAllWalletImplementations(IKintoWallet(_impl));
+        } else {
+            vm.broadcast(); // requires LEDGER_ADMIN
+            // vm.prank(vm.envAddress("LEDGER_ADMIN"));
+            UUPSUpgradeable(proxy).upgradeTo(_impl);
+        }
+    }
+
+    // utils for doing actions through EntryPoint
+
+    function _upgradeTo(address _proxy, address _newImpl, uint256 _signerPk) internal {
+        address payable _from = payable(_getChainDeployment("KintoWallet-admin"));
+
+        // prep upgradeTo user op
+        uint256 nonce = IKintoWallet(_from).getNonce();
+        uint256[] memory privateKeys = new uint256[](1);
+        privateKeys[0] = _signerPk;
+        UserOperation[] memory userOps = new UserOperation[](1);
+        userOps[0] = _createUserOperation(
+            block.chainid,
+            _from,
+            _proxy,
+            0,
+            nonce,
+            privateKeys,
+            abi.encodeWithSelector(UUPSUpgradeable.upgradeTo.selector, address(_newImpl)),
+            _getChainDeployment("SponsorPaymaster")
+        );
+
+        vm.broadcast(deployerPrivateKey);
+        IEntryPoint(_getChainDeployment("EntryPoint")).handleOps(userOps, payable(vm.addr(_signerPk)));
+    }
+
+    function _initialize(address _proxy, uint256 _signerPk) internal {
+        address payable _from = payable(_getChainDeployment("KintoWallet-admin"));
+
+        // fund _proxy in the paymaster
+        require(
+            ISponsorPaymaster(payable(_getChainDeployment("SponsorPaymaster"))).balances(_proxy) > 0,
+            "Need to fund proxy in paymaster"
+        );
+
+        // todo: move to a fund function
+        // ISponsorPaymaster _paymaster = ISponsorPaymaster(_getChainDeployment("SponsorPaymaster"));
+        // vm.broadcast(deployerPrivateKey);
+        // _paymaster.addDepositFor{value: 0.00000001 ether}(_proxy);
+        // assertEq(_paymaster.balances(_proxy), 0.00000001 ether);
+
+        // prep upgradeTo user op
+        uint256 nonce = IKintoWallet(_from).getNonce();
+        uint256[] memory privateKeys = new uint256[](1);
+        privateKeys[0] = _signerPk;
+
+        UserOperation[] memory userOps = new UserOperation[](1);
+        userOps[0] = _createUserOperation(
+            block.chainid,
+            _from,
+            _proxy,
+            0,
+            nonce,
+            privateKeys,
+            abi.encodeWithSelector(IInitialize.initialize.selector),
+            _getChainDeployment("SponsorPaymaster")
+        );
+
+        vm.broadcast(deployerPrivateKey);
+        IEntryPoint(_getChainDeployment("EntryPoint")).handleOps(userOps, payable(vm.addr(_signerPk)));
+    }
+
+    function _transferOwnership(address _proxy, uint256 _signerPk, address _newOwner) internal {
+        address payable _from = payable(_getChainDeployment("KintoWallet-admin"));
+
+        // prep upgradeTo user op
+        uint256 nonce = IKintoWallet(_from).getNonce();
+        uint256[] memory privateKeys = new uint256[](1);
+        privateKeys[0] = _signerPk;
+
+        UserOperation[] memory userOps = new UserOperation[](1);
+        userOps[0] = _createUserOperation(
+            block.chainid,
+            _from,
+            _proxy,
+            0,
+            nonce,
+            privateKeys,
+            abi.encodeWithSelector(Ownable.transferOwnership.selector, _newOwner),
+            _getChainDeployment("SponsorPaymaster")
+        );
+
+        vm.broadcast(deployerPrivateKey);
+        IEntryPoint(_getChainDeployment("EntryPoint")).handleOps(userOps, payable(vm.addr(_signerPk)));
+    }
+}

script/test.sol

@@ -156,7 +156,6 @@ contract KintoDeployTestCounter is AASetup, KYCSignature, UserOp {
         );
         UserOperation[] memory userOps = new UserOperation[](1);
         userOps[0] = userOp;
-        // Execute the transaction via the entry point
         _entryPoint.handleOps(userOps, payable(deployerPublicKey));
         console.log("After UserOp. Counter:", counter.count());
         vm.stopBroadcast();
@@ -234,7 +233,6 @@ contract KintoDeployETHPriceIsRight is AASetup, KYCSignature, UserOp {
         );
         UserOperation[] memory userOps = new UserOperation[](1);
         userOps[0] = userOp;
-        // Execute the transaction via the entry point
         _entryPoint.handleOps(userOps, payable(deployerPublicKey));
         console.log("After UserOp. ETHPriceIsRight guess count", ethpriceisright.guessCount());
         console.log("After UserOp. ETHPriceIsRight avg guess", ethpriceisright.avgGuess());

src/paymasters/SponsorPaymaster.sol

@@ -277,25 +277,6 @@ contract SponsorPaymaster is Initializable, BasePaymaster, UUPSUpgradeable, Reen
         );
     }
 
-    /// Return app's sponsor from the registry
-    /// @return sponsor - the sponsor address
-    /// @dev reverts if neither execute nor executeBatch
-    /// @dev ensures all targets are sponsored by the same app if executeBatch (todo: decouple this)
-    function _getSponsor(bytes calldata callData) internal view returns (address sponsor) {
-        bytes4 selector = bytes4(callData[:4]);
-        if (selector == IKintoWallet.execute.selector) {
-            (address target,,) = abi.decode(callData[4:], (address, uint256, bytes));
-            sponsor = appRegistry.getSponsor(target);
-        } else if (selector == IKintoWallet.executeBatch.selector) {
-            // last target is the sponsor
-            (address[] memory targets,,) = abi.decode(callData[4:], (address[], uint256[], bytes[]));
-            sponsor = appRegistry.getSponsor(targets[targets.length - 1]);
-        } else {
-            // handle unknown function or error
-            revert("SP: Unknown function selector");
-        }
-    }
-
     // @notice extracts `target` contract from callData
     // @dev the last op on a batch MUST always be a contract whose sponsor is the one we want to
     // bear with the gas cost of all ops

src/wallet/KintoWallet.sol

@@ -111,7 +111,7 @@ contract KintoWallet is Initializable, BaseAccount, TokenCallbackHandler, IKinto
         for (uint256 i = 0; i < dest.length; i++) {
             _executeInner(dest[i], values[i], func[i]);
         }
-        // If can transact, cancel recovery
+        // if can transact, cancel recovery
         inRecovery = 0;
     }
 
@@ -282,8 +282,8 @@ contract KintoWallet is Initializable, BaseAccount, TokenCallbackHandler, IKinto
         // if app key is not set or signature is not valid, verify signer policy
         if (
             (
-                signerPolicy == SINGLE_SIGNER && owners.length == 1
-                    && _verifySingleSignature(owners[0], hashData, userOp.signature) == SIG_VALIDATION_SUCCESS
+                signerPolicy == SINGLE_SIGNER
+                    && _verifySingleSignature(hashData, userOp.signature) == SIG_VALIDATION_SUCCESS
             )
                 || (
                     signerPolicy != SINGLE_SIGNER
@@ -335,6 +335,20 @@ contract KintoWallet is Initializable, BaseAccount, TokenCallbackHandler, IKinto
         return appRegistry.isSponsored(sponsor, target) || appRegistry.childToParentContract(target) == sponsor;
     }
 
+    // @notice ensures at least 1 signer has signed the hash
+    // @dev useful when owners.length > 1
+    function _verifySingleSignature(bytes32 hashData, bytes memory signature) private view returns (uint256) {
+        // ensure at least one signer has signed
+        address recovered = hashData.recover(signature);
+        for (uint256 i = 0; i < owners.length; i++) {
+            if (owners[i] == recovered) {
+                return _packValidationData(false, 0, 0);
+            }
+        }
+        return SIG_VALIDATION_FAILED;
+    }
+
+    // @notice ensures signer has signed the hash
     function _verifySingleSignature(address signer, bytes32 hashData, bytes memory signature)
         private
         pure
@@ -346,6 +360,7 @@ contract KintoWallet is Initializable, BaseAccount, TokenCallbackHandler, IKinto
         return _packValidationData(false, 0, 0);
     }
 
+    // @notice ensures required signers have signed the hash
     function _verifyMultipleSignatures(bytes32 hashData, bytes memory signature) private view returns (uint256) {
         // calculate required signers
         uint256 requiredSigners =
@@ -433,7 +448,7 @@ contract KintoWallet is Initializable, BaseAccount, TokenCallbackHandler, IKinto
 }
 
 // Upgradeable version of KintoWallet
-contract KintoWalletV4 is KintoWallet {
+contract KintoWalletV5 is KintoWallet {
     constructor(IEntryPoint _entryPoint, IKintoID _kintoID, IKintoAppRegistry _appRegistry)
         KintoWallet(_entryPoint, _kintoID, _appRegistry)
     {}

test/EngenCredits.t.sol

@@ -139,7 +139,6 @@ contract EngenCreditsTest is UserOp, AATestScaffolding {
             privateKeys, address(_kintoWallet), _kintoWallet.getNonce(), address(_engenCredits), address(_paymaster)
         );
         userOps[1] = userOp;
-        // Execute the transaction via the entry point
         _entryPoint.handleOps(userOps, payable(_owner));
         assertEq(_engenCredits.balanceOf(address(_kintoWallet)), 15);
         vm.stopPrank();
@@ -171,7 +170,6 @@ contract EngenCreditsTest is UserOp, AATestScaffolding {
             privateKeys, address(_kintoWallet), _kintoWallet.getNonce(), address(_engenCredits), address(_paymaster)
         );
         userOps[1] = userOp;
-        // Execute the transaction via the entry point
         _entryPoint.handleOps(userOps, payable(_owner));
         assertEq(_engenCredits.balanceOf(address(_kintoWallet)), 20);
         vm.stopPrank();
@@ -198,7 +196,6 @@ contract EngenCreditsTest is UserOp, AATestScaffolding {
             privateKeys, address(_kintoWallet), _kintoWallet.getNonce(), address(_engenCredits), address(_paymaster)
         );
         userOps[1] = userOp;
-        // Execute the transaction via the entry point
         _entryPoint.handleOps(userOps, payable(_owner));
         assertEq(_engenCredits.balanceOf(address(_kintoWallet)), 15);
         // call again

test/KYCViewer.t.sol

@@ -9,9 +9,8 @@ import "../src/wallet/KintoWalletFactory.sol";
 import "../src/KintoID.sol";
 import "../src/viewers/KYCViewer.sol";
 
-import "./helpers/UserOp.sol";
+import "./SharedSetup.t.sol";
 import "./helpers/UUPSProxy.sol";
-import {AATestScaffolding} from "./helpers/AATestScaffolding.sol";
 
 contract KYCViewerUpgraded is KYCViewer {
     function newFunction() external pure returns (uint256) {
@@ -21,55 +20,28 @@ contract KYCViewerUpgraded is KYCViewer {
     constructor(address _kintoWalletFactory) KYCViewer(_kintoWalletFactory) {}
 }
 
-contract KYCViewerTest is UserOp, AATestScaffolding {
-    uint256 _chainID = 1;
-
-    UUPSProxy _proxyViewer;
-    KYCViewer _implkycViewer;
-    KYCViewerUpgraded _implKYCViewerUpgraded;
-    KYCViewer _kycViewer;
-    KYCViewerUpgraded _kycViewer2;
-
-    function setUp() public {
-        vm.chainId(_chainID);
-        vm.startPrank(address(1));
-        _owner.transfer(1e18);
-        vm.stopPrank();
-        deployAAScaffolding(_owner, 1, _kycProvider, _recoverer);
-        vm.startPrank(_owner);
-        _implkycViewer = new KYCViewer{salt: 0}(address(_walletFactory));
-        // deploy _proxy contract and point it to _implementation
-        _proxyViewer = new UUPSProxy{salt: 0}(address(_implkycViewer), "");
-        // wrap in ABI to support easier calls
-        _kycViewer = KYCViewer(address(_proxyViewer));
-        // Initialize kyc viewer _proxy
-        _kycViewer.initialize();
-        vm.stopPrank();
-    }
-
-    function testUp() public {
-        console.log("address owner", address(_owner));
+contract KYCViewerTest is SharedSetup {
+    function testUp() public override {
+        super.testUp();
         assertEq(_kycViewer.owner(), _owner);
         assertEq(address(_entryPoint.walletFactory()), address(_kycViewer.walletFactory()));
-        address kintoID = address(_kycViewer.kintoID());
-        assertEq(address(_walletFactory.kintoID()), kintoID);
+        assertEq(address(_walletFactory.kintoID()), address(_kycViewer.kintoID()));
     }
 
     /* ============ Upgrade Tests ============ */
 
-    function testOwnerCanUpgradeViewer() public {
-        vm.startPrank(_owner);
+    function testUpgradeTo_WhenCallerIsOwner() public {
         KYCViewerUpgraded _implementationV2 = new KYCViewerUpgraded(address(_walletFactory));
+        vm.prank(_owner);
         _kycViewer.upgradeTo(address(_implementationV2));
-        // re-wrap the _proxy
-        _kycViewer2 = KYCViewerUpgraded(address(_kycViewer));
-        assertEq(_kycViewer2.newFunction(), 1);
-        vm.stopPrank();
+        assertEq(KYCViewerUpgraded(address(_kycViewer)).newFunction(), 1);
     }
 
-    function test_RevertWhen_OthersCannotUpgradeFactory() public {
+    function testUpgradeTo_RevertWhen_CallerIsNotOwner(address someone) public {
+        vm.assume(someone != _owner);
         KYCViewerUpgraded _implementationV2 = new KYCViewerUpgraded(address(_walletFactory));
         vm.expectRevert("only owner");
+        vm.prank(someone);
         _kycViewer.upgradeTo(address(_implementationV2));
     }
 
@@ -80,5 +52,7 @@ contract KYCViewerTest is UserOp, AATestScaffolding {
         assertEq(_kycViewer.isIndividual(address(_kintoWallet)), _kycViewer.isIndividual(_owner));
         assertEq(_kycViewer.isCompany(address(_kintoWallet)), false);
         assertEq(_kycViewer.hasTrait(address(_kintoWallet), 6), false);
+        assertEq(_kycViewer.isSanctionsSafe(address(_kintoWallet)), true);
+        assertEq(_kycViewer.isSanctionsSafeIn(address(_kintoWallet), 1), true);
     }
 }