Merge pull request #87 from worldcoin/karan-pbh-4337-module-2

PBH 4337 Module

.gitmodules

@@ -7,6 +7,9 @@
 [submodule "contracts/lib/world-id-contracts"]
 	path = contracts/lib/world-id-contracts
 	url = https://github.com/worldcoin/world-id-contracts
+[submodule "contracts/lib/safe-contracts"]
+	path = contracts/lib/safe-contracts
+	url = https://github.com/safe-global/safe-contracts
 [submodule "contracts/lib/openzeppelin-contracts"]
 	path = contracts/lib/openzeppelin-contracts
 	url = https://github.com/OpenZeppelin/openzeppelin-contracts
@@ -16,3 +19,6 @@
 [submodule "contracts/lib/openzeppelin-contracts-upgradeable"]
 	path = contracts/lib/openzeppelin-contracts-upgradeable
 	url = https://github.com/OpenZeppelin/openzeppelin-contracts-upgradeable
+[submodule "contracts/lib/safe-modules"]
+	path = contracts/lib/safe-modules
+	url = https://github.com/worldcoin/safe-modules

contracts/foundry.toml

@@ -3,13 +3,9 @@ evm_version = 'cancun'
 src = "src"
 out = "out"
 libs = ["lib"]
-via_ir = true
 optimizer = true
 optimizer_runs = 200
 
 [fuzz]
 runs = 5000
-max_test_rejects = 150000
-
-
-
+max_test_rejects = 150000

contracts/remappings.txt

@@ -5,5 +5,7 @@
 @BokkyPooBahsDateTimeLibrary/=lib/BokkyPooBahsDateTimeLibrary/contracts/
 @helpers/=src/helpers/
 openzeppelin-contracts/=lib/world-id-contracts/lib/openzeppelin-contracts/contracts/
+@4337=lib/safe-modules/modules/4337/contracts/
+@safe-global/safe-contracts/contracts/=lib/safe-contracts/contracts/
 @forge-std/=lib/forge-std/src/
 forge-std/=lib/forge-std/src/

contracts/src/PBH4337Module.sol

