test: signature aggregation fuzz/fail cases

contracts/test/PBHSignatureAggregator.t.sol

@@ -9,6 +9,7 @@ import {IPBHEntryPoint} from "../src/interfaces/IPBHEntryPoint.sol";
 import "@account-abstraction/contracts/interfaces/PackedUserOperation.sol";
 import "@BokkyPooBahsDateTimeLibrary/BokkyPooBahsDateTimeLibrary.sol";
 import "../src/helpers/PBHExternalNullifier.sol";
+import {PBHSignatureAggregator} from "../src/PBHSignatureAggregator.sol";
 
 contract PBHSignatureAggregatorTest is TestUtils, Setup {
     function setUp() public override {
@@ -40,7 +41,7 @@ contract PBHSignatureAggregatorTest is TestUtils, Setup {
         proofs[0] = abi.encode(proof0);
         proofs[1] = abi.encode(proof1);
 
-        PackedUserOperation[] memory uoTestFixture = createUOTestData(address(safe), proofs);
+        PackedUserOperation[] memory uoTestFixture = createUOTestData(address(safe), proofs, 1);
         bytes memory aggregatedSignature = pbhAggregator.aggregateSignatures(uoTestFixture);
 
         IEntryPoint.UserOpsPerAggregator[] memory userOpsPerAggregator = new IEntryPoint.UserOpsPerAggregator[](1);
@@ -64,7 +65,7 @@ contract PBHSignatureAggregatorTest is TestUtils, Setup {
         bytes[] memory proofs = new bytes[](2);
         proofs[0] = abi.encode(proof);
         proofs[1] = abi.encode(proof);
-        PackedUserOperation[] memory uoTestFixture = createUOTestData(address(safe), proofs);
+        PackedUserOperation[] memory uoTestFixture = createUOTestData(address(safe), proofs, 1);
         bytes memory aggregatedSignature = pbhAggregator.aggregateSignatures(uoTestFixture);
         IPBHEntryPoint.PBHPayload[] memory decodedProofs =
             abi.decode(aggregatedSignature, (IPBHEntryPoint.PBHPayload[]));
@@ -82,5 +83,61 @@ contract PBHSignatureAggregatorTest is TestUtils, Setup {
         assertEq(decodedProofs[1].nullifierHash, proof.nullifierHash, "Nullifier Hash should match");
     }
 
+    function testAggregateSignatures_VariableThreshold(
+        uint256 root,
+        uint256 pbhExternalNullifier,
+        uint256 nullifierHash,
+        uint8 threshold
+    ) public {
+        deploySafeAccount(address(pbhAggregator), threshold);
+        IPBHEntryPoint.PBHPayload memory proof = IPBHEntryPoint.PBHPayload({
+            root: root,
+            pbhExternalNullifier: pbhExternalNullifier,
+            nullifierHash: nullifierHash,
+            proof: [uint256(0), 0, 0, 0, 0, 0, 0, 0]
+        });
+
+        bytes[] memory proofs = new bytes[](2);
+        proofs[0] = abi.encode(proof);
+        proofs[1] = abi.encode(proof);
+        PackedUserOperation[] memory uoTestFixture = createUOTestData(address(safe), proofs, threshold);
+        bytes memory aggregatedSignature = pbhAggregator.aggregateSignatures(uoTestFixture);
+        IPBHEntryPoint.PBHPayload[] memory decodedProofs =
+            abi.decode(aggregatedSignature, (IPBHEntryPoint.PBHPayload[]));
+        assertEq(decodedProofs.length, 2, "Decoded proof length should be 1");
+        assertEq(decodedProofs[0].root, proof.root, "Root should match");
+        assertEq(
+            decodedProofs[0].pbhExternalNullifier, proof.pbhExternalNullifier, "PBH External Nullifier should match"
+        );
+        assertEq(decodedProofs[0].nullifierHash, proof.nullifierHash, "Nullifier Hash should match");
+
+        assertEq(decodedProofs[1].root, proof.root, "Root should match");
+        assertEq(
+            decodedProofs[1].pbhExternalNullifier, proof.pbhExternalNullifier, "PBH External Nullifier should match"
+        );
+        assertEq(decodedProofs[1].nullifierHash, proof.nullifierHash, "Nullifier Hash should match");
+    }
+
+    function testFailAggregateSignatures_InvalidSignatureLength(
+        uint256 root,
+        uint256 pbhExternalNullifier,
+        uint256 nullifierHash
+    ) public {
+        IPBHEntryPoint.PBHPayload memory proof = IPBHEntryPoint.PBHPayload({
+            root: root,
+            pbhExternalNullifier: pbhExternalNullifier,
+            nullifierHash: nullifierHash,
+            proof: [uint256(0), 0, 0, 0, 0, 0, 0, 0]
+        });
+
+        bytes[] memory proofs = new bytes[](2);
+        proofs[0] = abi.encode(proof);
+        proofs[1] = abi.encode(proof);
+        PackedUserOperation[] memory uoTestFixture = createUOTestData(address(safe), proofs, 1);
+        uoTestFixture[0].signature = new bytes(12);
+        pbhAggregator.aggregateSignatures(uoTestFixture);
+        vm.expectRevert(PBHSignatureAggregator.InvalidSignatureLength.selector);
+    }
+
     receive() external payable {}
 }

contracts/test/Setup.sol

@@ -47,7 +47,7 @@ contract Setup is Test {
         deployWorldIDGroups();
         deployPBHEntryPoint(worldIDGroups, entryPoint);
         deployPBHSignatureAggregator(address(pbhEntryPoint));
-        deploySafeAccount(address(pbhAggregator));
+        deploySafeAccount(address(pbhAggregator), 1);
 
         // Label the addresses for better errors.
         vm.label(address(entryPoint), "ERC-4337 Entry Point");
@@ -91,8 +91,8 @@ contract Setup is Test {
 
     /// @notice Initializes a new safe account.
     /// @dev It is constructed in the globals.
-    function deploySafeAccount(address _pbhSignatureAggregator) public {
-        safe = new MockAccount(_pbhSignatureAggregator);
+    function deploySafeAccount(address _pbhSignatureAggregator, uint256 threshold) public {
+        safe = new MockAccount(_pbhSignatureAggregator, threshold);
     }
 
     /// @notice Initializes a new World ID Groups contract.

contracts/test/TestUtils.sol

@@ -8,20 +8,24 @@ import "@forge-std/console.sol";
 import "@solady/LibBytes.sol";
 
 contract TestUtils {
-    function encodeSignature(bytes memory proofData) public pure returns (bytes memory res) {
-        bytes memory sigBuffer = new bytes(77);
+    function encodeSignature(bytes memory proofData, uint256 signatureThreshold)
+        public
+        pure
+        returns (bytes memory res)
+    {
+        bytes memory sigBuffer = new bytes(65 * signatureThreshold + 12);
         res = LibBytes.concat(sigBuffer, proofData);
     }
 
     /// @notice Create a test data for UserOperations.
-    function createUOTestData(address sender, bytes[] memory proofs)
+    function createUOTestData(address sender, bytes[] memory proofs, uint256 signatureThreshold)
         public
         pure
         returns (PackedUserOperation[] memory)
     {
         PackedUserOperation[] memory uOps = new PackedUserOperation[](proofs.length);
         for (uint256 i = 0; i < proofs.length; i++) {
-            bytes memory signature = encodeSignature(proofs[i]);
+            bytes memory signature = encodeSignature(proofs[i], signatureThreshold);
             PackedUserOperation memory uo = PackedUserOperation({
                 sender: sender,
                 nonce: i,

contracts/test/mocks/MockAccount.sol

@@ -8,9 +8,11 @@ import "@account-abstraction/contracts/interfaces/PackedUserOperation.sol";
 
 contract MockAccount is IAccount, IAccountExecute {
     address public pbhAggregator;
+    uint256 public immutable threshold;
 
-    constructor(address _pbhAggregator) {
+    constructor(address _pbhAggregator, uint256 _threshold) {
         pbhAggregator = _pbhAggregator;
+        threshold = _threshold;
     }
 
     function validateUserOp(PackedUserOperation calldata, bytes32, uint256)
@@ -27,7 +29,7 @@ contract MockAccount is IAccount, IAccountExecute {
         // Do nothing
     }
 
-    function getThreshold() external pure returns (uint256) {
-        return 1;
+    function getThreshold() external view returns (uint256) {
+        return threshold;
     }
 }