From c461ec0922ac9a6d49cd7fbe4cd105a54cd0ae61 Mon Sep 17 00:00:00 2001
From: Kirill Fedoseev <k.fedoseev@innopolis.university>
Date: Thu, 20 Feb 2020 00:05:50 +0300
Subject: [PATCH] Add restriction on claiming Chai token

---
 .../ForeignBridgeErcToNative.sol              |  2 +
 test/erc_to_native/foreign_bridge.test.js     | 42 +++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/contracts/upgradeable_contracts/erc20_to_native/ForeignBridgeErcToNative.sol b/contracts/upgradeable_contracts/erc20_to_native/ForeignBridgeErcToNative.sol
index 1ff1e53e6..4d5c457a4 100644
--- a/contracts/upgradeable_contracts/erc20_to_native/ForeignBridgeErcToNative.sol
+++ b/contracts/upgradeable_contracts/erc20_to_native/ForeignBridgeErcToNative.sol
@@ -65,6 +65,8 @@ contract ForeignBridgeErcToNative is BasicForeignBridge, ERC20Bridge, OtherSideB
 
     function claimTokens(address _token, address _to) public {
         require(_token != address(erc20token()));
+        // Chai token is not claimable if investing into Chai is enabled
+        require(_token != address(chaiToken()) || !isChaiTokenEnabled());
         if (_token == address(halfDuplexErc20token())) {
             // SCD is not claimable if the bridge accepts deposits of this token
             // solhint-disable-next-line not-rely-on-time
diff --git a/test/erc_to_native/foreign_bridge.test.js b/test/erc_to_native/foreign_bridge.test.js
index b22c799f0..429bf96c9 100644
--- a/test/erc_to_native/foreign_bridge.test.js
+++ b/test/erc_to_native/foreign_bridge.test.js
@@ -1897,5 +1897,47 @@ contract('ForeignBridge_ERC20_to_Native', async accounts => {
         expect(await foreignBridgeProxy.dsrBalance()).to.be.bignumber.gte(ether('0.4'))
       })
     })
+
+    describe('claimTokens', async () => {
+      let foreignBridgeProxy
+
+      beforeEach(async () => {
+        foreignBridgeProxy = await EternalStorageProxy.new({ from: accounts[2] }).should.be.fulfilled
+        await foreignBridgeProxy.upgradeTo('1', foreignBridge.address, { from: accounts[2] }).should.be.fulfilled
+        foreignBridgeProxy = await ForeignBridgeErcToNativeMock.at(foreignBridgeProxy.address)
+        foreignBridgeProxy.setChaiToken(chaiToken.address)
+        await foreignBridgeProxy.initialize(
+          validatorContract.address,
+          token.address,
+          requireBlockConfirmations,
+          gasPrice,
+          [dailyLimit, maxPerTx, minPerTx],
+          [homeDailyLimit, homeMaxPerTx],
+          owner,
+          decimalShiftZero,
+          otherSideBridge.address,
+          { from: accounts[2] }
+        )
+      })
+
+      it('should not allow to claim Chai, if it is enabled', async () => {
+        await foreignBridgeProxy.initializeChaiToken({ from: owner })
+        await token.mint(foreignBridgeProxy.address, halfEther)
+        await foreignBridgeProxy.setMinDaiTokenBalance(ether('0.1'), { from: owner })
+        await foreignBridgeProxy.convertDaiToChai()
+        expect(await foreignBridgeProxy.isChaiTokenEnabled()).to.be.equal(true)
+
+        await foreignBridgeProxy.claimTokens(chaiToken.address, accounts[2], { from: accounts[2] }).should.be.rejected
+      })
+
+      it('should allow to claim chai after it is disabled', async () => {
+        expect(await foreignBridgeProxy.isChaiTokenEnabled()).to.be.equal(false)
+        await token.mint(accounts[3], halfEther)
+        await token.approve(chaiToken.address, halfEther, { from: accounts[3] })
+        await chaiToken.join(accounts[3], halfEther, { from: accounts[3] }).should.be.fulfilled
+
+        await foreignBridgeProxy.claimTokens(chaiToken.address, accounts[2], { from: accounts[2] }).should.be.fulfilled
+      })
+    })
   })
 })