From 33281e832f4150e5a18877259f599fcac6c7fe10 Mon Sep 17 00:00:00 2001
From: Denis Fadeev <denis@fadeev.org>
Date: Thu, 21 Nov 2024 11:02:26 +0300
Subject: [PATCH] amount

---
 examples/router/contracts/Connected.sol       |  6 ++++--
 examples/router/contracts/ConnectedRouter.sol | 16 ++++++++--------
 examples/router/contracts/Universal.sol       | 10 ++++++----
 3 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/examples/router/contracts/Connected.sol b/examples/router/contracts/Connected.sol
index f94dcbc..74b38d3 100644
--- 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
index 9d279ec..0afaa72 100644
--- 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
index 062a036..f12c506 100644
--- 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)
         );