Override _deposit and _withdraw

brritoxyz · Nov 13, 2023 · 85c1ad2 · 85c1ad2
1 parent 1f0a5d2
commit 85c1ad2
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 61 deletions.
diff --git a/src/BrrETH.sol b/src/BrrETH.sol
@@ -25,6 +25,8 @@ contract BrrETH is Ownable, ERC4626 {
     event AddRebaseToken(address);
     event RemoveRebaseToken(address);
 
+    error InvalidAssets();
+
     constructor(address initialOwner) {
         _initializeOwner(initialOwner);
 
@@ -47,26 +49,37 @@ contract BrrETH is Ownable, ERC4626 {
         return _COMET;
     }
 
-    function deposit(
+    function _deposit(
+        address by,
+        address to,
         uint256 assets,
-        address to
-    ) public override returns (uint256 shares) {
-        if (assets > _COMET.balanceOf(msg.sender)) revert InsufficientBalance();
+        uint256 shares
+    ) internal override {
+        if (assets == type(uint256).max) revert InvalidAssets();
+
+        _COMET.safeTransferFrom(by, address(this), assets);
+        _mint(to, shares);
 
-        shares = previewDeposit(assets);
+        emit Deposit(by, to, assets, shares);
 
-        _deposit(msg.sender, to, assets, shares);
+        _afterDeposit(assets, shares);
     }
 
-    function mint(
-        uint256 shares,
-        address to
-    ) public override returns (uint256 assets) {
-        assets = previewMint(shares);
+    function _withdraw(
+        address by,
+        address to,
+        address owner,
+        uint256 assets,
+        uint256 shares
+    ) internal override {
+        if (assets == type(uint256).max) revert InvalidAssets();
+        if (by != owner) _spendAllowance(owner, by, shares);
 
-        if (assets > _COMET.balanceOf(msg.sender)) revert InsufficientBalance();
+        _beforeWithdraw(assets, shares);
+        _burn(owner, shares);
+        _COMET.safeTransfer(to, assets);
 
-        _deposit(msg.sender, to, assets, shares);
+        emit Withdraw(by, to, owner, assets, shares);
     }
 
     function rebaseTokens() external view returns (address[] memory) {

diff --git a/test/BrrETH.t.sol b/test/BrrETH.t.sol
@@ -167,52 +167,4 @@ contract BrrETHTest is Helper, Test {
     function testAsset() external {
         assertEq(_COMET, vault.asset());
     }
-
-    /*//////////////////////////////////////////////////////////////
-                             deposit
-    //////////////////////////////////////////////////////////////*/
-
-    function testCannotDepositInsufficientBalance() external {
-        _getCWETH(1 ether);
-
-        uint256 assets = type(uint256).max;
-        address to = address(this);
-
-        assertTrue(assets > _COMET.balanceOf(address(this)));
-
-        vm.expectRevert(ERC20.InsufficientBalance.selector);
-
-        vault.deposit(assets, to);
-    }
-
-    function testCannotDepositInsufficientBalanceFuzz(uint80 balance) external {
-        vm.assume(balance != 0);
-
-        _getCWETH(balance);
-
-        uint256 assets = uint256(balance) + 1;
-        address to = address(this);
-
-        vm.expectRevert(ERC20.InsufficientBalance.selector);
-
-        vault.deposit(assets, to);
-    }
-
-    /*//////////////////////////////////////////////////////////////
-                             mint
-    //////////////////////////////////////////////////////////////*/
-
-    function testCannotMintInsufficientBalance() external {
-        _getCWETH(1 ether);
-
-        uint256 shares = vault.previewDeposit(type(uint256).max);
-        uint256 assets = vault.previewMint(shares);
-        address to = address(this);
-
-        assertTrue(assets > _COMET.balanceOf(address(this)));
-
-        vm.expectRevert(ERC20.InsufficientBalance.selector);
-
-        vault.mint(shares, to);
-    }
 }