fix: extend signature aggregation to a variable number of signers

.gitmodules

@@ -22,3 +22,6 @@
 [submodule "contracts/lib/safe-modules"]
 	path = contracts/lib/safe-modules
 	url = https://github.com/worldcoin/safe-modules
+[submodule "contracts/lib/solady"]
+	path = contracts/lib/solady
+	url = https://github.com/vectorized/solady

contracts/remappings.txt

@@ -8,4 +8,5 @@ openzeppelin-contracts/=lib/world-id-contracts/lib/openzeppelin-contracts/contra
 @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/
+forge-std/=lib/forge-std/src/
+@solady=lib/solady/src/utils/

contracts/src/PBHSignatureAggregator.sol

@@ -4,6 +4,7 @@ pragma solidity ^0.8.28;
 import "@account-abstraction/contracts/interfaces/PackedUserOperation.sol";
 import {IPBHEntryPoint} from "./interfaces/IPBHEntryPoint.sol";
 import {IAggregator} from "@account-abstraction/contracts/interfaces/IAggregator.sol";
+import {ISafe} from "@4337/interfaces/Safe.sol";
 
 /// @title PBH Signature Aggregator
 /// @author Worldcoin
@@ -12,10 +13,20 @@ import {IAggregator} from "@account-abstraction/contracts/interfaces/IAggregator
 ///         Smart Accounts that return the `PBHSignatureAggregator` as the authorizer in `validationData`
 ///         will be considered as Priority User Operations, and will need to pack a World ID proof in the signature field.
 contract PBHSignatureAggregator is IAggregator {
+    /// @notice The length of an ECDSA signature.
+    uint256 constant ECDSA_SIGNATURE_LENGTH = 65;
+    /// @notice The length of the timestamp bytes.
+    uint256 constant TIMESTAMP_BYTES = 12; // 6 bytes each for validAfter and validUntil
+    /// @notice The length of the encoded proof data.
+    uint256 constant PROOF_DATA_LENGTH = 352;
+
     /// @notice Thrown when the Hash of the UserOperations is not
     ///         in transient storage of the `PBHVerifier`.
     error InvalidUserOperations();
 
+    /// @notice Thrown when the length of the signature is invalid.
+    error InvalidSignatureLength(uint256 expected, uint256 actual);
+
     /// @notice The PBHVerifier contract.
     IPBHEntryPoint public immutable pbhEntryPoint;
 
@@ -60,14 +71,18 @@ contract PBHSignatureAggregator is IAggregator {
      */
     function aggregateSignatures(PackedUserOperation[] calldata userOps)
         external
-        pure
+        view
         returns (bytes memory aggregatedSignature)
     {
         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:];
+            uint256 expectedLength =
+                TIMESTAMP_BYTES + (ISafe(payable(userOps[i].sender)).getThreshold() * ECDSA_SIGNATURE_LENGTH);
+            require(
+                userOps[i].signature.length == expectedLength + PROOF_DATA_LENGTH,
+                InvalidSignatureLength(expectedLength + PROOF_DATA_LENGTH, userOps[i].signature.length)
+            );
+            bytes memory proofData = userOps[i].signature[expectedLength:];
             pbhPayloads[i] = abi.decode(proofData, (IPBHEntryPoint.PBHPayload));
         }
         aggregatedSignature = abi.encode(pbhPayloads);

contracts/test/TestUtils.sol

@@ -5,26 +5,12 @@ import "@account-abstraction/contracts/interfaces/PackedUserOperation.sol";
 import {IEntryPoint} from "@account-abstraction/contracts/interfaces/IEntryPoint.sol";
 import {IAggregator} from "@account-abstraction/contracts/interfaces/IAggregator.sol";
 import "@forge-std/console.sol";
+import "@solady/LibBytes.sol";
 
 contract TestUtils {
-    function encodeSignature(bytes memory proofData) public pure returns (bytes memory) {
-        bytes memory sigBuffer = new bytes(65);
-        bytes memory signature = new bytes(417);
-        assembly {
-            mstore(signature, sigBuffer)
-            mstore(add(signature, 32), mload(add(sigBuffer, 32)))
-            mstore(add(signature, 64), mload(add(sigBuffer, 64)))
-            mstore(add(signature, 65), mload(proofData))
-            mstore(add(add(signature, 65), 32), mload(add(proofData, 32)))
-            mstore(add(add(signature, 65), 64), mload(add(proofData, 64)))
-            mstore(add(add(signature, 65), 96), mload(add(proofData, 96)))
-            mstore(add(add(signature, 65), 128), mload(add(proofData, 128)))
-            mstore(add(add(signature, 65), 160), mload(add(proofData, 160)))
-            mstore(add(add(signature, 65), 192), mload(add(proofData, 192)))
-            mstore(add(add(signature, 65), 224), mload(add(proofData, 224)))
-            mstore(add(add(signature, 65), 256), mload(add(proofData, 256)))
-        }
-        return signature;
+    function encodeSignature(bytes memory proofData) public pure returns (bytes memory res) {
+        bytes memory sigBuffer = new bytes(77);
+        res = LibBytes.concat(sigBuffer, proofData);
     }
 
     /// @notice Create a test data for UserOperations.
@@ -34,7 +20,7 @@ contract TestUtils {
         returns (PackedUserOperation[] memory)
     {
         PackedUserOperation[] memory uOps = new PackedUserOperation[](proofs.length);
-        for (uint256 i = 0; i < proofs.length; ++i) {
+        for (uint256 i = 0; i < proofs.length; i++) {
             bytes memory signature = encodeSignature(proofs[i]);
             PackedUserOperation memory uo = PackedUserOperation({
                 sender: sender,

contracts/test/mocks/MockAccount.sol

@@ -26,4 +26,8 @@ contract MockAccount is IAccount, IAccountExecute {
     function executeUserOp(PackedUserOperation calldata userOp, bytes32 userOpHash) external {
         // Do nothing
     }
+
+    function getThreshold() external pure returns (uint256) {
+        return 1;
+    }
 }