From ca431084696637575edffd568923878a4a70a51b Mon Sep 17 00:00:00 2001 From: Nemitari Ajienka Date: Fri, 14 Feb 2025 18:07:04 +0000 Subject: [PATCH] update SignatureSender logic and add upgradeability --- src/decryption-requests/DecryptionSender.sol | 4 +- src/interfaces/ISignatureSender.sol | 58 +++++++- src/libraries/TypesLib.sol | 2 + src/mocks/MockBN254SignatureScheme.sol | 2 +- src/mocks/RevertingReceiver.sol | 2 +- src/signature-requests/SignatureSender.sol | 143 ++++++++++++++++--- test/forge/Blocklock.t.sol | 4 - 7 files changed, 187 insertions(+), 28 deletions(-) diff --git a/src/decryption-requests/DecryptionSender.sol b/src/decryption-requests/DecryptionSender.sol index 55056f9..78e3748 100644 --- a/src/decryption-requests/DecryptionSender.sol +++ b/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); diff --git a/src/interfaces/ISignatureSender.sol b/src/interfaces/ISignatureSender.sol index 9def9cd..06149a8 100644 --- a/src/interfaces/ISignatureSender.sol +++ b/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); } diff --git a/src/libraries/TypesLib.sol b/src/libraries/TypesLib.sol index 1200880..a4d8ece 100644 --- a/src/libraries/TypesLib.sol +++ b/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 diff --git a/src/mocks/MockBN254SignatureScheme.sol b/src/mocks/MockBN254SignatureScheme.sol index de61d52..ea45796 100644 --- a/src/mocks/MockBN254SignatureScheme.sol +++ b/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})); } -} \ No newline at end of file +} diff --git a/src/mocks/RevertingReceiver.sol b/src/mocks/RevertingReceiver.sol index 9fe9ddf..7ed433d 100644 --- a/src/mocks/RevertingReceiver.sol +++ b/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 diff --git a/src/signature-requests/SignatureSender.sol b/src/signature-requests/SignatureSender.sol index 53f3aa6..5b4636a 100644 --- a/src/signature-requests/SignatureSender.sol +++ b/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"; } } diff --git a/test/forge/Blocklock.t.sol b/test/forge/Blocklock.t.sol index 37f844c..e4fbf15 100644 --- a/test/forge/Blocklock.t.sol +++ b/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)); } }