From 09b011abe4c99d1036044e47ca72387ce8ef48f4 Mon Sep 17 00:00:00 2001
From: Thrilok Kumar <thrilok2000@gmail.com>
Date: Tue, 15 Sep 2020 04:06:33 +0530
Subject: [PATCH] Code refactoring and added SafeERC20 lib

---
 contracts/flashloan/dydx.sol | 40 +++++++++++++++++++++++++++---------
 1 file changed, 30 insertions(+), 10 deletions(-)

diff --git a/contracts/flashloan/dydx.sol b/contracts/flashloan/dydx.sol
index ddcc234..6d61df1 100644
--- a/contracts/flashloan/dydx.sol
+++ b/contracts/flashloan/dydx.sol
@@ -5,17 +5,14 @@ import "@studydefi/money-legos/dydx/contracts/DydxFlashloanBase.sol";
 import "@studydefi/money-legos/dydx/contracts/ICallee.sol";
 
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "@openzeppelin/contracts/token/ERC20/SafeERC20.sol";
+import { DSMath } from "../common/math.sol";
 
 interface DSAInterface {
     function cast(address[] calldata _targets, bytes[] calldata _datas, address _origin) external payable;
 }
 
-contract DydxFlashloaner is ICallee, DydxFlashloanBase {
-
-    address public constant soloAddr = 0x1E0447b19BB6EcFdAe1e4AE1694b0C3659614e4e;
-    address public constant wethAddr = 0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2;
-    address public constant ethAddr = 0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE;
-
+contract Helper {
     struct CastData {
         address dsa;
         address token;
@@ -24,12 +21,35 @@ contract DydxFlashloaner is ICallee, DydxFlashloanBase {
         bytes[] data;
     }
 
-    // check9898 - add block re-entrance
+    function encodeDsaAddr(address dsa, bytes memory data) internal view returns (bytes memory _data) {
+        CastData memory cd;
+        (cd.token, cd.amount, cd.targets, cd.data) = abi.decode(data, (address, uint256, address[], bytes[]));
+        _data = abi.encode(dsa, cd.token, cd.amount, cd.targets, cd.data);
+    }
+}
+
+contract DydxFlashloaner is ICallee, DydxFlashloanBase, DSMath, Helper {
+    using SafeERC20 for IERC20;
+
+    // address public constant soloAddr = 0x1E0447b19BB6EcFdAe1e4AE1694b0C3659614e4e;
+    // address public constant wethAddr = 0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2;
+    // address public constant ethAddr = 0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE;
+
+    address public constant soloAddr = 0x4EC3570cADaAEE08Ae384779B0f3A45EF85289DE;
+    address public constant wethAddr = 0xd0A1E359811322d97991E03f863a0C30C2cF029C;
+    address public constant ethAddr = 0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE;
+
+    event LogDydxFlashLoan(
+        address indexed sender,
+        address indexed token,
+        uint amount
+    );
+
     function callFunction(
         address sender,
         Account.Info memory account,
         bytes memory data
-    ) public {
+    ) public override {
         require(sender == address(this), "not-same-sender");
         CastData memory cd;
         (cd.dsa, cd.token, cd.amount, cd.targets, cd.data) = abi.decode(
@@ -45,7 +65,7 @@ contract DydxFlashloaner is ICallee, DydxFlashloanBase {
             payable(cd.dsa).transfer(cd.amount);
         } else {
             tokenContract = IERC20(cd.token);
-            tokenContract.transfer(cd.dsa, cd.amount);
+            tokenContract.safeTransfer(cd.dsa, cd.amount);
         }
 
         DSAInterface(cd.dsa).cast(cd.targets, cd.data, 0xB7fA44c2E964B6EB24893f7082Ecc08c8d0c0F87);
@@ -66,7 +86,7 @@ contract DydxFlashloaner is ICallee, DydxFlashloanBase {
         Actions.ActionArgs[] memory operations = new Actions.ActionArgs[](3);
 
         operations[0] = _getWithdrawAction(marketId, _amount);
-        operations[1] = _getCallAction(data);
+        operations[1] = _getCallAction(encodeDsaAddr(data));
         operations[2] = _getDepositAction(marketId, _amount + 2);
 
         Account.Info[] memory accountInfos = new Account.Info[](1);