@zerodev/wallet-core 0.0.3 → 0.0.5

package/src/core/createZeroDevWalletCore.ts

@@ -7,11 +7,8 @@ import {
   type ZeroDevWalletClient,
   zeroDevWalletTransport,
 } from '../client/index.js'
-import {
-  DEFAULT_ORGANIZATION_ID,
-  DEFAULT_SESSION_EXPIRATION_IN_SECONDS,
-  KMS_SERVER_URL,
-} from '../constants.js'
+import { DEFAULT_ORGANIZATION_ID, KMS_SERVER_URL } from '../constants.js'
+import { RestRequestError } from '../errors/request.js'
 import { createNoopPasskeyStamper } from '../stampers/noopPasskeyStamper.js'
 import type { ApiKeyStamper, PasskeyStamper } from '../stampers/types.js'
 import {
@@ -102,11 +99,20 @@ export interface ZeroDevWalletSDK {
     sessionId?: string,
   ) => Promise<ZeroDevWalletSession | undefined>
 
-  logout: () => Promise<boolean>
+  /**
+   * Revoke the active remote key, then erase local state. `force` still tries
+   * remote revocation first but permits local recovery if that attempt fails;
+   * use it only when accepting that the remote key may remain until expiry.
+   */
+  logout: (options?: { force?: boolean }) => Promise<boolean>
 
   toAccount: () => Promise<LocalAccount>
 }
 
+// The browser/native key vault is a single physical slot shared by SDK
+// instances, so key transitions must be serialized across projects.
+let keyTransitionTail = Promise.resolve()
+
 export async function createZeroDevWalletCore(
   config: ZeroDevWalletConfigCore,
 ): Promise<ZeroDevWalletSDK> {
@@ -137,12 +143,119 @@ export async function createZeroDevWalletCore(
     fallback: DEFAULT_ORGANIZATION_ID,
   })
 
