update SignatureSender logic and add upgradeability

src/decryption-requests/DecryptionSender.sol

@@ -10,7 +10,6 @@ import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet
 
 import {BLS} from "../libraries/BLS.sol";
 import {TypesLib} from "../libraries/TypesLib.sol";
-import {console} from "forge-std/console.sol";
 import {BytesLib} from "../libraries/BytesLib.sol";
 
 import {Multicall} from "@openzeppelin/contracts/utils/Multicall.sol";
@@ -61,7 +60,6 @@ contract DecryptionSender is
         uint256 requestedAt
     );
     event DecryptionReceiverCallbackSuccess(uint256 indexed requestID, bytes decryptionKey, bytes signature);
-
     event DecryptionReceiverCallbackFailed(uint256 requestID);
 
     modifier onlyOwner() {
@@ -169,7 +167,9 @@ contract DecryptionSender is
         requests[requestID].decryptionKey = decryptionKey;
         requests[requestID].signature = signature;
         requests[requestID].isFulfilled = true;
+
         unfulfilledRequestIds.remove(requestID);
+
         if (!success) {
             erroredRequestIds.add(requestID);
             emit DecryptionReceiverCallbackFailed(requestID);

src/interfaces/ISignatureSender.sol

@@ -28,6 +28,22 @@ interface ISignatureSender {
      */
     function fulfilSignatureRequest(uint256 requestID, bytes calldata signature) external;
 
+    /**
+     * @notice Retry an request that has previously failed during callback
+     * @dev This function is intended to be called after a signature has been generated off-chain but failed to
+     * call back into the originating contract.
+     *
+     * @param requestID The unique identifier for the signature request. This should match the ID used
+     *                  when the signature was initially requested.
+     */
+    function retryCallback(uint256 requestID) external;
+
+    /**
+     * @notice Updates the signature scheme address provider contract address
+     * @param newSignatureSchemeAddressProvider The signature address provider address to set
+     */
+    function setSignatureSchemeAddressProvider(address newSignatureSchemeAddressProvider) external;
+
     /// Getters
 
     /**
@@ -39,11 +55,18 @@ interface ISignatureSender {
     function isInFlight(uint256 requestID) external view returns (bool);
 
     /**
-     * @notice Returns request data if a signature request is still in flight.
+     * @notice Returns request data.
      * @param requestID The unique identifier of the signature request.
-     * @return The corresponding SignatureRequest struct if the request is still in flight, otherwise struct with zero values.
+     * @return The corresponding SignatureRequest struct for the request Id.
      */
-    function getRequestInFlight(uint256 requestID) external view returns (TypesLib.SignatureRequest memory);
+    function getRequest(uint256 requestID) external view returns (TypesLib.SignatureRequest memory);
+
+    /**
+     * @notice returns whether a specific request errored during callback or not.
+     * @param requestID The ID of the request to check.
+     * @return boolean indicating whether the request has errored or not.
+     */
+    function hasErrored(uint256 requestID) external view returns (bool);
 
     /**
      * @notice Retrieves the public key associated with the signature process.
@@ -57,4 +80,33 @@ interface ISignatureSender {
      * @return Bytes string representing the public key points on the elliptic curve.
      */
     function getPublicKeyBytes() external view returns (bytes memory);
+
+    /**
+     * @notice Returns all the fulfilled request ids.
+     * @return A uint array representing a set containing all fulfilled request ids.
+     */
+    function getAllFulfilledRequestIds() external view returns (uint256[] memory);
+
+    /**
+     * @notice Returns all the request ids that are yet to be fulfilled.
+     * @return A uint array representing a set containing all request ids that are yet to be fulfilled.
+     */
+    function getAllUnfulfilledRequestIds() external view returns (uint256[] memory);
+
+    /**
+     * @notice Returns all the request ids where the callback reverted but a decryption key was provided, i.e., "fulfilled" but still in flight.
+     * @return A uint array representing a set containing all request ids with reverting callbacks.
+     */
+    function getAllErroredRequestIds() external view returns (uint256[] memory);
+
+    /**
+     * @notice Returns count of all the request ids that are yet to be fulfilled.
+     * @return A uint representing a count of all request ids that are yet to be fulfilled.
+     */
+    function getCountOfUnfulfilledRequestIds() external view returns (uint256);
+
+    /**
+     * @dev Returns the version number of the upgradeable contract.
+     */
+    function version() external pure returns (string memory);
 }

src/libraries/TypesLib.sol

@@ -11,6 +11,8 @@ library TypesLib {
         bytes condition; // optional condition, length can be zero for immediate message signing
         string schemeID; // signature scheme id, e.g., "BN254", "BLS12-381", "TESS"
         address callback; // the requester address to call back. Must implement ISignatureReceiver interface to support the required callback
+        bytes signature;
+        bool isFulfilled;
     }
 
     // Blocklock request stores details needed to generate blocklock decryption keys

src/mocks/MockBN254SignatureScheme.sol

@@ -43,4 +43,4 @@ contract MockBN254SignatureScheme is ISignatureScheme {
         (uint256 x, uint256 y) = hashToPoint(message);
         return BLS.g1Marshal(BLS.PointG1({x: x, y: y}));
     }
-}
+}

