Merge pull request #53 from ampleforth/flashloan-protection

Added tx_origin check

Author: aalavandhan

contracts/_mocks/MockConstructorRebaseCallerContract.sol

@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+pragma solidity 0.6.12;
+
+import "../_interfaces/IXCAmpleController.sol";
+
+contract MockConstructorRebaseCallerContract {
+    constructor(address policy) public {
+        // Take out a flash loan.
+        // Do something funky...
+        IXCAmpleController(policy).rebase(); // should fail
+        // pay back flash loan.
+    }
+}

contracts/_mocks/MockRebaseCallerContract.sol

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+pragma solidity 0.6.12;
+
+import "../_interfaces/IXCAmpleController.sol";
+
+contract MockRebaseCallerContract {
+    function callRebase(address policy) public returns (bool) {
+        // Take out a flash loan.
+        // Do something funky...
+        IXCAmpleController(policy).rebase(); // should fail
+        // pay back flash loan.
+        return true;
+    }
+}

contracts/satellite-chain/xc-ampleforth/XCAmpleController.sol

@@ -172,6 +172,8 @@ contract XCAmpleController is OwnableUpgradeable {
      *      on the rebase relayer.
      */
     function rebase() external {
+        require(msg.sender == tx.origin, "XCAmpleController: expected caller to be eoa");
+
         // recently reported epoch needs to be more than current globalEpoch in storage
         require(
             nextGlobalAmpleforthEpoch > globalAmpleforthEpoch,

helpers/deploy.js

@@ -211,10 +211,7 @@ async function deployChainBridgeContracts(
   };
 }
 
-async function deployTokenVault(
-  ethers,
-  deployer,
-  txParams = {}) {
+async function deployTokenVault(ethers, deployer, txParams = {}) {
   const tokenVault = await deployContract(
     ethers,
     'TokenVault',
@@ -255,7 +252,7 @@ async function deployChainBridgeBaseChainGatewayContracts(
     rebaseGateway,
   );
 
-  if(isAdmin){
+  if (isAdmin) {
     await (
       await bridge
         .connect(deployer)
@@ -281,7 +278,7 @@ async function deployChainBridgeBaseChainGatewayContracts(
   }
 
   const transferFnSig = CB_FUNCTION_SIG_baseChainTransfer(transferGateway);
-  if(isAdmin) {
+  if (isAdmin) {
     await (
       await bridge
         .connect(deployer)
@@ -306,7 +303,7 @@ async function deployChainBridgeBaseChainGatewayContracts(
     ]);
   }
 
-  if(await tokenVault.owner() == deployerAddress){
+  if ((await tokenVault.owner()) == deployerAddress) {
     await (
       await tokenVault
         .connect(deployer)
@@ -317,9 +314,7 @@ async function deployChainBridgeBaseChainGatewayContracts(
       'Failed to add whitelist transfer gateway to vault as deployer not vault owner',
     );
     console.log('Execute the following on-chain');
-    console.log('addBridgeGateway', [
-      transferGateway.address,
-    ]);
+    console.log('addBridgeGateway', [transferGateway.address]);
   }
 
   return { rebaseGateway, transferGateway };
@@ -364,7 +359,7 @@ async function deployChainBridgeSatelliteChainGatewayContracts(
   const reportRebaseFnSig = CB_FUNCTION_SIG_satelliteChainReportRebase(
     rebaseGateway,
   );
-  if(isAdmin) {
+  if (isAdmin) {
     await (
       await bridge.adminSetGenericResource(
         genericHandler.address,
@@ -388,7 +383,7 @@ async function deployChainBridgeSatelliteChainGatewayContracts(
   }
 
   const transferFnSig = CB_FUNCTION_SIG_satelliteChainTransfer(transferGateway);
-  if(isAdmin) {
+  if (isAdmin) {
     await (
       await bridge.adminSetGenericResource(
         genericHandler.address,

tasks/deploy/ampleforth.js

@@ -221,7 +221,6 @@ txTask('deploy:ampleforth_xc', 'Deploy cross chain ampleforth contract suite')
     );
   });
 
-
 txTask('deploy:token_vault', 'Deploy the token vault contract on base chain')
   .addParam('bridge', 'The bridge which secures the vault')
   .setAction(async (args, hre) => {
@@ -239,11 +238,7 @@ txTask('deploy:token_vault', 'Deploy the token vault contract on base chain')
 
     console.log('------------------------------------------------------------');
     console.log('Deploying TokenVault on base chain');
-    const tokenVault = await deployTokenVault(
-      hre.ethers,
-      deployer,
-      txParams,
-    );
+    const tokenVault = await deployTokenVault(hre.ethers, deployer, txParams);
 
     console.log('------------------------------------------------------------');
     console.log('Writing data to file');

tasks/deploy/chain_bridge.js

@@ -1,4 +1,3 @@
-
 const {
   types,
   task,
@@ -72,10 +71,7 @@ cbDeployTask(
     console.log('Deployer:', deployerAddress);
     console.log(txParams);
 
-    let bridge,
-      genericHandler,
-      erc20Handler,
-      erc721Handler;
+    let bridge, genericHandler, erc20Handler, erc721Handler;
 
     if (args.useDeployed) {
       console.log('Using deployed bridge');
@@ -89,7 +85,6 @@ cbDeployTask(
         'chainBridge/genericHandler',
         hre.ethers.provider,
       );
-
     } else {
       const chainBridge = await deployChainBridgeContracts(
         args,

test/integration/chain_bridge.js

@@ -571,10 +571,7 @@ describe('Transfers scenarios', function () {
     it('should revert', async function () {
       await baseChainAmplContracts.ampl
         .connect(userBBaseChainWallet)
-        .approve(
-          baseChainAmplContracts.tokenVault.address,
-          toAmplFixedPt('0'),
-        );
+        .approve(baseChainAmplContracts.tokenVault.address, toAmplFixedPt('0'));
       await expect(
         execXCSend(
           'base',

test/unit/satellite-chain/xc-ampleforth/xc_controller_rebase.js

@@ -69,6 +69,31 @@ describe('XCAmpleController:rebase:epoch', async () => {
   });
 });
 
+describe('XCAmpleController:rebase:called by a contract', function () {
+  it('should fail', async function () {
+    await controller.connect(bridge).reportRebase(2, 1000);
+    const rebaseCallerContract = await (
+      await ethers.getContractFactory('MockRebaseCallerContract')
+    )
+      .connect(deployer)
+      .deploy();
+    await expect(
+      rebaseCallerContract.callRebase(controller.address),
+    ).to.be.revertedWith('XCAmpleController: expected caller to be eoa');
+  });
+});
+
+describe('XCAmpleController:rebase:called by a contract which is being constructed', function () {
+  it('should fail', async function () {
+    await controller.connect(bridge).reportRebase(2, 1000);
+    await expect(
+      (await ethers.getContractFactory('MockConstructorRebaseCallerContract'))
+        .connect(deployer)
+        .deploy(controller.address),
+    ).to.be.revertedWith('XCAmpleController: expected caller to be eoa');
+  });
+});
+
 describe('XCAmpleController:rebase', async () => {
   beforeEach('setup XCAmpleController contract', async () => {
     await setupContracts();