From 41f01ba2375bf83d5fd968298e13ffc483de8f13 Mon Sep 17 00:00:00 2001
From: Alistair Singh <alistair.singh7@gmail.com>
Date: Tue, 24 Dec 2024 22:44:15 +0200
Subject: [PATCH] add reantrancy test

---
 contracts/test/Gateway.t.sol               | 26 +++++++++++++++
 contracts/test/mocks/ReantrantAttacker.sol | 39 ++++++++++++++++++++++
 2 files changed, 65 insertions(+)
 create mode 100644 contracts/test/mocks/ReantrantAttacker.sol

diff --git a/contracts/test/Gateway.t.sol b/contracts/test/Gateway.t.sol
index 951723d465..9af9ca7637 100644
--- a/contracts/test/Gateway.t.sol
+++ b/contracts/test/Gateway.t.sol
@@ -13,6 +13,7 @@ import {IUpgradable} from "../src/interfaces/IUpgradable.sol";
 import {Gateway} from "../src/Gateway.sol";
 import {MockGateway} from "./mocks/MockGateway.sol";
 import {MockGatewayV2} from "./mocks/MockGatewayV2.sol";
+import {ReantrantAttacker} from "./mocks/ReantrantAttacker.sol";
 import {GatewayProxy} from "../src/GatewayProxy.sol";
 
 import {AgentExecutor} from "../src/AgentExecutor.sol";
@@ -1022,4 +1023,29 @@ contract GatewayTest is Test {
         vm.expectRevert(Assets.TokenAlreadyRegistered.selector);
         IGateway(address(gateway)).registerToken{value: fee}(dotToken);
     }
+
+    function testReantrancyGuardReverts() public {
+        testRegisterToken();
+
+        ReantrantAttacker attacker = new ReantrantAttacker(address(gateway), address(token));
+        // Fund attacker
+        deal(address(attacker), 1 ether);
+        deal(address(token), address(attacker), 1);
+
+        uint128 amount = 1;
+        uint128 extra = 1;
+        uint128 destinationFee = 1;
+        ParaID paraID = ParaID.wrap(1000);
+
+        uint128 fee = uint128(IGateway(address(gateway)).quoteSendTokenFee(address(token), paraID, 0));
+
+        hoax(address(attacker));
+        token.approve(address(gateway), 1);
+
+        vm.expectRevert(NativeTransferFailed.selector);
+        hoax(address(attacker));
+        IGateway(address(gateway)).sendToken{value: fee + extra}(
+            address(token), paraID, recipientAddress32, destinationFee, amount
+        );
+    }
 }
diff --git a/contracts/test/mocks/ReantrantAttacker.sol b/contracts/test/mocks/ReantrantAttacker.sol
new file mode 100644
index 0000000000..3adbea953c
--- /dev/null
+++ b/contracts/test/mocks/ReantrantAttacker.sol
@@ -0,0 +1,39 @@
+// SPDX-License-Identifier: Apache-2.0
+pragma solidity 0.8.28;
+
+import {IGateway} from "../../src/interfaces/IGateway.sol";
+import {ParaID, MultiAddress, multiAddressFromBytes32} from "../../src/Types.sol";
+import {console} from "forge-std/console.sol";
+
+contract ReantrantAttacker {
+    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, address token) {
+        targetContract = IGateway(_targetAddr);
+        owner = msg.sender;
+        fee = targetContract.quoteSendTokenFee(token, assetHub, 0);
+    }
+
+    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 {
+        targetContract.sendToken{value: amount + fee + extra}(address(0), assetHub, recipientAddress32, 1, amount);
+    }
+}