@@ -0,0 +1,94 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.28;
+
+import {Safe4337Module} from "@4337/Safe4337Module.sol";
+import {PackedUserOperation} from "@account-abstraction/contracts/interfaces/PackedUserOperation.sol";
+import {ValidationData} from "@account-abstraction/contracts/core/Helpers.sol";
+import {_packValidationData} from "@account-abstraction/contracts/core/Helpers.sol";
+import {ISafe} from "@4337/interfaces/Safe.sol";
+
+contract PBHSafe4337Module is Safe4337Module {
+    uint256 constant ECDSA_SIGNATURE_LENGTH = 65;
+    uint256 constant TIMESTAMP_BYTES = 12; // 6 bytes each for validAfter and validUntil
+
+    address public immutable PBH_SIGNATURE_AGGREGATOR;
+    uint192 public immutable PBH_NONCE_KEY;
+
+    error InvalidProofSize();
+
+    constructor(address entryPoint, address _pbhSignatureAggregator, uint192 _pbhNonceKey) Safe4337Module(entryPoint) {
+        PBH_SIGNATURE_AGGREGATOR = _pbhSignatureAggregator;
+        PBH_NONCE_KEY = _pbhNonceKey;
+    }
+
+    // TODO: Fork the Safe4337Module dependency and add 'override' to _validateSignatures. It is manually updated currently and CI will fail
+    /**
+     * @dev Validates that the user operation is correctly signed and returns an ERC-4337 packed validation data
+     * of `validAfter || validUntil || authorizer`:
+     *  - `authorizer`: 20-byte address, 0 for valid signature or 1 to mark signature failure (this module does not make use of signature aggregators).
+     *  - `validUntil`: 6-byte timestamp value, or zero for "infinite". The user operation is valid only up to this time.
+     *  - `validAfter`: 6-byte timestamp. The user operation is valid only after this time.
+     * @param userOp User operation struct.
+     * @return validationData An integer indicating the result of the validation.
+     */
+    function _validateSignatures(PackedUserOperation calldata userOp)
+        internal
+        view
+        override
+        returns (uint256 validationData)
+    {
+        // Check if the userOp has the specified PBH key
+        // https://github.com/eth-infinitism/account-abstraction/blob/develop/contracts/core/NonceManager.sol#L38
+        uint192 key = uint192(userOp.nonce >> 64);
+
+        // This does NOT validate the proof
+        // It removes the first 12 bytes from the signature as it represents the validAfter and validUntil values
+        // operationData is not determined by the signature
+        (bytes memory operationData, uint48 validAfter, uint48 validUntil, bytes calldata signatures) =
+            _getSafeOp(userOp);
+
+        // If it is a PBH transaction, we need to handle two cases with the signature:
+        // 1. The bundler simulates the call with the proof appended
+        // 2. UserOp execution without proof appended
+        bool isPBH = (key == PBH_NONCE_KEY);
+
+        // Base signature length calculation:
+        // TIMESTAMP_BYTES (12) + (threshold * ECDSA_SIGNATURE_LENGTH)
+        uint256 expectedLength =
+            TIMESTAMP_BYTES + (ISafe(payable(userOp.sender)).getThreshold() * ECDSA_SIGNATURE_LENGTH);
+
+        // If the signature length is greater than the expected length, then we know that the bundler appended the proof
+        // We need to remove the proof from the signature before validation
+        if (isPBH && userOp.signature.length > expectedLength) {
+            if (userOp.signature.length - expectedLength != 352) {
+                revert InvalidProofSize();
+            }
+            // Remove the proof from the signature
+            signatures = userOp.signature[0:expectedLength];
+        }
+
+        // The `checkSignatures` function in the Safe contract does not force a fixed size on signature length.
+        // A malicious bundler can pad the Safe operation `signatures` with additional bytes, causing the account to pay
+        // more gas than needed for user operation validation (capped by `verificationGasLimit`).
+        // `_checkSignaturesLength` ensures that there are no additional bytes in the `signature` than are required.
+        bool validSignature = _checkSignaturesLength(signatures, ISafe(payable(userOp.sender)).getThreshold());
+
+        try ISafe(payable(userOp.sender)).checkSignatures(keccak256(operationData), operationData, signatures) {}
+        catch {
+            validSignature = false;
+        }
+
+        address authorizer;
+
+        // If the signature is valid and the userOp is a PBH userOp, return the PBH signature aggregator as the authorizer
+        // Else return 0 for valid signature and 1 for invalid signature
+        if (isPBH && validSignature) {
+            authorizer = PBH_SIGNATURE_AGGREGATOR;
+        } else {
+            authorizer = validSignature ? address(0) : address(1);
+        }
+
+        // The timestamps are validated by the entry point, therefore we will not check them again.
+        validationData = _packValidationData(ValidationData(authorizer, validUntil, validAfter));
+    }
+}

contracts/src/PBHEntryPoint.sol

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.8.20;
+pragma solidity ^0.8.28;
 
 import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
 

contracts/src/PBHEntryPointImplV1.sol

@@ -1,16 +1,19 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.8.20;
+pragma solidity ^0.8.28;
 
-import {PBHVerifier} from "./PBHVerifier.sol";
 import {IWorldIDGroups} from "@world-id-contracts/interfaces/IWorldIDGroups.sol";
 import {IEntryPoint} from "@account-abstraction/contracts/interfaces/IEntryPoint.sol";
 import {IPBHEntryPoint} from "./interfaces/IPBHEntryPoint.sol";
+import {ByteHasher} from "./helpers/ByteHasher.sol";
+import {PBHExternalNullifier} from "./helpers/PBHExternalNullifier.sol";
+import {WorldIDImpl} from "@world-id-contracts/abstract/WorldIDImpl.sol";
+import "@BokkyPooBahsDateTimeLibrary/BokkyPooBahsDateTimeLibrary.sol";
 
 /// @title PBH Entry Point Implementation V1
 /// @dev This contract is an implementation of the PBH Entry Point.
 ///      It is used to verify the signature of a Priority User Operation, and Relaying Priority Bundles to the EIP-4337 Entry Point.
 /// @author Worldcoin
