Invalidate sessionKeys when owner changed

openfort-xyz · Nov 15, 2023 · a6997d2 · a6997d2
1 parent be88daa
commit a6997d2
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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);