docs: add nat-spec doc comments

cspell-config/cspell-misc.txt

@@ -11,6 +11,7 @@ testid
 vueuse
 dockerized
 ethereum
+sepolia
 
 // examples/bank-demo
 ctap

src/SsoBeacon.sol

@@ -6,6 +6,9 @@ import { UpgradeableBeacon } from "@openzeppelin/contracts/proxy/beacon/Upgradea
 /// @title SsoBeacon
 /// @author Matter Labs
 /// @custom:security-contact [email protected]
+/// @dev This beacon stores the implementation address of SsoAccount contract,
+/// which every SSO account delegates to. This beacon's address is immutably stored
+/// in AAFactory contract, as it is required for deploying new SSO accounts.
 contract SsoBeacon is UpgradeableBeacon {
   constructor(address _implementation) UpgradeableBeacon(_implementation) {}
 }

src/helpers/TimestampAsserterLocator.sol

@@ -3,14 +3,22 @@ pragma solidity ^0.8.24;
 
 import { ITimestampAsserter } from "../interfaces/ITimestampAsserter.sol";
 
+/// @title Timestamp asserter locator
+/// @author Matter Labs
+/// @custom:security-contact [email protected]
+/// @notice This library is used to locate the TimestampAsserter contract on different networks.
+/// @dev Might be removed in the future, when TimestampAsserter is deployed via create2 to the same address on all networks.
 library TimestampAsserterLocator {
   function locate() internal view returns (ITimestampAsserter) {
+    // anvil-zksync (era-test-node)
     if (block.chainid == 260) {
       return ITimestampAsserter(address(0x00000000000000000000000000000000808012));
     }
+    // era sepolia testnet
     if (block.chainid == 300) {
       return ITimestampAsserter(address(0xa64EC71Ee812ac62923c85cf0796aA58573c4Cf3));
     }
+    // era mainnet
     if (block.chainid == 324) {
       revert("Timestamp asserter is not deployed on ZKsync mainnet yet");
     }

src/helpers/TokenCallbackHandler.sol

@@ -6,8 +6,9 @@ import { IERC721Receiver } from "@openzeppelin/contracts/token/ERC721/IERC721Rec
 import { IERC1155Receiver } from "@openzeppelin/contracts/token/ERC1155/IERC1155Receiver.sol";
 
 /**
- * Token callback handler.
- *   Handles supported tokens' callbacks, allowing account receiving these tokens.
+ * @title Token callback handler
+ * @notice Contract to handle ERC721 and ERC1155 token callbacks
+ * @author https://getclave.io
  */
 contract TokenCallbackHandler is IERC721Receiver, IERC1155Receiver {
   function onERC721Received(address, address, uint256, bytes calldata) external pure override returns (bytes4) {

src/interfaces/IHook.sol

@@ -5,15 +5,26 @@ import { Transaction } from "@matterlabs/zksync-contracts/l2/system-contracts/li
 import { IERC165 } from "@openzeppelin/contracts/utils/introspection/IERC165.sol";
 import { IModule } from "./IModule.sol";
 
-// Validation hooks are a non-standard way to always perform validation,
-// They can't expect any specific transaction data or signature, but can be used to enforce
-// additional restrictions on the account during the validation phase
+/// @title Validation hook interface for native AA
+/// @author getclave.io
+/// @notice Validation hooks trigger before each transaction,
+/// can be used to enforce additional restrictions on the account and/or transaction during the validation phase.
 interface IValidationHook is IModule, IERC165 {
+  /// @notice Hook that triggers before each transaction during the validation phase.
+  /// @param signedHash Hash of the transaction that is being validated.
+  /// @param transaction Transaction that is being validated.
+  /// @dev If reverts, the transaction is rejected from the mempool and not included in the block.
   function validationHook(bytes32 signedHash, Transaction calldata transaction) external;
 }
 
+/// @title Execution hook interface for native AA
+/// @author getclave.io
+/// @notice Execution hooks trigger before and after each transaction, during the execution phase.
 interface IExecutionHook is IModule, IERC165 {
-  function preExecutionHook(Transaction calldata transaction) external returns (bytes memory context);
+  /// @notice Hook that triggers before each transaction during the execution phase.
+  /// @param transaction Transaction that is being executed.
+  function preExecutionHook(Transaction calldata transaction) external;
 
+  /// @notice Hook that triggers after each transaction during the execution phase.
   function postExecutionHook() external;
 }

src/interfaces/IModule.sol

@@ -1,3 +1,7 @@
+/**
+ * @title Module interface
+ * @dev Interface for a module that can be added to SSO account (e.g. hook or validator).
+ */
 interface IModule {
   /**
    * @dev This function is called by the smart account during installation of the module

src/interfaces/IModuleValidator.sol

@@ -7,16 +7,28 @@ import { IModule } from "./IModule.sol";
 
 /**
  * @title Modular validator interface for native AA
- * @dev Add signature to module or validate existing signatures for acccount
+ * @notice Validators are used for custom validation of transactions and/or message signatures
  */
 interface IModuleValidator is IModule, IERC165 {
+  /**
+   * @notice Validate transaction for account
+   * @param signedHash Hash of the transaction
+   * @param signature Signature for the transaction
+   * @param transaction Transaction to validate
+   * @return bool True if transaction is valid
+   */
   function validateTransaction(
     bytes32 signedHash,
-    bytes memory signature,
+    bytes calldata signature,
     Transaction calldata transaction
   ) external returns (bool);
 
-  function validateSignature(bytes32 signedHash, bytes memory signature) external view returns (bool);
-
-  function addValidationKey(bytes memory key) external returns (bool);
+  /**
+   * @notice Validate signature for account (including via EIP-1271)
+   * @dev If module is not supposed to validate signatures, it MUST return false
+   * @param signedHash Hash of the message
+   * @param signature Signature of the message
+   * @return bool True if signature is valid
+   */
+  function validateSignature(bytes32 signedHash, bytes calldata signature) external view returns (bool);
 }

src/interfaces/ITimestampAsserter.sol

@@ -1,6 +1,10 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.24;
 
+/// @title Timestamp asserter interface
+/// @author Matter Labs
+/// @custom:security-contact [email protected]
+/// @notice Used to assert that the current timestamp is within a given range in AA validation context.
 interface ITimestampAsserter {
   function assertTimestampInRange(uint256 start, uint256 end) external view;
 }

src/libraries/Errors.sol

@@ -1,56 +1,25 @@
 // SPDX-License-Identifier: GPL-3.0
 pragma solidity ^0.8.24;
 
+/// @title Errors
+/// @notice Errors used by ZKsync SSO and its components
+/// @author getclave.io
 library Errors {
-  /*//////////////////////////////////////////////////////////////
-                               ACCOUNT
-    //////////////////////////////////////////////////////////////*/
-
+  // Account errors
   error INSUFFICIENT_FUNDS(uint256 required, uint256 available);
   error FEE_PAYMENT_FAILED();
   error METHOD_NOT_IMPLEMENTED();
 
-  /*//////////////////////////////////////////////////////////////
-                               LINKED LIST
-    //////////////////////////////////////////////////////////////*/
-
-  error INVALID_PREV_BYTES(bytes prevValue, bytes oldValue);
-  error INVALID_PREV_ADDR(address prevValue, address oldValue);
-  // Bytes
-  error INVALID_BYTES(uint256 length);
-  error BYTES_ALREADY_EXISTS(bytes length);
-  error BYTES_NOT_EXISTS(bytes lookup);
-  // Address
-  error INVALID_ADDRESS(address valid);
-  error ADDRESS_ALREADY_EXISTS(address exists);
-  error ADDRESS_NOT_EXISTS(address notExists);
-
-  /*//////////////////////////////////////////////////////////////
-                             VALIDATOR MANAGER
-    //////////////////////////////////////////////////////////////*/
-
+  // ERC165 module errors
   error VALIDATOR_ERC165_FAIL(address validator);
-
-  /*//////////////////////////////////////////////////////////////
-                              HOOK MANAGER
-    //////////////////////////////////////////////////////////////*/
-
-  error EMPTY_HOOK_ADDRESS(uint256 hookAndDataLength);
   error HOOK_ERC165_FAIL(address hookAddress, bool isValidation);
-  error INVALID_KEY(bytes32 key);
-
-  /*//////////////////////////////////////////////////////////////
-                              AUTH
-    //////////////////////////////////////////////////////////////*/
 
+  // Auth errors
   error NOT_FROM_BOOTLOADER(address notBootloader);
   error NOT_FROM_HOOK(address notHook);
   error NOT_FROM_SELF(address notSelf);
 
-  /*//////////////////////////////////////////////////////////////
-                            BatchCaller
-    //////////////////////////////////////////////////////////////*/
-
+  // Batch caller errors
   error CALL_FAILED(uint256 batchCallIndex);
   error MSG_VALUE_MISMATCH(uint256 actualValue, uint256 expectedValue);
 }

src/libraries/SessionLib.sol

@@ -1,4 +1,4 @@
-// SPDX-License-Identifier: GPL-3.0
+// SPDX-License-Identifier: MIT
 pragma solidity ^0.8.24;
 
 import { Transaction } from "@matterlabs/zksync-contracts/l2/system-contracts/libraries/TransactionHelper.sol";
@@ -7,26 +7,29 @@ import { TimestampAsserterLocator } from "../helpers/TimestampAsserterLocator.so
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import { LibBytes } from "solady/src/utils/LibBytes.sol";
 
+/// @title Session Library
+/// @author Matter Labs
+/// @notice Library for session management, used by SessionKeyValidator
+/// @custom:security-contact [email protected]
 library SessionLib {
   using SessionLib for SessionLib.Constraint;
   using SessionLib for SessionLib.UsageLimit;
   using LibBytes for bytes;
 
-  // We do not permit session keys to be reused to open multiple sessions
-  // (after one expires or is closed, e.g.).
-  // For each session key, its session status can only be changed
-  // from NotInitialized to Active, and from Active to Closed.
+  /// @notice We do not permit opening multiple identical sessions (even after one is closed, e.g.).
+  /// For each session key, its session status can only be changed
+  /// from NotInitialized to Active, and from Active to Closed.
   enum Status {
     NotInitialized,
     Active,
     Closed
   }
 
-  // This struct is used to track usage information for each session.
-  // Along with `status`, this is considered the session state.
-  // While everything else is considered the session spec, and is stored offchain.
-  // 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 => ...).
+  /// @notice This struct is used to track usage information for each session.
+  /// Along with `status`, this is considered the session state.
+  /// While everything else is considered the session spec, and is stored offchain.
+  /// @dev 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 {
     mapping(address => Status) status;
     UsageTracker fee;
@@ -76,6 +79,10 @@ library SessionLib {
     NotEqual
   }
 
+  /// @notice This struct is provided by the account to create a session.
+  /// It is used to define the session's policies, limits and constraints.
+  /// Only its hash is stored onchain, and the full struct is provided with
+  /// each transaction in calldata via `validatorData`, encoded in the signature.
   struct SessionSpec {
     address signer;
     uint256 expiresAt;
@@ -120,6 +127,12 @@ library SessionLib {
     LimitState[] callParams;
   }
 
+  /// @notice Checks if the limit is exceeded and updates the usage tracker.
+  /// @param limit The limit to check.
+  /// @param tracker The usage tracker to update.
+  /// @param value The tracked value to check the limit against.
+  /// @param period The period ID to check the limit against. Ignored if the limit is not of type Allowance.
+  /// @dev Reverts if the limit is exceeded or the period is invalid.
   function checkAndUpdate(
     UsageLimit memory limit,
     UsageTracker storage tracker,
@@ -136,6 +149,13 @@ library SessionLib {
     }
   }
 
+  /// @notice Checks if the constraint is met and update the usage tracker.
+  /// @param constraint The constraint to check.
+  /// @param tracker The usage tracker to update.
+  /// @param data The transaction data to check the constraint against.
+  /// @param period The period ID to check the allowances against.
+  /// @dev Reverts if the constraint is not met.
+  /// @dev Forwards the call to `checkAndUpdate(limit, ...)` on the limit of the constraint.
   function checkAndUpdate(
     Constraint memory constraint,
     UsageTracker storage tracker,
@@ -164,6 +184,15 @@ library SessionLib {
     constraint.limit.checkAndUpdate(tracker, uint256(param), period);
   }
 
+  /// @notice Finds the call policy, checks if it is violated and updates the usage trackers.
+  /// @param state The session storage to update.
+  /// @param data The transaction data to check the call policy against.
+  /// @param target The target address of the call.
+  /// @param selector The 4-byte selector of the call.
+  /// @param callPolicies The call policies to search through.
+  /// @param periodIds The period IDs to check the allowances against. The length has to be at least `periodIdsOffset + callPolicies.length`.
+  /// @param periodIdsOffset The offset in the `periodIds` array to start checking the constraints.
+  /// @return The call policy that was found, reverts if not found or if the call is not allowed.
   function checkCallPolicy(
     SessionStorage storage state,
     bytes memory data,
@@ -193,22 +222,28 @@ library SessionLib {
     return callPolicy;
   }
 
+  /// @notice Validates the fee limit of the session and updates the tracker.
+  /// Only performs the checks if the transaction is not using a paymaster.
+  /// @param state The session storage to update.
+  /// @param transaction The transaction to check the fee of.
+  /// @param spec The session spec to check the fee limit against.
+  /// @param periodId The period ID to check the fee limit against. Ignored if the limit is not of type Allowance.
+  /// @dev Reverts if the fee limit is exceeded.
+  /// @dev This is split from `validate` to prevent gas estimation failures.
+  /// When this check was part of `validate`, gas estimation could fail due to
+  /// fee limit being smaller than the upper bound of the gas estimation binary search.
+  /// By splitting this check, we can now have this order of operations in `validateTransaction`:
+  /// 1. session.validate()
+  /// 2. ECDSA.tryRecover()
+  /// 3. session.validateFeeLimit()
+  /// This way, gas estimation will exit on step 2 instead of failing, but will still run through
+  /// most of the computation needed to validate the session.
   function validateFeeLimit(
     SessionStorage storage state,
     Transaction calldata transaction,
     SessionSpec memory spec,
     uint64 periodId
   ) internal {
-    // This is split from `validate` to prevent gas estimation failures.
-    // When this check was part of `validate`, gas estimation could fail due to
-    // fee limit being smaller than the upper bound of the gas estimation binary search.
-    // By splitting this check, we can now have this order of operations in `validateTransaction`:
-    // 1. session.validate()
-    // 2. ECDSA.tryRecover()
-    // 3. session.validateFeeLimit()
-    // This way, gas estimation will exit on step 2 instead of failing, but will still run through
-    // most of the computation needed to validate the session.
-
     // TODO: update fee allowance with the gasleft/refund at the end of execution
     // If a paymaster is paying the fee, we don't need to check the fee limit
     if (transaction.paymaster == 0) {
@@ -217,22 +252,25 @@ library SessionLib {
     }
   }
 
+  /// @notice Validates the transaction against the session spec and updates the usage trackers.
+  /// @param state The session storage to update.
+  /// @param transaction The transaction to validate.
+  /// @param spec The session spec to validate against.
+  /// @param periodIds The period IDs to check the allowances against.
+  /// @dev periodId is defined as block.timestamp / limit.period if limitType == Allowance, and 0 otherwise (which will be ignored).
+  /// periodIds[0] is for fee limit (not used in this function),
+  /// periodIds[1] is for value limit,
+  /// peroidIds[2:2+n] are for `ERC20.approve()` constraints, where `n` is the number of constraints in the `ERC20.approve()` policy
+  ///   if an approval-based paymaster is used, 0 otherwise.
+  /// periodIds[2+n:] are for call constraints, if there are any.
+  /// It is required to pass them in (instead of computing via block.timestamp) since during validation
+  /// we can only assert the range of the timestamp, but not access its value.
   function validate(
     SessionStorage storage state,
     Transaction calldata transaction,
     SessionSpec memory spec,
     uint64[] memory periodIds
   ) internal {
-    // Here we additionally pass uint64[] periodId to check allowance limits
-    // periodId is defined as block.timestamp / limit.period if limitType == Allowance, and 0 otherwise (which will be ignored).
-    // periodIds[0] is for fee limit (not used in this function),
-    // periodIds[1] is for value limit,
-    // peroidIds[2:2+n] are for `ERC20.approve()` constraints, if an approval-based paymaster is used
-    //   where `n` is the number of constraints in the `ERC20.approve()` policy if an approval-based paymaster is used, 0 otherwise.
-    // periodIds[2+n:] are for call constraints, if there are any.
-    // It is required to pass them in (instead of computing via block.timestamp) since during validation
-    // we can only assert the range of the timestamp, but not access its value.
-
     require(state.status[msg.sender] == Status.Active, "Session is not active");
     TimestampAsserterLocator.locate().assertTimestampInRange(0, spec.expiresAt);
 
@@ -295,6 +333,11 @@ library SessionLib {
     }
   }
 
+  /// @notice Getter for the remainder of a usage limit.
+  /// @param limit The limit to check.
+  /// @param tracker The corresponding usage tracker to get the usage from.
+  /// @param account The account to get the usage for.
+  /// @return The remaining limit. If unlimited, returns `type(uint256).max`.
   function remainingLimit(
     UsageLimit memory limit,
     UsageTracker storage tracker,
@@ -314,6 +357,11 @@ library SessionLib {
     }
   }
 
+  /// @notice Getter for the session state.
+  /// @param session The session storage to get the state from.
+  /// @param account The account to get the state for.
+  /// @param spec The session spec to get the state for.
+  /// @return The session state: status, remaining fee limit, transfer limits, call value and call parameter limits.
   function getState(
     SessionStorage storage session,
     address account,