amount

zeta-chain · Nov 21, 2024 · 33281e8 · 33281e8
1 parent 5ee524f
commit 33281e8
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 14 deletions.
diff --git a/examples/router/contracts/Connected.sol b/examples/router/contracts/Connected.sol
@@ -15,14 +15,16 @@ contract Connected is ConnectedRouter {
 
     function onMessageReceive(
         bytes memory data,
-        address sender
+        address sender,
+        uint256 amount
     ) internal override {
         emit OnMessageReceiveEvent();
     }
 
     function onMessageRevert(
         bytes memory data,
-        address sender
+        address sender,
+        uint256 amount
     ) internal override {
         emit OnMessageRevertEvent();
     }

diff --git a/examples/router/contracts/ConnectedRouter.sol b/examples/router/contracts/ConnectedRouter.sol
@@ -61,17 +61,15 @@ contract ConnectedRouter is Ownable {
         bytes calldata message
     ) external payable onlyGateway returns (bytes4) {
         if (context.sender != router) revert Unauthorized();
-        (bytes memory data, address sender, bool isCall) = abi.decode(
-            message,
-            (bytes, address, bool)
-        );
+        (bytes memory data, address sender, uint256 amount, bool isCall) = abi
+            .decode(message, (bytes, address, uint256, bool));
 
         if (sender != counterparty) revert Unauthorized();
 
         if (isCall) {
-            onMessageReceive(data, sender);
+            onMessageReceive(data, sender, amount);
         } else {
-            onMessageRevert(data, sender);
+            onMessageRevert(data, sender, amount);
         }
         return "";
     }
@@ -86,14 +84,16 @@ contract ConnectedRouter is Ownable {
 
     function onMessageReceive(
         bytes memory data,
-        address sender
+        address sender,
+        uint256 amount
     ) internal virtual {
         // To be overridden in the child contract
     }
 
     function onMessageRevert(
         bytes memory data,
-        address sender
+        address sender,
+        uint256 amount
     ) internal virtual {
         // To be overridden in the child contract
     }

diff --git a/examples/router/contracts/Universal.sol b/examples/router/contracts/Universal.sol
@@ -97,13 +97,15 @@ contract Universal is UniversalContract, Ownable {
             callOptions.gasLimit
         );
 
+        uint256 withdrawAmount = out - gasFee;
+
         bytes memory m = callOptions.isArbitraryCall
-            ? abi.encodePacked(data, context.sender, true)
-            : abi.encode(data, context.sender, true);
+            ? abi.encodePacked(data, context.sender, withdrawAmount, true)
+            : abi.encode(data, context.sender, withdrawAmount, true);
 
         gateway.withdrawAndCall(
             receiver,
-            out - gasFee,
+            withdrawAmount,
             destination,
             m,
             callOptions,
@@ -141,7 +143,7 @@ contract Universal is UniversalContract, Ownable {
             abi.encodePacked(revertOptions.revertAddress),
             out - gasFee,
             destination,
-            abi.encode(data, receiver, false),
+            abi.encode(data, receiver, out - gasFee, false),
             CallOptions(onRevertGasLimit, false),
             RevertOptions(address(0), false, address(0), "", 0)
         );