✨ Minimal Pod

src/accounts/Pod.sol

@@ -0,0 +1,141 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+import {Receiver} from "./Receiver.sol";
+
+/// @notice Hyper minimal sub account contract that is intended to be controlled by a mothership contract.
+/// @author Solady (https://github.com/vectorized/solady/blob/main/src/accounts/LibERC6551.sol)
+abstract contract Pod is Receiver {
+    /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
+    /*                          STRUCTS                           */
+    /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
+    /// @dev Call struct for the `executeBatch` function.
+    struct Call {
+        address target;
+        uint256 value;
+        bytes data;
+    }
+
+    /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
+    /*                       CUSTOM ERRORS                        */
+    /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
+
+    /// @dev The function selector is not recognized.
+    error FnSelectorNotRecognized();
+
+    /// @dev The caller is not mothership.
+    error CallerNotMothership();
+
+    /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
+    /*                  CONSTANTS AND IMMUTABLES                  */
+    /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
+
+    /// @dev The Pod `Mothership` slot is given by:
+    ///         `bytes9(keccak256("_MOTHERSHIP_SLOT"))`.
+    uint256 internal constant _MOTHERSHIP_SLOT = 0xe40cb4b49e7f0723b2;
+
+    /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
+    /*                     INTERNAL FUNCTIONS                     */
+    /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
+
+    /// @dev Sets the mothership directly without any emitting event.
+    function _setMothership(address newMothership) internal virtual {
+        /// @solidity memory-safe-assembly
+        assembly {
+            sstore(_MOTHERSHIP_SLOT, shr(96, shl(96, newMothership)))
+        }
+    }
+
+    /// @dev Returns the mothership contract.
+    function mothership() public view virtual returns (address result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            result := sload(_MOTHERSHIP_SLOT)
+        }
+    }
+
+    /// @dev Requires that the caller is the mothership.
+    /// This override affects the `onlyMothership` modifier.
+    function _checkMothership() internal view virtual {
+        if (msg.sender != mothership()) revert CallerNotMothership();
+    }
+
+    /// @dev Requires that the current mothership.
+    modifier onlyMothership() virtual {
+        _checkMothership();
+        _;
+    }
+
+    /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
+    /*                    EXECUTION OPERATIONS                    */
+    /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
+
+    /// @dev Execute a call from this account.
+    /// Reverts and bubbles up error if call fails.
+    /// Returns the result of the call.
+    function execute(address target, uint256 value, bytes calldata data)
+        public
+        payable
+        virtual
+        onlyMothership
+        returns (bytes memory result)
+    {
+        /// @solidity memory-safe-assembly
+        assembly {
+            result := mload(0x40)
+            calldatacopy(result, data.offset, data.length)
+            if iszero(call(gas(), target, value, result, data.length, codesize(), 0x00)) {
+                // Bubble up the revert if the call reverts.
+                returndatacopy(result, 0x00, returndatasize())
+                revert(result, returndatasize())
+            }
+            mstore(result, returndatasize()) // Store the length.
+            let o := add(result, 0x20)
+            returndatacopy(o, 0x00, returndatasize()) // Copy the returndata.
+            mstore(0x40, add(o, returndatasize())) // Allocate the memory.
+        }
+    }
+
+    /// @dev Execute a sequence of calls from this account.
+    /// Reverts and bubbles up error if any call fails.
+    /// Returns the results of the calls.
+    function executeBatch(Call[] calldata calls)
+        public
+        payable
+        virtual
+        onlyMothership
+        returns (bytes[] memory results)
+    {
+        /// @solidity memory-safe-assembly
+        assembly {
+            results := mload(0x40)
+            mstore(results, calls.length)
+            let r := add(0x20, results)
+            let m := add(r, shl(5, calls.length))
+            calldatacopy(r, calls.offset, shl(5, calls.length))
+            for { let end := m } iszero(eq(r, end)) { r := add(r, 0x20) } {
+                let e := add(calls.offset, mload(r))
+                let o := add(e, calldataload(add(e, 0x40)))
+                calldatacopy(m, add(o, 0x20), calldataload(o))
+                // forgefmt: disable-next-item
+                if iszero(call(gas(), calldataload(e), calldataload(add(e, 0x20)),
+                    m, calldataload(o), codesize(), 0x00)) {
+                    // Bubble up the revert if the call reverts.
+                    returndatacopy(m, 0x00, returndatasize())
+                    revert(m, returndatasize())
+                }
+                mstore(r, m) // Append `m` into `results`.
+                mstore(m, returndatasize()) // Store the length,
+                let p := add(m, 0x20)
+                returndatacopy(p, 0x00, returndatasize()) // and copy the returndata.
+                m := add(p, returndatasize()) // Advance `m`.
+            }
+            mstore(0x40, m) // Allocate the memory.
+        }
+    }
+
+    /// @dev Handle token callbacks. Reverts If no token callback is triggered.
+    fallback() external payable virtual override(Receiver) receiverFallback {
+        revert FnSelectorNotRecognized();
+    }
+}

test/Pod.t.sol

