From b0e2f1b61bdaaed6805146176fb3c7abddad8b39 Mon Sep 17 00:00:00 2001
From: Ricardo Guilherme Schmidt <3esmit@gmail.com>
Date: Fri, 6 Oct 2023 12:05:20 -0300
Subject: [PATCH] Improve claimTokens

---
 contracts/SNTOptimismController.sol | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/contracts/SNTOptimismController.sol b/contracts/SNTOptimismController.sol
index 62cbc56..36ebdaa 100644
--- a/contracts/SNTOptimismController.sol
+++ b/contracts/SNTOptimismController.sol
@@ -53,18 +53,20 @@ contract SNTOptimismController is TokenController, Ownable2Step {
     }
 
     /**
-     * @notice Extract mistakenly sent tokens to this contract.
+     * @notice Send tokens or ether from this contract to owner.
      * @param _token Token contract to recover, 0 to extract ether.
      */
     function claimTokens(MiniMeBase _token) public onlyOwner {
+        uint256 balance;
         if (address(_token) == address(0)) {
-            payable(owner()).transfer(address(this).balance);
+            balance = address(this).balance;
+            payable(msg.sender).transfer(balance);
             return;
+        } else {
+            balance = _token.balanceOf(address(this));
+            _token.transfer(msg.sender, balance);
         }
-
-        uint256 balance = _token.balanceOf(address(this));
-        _token.transfer(owner(), balance);
-        emit ClaimedTokens(address(_token), owner(), balance);
+        emit ClaimedTokens(address(_token), msg.sender, balance);
     }
 
     event ClaimedTokens(address indexed _token, address indexed _controller, uint256 _amount);