Add ReentrancyGuard

contracts/src/Gateway.sol

@@ -3,6 +3,7 @@
 pragma solidity 0.8.25;
 
 import {MerkleProof} from "openzeppelin/utils/cryptography/MerkleProof.sol";
+import {ReentrancyGuard} from "openzeppelin/security/ReentrancyGuard.sol";
 import {Verification} from "./Verification.sol";
 
 import {Assets} from "./Assets.sol";
@@ -53,7 +54,7 @@ import {OperatorStorage} from "./storage/OperatorStorage.sol";
 
 import {UD60x18, ud60x18, convert} from "prb/math/src/UD60x18.sol";
 
-contract Gateway is IGateway, IInitializable, IUpgradable {
+contract Gateway is IGateway, IInitializable, IUpgradable, ReentrancyGuard {
     using Address for address;
     using SafeNativeTransfer for address payable;
 
@@ -469,7 +470,7 @@ contract Gateway is IGateway, IInitializable, IUpgradable {
         MultiAddress calldata destinationAddress,
         uint128 destinationFee,
         uint128 amount
-    ) external payable {
+    ) external payable nonReentrant {
         Ticket memory ticket = Assets.sendToken(
             token, msg.sender, destinationChain, destinationAddress, destinationFee, MAX_DESTINATION_FEE, amount
         );

contracts/test/Gateway.t.sol

@@ -28,6 +28,7 @@ import {PricingStorage} from "../src/storage/PricingStorage.sol";
 import {IERC20} from "../src/interfaces/IERC20.sol";
 import {TokenLib} from "../src/TokenLib.sol";
 import {Token} from "../src/Token.sol";
+import {Attacker} from "./mocks/Attacker.sol";
 
 import {
     UpgradeParams,
@@ -58,7 +59,9 @@ import {WETH9} from "canonical-weth/WETH9.sol";
 import {UD60x18, ud60x18, convert} from "prb/math/src/UD60x18.sol";
 
 contract GatewayTest is Test {
+    error NativeTransferFailed();
     // Emitted when token minted/burnt/transfered
+
     event Transfer(address indexed from, address indexed to, uint256 value);
 
     ParaID public bridgeHubParaID = ParaID.wrap(1013);
@@ -107,6 +110,8 @@ contract GatewayTest is Test {
     // tokenID for DOT
     bytes32 public dotTokenID;
 
+    Attacker public attacker;
+
     function setUp() public {
         AgentExecutor executor = new AgentExecutor();
         gatewayLogic = new MockGateway(
@@ -154,6 +159,8 @@ contract GatewayTest is Test {
         recipientAddress20 = multiAddressFromBytes20(bytes20(keccak256("recipient")));
 
         dotTokenID = bytes32(uint256(1));
+
+        attacker = new Attacker(address(gateway));
     }
 
     function makeCreateAgentCommand() public pure returns (Command, bytes memory) {
@@ -1244,4 +1251,25 @@ contract GatewayTest is Test {
         vm.expectRevert(Assets.InvalidToken.selector);
         IGateway(address(gateway)).registerToken{value: fee}(address(0));
     }
+
+    function testExploitSendingEthWithReentrancyWillFail() public {
+        uint128 amount = 1;
+        uint128 extra = 1;
+        uint128 destinationFee = 1;
+        ParaID paraID = ParaID.wrap(1000);
+        uint128 fee = uint128(IGateway(address(gateway)).quoteSendTokenFee(address(0), paraID, 1));
+        // Fund attacker
+        vm.deal(address(attacker), 1 ether);
+        assertTrue(address(gateway).balance == 0);
+        // vm.expectRevert("ReentrancyGuard: reentrant call");
+        vm.expectRevert(NativeTransferFailed.selector);
+        vm.prank(address(attacker));
+        IGateway(address(gateway)).sendToken{value: amount + fee + extra}(
+            address(0), paraID, recipientAddress32, destinationFee, amount
+        );
+        // console.log("%s:%d", "Attacker's balance", address(attacker).balance);
+        // console.log("%s:%d", "Gateway's balance", address(gateway).balance);
+        // assertTrue(address(attacker).balance < 0.9 ether);
+        // assertTrue(address(gateway).balance == 0.1 ether);
+    }
 }

contracts/test/mocks/Attacker.sol

@@ -0,0 +1,45 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.25;
+
+import {IGateway} from "../../src/interfaces/IGateway.sol";
+import {ParaID, MultiAddress, multiAddressFromBytes32} from "../../src/Types.sol";
+import {console} from "forge-std/console.sol";
+
+contract Attacker {
+    address public owner;
+    IGateway targetContract;
+    uint256 targetValue = 0.9 ether;
+    uint256 fee;
+    ParaID assetHub = ParaID.wrap(1000);
+    uint128 amount = 1;
+    uint128 extra = 1;
+    MultiAddress recipientAddress32 = multiAddressFromBytes32(keccak256("recipient"));
+
+    constructor(address _targetAddr) {
+        targetContract = IGateway(_targetAddr);
+        owner = msg.sender;
+        fee = targetContract.quoteSendTokenFee(address(0), assetHub, 1);
+    }
+
+    function balance() public view returns (uint256) {
+        return address(this).balance;
+    }
+
+    function withdrawAll() public returns (bool) {
+        require(msg.sender == owner, "my money!!");
+        uint256 totalBalance = address(this).balance;
+        (bool sent,) = msg.sender.call{value: totalBalance}("");
+        require(sent, "Failed to send Ether");
+        return sent;
+    }
+
+    receive() external payable {
+        uint256 currentBalance = address(this).balance;
+        console.log("%s:%d", "Attacker's current balance", currentBalance);
+        uint256 gatewayBalance = address(targetContract).balance;
+        console.log("%s:%d", "Gateway's current balance", gatewayBalance);
+        if (currentBalance >= targetValue) {
+            targetContract.sendToken{value: amount + fee + extra}(address(0), assetHub, recipientAddress32, 1, amount);
+        }
+    }
+}