From 1031c3493704b0e9d37beb64501737bcbcd5e336 Mon Sep 17 00:00:00 2001
From: Thrilok Kumar <thrilok2000@gmail.com>
Date: Sat, 20 Jun 2020 02:54:36 +0530
Subject: [PATCH] Added claim COMP function

---
 contracts/connectors/compound.sol | 32 +++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/contracts/connectors/compound.sol b/contracts/connectors/compound.sol
index 1cc7e2c..f42b88e 100644
--- a/contracts/connectors/compound.sol
+++ b/contracts/connectors/compound.sol
@@ -35,6 +35,7 @@ interface ComptrollerInterface {
     function exitMarket(address cTokenAddress) external returns (uint);
     function getAssetsIn(address account) external view returns (address[] memory);
     function getAccountLiquidity(address account) external view returns (uint, uint, uint);
+    function claimComp(address) external;
 }
 
 interface InstaMapping {
@@ -70,6 +71,10 @@ contract DSMath {
         z = add(mul(x, WAD), y / 2) / y;
     }
 
+    function sub(uint x, uint y) internal pure returns (uint z) {
+        require((z = x - y) <= x, "ds-math-sub-underflow");
+    }
+
 }
 
 
@@ -126,6 +131,13 @@ contract CompoundHelpers is Helpers {
         return 0x3d9819210A31b4961b30EF54bE2aeD79B9c9Cd3B;
     }
 
+    /**
+     * @dev Return COMP Token.
+     */
+    function getCompToken() internal pure returns (TokenInterface) {
+        return TokenInterface(0xc00e94Cb662C3520282E6f5717214004A7f26888);
+    }
+
     /**
      * @dev Return InstaDApp Mapping Addresses
      */
@@ -272,6 +284,7 @@ contract BasicResolver is CompoundHelpers {
 }
 
 contract ExtraResolver is BasicResolver {
+    event LogClaimedComp(uint256 CompAmt, uint256 setId);
     event LogDepositCToken(address indexed token, address cToken, uint256 tokenAmt, uint256 cTokenAmt,uint256 getId, uint256 setId);
     event LogWithdrawCToken(address indexed token, address cToken, uint256 cTokenAmt, uint256 getId, uint256 setId);
     event LogLiquidate(
@@ -283,6 +296,25 @@ contract ExtraResolver is BasicResolver {
         uint256 setId
     );
 
+    /**
+     * @dev Claim Accrued COMP Token.
+     * @param setId Set ctoken amount at this ID in `InstaMemory` Contract.
+    */
+    function ClaimComp(uint setId) external payable {
+        uint intialBal = getCompToken().balanceOf(address(this));
+        ComptrollerInterface(getComptrollerAddress()).claimComp(address(this));
+        uint finalBal = getCompToken().balanceOf(address(this));
+        uint amt = sub(finalBal, intialBal);
+
+        setUint(setId, amt);
+
+        emit LogClaimedComp(amt, setId);
+        bytes32 _eventCode = keccak256("LogClaimedComp(uint256,uint256)");
+        bytes memory _eventParam = abi.encode(amt, setId);
+        (uint _type, uint _id) = connectorID();
+        EventInterface(getEventAddr()).emitEvent(_type, _id, _eventCode, _eventParam);
+    }
+
     /**
      * @dev Deposit ETH/ERC20_Token.
      * @param token token address to depositCToken.(For ETH: 0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE)