Add rate limit to bridging

script/DeployGydL1CCIPEscrow.s.sol

@@ -38,6 +38,10 @@ contract DeployGydL1CCIPEscrow is Script {
 
   uint256 gasLimit = 200_000; // max 200k gas to complete the bridging
 
+  uint256 capacity = 100_000; // max 100k GYD at once
+
+  uint256 refillRate = 10; // 1 GYD per second
+
   // CREATE3 Factory
   ICREATE3Factory factory =
     ICREATE3Factory(0x93FEC2C00BfE902F733B57c5a6CeeD7CD1384AE1);
@@ -60,7 +64,9 @@ contract DeployGydL1CCIPEscrow is Script {
       chainSelector: arbitrumChainSelector,
       metadata: IGydBridge.ChainMetadata({
         targetAddress: l2Address,
-        gasLimit: gasLimit
+        gasLimit: gasLimit,
+        capacity: capacity,
+        refillRate: refillRate
       })
     });
 

src/CCIPHelpers.sol

@@ -5,6 +5,8 @@ import {Client} from "ccip/libraries/Client.sol";
 import {IRouterClient} from "ccip/interfaces/IRouterClient.sol";
 import {Address} from "oz/utils/Address.sol";
 
+import {IGydBridge} from "./IGydBridge.sol";
+
 library CCIPHelpers {
   using Address for address payable;
 
@@ -43,4 +45,25 @@ library CCIPHelpers {
       payable(msg.sender).sendValue(refund);
     }
   }
+
+  function validateRateLimit(
+    uint64 chainSelector,
+    uint256 amount,
+    IGydBridge.RateLimitData memory rateLimit,
+    IGydBridge.ChainMetadata memory chainMeta
+  ) internal view returns (IGydBridge.RateLimitData memory) {
+    uint256 ellapsedSinceReplenish = block.timestamp - rateLimit.lastRefill;
+    uint256 amountToReplenish = ellapsedSinceReplenish * chainMeta.refillRate;
+    uint256 available = rateLimit.available + amountToReplenish;
+    if (available > chainMeta.capacity) {
+      available = chainMeta.capacity;
+    }
+    if (available < amount) {
+      revert IGydBridge.RateLimitExceeded(chainSelector, amount, available);
+    }
+    return IGydBridge.RateLimitData({
+      available: uint192(available - amount),
+      lastRefill: uint64(block.timestamp)
+    });
+  }
 }

src/GydL1CCIPEscrow.sol

@@ -51,6 +51,9 @@ contract GydL1CCIPEscrow is
   /// @notice The total amount of GYD bridged per chain
   uint256 public totalBridgedGYD;
 
+  /// @notice Rate limit data per chain
+  mapping(uint64 => RateLimitData) public rateLimitData;
+
   /// @notice Disable initializer on deploy
   constructor() {
     _disableInitializers();
@@ -76,11 +79,7 @@ contract GydL1CCIPEscrow is
     router = IRouterClient(_routerAddress);
     for (uint256 i; i < chains.length; i++) {
       chainsMetadata[chains[i].chainSelector] = chains[i].metadata;
-      emit ChainAdded(
-        chains[i].chainSelector,
-        chains[i].metadata.targetAddress,
-        chains[i].metadata.gasLimit
-      );
+      emit ChainSet(chains[i].chainSelector, chains[i].metadata);
     }
   }
 
@@ -106,14 +105,19 @@ contract GydL1CCIPEscrow is
    * @notice Allows the owner to support a new chain
    * @param chainSelector the selector of the chain
    * https://docs.chain.link/ccip/supported-networks/v1_2_0/mainnet#configuration
-   * @param gydAddress the GYD contract address on the chain
+   * @param metadata the metadata for this chain
    */
