From ad9ac7c7a4c8794fede7e66aac5aaa0c4346e8c8 Mon Sep 17 00:00:00 2001
From: Gerardo Nardelli <gerardo.nardelli@altoros.com>
Date: Thu, 11 Oct 2018 16:24:16 -0300
Subject: [PATCH] Add check for bridge address on setBlockRewardContract

---
 contracts/IBlockReward.sol                        |  1 +
 contracts/test/BlockReward.sol                    |  9 +++++++++
 .../erc20_to_native/HomeBridgeErcToNative.sol     | 13 +++++++++++++
 test/erc_to_native/home_bridge.test.js            | 15 +++++++++++++++
 4 files changed, 38 insertions(+)

diff --git a/contracts/IBlockReward.sol b/contracts/IBlockReward.sol
index dfff214..e027e7f 100644
--- a/contracts/IBlockReward.sol
+++ b/contracts/IBlockReward.sol
@@ -5,4 +5,5 @@ interface IBlockReward {
     function addExtraReceiver(uint256 _amount, address _receiver) external;
     function mintedTotally() public view returns (uint256);
     function mintedTotallyByBridge(address _bridge) public view returns(uint256);
+    function bridgesAllowed() public pure returns(address[3]);
 }
diff --git a/contracts/test/BlockReward.sol b/contracts/test/BlockReward.sol
index 70ebb2f..94f3675 100644
--- a/contracts/test/BlockReward.sol
+++ b/contracts/test/BlockReward.sol
@@ -10,6 +10,7 @@ contract BlockReward is IBlockReward {
     uint256 public mintedCoins = 0;
     mapping(bytes32 => uint256) internal uintStorage;
     bytes32 internal constant MINTED_TOTALLY_BY_BRIDGE = "mintedTotallyByBridge";
+    uint256 public constant bridgesAllowedLength = 3;
 
     function () external payable {
     }
@@ -36,4 +37,12 @@ contract BlockReward is IBlockReward {
         bytes32 hash = keccak256(abi.encode(MINTED_TOTALLY_BY_BRIDGE, _bridge));
         uintStorage[hash] = uintStorage[hash].add(_amount);
     }
+
+    function bridgesAllowed() public pure returns(address[bridgesAllowedLength]) {
+        return([
+            0x0cDEE95B0Ed18FccEB4c344Df4dd1C1642798a8b,
+            0x320051BbD4eeE344Bb86F0A858d03595837463eF,
+            0xce42bdB34189a93c55De250E011c68FaeE374Dd3
+        ]);
+    }
 }
diff --git a/contracts/upgradeable_contracts/erc20_to_native/HomeBridgeErcToNative.sol b/contracts/upgradeable_contracts/erc20_to_native/HomeBridgeErcToNative.sol
index 463f4df..9f81066 100644
--- a/contracts/upgradeable_contracts/erc20_to_native/HomeBridgeErcToNative.sol
+++ b/contracts/upgradeable_contracts/erc20_to_native/HomeBridgeErcToNative.sol
@@ -67,6 +67,7 @@ contract HomeBridgeErcToNative is EternalStorage, BasicBridge, BasicHomeBridge {
 
     function setBlockRewardContract(address _blockReward) public onlyOwner {
         require(_blockReward != address(0) && isContract(_blockReward));
+        require(_isBridgeAllowed(_blockReward, this));
         addressStorage[keccak256(abi.encodePacked("blockRewardContract"))] = _blockReward;
     }
 
@@ -84,4 +85,16 @@ contract HomeBridgeErcToNative is EternalStorage, BasicBridge, BasicHomeBridge {
     function setTotalBurntCoins(uint256 _amount) internal {
         uintStorage[keccak256(abi.encodePacked("totalBurntCoins"))] = _amount;
     }
+
+    function _isBridgeAllowed(address _blockReward, address _addr) private pure returns(bool) {
+        IBlockReward RewardContract = IBlockReward(_blockReward);
+        address[3] memory bridges = RewardContract.bridgesAllowed();
+
+        for (uint256 i = 0; i < bridges.length; i++) {
+            if (_addr == bridges[i]) {
+                return true;
+            }
+        }
+        return false;
+    }
 }
diff --git a/test/erc_to_native/home_bridge.test.js b/test/erc_to_native/home_bridge.test.js
index 1abf2d9..5e0504b 100644
--- a/test/erc_to_native/home_bridge.test.js
+++ b/test/erc_to_native/home_bridge.test.js
@@ -73,6 +73,21 @@ contract('HomeBridge_ERC20_to_Native', async (accounts) => {
       secondBlockRewardContract.address.should.be.equal(await homeContract.blockRewardContract())
     })
 
+    it('can update block reward contract only if bridge is allowed', async () => {
+      ZERO_ADDRESS.should.be.equal(await homeContract.blockRewardContract())
+      await homeContract.initialize(validatorContract.address, '3', '2', '1', gasPrice, requireBlockConfirmations, blockRewardContract.address).should.be.fulfilled
+      blockRewardContract.address.should.be.equal(await homeContract.blockRewardContract())
+
+      await homeContract.setBlockRewardContract(blockRewardContract.address)
+      blockRewardContract.address.should.be.equal(await homeContract.blockRewardContract())
+
+      const notAllowedHomeBridge = await HomeBridge.new()
+
+      await notAllowedHomeBridge.setBlockRewardContract(blockRewardContract.address).should.be.rejectedWith(ERROR_MSG)
+
+      ZERO_ADDRESS.should.be.equal(await notAllowedHomeBridge.blockRewardContract())
+    })
+
     it('cant set maxPerTx > dailyLimit', async () => {
       false.should.be.equal(await homeContract.isInitialized())