From 691239d23044d22426b9cc9a60f94ba5b827a869 Mon Sep 17 00:00:00 2001 From: Eric Bailey Date: Fri, 28 Aug 2026 08:46:30 -0500 Subject: [PATCH] Yeet --- src/state/persisted/__tests__/session-test.ts | 250 ++++++++++ src/state/persisted/index.ts | 76 ++- src/state/persisted/index.web.ts | 147 ++++-- src/state/persisted/schema.ts | 9 + src/state/persisted/session-lock.ts | 18 + src/state/persisted/session-lock.web.ts | 37 ++ src/state/persisted/session.ts | 294 +++++++++++ src/state/persisted/types.ts | 9 + .../session/__tests__/provider-abort-test.tsx | 4 + .../__tests__/provider-clients-test.tsx | 4 + .../provider-refresh-session-test.tsx | 4 + .../provider-session-events-test.tsx | 104 +++- .../session/__tests__/session-core-test.ts | 21 +- src/state/session/index.tsx | 462 ++++++++++++------ src/state/session/reducer.ts | 8 +- src/state/session/session-core.ts | 30 +- 16 files changed, 1264 insertions(+), 213 deletions(-) create mode 100644 src/state/persisted/__tests__/session-test.ts create mode 100644 src/state/persisted/session-lock.ts create mode 100644 src/state/persisted/session-lock.web.ts create mode 100644 src/state/persisted/session.ts diff --git a/src/state/persisted/__tests__/session-test.ts b/src/state/persisted/__tests__/session-test.ts new file mode 100644 index 0000000000..2dafc87a4d --- /dev/null +++ b/src/state/persisted/__tests__/session-test.ts @@ -0,0 +1,250 @@ +import {describe, expect, it} from '@jest/globals' + +import {type PersistedAccount, type Schema} from '../schema' +import { + applySessionUpdate, + getCredentialState, + type SessionCredentialMutation, +} from '../session' + +const DID = 'did:plc:example123' + +function jwt({jti, issuedAt}: {jti: string; issuedAt: number}) { + const encode = (value: object) => + Buffer.from(JSON.stringify(value)).toString('base64url') + return `${encode({alg: 'none'})}.${encode({jti, iat: issuedAt})}.signature` +} + +function account({ + refreshJwt, + accessJwt = `access-${refreshJwt}`, +}: { + refreshJwt?: string + accessJwt?: string +}): PersistedAccount { + return { + service: 'https://bsky.social/', + did: DID, + handle: 'alice.test', + refreshJwt, + accessJwt, + } +} + +function session({ + account: sessionAccount, + credentialStates, +}: { + account?: PersistedAccount + credentialStates?: Schema['session']['credentialStates'] +}): Schema['session'] { + return { + accounts: sessionAccount ? [sessionAccount] : [], + currentAccount: sessionAccount, + credentialStates, + } +} + +function update({ + storedSession, + nextAccount, + mutation, +}: { + storedSession: Schema['session'] + nextAccount?: PersistedAccount + mutation?: SessionCredentialMutation +}) { + return applySessionUpdate({ + storedSession, + nextSession: session({ + account: nextAccount, + credentialStates: storedSession.credentialStates, + }), + credentialMutations: mutation ? [mutation] : [], + }) +} + +describe('versioned persisted sessions', () => { + const refreshA = jwt({jti: 'A', issuedAt: 1}) + const refreshB = jwt({jti: 'B', issuedAt: 2}) + const refreshBAlias = jwt({jti: 'B', issuedAt: 3}) + const refreshC = jwt({jti: 'C', issuedAt: 4}) + + it('advances the credential version when refresh moves to a new jti', () => { + const storedSession = session({account: account({refreshJwt: refreshA})}) + const next = account({refreshJwt: refreshB}) + + const result = update({ + storedSession, + nextAccount: next, + mutation: { + type: 'refresh', + accountDid: DID, + baseRefreshJwt: refreshA, + resultRefreshJwt: refreshB, + }, + }) + + expect(result.accounts[0].refreshJwt).toBe(refreshB) + expect(getCredentialState({session: result, accountDid: DID})).toEqual({ + credentialVersion: 1, + refreshJti: 'B', + status: 'active', + }) + }) + + it('keeps one version for convergent aliases with the same jti', () => { + const generationB = update({ + storedSession: session({account: account({refreshJwt: refreshA})}), + nextAccount: account({refreshJwt: refreshB}), + mutation: { + type: 'refresh', + accountDid: DID, + baseRefreshJwt: refreshA, + resultRefreshJwt: refreshB, + }, + }) + + const result = update({ + storedSession: generationB, + nextAccount: account({refreshJwt: refreshBAlias}), + mutation: { + type: 'refresh', + accountDid: DID, + baseRefreshJwt: refreshA, + resultRefreshJwt: refreshBAlias, + }, + }) + + expect(result.accounts[0].refreshJwt).toBe(refreshB) + expect( + getCredentialState({session: result, accountDid: DID}).credentialVersion, + ).toBe(1) + }) + + it('does not let a stale refresh overwrite a newer generation', () => { + const storedSession = session({ + account: account({refreshJwt: refreshB}), + credentialStates: { + [DID]: {credentialVersion: 8, refreshJti: 'B', status: 'active'}, + }, + }) + + const result = update({ + storedSession, + nextAccount: account({refreshJwt: refreshC}), + mutation: { + type: 'refresh', + accountDid: DID, + baseRefreshJwt: refreshA, + resultRefreshJwt: refreshC, + }, + }) + + expect(result.accounts[0].refreshJwt).toBe(refreshB) + expect( + getCredentialState({session: result, accountDid: DID}).credentialVersion, + ).toBe(8) + }) + + it('patches metadata without replacing authoritative credentials', () => { + const storedSession = session({ + account: account({refreshJwt: refreshB}), + credentialStates: { + [DID]: {credentialVersion: 8, refreshJti: 'B', status: 'active'}, + }, + }) + const staleAccount = { + ...account({refreshJwt: refreshA}), + handle: 'alice-renamed.test', + } + + const result = update({storedSession, nextAccount: staleAccount}) + + expect(result.accounts[0].handle).toBe('alice-renamed.test') + expect(result.accounts[0].refreshJwt).toBe(refreshB) + expect(result.accounts[0].accessJwt).toBe(`access-${refreshB}`) + }) + + it('does not let a stale expiry clear a newer generation', () => { + const storedSession = session({ + account: account({refreshJwt: refreshB}), + credentialStates: { + [DID]: {credentialVersion: 8, refreshJti: 'B', status: 'active'}, + }, + }) + + const result = update({ + storedSession, + nextAccount: account({refreshJwt: undefined, accessJwt: undefined}), + mutation: { + type: 'expire', + accountDid: DID, + baseRefreshJwt: refreshA, + }, + }) + + expect(result.accounts[0].refreshJwt).toBe(refreshB) + expect(getCredentialState({session: result, accountDid: DID}).status).toBe( + 'active', + ) + }) + + it('preserves logout and removal tombstones until an explicit login', () => { + const active = session({ + account: account({refreshJwt: refreshB}), + credentialStates: { + [DID]: {credentialVersion: 8, refreshJti: 'B', status: 'active'}, + }, + }) + const loggedOut = update({ + storedSession: active, + nextAccount: account({refreshJwt: undefined, accessJwt: undefined}), + mutation: {type: 'logout', accountDid: DID}, + }) + + const staleRefresh = update({ + storedSession: loggedOut, + nextAccount: account({refreshJwt: refreshC}), + mutation: { + type: 'refresh', + accountDid: DID, + baseRefreshJwt: refreshB, + resultRefreshJwt: refreshC, + }, + }) + expect(staleRefresh.accounts[0].refreshJwt).toBeUndefined() + expect( + getCredentialState({session: staleRefresh, accountDid: DID}), + ).toEqual({credentialVersion: 9, status: 'logged-out'}) + + const removed = update({ + storedSession: staleRefresh, + mutation: {type: 'remove', accountDid: DID}, + }) + const staleSnapshot = update({ + storedSession: removed, + nextAccount: account({refreshJwt: refreshB}), + }) + expect(staleSnapshot.accounts).toEqual([]) + expect( + getCredentialState({session: staleSnapshot, accountDid: DID}), + ).toEqual({credentialVersion: 10, status: 'removed'}) + + const loggedIn = update({ + storedSession: staleSnapshot, + nextAccount: account({refreshJwt: refreshC}), + mutation: { + type: 'login', + accountDid: DID, + resultRefreshJwt: refreshC, + }, + }) + expect(loggedIn.accounts[0].refreshJwt).toBe(refreshC) + expect(getCredentialState({session: loggedIn, accountDid: DID})).toEqual({ + credentialVersion: 11, + refreshJti: 'C', + status: 'active', + }) + }) +}) diff --git a/src/state/persisted/index.ts b/src/state/persisted/index.ts index 0659e7817b..583e9c3e85 100644 --- a/src/state/persisted/index.ts +++ b/src/state/persisted/index.ts @@ -8,15 +8,20 @@ import { tryStringify, } from '#/state/persisted/schema' import {device} from '#/storage' +import {applySessionUpdate} from './session' +import {runWithSessionCredentialLock} from './session-lock' import {type PersistedApi} from './types' import {normalizeData} from './util' +export type {SessionCredentialMutation} from './session' +export {runWithSessionCredentialLock} from './session-lock' export type {PersistedAccount, Schema} from '#/state/persisted/schema' export {defaults} from '#/state/persisted/schema' const BSKY_STORAGE = 'BSKY_STORAGE' let _state: Schema = defaults +let pendingWrite: Promise = Promise.resolve() export async function init() { const stored = await readFromStorage() @@ -43,18 +48,43 @@ export function readLatest(key: K): Schema[K] { } readLatest satisfies PersistedApi['readLatest'] -export async function write( +export function write( key: K, value: Schema[K], ): Promise { - _state = normalizeData({ - ..._state, - [key]: value, + return enqueueWrite(async () => { + const next = normalizeData({ + ..._state, + [key]: value, + }) + await persistWithRetry(next) + _state = next }) - await writeToStorage(_state) } write satisfies PersistedApi['write'] +export function updateSession({ + nextSession, + credentialMutations, +}: { + nextSession: Schema['session'] + credentialMutations: import('./session').SessionCredentialMutation[] +}): Promise { + return enqueueWrite(async () => { + const session = applySessionUpdate({ + storedSession: _state.session, + nextSession, + credentialMutations, + }) + const next = normalizeData({..._state, session}) + await persistWithRetry(next) + _state = next + return session + }) +} +updateSession satisfies PersistedApi['updateSession'] +runWithSessionCredentialLock satisfies PersistedApi['runWithSessionCredentialLock'] + export function onUpdate( _key: K, _cb: (v: Schema[K]) => void, @@ -73,16 +103,36 @@ export async function clearStorage() { } clearStorage satisfies PersistedApi['clearStorage'] +function enqueueWrite(operation: () => Promise): Promise { + const result = pendingWrite.then(operation, operation) + pendingWrite = result.then( + () => undefined, + () => undefined, + ) + return result +} + +async function persistWithRetry(value: Schema) { + try { + await writeToStorage(value) + } catch { + /* Retry the latest complete snapshot while the process still owns it. */ + await writeToStorage(value) + } +} + async function writeToStorage(value: Schema) { const rawData = tryStringify(value) - if (rawData) { - try { - await AsyncStorage.setItem(BSKY_STORAGE, rawData) - } catch (e) { - logger.error(`persisted state: failed writing root state to storage`, { - message: e, - }) - } + if (!rawData) { + throw new Error('Failed to serialize persisted state') + } + try { + await AsyncStorage.setItem(BSKY_STORAGE, rawData) + } catch (e) { + logger.error(`persisted state: failed writing root state to storage`, { + message: e, + }) + throw e } } diff --git a/src/state/persisted/index.web.ts b/src/state/persisted/index.web.ts index a4c5d1b247..27a1065d9a 100644 --- a/src/state/persisted/index.web.ts +++ b/src/state/persisted/index.web.ts @@ -8,9 +8,17 @@ import { tryParse, tryStringify, } from '#/state/persisted/schema' +import { + applySessionUpdate, + getCredentialState, + type SessionCredentialMutation, +} from './session' +import {runWithSessionCredentialLock} from './session-lock' import {type PersistedApi} from './types' import {normalizeData} from './util' +export type {SessionCredentialMutation} from './session' +export {runWithSessionCredentialLock} from './session-lock' export type {PersistedAccount, Schema} from '#/state/persisted/schema' export {defaults} from '#/state/persisted/schema' @@ -62,38 +70,73 @@ export function readLatest(key: K): Schema[K] { } readLatest satisfies PersistedApi['readLatest'] -// eslint-disable-next-line @typescript-eslint/require-await -export async function write( +export function write( key: K, value: Schema[K], ): Promise { - const next = readFromStorage() - if (next) { - // The storage could have been updated by a different tab before this tab is notified. - // Make sure this write is applied on top of the latest data in the storage as long as it's valid. - _state = next - // Don't fire the update listeners yet to avoid a loop. - // If there was a change, we'll receive the broadcast event soon enough which will do that. - } - try { - if (JSON.stringify({v: _state[key]}) === JSON.stringify({v: value})) { - // Fast path for updates that are guaranteed to be noops. - // This is good mostly because it avoids useless broadcasts to other tabs. - return - } - } catch (e) { - // Ignore and go through the normal path. - } - _state = normalizeData({ - ..._state, - [key]: value, + return runWithSessionCredentialLock({ + accountDids: [], + operation: () => { + const next = readFromStorage() + if (next) { + // The storage could have been updated by a different tab before this tab is notified. + // Make sure this write is applied on top of the latest data in the storage as long as it's valid. + _state = next + // Don't fire the update listeners yet to avoid a loop. + // If there was a change, we'll receive the broadcast event soon enough which will do that. + } + try { + if (JSON.stringify({v: _state[key]}) === JSON.stringify({v: value})) { + // Fast path for updates that are guaranteed to be noops. + // This is good mostly because it avoids useless broadcasts to other tabs. + return + } + } catch (e) { + // Ignore and go through the normal path. + } + const updated = normalizeData({ + ..._state, + [key]: value, + }) + writeToStorage(updated) + _state = updated + broadcastUpdate({key}) + }, }) - writeToStorage(_state) - broadcast.postMessage({event: {type: UPDATE_EVENT, key}}) - broadcast.postMessage({event: UPDATE_EVENT}) // Backcompat while upgrading } write satisfies PersistedApi['write'] +// eslint-disable-next-line @typescript-eslint/require-await +export async function updateSession({ + nextSession, + credentialMutations, +}: { + nextSession: Schema['session'] + credentialMutations: SessionCredentialMutation[] +}): Promise { + const stored = readFromStorage() ?? _state + const session = applySessionUpdate({ + storedSession: stored.session, + nextSession, + credentialMutations, + }) + const updated = normalizeData({...stored, session}) + writeToStorage(updated) + _state = updated + + const invalidations = credentialMutations.map(mutation => ({ + accountDid: mutation.accountDid, + credentialVersion: getCredentialState({ + session, + accountDid: mutation.accountDid, + }).credentialVersion, + })) + broadcastUpdate({key: 'session', invalidations}) + return session +} +updateSession satisfies PersistedApi['updateSession'] +runWithSessionCredentialLock satisfies PersistedApi['runWithSessionCredentialLock'] + export function onUpdate( key: K, cb: (v: Schema[K]) => void, @@ -108,13 +151,17 @@ export function onUpdate( } onUpdate satisfies PersistedApi['onUpdate'] -// eslint-disable-next-line @typescript-eslint/require-await -export async function clearStorage() { - try { - localStorage.removeItem(BSKY_STORAGE) - } catch (e: any) { - // Expected on the web in private mode. - } +export function clearStorage(): Promise { + return runWithSessionCredentialLock({ + accountDids: [], + operation: () => { + try { + localStorage.removeItem(BSKY_STORAGE) + } catch (e: any) { + // Expected on the web in private mode. + } + }, + }) } clearStorage satisfies PersistedApi['clearStorage'] @@ -156,14 +203,40 @@ async function onBroadcastMessage({data}: MessageEvent) { } } +function broadcastUpdate({ + key, + invalidations = [], +}: { + key: keyof Schema + invalidations?: {accountDid: string; credentialVersion: number}[] +}) { + if (invalidations.length > 0) { + for (const invalidation of invalidations) { + broadcast.postMessage({ + event: {type: UPDATE_EVENT, key, ...invalidation}, + }) + } + } else { + broadcast.postMessage({event: {type: UPDATE_EVENT, key}}) + } + broadcast.postMessage({event: UPDATE_EVENT}) // Backcompat while upgrading +} + function writeToStorage(value: Schema) { const rawData = tryStringify(value) - if (rawData) { - try { - localStorage.setItem(BSKY_STORAGE, rawData) - } catch (e) { - // Expected on the web in private mode. + if (!rawData) { + throw new Error('Failed to serialize persisted state') + } + try { + localStorage.setItem(BSKY_STORAGE, rawData) + if (localStorage.getItem(BSKY_STORAGE) !== rawData) { + throw new Error('Failed to verify persisted state') } + } catch (error) { + logger.error('persisted state: failed writing root state to storage', { + message: error, + }) + throw error } } diff --git a/src/state/persisted/schema.ts b/src/state/persisted/schema.ts index e5a3c0b372..dab551281a 100644 --- a/src/state/persisted/schema.ts +++ b/src/state/persisted/schema.ts @@ -57,12 +57,21 @@ const currentAccountSchema = accountSchema.extend({ }) export type PersistedCurrentAccount = z.infer +const credentialStateSchema = z.object({ + credentialVersion: z.number().int().nonnegative(), + refreshJti: z.string().optional(), + status: z.enum(['active', 'logged-out', 'removed']), +}) +export type PersistedCredentialState = z.infer + const schema = z.object({ colorMode: z.enum(['system', 'light', 'dark']), darkTheme: z.enum(['dim', 'dark']).optional(), session: z.object({ accounts: z.array(accountSchema), currentAccount: currentAccountSchema.optional(), + /** Per-account credential ordering, including logout and removal tombstones. */ + credentialStates: z.record(z.string(), credentialStateSchema).optional(), }), reminders: z.object({ lastEmailConfirm: z.string().optional(), diff --git a/src/state/persisted/session-lock.ts b/src/state/persisted/session-lock.ts new file mode 100644 index 0000000000..bd8ab5fc8d --- /dev/null +++ b/src/state/persisted/session-lock.ts @@ -0,0 +1,18 @@ +export function runWithSessionCredentialLock({ + accountDids, + operation, +}: { + accountDids: string[] + operation: () => T | Promise +}): Promise { + void accountDids + try { + return Promise.resolve(operation()) + } catch (error) { + return Promise.reject( + error instanceof Error + ? error + : new Error('Session credential operation failed', {cause: error}), + ) + } +} diff --git a/src/state/persisted/session-lock.web.ts b/src/state/persisted/session-lock.web.ts new file mode 100644 index 0000000000..26ce495cf7 --- /dev/null +++ b/src/state/persisted/session-lock.web.ts @@ -0,0 +1,37 @@ +const PERSISTED_STORAGE_LOCK = 'bsky-persisted-storage' +const SESSION_LOCK_PREFIX = 'bsky-session:' + +export function runWithSessionCredentialLock({ + accountDids, + operation, +}: { + accountDids: string[] + operation: () => T | Promise +}): Promise { + const lockManager = navigator.locks + if (!lockManager) { + try { + return Promise.resolve(operation()) + } catch (error) { + return Promise.reject( + error instanceof Error + ? error + : new Error('Session credential operation failed', {cause: error}), + ) + } + } + + const accountLockNames = [...new Set(accountDids)] + .sort() + .map(did => `${SESSION_LOCK_PREFIX}${did}`) + /* All values share one localStorage blob, so every write also takes its lock. */ + const lockNames = [PERSISTED_STORAGE_LOCK, ...accountLockNames] + + const run = (index: number): Promise => { + const lockName = lockNames[index] + if (!lockName) return Promise.resolve(operation()) + return lockManager.request(lockName, () => run(index + 1)) + } + + return run(0) +} diff --git a/src/state/persisted/session.ts b/src/state/persisted/session.ts new file mode 100644 index 0000000000..1d596a0dfa --- /dev/null +++ b/src/state/persisted/session.ts @@ -0,0 +1,294 @@ +import {jwtDecode} from 'jwt-decode' + +import { + type PersistedAccount, + type PersistedCredentialState, + type Schema, +} from './schema' + +export type SessionCredentialMutation = + | { + type: 'refresh' + accountDid: string + baseRefreshJwt: string | undefined + resultRefreshJwt: string | undefined + } + | { + type: 'expire' + accountDid: string + baseRefreshJwt: string | undefined + } + | { + type: 'login' + accountDid: string + resultRefreshJwt: string | undefined + } + | { + type: 'logout' + accountDid: string + } + | { + type: 'remove' + accountDid: string + } + +const CREDENTIAL_FIELDS = new Set([ + 'accessJwt', + 'refreshJwt', +]) + +/** Read the server-side refresh-token generation, tolerating legacy bad data. */ +export function getRefreshJti({ + refreshJwt, +}: { + refreshJwt: string | undefined +}): string | undefined { + if (!refreshJwt) return undefined + try { + const decoded = jwtDecode(refreshJwt) + if (typeof decoded.jti === 'string') return decoded.jti + } catch {} + /* A malformed legacy token still needs stable local identity. */ + return refreshJwt +} + +export function getCredentialState({ + session, + accountDid, +}: { + session: Schema['session'] + accountDid: string +}): PersistedCredentialState { + const stored = session.credentialStates?.[accountDid] + if (stored) return stored + const account = session.accounts.find( + candidate => candidate.did === accountDid, + ) + return { + credentialVersion: 0, + refreshJti: getRefreshJti({refreshJwt: account?.refreshJwt}), + status: account?.refreshJwt ? 'active' : 'logged-out', + } +} + +/** + * Apply a session snapshot to the latest persisted session without allowing its + * credential fields to overwrite newer credential generations. + */ +export function applySessionUpdate({ + storedSession, + nextSession, + credentialMutations, +}: { + storedSession: Schema['session'] + nextSession: Schema['session'] + credentialMutations: SessionCredentialMutation[] +}): Schema['session'] { + const incomingByDid = new Map( + nextSession.accounts.map(account => [account.did, account]), + ) + const storedByDid = new Map( + storedSession.accounts.map(account => [account.did, account]), + ) + + /* Incoming metadata wins, but stored credentials remain authoritative. */ + const accounts = nextSession.accounts.map(incoming => { + const stored = storedByDid.get(incoming.did) + return stored + ? mergeAccountMetadata({storedAccount: stored, incomingAccount: incoming}) + : incoming + }) + for (const stored of storedSession.accounts) { + if (!incomingByDid.has(stored.did)) accounts.push(stored) + } + + const credentialStates: Record = { + ...storedSession.credentialStates, + } + for (const account of storedSession.accounts) { + credentialStates[account.did] = getCredentialState({ + session: storedSession, + accountDid: account.did, + }) + } + + const result: Schema['session'] = { + accounts, + currentAccount: nextSession.currentAccount, + credentialStates, + } + + for (const mutation of credentialMutations) { + applyCredentialMutation({session: result, nextSession, mutation}) + } + + result.accounts = result.accounts + .filter( + account => + getCredentialState({session: result, accountDid: account.did}) + .status !== 'removed', + ) + .map(account => { + const credentialState = getCredentialState({ + session: result, + accountDid: account.did, + }) + return credentialState.status === 'logged-out' + ? {...account, accessJwt: undefined, refreshJwt: undefined} + : account + }) + + const currentDid = result.currentAccount?.did + const currentAccount = currentDid + ? result.accounts.find(account => account.did === currentDid) + : undefined + const currentCredentialState = currentDid + ? getCredentialState({session: result, accountDid: currentDid}) + : undefined + result.currentAccount = + currentAccount && currentCredentialState?.status === 'active' + ? currentAccount + : undefined + + return result +} + +function applyCredentialMutation({ + session, + nextSession, + mutation, +}: { + session: Schema['session'] + nextSession: Schema['session'] + mutation: SessionCredentialMutation +}) { + const previousState = getCredentialState({ + session, + accountDid: mutation.accountDid, + }) + const nextAccount = nextSession.accounts.find( + account => account.did === mutation.accountDid, + ) + + switch (mutation.type) { + case 'refresh': { + const baseRefreshJti = getRefreshJti({ + refreshJwt: mutation.baseRefreshJwt, + }) + const resultRefreshJti = getRefreshJti({ + refreshJwt: mutation.resultRefreshJwt, + }) + if ( + previousState.status !== 'active' || + previousState.refreshJti !== baseRefreshJti + ) { + /* The stored generation already advanced or became a tombstone. */ + return + } + if (!nextAccount || !resultRefreshJti) return + replaceAccount({session, account: nextAccount}) + session.credentialStates![mutation.accountDid] = { + credentialVersion: + resultRefreshJti === previousState.refreshJti + ? previousState.credentialVersion + : previousState.credentialVersion + 1, + refreshJti: resultRefreshJti, + status: 'active', + } + return + } + case 'expire': { + const baseRefreshJti = getRefreshJti({ + refreshJwt: mutation.baseRefreshJwt, + }) + if ( + previousState.status !== 'active' || + previousState.refreshJti !== baseRefreshJti + ) { + return + } + clearAccountCredentials({session, accountDid: mutation.accountDid}) + session.credentialStates![mutation.accountDid] = { + credentialVersion: previousState.credentialVersion + 1, + status: 'logged-out', + } + return + } + case 'login': { + const resultRefreshJti = getRefreshJti({ + refreshJwt: mutation.resultRefreshJwt, + }) + if (!nextAccount || !resultRefreshJti) return + replaceAccount({session, account: nextAccount}) + session.credentialStates![mutation.accountDid] = { + credentialVersion: previousState.credentialVersion + 1, + refreshJti: resultRefreshJti, + status: 'active', + } + return + } + case 'logout': { + clearAccountCredentials({session, accountDid: mutation.accountDid}) + session.credentialStates![mutation.accountDid] = { + credentialVersion: previousState.credentialVersion + 1, + status: 'logged-out', + } + return + } + case 'remove': { + session.accounts = session.accounts.filter( + account => account.did !== mutation.accountDid, + ) + session.credentialStates![mutation.accountDid] = { + credentialVersion: previousState.credentialVersion + 1, + status: 'removed', + } + return + } + } +} + +function mergeAccountMetadata({ + storedAccount, + incomingAccount, +}: { + storedAccount: PersistedAccount + incomingAccount: PersistedAccount +}): PersistedAccount { + const merged = {...storedAccount} + for (const key of Object.keys( + incomingAccount, + ) as (keyof PersistedAccount)[]) { + if (!CREDENTIAL_FIELDS.has(key)) { + Object.assign(merged, {[key]: incomingAccount[key]}) + } + } + return merged +} + +function replaceAccount({ + session, + account, +}: { + session: Schema['session'] + account: PersistedAccount +}) { + session.accounts = [ + account, + ...session.accounts.filter(candidate => candidate.did !== account.did), + ] +} + +function clearAccountCredentials({ + session, + accountDid, +}: { + session: Schema['session'] + accountDid: string +}) { + session.accounts = session.accounts.map(account => + account.did === accountDid + ? {...account, accessJwt: undefined, refreshJwt: undefined} + : account, + ) +} diff --git a/src/state/persisted/types.ts b/src/state/persisted/types.ts index 4b5a29d938..fef9921ca0 100644 --- a/src/state/persisted/types.ts +++ b/src/state/persisted/types.ts @@ -1,4 +1,5 @@ import {type Schema} from './schema' +import {type SessionCredentialMutation} from './session' export type PersistedApi = { init(): Promise @@ -13,6 +14,14 @@ export type PersistedApi = { */ readLatest(key: K): Schema[K] write(key: K, value: Schema[K]): Promise + updateSession(params: { + nextSession: Schema['session'] + credentialMutations: SessionCredentialMutation[] + }): Promise + runWithSessionCredentialLock(params: { + accountDids: string[] + operation: () => T | Promise + }): Promise onUpdate( key: K, cb: (v: Schema[K]) => void, diff --git a/src/state/session/__tests__/provider-abort-test.tsx b/src/state/session/__tests__/provider-abort-test.tsx index fdc1585981..ad91dc7e38 100644 --- a/src/state/session/__tests__/provider-abort-test.tsx +++ b/src/state/session/__tests__/provider-abort-test.tsx @@ -14,6 +14,10 @@ jest.mock('#/state/persisted', () => { defaults, get: (key: keyof typeof defaults) => defaults[key], write: () => Promise.resolve(), + updateSession: ({nextSession}: {nextSession: typeof defaults.session}) => + Promise.resolve(nextSession), + runWithSessionCredentialLock: ({operation}: {operation: () => unknown}) => + Promise.resolve(operation()), readLatest: (key: keyof typeof defaults) => defaults[key], onUpdate: () => () => {}, } diff --git a/src/state/session/__tests__/provider-clients-test.tsx b/src/state/session/__tests__/provider-clients-test.tsx index d10af7a12b..57d9610efc 100644 --- a/src/state/session/__tests__/provider-clients-test.tsx +++ b/src/state/session/__tests__/provider-clients-test.tsx @@ -18,6 +18,10 @@ jest.mock('#/state/persisted', () => { defaults, get: (key: keyof typeof defaults) => defaults[key], write: () => Promise.resolve(), + updateSession: ({nextSession}: {nextSession: typeof defaults.session}) => + Promise.resolve(nextSession), + runWithSessionCredentialLock: ({operation}: {operation: () => unknown}) => + Promise.resolve(operation()), readLatest: (key: keyof typeof defaults) => defaults[key], onUpdate: () => () => {}, } diff --git a/src/state/session/__tests__/provider-refresh-session-test.tsx b/src/state/session/__tests__/provider-refresh-session-test.tsx index 2b7bd1c4cb..b2c9a4e41e 100644 --- a/src/state/session/__tests__/provider-refresh-session-test.tsx +++ b/src/state/session/__tests__/provider-refresh-session-test.tsx @@ -17,6 +17,10 @@ jest.mock('#/state/persisted', () => { defaults, get: (key: keyof typeof defaults) => defaults[key], write: () => Promise.resolve(), + updateSession: ({nextSession}: {nextSession: typeof defaults.session}) => + Promise.resolve(nextSession), + runWithSessionCredentialLock: ({operation}: {operation: () => unknown}) => + Promise.resolve(operation()), readLatest: (key: keyof typeof defaults) => defaults[key], onUpdate: () => () => {}, } diff --git a/src/state/session/__tests__/provider-session-events-test.tsx b/src/state/session/__tests__/provider-session-events-test.tsx index cd7873b1ba..f47d7aa0c2 100644 --- a/src/state/session/__tests__/provider-session-events-test.tsx +++ b/src/state/session/__tests__/provider-session-events-test.tsx @@ -39,6 +39,27 @@ jest.mock('#/state/persisted', () => { ? mockPersisted.latest : defaults[key as keyof typeof defaults], write: () => Promise.resolve(), + updateSession: ({ + nextSession, + credentialMutations, + }: { + nextSession: Schema['session'] + credentialMutations: import('#/state/persisted/session').SessionCredentialMutation[] + }) => { + const { + applySessionUpdate, + }: typeof import('#/state/persisted/session') = require('#/state/persisted/session') + const committed = applySessionUpdate({ + storedSession: mockPersisted.latest, + nextSession, + credentialMutations, + }) + mockPersisted.session = committed + mockPersisted.latest = committed + return Promise.resolve(committed) + }, + runWithSessionCredentialLock: ({operation}: {operation: () => unknown}) => + Promise.resolve(operation()), onUpdate: (_key: string, cb: (value: Schema['session']) => void) => { mockPersistedListeners.push(cb) return () => {} @@ -308,7 +329,7 @@ describe('expiry rescue', () => { mockPersisted.latest = {accounts: [fresher], currentAccount: fresher} act(() => { - onSessionChange( + void onSessionChange( bundle as unknown as SessionBundle, DID, 'expired', @@ -337,7 +358,7 @@ describe('expiry rescue', () => { mockPersisted.latest = {accounts: [account], currentAccount: account} act(() => { - onSessionChange( + void onSessionChange( bundle as unknown as SessionBundle, DID, 'expired', @@ -365,7 +386,7 @@ describe('expiry rescue', () => { /* generation 1 dies and is rescued onto generation 2 */ act(() => { - onSessionChange( + void onSessionChange( bundle as unknown as SessionBundle, DID, 'expired', @@ -382,7 +403,7 @@ describe('expiry rescue', () => { */ mockPersisted.latest = {accounts: [account], currentAccount: account} act(() => { - onSessionChange( + void onSessionChange( rescued as unknown as SessionBundle, DID, 'expired', @@ -416,7 +437,7 @@ describe('refresh persistence', () => { ) act(() => { - onSessionChange( + void onSessionChange( bundle as unknown as SessionBundle, DID, 'update', @@ -437,7 +458,7 @@ describe('refresh persistence', () => { ) act(() => { - onSessionChange( + void onSessionChange( bundle as unknown as SessionBundle, DID, 'update', @@ -447,10 +468,81 @@ describe('refresh persistence', () => { expect(currentAccount()?.pdsUrl).toBe(`${DIDDOC_PDS_HOST}/`) }) + + it('adopts the committed generation when a stale refresh succeeds', async () => { + const account = makeAccount() + const bundle = makeBundle(account) + const {onSessionChange, currentAccount} = await renderLoggedIn( + account, + bundle, + ) + const committed = makeAccount({ + accessJwt: 'access-jwt-2', + refreshJwt: 'refresh-jwt-2', + }) + mockPersisted.latest = mockPersisted.session = { + accounts: [committed], + currentAccount: committed, + credentialStates: { + [DID]: { + credentialVersion: 8, + refreshJti: 'refresh-jwt-2', + status: 'active', + }, + }, + } + + await act(async () => { + await onSessionChange( + bundle as unknown as SessionBundle, + DID, + 'update', + refreshedData('refresh-jwt-3'), + ) + }) + + expect(mockRebuilds.at(-1)?.account.refreshJwt).toBe('refresh-jwt-2') + expect(currentAccount()?.refreshJwt).toBe('refresh-jwt-2') + }) + + it('preserves a logout tombstone against a delayed successful refresh', async () => { + const account = makeAccount() + const bundle = makeBundle(account) + const {onSessionChange, currentAccount, hasSession} = await renderLoggedIn( + account, + bundle, + ) + const loggedOut = makeAccount({ + accessJwt: undefined, + refreshJwt: undefined, + }) + mockPersisted.latest = mockPersisted.session = { + accounts: [loggedOut], + currentAccount: undefined, + credentialStates: { + [DID]: {credentialVersion: 8, status: 'logged-out'}, + }, + } + + await act(async () => { + await onSessionChange( + bundle as unknown as SessionBundle, + DID, + 'update', + refreshedData('refresh-jwt-2'), + ) + }) + + expect(hasSession()).toBe(false) + expect(currentAccount()).toBeUndefined() + expect(mockPersisted.latest.accounts[0].refreshJwt).toBeUndefined() + }) }) /** Deliver a cross-tab `persisted` update to the provider's newest listener. */ function emitSynced(session: Schema['session']) { + mockPersisted.session = session + mockPersisted.latest = session mockPersistedListeners[mockPersistedListeners.length - 1](session) } diff --git a/src/state/session/__tests__/session-core-test.ts b/src/state/session/__tests__/session-core-test.ts index 3503acd08a..c9eb1f976a 100644 --- a/src/state/session/__tests__/session-core-test.ts +++ b/src/state/session/__tests__/session-core-test.ts @@ -97,10 +97,12 @@ import { disposeBundle, finishPreparation, makeSessionHooks, + type OnSessionChange, registerBundleKillSwitch, sessionAccountToSessionData, type SessionBundle, sessionDataToSessionAccount, + takeSessionChangeError, } from '../session-core' import { asFetch, @@ -367,7 +369,7 @@ describe('createSessionBundleFromStoredAccount', () => { it('builds three clients over one session', async () => { const result = createSessionBundleFromStoredAccount( makeAccount(), - jest.fn(), + jest.fn(), )! /* every client reads its identity straight through the shared session */ @@ -384,7 +386,7 @@ describe('createSessionBundleFromStoredAccount', () => { }) it('disposes a bundle rejected by the activation guard', async () => { - const onSessionChange = jest.fn() + const onSessionChange = jest.fn() let rejectedBundle: SessionBundle | undefined const result = createSessionBundleFromStoredAccount( makeAccount(), @@ -607,7 +609,7 @@ describe('session-hook payload threading (pre-commit ordering)', () => { describe('disposeBundle kill-switch', () => { it('the injected fetch throws after disposeBundle', () => { const hooks = makeSessionHooks({ - onSessionChange: jest.fn(), + onSessionChange: jest.fn(), getBundle: () => ({}) as SessionBundle, getDid: () => DID, }) @@ -850,7 +852,7 @@ describe('factory account snapshot after preparation', () => { await withFreshFactory(asFetch(fetchMock), async core => { const {account, bundle} = await core.createSessionBundleAndResume( makeAccount({accessJwt: 'valid-access-jwt'}), - jest.fn(), + jest.fn(), ) /* the moderation prep step ran the refresh */ expect(mockConfigureModerationForAccount).toHaveBeenCalledTimes(1) @@ -873,7 +875,7 @@ describe('factory account snapshot after preparation', () => { await withFreshFactory(asFetch(fetchMock), async core => { const {account} = await core.createSessionBundleAndResume( makeAccount({accessJwt: 'valid-access-jwt'}), - jest.fn(), + jest.fn(), ) expect(account.accessJwt).toBe('valid-access-jwt') expect(account.refreshJwt).toBe('refresh-jwt') @@ -921,7 +923,7 @@ describe('a session destroyed or rejected during preparation', () => { await expect( core.createSessionBundleAndResume( makeAccount({accessJwt: 'valid-access-jwt'}), - jest.fn(), + jest.fn(), ), ).rejects.toThrow('Session was revoked while it was being prepared') @@ -946,7 +948,7 @@ describe('a session destroyed or rejected during preparation', () => { await expect( core.createSessionBundleAndResume( makeAccount({accessJwt: 'valid-access-jwt'}), - jest.fn(), + jest.fn(), ), ).rejects.toThrow('prefetch blew up') @@ -959,7 +961,7 @@ describe('a session destroyed or rejected during preparation', () => { it('finishPreparation disposes and rethrows without running the snapshot', async () => { const hooks = makeSessionHooks({ - onSessionChange: jest.fn(), + onSessionChange: jest.fn(), getBundle: () => bundle, getDid: () => DID, }) @@ -1024,6 +1026,9 @@ describe('a throwing onSessionChange does not brick the session', () => { expect(mockLoggerError.mock.calls[0][1]).toEqual({ message: "session: onSessionChange threw for a 'update' event", }) + expect(takeSessionChangeError({bundle})).toEqual( + new Error('reducer side effect exploded'), + ) /* the session still committed the rotation, and can still refresh again */ expect(session.session.accessJwt).toBe('access-jwt-2') diff --git a/src/state/session/index.tsx b/src/state/session/index.tsx index e8e3613ddc..7d57f41d8d 100644 --- a/src/state/session/index.tsx +++ b/src/state/session/index.tsx @@ -13,6 +13,7 @@ import {type Client} from '@atproto/lex' import {type SessionData} from '@atproto/lex-password-session' import * as persisted from '#/state/persisted' +import {type Schema, type SessionCredentialMutation} from '#/state/persisted' import {useCloseAllActiveElements} from '#/state/util' import {useGlobalDialogsControlContext} from '#/components/dialogs/Context' import {AnalyticsContext, useAnalyticsBase, utils} from '#/analytics' @@ -29,9 +30,11 @@ import { createSessionBundleAndResume, createSessionBundleFromStoredAccount, disposeBundle, + type OnSessionChange, type PublicSessionBundle, type SessionBundle, sessionDataToSessionAccount, + takeSessionChangeError, } from './session-core' export {isSignupQueued} from './session-data' import { @@ -102,10 +105,15 @@ class SessionStore { } } - dispatch = (action: Action) => { + dispatch = ( + action: Action, + credentialMutations: SessionCredentialMutation[] = [], + ): Promise => { const nextState = reducer(this.state, action) this.state = nextState - // Persist synchronously without waiting for the React render cycle. + let persistence: Promise = + Promise.resolve(undefined) + // Schedule persistence immediately without waiting for the React render cycle. if (nextState.needsPersist) { nextState.needsPersist = false const persistedData = { @@ -118,9 +126,13 @@ class SessionStore { type: 'persisted:broadcast', data: redactPersistedSession(persistedData), }) - void persisted.write('session', persistedData) + persistence = persisted.updateSession({ + nextSession: persistedData, + credentialMutations, + }) } this.listeners.forEach(listener => listener()) + return persistence } } @@ -140,15 +152,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) { * insertion effect below, which commits well before any session hook can * fire: hooks are armed only after an asynchronous session factory resolves. */ - const onSessionChangeRef = useRef< - | (( - bundle: SessionBundle, - accountDid: string, - sessionEvent: AtpSessionEvent, - sessionData?: SessionData, - ) => void) - | null - >(null) + const onSessionChangeRef = useRef(null) const onSessionChange = useCallback( ( @@ -157,106 +161,169 @@ export function Provider({children}: React.PropsWithChildren<{}>) { sessionEvent: AtpSessionEvent, sessionData?: SessionData, ) => { - /* - * Only the live bundle may reset the expiry-rescue bookkeeping: a stale - * bundle's late update would otherwise clear the failed-generation set - * that bounds the rescue loop. (Its dispatch below is separately dropped - * by the reducer's identity guard.) - */ - if ( - sessionEvent === 'update' && - sessionData && - (store.getState().currentBundleState.bundle as unknown as - SessionBundle | PublicSessionBundle) === bundle - ) { - failedExpiryTokensRef.current.get(accountDid)?.clear() - } + const baseRefreshJwt = + sessionEvent === 'update' && !bundle.session.destroyed + ? bundle.session.session.refreshJwt + : sessionData?.refreshJwt - /* - * PasswordSession invokes its hooks before updating its live getter. Use - * the delivered payload so a refresh persists the newly rotated tokens. - * - * A refresh payload carries no didDoc unless the server sends one, so the - * stored account's `pdsUrl` is threaded in as the fallback. Without it the - * refresh would persist `pdsUrl: undefined` and the next cold start would - * route pre-refresh requests to the entryway instead of the PDS. - */ - const refreshedAccount = - sessionEvent === 'update' && sessionData - ? sessionDataToSessionAccount( - sessionData, - sessionData.service, - store.getState().accounts.find(a => a.did === accountDid)?.pdsUrl, - ) - : undefined - - /* - * A stale tab may expire a token after another tab has already rotated it. - * Prefer a newer persisted or reducer generation over logging every tab - * out. Failed generations are recorded and bounded to guarantee that a - * repeatedly expiring session eventually falls through to logout. - */ - if (sessionEvent === 'expired') { - const current = store.getState() - const currentBundle = current.currentBundleState.bundle as unknown as - SessionBundle | PublicSessionBundle - const dyingRefreshJwt = sessionData?.refreshJwt - // Stale bundle events are handled by the reducer's identity guard. - if ( - currentBundle === bundle && - current.currentBundleState.did === accountDid && - dyingRefreshJwt - ) { - let failedSet = failedExpiryTokensRef.current.get(accountDid) - if (!failedSet) { - failedSet = new Set() - failedExpiryTokensRef.current.set(accountDid, failedSet) + return persisted.runWithSessionCredentialLock({ + accountDids: [accountDid], + operation: async () => { + /* + * Only the live bundle may reset the expiry-rescue bookkeeping: a stale + * bundle's late update would otherwise clear the failed-generation set + * that bounds the rescue loop. (Its dispatch below is separately dropped + * by the reducer's identity guard.) + */ + if ( + sessionEvent === 'update' && + sessionData && + (store.getState().currentBundleState.bundle as unknown as + SessionBundle | PublicSessionBundle) === bundle + ) { + failedExpiryTokensRef.current.get(accountDid)?.clear() } - failedSet.add(dyingRefreshJwt) - const persistedCandidate = persisted - .readLatest('session') - .accounts.find(a => a.did === accountDid) - const reducerCandidate = current.accounts.find( - a => a.did === accountDid, - ) - const candidate = pickExpiryRescueCandidate({ - dyingRefreshJwt, - candidates: [persistedCandidate, reducerCandidate], - failedRefreshJwts: failedSet, - }) + /* + * PasswordSession invokes its hooks before updating its live getter. Use + * the delivered payload so a refresh persists the newly rotated tokens. + * + * A refresh payload carries no didDoc unless the server sends one, so the + * stored account's `pdsUrl` is threaded in as the fallback. Without it the + * refresh would persist `pdsUrl: undefined` and the next cold start would + * route pre-refresh requests to the entryway instead of the PDS. + */ + const refreshedAccount = + sessionEvent === 'update' && sessionData + ? sessionDataToSessionAccount( + sessionData, + sessionData.service, + store.getState().accounts.find(a => a.did === accountDid) + ?.pdsUrl, + ) + : undefined - if (candidate) { - const rebuilt = createSessionBundleFromStoredAccount( - candidate, - onSessionChangeRef.current!, - ) - if (rebuilt) { - store.dispatch({ - type: 'replaced-current-bundle', - newBundle: rebuilt.bundle, - newAccount: rebuilt.account, + /* + * A stale tab may expire a token after another tab has already rotated it. + * Prefer a newer persisted or reducer generation over logging every tab + * out. Failed generations are recorded and bounded to guarantee that a + * repeatedly expiring session eventually falls through to logout. + */ + if (sessionEvent === 'expired') { + const current = store.getState() + const currentBundle = current.currentBundleState + .bundle as unknown as SessionBundle | PublicSessionBundle + const dyingRefreshJwt = sessionData?.refreshJwt + // Stale bundle events are handled by the reducer's identity guard. + if ( + currentBundle === bundle && + current.currentBundleState.did === accountDid && + dyingRefreshJwt + ) { + let failedSet = failedExpiryTokensRef.current.get(accountDid) + if (!failedSet) { + failedSet = new Set() + failedExpiryTokensRef.current.set(accountDid, failedSet) + } + failedSet.add(dyingRefreshJwt) + + const persistedCandidate = persisted + .readLatest('session') + .accounts.find(a => a.did === accountDid) + const reducerCandidate = current.accounts.find( + a => a.did === accountDid, + ) + const candidate = pickExpiryRescueCandidate({ + dyingRefreshJwt, + candidates: [persistedCandidate, reducerCandidate], + failedRefreshJwts: failedSet, }) - return + + if (candidate) { + const rebuilt = createSessionBundleFromStoredAccount( + candidate, + onSessionChangeRef.current!, + ) + if (rebuilt) { + await store.dispatch({ + type: 'replaced-current-bundle', + newBundle: rebuilt.bundle, + newAccount: rebuilt.account, + }) + return + } + } } } - } - } - // Only the current bundle may report that its session was dropped. - if ( - sessionEvent === 'expired' && - store.getState().currentBundleState.bundle === bundle - ) { - emitSessionDropped() - } - // Bundle identity prevents stale sessions from changing the active account. - store.dispatch({ - type: 'received-session-event', - bundle, - refreshedAccount, - accountDid, - sessionEvent, + // Only the current bundle may report that its session was dropped. + if ( + sessionEvent === 'expired' && + store.getState().currentBundleState.bundle === bundle + ) { + emitSessionDropped() + } + + const credentialMutations: SessionCredentialMutation[] = [] + if (sessionEvent === 'update' && refreshedAccount) { + credentialMutations.push({ + type: 'refresh', + accountDid, + baseRefreshJwt, + resultRefreshJwt: refreshedAccount.refreshJwt, + }) + } else if (sessionEvent === 'expired') { + credentialMutations.push({ + type: 'expire', + accountDid, + baseRefreshJwt, + }) + } + + // Bundle identity prevents stale sessions from changing the active account. + const committedSession = await store.dispatch( + { + type: 'received-session-event', + bundle, + refreshedAccount, + accountDid, + sessionEvent, + }, + credentialMutations, + ) + + if (sessionEvent === 'update' && committedSession) { + const committedAccount = committedSession.accounts.find( + account => account.did === accountDid, + ) + if ( + committedAccount?.refreshJwt && + committedAccount.refreshJwt !== refreshedAccount?.refreshJwt && + store.getState().currentBundleState.bundle === bundle + ) { + const rebuilt = createSessionBundleFromStoredAccount( + committedAccount, + onSessionChangeRef.current!, + ) + if (rebuilt) { + await store.dispatch({ + type: 'replaced-current-bundle', + newBundle: rebuilt.bundle, + newAccount: rebuilt.account, + }) + } + } else if ( + !committedAccount?.refreshJwt && + store.getState().currentBundleState.bundle === bundle + ) { + await store.dispatch({ + type: 'synced-accounts', + syncedAccounts: committedSession.accounts, + syncedCurrentDid: committedSession.currentAccount?.did, + }) + } + } + }, }) }, [store], @@ -286,10 +353,23 @@ export function Provider({children}: React.PropsWithChildren<{}>) { disposeBundle(bundle) return } - store.dispatch({ - type: 'switched-to-account', - newBundle: bundle, - newAccount: account, + await persisted.runWithSessionCredentialLock({ + accountDids: [account.did], + operation: () => + store.dispatch( + { + type: 'switched-to-account', + newBundle: bundle, + newAccount: account, + }, + [ + { + type: 'login', + accountDid: account.did, + resultRefreshJwt: account.refreshJwt, + }, + ], + ), }) ax.metric('account:create:success', metrics, { session: utils.accountToSessionMetadata(account), @@ -317,10 +397,23 @@ export function Provider({children}: React.PropsWithChildren<{}>) { disposeBundle(bundle) return } - store.dispatch({ - type: 'switched-to-account', - newBundle: bundle, - newAccount: account, + await persisted.runWithSessionCredentialLock({ + accountDids: [account.did], + operation: () => + store.dispatch( + { + type: 'switched-to-account', + newBundle: bundle, + newAccount: account, + }, + [ + { + type: 'login', + accountDid: account.did, + resultRefreshJwt: account.refreshJwt, + }, + ], + ), }) ax.metric( 'account:loggedIn', @@ -343,9 +436,18 @@ export function Provider({children}: React.PropsWithChildren<{}>) { addSessionDebugLog({type: 'method:start', method: 'logout'}) cancelPendingTask() const prevState = store.getState() - store.dispatch({ - type: 'logged-out-current-account', - }) + const accountDid = prevState.currentBundleState.did + if (accountDid) { + void persisted + .runWithSessionCredentialLock({ + accountDids: [accountDid], + operation: () => + store.dispatch({type: 'logged-out-current-account', accountDid}, [ + {type: 'logout', accountDid}, + ]), + }) + .catch(() => {}) + } ax.metric( 'account:loggedOut', {logContext, scope: 'current'}, @@ -377,9 +479,27 @@ export function Provider({children}: React.PropsWithChildren<{}>) { addSessionDebugLog({type: 'method:start', method: 'logout'}) cancelPendingTask() const prevState = store.getState() - store.dispatch({ - type: 'logged-out-every-account', - }) + const accountDids = [ + ...new Set([ + ...prevState.accounts.map(account => account.did), + ...persisted + .readLatest('session') + .accounts.map(account => account.did), + ]), + ] + void persisted + .runWithSessionCredentialLock({ + accountDids, + operation: () => + store.dispatch( + {type: 'logged-out-every-account'}, + accountDids.map(accountDid => ({ + type: 'logout' as const, + accountDid, + })), + ), + }) + .catch(() => {}) ax.metric( 'account:loggedOut', {logContext, scope: 'every'}, @@ -410,8 +530,12 @@ export function Provider({children}: React.PropsWithChildren<{}>) { account: redactAccount(storedAccount), }) const signal = cancelPendingTask() + const latestStoredAccount = persisted + .readLatest('session') + .accounts.find(account => account.did === storedAccount.did) + if (!latestStoredAccount?.refreshJwt) return const {bundle, account} = await createSessionBundleAndResume( - storedAccount, + latestStoredAccount, onSessionChange, ) @@ -431,11 +555,48 @@ export function Provider({children}: React.PropsWithChildren<{}>) { disposeBundle(bundle) return } - store.dispatch({ - type: 'switched-to-account', - newBundle: bundle, - newAccount: account, + const committedSession = await persisted.runWithSessionCredentialLock({ + accountDids: [account.did], + operation: () => + store.dispatch( + { + type: 'switched-to-account', + newBundle: bundle, + newAccount: account, + }, + [ + { + type: 'refresh', + accountDid: account.did, + baseRefreshJwt: latestStoredAccount.refreshJwt, + resultRefreshJwt: account.refreshJwt, + }, + ], + ), }) + const committedAccount = committedSession?.accounts.find( + candidate => candidate.did === account.did, + ) + if (!committedAccount?.refreshJwt) { + await store.dispatch({ + type: 'synced-accounts', + syncedAccounts: committedSession?.accounts ?? [], + syncedCurrentDid: committedSession?.currentAccount?.did, + }) + return + } + if (committedAccount.refreshJwt !== account.refreshJwt) { + const rebuilt = createSessionBundleFromStoredAccount( + committedAccount, + onSessionChange, + ) + if (!rebuilt) return + await store.dispatch({ + type: 'replaced-current-bundle', + newBundle: rebuilt.bundle, + newAccount: rebuilt.account, + }) + } addSessionDebugLog({ type: 'method:end', method: 'resumeSession', @@ -464,18 +625,22 @@ export function Provider({children}: React.PropsWithChildren<{}>) { /* getSession targets the PDS; only the persisted account fields are patched. */ const data = await bundle.pdsClient.call(com.atproto.server.getSession, {}) if (signal.aborted) return - store.dispatch({ - type: 'partial-refresh-session', - /* - * Read the did off the response rather than the session: the bundle may - * have been disposed while the request was in flight, and the live - * getters throw in that state. - */ - accountDid: data.did, - patch: { - emailConfirmed: data.emailConfirmed, - emailAuthFactor: data.emailAuthFactor, - }, + await persisted.runWithSessionCredentialLock({ + accountDids: [data.did], + operation: () => + store.dispatch({ + type: 'partial-refresh-session', + /* + * Read the did off the response rather than the session: the bundle may + * have been disposed while the request was in flight, and the live + * getters throw in that state. + */ + accountDid: data.did, + patch: { + emailConfirmed: data.emailConfirmed, + emailAuthFactor: data.emailAuthFactor, + }, + }), }) }, [store, cancelPendingTask]) @@ -511,6 +676,12 @@ export function Provider({children}: React.PropsWithChildren<{}>) { if (!bundle.session) return undefined // logged out: nothing to refresh const before = bundle.session.session const after = await bundle.session.refresh() + const sessionChangeError = takeSessionChangeError({bundle}) + if (sessionChangeError) { + throw sessionChangeError instanceof Error + ? sessionChangeError + : new Error('Session persistence failed', {cause: sessionChangeError}) + } if (after === before) { throw new Error('Failed to refresh session') } @@ -541,10 +712,19 @@ export function Provider({children}: React.PropsWithChildren<{}>) { account: redactAccount(account), }) cancelPendingTask() - store.dispatch({ - type: 'removed-account', - accountDid: account.did, - }) + void persisted + .runWithSessionCredentialLock({ + accountDids: [account.did], + operation: () => + store.dispatch( + { + type: 'removed-account', + accountDid: account.did, + }, + [{type: 'remove', accountDid: account.did}], + ), + }) + .catch(() => {}) addSessionDebugLog({ type: 'method:end', method: 'removeAccount', @@ -561,7 +741,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) { type: 'persisted:receive', data: redactPersistedSession(synced), }) - store.dispatch({ + void store.dispatch({ type: 'synced-accounts', syncedAccounts: synced.accounts, syncedCurrentDid: synced.currentAccount?.did, @@ -636,7 +816,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) { return } const {bundle: newBundle, account: newAccount} = rebuilt - store.dispatch({ + void store.dispatch({ type: 'replaced-current-bundle', newBundle, newAccount, diff --git a/src/state/session/reducer.ts b/src/state/session/reducer.ts index 06ceb8ece6..85f2b4be92 100644 --- a/src/state/session/reducer.ts +++ b/src/state/session/reducer.ts @@ -46,6 +46,7 @@ export type Action = } | { type: 'logged-out-current-account' + accountDid?: string } | { type: 'logged-out-every-account' @@ -179,7 +180,7 @@ let reducer = (state: State, action: Action): State => { } case 'logged-out-current-account': { const {currentBundleState} = state - const accountDid = currentBundleState.did + const accountDid = action.accountDid ?? currentBundleState.did // side effect const account = state.accounts.find(a => a.did === accountDid) if (account && accountDid) { @@ -206,7 +207,10 @@ let reducer = (state: State, action: Action): State => { } : a, ), - currentBundleState: createPublicBundleState(), + currentBundleState: + currentBundleState.did === accountDid + ? createPublicBundleState() + : currentBundleState, needsPersist: true, } } diff --git a/src/state/session/session-core.ts b/src/state/session/session-core.ts index dbed9bfa2f..d5c5ee1024 100644 --- a/src/state/session/session-core.ts +++ b/src/state/session/session-core.ts @@ -68,6 +68,7 @@ export type SessionBundle = { * state private and tied to bundle identity. */ const bundleKillSwitches = new WeakMap void>() +const bundleSessionChangeErrors = new WeakMap() /** * Register the lifecycle closure used by {@link disposeBundle}. @@ -83,6 +84,17 @@ export function registerBundleKillSwitch( bundleKillSwitches.set(bundle, kill) } +/** Return and clear the latest persistence/session-change hook failure. */ +export function takeSessionChangeError({ + bundle, +}: { + bundle: SessionBundle +}): unknown { + const error = bundleSessionChangeErrors.get(bundle) + bundleSessionChangeErrors.delete(bundle) + return error +} + /** * Build the three clients over a session. * @@ -126,7 +138,7 @@ export type OnSessionChange = ( did: string, event: AtpSessionEvent, sessionData?: SessionData, -) => void +) => void | Promise /** * Hooks stay inert during initial session preparation. `kill()` disarms them @@ -146,7 +158,10 @@ export function makeSessionHooks({ }) { let armed = false let killed = false - const dispatch = (event: AtpSessionEvent, sessionData?: SessionData) => { + const dispatch = async ( + event: AtpSessionEvent, + sessionData?: SessionData, + ) => { if (!armed) { return } @@ -159,12 +174,15 @@ export function makeSessionHooks({ * effects and event emitters, so treat it as capable of throwing. */ try { + const bundle = getBundle() + bundleSessionChangeErrors.delete(bundle) const did = getDid() - onSessionChange(getBundle(), did, event, sessionData) + await onSessionChange(bundle, did, event, sessionData) if (event !== 'update') { addSessionErrorLog(did, event) } } catch (e) { + bundleSessionChangeErrors.set(getBundle(), e) logger.error(e instanceof Error ? e : String(e), { message: `session: onSessionChange threw for a '${event}' event`, }) @@ -178,13 +196,13 @@ export function makeSessionHooks({ return networkAwareFetch(input, init) }, onUpdated(data) { - dispatch('update', data) + return dispatch('update', data) }, onDeleted(data) { - dispatch('expired', data) + return dispatch('expired', data) }, onUpdateFailure() { - dispatch('network-error') + return dispatch('network-error') }, } return Object.assign(hooks, {