fix[L04] Enforce chainId requirements in PolygonTokenBridger (#115)

Author: mrice32

contracts/PolygonTokenBridger.sol

@@ -41,14 +41,36 @@ contract PolygonTokenBridger is Lockable {
     // WETH contract on Ethereum.
     WETH9 public immutable l1Weth;
 
+    // Chain id for the L1 that this contract is deployed on or communicates with.
+    // For example: if this contract were meant to facilitate transfers from polygon to mainnet, this value would be
+    // the mainnet chainId 1.
+    uint256 public immutable l1ChainId;
+
+    // Chain id for the L2 that this contract is deployed on or communicates with.
+    // For example: if this contract were meant to facilitate transfers from polygon to mainnet, this value would be
+    // the polygon chainId 137.
+    uint256 public immutable l2ChainId;
+
+    modifier onlyChainId(uint256 chainId) {
+        _requireChainId(chainId);
+        _;
+    }
+
     /**
      * @notice Constructs Token Bridger contract.
      * @param _destination Where to send tokens to for this network.
      * @param _l1Weth Ethereum WETH address.
      */
-    constructor(address _destination, WETH9 _l1Weth) {
+    constructor(
+        address _destination,
+        WETH9 _l1Weth,
+        uint256 _l1ChainId,
+        uint256 _l2ChainId
+    ) {
         destination = _destination;
         l1Weth = _l1Weth;
+        l1ChainId = _l1ChainId;
+        l2ChainId = _l2ChainId;
     }
 
     /**
@@ -62,7 +84,7 @@ contract PolygonTokenBridger is Lockable {
         PolygonIERC20 token,
         uint256 amount,
         bool isWrappedMatic
-    ) public nonReentrant {
+    ) public nonReentrant onlyChainId(l2ChainId) {
         token.safeTransferFrom(msg.sender, address(this), amount);
 
         // In the wMatic case, this unwraps. For other ERC20s, this is the burn/send action.
@@ -76,12 +98,22 @@ contract PolygonTokenBridger is Lockable {
      * @notice Called by someone to send tokens to the destination, which should be set to the HubPool.
      * @param token Token to send to destination.
      */
-    function retrieve(IERC20 token) public nonReentrant {
+    function retrieve(IERC20 token) public nonReentrant onlyChainId(l1ChainId) {
         token.safeTransfer(destination, token.balanceOf(address(this)));
     }
 
     receive() external payable {
-        // Note: this should only happen on the mainnet side where ETH is sent to the contract directly by the bridge.
-        if (functionCallStackOriginatesFromOutsideThisContract()) l1Weth.deposit{ value: address(this).balance }();
+        if (functionCallStackOriginatesFromOutsideThisContract()) {
+            // This should only happen on the mainnet side where ETH is sent to the contract directly by the bridge.
+            _requireChainId(l1ChainId);
+            l1Weth.deposit{ value: address(this).balance }();
+        } else {
+            // This should only happen on the l2 side where matic is unwrapped by this contract.
+            _requireChainId(l2ChainId);
+        }
+    }
+
+    function _requireChainId(uint256 chainId) internal view {
+        require(block.chainid == chainId, "Cannot run method on this chain");
     }
 }

deploy/008_deploy_polygon_token_bridger_mainnet.ts

@@ -2,7 +2,7 @@
 import "hardhat-deploy";
 import { HardhatRuntimeEnvironment } from "hardhat/types/runtime";
 
-import { L1_ADDRESS_MAP } from "./consts";
+import { L1_ADDRESS_MAP, POLYGON_CHAIN_IDS } from "./consts";
 
 const func = async function (hre: HardhatRuntimeEnvironment) {
   const { deployments, getNamedAccounts, getChainId } = hre;
@@ -17,7 +17,7 @@ const func = async function (hre: HardhatRuntimeEnvironment) {
     from: deployer,
     log: true,
     skipIfAlreadyDeployed: true,
-    args: [hubPool.address, L1_ADDRESS_MAP[chainId].weth],
+    args: [hubPool.address, L1_ADDRESS_MAP[chainId].weth, chainId, POLYGON_CHAIN_IDS[chainId]],
     deterministicDeployment: "0x1234", // Salt for the create2 call.
   });
 };

deploy/010_deploy_polygon_token_bridger_polygon.ts

@@ -5,19 +5,20 @@ import { HardhatRuntimeEnvironment } from "hardhat/types/runtime";
 import { L1_ADDRESS_MAP } from "./consts";
 
 const func = async function (hre: HardhatRuntimeEnvironment) {
-  const { deployments, getNamedAccounts } = hre;
+  const { deployments, getNamedAccounts, getChainId } = hre;
   const { deploy } = deployments;
 
   const { deployer } = await getNamedAccounts();
 
+  const chainId = parseInt(await getChainId());
   const l1ChainId = parseInt(await hre.companionNetworks.l1.getChainId());
   const l1HubPool = await hre.companionNetworks.l1.deployments.get("HubPool");
 
   await deploy("PolygonTokenBridger", {
     from: deployer,
     log: true,
     skipIfAlreadyDeployed: true,
-    args: [l1HubPool.address, L1_ADDRESS_MAP[l1ChainId].weth],
+    args: [l1HubPool.address, L1_ADDRESS_MAP[l1ChainId].weth, l1ChainId, chainId],
     deterministicDeployment: "0x1234", // Salt for the create2 call.
   });
 };

deploy/consts.ts

@@ -60,3 +60,8 @@ export const L2_ADDRESS_MAP: { [key: number]: { [contractName: string]: string }
     fxChild: "0xCf73231F28B7331BBe3124B907840A94851f9f11",
   },
 };
+
+export const POLYGON_CHAIN_IDS: { [l1ChainId: number]: number } = {
+  1: 137,
+  5: 80001,
+};

test/chain-specific-spokepools/Polygon_SpokePool.ts

@@ -1,6 +1,16 @@
 import { TokenRolesEnum, ZERO_ADDRESS } from "@uma/common";
 import { mockTreeRoot, amountToReturn, amountHeldByPool } from "../constants";
-import { ethers, expect, Contract, SignerWithAddress, getContractFactory, seedContract, toWei } from "../utils";
+import {
+  ethers,
+  expect,
+  Contract,
+  SignerWithAddress,
+  getContractFactory,
+  seedContract,
+  toWei,
+  randomBigNumber,
+  seedWallet,
+} from "../utils";
 import { hubPoolFixture } from "../fixtures/HubPool.Fixture";
 import { constructSingleRelayerRefundTree } from "../MerkleLib.utils";
 
@@ -13,9 +23,13 @@ describe("Polygon Spoke Pool", function () {
     [owner, relayer, fxChild, rando] = await ethers.getSigners();
     ({ weth, hubPool, timer, l2Dai } = await hubPoolFixture());
 
+    // The spoke pool exists on l2, so add a random chainId for L1 to ensure that the L2's block.chainid will not match.
+    const l1ChainId = randomBigNumber();
+    const l2ChainId = await owner.getChainId();
+
     const polygonTokenBridger = await (
       await getContractFactory("PolygonTokenBridger", owner)
-    ).deploy(hubPool.address, weth.address);
+    ).deploy(hubPool.address, weth.address, l1ChainId, l2ChainId);
 
     dai = await (await getContractFactory("PolygonERC20Test", owner)).deploy();
     await dai.addMember(TokenRolesEnum.MINTER, owner.address);
@@ -25,6 +39,7 @@ describe("Polygon Spoke Pool", function () {
     ).deploy(polygonTokenBridger.address, owner.address, hubPool.address, weth.address, fxChild.address, timer.address);
 
     await seedContract(polygonSpokePool, relayer, [dai], weth, amountHeldByPool);
+    await seedWallet(owner, [], weth, toWei("1"));
   });
 
   it("Only correct caller can set the cross domain admin", async function () {
@@ -196,9 +211,13 @@ describe("Polygon Spoke Pool", function () {
   });
 
   it("PolygonTokenBridger retrieves and unwraps tokens correctly", async function () {
+    const l1ChainId = await owner.getChainId();
+
+    // Retrieve can only be performed on L1, so seed the L2 chainId with a non matching value.
+    const l2ChainId = randomBigNumber();
     const polygonTokenBridger = await (
       await getContractFactory("PolygonTokenBridger", owner)
-    ).deploy(hubPool.address, weth.address);
+    ).deploy(hubPool.address, weth.address, l1ChainId, l2ChainId);
 
     await expect(() =>
       owner.sendTransaction({ to: polygonTokenBridger.address, value: toWei("1") })
@@ -210,4 +229,43 @@ describe("Polygon Spoke Pool", function () {
       [toWei("1").mul(-1), toWei("1")]
     );
   });
+
+  it("PolygonTokenBridger doesn't allow L1 actions on L2", async function () {
+    // Make sure the L1 chain is different from the chainId where this is deployed.
+    const l1ChainId = randomBigNumber();
+    const l2ChainId = await owner.getChainId();
+
+    const polygonTokenBridger = await (
+      await getContractFactory("PolygonTokenBridger", owner)
+    ).deploy(hubPool.address, weth.address, l1ChainId, l2ChainId);
+
+    // Cannot send ETH directly into the contract on L2.
+    await expect(owner.sendTransaction({ to: polygonTokenBridger.address, value: toWei("1") })).to.be.revertedWith(
+      "Cannot run method on this chain"
+    );
+
+    // Cannot call retrieve on the contract on L2.
+    await weth.connect(owner).transfer(polygonTokenBridger.address, toWei("1"));
+    await expect(polygonTokenBridger.connect(owner).retrieve(weth.address)).to.be.revertedWith(
+      "Cannot run method on this chain"
+    );
+  });
+
+  it("PolygonTokenBridger doesn't allow L2 actions on L1", async function () {
+    const l1ChainId = await owner.getChainId();
+
+    // Make sure the L1 chain is different from the chainId where this is deployed.
+    const l2ChainId = randomBigNumber();
+
+    const polygonTokenBridger = await (
+      await getContractFactory("PolygonTokenBridger", owner)
+    ).deploy(hubPool.address, weth.address, l1ChainId, l2ChainId);
+
+    await weth.connect(owner).approve(polygonTokenBridger.address, toWei("1"));
+
+    // Cannot call send on the contract on L1.
+    await expect(polygonTokenBridger.connect(owner).send(weth.address, toWei("1"), false)).to.be.revertedWith(
+      "Cannot run method on this chain"
+    );
+  });
 });