From ca180f2f0dd1574ffc4f789995782017765b8778 Mon Sep 17 00:00:00 2001 From: vimageDE Date: Mon, 4 Nov 2024 15:54:24 +0100 Subject: [PATCH] register attest via signature + tests --- src/examples/allocator/ServerAllocator.sol | 108 +++++-- test/ServerAllocator.t.sol | 312 ++++++++++++++++++++- 2 files changed, 389 insertions(+), 31 deletions(-) diff --git a/src/examples/allocator/ServerAllocator.sol b/src/examples/allocator/ServerAllocator.sol index 6bfddab..f775bfc 100644 --- a/src/examples/allocator/ServerAllocator.sol +++ b/src/examples/allocator/ServerAllocator.sol @@ -13,6 +13,13 @@ import {IERC1271} from "lib/openzeppelin-contracts/contracts/interfaces/IERC1271 contract ServerAllocator is Ownable2Step, EIP712, IAllocator { using ECDSA for bytes32; + struct RegisterAttest { + address signer; + bytes32 attestHash; + uint256 expiration; + uint256 nonce; + } + struct NonceConsumption { address signer; uint256[] nonces; @@ -22,6 +29,10 @@ contract ServerAllocator is Ownable2Step, EIP712, IAllocator { // keccak256("Attest(address,address,address,uint256,uint256)") bytes4 private constant _ATTEST_SELECTOR = 0x1a808f91; + // keccak256("RegisterAttest(address signer,bytes32 attestHash,uint256 expiration,uint256 nonce)") + bytes32 private constant _ATTEST_TYPE_HASH = + 0xaf2dfd3fe08723f490d203be627da2725f4ad38681e455221da2fc1a633bbb18; + // keccak256("Allocator(bytes32 hash)") bytes32 private constant _ALLOCATOR_TYPE_HASH = 0xcdf324dc7c3490a07fbbb105911393dcbc0676ac7c6c1c32c786721de6179e70; @@ -34,8 +45,10 @@ contract ServerAllocator is Ownable2Step, EIP712, IAllocator { mapping(address => uint256) private _signers; address[] private _activeSigners; + mapping(bytes32 => uint256) private _attestExpirations; mapping(bytes32 => uint256) private _attestCounts; + mapping(bytes32 => bool) private _attestSignatures; event SignerAdded(address signer_); event SignerRemoved(address signer_); @@ -49,6 +62,7 @@ contract ServerAllocator is Ownable2Step, EIP712, IAllocator { error InvalidCaller(address caller_, address expected_); error InvalidSigner(address signer_); error InvalidSignature(bytes signature_, address signer_); + error AlreadyUsedSig(bytes32 attest_, uint256 nonce); error InvalidInput(); modifier isSigner(address signer_) { @@ -95,15 +109,34 @@ contract ServerAllocator is Ownable2Step, EIP712, IAllocator { bytes32 attest_, uint256 expiration_ ) external isSigner(msg.sender) { - if (expiration_ < block.timestamp) { - revert Expired(expiration_, block.timestamp); - } - uint256 count = ++_attestCounts[attest_]; - bytes32 countedAttest = keccak256(abi.encode(attest_, count)); + _registerAttest(attest_, expiration_); + } - _attestExpirations[countedAttest] = expiration_; + /// @dev Nonce management in the RegisterAttest is only required for multiple registers of the same attest with the same expiration. + function registerAttestViaSignature( + RegisterAttest calldata attest_, + bytes calldata signature_ + ) external { + bytes32 _attestWithNonce = keccak256( + abi.encode(attest_.attestHash, attest_.expiration, attest_.nonce) + ); + if (_attestSignatures[_attestWithNonce]) { + revert AlreadyUsedSig(attest_.attestHash, attest_.nonce); + } + address signer = _validateSignedAttest( + attest_.signer, + attest_.attestHash, + attest_.expiration, + attest_.nonce, + signature_ + ); + if (signer != attest_.signer || !_containsSigner(signer)) { + revert InvalidSignature(signature_, signer); + } - emit AttestRegistered(attest_, expiration_); + // Invalidate signature + _attestSignatures[_attestWithNonce] = true; + _registerAttest(attest_.attestHash, attest_.expiration); } /// @dev There is no way to uniquely identify a transfer, so the contract relies on its own accounting of registered attests. @@ -186,21 +219,6 @@ contract ServerAllocator is Ownable2Step, EIP712, IAllocator { ) external view returns (bytes4 magicValue) { address signer = _validateSignedHash(hash_, signature_); if (!_containsSigner(signer)) { - // Check registered attests as fallback - /// TODO: This fallback must modify state to not be a source of endless verifications. - // uint256 count = _attestCounts[hash_]; - // if (count != 0) { - // for (uint256 i = count; i > 0; --i) { - // bytes32 countedAttest = keccak256(abi.encode(hash_, i)); - // if (_attestExpirations[countedAttest] >= block.timestamp) { - // // Found a valid registered attest - - // // _attestCounts[hash_] = --count; - // // delete _attestExpirations[countedAttest]; - // return IERC1271.isValidSignature.selector; - // } - // } - // } revert InvalidSignature(signature_, signer); } return IERC1271.isValidSignature.selector; @@ -235,6 +253,18 @@ contract ServerAllocator is Ownable2Step, EIP712, IAllocator { return _COMPACT_CONTRACT; } + function _registerAttest(bytes32 attest_, uint256 expiration_) internal { + if (expiration_ < block.timestamp) { + revert Expired(expiration_, block.timestamp); + } + uint256 count = ++_attestCounts[attest_]; + bytes32 countedAttest = keccak256(abi.encode(attest_, count)); + + _attestExpirations[countedAttest] = expiration_; + + emit AttestRegistered(attest_, expiration_); + } + /// Todo: This will lead to always the last registered hash being consumed. function _consumeNonces( uint256[] calldata nonces_, @@ -258,6 +288,37 @@ contract ServerAllocator is Ownable2Step, EIP712, IAllocator { emit NoncesConsumed(nonces_); } + function _validateSignedAttest( + address signer_, + bytes32 hash_, + uint256 expiration_, + uint256 nonce, + bytes calldata signature_ + ) internal view returns (address) { + bytes32 message = _hashAttest(signer_, hash_, expiration_, nonce); + return message.recover(signature_); + } + + function _hashAttest( + address signer_, + bytes32 hash_, + uint256 expiration_, + uint256 nonce_ + ) internal view returns (bytes32) { + return + _hashTypedDataV4( + keccak256( + abi.encode( + _ATTEST_TYPE_HASH, + signer_, + hash_, + expiration_, + nonce_ + ) + ) + ); + } + function _validateSignedHash( bytes32 hash_, bytes calldata signature_ @@ -290,7 +351,8 @@ contract ServerAllocator is Ownable2Step, EIP712, IAllocator { abi.encode( _NONCE_CONSUMPTION_TYPE_HASH, data_.signer, - data_.nonces + data_.nonces, + data_.attests ) ) ); diff --git a/test/ServerAllocator.t.sol b/test/ServerAllocator.t.sol index b8f0bd1..761d9fb 100644 --- a/test/ServerAllocator.t.sol +++ b/test/ServerAllocator.t.sol @@ -8,11 +8,14 @@ import {ServerAllocator} from "src/examples/allocator/ServerAllocator.sol"; import {TheCompactMock} from "src/test/TheCompactMock.sol"; import {ERC20Mock} from "src/test/ERC20Mock.sol"; import {console} from "forge-std/console.sol"; +import {IERC1271} from "lib/permit2/src/interfaces/IERC1271.sol"; abstract contract MocksSetup is Test { address owner = makeAddr("owner"); - address signer = makeAddr("signer"); - address attacker = makeAddr("attacker"); + address signer; + uint256 signerPK; + address attacker; + uint256 attackerPK; ERC20Mock usdc; TheCompactMock compactContract; ServerAllocator serverAllocator; @@ -26,6 +29,8 @@ abstract contract MocksSetup is Test { address(usdc), address(serverAllocator) ); + (signer, signerPK) = makeAddrAndKey("signer"); + (attacker, attackerPK) = makeAddrAndKey("attacker"); } } @@ -41,11 +46,17 @@ abstract contract AttestSetup { } } -abstract contract CreateHash { +abstract contract CreateHash is Test { + struct Allocator { + bytes32 hash; + } + // stringified types string EIP712_DOMAIN_TYPE = "EIP712Domain(string name,string version,uint256 chainId,address verifyingContract)"; // Hashed inside the funcion string ALLOCATOR_TYPE = "Allocator(bytes32 hash)"; // Hashed inside the funcion + string REGISTER_ATTEST_TYPE = + "RegisterAttest(address signer,bytes32 attestHash,uint256 expiration,uint256 nonce)"; // Hashed inside the funcion string NONCE_CONSUMPTION_TYPE = "NonceConsumption(address signer,uint256[] nonces,bytes32[] attests)"; // Hashed inside the funcion // EIP712 domain type @@ -53,7 +64,7 @@ abstract contract CreateHash { string version = "1"; function _hashAllocator( - bytes32 data, + Allocator memory data, address verifyingContract ) internal view returns (bytes32) { // hash typed data @@ -63,7 +74,29 @@ abstract contract CreateHash { "\x19\x01", // backslash is needed to escape the character _domainSeperator(verifyingContract), keccak256( - abi.encode(keccak256(bytes(ALLOCATOR_TYPE)), data) + abi.encode(keccak256(bytes(ALLOCATOR_TYPE)), data.hash) + ) + ) + ); + } + + function _hashRegisterAttest( + ServerAllocator.RegisterAttest memory data, + address verifyingContract + ) internal view returns (bytes32) { + return + keccak256( + abi.encodePacked( + "\x19\x01", // backslash is needed to escape the character + _domainSeperator(verifyingContract), + keccak256( + abi.encode( + keccak256(bytes(REGISTER_ATTEST_TYPE)), + data.signer, + data.attestHash, + data.expiration, + data.nonce + ) ) ) ); @@ -105,6 +138,14 @@ abstract contract CreateHash { ) ); } + + function _signMessage( + bytes32 hash_, + uint256 signerPK_ + ) internal pure returns (bytes memory) { + (uint8 v, bytes32 r, bytes32 s) = vm.sign(signerPK_, hash_); + return abi.encodePacked(r, s, v); + } } abstract contract SignerSet is MocksSetup, CreateHash, AttestSetup { @@ -214,7 +255,7 @@ contract ServerAllocator_ManageSigners is MocksSetup, CreateHash { } contract ServerAllocator_Attest is SignerSet { - function test_fuzz_onlySignerCanRegisterAttest(address attacker_) public { + function test_fuzz_RegisterAttest_onlySigner(address attacker_) public { vm.assume(attacker_ != signer); vm.prank(attacker_); @@ -230,7 +271,9 @@ contract ServerAllocator_Attest is SignerSet { ); } - function test_fuzz_attestExpired(uint256 expiration_) public { + function test_fuzz_registerAttest_attestExpired( + uint256 expiration_ + ) public { vm.assume(expiration_ < vm.getBlockTimestamp()); vm.prank(signer); @@ -247,7 +290,7 @@ contract ServerAllocator_Attest is SignerSet { ); } - function test_registerAttest() public { + function test_registerAttest_successful() public { vm.prank(signer); bytes32 attest = createAttest(signer, usdcId, 100); uint256 expiration = vm.getBlockTimestamp() + 1 days; @@ -258,6 +301,75 @@ contract ServerAllocator_Attest is SignerSet { assertEq(serverAllocator.checkAttestExpirations(attest)[0], expiration); } + function test_registerAttestViaSignature_InvalidSignature() public { + bytes32 attest = createAttest(signer, usdcId, 100); + uint256 expiration = vm.getBlockTimestamp() + 1 days; + + ServerAllocator.RegisterAttest memory attestData = ServerAllocator + .RegisterAttest(signer, attest, expiration, 0); + bytes32 message = _hashRegisterAttest( + attestData, + address(serverAllocator) + ); + bytes memory signature = _signMessage(message, attackerPK); + + vm.prank(attacker); + vm.expectRevert( + abi.encodeWithSelector( + ServerAllocator.InvalidSignature.selector, + signature, + attacker + ) + ); + serverAllocator.registerAttestViaSignature(attestData, signature); + } + + function test_registerAttestViaSignature_successful() public { + bytes32 attest = createAttest(signer, usdcId, 100); + uint256 expiration = vm.getBlockTimestamp() + 1 days; + + ServerAllocator.RegisterAttest memory attestData = ServerAllocator + .RegisterAttest(signer, attest, expiration, 0); + bytes32 message = _hashRegisterAttest( + attestData, + address(serverAllocator) + ); + bytes memory signature = _signMessage(message, signerPK); + + vm.prank(attacker); + vm.expectEmit(address(serverAllocator)); + emit ServerAllocator.AttestRegistered(attest, expiration); + serverAllocator.registerAttestViaSignature(attestData, signature); + } + + function test_registerAttestViaSignature_AlreadyUsedSig() public { + bytes32 attest = createAttest(signer, usdcId, 100); + uint256 expiration = vm.getBlockTimestamp() + 1 days; + + ServerAllocator.RegisterAttest memory attestData = ServerAllocator + .RegisterAttest(signer, attest, expiration, 0); + bytes32 message = _hashRegisterAttest( + attestData, + address(serverAllocator) + ); + bytes memory signature = _signMessage(message, signerPK); + + vm.prank(attacker); + vm.expectEmit(address(serverAllocator)); + emit ServerAllocator.AttestRegistered(attest, expiration); + serverAllocator.registerAttestViaSignature(attestData, signature); + + vm.prank(attacker); + vm.expectRevert( + abi.encodeWithSelector( + ServerAllocator.AlreadyUsedSig.selector, + attest, + 0 + ) + ); + serverAllocator.registerAttestViaSignature(attestData, signature); + } + function test_registerSameAttestTwice() public { vm.startPrank(signer); bytes32 attest = createAttest(signer, usdcId, 100); @@ -491,4 +603,188 @@ contract ServerAllocator_Consume is SignerSet { ); serverAllocator.checkAttestExpirations(attests[2]); } + + function test_consumeViaSignature_requiresNoncesAndAttestsToBeOfSameLength() + public + { + bytes32 message = _hashNonceConsumption( + ServerAllocator.NonceConsumption( + signer, + new uint256[](0), + new bytes32[](1) + ), + address(serverAllocator) + ); + bytes memory signature = _signMessage(message, signerPK); + + vm.prank(signer); + vm.expectRevert( + abi.encodeWithSelector(ServerAllocator.InvalidInput.selector) + ); + serverAllocator.consumeViaSignature( + ServerAllocator.NonceConsumption( + signer, + new uint256[](0), + new bytes32[](1) + ), + signature + ); + } + + function test_consumeViaSignature_requireValidSignature() public { + bytes32 message = _hashNonceConsumption( + ServerAllocator.NonceConsumption( + signer, + new uint256[](1), + new bytes32[](1) + ), + address(serverAllocator) + ); + bytes memory signature = _signMessage(message, attackerPK); + + vm.prank(signer); + vm.expectRevert( + abi.encodeWithSelector( + ServerAllocator.InvalidSigner.selector, + attacker + ) + ); + serverAllocator.consumeViaSignature( + ServerAllocator.NonceConsumption( + signer, + new uint256[](1), + new bytes32[](1) + ), + signature + ); + } + + function test_consumeViaSignature_successfulWithoutAttests() public { + uint256[] memory nonces = new uint256[](3); + nonces[0] = 1; + nonces[1] = 2; + nonces[2] = 3; + + bytes32 message = _hashNonceConsumption( + ServerAllocator.NonceConsumption(signer, nonces, new bytes32[](3)), + address(serverAllocator) + ); + bytes memory signature = _signMessage(message, signerPK); + + vm.prank(attacker); + vm.expectEmit(address(serverAllocator)); + emit ServerAllocator.NoncesConsumed(nonces); + serverAllocator.consumeViaSignature( + ServerAllocator.NonceConsumption(signer, nonces, new bytes32[](3)), + signature + ); + } + + function test_consumeViaSignature_successfulWithAttests() public { + uint256[] memory nonces = new uint256[](3); + nonces[0] = 1; + nonces[1] = 2; + nonces[2] = 3; + + bytes32[] memory attests = new bytes32[](3); + attests[0] = createAttest(signer, usdcId, 100); + attests[1] = createAttest(signer, usdcId, 200); + attests[2] = createAttest(signer, usdcId, 300); + + vm.startPrank(signer); + // register attests + serverAllocator.registerAttest(attests[0], vm.getBlockTimestamp()); + serverAllocator.registerAttest(attests[1], vm.getBlockTimestamp()); + serverAllocator.registerAttest(attests[2], vm.getBlockTimestamp()); + vm.stopPrank(); + + assertEq( + serverAllocator.checkAttestExpirations(attests[0])[0], + vm.getBlockTimestamp() + ); + assertEq( + serverAllocator.checkAttestExpirations(attests[1])[0], + vm.getBlockTimestamp() + ); + assertEq( + serverAllocator.checkAttestExpirations(attests[2])[0], + vm.getBlockTimestamp() + ); + + bytes32 message = _hashNonceConsumption( + ServerAllocator.NonceConsumption(signer, nonces, attests), + address(serverAllocator) + ); + bytes memory signature = _signMessage(message, signerPK); + + vm.prank(attacker); + vm.expectEmit(address(serverAllocator)); + emit ServerAllocator.NoncesConsumed(nonces); + serverAllocator.consumeViaSignature( + ServerAllocator.NonceConsumption(signer, nonces, attests), + signature + ); + + assertEq(compactContract.consumedNonces(0), false); + for (uint256 i = 0; i < nonces.length; ++i) { + assertEq(compactContract.consumedNonces(nonces[i]), true); + } + + // check attests were consumed + vm.expectRevert( + abi.encodeWithSelector( + ServerAllocator.UnregisteredAttest.selector, + attests[0] + ) + ); + serverAllocator.checkAttestExpirations(attests[0]); + vm.expectRevert( + abi.encodeWithSelector( + ServerAllocator.UnregisteredAttest.selector, + attests[1] + ) + ); + serverAllocator.checkAttestExpirations(attests[1]); + vm.expectRevert( + abi.encodeWithSelector( + ServerAllocator.UnregisteredAttest.selector, + attests[2] + ) + ); + serverAllocator.checkAttestExpirations(attests[2]); + } +} + +contract ServerAllocator_isValidSignature is SignerSet { + function test_isValidSignature_revertInvalidSig() public { + bytes32 attest = createAttest(signer, usdcId, 100); + bytes32 message = _hashAllocator( + Allocator(attest), + address(serverAllocator) + ); + bytes memory signature = _signMessage(message, attackerPK); + + vm.prank(attacker); + vm.expectRevert( + abi.encodeWithSelector( + ServerAllocator.InvalidSignature.selector, + signature, + attacker + ) + ); + serverAllocator.isValidSignature(attest, signature); + } + + function test_isValidSignature_successful() public { + bytes32 attest = createAttest(signer, usdcId, 100); + bytes32 message = _hashAllocator( + Allocator(attest), + address(serverAllocator) + ); + bytes memory signature = _signMessage(message, signerPK); + + vm.prank(attacker); + bytes4 magicValue = serverAllocator.isValidSignature(attest, signature); + assertEq(magicValue, IERC1271.isValidSignature.selector); + } }