draft with notes

contracts/smart-account/modules/SessionKeyManagers/DANSessionKeyManagerModule.sol

@@ -0,0 +1,159 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.23;
+
+import {BaseAuthorizationModule} from "../BaseAuthorizationModule.sol";
+import {MerkleProof} from "@openzeppelin/contracts/utils/cryptography/MerkleProof.sol";
+import {_packValidationData} from "@account-abstraction/contracts/core/Helpers.sol";
+import {UserOperation} from "@account-abstraction/contracts/interfaces/UserOperation.sol";
+import {ISessionValidationModule} from "../../interfaces/modules/ISessionValidationModule.sol";
+import {ISessionKeyManagerModule} from "../../interfaces/modules/SessionKeyManagers/ISessionKeyManagerModule.sol";
+import {IAuthorizationModule} from "../../interfaces/IAuthorizationModule.sol";
+import {ISignatureValidator} from "../../interfaces/ISignatureValidator.sol";
+import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
+
+/**
+ * @title Session Key Manager module for Biconomy Modular Smart Accounts.
+ * @dev Performs basic verifications for every session key signed userOp.
+ * Checks if the session key has been enabled, that it is not due and has not yet expired
+ * Then passes the validation flow to appropriate Session Validation module
+ *         - For the sake of efficiency and flexibility, doesn't limit what operations
+ *           Session Validation modules can perform
+ *         - Should be used with carefully verified and audited Session Validation Modules only
+ *         - Compatible with Biconomy Modular Interface v 0.1
+ * @author Fil Makarov - <[email protected]>
+ */
+
+contract DANSessionKeyManager is
+    BaseAuthorizationModule,
+    ISessionKeyManagerModule
+{
+    string public constant NAME = "DAN Session Manager";
+    string public constant VERSION = "1.1.0";
+
+    uint256 private constant MODULE_SIGNATURE_OFFSET = 96;
+
+    /**
+     * @dev mapping of Smart Account to a SessionStorage
+     * Info about session keys is stored as root of the merkle tree built over the session keys
+     */
+    mapping(address => SessionStorage) internal _userSessions;
+
+    // TODO // Review
+    // What if we could take some inspiration from Session Key Manager Hybrid module.
+
+    /// @inheritdoc ISessionKeyManagerModule
+    function setMerkleRoot(bytes32 _merkleRoot) external override {
+        _userSessions[msg.sender].merkleRoot = _merkleRoot;
+        emit MerkleRootUpdated(msg.sender, _merkleRoot);
+    }
+
+    // TODO // Review
+    // We could also remove sessionValidationModule everywhere
+
+    /// @inheritdoc IAuthorizationModule
+    function validateUserOp(
+        UserOperation calldata userOp,
+        bytes32 userOpHash
+    ) external virtual returns (uint256) {
+        (
+            uint48 validUntil,
+            uint48 validAfter,
+            address sessionValidationModule,
+            bytes memory sessionKeyData,
+            bytes32[] memory merkleProof,
+            bytes memory sessionKeySignature
+        ) = abi.decode(
+                userOp.signature[MODULE_SIGNATURE_OFFSET:],
+                (uint48, uint48, address, bytes, bytes32[], bytes)
+            );
+
+        validateSessionKey(
+            userOp.sender,
+            validUntil,
+            validAfter,
+            sessionValidationModule,
+            sessionKeyData,
+            merkleProof
+        );
+
+        (address sessionKey, , , , ) = abi.decode(
+            sessionKeyData,
+            (address, address, address, uint256, uint256)
+        );
+
+        bool isValidSignatureFromDAN = ECDSA.recover(
+            ECDSA.toEthSignedMessageHash(userOpHash),
+            sessionKeySignature
+        ) == sessionKey;
+
+        return
+            _packValidationData(
+                !isValidSignatureFromDAN,
+                validUntil,
+                validAfter
+            );
+    }
+
+    /// @inheritdoc ISessionKeyManagerModule
+    function getSessionKeys(
+        address smartAccount
+    ) external view override returns (SessionStorage memory) {
+        return _userSessions[smartAccount];
+    }
+
+    /// @inheritdoc ISessionKeyManagerModule
+    function validateSessionKey(
+        address smartAccount,
+        uint48 validUntil,
+        uint48 validAfter,
+        address sessionValidationModule,
+        bytes memory sessionKeyData,
+        bytes32[] memory merkleProof
+    ) public virtual override {
+        SessionStorage storage sessionKeyStorage = _getSessionData(
+            smartAccount
+        );
+        bytes32 leaf = keccak256(
+            abi.encodePacked(
+                validUntil,
+                validAfter,
+                sessionValidationModule,
+                sessionKeyData
+            )
+        );
+        if (
+            !MerkleProof.verify(merkleProof, sessionKeyStorage.merkleRoot, leaf)
+        ) {
+            revert("SessionNotApproved");
+        }
+    }
+
+    /// @inheritdoc ISignatureValidator
+    function isValidSignature(
+        bytes32 _dataHash,
+        bytes memory _signature
+    ) public pure override returns (bytes4) {
+        (_dataHash, _signature);
+        return 0xffffffff; // do not support it here
+    }
+
+    /// @inheritdoc ISignatureValidator
+    function isValidSignatureUnsafe(
+        bytes32 _dataHash,
+        bytes memory _signature
+    ) public pure override returns (bytes4) {
+        (_dataHash, _signature);
+        return 0xffffffff; // do not support it here
+    }
+
+    /**
+     * @dev returns the SessionStorage object for a given smartAccount
+     * @param _account Smart Account address
+     * @return sessionKeyStorage SessionStorage object at storage
+     */
+    function _getSessionData(
+        address _account
+    ) internal view returns (SessionStorage storage sessionKeyStorage) {
+        sessionKeyStorage = _userSessions[_account];
+    }
+}

