fraxbp module modifications

src/modules/fraxBPInitModule/FraxBPInitModule.sol

@@ -10,6 +10,7 @@ import "@protocol/interfaces/IPHO.sol";
 import "@protocol/interfaces/IModuleManager.sol";
 import "@external/curve/ICurvePool.sol";
 import "@oracle/IPriceOracle.sol";
+import "./FraxBPInitModuleAMO.sol";
 
 /// @title FraxBP Init module
 /// @author Ekonomia: https://github.com/ekonomia-tech
@@ -44,6 +45,9 @@ contract FraxBPInitModule is Ownable, ReentrancyGuard {
     uint256 public saleEndDate; // when sale ends
     uint256 public redemptionStartDate; // when redemptions are available
     IPHO public pho;
+    address public fraxBPInitModuleAMO;
+    address public stakingToken = 0x5f98805A4E8be255a32880FDeC7F6728C6568bA0; // LUSD
+    address rewardToken = 0xD533a949740bb3306d119CC777fa900bA034cd52; // CRV
 
     mapping(address => uint256) public metapoolBalance;
 
@@ -59,11 +63,12 @@ contract FraxBPInitModule is Ownable, ReentrancyGuard {
         address _pho,
         address _priceOracle,
         uint256 _saleEndDate,
-        uint256 _redemptionStartDate
+        uint256 _redemptionStartDate,
+        address _gauge
     ) {
         if (
             _moduleManager == address(0) || _fraxBPPHOMetapool == address(0) || _pho == address(0)
-                || _priceOracle == address(0)
+                || _priceOracle == address(0) || _gauge == address(0)
         ) {
             revert ZeroAddressDetected();
         }
@@ -79,6 +84,21 @@ contract FraxBPInitModule is Ownable, ReentrancyGuard {
 
         usdc.approve(address(fraxBPPool), type(uint256).max);
         frax.approve(address(fraxBPPool), type(uint256).max);
+
+        FraxBPInitModuleAMO fraxBPModuleAMO = new FraxBPInitModuleAMO(
+            "FRAXBPPHO Module AMO",
+            "FBPPHO-AMO",
+            stakingToken,
+            rewardToken,
+            msg.sender,
+            address(this),
+            address(fraxBPPHOMetapool),
+            _gauge
+        );
+
+        fraxBPInitModuleAMO = address(fraxBPModuleAMO);
+
+        IERC20(_fraxBPPHOMetapool).approve(fraxBPInitModuleAMO, type(uint256).max);
     }
 
     /// @notice Helper for user depositing both FRAX and USDC
@@ -112,41 +132,45 @@ contract FraxBPInitModule is Ownable, ReentrancyGuard {
 
         // call _depositFor() for user based on FraxBP LP received
         uint256 fraxBPLpAmount = fraxBPLpBalanceAfter - fraxBPLpBalanceBefore;
-        _depositFor(msg.sender, fraxBPLpAmount);
+        uint256 fraxBPPHOLpBalanceIssued = _depositFor(msg.sender, fraxBPLpAmount);
+        IModuleAMO(fraxBPInitModuleAMO).stakeFor(msg.sender, fraxBPPHOLpBalanceIssued);
     }
 
     /// @notice Helper function for deposits in FraxBP LP token for user
     /// @param depositor Depositor
     /// @param amount Amount in FraxBP LP
-    function _depositFor(address depositor, uint256 amount) private {
+    function _depositFor(address depositor, uint256 amount) private returns (uint256) {
         uint256 usdPerFraxBP = getUSDPerFraxBP();
         uint256 phoAmount = (usdPerFraxBP * amount) / 10 ** 18;
 
         moduleManager.mintPHO(address(this), phoAmount);
-        uint256 fraxBPPHOLpBalance = _addFraxBPPHOLiquidity(amount, phoAmount);
+        uint256 fraxBPPHOLpBalanceIssued = _addFraxBPPHOLiquidity(amount, phoAmount);
 
-        metapoolBalance[depositor] += fraxBPPHOLpBalance;
+        metapoolBalance[depositor] += fraxBPPHOLpBalanceIssued;
         emit Deposited(depositor, amount, phoAmount);
+        return fraxBPPHOLpBalanceIssued;
     }
 
     /// @notice Places deposits in FraxBP LP token for user
     /// @param depositor Depositor
     /// @param amount Amount in FraxBP LP
-    function depositFor(address depositor, uint256 amount) public {
+    function depositFor(address depositor, uint256 amount) public returns (uint256) {
         if (block.timestamp > saleEndDate) {
             revert CannotDepositAfterSaleEnded();
         }
         if (amount == 0) {
             revert CannotDepositZero();
         }
         fraxBPLp.safeTransferFrom(depositor, address(this), amount);
-        _depositFor(depositor, amount);
+        uint256 fraxBPPHOLPNetBalance = _depositFor(depositor, amount);
+        return fraxBPPHOLPNetBalance;
     }
 
     /// @notice Accept deposits in FraxBP LP token
     /// @param amount Amount in FraxBP LP
     function deposit(uint256 amount) external nonReentrant {
-        depositFor(msg.sender, amount);
+        uint256 fraxBPPHOLpBalanceIssued = depositFor(msg.sender, amount);
+        IModuleAMO(fraxBPInitModuleAMO).stakeFor(msg.sender, fraxBPPHOLpBalanceIssued);
     }
 
     /// @notice Adds FraxBP LP and PHO to FraxBP/PHO pool
@@ -185,7 +209,9 @@ contract FraxBPInitModule is Ownable, ReentrancyGuard {
         }
 
         delete metapoolBalance[msg.sender];
-        fraxBPPHOMetapool.transfer(msg.sender, redeemAmount / 2);
+        // Note: Always a full withdrawal
+        IModuleAMO(fraxBPInitModuleAMO).withdrawAllFor(msg.sender);
+        //fraxBPPHOMetapool.transfer(msg.sender, redeemAmount / 2);
         emit Redeemed(msg.sender, redeemAmount);
     }
 

src/modules/fraxBPInitModule/FraxBPInitModuleAMO.sol

@@ -0,0 +1,230 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+pragma solidity ^0.8.13;
+
+import "@openzeppelin/contracts/utils/math/SafeMath.sol";
+import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "@openzeppelin/contracts/utils/Address.sol";
+import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import "../interfaces/IModuleAMO.sol";
+import "./interfaces/IGauge.sol";
+import "forge-std/console2.sol";
+
+/// @title FraxBPInitModuleAMO
+/// @notice FraxBPInit Module AMO
+/// @author Ekonomia: https://github.com/Ekonomia
+contract FraxBPInitModuleAMO is IModuleAMO, ERC20 {
+    using SafeMath for uint256;
+    using SafeERC20 for IERC20;
+
+    // Errors
+    error ZeroAddressDetected();
+    error CannotStakeZero();
+    error CannotWithdrawMoreThanDeposited();
+
+    /// State vars
+    address public rewardToken;
+    address public stakingToken;
+    address public operator;
+    address public module;
+    uint256 private _totalDeposits;
+    uint256 private _totalShares;
+    uint256 private _totalRewards; // rewards in CRV
+
+    mapping(address => uint256) public depositedAmount; // MPL deposited
+    mapping(address => uint256) public stakedAmount; // MPL staked
+    mapping(address => uint256) public claimedRewards; // rewards claimed
+    mapping(address => uint256) private _shares;
+
+    // Needed for interactions w/ external contracts
+    address public depositToken;
+    IERC20 public crv = IERC20(0xD533a949740bb3306d119CC777fa900bA034cd52);
+    IGauge public curveGauge; // TODO: gauge address
+
+    // Events
+    event RewardAdded(uint256 reward);
+    event Staked(address indexed user, uint256 amount, uint256 shares);
+    event Withdrawn(address indexed user, uint256 amount, uint256 shares);
+    event RewardPaid(address indexed user, uint256 reward);
+    event FraxBPInitRewardsReceived(uint256 totalRewards);
+
+    modifier onlyOperator() {
+        require(msg.sender == address(operator), "Only Operator");
+        _;
+    }
+
+    modifier onlyModule() {
+        require(msg.sender == address(module), "Only Module");
+        _;
+    }
+
+    /// Constructor
+    constructor(
+        string memory _name,
+        string memory _symbol,
+        address _stakingToken,
+        address _rewardToken,
+        address _operator,
+        address _module,
+        address _depositToken,
+        address _curveGauge
+    ) ERC20(_name, _symbol) {
+        if (
+            _stakingToken == address(0) || _rewardToken == address(0) || _operator == address(0)
+                || _module == address(0) || _depositToken == address(0) || _curveGauge == address(0)
+        ) {
+            revert ZeroAddressDetected();
+        }
+        stakingToken = _stakingToken;
+        rewardToken = _rewardToken;
+        operator = _operator;
+        module = _module;
+        depositToken = _depositToken;
+        curveGauge = IGauge(_curveGauge);
+
+        IERC20(depositToken).approve(address(curveGauge), type(uint256).max);
+    }
+
+    /// @notice Get total shares
+    function totalShares() public view returns (uint256) {
+        return _totalShares;
+    }
+
+    /// @notice Get shares of account
+    /// @param account Account
+    function sharesOf(address account) public view returns (uint256) {
+        return _shares[account];
+    }
+
+    /// @notice Get earned amount by account (total available - claimed)
+    /// @param account Account
+    function earned(address account) public view returns (uint256) {
+        uint256 ts = totalSupply();
+        if (ts == 0) {
+            return 0;
+        }
+        uint256 earnedRewards = (balanceOf(account) * _totalRewards) / ts - claimedRewards[account];
+        return earnedRewards;
+    }
+
+    /// @notice Convert a deposit/withdraw amount into shares
+    function _toShares(uint256 amount) private view returns (uint256) {
+        if (_totalShares == 0) {
+            return amount;
+        }
+        return (amount * _totalShares) / _totalDeposits;
+    }
+
+    /// @notice Tracks shares for deposits
+    function _trackDepositShares(address account, uint256 amount) private returns (uint256) {
+        uint256 shares = _toShares(amount);
+        _shares[account] += shares;
+        _totalShares += shares;
+        _mint(account, shares);
+        return shares;
+    }
+
+    /// @notice Tracks shares for withdrawals
+    function _trackWithdrawShares(address account) private returns (uint256) {
+        uint256 shares = _shares[account];
+        _shares[account] = 0;
+        _totalShares -= shares;
+        _burn(account, shares);
+        return shares;
+    }
+
+    /// @notice Stake for
+    /// @param account For
+    /// @param amount Amount
+    function stakeFor(address account, uint256 amount) public onlyModule returns (bool) {
+        if (amount == 0) {
+            revert CannotStakeZero();
+        }
+
+        // Get depositToken from module
+        IERC20(depositToken).safeTransferFrom(module, address(this), amount);
+
+        console2.log("In FraxBPInitModuleAMO.. about to approve()..");
+
+        //IERC20(depositToken).approve(address(curveGauge), type(uint256).max);
+
+        console2.log("In FraxBPInitModuleAMO.. approved, now about to call curveGauge.deposit()..");
+
+        // Deposit to gauge
+        //curveGauge.deposit(amount);
+        curveGauge.deposit(amount, address(this));
+
+        depositedAmount[account] += amount;
+        stakedAmount[account] += amount;
+        _totalDeposits += amount;
+
+        uint256 shares = _trackDepositShares(account, amount);
+        emit Staked(account, amount, shares);
+        return true;
+    }
+
+    /// @notice Withdraw
+    /// @param account Account
+    /// @param amount amount
+    function withdrawFor(address account, uint256 amount) public onlyModule returns (bool) {
+        uint256 depositAmount = depositedAmount[account];
+        uint256 stakedPoolTokenAmount = stakedAmount[account];
+        if (amount > depositAmount) {
+            revert CannotWithdrawMoreThanDeposited();
+        }
+        depositedAmount[account] -= depositAmount;
+        stakedAmount[account] -= stakedPoolTokenAmount;
+
+        // Withdraw from gauge
+        curveGauge.withdraw(amount);
+
+        uint256 shares = _trackWithdrawShares(account);
+        _totalDeposits -= depositAmount;
+
+        console2.log(
+            "this is in withdrawFor, this is balance of depositToken: ",
+            IERC20(depositToken).balanceOf(address(this))
+        );
+
+        // Transfer depositToken to caller
+        IERC20(depositToken).transfer(account, depositAmount / 2);
+        emit Withdrawn(account, amount, shares);
+        return true;
+    }
+
+    /// @notice Withdraw all for
+    /// @param account Account
+    function withdrawAllFor(address account) external returns (bool) {
+        return withdrawFor(account, depositedAmount[account]);
+    }
+
+    /// @notice Gets reward
+    function getRewardFraxBPInit() external onlyOperator returns (uint256) {
+        uint256 crvBalanceBefore = crv.balanceOf(address(this));
+        // Get rewards
+        //curveGauge.claim_rewards();
+        uint256 crvBalanceAfter = crv.balanceOf(address(this));
+        _totalRewards = crvBalanceAfter - crvBalanceBefore;
+        emit FraxBPInitRewardsReceived(_totalRewards);
+        return _totalRewards;
+    }
+
+    /// @notice Get reward
+    /// @param account Account
+    function getReward(address account) public returns (bool) {
+        uint256 reward = earned(account);
+        if (reward > 0) {
+            claimedRewards[account] += reward;
+            IERC20(rewardToken).safeTransfer(account, reward);
+            emit RewardPaid(account, reward);
+        }
+        return true;
+    }
+
+    /// @notice Get reward for msg.sender
+    function getReward() external returns (bool) {
+        getReward(msg.sender);
+        return true;
+    }
+}