+  let preparedOAuthPublicKey: string | undefined
+
+  const prepareFreshKeyRotation = async () => {
+    if (preparedOAuthPublicKey) {
+      await client.apiKeyStamper.discardKeyRotation()
+      preparedOAuthPublicKey = undefined
+    }
+    return client.apiKeyStamper.prepareKeyRotation()
+  }
+
+  const withKeyTransition = async <T>(transition: () => Promise<T>) => {
+    const result = keyTransitionTail.then(transition)
+    keyTransitionTail = result.then(
+      () => undefined,
+      () => undefined,
+    )
+    return result
+  }
+
+  const normalizeKey = (key: string) => key.replace(/^0x/, '').toLowerCase()
+
+  const createReplacementSession = (
+    token: string,
+    expectedPublicKey: string,
+    idPrefix: 'indexedDb' | 'oauth' | 'otp',
+    sessionType?: SessionType,
+  ): ZeroDevWalletSession => {
+    const parsed = parseSession(token)
+    if (
+      !parsed.publicKey ||
+      normalizeKey(parsed.publicKey) !== normalizeKey(expectedPublicKey)
+    ) {
+      throw new Error(
+        'Session public key does not match the locally prepared signing key',
+      )
+    }
+    const now = Date.now()
+    return {
+      id: `session_${idPrefix}_${now}`,
+      userId: parsed.userId,
+      organizationId: parsed.organizationId,
+      stamperType: 'apiKey',
+      sessionType: sessionType ?? parsed.sessionType ?? SessionType.READ_WRITE,
+      token,
+      publicKey: expectedPublicKey,
+      expiry: parsed.expiry,
+      createdAt: now,
+    }
+  }
+
+  const commitReplacementSession = async (
+    session: ZeroDevWalletSession,
+    targetPublicKey: string,
+  ) => {
+    await sessionStorageManager.stageSessionTransition(session, targetPublicKey)
+    try {
+      await client.apiKeyStamper.commitKeyRotation()
+    } catch (error) {
+      await client.apiKeyStamper.discardKeyRotation()
+      const activePublicKey = await client.apiKeyStamper.getPublicKey()
+      const recovered =
+        await sessionStorageManager.recoverSessionTransition(activePublicKey)
+      if (
+        !activePublicKey ||
+        normalizeKey(activePublicKey) !== normalizeKey(targetPublicKey) ||
+        !recovered
+      ) {
+        throw error
+      }
+      return
+    }
+
+    try {
+      const committed = await sessionStorageManager.commitSessionTransition()
+      if (!committed) {
+        throw new Error('Session transition journal disappeared before commit')
+      }
+    } catch (error) {
+      const recovered = await sessionStorageManager.recoverSessionTransition(
+        await client.apiKeyStamper.getPublicKey(),
+      )
+      if (!recovered) throw error
+    }
+  }
+
+  await withKeyTransition(async () => {
+    const activePublicKey = await client.apiKeyStamper.getPublicKey()
+    await sessionStorageManager.recoverSessionTransition(activePublicKey)
+    const restoredSession = await sessionStorageManager.getActiveSession()
+    if (restoredSession) {
+      try {
+        const restoredPublicKey = parseSession(restoredSession.token).publicKey
+        if (
+          !activePublicKey ||
+          !restoredPublicKey ||
+          normalizeKey(activePublicKey) !== normalizeKey(restoredPublicKey)
+        ) {
+          await sessionStorageManager.clearAllSessions()
+        }
+      } catch {
+        await sessionStorageManager.clearAllSessions()
+      }
+    }
+  })
+
   return {
     client,
     async getPublicKey() {
-      await client.apiKeyStamper.resetKeyPair()
-      const compressedPublicKey = await client.apiKeyStamper.getPublicKey()
-      return compressedPublicKey
+      return withKeyTransition(async () => {
+        if (preparedOAuthPublicKey) return preparedOAuthPublicKey
+        preparedOAuthPublicKey = await client.apiKeyStamper.prepareKeyRotation()
+        return preparedOAuthPublicKey
+      })
     },
 
     async getSession() {
@@ -159,87 +272,113 @@ export async function createZeroDevWalletCore(
     },
 
     async switchSession(sessionId: string) {
-      await sessionStorageManager.setActiveSession(sessionId)
-      const session = await sessionStorageManager.getActiveSession()
-
-      return session
+      const activeSessionKey = await sessionStorageManager.getActiveSessionKey()
+      if (sessionId !== activeSessionKey) {
+        throw new Error(
+          'Session switching is not supported with a single active key vault',
+        )
+      }
+      return sessionStorageManager.getActiveSession()
     },
 
     async clearSession(sessionId: string) {
+      if ((await sessionStorageManager.getActiveSessionKey()) === sessionId) {
+        throw new Error(
+          'Refusing to clear the active session locally; use logout() to revoke its remote key first.',
+        )
+      }
       await sessionStorageManager.clearSession(sessionId)
     },
 
     async clearAllSessions() {
+      if (await sessionStorageManager.getActiveSession()) {
+        throw new Error(
+          'Refusing to clear the active session locally; use logout() to revoke its remote key first.',
+        )
+      }
       await sessionStorageManager.clearAllSessions()
     },
 
     async refreshSession(sessionId?: string) {
-      const activeSession = sessionId
-        ? await sessionStorageManager.getSession(sessionId)
-        : await sessionStorageManager.getActiveSession()
-      if (!activeSession) {
-        throw new Error('No active session')
-      }
-      if (activeSession.stamperType === 'apiKey') {
-        const compressedPublicKeyHex =
-          await client.apiKeyStamper.prepareKeyRotation()
-        const data = await client.loginWithStamp({
-          targetPublicKey: compressedPublicKeyHex,
-          projectId,
-          // Stamp-login is signed against the Turnkey parent org; the backend
-          // resolves the sub-org from the stamped credential. Signing the
-          // sub-org here makes the relayed payload's org mismatch the
-          // signature → Turnkey SIGNATURE_INVALID.
-          organizationId: await resolveOrganizationId(),
-          stampWith: 'apiKey',
-        })
-        await client.apiKeyStamper.commitKeyRotation()
-        const parsedSession = parseSession(data.session)
-        const session: ZeroDevWalletSession = {
-          id: `session_indexedDb_${Date.now()}`,
-          userId: parsedSession.userId,
-          organizationId: parsedSession.organizationId,
-          stamperType: 'apiKey',
-          sessionType: SessionType.READ_WRITE,
-          token: data.session,
-          expiry: parsedSession.expiry,
-          createdAt: Date.now(),
+      return withKeyTransition(async () => {
+        const activeSessionKey =
+          await sessionStorageManager.getActiveSessionKey()
+        if (sessionId && sessionId !== activeSessionKey) {
+          throw new Error(
+            'Refreshing an inactive session is not supported with a single active key vault',
+          )
         }
-        await sessionStorageManager.clearSession(activeSession.id)
-        await sessionStorageManager.storeSession(session, session.id)
-        return session
-      }
-      throw new Error('Invalid session type')
+        const activeSession = await sessionStorageManager.getActiveSession()
+        if (!activeSession) {
+          throw new Error('No active session')
+        }
+        if (activeSession.stamperType !== 'apiKey') {
+          throw new Error('Invalid session type')
+        }
+
+        const targetPublicKey = await prepareFreshKeyRotation()
+        try {
+          const data = await client.loginWithStamp({
+            targetPublicKey,
+            projectId,
+            // Stamp-login is signed against the Turnkey parent org; the backend
+            // resolves the sub-org from the stamped credential. Signing the
+            // sub-org here makes the relayed payload's org mismatch the
+            // signature → Turnkey SIGNATURE_INVALID.
+            organizationId: await resolveOrganizationId(),
+            stampWith: 'apiKey',
+          })
+          const session = createReplacementSession(
+            data.session,
+            targetPublicKey,
+            'indexedDb',
+            SessionType.READ_WRITE,
+          )
+          await commitReplacementSession(session, targetPublicKey)
+          return session
+        } catch (error) {
+          await client.apiKeyStamper.discardKeyRotation()
+          throw error
+        }
+      })
     },
 
     // [TODO] refactor to smaller utils/actions
     async auth(params: AuthParams) {
       switch (params.type) {
         case 'oauth': {
-          const popSignature = await client.apiKeyStamper.sign(params.sessionId)
-          const data = await client.authenticateWithOAuth({
-            provider: params.provider,
-            projectId,
-            sessionId: params.sessionId,
-            popSignature,
-          })
+          return withKeyTransition(async () => {
+            const targetPublicKey =
+              preparedOAuthPublicKey ??
+              (await client.apiKeyStamper.prepareKeyRotation())
+            preparedOAuthPublicKey = undefined
+            try {
+              const popSignature = await client.apiKeyStamper.signPending(
+                params.sessionId,
+              )
+              const data = await client.authenticateWithOAuth({
+                provider: params.provider,
+                projectId,
+                sessionId: params.sessionId,
+                popSignature,
+              })
 
-          if (data.session) {
-            // Parse the JWT to get session data
-            const parsedSession = parseSession(data.session)
-            const session: ZeroDevWalletSession = {
-              id: `session_oauth_${Date.now()}`,
-              userId: parsedSession.userId,
-              organizationId: parsedSession.organizationId,
-              stamperType: 'apiKey',
-              sessionType: parsedSession.sessionType || SessionType.READ_WRITE,
-              token: data.session,
-              expiry: parsedSession.expiry,
-              createdAt: Date.now(),
+              if (!data.session) {
+                await client.apiKeyStamper.discardKeyRotation()
+                return data
+              }
+              const session = createReplacementSession(
+                data.session,
+                targetPublicKey,
+                'oauth',
+              )
+              await commitReplacementSession(session, targetPublicKey)
+              return data
+            } catch (error) {
+              await client.apiKeyStamper.discardKeyRotation()
+              throw error
             }
-            await sessionStorageManager.storeSession(session, session.id)
-          }
-          return data
+          })
         }
         case 'passkey': {
           const { type } = params
@@ -248,83 +387,80 @@ export async function createZeroDevWalletCore(
             'mode' in params &&
             params.mode === 'register'
           ) {
-            await client.apiKeyStamper.resetKeyPair()
-            const tempPublicKey = await client.apiKeyStamper.getPublicKey()
-            if (!tempPublicKey) {
-              throw new Error('Failed to get public key')
-            }
-            const name = `ZeroDevWallet-${humanReadableDateTime()}`
-            const { attestation, encodedChallenge } =
-              await passkeyStamper.register({
-                rp: { id: rpId, name: '' },
-                userName: name,
-              })
-            const data = await client.registerWithPasskey({
-              attestation,
-              challenge: encodedChallenge,
-              projectId,
-              encodedPublicKey: tempPublicKey,
-            })
-            const compressedPublicKeyHex =
-              await client.apiKeyStamper.prepareKeyRotation()
-            const loginData = await client.loginWithStamp({
-              projectId,
-              targetPublicKey: compressedPublicKeyHex,
-              // Sign against the parent org (see refreshSession note) — the
-              // backend derives the sub-org from the stamped credential.
-              organizationId: await resolveOrganizationId(),
+            return withKeyTransition(async () => {
+              if (await sessionStorageManager.getActiveSession()) {
+                throw new Error(
+                  'Passkey registration requires an unauthenticated wallet; logout before registering a new wallet.',
+                )
+              }
+              const registrationPublicKey = await prepareFreshKeyRotation()
+              try {
+                const name = `ZeroDevWallet-${humanReadableDateTime()}`
+                const { attestation, encodedChallenge } =
+                  await passkeyStamper.register({
+                    rp: { id: rpId, name: 'ZeroDev Wallet' },
+                    userName: name,
+                  })
+                const data = await client.registerWithPasskey({
+                  attestation,
+                  challenge: encodedChallenge,
+                  projectId,
+                  encodedPublicKey: registrationPublicKey,
+                })
+
+                // Registration proves the temporary key exists remotely. Make
+                // it active so it can authorize the durable replacement key.
+                await client.apiKeyStamper.commitKeyRotation()
+                const targetPublicKey = await prepareFreshKeyRotation()
+                const loginData = await client.loginWithStamp({
+                  projectId,
+                  targetPublicKey,
+                  // Sign against the parent org (see refreshSession note) — the
+                  // backend derives the sub-org from the stamped credential.
+                  organizationId: await resolveOrganizationId(),
+                })
+                const session = createReplacementSession(
+                  loginData.session,
+                  targetPublicKey,
+                  'indexedDb',
+                )
+                await commitReplacementSession(session, targetPublicKey)
+                return data
+              } catch (error) {
+                await client.apiKeyStamper.discardKeyRotation()
+                throw error
+              }
             })
-            await client.apiKeyStamper.commitKeyRotation()
-            const parsedSession = parseSession(loginData.session)
-            const session: ZeroDevWalletSession = {
-              id: `session_indexedDb_${Date.now()}`,
-              stamperType: 'apiKey',
-              createdAt: Date.now(),
-              sessionType: SessionType.READ_WRITE,
-              userId: parsedSession.userId,
-              organizationId: parsedSession.organizationId,
-              expiry:
-                Date.now() +
-                Number(DEFAULT_SESSION_EXPIRATION_IN_SECONDS) * 1000,
-              token: loginData.session,
-            }
-            await sessionStorageManager.storeSession(session, session.id)
-            return data
           }
           if (
             type === 'passkey' &&
             'mode' in params &&
             params.mode === 'login'
           ) {
-            await client.apiKeyStamper.resetKeyPair()
-            const generatedPublicKey = await client.apiKeyStamper.getPublicKey()
-            if (!generatedPublicKey) {
-              throw new Error('Failed to get public key')
-            }
-            const loginData = await client.loginWithStamp({
-              targetPublicKey: generatedPublicKey,
-              projectId,
-              // Sign against the parent org, not the user's sub-org (see the
-              // refreshSession note). The backend derives the sub-org from the
-              // stamped passkey credential.
-              organizationId: await resolveOrganizationId(),
-              stampWith: 'passkey',
+            return withKeyTransition(async () => {
+              const targetPublicKey = await prepareFreshKeyRotation()
+              try {
+                const loginData = await client.loginWithStamp({
+                  targetPublicKey,
+                  projectId,
+                  // Sign against the parent org, not the user's sub-org (see the
+                  // refreshSession note). The backend derives the sub-org from the
+                  // stamped passkey credential.
+                  organizationId: await resolveOrganizationId(),
+                  stampWith: 'passkey',
+                })
+                const session = createReplacementSession(
+                  loginData.session,
+                  targetPublicKey,
+                  'indexedDb',
+                )
+                await commitReplacementSession(session, targetPublicKey)
+                return loginData
+              } catch (error) {
+                await client.apiKeyStamper.discardKeyRotation()
+                throw error
+              }
             })
-            const parsedSession = parseSession(loginData.session)
-            const session: ZeroDevWalletSession = {
-              id: `session_indexedDb_${Date.now()}`,
-              stamperType: 'apiKey',
-              createdAt: Date.now(),
-              sessionType: SessionType.READ_WRITE,
-              userId: parsedSession.userId,
-              organizationId: parsedSession.organizationId,
-              expiry:
-                Date.now() +
-                Number(DEFAULT_SESSION_EXPIRATION_IN_SECONDS) * 1000,
-              token: loginData.session,
-            }
-            await sessionStorageManager.storeSession(session, session.id)
-            return loginData
           }
           throw new Error('Passkey authentication requires passkey parameter')
         }
@@ -369,69 +505,64 @@ export async function createZeroDevWalletCore(
           }
 
           if (otpParams.mode === 'verifyOtp') {
-            const { otpId, otpCode, otpEncryptionTargetBundle } = otpParams
-
-            // Step 1: Generate new key pair
-            await client.apiKeyStamper.resetKeyPair()
-            const targetPublicKey = await client.apiKeyStamper.getPublicKey()
-
-            if (!targetPublicKey) {
-              throw new Error('Failed to get public key')
-            }
-
-            // Step 2a: HPKE-seal the OTP attempt to the enclave's per-session
-            // target key. The auth proxy never sees the plaintext OTP code.
-            const encryptedOtpBundle = await encryptOtpAttempt({
-              otpCode,
-              publicKey: targetPublicKey,
-              encryptionTargetBundle: otpEncryptionTargetBundle,
-            })
-
-            // Step 2b: Verify OTP via Auth Proxy
-            if (!cachedAuthProxyConfigId) {
-              const { authProxyConfigId } = await client.getAuthProxyConfigId()
-              cachedAuthProxyConfigId = authProxyConfigId
-            }
-            const authProxyClient = createAuthProxyClient({
-              authProxyConfigId: cachedAuthProxyConfigId,
-            })
-
-            const { verificationToken } = await authProxyClient.verifyOtp({
-              otpId,
-              encryptedOtpBundle,
-            })
-
-            // Step 3: Build client signature
-            const clientSignature = await buildClientSignature({
-              verificationToken,
-              publicKey: targetPublicKey,
-              stamper: client.apiKeyStamper,
-            })
-
-            // Step 4: Login via backend (not Auth Proxy!)
-            const data = await client.loginWithOTP({
-              verificationToken,
-              clientSignature,
-              projectId,
-            })
-
-            if (data.session) {
-              // Parse the JWT to get session data
-              const parsedSession = parseSession(data.session)
-              const session: ZeroDevWalletSession = {
-                id: `session_otp_${Date.now()}`,
-                userId: parsedSession.userId,
-                organizationId: parsedSession.organizationId,
-                stamperType: 'apiKey',
-                sessionType:
-                  parsedSession.sessionType || SessionType.READ_WRITE,
-                token: data.session,
-                expiry: parsedSession.expiry,
-                createdAt: Date.now(),
+            return withKeyTransition(async () => {
+              const { otpId, otpCode, otpEncryptionTargetBundle } = otpParams
+              const targetPublicKey = await prepareFreshKeyRotation()
+
+              try {
+                // HPKE-seal the OTP attempt to the enclave's per-session target
+                // key. The auth proxy never sees the plaintext OTP code.
+                const encryptedOtpBundle = await encryptOtpAttempt({
+                  otpCode,
+                  publicKey: targetPublicKey,
+                  encryptionTargetBundle: otpEncryptionTargetBundle,
+                })
+
+                if (!cachedAuthProxyConfigId) {
+                  const { authProxyConfigId } =
+                    await client.getAuthProxyConfigId()
+                  cachedAuthProxyConfigId = authProxyConfigId
+                }
+                const authProxyClient = createAuthProxyClient({
+                  authProxyConfigId: cachedAuthProxyConfigId,
+                })
+
+                const { verificationToken } = await authProxyClient.verifyOtp({
+                  otpId,
+                  encryptedOtpBundle,
+                })
+
+                const clientSignature = await buildClientSignature({
+                  verificationToken,
+                  publicKey: targetPublicKey,
+                  stamper: {
+                    stamp: (payload) =>
+                      client.apiKeyStamper.stampPending(payload),
+                  },
+                })
+
+                const data = await client.loginWithOTP({
+                  verificationToken,
+                  clientSignature,
+                  projectId,
+                })
+
+                if (!data.session) {
+                  await client.apiKeyStamper.discardKeyRotation()
+                  return data
+                }
+                const session = createReplacementSession(
+                  data.session,
+                  targetPublicKey,
+                  'otp',
+                )
+                await commitReplacementSession(session, targetPublicKey)
+                return data
+              } catch (error) {
+                await client.apiKeyStamper.discardKeyRotation()
+                throw error
               }
-              await sessionStorageManager.storeSession(session, session.id)
-            }
-            return data
+            })
           }
 
           throw new Error('OTP authentication requires mode parameter')
@@ -441,10 +572,67 @@ export async function createZeroDevWalletCore(
       }
     },
 
-    async logout() {
-      await sessionStorageManager.clearAllSessions()
-      await client.apiKeyStamper.resetKeyPair()
-      return true
+    async logout(options) {
+      return withKeyTransition(async () => {
+        const session = await sessionStorageManager.getActiveSession()
+        if (session) {
+          const publicKey = await client.apiKeyStamper.getPublicKey()
+          if (publicKey) {
+            try {
+              const authenticators = await client.getAuthenticators({
+                subOrganizationId: session.organizationId,
+                projectId,
+                token: session.token,
+              })
+              const sessionKey = authenticators.sessionKeys?.find((key) => {
+                const candidate = key.ApiKey ?? key.apiKey
+                return (
+                  candidate !== undefined &&
+                  normalizeKey(candidate) === normalizeKey(publicKey)
+                )
+              })
+              const apiKeyId = sessionKey?.TurnkeyId ?? sessionKey?.turnkeyId
+              if (!apiKeyId) {
+                throw new Error(
+                  'Current session key was not found in the authenticated key list; refusing to erase local credentials without confirmed remote revocation.',
+                )
+              }
+              await client.logout({
+                projectId,
+                organizationId: session.organizationId,
+                userId: session.userId,
+                apiKeyId,
+              })
+            } catch (error) {
+              // Only a confirmed rejection (401) or an explicit `force` clears
+              // local credentials below. Any other failure (403, 5xx, network)
+              // is ambiguous — erasing the only key on an inconclusive signal
+              // could lock the user out unrecoverably — so preserve it and
+              // return without throwing, so callers (e.g. a logout button)
+              // don't crash. Pass `logout({ force: true })` to clear locally
+              // regardless when the caller knows the remote is gone.
+              const credentialIsRejected =
+                error instanceof RestRequestError && error.status === 401
+              if (!options?.force && !credentialIsRejected) return false
+            }
+          }
+        }
+
+        preparedOAuthPublicKey = undefined
+        let cleanupError: unknown
+        try {
+          await client.apiKeyStamper.clear()
+        } catch (error) {
+          cleanupError = error
+        }
+        try {
+          await sessionStorageManager.clearAllSessions()
+        } catch (error) {
+          cleanupError ??= error
+        }
+        if (cleanupError) throw cleanupError
+        return true
+      })
     },
 
     async toAccount(): Promise<LocalAccount> {