Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

✨ Minimal Pod #1087

Merged
merged 4 commits into from
Sep 19, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ The Solidity smart contracts are located in the `src` directory.
```ml
accounts
├─ Receiver — "Receiver mixin for ETH and safe-transferred ERC721 and ERC1155 tokens"
├─ Pod — "Minimal account to be spawned and controlled by a mothership"
├─ ERC1271 — "ERC1271 mixin with nested EIP-712 approach"
├─ ERC4337 — "Simple ERC4337 account implementation"
├─ ERC4337Factory — "Simple ERC4337 account factory implementation"
Expand Down
1 change: 1 addition & 0 deletions src/Milady.sol
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import "./accounts/ERC4337Factory.sol";
import "./accounts/ERC6551.sol";
import "./accounts/ERC6551Proxy.sol";
import "./accounts/LibERC6551.sol";
import "./accounts/Pod.sol";
import "./accounts/Receiver.sol";
import "./auth/Ownable.sol";
import "./auth/OwnableRoles.sol";
Expand Down
166 changes: 166 additions & 0 deletions src/accounts/Pod.sol
Original file line number Diff line number Diff line change
@@ -0,0 +1,166 @@
// SPDX-License-Identifier: MIT
pragma solidity ^0.8.4;

import {Receiver} from "./Receiver.sol";

/// @notice Minimal account to be spawned and controlled by a mothership.
/// @author Solady (https://github.com/vectorized/solady/blob/main/src/accounts/Pod.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();

/// @dev The mothership is already been initialized.
error MothershipAlreadyInitialized();

/*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
/* CONSTANTS AND IMMUTABLES */
/*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/

/// @dev The Pod `Mothership` slot is given by:
/// `uint72(bytes9(keccak256("_MOTHERSHIP_SLOT")))`.
uint256 internal constant _MOTHERSHIP_SLOT = 0xe40cb4b49e7f0723b2;

/*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
/* MOTHERSHIP OPERATIONS */
/*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/

/// @dev Returns the mothership contract.
function mothership() public view virtual returns (address result) {
/// @solidity memory-safe-assembly
assembly {
result := shr(96, sload(_MOTHERSHIP_SLOT))
}
}

/// @dev Sets the mothership directly without any emitting event.
/// Call this function in the initializer or constructor, if any.
/// Reverts if the mothership has already been initialized.
function _initializeMothership(address initialMothership) internal virtual {
/// @solidity memory-safe-assembly
assembly {
let s := _MOTHERSHIP_SLOT
if sload(s) {
mstore(0x00, 0xcc62e56e) // `MothershipAlreadyInitialized()`.
revert(0x1c, 0x04)
}
sstore(s, or(s, shl(96, initialMothership)))
}
}

/// @dev Sets the mothership directly without any emitting event.
/// Expose this is a guarded public function if needed.
function _setMothership(address newMothership) internal virtual {
/// @solidity memory-safe-assembly
assembly {
let s := _MOTHERSHIP_SLOT
sstore(s, or(s, shl(96, newMothership)))
}
}

/// @dev Requires that the caller is the mothership.
/// This is called in the `onlyMothership` modifier.
function _checkMothership() internal view virtual {
if (msg.sender != mothership()) revert CallerNotMothership();
}

/// @dev Requires that the caller is the 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.
}
}

/*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
/* OVERRIDES */
/*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/

/// @dev Handle token callbacks. Reverts If no token callback is triggered.
fallback() external payable virtual override(Receiver) receiverFallback {
revert FnSelectorNotRecognized();
}
}
149 changes: 149 additions & 0 deletions test/Pod.t.sol
Original file line number Diff line number Diff line change
@@ -0,0 +1,149 @@
// 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();
pod.initializeMothership(address(this));
}

function testSetMothership() public {
assertEq(pod.mothership(), address(this));
pod.setMothership(address(0xABCD));
assertEq(pod.mothership(), address(0xABCD));
}

function testInitializeMothership() public {
vm.expectRevert(Pod.MothershipAlreadyInitialized.selector);
pod.initializeMothership(address(this));
}

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))
}
}
}
}
32 changes: 32 additions & 0 deletions test/utils/mocks/MockPod.sol
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
// 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 {
function initializeMothership(address initialMothership) public {
_initializeMothership(_brutalized(initialMothership));
}

function setMothership(address newMothership) external {
_setMothership(_brutalized(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);
}
}