added safe mint and reentrancy guard on the lock contract

aragon · Sep 25, 2024 · feba2fa · feba2fa
1 parent 98912bf
commit feba2fa
Show file tree

Hide file tree

Showing 5 changed files with 48 additions and 5 deletions.
diff --git a/src/escrow/increasing/Lock.sol b/src/escrow/increasing/Lock.sol
@@ -3,12 +3,13 @@ pragma solidity ^0.8.17;
 
 import {ILock} from "@escrow-interfaces/ILock.sol";
 import {ERC721EnumerableUpgradeable as ERC721Enumerable} from "@openzeppelin/contracts-upgradeable/token/ERC721/extensions/ERC721EnumerableUpgradeable.sol";
+import {ReentrancyGuardUpgradeable as ReentrancyGuard} from "@openzeppelin/contracts-upgradeable/security/ReentrancyGuardUpgradeable.sol";
 import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
 import {DaoAuthorizableUpgradeable as DaoAuthorizable} from "@aragon/osx/core/plugin/dao-authorizable/DaoAuthorizableUpgradeable.sol";
 import {IDAO} from "@aragon/osx/core/dao/IDAO.sol";
 
 /// @title NFT representation of an escrow locking mechanism
-contract Lock is ILock, ERC721Enumerable, UUPSUpgradeable, DaoAuthorizable {
+contract Lock is ILock, ERC721Enumerable, UUPSUpgradeable, DaoAuthorizable, ReentrancyGuard {
     /// @dev enables transfers without whitelisting
     address public constant WHITELIST_ANY_ADDRESS =
         address(uint160(uint256(keccak256("WHITELIST_ANY_ADDRESS"))));
@@ -57,6 +58,7 @@ contract Lock is ILock, ERC721Enumerable, UUPSUpgradeable, DaoAuthorizable {
     ) external initializer {
         __ERC721_init(_name, _symbol);
         __DaoAuthorizableUpgradeable_init(IDAO(_dao));
+        __ReentrancyGuard_init();
         escrow = _escrow;
 
         // allow sending nfts to the escrow
@@ -97,12 +99,13 @@ contract Lock is ILock, ERC721Enumerable, UUPSUpgradeable, DaoAuthorizable {
     }
 
     /// @notice Minting and burning functions that can only be called by the escrow contract
-    function mint(address _to, uint256 _tokenId) external onlyEscrow {
-        _mint(_to, _tokenId);
+    /// @dev Safe mint ensures contract addresses are ERC721 Receiver contracts
+    function mint(address _to, uint256 _tokenId) external onlyEscrow nonReentrant {
+        _safeMint(_to, _tokenId);
     }
 
     /// @notice Minting and burning functions that can only be called by the escrow contract
-    function burn(uint256 _tokenId) external onlyEscrow {
+    function burn(uint256 _tokenId) external onlyEscrow nonReentrant {
         _burn(_tokenId);
     }
 

diff --git a/test/escrow/escrow/EscrowCreateLock.t.sol b/test/escrow/escrow/EscrowCreateLock.t.sol
@@ -46,7 +46,7 @@ contract TestCreateLock is EscrowBase, IEscrowCurveUserStorage {
     /// @param _time is bound to 128 bits to avoid overflow - seems reasonable as is not a user input
     function testFuzz_createLock(uint128 _value, address _depositor, uint128 _time) public {
         vm.assume(_value > 0);
-        vm.assume(_depositor != address(0));
+        vm.assume(_depositor != address(0) && _depositor != address(vm));
 
         // set zero warmup for this test
         curve.setWarmupPeriod(0);

diff --git a/test/escrow/escrow/EscrowTransfers.t.sol b/test/escrow/escrow/EscrowTransfers.t.sol
@@ -48,4 +48,13 @@ contract TestEscrowTransfers is EscrowBase, IEscrowCurveUserStorage {
         assertEq(nftLock.balanceOf(address(123)), 1);
         assertEq(nftLock.balanceOf(address(this)), 0);
     }
+
+    function onERC721Received(
+        address,
+        address,
+        uint256,
+        bytes calldata
+    ) external pure returns (bytes4) {
+        return this.onERC721Received.selector;
+    }
 }
diff --git a/test/escrow/escrow/EscrowWithdraw.t.sol b/test/escrow/escrow/EscrowWithdraw.t.sol
@@ -33,6 +33,7 @@ contract TestWithdraw is EscrowBase, IEscrowCurveUserStorage, IGaugeVote {
     // setup a fee withdrawal
     function testFuzz_feeWithdrawal(uint64 _fee, uint128 _dep, address _who) public {
         vm.assume(_who != address(0) && _who != address(queue) && _who != address(escrow));
+        vm.assume(_who != address(vm)); // causes problems with erc721 receiver
         vm.assume(_dep > 0);
 
         if (_fee > 1e18) _fee = 1e18;

diff --git a/test/escrow/escrow/Lock.t.sol b/test/escrow/escrow/Lock.t.sol
@@ -69,4 +69,34 @@ contract TestLockMintBurn is EscrowBase, IEscrowCurveUserStorage, IGaugeVote {
         assertEq(nftLock.balanceOf(address(123)), 0);
         assertEq(nftLock.totalSupply(), 0);
     }
+
+    // HAL-14 receiver must implement ERC721Receiver
+    function testCannotMintToNonReceiver() public {
+        vm.prank(address(escrow));
+        vm.expectRevert("ERC721: transfer to non ERC721Receiver implementer");
+        nftLock.mint(address(this), 1);
+    }
+
+    // HAL-14 test reentrancy with safe mint
+    function testReentrantCantCallMint() public {
+        NFTReentrant reentrant = new NFTReentrant();
+
+        Lock newLock = _deployLock(address(reentrant), "name", "symbol", address(dao));
+
+        vm.prank(address(reentrant));
+        vm.expectRevert("revert");
+        newLock.mint(address(reentrant), 1);
+    }
+}
+
+contract NFTReentrant {
+    function onERC721Received(address, address, uint256, bytes memory) public returns (bytes4) {
+        (bool success, ) = msg.sender.call(
+            abi.encodeWithSignature("mint(address,uint256)", address(this), 1)
+        );
+        if (!success) {
+            revert("revert");
+        }
+        return this.onERC721Received.selector;
+    }
 }