Tidy and tests

src/MessageHubLibV1.sol

@@ -18,7 +18,7 @@ library MessageHubLibV1 {
     /*-»-»-»-»-»-»-»-»-»-»-»-»-»-»-»-»-»-»-»-»-»-»-»-»-»-»-»-»-»-»*/
 
     /// @dev Returns the sender or signer.
-    ///      Note: this function will never return a zero address.
+    /// Note: this function will never return a zero address.
     function senderOrSigner() internal view returns (address result) {
         result = LibMulticaller.senderOrSigner();
         if (result == MESSAGE_HUB) {

src/MessageHubV1.sol

@@ -4,6 +4,7 @@ pragma solidity ^0.8.23;
 import {UUPSUpgradeable} from "solady/utils/UUPSUpgradeable.sol";
 import {SafeTransferLib} from "solady/utils/SafeTransferLib.sol";
 import {LibClone} from "solady/utils/LibClone.sol";
+import {LibBytes} from "solady/utils/LibBytes.sol";
 import {ERC7821} from "solady/accounts/ERC7821.sol";
 import {Origin, OAppReceiverUpgradeable} from "layerzero-oapp/contracts/oapp-upgradeable/OAppReceiverUpgradeable.sol";
 
@@ -108,13 +109,21 @@ contract MessageHubV1 is UUPSUpgradeable, OAppReceiverUpgradeable {
         return LibClone.predictDeterministicAddress(_podImplementation, args, keccak256(args), address(this));
     }
 
-    /// @dev For efficiency, we will use the receive function
-    ///      to return the original sender and the sender.
+    /// @dev Returns `abi.encode(sender, originalSender, originalSenderType)`.
+    /// This is queried by the MessageHubPodV1 or contracts that use MessageHubLibV1.
     receive() external payable {
         assembly ("memory-safe") {
-            mstore(0x40, tload(_ORIGINAL_SENDER_TYPE_TRANSIENT_SLOT))
-            mstore(0x20, tload(_ORIGINAL_SENDER_TRANSIENT_SLOT))
+            // The `sender` can be either:
+            // - A MessageHubPodV1 that is deployed on-the-fly if the
+            //   message was initiated from a non-Ethereum address.
+            // - An Ethereum address that initiated the message.
+            //
+            // The `originalSender` can be either:
+            // - A non-Ethereum address.
+            // - An Ethereum address that matches `sender`.
             mstore(0x00, tload(_SENDER_TRANSIENT_SLOT))
+            mstore(0x20, tload(_ORIGINAL_SENDER_TRANSIENT_SLOT))
+            mstore(0x40, tload(_ORIGINAL_SENDER_TYPE_TRANSIENT_SLOT))
             return(0x00, 0x60)
         }
     }
@@ -143,6 +152,7 @@ contract MessageHubV1 is UUPSUpgradeable, OAppReceiverUpgradeable {
     }
 
     /// @dev Returns the immutable arguments for the sub account.
+    /// `abi.encodePacked(keccak256(abi.encode(originalSender, originalSenderType)), address(this))`.
     function _subAccountArgs(bytes32 originalSender, uint256 originalSenderType)
         internal
         view
@@ -151,8 +161,8 @@ contract MessageHubV1 is UUPSUpgradeable, OAppReceiverUpgradeable {
         if (originalSender == bytes32(0)) revert OriginalSenderIsZero();
         assembly ("memory-safe") {
             result := mload(0x40)
-            mstore(0x20, originalSenderType)
             mstore(0x00, originalSender)
+            mstore(0x20, originalSenderType)
             // Address of the mothership.
             mstore(add(result, 0x34), address())
             // Hash of the original sender and sender type.
@@ -163,7 +173,7 @@ contract MessageHubV1 is UUPSUpgradeable, OAppReceiverUpgradeable {
     }
 
     /// @dev Decodes and forwards the message to the target.
-    ///      This modifier is to be attached onto the `_lzReceive` function.
+    /// This modifier is to be attached onto the `_lzReceive` function.
     modifier forwardMessage(bytes calldata message) {
         bytes32 originalSender;
         uint256 originalSenderType;
@@ -222,6 +232,7 @@ contract MessageHubV1 is UUPSUpgradeable, OAppReceiverUpgradeable {
 
             toRefund := mul(sub(selfbalance(), balanceBefore), gt(selfbalance(), balanceBefore))
         }
+        // If the balance has somehow increased, refunds to the sender.
         if (toRefund != 0) {
             SafeTransferLib.forceSafeTransferETH(sender, toRefund);
         }

test/MessageHubV1.t.sol

@@ -8,15 +8,47 @@ import "devtools/mocks/EndpointV2Mock.sol";
 import {LibClone} from "solady/utils/LibClone.sol";
 import {SafeCastLib} from "solady/utils/SafeCastLib.sol";
 
+contract Target {
+    address public hub;
+    uint256 public x;
+    uint256 public msgValueDuringSetX;
+    address public msgSenderDuringSetX;
+    address public senderDuringSetX;
+    bytes32 public originalSenderDuringSetX;
+    uint256 public originalSenderTypeDuringSetX;
+    uint256 public refundAmount;
+
+    function depositRefund() public payable {
+        refundAmount += msg.value;
+    }
+
+    function setHub(address newHub) public {
+        hub = newHub;
+    }
+
+    function setX(uint256 newX) public payable {
+        msgValueDuringSetX = msg.value;
+        msgSenderDuringSetX = msg.sender;
+        x = newX;
+        (bool success, bytes memory results) = hub.staticcall("");
+        (senderDuringSetX, originalSenderDuringSetX, originalSenderTypeDuringSetX) =
+            abi.decode(results, (address, bytes32, uint256));
+        require(success);
+
+        (success,) = msg.sender.call{value: refundAmount}("");
+        require(success);
+        refundAmount = 0;
+    }
+}
+
 contract MessagehubV1Test is SoladyTest {
     MockMessageHubV1 hub;
     MessageInitiatorV1 initiator;
 
     EndpointV2Mock eid1;
     EndpointV2Mock eid2;
 
-    uint256 x;
-    uint256 valueDuringSetX;
+    Target target;
 
     address constant ALICE = address(111);
 
@@ -45,65 +77,88 @@ contract MessagehubV1Test is SoladyTest {
 
         eid1.setDestLzEndpoint(address(initiator), address(eid2));
         eid2.setDestLzEndpoint(address(hub), address(eid1));
+
+        target = new Target();
+        target.setHub(address(hub));
+    }
+
+    function testSubAccountArgsDifferential(bytes32 originalSender, uint256 originalSenderType) public view {
+        bytes memory expected =
+            abi.encodePacked(keccak256(abi.encode(originalSender, originalSenderType)), address(hub));
+        assertEq(hub.subAccountArgs(originalSender, originalSenderType), expected);
     }
 
     function testCrosschain() public {
-        uint256 gas = 100000;
+        uint256 gas = 1000000;
         uint256 value = 1 ether;
         uint256 newX = _random();
 
         bytes memory data = abi.encodeWithSignature("setX(uint256)", newX);
 
-        uint256 nativeFee = initiator.quoteWithDefaultOptions(address(this), data, gas, value);
+        uint256 nativeFee = initiator.quoteWithDefaultOptions(address(target), data, gas, value);
 
         vm.deal(ALICE, 10 ether);
 
         vm.prank(ALICE);
-        initiator.sendWithDefaultOptions{value: nativeFee}(address(this), data, gas, value);
+        initiator.sendWithDefaultOptions{value: nativeFee}(address(target), data, gas, value);
 
-        assertEq(x, newX);
-        assertEq(valueDuringSetX, 1 ether);
+        assertEq(target.x(), newX);
+        assertEq(target.msgValueDuringSetX(), 1 ether);
+        assertEq(target.msgSenderDuringSetX(), address(hub));
+
+        assertEq(target.originalSenderDuringSetX(), bytes32(uint256(uint160(ALICE))));
+        assertEq(target.originalSenderTypeDuringSetX(), 0);
+        assertEq(target.senderDuringSetX(), ALICE);
     }
 
-    function testMothership() public {
-        bytes32 originalSender = bytes32(_random());
-        while (originalSender == bytes32(0)) originalSender = bytes32(_random());
+    struct _TestMothershipTemps {
+        bool originalSenderIsEthereumAddress;
+        bytes32 originalSender;
+        uint256 originalSenderType;
+        address sender;
+        bytes message;
+        uint256 newX;
+        uint256 refundAmount;
+    }
 
-        uint256 senderType = _random();
-        while (senderType == 0) senderType = _random();
+    function testMothership(bytes32) public {
+        _TestMothershipTemps memory t;
 
-        address expected = hub.predictSubAccount(originalSender, senderType);
-        assertEq(hub.createSubAccount(originalSender, senderType), expected);
+        t.originalSenderIsEthereumAddress = _randomChance(2);
 
-        uint256 newX = _random();
+        if (t.originalSenderIsEthereumAddress) {
+            t.sender = _randomNonZeroAddress();
+            t.originalSenderType = 0;
+            t.originalSender = bytes32(uint256(uint160(t.sender)));
+        } else {
+            t.originalSender = bytes32(_random());
+            t.originalSenderType = 1 | _random();
+            while (t.originalSender == bytes32(0)) t.originalSender = bytes32(_random());
+            t.sender = hub.predictSubAccount(t.originalSender, t.originalSenderType);
+            assertEq(hub.createSubAccount(t.originalSender, t.originalSenderType), t.sender);
+        }
+        t.newX = _random();
 
         vm.deal(address(this), 10 ether);
+        vm.deal(address(hub), _random() % 0.1 ether);
 
-        bytes memory message = abi.encode(
-            originalSender,
-            senderType,
-            expected,
-            _encodeExecuteCalldata(address(this), 1 ether, abi.encodeWithSignature("setX(uint256)", newX))
+        t.message = abi.encode(
+            t.originalSender, t.originalSenderType, address(target), abi.encodeWithSignature("setX(uint256)", t.newX)
         );
 
-        hub.forward{value: 1 ether}(message);
+        vm.deal(address(this), 10 ether);
+        t.refundAmount = _random() % 1 ether;
+        target.depositRefund{value: t.refundAmount}();
 
-        assertEq(x, newX);
-        assertEq(valueDuringSetX, 1 ether);
-    }
+        hub.forward{value: 1 ether}(t.message);
 
-    function _encodeExecuteCalldata(address target, uint256 value, bytes memory data) internal returns (bytes memory) {
-        Call[] memory calls = new Call[](1);
-        calls[0].target = target;
-        calls[0].value = value;
-        calls[0].data = data;
-        bytes memory executionData = abi.encode(calls);
-        bytes32 mode = bytes32(bytes1(0x01));
-        return abi.encodeWithSignature("execute(bytes32,bytes)", mode, executionData);
-    }
+        assertEq(target.msgValueDuringSetX(), 1 ether);
+        assertEq(target.msgSenderDuringSetX(), address(hub));
+        assertEq(target.originalSenderDuringSetX(), t.originalSender);
+        assertEq(target.originalSenderTypeDuringSetX(), t.originalSenderType);
+        assertEq(target.senderDuringSetX(), t.sender);
 
-    function setX(uint256 value) public payable {
-        valueDuringSetX = msg.value;
-        x = value;
+        assertEq(target.x(), t.newX);
+        assertEq(t.sender.balance, t.refundAmount);
     }
 }

test/mocks/MockMessageHubV1.sol

@@ -9,4 +9,8 @@ contract MockMessageHubV1 is MessageHubV1 {
     }
 
     function forward(bytes calldata message) public payable forwardMessage(message) {}
+
+    function subAccountArgs(bytes32 originalSender, uint256 originalSenderType) public view returns (bytes memory) {
+        return _subAccountArgs(originalSender, originalSenderType);
+    }
 }