refactor: bytes signature library

src/libraries/ByteSignature.sol

@@ -2,46 +2,31 @@
 pragma solidity ^0.8.12;
 
 library ByteSignature {
-    function extractTwoSignatures(bytes memory _fullSignature)
+    function extractSignatures(bytes memory _fullSignature, uint256 count)
         internal
         pure
-        returns (bytes memory signature1, bytes memory signature2)
+        returns (bytes[] memory signatures)
     {
-        signature1 = new bytes(65);
-        signature2 = new bytes(65);
-        return (extractECDASignatureFromBytes(_fullSignature, 0), extractECDASignatureFromBytes(_fullSignature, 1));
-    }
-
-    function extractThreeSignatures(bytes memory _fullSignature)
-        internal
-        pure
-        returns (bytes memory signature1, bytes memory signature2, bytes memory signature3)
-    {
-        signature1 = new bytes(65);
-        signature2 = new bytes(65);
-        signature3 = new bytes(65);
-        return (
-            extractECDASignatureFromBytes(_fullSignature, 0),
-            extractECDASignatureFromBytes(_fullSignature, 1),
-            extractECDASignatureFromBytes(_fullSignature, 2)
-        );
+        require(_fullSignature.length >= count * 65, "ByteSignature: Invalid signature length");
+        signatures = new bytes[](count);
+        for (uint256 i = 0; i < count; i++) {
+            signatures[i] = extractECDASignatureFromBytes(_fullSignature, i);
+        }
     }
 
     function extractECDASignatureFromBytes(bytes memory _fullSignature, uint256 position)
         internal
         pure
         returns (bytes memory signature)
     {
+        uint256 offset = (position * 0x40) + position;
         signature = new bytes(65);
         // Copying the first signature. Note, that we need an offset of 0x20
         // since it is where the length of the `_fullSignature` is stored
-        uint256 firstIndex = (position * 0x40) + 0x20 + position;
-        uint256 secondIndex = (position * 0x40) + 0x40 + position;
-        uint256 thirdIndex = (position * 0x40) + 0x41 + position;
         assembly {
-            let r := mload(add(_fullSignature, firstIndex))
-            let s := mload(add(_fullSignature, secondIndex))
-            let v := and(mload(add(_fullSignature, thirdIndex)), 0xff)
+            let r := mload(add(_fullSignature, add(offset, 0x20)))
+            let s := mload(add(_fullSignature, add(offset, 0x40)))
+            let v := and(mload(add(_fullSignature, add(offset, 0x41))), 0xff)
 
             mstore(add(signature, 0x20), r)
             mstore(add(signature, 0x40), s)

src/wallet/KintoWallet.sol

@@ -328,12 +328,10 @@ contract KintoWallet is Initializable, BaseAccount, TokenCallbackHandler, IKinto
         }
         bytes[] memory signatures = new bytes[](owners.length);
         // Split signature from userOp.signature
-        if (requiredSigners == 2) {
-            (signatures[0], signatures[1]) = ByteSignature.extractTwoSignatures(userOp.signature);
-        } else if (requiredSigners == 3) {
-            (signatures[0], signatures[1], signatures[2]) = ByteSignature.extractThreeSignatures(userOp.signature);
-        } else {
+        if (requiredSigners == 1) {
             signatures[0] = userOp.signature;
+        } else {
+            signatures = ByteSignature.extractSignatures(userOp.signature, requiredSigners);
         }
         for (uint256 i = 0; i < owners.length; i++) {
             if (owners[i] == hash.recover(signatures[i])) {

test/BytesSignature.t.sol

@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.12;
+
+import "../src/libraries/ByteSignature.sol";
+import {UserOp} from "./helpers/UserOp.sol";
+
+contract ByteSignatureTest is UserOp {
+    function testExtractSingleSignature() public {
+        // dummy signature
+        bytes memory signature =
+            hex"ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff01";
+        bytes[] memory extracted = ByteSignature.extractSignatures(signature, 1);
+        assertEq(extracted[0], signature, "The extracted signature does not match the original signature.");
+    }
+
+    function testExtractMultipleSignatures() public {
+        // create 2 dummy signatures
+        bytes memory signature1 =
+            hex"ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff01";
+        bytes memory signature2 =
+            hex"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb1c";
+        bytes memory fullSignature = abi.encodePacked(signature1, signature2);
+
+        bytes[] memory extracted = ByteSignature.extractSignatures(fullSignature, 2);
+        assertEq(extracted[0], signature1, "The first signature does not match the original signature.");
+        assertEq(extracted[1], signature2, "The second signature does not match the original signature.");
+    }
+
+    function testExtract_RevertWhen_InvalidLength() public {
+        bytes memory invalidSignature = hex"abcd";
+        vm.expectRevert("ByteSignature: Invalid signature length");
+        ByteSignature.extractSignatures(invalidSignature, 1);
+    }
+}