diff --git a/src/Vault.sol b/src/Vault.sol index 6281f86..b3ea1dd 100644 --- a/src/Vault.sol +++ b/src/Vault.sol @@ -24,6 +24,8 @@ contract Vault is Initializable { address tokenCollection; uint256 tokenId; + mapping(address => uint256) unlockTimestamp; + function initialize( address _vaultRegistry, address _tokenCollection, @@ -59,11 +61,20 @@ contract Vault is Initializable { _; } + function lock(uint256 _unlockTimestamp) public payable onlyVault onlyOwner { + unlockTimestamp[ + IERC721(tokenCollection).ownerOf(tokenId) + ] = _unlockTimestamp; + } + function execTransaction( address payable to, uint256 value, bytes calldata data ) public payable onlyVault onlyOwner { + address owner = IERC721(tokenCollection).ownerOf(tokenId); + require(unlockTimestamp[owner] < block.timestamp, "Vault is locked"); + (bool success, bytes memory result) = to.call{value: value}(data); if (!success) { assembly { @@ -78,13 +89,15 @@ contract Vault is Initializable { onlyVault returns (bytes4 magicValue) { + address owner = IERC721(tokenCollection).ownerOf(tokenId); + bool isValid = SignatureChecker.isValidSignatureNow( - IERC721(tokenCollection).ownerOf(tokenId), + owner, _hash, _signature ); - if (isValid) { + if (isValid && unlockTimestamp[owner] < block.timestamp) { return IERC1271.isValidSignature.selector; } } diff --git a/test/Vault.t.sol b/test/Vault.t.sol index e9eb77c..1be4815 100644 --- a/test/Vault.t.sol +++ b/test/Vault.t.sol @@ -447,12 +447,12 @@ contract VaultCollectionTest is Test { bytes4 returnValue1 = vault.isValidSignature(hash, signature1); - assertEq(returnValue1, vault.isValidSignature.selector); + assertEq(returnValue1, IERC1271.isValidSignature.selector); } function testMessageSigningAndVerificationForUnauthorizedUser() public { address user1 = vm.addr(1); - uint256 tokenId = 11; + uint256 tokenId = 12; tokenCollection.mint(user1, tokenId); assertEq(tokenCollection.ownerOf(tokenId), user1); @@ -473,6 +473,112 @@ contract VaultCollectionTest is Test { assertEq(returnValue2, 0); } + + function testVaultLocksAndUnlocks() public { + address user1 = vm.addr(1); + uint256 tokenId = 13; + + tokenCollection.mint(user1, tokenId); + assertEq(tokenCollection.ownerOf(tokenId), user1); + + address payable vaultAddress = vaultRegistry.deployVault( + address(tokenCollection), + tokenId + ); + + vm.deal(vaultAddress, 1 ether); + + Vault vault = Vault(vaultAddress); + + // lock vault for 10 days + uint256 unlockTimestamp = block.timestamp + 10 days; + vm.prank(user1); + vault.lock(unlockTimestamp); + + // transaction should fail if vault is locked + vm.prank(user1); + vm.expectRevert(bytes("Vault is locked")); + vault.execTransaction(payable(user1), 0.1 ether, ""); + + // signing should fail if vault is locked + bytes32 hash = keccak256("This is a signed message"); + (uint8 v1, bytes32 r1, bytes32 s1) = vm.sign(2, hash); + bytes memory signature1 = abi.encodePacked(r1, s1, v1); + bytes4 returnValue = vault.isValidSignature(hash, signature1); + assertEq(returnValue, 0); + + // warp to timestamp after vault is unlocked + vm.warp(unlockTimestamp + 1 days); + + // transaction succeed now that vault lock has expired + vm.prank(user1); + vault.execTransaction(payable(user1), 1 ether, ""); + assertEq(user1.balance, 1 ether); + + // signing should now that vault lock has expired + bytes32 hashAfterUnlock = keccak256("This is a signed message"); + (uint8 v2, bytes32 r2, bytes32 s2) = vm.sign(1, hashAfterUnlock); + bytes memory signature2 = abi.encodePacked(r2, s2, v2); + bytes4 returnValue1 = vault.isValidSignature( + hashAfterUnlock, + signature2 + ); + assertEq(returnValue1, IERC1271.isValidSignature.selector); + } + + function testVaultUnlocksAfterTransfer() public { + address user1 = vm.addr(1); + address user2 = vm.addr(2); + uint256 tokenId = 14; + + tokenCollection.mint(user1, tokenId); + assertEq(tokenCollection.ownerOf(tokenId), user1); + + address payable vaultAddress = vaultRegistry.deployVault( + address(tokenCollection), + tokenId + ); + + vm.deal(vaultAddress, 1 ether); + + Vault vault = Vault(vaultAddress); + + // lock vault for 10 days + uint256 unlockTimestamp = block.timestamp + 10 days; + vm.prank(user1); + vault.lock(unlockTimestamp); + + // transaction should fail if vault is locked + vm.prank(user1); + vm.expectRevert(bytes("Vault is locked")); + vault.execTransaction(payable(user1), 0.1 ether, ""); + + // signing should fail if vault is locked + bytes32 hash = keccak256("This is a signed message"); + (uint8 v1, bytes32 r1, bytes32 s1) = vm.sign(2, hash); + bytes memory signature1 = abi.encodePacked(r1, s1, v1); + bytes4 returnValue = vault.isValidSignature(hash, signature1); + assertEq(returnValue, 0); + + // transfer vault to new owner + vm.prank(user1); + tokenCollection.safeTransferFrom(user1, user2, tokenId); + + // transaction succeed now that vault ownership has transferred + vm.prank(user2); + vault.execTransaction(payable(user2), 1 ether, ""); + assertEq(user2.balance, 1 ether); + + // signing should now that vault vault ownership has transferred + bytes32 hashAfterUnlock = keccak256("This is a signed message"); + (uint8 v2, bytes32 r2, bytes32 s2) = vm.sign(2, hashAfterUnlock); + bytes memory signature2 = abi.encodePacked(r2, s2, v2); + bytes4 returnValue1 = vault.isValidSignature( + hashAfterUnlock, + signature2 + ); + assertEq(returnValue1, IERC1271.isValidSignature.selector); + } } contract TokenCollection is ERC721 {