Warp/apps/web/src/stores/session.ts

import { HttpApiEvent, MatrixError, type AuthDict, type LoginResponse, type RegisterResponse } from 'matrix-js-sdk'
import { defineStore } from 'pinia'
import { computed, ref } from 'vue'
import { config } from '@/config'
import { anonymousClient, client, destroyClient, newStorageKey, startClient, stopClient, type Credentials } from '@/matrix/client'
import { cryptoDatabasePrefix } from '@/matrix/crypto'
import { clearMediaCache } from '@/matrix/media'

// The person's Matrix session. It is kept in localStorage while "Stay
// signed in" is on (the default for a chat client), otherwise in
// sessionStorage, which goes with the tab. Signing out logs the session out
// on the homeserver and forgets it here, its encryption keys included.

const storageKey = 'warp.web.session'
const deviceName = 'Warp Web'

function load(): Credentials | null {
  for (const store of [localStorage, sessionStorage]) {
    try {
      const raw = store.getItem(storageKey)
      if (raw) return JSON.parse(raw) as Credentials
    } catch {
      // Unreadable or blocked: treat as signed out.
    }
  }
  return null
}

/** The crypto stores of the sessions this browser has credentials for, in either storage. */
export function savedCryptoStores(): string[] {
  const out: string[] = []
  for (const store of [localStorage, sessionStorage]) {
    try {
      const raw = store.getItem(storageKey)
      if (!raw) continue
      const c = JSON.parse(raw) as Credentials
      if (c.userId && c.deviceId) out.push(cryptoDatabasePrefix(c.userId, c.deviceId))
    } catch {
      // Unreadable: nothing to keep.
    }
  }
  return out
}

function save(c: Credentials | null, remember: boolean) {
  for (const store of [localStorage, sessionStorage]) {
    try {
      store.removeItem(storageKey)
    } catch {
      // Ignore.
    }
  }
  if (!c) return
  try {
    ;(remember ? localStorage : sessionStorage).setItem(storageKey, JSON.stringify(c))
  } catch {
    // Not kept across reloads; the session still works in this page.
  }
}

/** What registration on this homeserver asks for. */
export interface RegistrationInfo {
  enabled: boolean
  /** Every flow on offer needs a registration token. */
  tokenRequired: boolean
  /** Some flow takes a token, so offering the field makes sense. */
  tokenAccepted: boolean
  /** The server's words when it's off. */
  reason?: string
}

interface UiaFlow {
  stages: string[]
}

interface UiaState {
  session?: string
  flows?: UiaFlow[]
  completed?: string[]
}

function uiaState(e: unknown): UiaState | null {
  if (e instanceof MatrixError && e.httpStatus === 401 && Array.isArray(e.data?.flows)) {
    return e.data as UiaState
  }
  return null
}

const supportedStages = new Set(['m.login.dummy', 'm.login.registration_token'])

