Skip to content

Commit

Permalink
feat(v2): more flexible access controls (#37)
Browse files Browse the repository at this point in the history
  • Loading branch information
@jaypaik
jaypaik authored Mar 14, 2024
1 parent af7435d commit 2896be7
Show file tree
Hide file tree
Showing 5 changed files with 267 additions and 58 deletions.
2 changes: 1 addition & 1 deletion src/LightAccount.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ contract LightAccount is BaseLightAccount, CustomSlotInitializable {
/// @notice Transfers ownership of the contract to a new account (`newOwner`). Can only be called by the current
/// owner or from the entry point via a user operation signed by the current owner.
/// @param newOwner The new owner.
function transferOwnership(address newOwner) external virtual onlyOwner {
function transferOwnership(address newOwner) external virtual onlyAuthorized {
if (newOwner == address(0) || newOwner == address(this)) {
revert InvalidOwner(newOwner);
}
Expand Down
6 changes: 5 additions & 1 deletion src/MultiOwnerLightAccount.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,11 @@ contract MultiOwnerLightAccount is BaseLightAccount, CustomSlotInitializable {
/// array cannot have 0 or duplicate addresses.
/// @param ownersToAdd The address array of owners to be added.
/// @param ownersToRemove The address array of owners to be removed.
function updateOwners(address[] memory ownersToAdd, address[] memory ownersToRemove) external virtual onlyOwner {
function updateOwners(address[] memory ownersToAdd, address[] memory ownersToRemove)
external
virtual
onlyAuthorized
{
_updateOwners(ownersToAdd, ownersToRemove);
}

Expand Down
42 changes: 17 additions & 25 deletions src/common/BaseLightAccount.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,8 @@ abstract contract BaseLightAccount is BaseAccount, TokenCallbackHandler, UUPSUpg
/// @dev The caller is not authorized.
error NotAuthorized(address caller);

modifier onlyOwner() {
_onlyOwner();
modifier onlyAuthorized() {
_onlyAuthorized();
_;
}

Expand All @@ -33,17 +33,15 @@ abstract contract BaseLightAccount is BaseAccount, TokenCallbackHandler, UUPSUpg
/// @param dest The target of the transaction.
/// @param value The amount of wei sent in the transaction.
/// @param func The transaction's calldata.
function execute(address dest, uint256 value, bytes calldata func) external virtual {
_onlyOwnerOrEntryPoint();
function execute(address dest, uint256 value, bytes calldata func) external virtual onlyAuthorized {
_call(dest, value, func);
}

/// @notice Execute a sequence of transactions.
/// @param dest An array of the targets for each transaction in the sequence.
/// @param func An array of calldata for each transaction in the sequence. Must be the same length as `dest`, with
/// corresponding elements representing the parameters for each transaction.
function executeBatch(address[] calldata dest, bytes[] calldata func) external virtual {
_onlyOwnerOrEntryPoint();
function executeBatch(address[] calldata dest, bytes[] calldata func) external virtual onlyAuthorized {
if (dest.length != func.length) {
revert ArrayLengthMismatch();
}
Expand All @@ -58,8 +56,11 @@ abstract contract BaseLightAccount is BaseAccount, TokenCallbackHandler, UUPSUpg
/// @param value An array of value for each transaction in the sequence.
/// @param func An array of calldata for each transaction in the sequence. Must be the same length as `dest`, with
/// corresponding elements representing the parameters for each transaction.
function executeBatch(address[] calldata dest, uint256[] calldata value, bytes[] calldata func) external virtual {
_onlyOwnerOrEntryPoint();
function executeBatch(address[] calldata dest, uint256[] calldata value, bytes[] calldata func)
external
virtual
onlyAuthorized
{
if (dest.length != func.length || dest.length != value.length) {
revert ArrayLengthMismatch();
}
Expand All @@ -77,7 +78,7 @@ abstract contract BaseLightAccount is BaseAccount, TokenCallbackHandler, UUPSUpg
/// @notice Withdraw value from the account's deposit.
/// @param withdrawAddress Target to send to.
/// @param amount Amount to withdraw.
function withdrawDepositTo(address payable withdrawAddress, uint256 amount) public onlyOwner {
function withdrawDepositTo(address payable withdrawAddress, uint256 amount) public onlyAuthorized {
entryPoint().withdrawTo(withdrawAddress, amount);
}

Expand All @@ -95,20 +96,12 @@ abstract contract BaseLightAccount is BaseAccount, TokenCallbackHandler, UUPSUpg
/// @dev Must override to allow calls to protected functions.
function _isFromOwner() internal view virtual returns (bool);

function _isFromEntryPoint() internal view returns (bool) {
return msg.sender == address(entryPoint());
}

/// @dev Revert if the caller is not an owner or the account itself (when redirected through `execute`).
function _onlyOwner() internal view {
if (msg.sender != address(this) && !_isFromOwner()) {
revert NotAuthorized(msg.sender);
}
}

/// @dev Require that the call is from the entry point or an owner.
function _onlyOwnerOrEntryPoint() internal view {
if (!_isFromEntryPoint() && !_isFromOwner()) {
/// @dev Revert if the caller is not any of:
/// 1. The entry point
/// 2. The account itself (when redirected through `execute`, etc.)
/// 3. An owner
function _onlyAuthorized() internal view {
if (msg.sender != address(entryPoint()) && msg.sender != address(this) && !_isFromOwner()) {
revert NotAuthorized(msg.sender);
}
}
Expand All @@ -122,8 +115,7 @@ abstract contract BaseLightAccount is BaseAccount, TokenCallbackHandler, UUPSUpg
}
}

function _authorizeUpgrade(address newImplementation) internal view override {
function _authorizeUpgrade(address newImplementation) internal view override onlyAuthorized {
(newImplementation);
_onlyOwner();
}
}
134 changes: 119 additions & 15 deletions test/LightAccount.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,8 +59,10 @@ contract LightAccountTest is Test {
}

function testExecuteCanBeCalledByEntryPointWithExternalOwner() public {
PackedUserOperation memory op =
_getSignedOp(address(lightSwitch), abi.encodeCall(LightSwitch.turnOn, ()), EOA_PRIVATE_KEY);
PackedUserOperation memory op = _getSignedOp(
abi.encodeCall(BaseLightAccount.execute, (address(lightSwitch), 0, abi.encodeCall(LightSwitch.turnOn, ()))),
EOA_PRIVATE_KEY
);
PackedUserOperation[] memory ops = new PackedUserOperation[](1);
ops[0] = op;
entryPoint.handleOps(ops, BENEFICIARY);
Expand All @@ -69,7 +71,9 @@ contract LightAccountTest is Test {

function testExecutedCanBeCalledByEntryPointWithContractOwner() public {
_useContractOwner();
PackedUserOperation memory op = _getUnsignedOp(address(lightSwitch), abi.encodeCall(LightSwitch.turnOn, ()));
PackedUserOperation memory op = _getUnsignedOp(
abi.encodeCall(BaseLightAccount.execute, (address(lightSwitch), 0, abi.encodeCall(LightSwitch.turnOn, ())))
);
op.signature = contractOwner.sign(entryPoint.getUserOpHash(op));
PackedUserOperation[] memory ops = new PackedUserOperation[](1);
ops[0] = op;
Expand All @@ -78,7 +82,10 @@ contract LightAccountTest is Test {
}

function testRejectsUserOpsWithInvalidSignature() public {
PackedUserOperation memory op = _getSignedOp(address(lightSwitch), abi.encodeCall(LightSwitch.turnOn, ()), 1234);
PackedUserOperation memory op = _getSignedOp(
abi.encodeCall(BaseLightAccount.execute, (address(lightSwitch), 0, abi.encodeCall(LightSwitch.turnOn, ()))),
1234
);
PackedUserOperation[] memory ops = new PackedUserOperation[](1);
ops[0] = op;
vm.expectRevert(abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error"));
Expand Down Expand Up @@ -171,6 +178,37 @@ contract LightAccountTest is Test {
assertEq(entryPoint.balanceOf(address(account)), 5);
}

function testWithdrawDepositCanBeCalledByEntryPointWithExternalOwner() public {
account.addDeposit{value: 1 ether}();
address payable withdrawalAddress = payable(address(1));

PackedUserOperation memory op =
_getSignedOp(abi.encodeCall(BaseLightAccount.withdrawDepositTo, (withdrawalAddress, 5)), EOA_PRIVATE_KEY);
PackedUserOperation[] memory ops = new PackedUserOperation[](1);
ops[0] = op;
entryPoint.handleOps(ops, BENEFICIARY);

assertEq(withdrawalAddress.balance, 5);
}

function testWithdrawDepositCanBeCalledBySelf() public {
account.addDeposit{value: 1 ether}();
address payable withdrawalAddress = payable(address(1));

PackedUserOperation memory op = _getSignedOp(
abi.encodeCall(
BaseLightAccount.execute,
(address(account), 0, abi.encodeCall(BaseLightAccount.withdrawDepositTo, (withdrawalAddress, 5)))
),
EOA_PRIVATE_KEY
);
PackedUserOperation[] memory ops = new PackedUserOperation[](1);
ops[0] = op;
entryPoint.handleOps(ops, BENEFICIARY);

assertEq(withdrawalAddress.balance, 5);
}

function testWithdrawDepositToCannotBeCalledByRandos() public {
account.addDeposit{value: 10}();
vm.expectRevert(abi.encodeWithSelector(BaseLightAccount.NotAuthorized.selector, (address(this))));
Expand All @@ -189,7 +227,24 @@ contract LightAccountTest is Test {
function testEntryPointCanTransferOwnership() public {
address newOwner = address(0x100);
PackedUserOperation memory op =
_getSignedOp(address(account), abi.encodeCall(LightAccount.transferOwnership, (newOwner)), EOA_PRIVATE_KEY);
_getSignedOp(abi.encodeCall(LightAccount.transferOwnership, (newOwner)), EOA_PRIVATE_KEY);
PackedUserOperation[] memory ops = new PackedUserOperation[](1);
ops[0] = op;
vm.expectEmit(true, true, false, false);
emit OwnershipTransferred(eoaAddress, newOwner);
entryPoint.handleOps(ops, BENEFICIARY);
assertEq(account.owner(), newOwner);
}

function testSelfCanTransferOwnership() public {
address newOwner = address(0x100);
PackedUserOperation memory op = _getSignedOp(
abi.encodeCall(
BaseLightAccount.execute,
(address(account), 0, abi.encodeCall(LightAccount.transferOwnership, (newOwner)))
),
EOA_PRIVATE_KEY
);
PackedUserOperation[] memory ops = new PackedUserOperation[](1);
ops[0] = op;
vm.expectEmit(true, true, false, false);
Expand Down Expand Up @@ -302,10 +357,63 @@ contract LightAccountTest is Test {
// Upgrade to a normal SimpleAccount with a different entry point.
IEntryPoint newEntryPoint = IEntryPoint(address(0x2000));
SimpleAccount newImplementation = new SimpleAccount(newEntryPoint);

vm.prank(eoaAddress);
vm.expectEmit(true, true, false, false);
emit SimpleAccountInitialized(newEntryPoint, address(this));
vm.prank(eoaAddress);
account.upgradeToAndCall(address(newImplementation), abi.encodeCall(SimpleAccount.initialize, (address(this))));

SimpleAccount upgradedAccount = SimpleAccount(payable(account));
assertEq(address(upgradedAccount.entryPoint()), address(newEntryPoint));
}

function testEntryPointCanUpgrade() public {
// Upgrade to a normal SimpleAccount with a different entry point.
IEntryPoint newEntryPoint = IEntryPoint(address(0x2000));
SimpleAccount newImplementation = new SimpleAccount(newEntryPoint);
PackedUserOperation memory op = _getSignedOp(
abi.encodeCall(
account.upgradeToAndCall,
(address(newImplementation), abi.encodeCall(SimpleAccount.initialize, (address(this))))
),
EOA_PRIVATE_KEY
);
PackedUserOperation[] memory ops = new PackedUserOperation[](1);
ops[0] = op;

vm.expectEmit(true, true, false, false);
emit SimpleAccountInitialized(newEntryPoint, address(this));
entryPoint.handleOps(ops, BENEFICIARY);

SimpleAccount upgradedAccount = SimpleAccount(payable(account));
assertEq(address(upgradedAccount.entryPoint()), address(newEntryPoint));
}

function testSelfCanUpgrade() public {
// Upgrade to a normal SimpleAccount with a different entry point.
IEntryPoint newEntryPoint = IEntryPoint(address(0x2000));
SimpleAccount newImplementation = new SimpleAccount(newEntryPoint);
PackedUserOperation memory op = _getSignedOp(
abi.encodeCall(
BaseLightAccount.execute,
(
address(account),
0,
abi.encodeCall(
account.upgradeToAndCall,
(address(newImplementation), abi.encodeCall(SimpleAccount.initialize, (address(this))))
)
)
),
EOA_PRIVATE_KEY
);
PackedUserOperation[] memory ops = new PackedUserOperation[](1);
ops[0] = op;

vm.expectEmit(true, true, false, false);
emit SimpleAccountInitialized(newEntryPoint, address(this));
entryPoint.handleOps(ops, BENEFICIARY);

SimpleAccount upgradedAccount = SimpleAccount(payable(account));
assertEq(address(upgradedAccount.entryPoint()), address(newEntryPoint));
}
Expand Down Expand Up @@ -343,7 +451,7 @@ contract LightAccountTest is Test {
bytes32(uint256(uint160(0x0000000071727De22E5E9d8BAf0edAc6f37da032)))
)
),
0x56aa27383cd945ea0c90683be2dec26099099eb0cb138033e62c3e2e6c613a7c
0x7448a966519c6db16b5148dde40a37c19b5fe204acc51ef0cbfe865ea110a15d
);
}

Expand All @@ -352,11 +460,7 @@ contract LightAccountTest is Test {
account.transferOwnership(address(contractOwner));
}

function _getUnsignedOp(address target, bytes memory innerCallData)
internal
view
returns (PackedUserOperation memory)
{
function _getUnsignedOp(bytes memory callData) internal view returns (PackedUserOperation memory) {
uint128 verificationGasLimit = 1 << 24;
uint128 callGasLimit = 1 << 24;
uint128 maxPriorityFeePerGas = 1 << 8;
Expand All @@ -365,7 +469,7 @@ contract LightAccountTest is Test {
sender: address(account),
nonce: 0,
initCode: "",
callData: abi.encodeCall(BaseLightAccount.execute, (target, 0, innerCallData)),
callData: callData,
accountGasLimits: bytes32(uint256(verificationGasLimit) << 128 | callGasLimit),
preVerificationGas: 1 << 24,
gasFees: bytes32(uint256(maxPriorityFeePerGas) << 128 | maxFeePerGas),
Expand All @@ -374,12 +478,12 @@ contract LightAccountTest is Test {
});
}

function _getSignedOp(address target, bytes memory innerCallData, uint256 privateKey)
function _getSignedOp(bytes memory callData, uint256 privateKey)
internal
view
returns (PackedUserOperation memory)
{
PackedUserOperation memory op = _getUnsignedOp(target, innerCallData);
PackedUserOperation memory op = _getUnsignedOp(callData);
op.signature = _sign(privateKey, entryPoint.getUserOpHash(op).toEthSignedMessageHash());
return op;
}
Expand Down
Loading

0 comments on commit 2896be7

Please sign in to comment.