Merge pull request #742 from 0xsequence/session-refactor

Use wallets configuration for sessions management

Author: Agusx1211

packages/wallet/wdk/src/index.ts

@@ -1,3 +1,2 @@
 export * as Identity from './identity/index.js'
-export * as Session from './session/index.js'
 export * as Sequence from './sequence/index.js'

packages/wallet/wdk/src/sequence/manager.ts

@@ -473,6 +473,7 @@ export class Manager {
     args: AuthorizeImplicitSessionArgs,
   ): Promise<string> {
     return this.shared.modules.sessions.prepareAuthorizeImplicitSession(walletAddress, sessionAddress, args)
+    // Run completeAuthorizeImplicitSession next
   }
 
   public async completeAuthorizeImplicitSession(requestId: string): Promise<{
@@ -488,26 +489,26 @@ export class Manager {
     permissions: CoreSigners.Session.ExplicitParams,
   ): Promise<string> {
     return this.shared.modules.sessions.addExplicitSession(walletAddress, sessionAddress, permissions)
+    // Run completeSessionUpdate next
   }
 
   public async removeExplicitSession(walletAddress: Address.Address, sessionAddress: Address.Address): Promise<string> {
     return this.shared.modules.sessions.removeExplicitSession(walletAddress, sessionAddress)
+    // Run completeSessionUpdate next
   }
 
   public async addBlacklistAddress(walletAddress: Address.Address, address: Address.Address): Promise<string> {
     return this.shared.modules.sessions.addBlacklistAddress(walletAddress, address)
+    // Run completeSessionUpdate next
   }
 
   public async removeBlacklistAddress(walletAddress: Address.Address, address: Address.Address): Promise<string> {
     return this.shared.modules.sessions.removeBlacklistAddress(walletAddress, address)
+    // Run completeSessionUpdate next
   }
 
   public async completeSessionUpdate(requestId: string) {
-    const sigRequest = await this.shared.modules.signatures.get(requestId)
-    if (sigRequest.action !== 'session-update' || !Payload.isConfigUpdate(sigRequest.envelope.payload)) {
-      throw new Error('Invalid action')
-    }
-    return this.shared.modules.sessions.completeSessionUpdate(sigRequest.wallet, requestId)
+    return this.shared.modules.sessions.completeSessionUpdate(requestId)
   }
 
   // Recovery

packages/wallet/wdk/src/sequence/sessions.ts

@@ -1,66 +1,47 @@
-import { Signers as CoreSigners, Envelope, Wallet } from '@0xsequence/wallet-core'
+import { Signers as CoreSigners, Envelope } from '@0xsequence/wallet-core'
 import {
   Attestation,
+  Config,
+  Constants,
+  GenericTree,
   Payload,
-  SessionConfig,
   Signature as SequenceSignature,
-  Config,
+  SessionConfig,
 } from '@0xsequence/wallet-primitives'
-import { Address, Bytes, Hex, Provider, RpcTransport } from 'ox'
-import { SessionController } from '../session/index.js'
-import { IdentityHandler, identityTypeToHex } from './handlers/identity.js'
-import { Shared } from './manager.js'
-import { AuthCodePkceHandler } from './handlers/authcode-pkce.js'
+import { Address, Bytes, Hex } from 'ox'
 import { IdentityType } from '../identity/index.js'
+import { AuthCodePkceHandler } from './handlers/authcode-pkce.js'
+import { IdentityHandler, identityTypeToHex } from './handlers/identity.js'
+import { ManagerOptionsDefaults, Shared } from './manager.js'
+import { Actions } from './types/signature-request.js'
 
 export type AuthorizeImplicitSessionArgs = {
   target: string
   applicationData?: Hex.Hex
 }
 
-export class Sessions {
-  private readonly _sessionControllers: Map<Address.Address, SessionController> = new Map()
+const DefaultSessionManagerAddresses = [Constants.DefaultSessionManager]
 
-  constructor(private readonly shared: Shared) {}
-
-  getCoreWallet(walletAddress: Address.Address): Wallet {
-    return new Wallet(walletAddress, {
-      context: this.shared.sequence.context,
-      guest: this.shared.sequence.guest,
-      stateProvider: this.shared.sequence.stateProvider,
-    })
-  }
+export class Sessions {
+  constructor(
+    private readonly shared: Shared,
+    private readonly sessionManagerAddresses: Address.Address[] = DefaultSessionManagerAddresses,
+  ) {}
 
-  async getControllerForWallet(walletAddress: Address.Address, chainId?: bigint): Promise<SessionController> {
-    if (this._sessionControllers.has(walletAddress)) {
-      return this._sessionControllers.get(walletAddress)!
+  async getSessionTopology(walletAddress: Address.Address): Promise<SessionConfig.SessionsTopology> {
+    const { modules } = await this.shared.modules.wallets.getConfigurationParts(walletAddress)
+    const managerLeaf = modules.find((leaf) =>
+      this.sessionManagerAddresses.some((addr) => Address.isEqual(addr, leaf.address)),
+    )
+    if (!managerLeaf) {
+      throw new Error('Session manager not found')
     }
-
-    // Construct the wallet
-    const wallet = this.getCoreWallet(walletAddress)
-
-    // Get the provider if available
-    let provider: Provider.Provider | undefined
-    if (chainId) {
-      const network = this.shared.sequence.networks.find((network) => network.chainId === chainId)
-      if (network) {
-        provider = Provider.from(RpcTransport.fromHttp(network.rpc))
-      }
+    const imageHash = managerLeaf.imageHash
+    const tree = await this.shared.sequence.stateProvider.getTree(imageHash)
+    if (!tree) {
+      throw new Error('Session topology not found')
     }
-
-    // Create the controller
-    const controller = new SessionController({
-      wallet,
-      provider,
-      stateProvider: this.shared.sequence.stateProvider,
-    })
-    this._sessionControllers.set(walletAddress, controller)
-    return controller
-  }
-
-  async getSessionTopology(walletAddress: Address.Address): Promise<SessionConfig.SessionsTopology> {
-    const controller = await this.getControllerForWallet(walletAddress)
-    return controller.getTopology()
+    return SessionConfig.configurationTreeToSessionsTopology(tree)
   }
 
   async prepareAuthorizeImplicitSession(
@@ -171,59 +152,89 @@ export class Sessions {
     permissions: CoreSigners.Session.ExplicitParams,
     origin?: string,
   ): Promise<string> {
-    const controller = await this.getControllerForWallet(walletAddress)
-    const envelope = await controller.addExplicitSession(sessionAddress, permissions)
-    return this.prepareSessionUpdate(envelope, origin)
+    const topology = await this.getSessionTopology(walletAddress)
+    const newTopology = SessionConfig.addExplicitSession(topology, {
+      ...permissions,
+      signer: sessionAddress,
+    })
+    return this.prepareSessionUpdate(walletAddress, newTopology, origin)
   }
 
   async removeExplicitSession(
     walletAddress: Address.Address,
     sessionAddress: Address.Address,
     origin?: string,
   ): Promise<string> {
-    const controller = await this.getControllerForWallet(walletAddress)
-    const envelope = await controller.removeExplicitSession(sessionAddress)
-    return this.prepareSessionUpdate(envelope, origin)
+    const topology = await this.getSessionTopology(walletAddress)
+    const newTopology = SessionConfig.removeExplicitSession(topology, sessionAddress)
+    if (!newTopology) {
+      throw new Error('Session not found')
+    }
+    return this.prepareSessionUpdate(walletAddress, newTopology, origin)
   }
 
   async addBlacklistAddress(
     walletAddress: Address.Address,
     address: Address.Address,
     origin?: string,
   ): Promise<string> {
-    const controller = await this.getControllerForWallet(walletAddress)
-    const envelope = await controller.addBlacklistAddress(address)
-    return this.prepareSessionUpdate(envelope, origin)
+    const topology = await this.getSessionTopology(walletAddress)
+    const newTopology = SessionConfig.addToImplicitBlacklist(topology, address)
+    return this.prepareSessionUpdate(walletAddress, newTopology, origin)
   }
 
   async removeBlacklistAddress(
     walletAddress: Address.Address,
     address: Address.Address,
     origin?: string,
   ): Promise<string> {
-    const controller = await this.getControllerForWallet(walletAddress)
-    const envelope = await controller.removeBlacklistAddress(address)
-    return this.prepareSessionUpdate(envelope, origin)
+    const topology = await this.getSessionTopology(walletAddress)
+    const newTopology = SessionConfig.removeFromImplicitBlacklist(topology, address)
+    return this.prepareSessionUpdate(walletAddress, newTopology, origin)
   }
 
   private async prepareSessionUpdate(
-    envelope: Envelope.Envelope<Payload.ConfigUpdate>,
+    walletAddress: Address.Address,
+    topology: SessionConfig.SessionsTopology,
     origin: string = 'wallet-webapp',
   ): Promise<string> {
-    return await this.shared.modules.signatures.request(envelope, 'session-update', {
+    // Store the new configuration
+    const tree = SessionConfig.sessionsTopologyToConfigurationTree(topology)
+    await this.shared.sequence.stateProvider.saveTree(tree)
+    const newImageHash = GenericTree.hash(tree)
+
+    // Find the session manager in the old configuration
+    const { modules } = await this.shared.modules.wallets.getConfigurationParts(walletAddress)
+    const managerLeaf = modules.find((leaf) =>
+      this.sessionManagerAddresses.some((addr) => Address.isEqual(addr, leaf.address)),
+    )
+    if (!managerLeaf) {
+      // Missing. Add it
+      modules.push({
+        ...ManagerOptionsDefaults.defaultSessionsTopology,
+        imageHash: newImageHash,
+      })
+    } else {
+      // Update the configuration to use the new session manager image hash
+      managerLeaf.imageHash = newImageHash
+    }
+
+    return this.shared.modules.wallets.requestConfigurationUpdate(
+      walletAddress,
+      {
+        modules,
+      },
+      Actions.SessionUpdate,
       origin,
-    })
+    )
   }
 
-  async completeSessionUpdate(walletAddress: Address.Address, requestId: string) {
-    const controller = await this.getControllerForWallet(walletAddress)
+  async completeSessionUpdate(requestId: string) {
     const sigRequest = await this.shared.modules.signatures.get(requestId)
-    const envelope = sigRequest.envelope
-    if (sigRequest.action !== 'session-update' || !Payload.isConfigUpdate(envelope.payload)) {
+    if (sigRequest.action !== 'session-update' || !Payload.isConfigUpdate(sigRequest.envelope.payload)) {
       throw new Error('Invalid action')
     }
-    console.log('Completing session update:', requestId)
-    await controller.completeUpdateConfiguration(envelope as Envelope.Signed<Payload.ConfigUpdate>)
-    return this.shared.modules.signatures.complete(requestId)
+
+    return this.shared.modules.wallets.completeConfigurationUpdate(requestId)
   }
 }

packages/wallet/wdk/src/sequence/wallets.ts

@@ -1,15 +1,14 @@
 import { Wallet as CoreWallet, Envelope, Signers, State } from '@0xsequence/wallet-core'
-import { Config, Extensions, GenericTree, Payload, SessionConfig } from '@0xsequence/wallet-primitives'
+import { Config, GenericTree, Payload, SessionConfig } from '@0xsequence/wallet-primitives'
 import { Address, Hex, Provider, RpcTransport } from 'ox'
 import { AuthCommitment } from '../dbs/auth-commitments.js'
 import { AuthCodePkceHandler } from './handlers/authcode-pkce.js'
 import { MnemonicHandler } from './handlers/mnemonic.js'
 import { OtpHandler } from './handlers/otp.js'
 import { Shared } from './manager.js'
-import { Wallet } from './types/wallet.js'
-import { Kinds, RecoverySigner, WitnessExtraSignerKind } from './types/signer.js'
-import { WalletSelectionUiHandler } from './types/wallet.js'
 import { Action } from './types/index.js'
+import { Kinds, WitnessExtraSignerKind } from './types/signer.js'
+import { Wallet, WalletSelectionUiHandler } from './types/wallet.js'
 
 export type StartSignUpWithRedirectArgs = {
   kind: 'google-pkce' | 'apple-pkce'