-contract PBHEntryPointImplV1 is IPBHEntryPoint, PBHVerifier {
+contract PBHEntryPointImplV1 is IPBHEntryPoint, WorldIDImpl {
     ///////////////////////////////////////////////////////////////////////////////
     ///                   A NOTE ON IMPLEMENTATION CONTRACTS                    ///
     ///////////////////////////////////////////////////////////////////////////////
@@ -48,6 +51,15 @@ contract PBHEntryPointImplV1 is IPBHEntryPoint, PBHVerifier {
     // these variables takes place. If reordering happens, a storage clash will occur (effectively a
     // memory safety error).
 
+    using ByteHasher for bytes;
+
+    ///////////////////////////////////////////////////////////////////////////////
+    ///                                  ERRORS                                ///
+    //////////////////////////////////////////////////////////////////////////////
+
+    /// @notice Thrown when attempting to reuse a nullifier
+    error InvalidNullifier();
+
     ///////////////////////////////////////////////////////////////////////////////
     ///                                  Events                                ///
     //////////////////////////////////////////////////////////////////////////////
@@ -56,6 +68,23 @@ contract PBHEntryPointImplV1 is IPBHEntryPoint, PBHVerifier {
         IWorldIDGroups indexed worldId, IEntryPoint indexed entryPoint, uint8 indexed numPbhPerMonth
     );
 
+    /// @notice Emitted once for each successful PBH verification.
+    ///
+    /// @param sender The sender of this particular transaction or UserOp.
+    /// @param nonce Transaction/UserOp nonce.
+    /// @param payload The zero-knowledge proof that demonstrates the claimer is registered with World ID.
+    event PBH(address indexed sender, uint256 indexed nonce, PBHPayload payload);
+
+    /// @notice Emitted when the World ID address is set.
+    ///
+    /// @param worldId The World ID instance that will be used for verifying proofs.
+    event WorldIdSet(address indexed worldId);
+
+    /// @notice Emitted when the number of PBH transactions allowed per month is set.
+    ///
+    /// @param numPbhPerMonth The number of allowed PBH transactions per month.
+    event NumPbhPerMonthSet(uint8 indexed numPbhPerMonth);
+
     ///////////////////////////////////////////////////////////////////////////////
     ///                                  Vars                                  ///
     //////////////////////////////////////////////////////////////////////////////
@@ -65,6 +94,22 @@ contract PBHEntryPointImplV1 is IPBHEntryPoint, PBHVerifier {
     ///     The PBHVerifier is always the proxy to the EntryPoint for PBH Bundles.
     bytes32 internal _hashedOps;
 
+    /// @dev The World ID group ID (always 1)
+    uint256 internal constant _GROUP_ID = 1;
+
+    /// @dev The World ID instance that will be used for verifying proofs
+    IWorldIDGroups public worldId;
+
+    /// @dev The EntryPoint where Aggregated PBH Bundles will be proxied to.
+    IEntryPoint public entryPoint;
+
+    /// @notice The number of PBH transactions that may be used by a single
+    ///         World ID in a given month.
+    uint8 public numPbhPerMonth;
+
+    /// @dev Whether a nullifier hash has been used already. Used to guarantee an action is only performed once by a single person
+    mapping(uint256 => bool) public nullifierHashes;
+
     ///////////////////////////////////////////////////////////////////////////////
     ///                             INITIALIZATION                              ///
     ///////////////////////////////////////////////////////////////////////////////
@@ -158,4 +203,60 @@ contract PBHEntryPointImplV1 is IPBHEntryPoint, PBHVerifier {
             if iszero(eq(tload(hashedOps), hashedOps)) { revert(0, 0) }
         }
     }
+
+    /// @param sender The sender of this particular transaction or UserOp.
+    /// @param nonce Transaction/UserOp nonce.
+    /// @param callData Transaction/UserOp call data.
+    /// @param pbhPayload The PBH payload containing the proof data.
+    function verifyPbh(address sender, uint256 nonce, bytes calldata callData, PBHPayload memory pbhPayload)
+        public
+        virtual
+        onlyInitialized
+        onlyProxy
+    {
+        // First, we make sure this nullifier has not been used before.
+        if (nullifierHashes[pbhPayload.nullifierHash]) {
+            revert InvalidNullifier();
+        }
+
+        // Verify the external nullifier
+        PBHExternalNullifier.verify(pbhPayload.pbhExternalNullifier, numPbhPerMonth);
+
+        // If worldId address is set, proceed with on chain verification,
+        // otherwise assume verification has been done off chain by the builder.
+        if (address(worldId) != address(0)) {
+            // We now generate the signal hash from the sender, nonce, and calldata
+            uint256 signalHash = abi.encodePacked(sender, nonce, callData).hashToField();
+
+            // We now verify the provided proof is valid and the user is verified by World ID
+            worldId.verifyProof(
+                pbhPayload.root,
+                _GROUP_ID,
+                signalHash,
+                pbhPayload.nullifierHash,
+                pbhPayload.pbhExternalNullifier,
+                pbhPayload.proof
+            );
+        }
+
+        // We now record the user has done this, so they can't do it again (proof of uniqueness)
+        nullifierHashes[pbhPayload.nullifierHash] = true;
+
+        emit PBH(sender, nonce, pbhPayload);
+    }
+
+    /// @notice Sets the number of PBH transactions allowed per month.
+    /// @param _numPbhPerMonth The number of allowed PBH transactions per month.
+    function setNumPbhPerMonth(uint8 _numPbhPerMonth) external virtual onlyOwner onlyProxy onlyInitialized {
+        numPbhPerMonth = _numPbhPerMonth;
+        emit NumPbhPerMonthSet(_numPbhPerMonth);
+    }
+
+    /// @dev If the World ID address is set to 0, then it is assumed that verification will take place off chain.
+    /// @notice Sets the World ID instance that will be used for verifying proofs.
+    /// @param _worldId The World ID instance that will be used for verifying proofs.
+    function setWorldId(address _worldId) external virtual onlyOwner onlyProxy onlyInitialized {
+        worldId = IWorldIDGroups(_worldId);
+        emit WorldIdSet(_worldId);
+    }
 }

contracts/src/PBHSignatureAggregator.sol

@@ -1,9 +1,8 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.8.20;
+pragma solidity ^0.8.28;
 
 import "@account-abstraction/contracts/interfaces/PackedUserOperation.sol";
 import {IPBHEntryPoint} from "./interfaces/IPBHEntryPoint.sol";
-import {IPBHVerifier} from "./interfaces/IPBHVerifier.sol";
 import {IAggregator} from "@account-abstraction/contracts/interfaces/IAggregator.sol";
 
 /// @title PBH Signature Aggregator
@@ -64,12 +63,12 @@ contract PBHSignatureAggregator is IAggregator {
         pure
         returns (bytes memory aggregatedSignature)
     {
-        IPBHVerifier.PBHPayload[] memory pbhPayloads = new IPBHVerifier.PBHPayload[](userOps.length);
+        IPBHEntryPoint.PBHPayload[] memory pbhPayloads = new IPBHEntryPoint.PBHPayload[](userOps.length);
         for (uint256 i = 0; i < userOps.length; ++i) {
             // Bytes (0:65) - UserOp Signature
             // Bytes (65:65 + 352) - Packed Proof Data
             bytes memory proofData = userOps[i].signature[65:];
-            pbhPayloads[i] = abi.decode(proofData, (IPBHVerifier.PBHPayload));
+            pbhPayloads[i] = abi.decode(proofData, (IPBHEntryPoint.PBHPayload));
         }
         aggregatedSignature = abi.encode(pbhPayloads);
     }