src/mocks/RevertingReceiver.sol

@@ -22,7 +22,7 @@ contract RevertingReceiver is AbstractBlocklockReceiver {
         return requestId;
     }
 
-    function receiveBlocklock(uint256 /*requestID*/, bytes calldata /*decryptionKey*/)
+    function receiveBlocklock(uint256, /*requestID*/ bytes calldata /*decryptionKey*/ )
         external
         view
         override

src/signature-requests/SignatureSender.sol

@@ -1,12 +1,20 @@
 // SPDX-License-Identifier: MIT
 pragma solidity 0.8.24;
 
+import {AccessControlEnumerableUpgradeable} from
+    "@openzeppelin/contracts-upgradeable/access/extensions/AccessControlEnumerableUpgradeable.sol";
+import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+import {ContextUpgradeable} from "@openzeppelin/contracts-upgradeable/utils/ContextUpgradeable.sol";
+import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
+
 import {BLS} from "../libraries/BLS.sol";
 import {TypesLib} from "../libraries/TypesLib.sol";
 import {BytesLib} from "../libraries/BytesLib.sol";
 
-import {AccessControl} from "@openzeppelin/contracts/access/AccessControl.sol";
 import {Multicall} from "@openzeppelin/contracts/utils/Multicall.sol";
+import {Context} from "@openzeppelin/contracts/utils/Context.sol";
+import {AccessControl} from "@openzeppelin/contracts/access/AccessControl.sol";
 
 import {ISignatureReceiver} from "../interfaces/ISignatureReceiver.sol";
 import {ISignatureSender} from "../interfaces/ISignatureSender.sol";
@@ -16,17 +24,30 @@ import {ISignatureSchemeAddressProvider} from "../interfaces/ISignatureSchemeAdd
 /// @notice Smart Contract for Conditional Threshold Signing of messages sent within signature requests.
 /// @notice Signatures are sent in callbacks to contract addresses implementing the SignatureReceiverBase abstract contract which implements the ISignatureReceiver interface.
 /// @notice Signature requests can also be made for requests requiring immediate signing of messages as the conditions are optional.
