Skip to content

Commit

Permalink
huge refactor sessionKeys and ERC6551
Browse files Browse the repository at this point in the history
  • Loading branch information
@eloi010
eloi010 committed Nov 13, 2023
1 parent dee8158 commit 9c439f2
Show file tree
Hide file tree
Showing 11 changed files with 943 additions and 920 deletions.
59 changes: 10 additions & 49 deletions contracts/core/BaseOpenfortAccount.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -285,52 +285,6 @@ abstract contract BaseOpenfortAccount is
return SIG_VALIDATION_FAILED;
}

/**
* Register a master session key to the account
* @param _key session key to register
* @param _validAfter - this session key is valid only after this timestamp.
* @param _validUntil - this session key is valid only up to this timestamp.
* @notice using this function will automatically set the sessionkey as a
* master session key because no further restriction was set.
* @notice default limit set to 100.
*/
function registerSessionKey(address _key, uint48 _validAfter, uint48 _validUntil) public {
registerSessionKey(_key, _validAfter, _validUntil, DEFAULT_LIMIT);
sessionKeys[_key].masterSessionKey = true;
}

/**
* Register a master session key to the account
* @param _key session key to register
* @param _validAfter - this session key is valid only after this timestamp.
* @param _validUntil - this session key is valid only up to this timestamp.
* @param _limit - limit of uses remaining.
* @notice using this function will automatically set the sessionkey as a
* master session key because no further restriction was set.
*/
function registerSessionKey(address _key, uint48 _validAfter, uint48 _validUntil, uint48 _limit) public {
_requireFromEntryPointOrOwnerorSelf();
sessionKeys[_key].validAfter = _validAfter;
sessionKeys[_key].validUntil = _validUntil;
sessionKeys[_key].limit = _limit;
sessionKeys[_key].masterSessionKey = false;
sessionKeys[_key].whitelising = false;
emit SessionKeyRegistered(_key);
}

/**
* Register a session key to the account
* @param _key session key to register
* @param _validAfter - this session key is valid only after this timestamp.
* @param _validUntil - this session key is valid only up to this timestamp.
* @param _whitelist - this session key can only interact with the addresses in the _whitelist.
*/
function registerSessionKey(address _key, uint48 _validAfter, uint48 _validUntil, address[] calldata _whitelist)
public
{
registerSessionKey(_key, _validAfter, _validUntil, DEFAULT_LIMIT, _whitelist);
}

