Merge pull request #118 from getwax/wax-89-fix-unauthorized-access

Fix unauthorized access
getwax · Oct 12, 2023 · 45c815d · 45c815d
2 parents 0fb682b + 454f7fc
commit 45c815d
Show file tree

Hide file tree

Showing 4 changed files with 100 additions and 24 deletions.
diff --git a/account-integrations/safe/src/SafeBlsPlugin.sol b/account-integrations/safe/src/SafeBlsPlugin.sol
@@ -2,6 +2,8 @@
 pragma solidity >=0.7.0 <0.9.0;
 pragma abicoder v2;
 
+import {HandlerContext} from "safe-contracts/contracts/handler/HandlerContext.sol";
+
 import {BaseAccount} from "account-abstraction/contracts/core/BaseAccount.sol";
 import {IEntryPoint, UserOperation} from "account-abstraction/contracts/interfaces/IEntryPoint.sol";
 import {BLS} from "account-abstraction/contracts/samples/bls/lib/hubble-contracts/contracts/libs/BLS.sol";
@@ -19,7 +21,7 @@ interface ISafe {
 
 error IncorrectSignatureLength(uint256 length);
 
-contract SafeBlsPlugin is BaseAccount {
+contract SafeBlsPlugin is BaseAccount, HandlerContext {
     // TODO: Use EIP 712 for domain separation
     bytes32 public constant BLS_DOMAIN = keccak256("eip4337.bls.domain");
     address public immutable myAddress;
@@ -50,7 +52,7 @@ contract SafeBlsPlugin is BaseAccount {
         address to,
         uint256 value,
         bytes calldata data
-    ) external payable {
+    ) external payable fromThisOrEntryPoint {
         address payable safeAddress = payable(msg.sender);
         ISafe safe = ISafe(safeAddress);
         require(
@@ -124,4 +126,12 @@ contract SafeBlsPlugin is BaseAccount {
             );
         }
     }
+
+    modifier fromThisOrEntryPoint() {
+        require(
+            _msgSender() == address(this) ||
+            _msgSender() == _entryPoint
+        );
+        _;
+    }
 }
diff --git a/account-integrations/safe/src/SafeECDSAPlugin.sol b/account-integrations/safe/src/SafeECDSAPlugin.sol
@@ -2,6 +2,8 @@
 pragma solidity >=0.7.0 <0.9.0;
 pragma abicoder v2;
 
+import {HandlerContext} from "safe-contracts/contracts/handler/HandlerContext.sol";
+
 import {BaseAccount} from "account-abstraction/contracts/core/BaseAccount.sol";
 import {IEntryPoint, UserOperation} from "account-abstraction/contracts/interfaces/IEntryPoint.sol";
 import {UserOperation} from "account-abstraction/contracts/interfaces/IEntryPoint.sol";
@@ -23,7 +25,7 @@ struct ECDSAOwnerStorage {
     address owner;
 }
 
-contract SafeECDSAPlugin is BaseAccount {
+contract SafeECDSAPlugin is BaseAccount, HandlerContext {
     using ECDSA for bytes32;
 
     mapping(address => ECDSAOwnerStorage) public ecdsaOwnerStorage;
@@ -55,7 +57,7 @@ contract SafeECDSAPlugin is BaseAccount {
         address to,
         uint256 value,
         bytes calldata data
-    ) external payable {
+    ) external payable fromThisOrEntryPoint {
         address payable safeAddress = payable(msg.sender);
         ISafe safe = ISafe(safeAddress);
         require(
@@ -123,4 +125,12 @@ contract SafeECDSAPlugin is BaseAccount {
             );
         }
     }
+
+    modifier fromThisOrEntryPoint() {
+        require(
+            _msgSender() == address(this) ||
+            _msgSender() == _entryPoint
+        );
+        _;
+    }
 }
diff --git a/account-integrations/safe/src/SafeWebAuthnPlugin.sol b/account-integrations/safe/src/SafeWebAuthnPlugin.sol
@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: LGPL-3.0-only
 pragma solidity >=0.8.0 <0.9.0;
 
+import {HandlerContext} from "safe-contracts/contracts/handler/HandlerContext.sol";
+
 import {BaseAccount} from "account-abstraction/contracts/core/BaseAccount.sol";
 import {IEntryPoint, UserOperation} from "account-abstraction/contracts/interfaces/IEntryPoint.sol";
 import {WebAuthn} from "wax/primitives/src/WebAuthn.sol";
@@ -16,7 +18,7 @@ interface ISafe {
     ) external returns (bool success);
 }
 
-contract SafeWebAuthnPlugin is BaseAccount, WebAuthn {
+contract SafeWebAuthnPlugin is BaseAccount, WebAuthn, HandlerContext {
     address public immutable myAddress;
     address private immutable _entryPoint;
     uint256[2] private _publicKey;
@@ -45,7 +47,7 @@ contract SafeWebAuthnPlugin is BaseAccount, WebAuthn {
         address to,
         uint256 value,
         bytes calldata data
-    ) external payable {
+    ) external payable fromThisOrEntryPoint {
         address payable safeAddress = payable(msg.sender);
         ISafe safe = ISafe(safeAddress);
         require(
@@ -189,4 +191,12 @@ contract SafeWebAuthnPlugin is BaseAccount, WebAuthn {
             );
         }
     }
+
+    modifier fromThisOrEntryPoint() {
+        require(
+            _msgSender() == address(this) ||
+            _msgSender() == _entryPoint
+        );
+        _;
+    }
 }
diff --git a/account-integrations/safe/test/hardhat/integration/SafeECDSAPlugin.test.ts b/account-integrations/safe/test/hardhat/integration/SafeECDSAPlugin.test.ts
@@ -1,10 +1,7 @@
-import hre from "hardhat";
 import { expect } from "chai";
-import { AddressZero } from "@ethersproject/constants";
-import { getBytes, concat, resolveProperties, ethers } from "ethers";
+import { getBytes, resolveProperties, ethers } from "ethers";
 import { UserOperationStruct } from "@account-abstraction/contracts";
 import { getUserOpHash } from "@account-abstraction/utils";
-import { calculateProxyAddress } from "../utils/calculateProxyAddress";
 import {
   SafeECDSAFactory__factory,
   SafeECDSAPlugin__factory,
@@ -26,6 +23,8 @@ const BUNDLER_URL = process.env.ERC4337_TEST_BUNDLER_URL;
 const NODE_URL = process.env.ERC4337_TEST_NODE_URL;
 const MNEMONIC = process.env.MNEMONIC;
 
+const oneEther = ethers.parseEther("1");
+
 describe("SafeECDSAPlugin", () => {
   const setupTests = async () => {
     const bundlerProvider = new ethers.JsonRpcProvider(BUNDLER_URL);
@@ -54,14 +53,11 @@ describe("SafeECDSAPlugin", () => {
     };
   };
 
-  /**
-   * This test verifies a ERC4337 transaction succeeds when sent via a plugin
-   * The user operation deploys a Safe with the ERC4337 plugin and a handler
-   * and executes a transaction, thus verifying two things:
-   * 1. Deployment of the Safe with the ERC4337 plugin and handler is possible
-   * 2. Executing a transaction is possible
-   */
-  itif("should pass the ERC4337 validation", async () => {
+  async function setupDeployedAccount(
+    to: ethers.AddressLike,
+    value: ethers.BigNumberish,
+    data: ethers.BytesLike,
+  ) {
     const { singleton, provider, bundlerProvider, userWallet, entryPoints } =
       await setupTests();
     const ENTRYPOINT_ADDRESS = entryPoints[0];
@@ -106,7 +102,7 @@ describe("SafeECDSAPlugin", () => {
     const userOpCallData =
       SafeECDSAPlugin__factory.createInterface().encodeFunctionData(
         "execTransaction",
-        [recipient.address, transferAmount, "0x00"],
+        [to, value, data],
       );
 
     // Native tokens for the pre-fund 💸
@@ -160,13 +156,63 @@ describe("SafeECDSAPlugin", () => {
     //     `;
     // console.log(DEBUG_MESSAGE);
 
-    const recipientBalanceBefore = await provider.getBalance(recipient.address);
-
     await sendUserOpAndWait(userOperation, ENTRYPOINT_ADDRESS, bundlerProvider);
 
-    const recipientBalanceAfter = await provider.getBalance(recipient.address);
+    return {
+      provider,
+      bundlerProvider,
+      entryPoint: ENTRYPOINT_ADDRESS,
+      userWallet,
+      accountAddress,
+    };
+  }
+
+  /**
+   * This test verifies a ERC4337 transaction succeeds when sent via a plugin
+   * The user operation deploys a Safe with the ERC4337 plugin and a handler
+   * and executes a transaction, thus verifying two things:
+   * 1. Deployment of the Safe with the ERC4337 plugin and handler is possible
+   * 2. Executing a transaction is possible
+   */
+  itif("should pass the ERC4337 validation", async () => {
+    const recipient = ethers.Wallet.createRandom();
+
+    const { provider } = await setupDeployedAccount(
+      recipient.address,
+      oneEther,
+      "0x",
+    );
+
+    expect(await provider.getBalance(recipient.address)).to.equal(oneEther);
+  });
+
+  itif("should not allow execTransaction from unrelated address", async () => {
+    const { accountAddress, userWallet, provider } = await setupDeployedAccount(
+      ethers.ZeroAddress,
+      0,
+      "0x",
+    );
+
+    const unrelatedWallet = ethers.Wallet.createRandom(provider);
+
+    await receiptOf(
+      userWallet.sendTransaction({
+        to: unrelatedWallet.address,
+        value: 100n * oneEther,
+      }),
+    );
+
+    const account = SafeECDSAPlugin__factory.connect(
+      accountAddress,
+      unrelatedWallet,
+    );
+
+    const recipient = ethers.Wallet.createRandom(provider);
+
+    await expect(
+      receiptOf(account.execTransaction(recipient.address, oneEther, "0x")),
+    ).to.eventually.rejected;
 
-    const expectedRecipientBalance = recipientBalanceBefore + transferAmount;
-    expect(recipientBalanceAfter).to.equal(expectedRecipientBalance);
+    await expect(provider.getBalance(recipient)).to.eventually.equal(0n);
   });
 });