diff --git a/contracts/core/BaseOpenfortAccount.sol b/contracts/core/BaseOpenfortAccount.sol index dd6b58c..58da9ba 100644 --- a/contracts/core/BaseOpenfortAccount.sol +++ b/contracts/core/BaseOpenfortAccount.sol @@ -55,6 +55,7 @@ abstract contract BaseOpenfortAccount is bool masterSessionKey; bool whitelising; mapping(address => bool) whitelist; + address registrarAddress; } mapping(address => SessionKeyStruct) public sessionKeys; @@ -111,6 +112,9 @@ abstract contract BaseOpenfortAccount is // If not owner and the session key is revoked, return false if (sessionKey.validUntil == 0) return false; + // If the sessionKey was not registered by the owner, return false + if (sessionKey.registrarAddress != owner()) return false; + // If the signer is a session key that is still valid // Let's first get the selector of the function that the caller is using bytes4 funcSelector = @@ -189,7 +193,9 @@ abstract contract BaseOpenfortAccount is ) { return 0xffffffff; } // Not owner or session key revoked - else { + else if (sessionKey.registrarAddress != owner()) { + return 0xffffffff; + } else { return MAGICVALUE; } } @@ -292,7 +298,6 @@ abstract contract BaseOpenfortAccount is ) public { _requireFromEntryPointOrOwner(); - // Not sure why changing this for a custom error increases gas dramatically require(_whitelist.length < 11, "Whitelist too big"); uint256 i = 0; for (i; i < _whitelist.length;) { @@ -302,6 +307,7 @@ abstract contract BaseOpenfortAccount is } } if (i > 0) { + // If there was some whitelisting, it is not a masterSessionKey sessionKeys[_key].whitelising = true; sessionKeys[_key].masterSessionKey = false; } else { @@ -312,6 +318,7 @@ abstract contract BaseOpenfortAccount is sessionKeys[_key].validAfter = _validAfter; sessionKeys[_key].validUntil = _validUntil; sessionKeys[_key].limit = _limit; + sessionKeys[_key].registrarAddress = owner(); emit SessionKeyRegistered(_key); } @@ -325,6 +332,7 @@ abstract contract BaseOpenfortAccount is if (sessionKeys[_key].validUntil != 0) { sessionKeys[_key].validUntil = 0; sessionKeys[_key].masterSessionKey = false; + sessionKeys[_key].registrarAddress = address(0); emit SessionKeyRevoked(_key); } } diff --git a/test/foundry/core/managed/ManagedOpenfortAccountTest.t.sol b/test/foundry/core/managed/ManagedOpenfortAccountTest.t.sol index 68e55ce..bca7926 100644 --- a/test/foundry/core/managed/ManagedOpenfortAccountTest.t.sol +++ b/test/foundry/core/managed/ManagedOpenfortAccountTest.t.sol @@ -700,7 +700,7 @@ contract ManagedOpenfortAccountTest is Test { // Verify that the registered key is not a MasterKey but has whitelisting bool isMasterKey; bool isWhitelisted; - (,,, isMasterKey, isWhitelisted) = ManagedOpenfortAccount(payable(account)).sessionKeys(sessionKey); + (,,, isMasterKey, isWhitelisted,) = ManagedOpenfortAccount(payable(account)).sessionKeys(sessionKey); assert(!isMasterKey); assert(isWhitelisted); @@ -746,7 +746,7 @@ contract ManagedOpenfortAccountTest is Test { // Verify that the registered key is not a MasterKey but has whitelisting bool isMasterKey; bool isWhitelisted; - (,,, isMasterKey, isWhitelisted) = ManagedOpenfortAccount(payable(account)).sessionKeys(sessionKey); + (,,, isMasterKey, isWhitelisted,) = ManagedOpenfortAccount(payable(account)).sessionKeys(sessionKey); assert(!isMasterKey); assert(isWhitelisted); diff --git a/test/foundry/core/recoverable/RecoverableOpenfortAccountTest.t.sol b/test/foundry/core/recoverable/RecoverableOpenfortAccountTest.t.sol index 06a1aca..3c3a9d3 100644 --- a/test/foundry/core/recoverable/RecoverableOpenfortAccountTest.t.sol +++ b/test/foundry/core/recoverable/RecoverableOpenfortAccountTest.t.sol @@ -776,7 +776,7 @@ contract RecoverableOpenfortAccountTest is Test { // Verify that the registered key is not a MasterKey but has whitelisting bool isMasterKey; bool isWhitelisted; - (,,, isMasterKey, isWhitelisted) = RecoverableOpenfortAccount(payable(account)).sessionKeys(sessionKey); + (,,, isMasterKey, isWhitelisted,) = RecoverableOpenfortAccount(payable(account)).sessionKeys(sessionKey); assert(!isMasterKey); assert(isWhitelisted); diff --git a/test/foundry/core/upgradeable/UpgradeableOpenfortAccountTest.t.sol b/test/foundry/core/upgradeable/UpgradeableOpenfortAccountTest.t.sol index e6a06a5..e2f556f 100644 --- a/test/foundry/core/upgradeable/UpgradeableOpenfortAccountTest.t.sol +++ b/test/foundry/core/upgradeable/UpgradeableOpenfortAccountTest.t.sol @@ -507,7 +507,7 @@ contract UpgradeableOpenfortAccountTest is Test { // Verify that the registered key is not a MasterKey nor has whitelisting bool isMasterKey; bool isWhitelisted; - (,,, isMasterKey, isWhitelisted) = UpgradeableOpenfortAccount(payable(account)).sessionKeys(sessionKey); + (,,, isMasterKey, isWhitelisted,) = UpgradeableOpenfortAccount(payable(account)).sessionKeys(sessionKey); assert(!isMasterKey); assert(!isWhitelisted); @@ -774,7 +774,7 @@ contract UpgradeableOpenfortAccountTest is Test { // Verify that the registered key is not a MasterKey but has whitelisting bool isMasterKey; bool isWhitelisted; - (,,, isMasterKey, isWhitelisted) = UpgradeableOpenfortAccount(payable(account)).sessionKeys(sessionKey); + (,,, isMasterKey, isWhitelisted,) = UpgradeableOpenfortAccount(payable(account)).sessionKeys(sessionKey); assert(!isMasterKey); assert(isWhitelisted);