@@ -0,0 +1,143 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+import "./utils/SoladyTest.sol";
+import {Pod, MockPod} from "./utils/mocks/MockPod.sol";
+
+contract Target {
+    error TargetError(bytes data);
+
+    bytes32 public datahash;
+
+    bytes public data;
+
+    function setData(bytes memory data_) public payable returns (bytes memory) {
+        data = data_;
+        datahash = keccak256(data_);
+        return data_;
+    }
+
+    function revertWithTargetError(bytes memory data_) public payable {
+        revert TargetError(data_);
+    }
+
+    function changeOwnerSlotValue(bool change) public payable {
+        /// @solidity memory-safe-assembly
+        assembly {
+            if change { sstore(not(0x8b78c6d8), 0x112233) }
+        }
+    }
+}
+
+contract PodTest is SoladyTest {
+    MockPod pod;
+
+    function setUp() public {
+        pod = new MockPod(address(this));
+    }
+
+    function testSetMothership() public {
+        assertEq(pod.mothership(), address(this));
+        pod.setMothership(address(0xABCD));
+        assertEq(pod.mothership(), address(0xABCD));
+    }
+
+    function testExecute() public {
+        vm.deal(address(pod), 1 ether);
+
+        address target = address(new Target());
+        bytes memory data = _randomBytes();
+        pod.execute(target, 123, abi.encodeWithSignature("setData(bytes)", data));
+        assertEq(Target(target).datahash(), keccak256(data));
+        assertEq(target.balance, 123);
+
+        vm.prank(_randomNonZeroAddress());
+        vm.expectRevert(Pod.CallerNotMothership.selector);
+        pod.execute(target, 123, abi.encodeWithSignature("setData(bytes)", data));
+
+        vm.expectRevert(abi.encodeWithSignature("TargetError(bytes)", data));
+        pod.execute(target, 123, abi.encodeWithSignature("revertWithTargetError(bytes)", data));
+    }
+
+    function testExecuteBatch() public {
+        vm.deal(address(pod), 1 ether);
+
+        Pod.Call[] memory calls = new Pod.Call[](2);
+        calls[0].target = address(new Target());
+        calls[1].target = address(new Target());
+        calls[0].value = 123;
+        calls[1].value = 456;
+        calls[0].data = abi.encodeWithSignature("setData(bytes)", _randomBytes(123));
+        calls[1].data = abi.encodeWithSignature("setData(bytes)", _randomBytes(345));
+
+        pod.executeBatch(calls);
+        assertEq(Target(calls[0].target).datahash(), keccak256(_randomBytes(123)));
+        assertEq(Target(calls[1].target).datahash(), keccak256(_randomBytes(345)));
+        assertEq(calls[0].target.balance, 123);
+        assertEq(calls[1].target.balance, 456);
+
+        calls[1].data = abi.encodeWithSignature("revertWithTargetError(bytes)", _randomBytes(111));
+        vm.expectRevert(abi.encodeWithSignature("TargetError(bytes)", _randomBytes(111)));
+        pod.executeBatch(calls);
+    }
+
+    function testExecuteBatch(uint256 r) public {
+        vm.deal(address(pod), 1 ether);
+
+        unchecked {
+            uint256 n = r & 3;
+            Pod.Call[] memory calls = new Pod.Call[](n);
+
+            for (uint256 i; i != n; ++i) {
+                uint256 v = _random() & 0xff;
+                calls[i].target = address(new Target());
+                calls[i].value = v;
+                calls[i].data = abi.encodeWithSignature("setData(bytes)", _randomBytes(v));
+            }
+
+            bytes[] memory results;
+            if (_random() & 1 == 0) {
+                results = pod.executeBatch(_random(), calls);
+            } else {
+                results = pod.executeBatch(calls);
+            }
+
+            assertEq(results.length, n);
+            for (uint256 i; i != n; ++i) {
+                uint256 v = calls[i].value;
+                assertEq(Target(calls[i].target).datahash(), keccak256(_randomBytes(v)));
+                assertEq(calls[i].target.balance, v);
+                assertEq(abi.decode(results[i], (bytes)), _randomBytes(v));
+            }
+        }
+    }
+
+    function testFallback(bytes4 selector) public {
+        if (
+            selector == bytes4(0xf23a6e61) || selector == bytes4(0x150b7a02)
+                || selector == bytes4(0xbc197c81)
+        ) {
+            (, bytes memory rD) = address(pod).call(abi.encodePacked(selector));
+            assertEq(abi.decode(rD, (bytes4)), selector);
+        } else {
+            vm.expectRevert(Pod.FnSelectorNotRecognized.selector);
+            (bool s,) = address(pod).call(abi.encodePacked(selector));
+            s; // suppressed compiler warning
+        }
+    }
+
+    function _randomBytes(uint256 seed) internal pure returns (bytes memory result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            mstore(0x00, seed)
+            let r := keccak256(0x00, 0x20)
+            if lt(byte(2, r), 0x20) {
+                result := mload(0x40)
+                let n := and(r, 0x7f)
+                mstore(result, n)
+                codecopy(add(result, 0x20), byte(1, r), add(n, 0x40))
+                mstore(0x40, add(add(result, 0x40), n))
+            }
+        }
+    }
+}

test/utils/mocks/MockPod.sol

@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+import {Pod} from "../../../src/accounts/Pod.sol";
+import {Brutalizer} from "../Brutalizer.sol";
+
+/// @dev WARNING! This mock is strictly intended for testing purposes only.
+/// Do NOT copy anything here into production code unless you really know what you are doing.
+contract MockPod is Pod, Brutalizer {
+    constructor(address _mothership) {
+        _setMothership(_mothership);
+    }
+
+    function setMothership(address newMothership) external {
+        newMothership = _brutalized(newMothership);
+        _setMothership(newMothership);
+    }
+
+    function executeBatch(uint256 filler, Call[] calldata calls)
+        public
+        payable
+        virtual
+        onlyMothership
+        returns (bytes[] memory results)
+    {
+        _brutalizeMemory();
+        /// @solidity memory-safe-assembly
+        assembly {
+            mstore(0x40, add(mload(0x40), mod(filler, 0x40)))
+        }
+        return super.executeBatch(calls);
+    }
+}