-  function addChain(uint64 chainSelector, address gydAddress, uint256 gasLimit)
+  function setChain(uint64 chainSelector, ChainMetadata memory metadata)
     external
     onlyRole(DEFAULT_ADMIN_ROLE)
   {
-    chainsMetadata[chainSelector] = ChainMetadata(gydAddress, gasLimit);
-    emit ChainAdded(chainSelector, gydAddress, gasLimit);
+    chainsMetadata[chainSelector] = metadata;
+    rateLimitData[chainSelector] = RateLimitData({
+      available: uint192(metadata.capacity),
+      lastRefill: uint64(block.timestamp)
+    });
+
+    emit ChainSet(chainSelector, metadata);
   }
 
   /**
@@ -213,18 +217,22 @@ contract GydL1CCIPEscrow is
     internal
     override
   {
-    address expectedSender =
-      chainsMetadata[any2EvmMessage.sourceChainSelector].targetAddress;
-    if (expectedSender == address(0)) {
-      revert ChainNotSupported(any2EvmMessage.sourceChainSelector);
-    }
+    uint64 chainSelector = any2EvmMessage.sourceChainSelector;
+    ChainMetadata memory chainMeta = chainsMetadata[chainSelector];
+    address expectedSender = chainMeta.targetAddress;
+
+    if (expectedSender == address(0)) revert ChainNotSupported(chainSelector);
+
     address actualSender = abi.decode(any2EvmMessage.sender, (address));
-    if (expectedSender != actualSender) {
-      revert MessageInvalid();
-    }
+    if (expectedSender != actualSender) revert MessageInvalid();
 
     (address recipient, uint256 amount, bytes memory data) =
       abi.decode(any2EvmMessage.data, (address, uint256, bytes));
+
+    rateLimitData[chainSelector] = CCIPHelpers.validateRateLimit(
+      chainSelector, amount, rateLimitData[chainSelector], chainMeta
+    );
+
     uint256 bridged = totalBridgedGYD;
     bridged -= amount;
     totalBridgedGYD = bridged;

src/IGydBridge.sol

@@ -5,19 +5,22 @@ interface IGydBridge {
   struct ChainMetadata {
     address targetAddress;
     uint256 gasLimit;
+    uint256 capacity;
+    uint256 refillRate;
   }
 
   struct ChainData {
     uint64 chainSelector;
     ChainMetadata metadata;
   }
 
-  /// @notice This event is emitted when a new chain is added
-  event ChainAdded(
-    uint64 indexed chainSelector,
-    address indexed targetAddress,
-    uint256 gasLimit
-  );
+  struct RateLimitData {
+    uint192 available;
+    uint64 lastRefill;
+  }
+
+  /// @notice This event is emitted when a new chain is set
+  event ChainSet(uint64 indexed chainSelector, ChainMetadata metadata);
 
   /// @notice This event is emitted when the gas limit is updated
   event GasLimitUpdated(uint64 indexed chainSelector, uint256 gasLimit);
@@ -46,4 +49,9 @@ interface IGydBridge {
 
   /// @notice This error is raised if the msg value is not enough for the fees
   error FeesNotCovered(uint256 fees);
+
+  /// @notice This error is raised if the rate limit is exceeded
+  error RateLimitExceeded(
+    uint64 chainSelector, uint256 requested, uint256 available
+  );
 }

src/L2Gyd.sol

@@ -38,6 +38,9 @@ contract L2Gyd is
   /// Only chains in this mapping can be bridged to
   mapping(uint64 => ChainMetadata) public chainsMetadata;
 
+  /// @notice Rate limit data per chain
+  mapping(uint64 => RateLimitData) public rateLimitData;
+
   /// @notice This error is raised if ownership is renounced
   error RenounceInvalid();
 
@@ -83,17 +86,20 @@ contract L2Gyd is
    * @notice Allows the owner to support a new chain
    * @param chainSelector the selector of the chain
    * https://docs.chain.link/ccip/supported-networks/v1_2_0/mainnet#configuration
-   * @param targetAddress the target address on the chain
+   * @param metadata the metadata for the chain to add
    * For Ethereum mainnet, this will be the address of the GYD escrow
    * For other L2s, it will be the L2Gyd contract
    */
