diff --git a/src/LlamaLocker.sol b/src/LlamaLocker.sol index 544ad51..6c96452 100644 --- a/src/LlamaLocker.sol +++ b/src/LlamaLocker.sol @@ -6,6 +6,7 @@ import {IERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol"; import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol"; import {IERC721} from "@openzeppelin/contracts/token/ERC721/ERC721.sol"; import {ERC721Holder} from "@openzeppelin/contracts/token/ERC721/utils/ERC721Holder.sol"; +import {ReentrancyGuard} from "@openzeppelin/contracts/utils/ReentrancyGuard.sol"; import {MerkleProof} from "@openzeppelin/contracts/utils/cryptography/MerkleProof.sol"; import {Math} from "@openzeppelin/contracts/utils/math/Math.sol"; import {SafeCast} from "@openzeppelin/contracts/utils/math/SafeCast.sol"; @@ -15,7 +16,7 @@ import {SafeCast} from "@openzeppelin/contracts/utils/math/SafeCast.sol"; * @author sepyke.eth * @dev Lock LLAMA to claim share of the yield generated by the treasury */ -contract LlamaLocker is ERC721Holder, Ownable2Step { +contract LlamaLocker is ERC721Holder, Ownable2Step, ReentrancyGuard { using SafeERC20 for IERC20; using SafeCast for uint256; using Math for uint256; @@ -154,6 +155,7 @@ contract LlamaLocker is ERC721Holder, Ownable2Step { uint256 modulo = lockedDurationInEpoch % LOCK_DURATION_IN_EPOCH; if (modulo != 0) revert InvalidUnlockWindow(); + locks[tokenId_] = NFTLock({owner: address(0), lockedAt: 0}); nft.safeTransferFrom(address(this), owner_, tokenId_); } @@ -222,7 +224,7 @@ contract LlamaLocker is ERC721Holder, Ownable2Step { } } - function claim(address recipient_) public returns (Claimable[] memory results) { + function claim(address recipient_) public nonReentrant returns (Claimable[] memory results) { if (recipient_ == address(0)) revert InvalidRecipient(); results = claimables(msg.sender); @@ -230,8 +232,8 @@ contract LlamaLocker is ERC721Holder, Ownable2Step { for (uint256 i = 0; i < resultsCount; ++i) { Claimable memory result = results[i]; if (result.amount > 0) { - IERC20(result.token).safeTransfer(recipient_, result.amount); claimedRewards[msg.sender][result.token] += result.amount; + IERC20(result.token).safeTransfer(recipient_, result.amount); } } }