/**
* Register a session key to the account
* @param _key session key to register
Expand All @@ -350,18 +304,24 @@ abstract contract BaseOpenfortAccount is

// Not sure why changing this for a custom error increases gas dramatically
require(_whitelist.length < 11, "Whitelist too big");
for (uint256 i = 0; i < _whitelist.length;) {
uint256 i = 0;
for (i; i < _whitelist.length;) {
sessionKeys[_key].whitelist[_whitelist[i]] = true;
unchecked {
++i; // gas optimization
}
}
if (i > 0) {
sessionKeys[_key].whitelising = true;
sessionKeys[_key].masterSessionKey = false;
} else {
if (_limit == ((2 ** 48) - 1)) sessionKeys[_key].masterSessionKey = true;
else sessionKeys[_key].masterSessionKey = false;
}

sessionKeys[_key].validAfter = _validAfter;
sessionKeys[_key].validUntil = _validUntil;
sessionKeys[_key].limit = _limit;
sessionKeys[_key].masterSessionKey = false;
sessionKeys[_key].whitelising = true;

emit SessionKeyRegistered(_key);
}
Expand All @@ -374,6 +334,7 @@ abstract contract BaseOpenfortAccount is
_requireFromEntryPointOrOwnerorSelf();
if (sessionKeys[_key].validUntil != 0) {
sessionKeys[_key].validUntil = 0;
sessionKeys[_key].masterSessionKey = false;
emit SessionKeyRevoked(_key);
}
}
Expand Down
92 changes: 31 additions & 61 deletions contracts/core/eip6551/EIP6551OpenfortAccount.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,8 @@ import {IERC1271Upgradeable} from "@openzeppelin/contracts-upgradeable/interface
import {IERC721} from "@openzeppelin/contracts/token/ERC721/IERC721.sol";
import {SafeCastUpgradeable} from "@openzeppelin/contracts-upgradeable/utils/math/SafeCastUpgradeable.sol";

import "./interfaces/IERC6551Account.sol";
import "./utils/Bytecode.sol";
import "erc6551/src/interfaces/IERC6551Account.sol";
import {ERC6551AccountLib} from "erc6551/src/lib/ERC6551AccountLib.sol";
import {BaseAccount, UserOperation, IEntryPoint} from "account-abstraction/core/BaseAccount.sol";
import {TokenCallbackHandler} from "account-abstraction/samples/callback/TokenCallbackHandler.sol";
import "account-abstraction/core/Helpers.sol" as Helpers;
Expand Down Expand Up @@ -45,6 +45,8 @@ contract EIP6551OpenfortAccount is
bytes4 internal constant EXECUTEBATCH_SELECTOR = 0x47e1da2a;
uint48 internal constant DEFAULT_LIMIT = 100;

uint256 public state;

/**
* Struct like ValidationData (from the EIP-4337) - alpha solution - to keep track of session keys' data
* @param validAfter this sessionKey is valid only after this timestamp.
Expand Down Expand Up @@ -93,6 +95,7 @@ contract EIP6551OpenfortAccount is
emit EntryPointUpdated(entrypointContract, _entrypoint);
entrypointContract = _entrypoint;
__EIP712_init("Openfort", "0.4");
state = 1;
}

/**
Expand Down Expand Up @@ -122,11 +125,10 @@ contract EIP6551OpenfortAccount is
}
}

function owner() public view returns (address) {
(uint256 chainId, address tokenContract, uint256 tokenId) = this.token();
function owner() public view virtual returns (address) {
(uint256 chainId, address contractAddress, uint256 tokenId) = token();
if (chainId != block.chainid) return address(0);

return IERC721(tokenContract).ownerOf(tokenId);
return IERC721(contractAddress).ownerOf(tokenId);
}

/**
Expand All @@ -147,11 +149,14 @@ contract EIP6551OpenfortAccount is
revert(add(result, 32), mload(result))
}
}
++state;
}

function token() external view returns (uint256 chainId, address tokenContract, uint256 tokenId) {
uint256 length = address(this).code.length;
return abi.decode(Bytecode.codeAt(address(this), length - 0x60, length), (uint256, address, uint256));
/**
* @dev {See IERC6551Account-token}
*/
function token() public view virtual override returns (uint256, address, uint256) {
return ERC6551AccountLib.token();
}

/**
Expand Down Expand Up @@ -272,6 +277,7 @@ contract EIP6551OpenfortAccount is
function execute(address dest, uint256 value, bytes calldata func) external {
_requireFromEntryPointOrOwner();
_call(dest, value, func);
++state;
}

/**
Expand All @@ -287,6 +293,7 @@ contract EIP6551OpenfortAccount is
unchecked {
++i; // gas optimization
}
++state;
}
}

Expand All @@ -305,6 +312,7 @@ contract EIP6551OpenfortAccount is
*/
function withdrawDepositTo(address payable withdrawAddress, uint256 amount) public {
_requireFromOwner();
++state;
entryPoint().withdrawTo(withdrawAddress, amount);
}

Expand All @@ -318,6 +326,7 @@ contract EIP6551OpenfortAccount is
revert(add(result, 32), mload(result))
}
}
++state;
}

/**
Expand All @@ -344,50 +353,13 @@ contract EIP6551OpenfortAccount is
return SIG_VALIDATION_FAILED;
}

/**
* Register a master session key to the account
* @param _key session key to register
* @param _validAfter - this session key is valid only after this timestamp.
* @param _validUntil - this session key is valid only up to this timestamp.
* @notice using this function will automatically set the sessionkey as a
* master session key because no further restriction was set.
* @notice default limit set to 100.
*/
function registerSessionKey(address _key, uint48 _validAfter, uint48 _validUntil) public {
registerSessionKey(_key, _validAfter, _validUntil, DEFAULT_LIMIT);
sessionKeys[_key].masterSessionKey = true;
}

