Session Key Audit Remediations (#8)

* [M-04] * [M-05] * [M-03] - Move check of session counter to init * [L-03] add expired state for sessions * [L-04] prevent managing session keys with a session key. * Add OperationType to distinguish between transaction and signature operations to block 1271 signatures from session keys * [L-03] only override status to expired for active sessions * [M-02]
Abstract-Foundation · Jan 6, 2025 · 5a94951 · 5a94951
1 parent b4c84c2
commit 5a94951
Show file tree

Hide file tree

Showing 15 changed files with 69 additions and 33 deletions.
diff --git a/contracts/AGWAccount.sol b/contracts/AGWAccount.sol
@@ -23,6 +23,7 @@ import {ERC1271Handler} from './handlers/ERC1271Handler.sol';
 import {Call} from './batch/BatchCaller.sol';
 
 import {IAGWAccount} from './interfaces/IAGWAccount.sol';
+import {OperationType} from './interfaces/IValidator.sol';
 import {AccountFactory} from './AccountFactory.sol';
 import {BatchCaller} from './batch/BatchCaller.sol';
 
@@ -272,7 +273,7 @@ contract AGWAccount is
         bool hookSuccess = runValidationHooks(signedHash, transaction, hookData);
 
         // Handle validation
-        bool valid = _handleValidation(validator, signedHash, signature);
+        bool valid = _handleValidation(validator, OperationType.Transaction, signedHash, signature);
 
         magicValue = (hookSuccess && valid) ? ACCOUNT_VALIDATION_SUCCESS_MAGIC : bytes4(0);
     }

diff --git a/contracts/handlers/ERC1271Handler.sol b/contracts/handlers/ERC1271Handler.sol
@@ -6,6 +6,7 @@ import {IERC1271} from '@openzeppelin/contracts/interfaces/IERC1271.sol';
 import {SignatureDecoder} from '../libraries/SignatureDecoder.sol';
 import {ValidationHandler} from './ValidationHandler.sol';
 import {EIP712Upgradeable} from '@openzeppelin/contracts-upgradeable/utils/cryptography/EIP712Upgradeable.sol';
+import {OperationType} from '../interfaces/IValidator.sol';
 
 /**
  * @title ERC1271Handler
@@ -47,7 +48,7 @@ abstract contract ERC1271Handler is
 
         bytes32 eip712Hash = _hashTypedDataV4(_agwMessageHash(AGWMessage(signedHash)));
 
-        bool valid = _handleValidation(validator, eip712Hash, signature);
+        bool valid = _handleValidation(validator, OperationType.Signature, eip712Hash, signature);
 
         magicValue = valid ? _ERC1271_MAGIC : bytes4(0);
     }

diff --git a/contracts/handlers/ValidationHandler.sol b/contracts/handlers/ValidationHandler.sol
@@ -8,6 +8,7 @@ import {ValidatorManager} from '../managers/ValidatorManager.sol';
 
 import {IK1Validator, IR1Validator} from '../interfaces/IValidator.sol';
 import {IModuleValidator} from '../interfaces/IModuleValidator.sol';
+import {OperationType} from '../interfaces/IValidator.sol';
 
 /**
  * @title ValidationHandler
@@ -17,6 +18,7 @@ import {IModuleValidator} from '../interfaces/IModuleValidator.sol';
 abstract contract ValidationHandler is OwnerManager, ValidatorManager {
     function _handleValidation(
         address validator,
+        OperationType operationType,
         bytes32 signedHash,
         bytes memory signature
     ) internal view returns (bool) {
@@ -31,6 +33,7 @@ abstract contract ValidationHandler is OwnerManager, ValidatorManager {
                 bytes32[2] memory pubKey = abi.decode(cursor, (bytes32[2]));
 
                 bool _success = IR1Validator(validator).validateSignature(
+                    operationType,
                     signedHash,
                     signature,
                     pubKey
@@ -44,6 +47,7 @@ abstract contract ValidationHandler is OwnerManager, ValidatorManager {
             }
         } else if (_k1IsValidator(validator)) {
             address recoveredAddress = IK1Validator(validator).validateSignature(
+                operationType,
                 signedHash,
                 signature
             );
@@ -55,8 +59,8 @@ abstract contract ValidationHandler is OwnerManager, ValidatorManager {
             if (OwnerManager._k1IsOwner(recoveredAddress)) {
                 return true;
             }
-        } else if (_isModuleValidator(validator)) {
-            return IModuleValidator(validator).handleValidation(signedHash, signature);
+        } else if ( _isModuleValidator(validator)) {
+            return IModuleValidator(validator).handleValidation(operationType, signedHash, signature);
         }
 
         return false;

diff --git a/contracts/helpers/TimestampAsserterLocator.sol b/contracts/helpers/TimestampAsserterLocator.sol
@@ -8,15 +8,8 @@ library TimestampAsserterLocator {
     if (block.chainid == 260) {
       return ITimestampAsserter(address(0x00000000000000000000000000000000808012));
     }
-    if (block.chainid == 11124) {
-        return ITimestampAsserter(address(0x27570660a298db7373EaA50c1a728DA93b5BC969));
+    else {
+      return ITimestampAsserter(address(0x958F70e4Fd676c9CeAaFe5c48cB78CDD08b4880d));
     }
-    if (block.chainid == 300) {
-      revert("Timestamp asserter is not deployed on ZKsync Sepolia testnet yet");
-    }
-    if (block.chainid == 324) {
-      revert("Timestamp asserter is not deployed on ZKsync mainnet yet");
-    }
-    revert("Timestamp asserter is not deployed on this network");
   }
 }
diff --git a/contracts/interfaces/IModuleValidator.sol b/contracts/interfaces/IModuleValidator.sol
@@ -1,12 +1,14 @@
 // SPDX-License-Identifier: GPL-3.0
 pragma solidity ^0.8.24;
 
+import {OperationType} from './IValidator.sol';
+
 /**
  * @title Modular validator interface for native AA
  * @dev Add signature to module or validate existing signatures for acccount
  */
 interface IModuleValidator {
-  function handleValidation(bytes32 signedHash, bytes memory signature) external view returns (bool);
+  function handleValidation(OperationType operationType, bytes32 signedHash, bytes memory signature) external view returns (bool);
 
   function addValidationKey(bytes memory key) external returns (bool);
 }
