refactor: use solmate's SafeTransferLib, use Ownable (#4)

.github/workflows/ci.yml

@@ -69,4 +69,4 @@ jobs:
 
       # Too slow to run regularly
       #- name: Run audit
-      #  run: poetry run slither --solc-remaps @openzeppelin=lib/openzeppelin-contracts --solc-args optimize src/
+      #  run: poetry run slither --solc-remaps @openzeppelin=lib/openzeppelin-contracts --solc-args optimize src/

src/Position.sol

@@ -4,6 +4,7 @@ pragma solidity ^0.8.21;
 // Local Imports
 import { DebtService } from "src/services/DebtService.sol";
 import { SwapService } from "src/services/SwapService.sol";
+import { SafeTransferLib, ERC20 } from "solmate/utils/SafeTransferLib.sol";
 import { IERC20 } from "src/interfaces/token/IERC20.sol";
 
 /// @title Position
@@ -32,7 +33,7 @@ contract Position is DebtService, SwapService {
      */
     function short(uint256 _cAmt, uint256 _ltv, uint256 _swapAmtOutMin, uint24 _poolFee) public payable onlyOwner {
         // 1. Transfer collateral to this contract
-        IERC20(C_TOKEN).transferFrom(msg.sender, address(this), _cAmt);
+        SafeTransferLib.safeTransferFrom(ERC20(C_TOKEN), msg.sender, address(this), _cAmt);
 
         // 2. Borrow debt token
         uint256 dAmt = _borrow(_cAmt, _ltv);
@@ -77,7 +78,7 @@ contract Position is DebtService, SwapService {
         uint256 gains = bTokenBalance - bAmtIn;
 
         if (gains > 0) {
-            IERC20(B_TOKEN).transfer(OWNER, gains);
+            SafeTransferLib.safeTransfer(ERC20(B_TOKEN), OWNER, gains);
         }
 
         // 5. Emit event

src/PositionAdmin.sol

@@ -3,6 +3,7 @@ pragma solidity ^0.8.21;
 
 // Local Imports
 import { IERC20 } from "src/interfaces/token/IERC20.sol";
+import { SafeTransferLib, ERC20 } from "solmate/utils/SafeTransferLib.sol";
 
 /// @title Position Admin
 /// @author Chain Rule, LLC
@@ -32,7 +33,7 @@ contract PositionAdmin {
     function extractERC20(address _token) public payable onlyOwner {
         uint256 balance = IERC20(_token).balanceOf(address(this));
 
-        IERC20(_token).transfer(msg.sender, balance);
+        SafeTransferLib.safeTransfer(ERC20(_token), msg.sender, balance);
     }
 
     /**

src/PositionFactory.sol

@@ -3,31 +3,28 @@ pragma solidity ^0.8.21;
 
 // Local
 import { Position } from "src/Position.sol";
+import { Ownable } from "src/dependencies/access/Ownable.sol";
+import { SafeTransferLib, ERC20 } from "solmate/utils/SafeTransferLib.sol";
 import { IERC20 } from "src/interfaces/token/IERC20.sol";
 
 /// @title Position Factory
 /// @author Chain Rule, LLC
 /// @notice Creates user position contracts and stores their addresses
-contract PositionFactory {
+contract PositionFactory is Ownable {
     // Constants: no SLOAD to save gas
     address private constant CONTRACT_DEPLOYER = 0x0a5B347509621337cDDf44CBCf6B6E7C9C908CD2;
 
-    // Immutables: no SLOAD to save gas
-    address public immutable OWNER;
-
     // Factory Storage
-    //      Mapping from owner to cToken to dToken to bToken to position
+    /// @dev Mapping from owner to cToken to dToken to bToken to position
     mapping(address => mapping(address => mapping(address => mapping(address => address)))) public positions;
-    //      Mapping from owner to list of positions
     mapping(address => address[]) public positionsLookup;
 
     // Errors
     error Unauthorized();
     error PositionExists();
 
-    constructor(address _owner) {
+    constructor(address _owner) Ownable(_owner) {
         if (msg.sender != CONTRACT_DEPLOYER) revert Unauthorized();
-        OWNER = _owner;
     }
 
     /**
@@ -71,7 +68,7 @@ contract PositionFactory {
     function extractERC20(address _token) public payable onlyOwner {
         uint256 balance = IERC20(_token).balanceOf(address(this));
 
-        IERC20(_token).transfer(msg.sender, balance);
+        SafeTransferLib.safeTransfer(ERC20(_token), msg.sender, balance);
     }
 
     /**
@@ -83,10 +80,4 @@ contract PositionFactory {
      * @notice Executes when native is sent to this contract with a plain transaction.
      */
     receive() external payable { } // solhint-disable-line no-empty-blocks
-
-    /// @notice Check if the caller is the owner of the PositionFactory contract
-    modifier onlyOwner() {
-        if (msg.sender != OWNER) revert Unauthorized();
-        _;
-    }
 }

src/dependencies/access/Context.sol

@@ -0,0 +1,28 @@
+// SPDX-License-Identifier: MIT
+// OpenZeppelin Contracts (last updated v5.0.1) (utils/Context.sol)
+
+pragma solidity ^0.8.21;
+
+/**
+ * @dev Provides information about the current execution context, including the
+ * sender of the transaction and its data. While these are generally available
+ * via msg.sender and msg.data, they should not be accessed in such a direct
+ * manner, since when dealing with meta-transactions the account sending and
+ * paying for execution may not be the actual sender (as far as an application
+ * is concerned).
+ *
+ * This contract is only required for intermediate, library-like contracts.
+ */
+abstract contract Context {
+    function _msgSender() internal view virtual returns (address) {
+        return msg.sender;
+    }
+
+    function _msgData() internal view virtual returns (bytes calldata) {
+        return msg.data;
+    }
+
+    function _contextSuffixLength() internal view virtual returns (uint256) {
+        return 0;
+    }
+}

src/dependencies/access/Ownable.sol

@@ -0,0 +1,100 @@
+// SPDX-License-Identifier: MIT
+// OpenZeppelin Contracts (last updated v5.0.0) (access/Ownable.sol)
+
+pragma solidity ^0.8.21;
+
+import { Context } from "src/dependencies/access/Context.sol";
+
+/**
+ * @dev Contract module which provides a basic access control mechanism, where
+ * there is an account (an owner) that can be granted exclusive access to
+ * specific functions.
+ *
+ * The initial owner is set to the address provided by the deployer. This can
+ * later be changed with {transferOwnership}.
+ *
+ * This module is used through inheritance. It will make available the modifier
+ * `onlyOwner`, which can be applied to your functions to restrict their use to
+ * the owner.
+ */
+abstract contract Ownable is Context {
+    address private _owner;
+
+    /**
+     * @dev The caller account is not authorized to perform an operation.
+     */
+    error OwnableUnauthorizedAccount(address account);
+
+    /**
+     * @dev The owner is not a valid owner account. (eg. `address(0)`)
+     */
+    error OwnableInvalidOwner(address owner);
+
+    event OwnershipTransferred(address indexed previousOwner, address indexed newOwner);
+
+    /**
+     * @dev Initializes the contract setting the address provided by the deployer as the initial owner.
+     */
+    constructor(address initialOwner) {
+        if (initialOwner == address(0)) {
+            revert OwnableInvalidOwner(address(0));
+        }
+        _transferOwnership(initialOwner);
+    }
+
+    /**
+     * @dev Throws if called by any account other than the owner.
+     */
+    modifier onlyOwner() {
+        _checkOwner();
+        _;
+    }
+
+    /**
+     * @dev Returns the address of the current owner.
+     */
+    function owner() public view virtual returns (address) {
+        return _owner;
+    }
+
+    /**
+     * @dev Throws if the sender is not the owner.
+     */
+    function _checkOwner() internal view virtual {
+        if (owner() != _msgSender()) {
+            revert OwnableUnauthorizedAccount(_msgSender());
+        }
+    }
+
+    /**
+     * @dev Leaves the contract without owner. It will not be possible to call
+     * `onlyOwner` functions. Can only be called by the current owner.
+     *
+     * NOTE: Renouncing ownership will leave the contract without an owner,
+     * thereby disabling any functionality that is only available to the owner.
+     */
+    function renounceOwnership() public virtual onlyOwner {
+        _transferOwnership(address(0));
+    }
+
+    /**
+     * @dev Transfers ownership of the contract to a new account (`newOwner`).
+     * Can only be called by the current owner.
+     */
+    function transferOwnership(address newOwner) public virtual onlyOwner {
+        if (newOwner == address(0)) {
+            revert OwnableInvalidOwner(address(0));
+        }
+        _transferOwnership(newOwner);
+    }
+
+    /**
+     * @dev Transfers ownership of the contract to a new account (`newOwner`).
+     * Internal function without access restriction.
+     */
+    function _transferOwnership(address newOwner) internal virtual {
+        address oldOwner = _owner;
+        _owner = newOwner;
+        emit OwnershipTransferred(oldOwner, newOwner);
+    }
+}

src/services/DebtService.sol

@@ -3,6 +3,7 @@ pragma solidity ^0.8.21;
 
 // Local Imports
 import { PositionAdmin } from "src/PositionAdmin.sol";
+import { SafeTransferLib, ERC20 } from "solmate/utils/SafeTransferLib.sol";
 import { IPool } from "src/interfaces/aave/IPool.sol";
 import { IERC20 } from "src/interfaces/token/IERC20.sol";
 import { IAaveOracle } from "src/interfaces/aave/IAaveOracle.sol";
@@ -47,7 +48,7 @@ contract DebtService is PositionAdmin {
      */
     function _borrow(uint256 _cAmt, uint256 _ltv) internal returns (uint256 dAmt) {
         // 1. Supply collateral to Aave
-        IERC20(C_TOKEN).approve(AAVE_POOL, _cAmt);
+        SafeTransferLib.safeApprove(ERC20(C_TOKEN), AAVE_POOL, _cAmt);
         IPool(AAVE_POOL).supply(C_TOKEN, _cAmt, address(this), 0);
 
         // 2. Get asset prices USD
@@ -66,7 +67,7 @@ contract DebtService is PositionAdmin {
      * @param _dAmt The amount of debt token to repay to Aave.
      */
     function _repay(uint256 _dAmt) internal {
-        IERC20(D_TOKEN).approve(AAVE_POOL, _dAmt);
+        SafeTransferLib.safeApprove(ERC20(D_TOKEN), AAVE_POOL, _dAmt);
         IPool(AAVE_POOL).repay(D_TOKEN, _dAmt, 2, address(this));
     }
 
@@ -114,10 +115,10 @@ contract DebtService is PositionAdmin {
      */
     function addCollateral(uint256 _cAmt) public payable onlyOwner {
         // 1. Transfer collateral from owner to this contract
-        IERC20(C_TOKEN).transferFrom(msg.sender, address(this), _cAmt);
+        SafeTransferLib.safeTransferFrom(ERC20(C_TOKEN), msg.sender, address(this), _cAmt);
 
         // 2. Approve Aave to spend _cAmt of this contract's C_TOKEN
-        IERC20(C_TOKEN).approve(AAVE_POOL, _cAmt);
+        SafeTransferLib.safeApprove(ERC20(C_TOKEN), AAVE_POOL, _cAmt);
 
         // 3. Supply collateral to Aave
         IPool(AAVE_POOL).supply(C_TOKEN, _cAmt, address(this), 0);
@@ -130,7 +131,7 @@ contract DebtService is PositionAdmin {
      * @param _withdrawBuffer The amount of collateral left as safety buffer for tx to go through (default = 100_000, units: 8 decimals).
      */
     function repayAfterClose(uint256 _dAmt, uint256 _withdrawBuffer) public payable onlyOwner {
-        IERC20(D_TOKEN).transferFrom(msg.sender, address(this), _dAmt);
+        SafeTransferLib.safeTransferFrom(ERC20(D_TOKEN), msg.sender, address(this), _dAmt);
 
         _repay(_dAmt);
 

test/PositionFactory.t.sol

@@ -22,6 +22,9 @@ contract PositionFactoryTest is Test, TokenUtils {
     uint256 public mainnetFork;
     address public positionOwner = address(this);
 
+    // Errors
+    error OwnableUnauthorizedAccount(address account);
+
     function setUp() public {
         // Setup: use mainnet fork
         mainnetFork = vm.createFork(vm.envString("RPC_URL"));
@@ -207,7 +210,7 @@ contract PositionFactoryTest is Test, TokenUtils {
 
         // Act: attempt to extract native
         vm.prank(_sender);
-        vm.expectRevert(PositionFactory.Unauthorized.selector);
+        vm.expectRevert(abi.encodeWithSelector(OwnableUnauthorizedAccount.selector, _sender));
         positionFactory.extractNative();
     }
 
@@ -262,7 +265,7 @@ contract PositionFactoryTest is Test, TokenUtils {
 
             // Act
             vm.prank(_sender);
-            vm.expectRevert(PositionFactory.Unauthorized.selector);
+            vm.expectRevert(abi.encodeWithSelector(OwnableUnauthorizedAccount.selector, _sender));
             positionFactory.extractERC20(supportedAssets[i]);
         }
     }