-  function addChain(
-    uint64 chainSelector,
-    address targetAddress,
-    uint256 gasLimit
-  ) external onlyOwner {
-    chainsMetadata[chainSelector] = ChainMetadata(targetAddress, gasLimit);
-    emit ChainAdded(chainSelector, targetAddress, gasLimit);
+  function setChain(uint64 chainSelector, ChainMetadata memory metadata)
+    external
+    onlyOwner
+  {
+    chainsMetadata[chainSelector] = metadata;
+    rateLimitData[chainSelector] = RateLimitData({
+      available: uint192(metadata.capacity),
+      lastRefill: uint64(block.timestamp)
+    });
+    emit ChainSet(chainSelector, metadata);
   }
 
   /**
@@ -178,22 +184,23 @@ contract L2Gyd is
     internal
     override
   {
-    ChainMetadata memory chainMeta =
-      chainsMetadata[any2EvmMessage.sourceChainSelector];
+    uint64 chainSelector = any2EvmMessage.sourceChainSelector;
+    ChainMetadata memory chainMeta = chainsMetadata[chainSelector];
     address actualSender = abi.decode(any2EvmMessage.sender, (address));
-    if (actualSender != chainMeta.targetAddress) {
-      revert MessageInvalid();
-    }
+    if (actualSender != chainMeta.targetAddress) revert MessageInvalid();
 
     (address recipient, uint256 amount, bytes memory data) =
       abi.decode(any2EvmMessage.data, (address, uint256, bytes));
+
+    rateLimitData[chainSelector] = CCIPHelpers.validateRateLimit(
+      chainSelector, amount, rateLimitData[chainSelector], chainMeta
+    );
+
     _mint(recipient, amount);
     if (data.length > 0) {
       recipient.functionCall(data);
     }
 
-    emit GYDClaimed(
-      any2EvmMessage.sourceChainSelector, recipient, amount, totalSupply()
-    );
+    emit GYDClaimed(chainSelector, recipient, amount, totalSupply());
   }
 }

test/GydL1Escrow.t.sol

@@ -50,6 +50,7 @@ contract GydL1EscrowTest is Test {
   address gyd = address(0xe07F9D810a48ab5c3c914BA3cA53AF14E4491e8A);
   address ccipRouterAddress =
     address(0x80226fc0Ee2b096224EeAc085Bb9a8cba1146f7D);
+  uint256 capacity = 10_000_000e18;
 
   GydL1CCIPEscrow v1;
   GydL1CCIPEscrow proxyV1;
@@ -73,7 +74,13 @@ contract GydL1EscrowTest is Test {
     UUPSProxy proxy = new UUPSProxy(address(v1), v1Data);
     proxyV1 = GydL1CCIPEscrow(address(proxy));
     vm.prank(admin);
-    proxyV1.addChain(arbitrumChainSelector, gyd, gasLimit);
+    IGydBridge.ChainMetadata memory metadata = IGydBridge.ChainMetadata({
+      targetAddress: gyd,
+      gasLimit: gasLimit,
+      capacity: capacity,
+      refillRate: 100e18
+    });
+    proxyV1.setChain(arbitrumChainSelector, metadata);
 
     mockedV1 = new GydL1CCIPEscrow();
     router = new RouterMock();
@@ -87,7 +94,7 @@ contract GydL1EscrowTest is Test {
     UUPSProxy mockedProxy = new UUPSProxy(address(v1), mockedV1Data);
     mockedProxyV1 = GydL1CCIPEscrow(address(mockedProxy));
     vm.prank(admin);
-    mockedProxyV1.addChain(arbitrumChainSelector, gyd, gasLimit);
+    mockedProxyV1.setChain(arbitrumChainSelector, metadata);
     v2 = new GydL1EscrowV2Mock();
     proxyV2 = GydL1EscrowV2Mock(address(proxyV1));
 
@@ -138,7 +145,7 @@ contract GydL1EscrowTest is Test {
   /// @notice Make sure GydL1CCIPEscrow submit correct message to the bridge
   function testBridgeWithMockedBridge(uint256 bridgeAmount) public {
     vm.assume(bridgeAmount > 1 ether);
-    vm.assume(bridgeAmount < 1_000_000_000 ether);
+    vm.assume(bridgeAmount < capacity);
 
     vm.startPrank(alice);
     uint256 fees =
@@ -166,7 +173,7 @@ contract GydL1EscrowTest is Test {
   /// @notice Make sure GydL1CCIPEscrow can interact with the router
   function testBridgeWithRealBridge(uint256 bridgeAmount) public {
     vm.assume(bridgeAmount > 1 ether);
-    vm.assume(bridgeAmount < 1_000_000_000 ether);
+    vm.assume(bridgeAmount < capacity);
 
     vm.startPrank(alice);
     uint256 fees =
@@ -194,7 +201,7 @@ contract GydL1EscrowTest is Test {
   /// @notice Make sure to revert if message is invalid
   function testOnMessageReceivedInvalidMessage(uint256 bridgeAmount) public {
     vm.assume(bridgeAmount > 1 ether);
-    vm.assume(bridgeAmount < 1_000_000_000 ether);
+    vm.assume(bridgeAmount < capacity);
 
     vm.startPrank(alice);
     uint256 fees = proxyV1.getFee(arbitrumChainSelector, alice, bridgeAmount);
@@ -207,7 +214,7 @@ contract GydL1EscrowTest is Test {
     vm.stopPrank();
 
     address routerAddress = address(proxyV1.router());
-    (address originAddress,) = proxyV1.chainsMetadata(arbitrumChainSelector);
+    (address originAddress,,,) = proxyV1.chainsMetadata(arbitrumChainSelector);
     uint64 chainSelector = arbitrumChainSelector;
     bytes memory data = abi.encode(bob, 1 ether, "");
 
@@ -245,7 +252,7 @@ contract GydL1EscrowTest is Test {
   /// @notice Make sure user can claim the GYD
   function testOnMessageReceivedValidMessage(uint256 bridgeAmount) public {
     vm.assume(bridgeAmount > 1 ether);
-    vm.assume(bridgeAmount < 1_000_000_000 ether);
+    vm.assume(bridgeAmount < capacity);
 
     vm.startPrank(alice);
     uint256 fees = proxyV1.getFee(arbitrumChainSelector, alice, bridgeAmount);
@@ -256,7 +263,7 @@ contract GydL1EscrowTest is Test {
     vm.stopPrank();
 
     address routerAddress = address(proxyV1.router());
-    (address originAddress,) = proxyV1.chainsMetadata(arbitrumChainSelector);
+    (address originAddress,,,) = proxyV1.chainsMetadata(arbitrumChainSelector);
     uint64 chainSelector = arbitrumChainSelector;
     bytes memory messageData = abi.encode(bob, bridgeAmount, "");
 
@@ -270,13 +277,44 @@ contract GydL1EscrowTest is Test {
     assertEq(proxyV1.totalBridgedGYD(), 0);
   }
 
+  function testOnMessageReceivedValidMessageOverLimit() public {
+    uint256 bridgeAmount = capacity + 1;
+
+    vm.startPrank(alice);
+    uint256 fees = proxyV1.getFee(arbitrumChainSelector, alice, bridgeAmount);
+    deal(alice, fees);
+    deal(gyd, alice, bridgeAmount);
+    IERC20(gyd).safeIncreaseAllowance(address(proxyV1), bridgeAmount);
+    proxyV1.bridgeToken{value: fees}(arbitrumChainSelector, bob, bridgeAmount);
+    vm.stopPrank();
+
+    address routerAddress = address(proxyV1.router());
+    (address originAddress,,,) = proxyV1.chainsMetadata(arbitrumChainSelector);
+    uint64 chainSelector = arbitrumChainSelector;
+    bytes memory messageData = abi.encode(bob, bridgeAmount, "");
+
+    vm.startPrank(routerAddress);
+    vm.expectRevert(
+      abi.encodeWithSelector(
+        IGydBridge.RateLimitExceeded.selector,
+        chainSelector,
+        bridgeAmount,
+        capacity
+      )
+    );
+    proxyV1.ccipReceive(
+      _receivedMessage(chainSelector, originAddress, messageData)
+    );
+    vm.stopPrank();
+  }
+
   function testUpdateGasLimit() public {
     uint256 newGasLimit = 100_000;
 
     vm.prank(admin);
     proxyV1.updateGasLimit(arbitrumChainSelector, newGasLimit);
 
-    (, uint256 gasLimit_) = proxyV1.chainsMetadata(arbitrumChainSelector);
+    (, uint256 gasLimit_,,) = proxyV1.chainsMetadata(arbitrumChainSelector);
     assertEq(gasLimit_, newGasLimit);
   }