fix: deposit everything to AffiliateToken (yearn#261)

sambacha · Sep 7, 2021 · e0f1235 · e0f1235
1 parent a0a7729
commit e0f1235
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 4 deletions.
diff --git a/contracts/BaseWrapper.sol b/contracts/BaseWrapper.sol
@@ -110,7 +110,11 @@ abstract contract BaseWrapper {
         VaultAPI _bestVault = bestVault();
 
         if (pullFunds) {
-            token.safeTransferFrom(depositor, address(this), amount);
+            if (amount != DEPOSIT_EVERYTHING) {
+                token.safeTransferFrom(depositor, address(this), amount);
+            } else {
+                token.safeTransferFrom(depositor, address(this), token.balanceOf(depositor));
+            }
         }
 
         if (token.allowance(address(this), address(_bestVault)) < amount) {

diff --git a/contracts/test/AffiliateToken.sol b/contracts/test/AffiliateToken.sol
@@ -64,11 +64,12 @@ contract AffiliateToken is ERC20, BaseWrapper {
     }
 
     function pricePerShare() external view returns (uint256) {
-        return 10**uint256(decimals());
+        return totalVaultBalance(address(this)).mul(10**uint256(decimals())).div(totalSupply());
     }
 
     function _sharesForValue(uint256 amount) internal view returns (uint256) {
-        uint256 totalWrapperAssets = totalVaultBalance(address(this));
+        // total wrapper assets before deposit (assumes deposit already occured)
+        uint256 totalWrapperAssets = totalVaultBalance(address(this)).sub(amount);
 
         if (totalWrapperAssets > 0) {
             return totalSupply().mul(amount).div(totalWrapperAssets);
@@ -82,8 +83,8 @@ contract AffiliateToken is ERC20, BaseWrapper {
     }
 
     function deposit(uint256 amount) public returns (uint256 deposited) {
-        uint256 shares = _sharesForValue(amount); // NOTE: Must be calculated before deposit is handled
         deposited = _deposit(msg.sender, address(this), amount, true); // `true` = pull from `msg.sender`
+        uint256 shares = _sharesForValue(deposited);  // NOTE: Must be calculated after deposit is handled
         _mint(msg.sender, shares);
     }
 

diff --git a/tests/functional/wrappers/test_affliate.py b/tests/functional/wrappers/test_affliate.py
@@ -74,6 +74,20 @@ def test_deposit(token, registry, vault, affiliate_token, gov, rando):
     assert vault.balanceOf(rando) == 0
 
 
+def test_deposit_max_uint256(token, registry, vault, affiliate_token, gov, rando):
+    registry.newRelease(vault, {"from": gov})
+    registry.endorseVault(vault, {"from": gov})
+    token.transfer(rando, 10000, {"from": gov})
+    assert affiliate_token.balanceOf(rando) == vault.balanceOf(rando) == 0
+
+    # NOTE: Must approve affiliate_token to deposit
+    token.approve(affiliate_token, 2 ** 256 - 1, {"from": rando})
+
+    affiliate_token.deposit({"from": rando})
+    assert affiliate_token.balanceOf(rando) == 10000
+    assert vault.balanceOf(rando) == 0
+
+
 def test_migrate(token, registry, create_vault, affiliate_token, gov, rando, affiliate):
     vault1 = create_vault(version="1.0.0", token=token)
     registry.newRelease(vault1, {"from": gov})