test/bundler-integration/module/DAN.SKM.Specs.ts

@@ -0,0 +1,235 @@
+import { expect } from "chai";
+import { ethers, deployments } from "hardhat";
+import { makeEcdsaModuleUserOp } from "../../utils/userOp";
+import {
+  makeEcdsaSessionKeySignedUserOp,
+  enableNewTreeForSmartAccountViaEcdsa,
+} from "../../utils/sessionKey";
+import { encodeTransfer } from "../../utils/testUtils";
+import { hexZeroPad, hexConcat } from "ethers/lib/utils";
+import {
+  getEntryPoint,
+  getSmartAccountImplementation,
+  getSmartAccountFactory,
+  getMockToken,
+  getEcdsaOwnershipRegistryModule,
+  getSmartAccountWithModule,
+} from "../../utils/setupHelper";
+import { keccak256 } from "ethereumjs-util";
+import { MerkleTree } from "merkletreejs";
+import { SignerWithAddress } from "@nomiclabs/hardhat-ethers/signers";
+import { BundlerTestEnvironment } from "../environment/bundlerEnvironment";
+
+describe("SessionKey: DAN Session Manager Module (with Bundler)", async () => {
+  let [deployer, smartAccountOwner, charlie, verifiedSigner, sessionKey] =
+    [] as SignerWithAddress[];
+
+  let environment: BundlerTestEnvironment;
+
+  before(async function () {
+    const chainId = (await ethers.provider.getNetwork()).chainId;
+    if (chainId !== BundlerTestEnvironment.BUNDLER_ENVIRONMENT_CHAIN_ID) {
+      this.skip();
+    }
+
+    environment = await BundlerTestEnvironment.getDefaultInstance();
+  });
+
+  beforeEach(async function () {
+    [deployer, smartAccountOwner, charlie, verifiedSigner, sessionKey] =
+      await ethers.getSigners();
+  });
+
+  afterEach(async function () {
+    const chainId = (await ethers.provider.getNetwork()).chainId;
+    if (chainId !== BundlerTestEnvironment.BUNDLER_ENVIRONMENT_CHAIN_ID) {
+      this.skip();
+    }
+
+    await Promise.all([
+      environment.revert(environment.defaultSnapshot!),
+      environment.resetBundler(),
+    ]);
+  });
+
+  const setupTests = deployments.createFixture(async ({ deployments }) => {
+    await deployments.fixture();
+
+    const entryPoint = await getEntryPoint();
+    const mockToken = await getMockToken();
+    const ecdsaModule = await getEcdsaOwnershipRegistryModule();
+    const EcdsaOwnershipRegistryModule = await ethers.getContractFactory(
+      "EcdsaOwnershipRegistryModule"
+    );
+    const ecdsaOwnershipSetupData =
+      EcdsaOwnershipRegistryModule.interface.encodeFunctionData(
+        "initForSmartAccount",
+        [await smartAccountOwner.getAddress()]
+      );
+    const smartAccountDeploymentIndex = 0;
+    const userSA = await getSmartAccountWithModule(
+      ecdsaModule.address,
+      ecdsaOwnershipSetupData,
+      smartAccountDeploymentIndex
+    );
+
+    // send funds to userSA and mint tokens
+    await deployer.sendTransaction({
+      to: userSA.address,
+      value: ethers.utils.parseEther("10"),
+    });
+    await mockToken.mint(userSA.address, ethers.utils.parseEther("1000000"));
+
+    const dANSessionKeyManager = await (
+      await ethers.getContractFactory("DANSessionKeyManager")
+    ).deploy();
+    const userOp = await makeEcdsaModuleUserOp(
+      "enableModule",
+      [dANSessionKeyManager.address],
+      userSA.address,
+      smartAccountOwner,
+      entryPoint,
+      ecdsaModule.address,
+      {
+        preVerificationGas: 50000,
+      }
+    );
+    await environment.sendUserOperation(userOp, entryPoint.address);
+
+    const mockSessionValidationModule = await (
+      await ethers.getContractFactory("MockSessionValidationModule")
+    ).deploy();
+
+    const validUntil = 0;
+    const validAfter = 0;
+    const sessionKeyData = hexZeroPad(sessionKey.address, 20);
+    const leafData = hexConcat([
+      hexZeroPad(ethers.utils.hexlify(validUntil), 6),
+      hexZeroPad(ethers.utils.hexlify(validAfter), 6),
+      hexZeroPad(mockSessionValidationModule.address, 20),
+      sessionKeyData,
+    ]);
+
+    const merkleTree = await enableNewTreeForSmartAccountViaEcdsa(
+      [ethers.utils.keccak256(leafData)],
+      dANSessionKeyManager,
+      userSA.address,
+      smartAccountOwner,
+      entryPoint,
+      ecdsaModule.address
+    );
+
+    return {
+      entryPoint: entryPoint,
+      smartAccountImplementation: await getSmartAccountImplementation(),
+      smartAccountFactory: await getSmartAccountFactory(),
+      mockToken: mockToken,
+      ecdsaModule: ecdsaModule,
+      userSA: userSA,
+      sessionKeyManager: dANSessionKeyManager,
+      mockSessionValidationModule: mockSessionValidationModule,
+      sessionKeyData: sessionKeyData,
+      leafData: leafData,
+      merkleTree: merkleTree,
+    };
+  });
+
+  describe("setMerkleRoot", async () => {
+    it("should add new session key by setting new merkle tree root", async () => {
+      const {
+        userSA,
+        sessionKeyManager,
+        mockSessionValidationModule,
+        entryPoint,
+        ecdsaModule,
+      } = await setupTests();
+
+      const data = hexConcat([
+        hexZeroPad("0x00", 6),
+        hexZeroPad("0x00", 6),
+        hexZeroPad(mockSessionValidationModule.address, 20),
+        hexZeroPad(sessionKey.address, 20),
+      ]);
+
+      const merkleTree = new MerkleTree(
+        [ethers.utils.keccak256(data)],
+        keccak256,
+        { sortPairs: false, hashLeaves: false }
+      );
+      const addMerkleRootUserOp = await makeEcdsaModuleUserOp(
+        "execute_ncC",
+        [
+          sessionKeyManager.address,
+          ethers.utils.parseEther("0"),
+          sessionKeyManager.interface.encodeFunctionData("setMerkleRoot", [
+            merkleTree.getHexRoot(),
+          ]),
+        ],
+        userSA.address,
+        smartAccountOwner,
+        entryPoint,
+        ecdsaModule.address,
+        {
+          preVerificationGas: 50000,
+        }
+      );
+
+      await environment.sendUserOperation(
+        addMerkleRootUserOp,
+        entryPoint.address
+      );
+
+      expect(
+        (await sessionKeyManager.getSessionKeys(userSA.address)).merkleRoot
+      ).to.equal(merkleTree.getHexRoot());
+    });
+  });
+
+  describe("validateUserOp", async () => {
+    // Review : failing
+    it("should be able to process Session Key signed userOp via Mock session validation module", async () => {
+      const {
+        entryPoint,
+        userSA,
+        sessionKeyManager,
+        mockSessionValidationModule,
+        mockToken,
+        sessionKeyData,
+        leafData,
+        merkleTree,
+      } = await setupTests();
+      const tokenAmountToTransfer = ethers.utils.parseEther("0.834");
+
+      const transferUserOp = await makeEcdsaSessionKeySignedUserOp(
+        "execute_ncC",
+        [
+          mockToken.address,
+          ethers.utils.parseEther("0"),
+          encodeTransfer(charlie.address, tokenAmountToTransfer.toString()),
+        ],
+        userSA.address,
+        sessionKey,
+        entryPoint,
+        sessionKeyManager.address,
+        0,
+        0,
+        mockSessionValidationModule.address,
+        sessionKeyData,
+        merkleTree.getHexProof(ethers.utils.keccak256(leafData)),
+        {
+          preVerificationGas: 50000,
+        }
+      );
+
+      const charlieTokenBalanceBefore = await mockToken.balanceOf(
+        charlie.address
+      );
+
+      await environment.sendUserOperation(transferUserOp, entryPoint.address);
+
+      expect(await mockToken.balanceOf(charlie.address)).to.equal(
+        charlieTokenBalanceBefore.add(tokenAmountToTransfer)
+      );
+    });
+  });
+});