From 259aab24751159bf1ff00646fd8c470e84e81561 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alberto=20Cuesta=20Ca=C3=B1ada?=
 <38806121+albertocuestacanada@users.noreply.github.com>
Date: Thu, 22 Oct 2020 05:34:12 +0100
Subject: [PATCH] fix: get as much Dai as possible after removing liquidity
 (#375)

* fix: get as much Dai as possible after removing liquidity

* fix: remove unused parameter
---
 contracts/peripheral/YieldProxy.sol   | 20 +++++++++-----------
 test/peripheral/551_liquidityProxy.ts |  4 ++--
 2 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/contracts/peripheral/YieldProxy.sol b/contracts/peripheral/YieldProxy.sol
index 54ea6be9..1a9d30d6 100644
--- a/contracts/peripheral/YieldProxy.sol
+++ b/contracts/peripheral/YieldProxy.sol
@@ -243,7 +243,7 @@ contract YieldProxy is DecimalMath {
                 "YieldProxy: minimumFYDaiPrice not reached"
             );
         }
-        withdrawAssets(fyDai);
+        withdrawAssets();
     }
 
     /// @dev Burns tokens and repays debt with proceedings. Sells any excess fyDai for Dai, then returns all Dai, and if there is no debt in the Controller, all posted Chai.
@@ -273,7 +273,7 @@ contract YieldProxy is DecimalMath {
                 "YieldProxy: minimumFYDaiPrice not reached"
             );
         }
-        withdrawAssets(fyDai);
+        withdrawAssets();
     }
 
     /// @dev Burns tokens and repays fyDai debt after Maturity. 
@@ -293,18 +293,16 @@ contract YieldProxy is DecimalMath {
         if (daiObtained > 0 && controller.debtFYDai(CHAI, maturity, msg.sender) > 0) {
             controller.repayDai(CHAI, maturity, address(this), msg.sender, daiObtained);
         }
-        withdrawAssets(fyDai);
+        withdrawAssets();
     }
 
     /// @dev Return to caller all posted chai if there is no debt, converted to dai, plus any dai remaining in the contract.
-    function withdrawAssets(IFYDai fyDai) internal {
-        if (controller.debtFYDai(CHAI, fyDai.maturity(), msg.sender) == 0) {
-            uint256 posted = controller.posted(CHAI, msg.sender);
-            uint256 locked = controller.locked(CHAI, msg.sender);
-            require (posted >= locked, "YieldProxy: Undercollateralized");
-            controller.withdraw(CHAI, msg.sender, address(this), posted - locked);
-            chai.exit(address(this), chai.balanceOf(address(this)));
-        }
+    function withdrawAssets() internal {
+        uint256 posted = controller.posted(CHAI, msg.sender);
+        uint256 locked = controller.locked(CHAI, msg.sender);
+        require (posted >= locked, "YieldProxy: Undercollateralized");
+        controller.withdraw(CHAI, msg.sender, address(this), posted - locked);
+        chai.exit(address(this), chai.balanceOf(address(this)));
         require(dai.transfer(msg.sender, dai.balanceOf(address(this))), "YieldProxy: Dai Transfer Failed");
     }
 
diff --git a/test/peripheral/551_liquidityProxy.ts b/test/peripheral/551_liquidityProxy.ts
index 0ba50e06..915d6af6 100644
--- a/test/peripheral/551_liquidityProxy.ts
+++ b/test/peripheral/551_liquidityProxy.ts
@@ -264,8 +264,8 @@ contract('YieldProxy - LiquidityProxy', async (accounts) => {
         expect(await pool0.balanceOf(user2)).to.be.bignumber.eq(ZERO)
         // Has less fyDai debt
         expect(await controller.debtFYDai(CHAI, maturity0, user2)).to.be.bignumber.lt(debt)
-        // Doesn't have dai
-        expect(await dai.balanceOf(user2)).to.be.bignumber.eq(ZERO)
+        // Got some dai
+        expect(await dai.balanceOf(user2)).to.be.bignumber.gt(ZERO)
         // Has the same fyDai
         expect(await fyDai0.balanceOf(user2)).to.be.bignumber.eq(toBorrow)
         // Proxy doesn't keep dai (beyond rounding)