export const useSession = defineStore('session', () => {
  const saved = load()
  const userId = ref(saved?.userId ?? '')
  const deviceId = ref(saved?.deviceId ?? '')
  const phase = ref<'signedOut' | 'starting' | 'ready'>(saved ? 'starting' : 'signedOut')
  /** Set when the session ended on its own, shown on the sign-in page. */
  const endedReason = ref('')

  const signedIn = computed(() => phase.value !== 'signedOut')
  const serverName = computed(() => userId.value.slice(userId.value.indexOf(':') + 1))

  function clear(reason = ''): Promise<void> {
    // The session is gone for good: its keys in IndexedDB go with it.
    const done = userId.value && deviceId.value ? destroyClient({ userId: userId.value, deviceId: deviceId.value }) : Promise.resolve(stopClient())
    clearMediaCache()
    userId.value = ''
    deviceId.value = ''
    phase.value = 'signedOut'
    endedReason.value = reason
    save(null, false)
    return done
  }

  async function begin(creds: Credentials) {
    userId.value = creds.userId
    deviceId.value = creds.deviceId
    phase.value = 'starting'
    const c = await startClient(creds)
    c.on(HttpApiEvent.SessionLoggedOut, () => clear('Your session ended. Sign in again to continue.'))
    phase.value = 'ready'
  }

  /** Picks the saved session back up on page load. */
  async function restore() {
    if (!saved || phase.value === 'ready') return
    try {
      await begin(saved)
    } catch (e) {
      if (e instanceof MatrixError && e.errcode === 'M_UNKNOWN_TOKEN') {
        clear('Your session ended. Sign in again to continue.')
        return
      }
      // Anything else (the homeserver unreachable, say) keeps the session
      // and its encryption keys: the shell offers to try again.
      stopClient()
      throw e
    }
  }

  async function adopt(res: LoginResponse | RegisterResponse, keep: boolean) {
    if (!res.access_token || !res.device_id) throw new Error('The homeserver didn’t start a session.')
    const creds: Credentials = {
      baseUrl: config.homeserverUrl,
      userId: res.user_id,
      deviceId: res.device_id,
      accessToken: res.access_token,
      // Kept only for the tab: its encryption keys are unreadable once it closes.
      storageKey: keep ? undefined : newStorageKey(),
    }
    save(creds, keep)
    await begin(creds)
  }

  async function signIn(username: string, password: string, keep: boolean) {
    endedReason.value = ''
    const res = await anonymousClient().loginRequest({
      type: 'm.login.password',
      identifier: { type: 'm.id.user', user: username.trim() },
      password,
      initial_device_display_name: deviceName,
    })
    await adopt(res, keep)
  }

  async function registrationInfo(): Promise<RegistrationInfo> {
    try {
      await anonymousClient().registerRequest({})
      // A server that registers without any flow at all.
      return { enabled: true, tokenRequired: false, tokenAccepted: false }
    } catch (e) {
      const uia = uiaState(e)
      if (uia) {
        const flows = (uia.flows ?? []).filter((f) => f.stages.every((s) => supportedStages.has(s)))
        if (!flows.length) {
          return {
            enabled: false,
            tokenRequired: false,
            tokenAccepted: false,
            reason: 'This homeserver asks for a sign-up step this client doesn’t support yet.',
          }
        }
        const withToken = flows.filter((f) => f.stages.includes('m.login.registration_token'))
        return { enabled: true, tokenRequired: withToken.length === flows.length, tokenAccepted: withToken.length > 0 }
      }
      if (e instanceof MatrixError && e.httpStatus === 403) {
        return { enabled: false, tokenRequired: false, tokenAccepted: false, reason: e.data?.error }
      }
      throw e
    }
  }

  /** Registers with the user-interactive auth flow: dummy and token stages. */
  async function register(username: string, password: string, token: string, keep: boolean) {
    endedReason.value = ''
    const c = anonymousClient()
    const body = {
      username: username.trim(),
      password,
      initial_device_display_name: deviceName,
    }
    let auth: AuthDict | undefined
    for (let step = 0; step < 6; step++) {
      try {
        const res = await c.registerRequest({ ...body, auth })
        await adopt(res, keep)
        return
      } catch (e) {
        const uia = uiaState(e)
        if (!uia) throw e
        // A stage that failed comes back with its error alongside the flows.
        if (auth && e instanceof MatrixError && e.data?.errcode) {
          throw new Error(e.data.error ?? 'The homeserver refused that sign-up step.')
        }
        const done = new Set(uia.completed ?? [])
        const flows = (uia.flows ?? []).filter((f) => f.stages.every((s) => supportedStages.has(s)))
        const flow =
          flows.find((f) => (token ? f.stages.includes('m.login.registration_token') : !f.stages.includes('m.login.registration_token'))) ??
          flows[0]
        if (!flow) throw new Error('This homeserver asks for a sign-up step this client doesn’t support yet.')
        const next = flow.stages.find((s) => !done.has(s))
        if (!next) throw e
        if (next === 'm.login.registration_token') {
          if (!token) throw new Error('This homeserver needs a registration token to sign up.')
          auth = { type: next, token, session: uia.session }
        } else {
          auth = { type: next, session: uia.session }
        }
      }
    }
    throw new Error('Registration didn’t complete. Try again.')
  }

  async function signOut() {
    try {
      await client().logout(true)
    } catch {
      // The token may already be gone; forget it either way.
    }
    await clear()
    // A fresh page drops the crypto stack's memory and the store it still
    // holds open, so the store is deleted now rather than on the next visit.
    window.location.replace('/login')
  }

  return {
    userId,
    deviceId,
    serverName,
    phase,
    signedIn,
    endedReason,
    restore,
    signIn,
    register,
    registrationInfo,
    signOut,
    clear,
  }
})