-contract SignatureSender is ISignatureSender, AccessControl, Multicall {
+contract SignatureSender is
+    ISignatureSender,
+    Multicall,
+    Initializable,
+    UUPSUpgradeable,
+    AccessControlEnumerableUpgradeable
+{
     using BytesLib for bytes;
+    using EnumerableSet for EnumerableSet.UintSet;
 
     bytes32 public constant ADMIN_ROLE = keccak256("ADMIN_ROLE");
 
     uint256 public lastRequestID = 0;
     BLS.PointG2 private publicKey = BLS.PointG2({x: [uint256(0), uint256(0)], y: [uint256(0), uint256(0)]});
-    mapping(uint256 => TypesLib.SignatureRequest) public requestsInFlight;
+    // mapping request ids to signature request structs
+    mapping(uint256 => TypesLib.SignatureRequest) public requests;
 
-    ISignatureSchemeAddressProvider public immutable signatureSchemeAddressProvider;
+    ISignatureSchemeAddressProvider public signatureSchemeAddressProvider;
 
+    EnumerableSet.UintSet private fulfilledRequestIds;
+    EnumerableSet.UintSet private unfulfilledRequestIds;
+    EnumerableSet.UintSet private erroredRequestIds;
+
+    event SignatureSchemeAddressProviderUpdated(address indexed newSignatureSchemeAddressProvider);
     event SignatureRequested(
         uint256 indexed requestID,
         address indexed callback,
@@ -36,16 +57,28 @@ contract SignatureSender is ISignatureSender, AccessControl, Multicall {
         bytes condition,
         uint256 requestedAt
     );
-    event SignatureRequestFulfilled(uint256 indexed requestID);
-
-    error SignatureCallbackFailed(uint256 requestID);
+    event SignatureRequestFulfilled(uint256 indexed requestID, bytes signature);
+    event SignatureCallbackFailed(uint256 requestID);
 
     modifier onlyOwner() {
         _checkRole(ADMIN_ROLE);
         _;
     }
 
-    constructor(uint256[2] memory x, uint256[2] memory y, address owner, address _signatureSchemeAddressProvider) {
+    /// @custom:oz-upgrades-unsafe-allow constructor
+    constructor() {
+        _disableInitializers();
+    }
+
+    function initialize(
+        uint256[2] memory x,
+        uint256[2] memory y,
+        address owner,
+        address _signatureSchemeAddressProvider
+    ) public initializer {
+        __UUPSUpgradeable_init();
+        __AccessControlEnumerable_init();
+
         publicKey = BLS.PointG2({x: x, y: y});
         require(_grantRole(ADMIN_ROLE, owner), "Grant role failed");
         require(_grantRole(DEFAULT_ADMIN_ROLE, owner), "Grant role reverts");
@@ -56,6 +89,21 @@ contract SignatureSender is ISignatureSender, AccessControl, Multicall {
         signatureSchemeAddressProvider = ISignatureSchemeAddressProvider(_signatureSchemeAddressProvider);
     }
 
+    // OVERRIDDEN UPGRADE FUNCTIONS
+    function _authorizeUpgrade(address newImplementation) internal override onlyOwner {}
+
+    function _msgSender() internal view override(Context, ContextUpgradeable) returns (address) {
+        return msg.sender;
+    }
+
+    function _msgData() internal pure override(Context, ContextUpgradeable) returns (bytes calldata) {
+        return msg.data;
+    }
+
+    function _contextSuffixLength() internal pure override(Context, ContextUpgradeable) returns (uint256) {
+        return 0;
+    }
+
     /**
      * @dev See {ISignatureSender-requestSignature}.
      */
@@ -78,15 +126,20 @@ contract SignatureSender is ISignatureSender, AccessControl, Multicall {
         ISignatureScheme sigScheme = ISignatureScheme(schemeContractAddress);
         bytes memory messageHash = sigScheme.hashToBytes(message);
 
-        requestsInFlight[lastRequestID] = TypesLib.SignatureRequest({
+        requests[lastRequestID] = TypesLib.SignatureRequest({
             callback: msg.sender,
             message: message,
             messageHash: messageHash,
             condition: condition,
-            schemeID: schemeID
+            schemeID: schemeID,
+            signature: hex"",
+            isFulfilled: false
         });
 
+        unfulfilledRequestIds.add(lastRequestID);
+
         emit SignatureRequested(lastRequestID, msg.sender, schemeID, message, messageHash, condition, block.timestamp);
+
         return lastRequestID;
     }
 
@@ -95,7 +148,7 @@ contract SignatureSender is ISignatureSender, AccessControl, Multicall {
      */
     function fulfilSignatureRequest(uint256 requestID, bytes calldata signature) external onlyOwner {
         require(isInFlight(requestID), "No request with specified requestID");
-        TypesLib.SignatureRequest memory request = requestsInFlight[requestID];
+        TypesLib.SignatureRequest memory request = requests[requestID];
 
         string memory schemeID = request.schemeID;
 
@@ -110,14 +163,43 @@ contract SignatureSender is ISignatureSender, AccessControl, Multicall {
         (bool success,) = request.callback.call(
             abi.encodeWithSelector(ISignatureReceiver.receiveSignature.selector, requestID, signature)
         );
+
+        requests[requestID].signature = signature;
+        requests[requestID].isFulfilled = true;
+
+        unfulfilledRequestIds.remove(requestID);
+
         if (!success) {
-            revert SignatureCallbackFailed(requestID);
+            erroredRequestIds.add(requestID);
+            emit SignatureCallbackFailed(requestID);
         } else {
-            emit SignatureRequestFulfilled(requestID);
-            delete requestsInFlight[requestID];
+            fulfilledRequestIds.add(requestID);
+            emit SignatureRequestFulfilled(requestID, signature);
         }
     }
 
+    function retryCallback(uint256 requestID) external {
+        require(hasErrored(requestID), "No request with specified requestID");
+        TypesLib.SignatureRequest memory request = requests[requestID];
+
+        (bool success,) = request.callback.call(
+            abi.encodeWithSelector(ISignatureReceiver.receiveSignature.selector, requestID, request.signature)
+        );
+
+        if (!success) {
+            emit SignatureCallbackFailed(requestID);
+        } else {
+            erroredRequestIds.remove(requestID);
+            fulfilledRequestIds.add(requestID);
+            emit SignatureRequestFulfilled(requestID, request.signature);
+        }
+    }
+
+    function setSignatureSchemeAddressProvider(address newSignatureSchemeAddressProvider) external onlyOwner {
+        signatureSchemeAddressProvider = ISignatureSchemeAddressProvider(newSignatureSchemeAddressProvider);
+        emit SignatureSchemeAddressProviderUpdated(newSignatureSchemeAddressProvider);
+    }
+
     /**
      * @dev See {ISignatureSender-getPublicKey}.
      */
@@ -136,13 +218,40 @@ contract SignatureSender is ISignatureSender, AccessControl, Multicall {
      * @dev See {ISignatureSender-isInFlight}.
      */
     function isInFlight(uint256 requestID) public view returns (bool) {
-        return requestsInFlight[requestID].callback != address(0);
+        return unfulfilledRequestIds.contains(requestID) || erroredRequestIds.contains(requestID);
+    }
+
+    function hasErrored(uint256 requestID) public view returns (bool) {
+        return erroredRequestIds.contains(requestID);
     }
 
     /**
      * @dev See {ISignatureSender-getRequestInFlight}.
      */
-    function getRequestInFlight(uint256 requestID) external view returns (TypesLib.SignatureRequest memory) {
-        return requestsInFlight[requestID];
+    function getRequest(uint256 requestID) external view returns (TypesLib.SignatureRequest memory) {
+        return requests[requestID];
+    }
+
+    function getAllFulfilledRequestIds() external view returns (uint256[] memory) {
+        return fulfilledRequestIds.values();
+    }
+
+    function getAllUnfulfilledRequestIds() external view returns (uint256[] memory) {
+        return unfulfilledRequestIds.values();
+    }
+
+    function getAllErroredRequestIds() external view returns (uint256[] memory) {
+        return erroredRequestIds.values();
+    }
+
+    function getCountOfUnfulfilledRequestIds() external view returns (uint256) {
+        return unfulfilledRequestIds.length();
+    }
+
+    /**
+     * @dev Returns the version number of the upgradeable contract.
+     */
+    function version() external pure returns (string memory) {
+        return "0.0.1";
     }
 }

test/forge/Blocklock.t.sol

@@ -4,7 +4,6 @@ pragma solidity 0.8.24;
 import {Test, console} from "forge-std/Test.sol";
 
 import {SignatureSchemeAddressProvider} from "../../src/signature-schemes/SignatureSchemeAddressProvider.sol";
-import {SignatureSender} from "../../src/signature-requests/SignatureSender.sol";
 import {BlocklockSender} from "../../src/blocklock/BlocklockSender.sol";
 import {BlocklockSignatureScheme} from "../../src/blocklock/BlocklockSignatureScheme.sol";
 import {DecryptionSender} from "../../src/decryption-requests/DecryptionSender.sol";
@@ -16,7 +15,6 @@ contract SimpleAuctionTest is Test {
     UUPSProxy decryptionSenderProxy;
     UUPSProxy blocklockSenderProxy;
 
-    SignatureSender public sigSender;
     DecryptionSender public decryptionSender;
     BlocklockSender public tlock;
 
@@ -65,7 +63,6 @@ contract SimpleAuctionTest is Test {
                 8044854403167346152897273335539146380878155193886184396711544300199836788154
             ]
         });
-        sigSender = new SignatureSender(pk.x, pk.y, owner, address(sigAddrProvider));
 
         // deploy implementation contracts for decryption and blocklock senders
         DecryptionSender decryptionSenderImplementationV1 = new DecryptionSender();
@@ -93,6 +90,5 @@ contract SimpleAuctionTest is Test {
         assertTrue(decryptionSender.hasRole(ADMIN_ROLE, owner));
         assert(address(tlock) != address(0));
         assert(address(decryptionSender) != address(0));
-        assert(address(sigSender) != address(0));
     }
 }