/**
* Register a master session key to the account
* @param _key session key to register
* @param _validAfter - this session key is valid only after this timestamp.
* @param _validUntil - this session key is valid only up to this timestamp.
* @param _limit - limit of uses remaining.
* @notice using this function will automatically set the sessionkey as a
* master session key because no further restriction was set.
*/
function registerSessionKey(address _key, uint48 _validAfter, uint48 _validUntil, uint48 _limit) public {
_requireFromEntryPointOrOwnerorSelf();
sessionKeys[_key].validAfter = _validAfter;
sessionKeys[_key].validUntil = _validUntil;
sessionKeys[_key].limit = _limit;
sessionKeys[_key].masterSessionKey = false;
sessionKeys[_key].whitelising = false;
emit SessionKeyRegistered(_key);
function isValidSigner(address signer, bytes calldata) external view virtual returns (bytes4) {
if (_isValidSigner(signer)) return IERC6551Account.isValidSigner.selector;
return bytes4(0);
}

/**
* Register a session key to the account
* @param _key session key to register
* @param _validAfter - this session key is valid only after this timestamp.
* @param _validUntil - this session key is valid only up to this timestamp.
* @param _whitelist - this session key can only interact with the addresses in the _whitelist.
*/
function registerSessionKey(address _key, uint48 _validAfter, uint48 _validUntil, address[] calldata _whitelist)
public
{
registerSessionKey(_key, _validAfter, _validUntil, DEFAULT_LIMIT, _whitelist);
function _isValidSigner(address signer) internal view virtual returns (bool) {
return signer == owner();
}

/**
Expand All @@ -409,19 +381,21 @@ contract EIP6551OpenfortAccount is

// Not sure why changing this for a custom error increases gas dramatically
require(_whitelist.length < 11, "Whitelist too big");
for (uint256 i = 0; i < _whitelist.length;) {
uint256 i = 0;
for (i; i < _whitelist.length;) {
sessionKeys[_key].whitelist[_whitelist[i]] = true;
unchecked {
++i; // gas optimization
}
}
if (i != 0) sessionKeys[_key].whitelising = true;
else if (_limit == 2 ** 48 - 1) sessionKeys[_key].masterSessionKey = true;

sessionKeys[_key].validAfter = _validAfter;
sessionKeys[_key].validUntil = _validUntil;
sessionKeys[_key].limit = _limit;
sessionKeys[_key].masterSessionKey = false;
sessionKeys[_key].whitelising = true;

++state;
emit SessionKeyRegistered(_key);
}

Expand All @@ -431,8 +405,10 @@ contract EIP6551OpenfortAccount is
*/
function revokeSessionKey(address _key) external {
_requireFromEntryPointOrOwnerorSelf();
++state;
if (sessionKeys[_key].validUntil != 0) {
sessionKeys[_key].validUntil = 0;
sessionKeys[_key].masterSessionKey = false;
emit SessionKeyRevoked(_key);
}
}
Expand All @@ -456,14 +432,8 @@ contract EIP6551OpenfortAccount is
revert ZeroAddressNotAllowed();
}
_requireFromOwner();
++state;
emit EntryPointUpdated(entrypointContract, _newEntrypoint);
entrypointContract = _newEntrypoint;
}

/**
* Like getNonce() from EIP4337. Needed to comply with EIP6551.
*/
function nonce() external view returns (uint256) {
return getNonce();
}
}
2 changes: 1 addition & 1 deletion contracts/core/upgradeable/UpgradeableOpenfortAccount.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ contract UpgradeableOpenfortAccount is BaseOpenfortAccount, UUPSUpgradeable {
emit EntryPointUpdated(entrypointContract, _entrypoint);
_transferOwnership(_defaultAdmin);
entrypointContract = _entrypoint;
__EIP712_init("Openfort", "0.4");
__EIP712_init("Openfort", "0.5");
}

function _authorizeUpgrade(address) internal override onlyOwner {}
Expand Down
1 change: 1 addition & 0 deletions remappings.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,3 +4,4 @@ ds-test/=lib/forge-std/lib/ds-test/src/
eth-gas-reporter/=node_modules/eth-gas-reporter/
forge-std/=lib/forge-std/src/
hardhat/=node_modules/hardhat/
erc6551/=lib/erc6551/
2 changes: 1 addition & 1 deletion script/UserOpTestCounter.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,7 @@ contract UserOpTestCounter is Script {
function run() public {
vm.startBroadcast(deployPrivKey);

// Verifiy that the counter is still set to 0
// Verify that the counter is still set to 0
assert(testCounter.counters(deployAddress) == 0);
// Count using deployPrivKey
testCounter.count();
Expand Down
Loading

0 comments on commit 9c439f2

Please sign in to comment.