diff --git a/contracts/interfaces/IValidator.sol b/contracts/interfaces/IValidator.sol
@@ -3,6 +3,11 @@ pragma solidity ^0.8.17;
 
 import {IERC165} from '@openzeppelin/contracts/utils/introspection/IERC165.sol';
 
+enum OperationType {
+    Signature,
+    Transaction
+}
+
 /**
  * @title secp256r1 ec keys' signature validator interface
  * @author https://getclave.io
@@ -16,6 +21,7 @@ interface IR1Validator is IERC165 {
      * @return valid bool                 - validation result
      */
     function validateSignature(
+        OperationType operationType,
         bytes32 signedHash,
         bytes calldata signature,
         bytes32[2] calldata pubKey
@@ -34,6 +40,7 @@ interface IK1Validator is IERC165 {
      * @return signer address             - recovered signer address
      */
     function validateSignature(
+        OperationType operationType,
         bytes32 signedHash,
         bytes calldata signature
     ) external view returns (address signer);

diff --git a/contracts/libraries/SessionLib.sol b/contracts/libraries/SessionLib.sol
@@ -16,7 +16,8 @@ library SessionLib {
   enum Status {
     NotInitialized,
     Active,
-    Closed
+    Closed,
+    Expired
   }
 
   // This struct is used to track usage information for each session.
@@ -25,6 +26,7 @@ library SessionLib {
   // Storage layout of this struct is weird to conform to ERC-7562 storage access restrictions during validation.
   // Each innermost mapping is always mapping(address account => ...).
   struct SessionStorage {
+    uint256 expiresAt;
     mapping(address => Status) status;
     UsageTracker fee;
     // (target) => transfer value tracker
@@ -110,6 +112,7 @@ library SessionLib {
 
   // Info about remaining session limits and its status
   struct SessionState {
+    uint256 expiresAt;
     Status status;
     uint256 feesRemaining;
     LimitState[] transferValue;
@@ -179,7 +182,7 @@ library SessionLib {
     // most of the computation needed to validate the session.
 
     // TODO: update fee allowance with the gasleft/refund at the end of execution
-    if (transaction.paymaster != 0) {
+    if (transaction.paymaster == 0) {
       // If a paymaster is paying the fee, we don't need to check the fee limit
       uint256 fee = transaction.maxFeePerGas * transaction.gasLimit;
       spec.feeLimit.checkAndUpdate(state.fee, fee, periodId);
@@ -217,6 +220,7 @@ library SessionLib {
       // Disallow self-targeting transactions with session keys as these have the ability to administer
       // the smart account.
       require(address(uint160(transaction.to)) != msg.sender, "Can not target self");
+      require(address(uint160(transaction.to)) != address(this), "Can not manage session keys");
 
       bytes4 selector = bytes4(transaction.data[:4]);
       CallSpec memory callPolicy;
@@ -329,9 +333,17 @@ library SessionLib {
       mstore(callParams, paramLimitIndex)
     }
 
+    Status status = session.status[account];
+    if (status == Status.Active) {
+      if (block.timestamp > session.expiresAt) {
+        status = Status.Expired;
+      }
+    }
+
     return
       SessionState({
-        status: session.status[account],
+        expiresAt: session.expiresAt,
+        status: status,
         feesRemaining: remainingLimit(spec.feeLimit, session.fee, account),
         transferValue: transferValue,
         callValue: callValue,

diff --git a/contracts/managers/ModuleManager.sol b/contracts/managers/ModuleManager.sol
@@ -84,7 +84,6 @@ abstract contract ModuleManager is IModuleManager, Auth {
     }
 
     function _removeModule(address module) internal {
-        _modulesLinkedList().remove(module);
 
         (bool success, ) = module.excessivelySafeCall(
             gasleft(),
@@ -94,6 +93,8 @@ abstract contract ModuleManager is IModuleManager, Auth {
         );
         (success); // silence unused local variable warning
 
+        _modulesLinkedList().remove(module);
+
         emit RemoveModule(module);
     }
 

diff --git a/contracts/test/MockValidator.sol b/contracts/test/MockValidator.sol
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: GPL-3.0
 pragma solidity ^0.8.17;
 
-import {IR1Validator, IERC165} from '../interfaces/IValidator.sol';
+import {IR1Validator, IERC165, OperationType} from '../interfaces/IValidator.sol';
 import {Transaction} from '@matterlabs/zksync-contracts/l2/system-contracts/libraries/TransactionHelper.sol';
 
 /**
@@ -10,6 +10,7 @@ import {Transaction} from '@matterlabs/zksync-contracts/l2/system-contracts/libr
  */
 contract MockValidator is IR1Validator {
     function validateSignature(
+        OperationType, // unused
         bytes32,
         bytes calldata,
         bytes32[2] calldata

diff --git a/contracts/test/PasskeyValidatorTest.sol b/contracts/test/PasskeyValidatorTest.sol
@@ -2,7 +2,7 @@
 pragma solidity ^0.8.17;
 
 import {Base64} from "@openzeppelin/contracts/utils/Base64.sol";
-import {IR1Validator, IERC165} from '../interfaces/IValidator.sol';
+import {IR1Validator, IERC165, OperationType} from '../interfaces/IValidator.sol';
 import {Errors} from '../libraries/Errors.sol';
 import {VerifierCaller} from '../helpers/VerifierCaller.sol';
 
@@ -34,6 +34,7 @@ contract PasskeyValidatorTest is IR1Validator, VerifierCaller {
 
     /// @inheritdoc IR1Validator
     function validateSignature(
+        OperationType, // unused
         bytes32 challenge,
         bytes calldata signature,
         bytes32[2] calldata pubKey

diff --git a/contracts/test/TEEValidatorTest.sol b/contracts/test/TEEValidatorTest.sol
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: GPL-3.0
 pragma solidity ^0.8.17;
 
-import {IR1Validator, IERC165} from '../interfaces/IValidator.sol';
+import {IR1Validator, IERC165, OperationType} from '../interfaces/IValidator.sol';
 import {Errors} from '../libraries/Errors.sol';
 import {VerifierCaller} from '../helpers/VerifierCaller.sol';
 
@@ -23,6 +23,7 @@ contract TEEValidatorTest is IR1Validator, VerifierCaller {
 
     /// @inheritdoc IR1Validator
     function validateSignature(
+        OperationType, // unused
         bytes32 signedHash,
         bytes calldata signature,
         bytes32[2] calldata pubKey

diff --git a/contracts/validators/EOAValidator.sol b/contracts/validators/EOAValidator.sol
@@ -1,9 +1,8 @@
 // SPDX-License-Identifier: GPL-3.0
 pragma solidity ^0.8.17;
 
-import {IK1Validator, IERC165} from '../interfaces/IValidator.sol';
+import {IK1Validator, IERC165, OperationType} from '../interfaces/IValidator.sol';
 import {ECDSA} from '@openzeppelin/contracts/utils/cryptography/ECDSA.sol';
-
 /**
  * @title secp256k1 ec keys' signature validator contract implementing its interface
  * @author https://getclave.io
@@ -14,6 +13,7 @@ contract EOAValidator is IK1Validator {
 
     /// @inheritdoc IK1Validator
     function validateSignature(
+        OperationType operationType,
         bytes32 signedHash,
         bytes calldata signature
     ) external pure override returns (address signer) {

diff --git a/contracts/validators/PasskeyValidator.sol b/contracts/validators/PasskeyValidator.sol
@@ -2,7 +2,7 @@
 pragma solidity ^0.8.17;
 
 import {Base64} from '@openzeppelin/contracts/utils/Base64.sol';
-import {IR1Validator, IERC165} from '../interfaces/IValidator.sol';
+import {IR1Validator, IERC165, OperationType} from '../interfaces/IValidator.sol';
 import {Errors} from '../libraries/Errors.sol';
 import {VerifierCaller} from '../helpers/VerifierCaller.sol';
 
@@ -26,6 +26,7 @@ contract PasskeyValidator is IR1Validator, VerifierCaller {
 
     /// @inheritdoc IR1Validator
     function validateSignature(
+        OperationType, // unused
         bytes32 challenge,
         bytes calldata signature,
         bytes32[2] calldata pubKey

diff --git a/contracts/validators/SessionKeyValidator.sol b/contracts/validators/SessionKeyValidator.sol
@@ -7,7 +7,7 @@ import { IERC165 } from "@openzeppelin/contracts/utils/introspection/IERC165.sol
 import { IModule } from "../interfaces/IModule.sol";
 import { IValidationHook } from "../interfaces/IHook.sol";
 import { IModuleValidator } from "../interfaces/IModuleValidator.sol";
-
+import { OperationType } from "../interfaces/IValidator.sol";
 import { Transaction } from "@matterlabs/zksync-contracts/l2/system-contracts/libraries/TransactionHelper.sol";
 import { ECDSA } from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
 
@@ -36,10 +36,19 @@ contract SessionKeyValidator is IValidationHook, IModuleValidator, IModule {
   }
 
   function sessionStatus(address account, bytes32 sessionHash) external view returns (SessionLib.Status) {
-    return sessions[sessionHash].status[account];
+    SessionLib.Status status = sessions[sessionHash].status[account];
+    if (status == SessionLib.Status.Active) {
+      if (block.timestamp > sessions[sessionHash].expiresAt) {
+        return SessionLib.Status.Expired;
+      }
+    }
+    return status;
   }
 
-  function handleValidation(bytes32 signedHash, bytes memory signature) external view returns (bool) {
+  function handleValidation(OperationType operationType, bytes32 signedHash, bytes memory signature) external view returns (bool) {
+    if (operationType != OperationType.Transaction) {
+      return false;
+    }
     // This only succeeds if the validationHook has previously succeeded for this hash.
     uint256 slot = uint256(signedHash);
     uint256 hookResult;
@@ -67,24 +76,25 @@ contract SessionKeyValidator is IValidationHook, IModuleValidator, IModule {
     require(sessionSpec.feeLimit.limitType != SessionLib.LimitType.Unlimited, "Unlimited fee allowance is not safe");
     sessionCounter[msg.sender]++;
     sessions[sessionHash].status[msg.sender] = SessionLib.Status.Active;
+    sessions[sessionHash].expiresAt = sessionSpec.expiresAt;
     emit SessionCreated(msg.sender, sessionHash, sessionSpec);
   }
 
   function init(bytes calldata data) external {
     // to prevent recursion, since addHook also calls init
     if (!_isInitialized(msg.sender)) {
+      // Ensure that all keys are revoked before installing the module again.
+      // This is to prevent the module from being installed, used and installed
+      // again later with dormant keys that could be used to execute transactions.
+      require(sessionCounter[msg.sender] == 0, "Revoke all keys first");
+
       IHookManager(msg.sender).addHook(abi.encodePacked(address(this)), true);
       IValidatorManager(msg.sender).addModuleValidator(address(this), data);
     }
   }
 
   function disable() external {
     if (_isInitialized(msg.sender)) {
-      // Here we have to revoke all keys, so that if the module
-      // is installed again later, there will be no active sessions from the past.
-      // Problem: if there are too many keys, this will run out of gas.
-      // Solution: before uninstalling, require that all keys are revoked manually.
-      require(sessionCounter[msg.sender] == 0, "Revoke all keys first");
       IValidatorManager(msg.sender).removeModuleValidator(address(this));
       IHookManager(msg.sender).removeHook(address(this), true);
     }

diff --git a/contracts/validators/TEEValidator.sol b/contracts/validators/TEEValidator.sol
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: GPL-3.0
 pragma solidity ^0.8.17;
 
-import {IR1Validator, IERC165} from '../interfaces/IValidator.sol';
+import {IR1Validator, IERC165, OperationType} from '../interfaces/IValidator.sol';
 import {Errors} from '../libraries/Errors.sol';
 import {VerifierCaller} from '../helpers/VerifierCaller.sol';
 
@@ -15,6 +15,7 @@ contract TEEValidator is IR1Validator, VerifierCaller {
 
     /// @inheritdoc IR1Validator
     function validateSignature(
+        OperationType, // unused
         bytes32 signedHash,
         bytes calldata signature,
         bytes32[2] calldata pubKey