add updateAndWithdraw function to OmnichainStaking

zerolend · Aug 6, 2024 · eb8c73b · eb8c73b
1 parent 24bb916
commit eb8c73b
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 22 deletions.
diff --git a/contracts/interfaces/IOmnichainStaking.sol b/contracts/interfaces/IOmnichainStaking.sol
@@ -19,7 +19,7 @@ interface IOmnichainStaking is IVotes {
     // An omni-chain staking contract that allows users to stake their veNFT
     // and get some voting power. Once staked the voting power is available cross-chain.
 
-    function unstakeToken(uint256 tokenId) external;
+    function unstakeAndWithdraw(uint256 tokenId) external;
 
     // function totalSupply() external view returns (uint256);
 

diff --git a/contracts/locker/BaseLocker.sol b/contracts/locker/BaseLocker.sol
@@ -292,7 +292,8 @@ contract BaseLocker is
     /// @dev Only possible if the lock has expired
     function withdraw(uint256 _tokenId) public virtual nonReentrant {
         require(
-            _isAuthorized(ownerOf(_tokenId), msg.sender, _tokenId),
+            _isAuthorized(ownerOf(_tokenId), msg.sender, _tokenId) ||
+            address(staking) == msg.sender,
             "caller is not owner nor approved"
         );
 

diff --git a/contracts/locker/staking/OmnichainStakingBase.sol b/contracts/locker/staking/OmnichainStakingBase.sol
@@ -165,27 +165,13 @@ abstract contract OmnichainStakingBase is
     }
 
     /**
-     * @dev Unstakes a regular token NFT and transfers it back to the user.
-     * @param tokenId The ID of the regular token NFT to unstake.
+     * @notice A single withdraw function to unstake NFT, transfer it back to
+     * the user, and also withdraw all tokens for `tokenId` from the locker.
+     * @param tokenId The ID of the regular token NFT to unstake and withdraw.
      */
-    function unstakeToken(uint256 tokenId) external updateReward(msg.sender) {
-        require(lockedByToken[tokenId] != address(0), "!tokenId");
-        address lockedBy_ = lockedByToken[tokenId];
-        if (_msgSender() != lockedBy_)
-            revert InvalidUnstaker(_msgSender(), lockedBy_);
-
-        delete lockedByToken[tokenId];
-        lockedTokenIdNfts[_msgSender()] = deleteAnElement(
-            lockedTokenIdNfts[_msgSender()],
-            tokenId
-        );
-
-        // reset and burn voting power
-        _burn(msg.sender, tokenPower[tokenId]);
-        tokenPower[tokenId] = 0;
-        votingPowerCombined.reset(msg.sender);
-
-        locker.safeTransferFrom(address(this), msg.sender, tokenId);
+    function unstakeAndWithdraw(uint256 tokenId) external {
+        _unstakeToken(tokenId);
+        locker.withdraw(tokenId);
     }
 
     /**
@@ -408,6 +394,30 @@ abstract contract OmnichainStakingBase is
         }
     }
 
+    /**
+     * @dev Unstakes a regular token NFT and transfers it back to the user.
+     * @param tokenId The ID of the regular token NFT to unstake.
+     */
+    function _unstakeToken(uint256 tokenId) internal updateReward(msg.sender) {
+        require(lockedByToken[tokenId] != address(0), "!tokenId");
+        address lockedBy_ = lockedByToken[tokenId];
+        if (_msgSender() != lockedBy_)
+            revert InvalidUnstaker(_msgSender(), lockedBy_);
+
+        delete lockedByToken[tokenId];
+        lockedTokenIdNfts[_msgSender()] = deleteAnElement(
+            lockedTokenIdNfts[_msgSender()],
+            tokenId
+        );
+
+        // reset and burn voting power
+        _burn(msg.sender, tokenPower[tokenId]);
+        tokenPower[tokenId] = 0;
+        votingPowerCombined.reset(msg.sender);
+
+        locker.safeTransferFrom(address(this), msg.sender, tokenId);
+    }
+
     /**
      * @dev Deletes an element from an array.
      * @param elements The array to delete from.