From 7d6518cd49ba24da74c7f77e3d9b4559a3054786 Mon Sep 17 00:00:00 2001 From: hailey Date: Thu, 27 Aug 2026 23:39:03 +0000 Subject: [PATCH] Fix stale cross-tab session writes --- src/state/persisted/index.ts | 12 +- src/state/persisted/index.web.ts | 15 +- src/state/persisted/types.ts | 11 +- src/state/session/__tests__/session-test.ts | 233 +++++++++++++++++++- src/state/session/index.tsx | 39 +++- src/state/session/reducer.ts | 165 ++++++++++++++ 6 files changed, 460 insertions(+), 15 deletions(-) diff --git a/src/state/persisted/index.ts b/src/state/persisted/index.ts index 0659e7817b..f3516c23a8 100644 --- a/src/state/persisted/index.ts +++ b/src/state/persisted/index.ts @@ -8,7 +8,7 @@ import { tryStringify, } from '#/state/persisted/schema' import {device} from '#/storage' -import {type PersistedApi} from './types' +import {type PersistedApi, type PersistedWriteValue} from './types' import {normalizeData} from './util' export type {PersistedAccount, Schema} from '#/state/persisted/schema' @@ -45,16 +45,22 @@ readLatest satisfies PersistedApi['readLatest'] export async function write( key: K, - value: Schema[K], + value: PersistedWriteValue, ): Promise { _state = normalizeData({ ..._state, - [key]: value, + [key]: resolveWriteValue(value, _state[key]), }) await writeToStorage(_state) } write satisfies PersistedApi['write'] +function resolveWriteValue(value: PersistedWriteValue, latest: T): T { + return typeof value === 'function' + ? (value as (current: T) => T)(latest) + : value +} + export function onUpdate( _key: K, _cb: (v: Schema[K]) => void, diff --git a/src/state/persisted/index.web.ts b/src/state/persisted/index.web.ts index a4c5d1b247..29d8729c5a 100644 --- a/src/state/persisted/index.web.ts +++ b/src/state/persisted/index.web.ts @@ -8,7 +8,7 @@ import { tryParse, tryStringify, } from '#/state/persisted/schema' -import {type PersistedApi} from './types' +import {type PersistedApi, type PersistedWriteValue} from './types' import {normalizeData} from './util' export type {PersistedAccount, Schema} from '#/state/persisted/schema' @@ -65,7 +65,7 @@ readLatest satisfies PersistedApi['readLatest'] // eslint-disable-next-line @typescript-eslint/require-await export async function write( key: K, - value: Schema[K], + value: PersistedWriteValue, ): Promise { const next = readFromStorage() if (next) { @@ -75,8 +75,9 @@ export async function write( // 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. } + const nextValue = resolveWriteValue(value, _state[key]) try { - if (JSON.stringify({v: _state[key]}) === JSON.stringify({v: value})) { + if (JSON.stringify({v: _state[key]}) === JSON.stringify({v: nextValue})) { // Fast path for updates that are guaranteed to be noops. // This is good mostly because it avoids useless broadcasts to other tabs. return @@ -86,7 +87,7 @@ export async function write( } _state = normalizeData({ ..._state, - [key]: value, + [key]: nextValue, }) writeToStorage(_state) broadcast.postMessage({event: {type: UPDATE_EVENT, key}}) @@ -94,6 +95,12 @@ export async function write( } write satisfies PersistedApi['write'] +function resolveWriteValue(value: PersistedWriteValue, latest: T): T { + return typeof value === 'function' + ? (value as (current: T) => T)(latest) + : value +} + export function onUpdate( key: K, cb: (v: Schema[K]) => void, diff --git a/src/state/persisted/types.ts b/src/state/persisted/types.ts index 4b5a29d938..4ce0555169 100644 --- a/src/state/persisted/types.ts +++ b/src/state/persisted/types.ts @@ -1,5 +1,11 @@ import {type Schema} from './schema' +/** + * A value to persist, or a function that derives one from the freshest + * persisted value while the write operation is in progress. + */ +export type PersistedWriteValue = T | ((latest: T) => T) + export type PersistedApi = { init(): Promise get(key: K): Schema[K] @@ -12,7 +18,10 @@ export type PersistedApi = { * it is identical to {@link get} (single-instance, no other writer). */ readLatest(key: K): Schema[K] - write(key: K, value: Schema[K]): Promise + write( + key: K, + value: PersistedWriteValue, + ): Promise onUpdate( key: K, cb: (v: Schema[K]) => void, diff --git a/src/state/session/__tests__/session-test.ts b/src/state/session/__tests__/session-test.ts index 578b4d26e7..5b47271e8b 100644 --- a/src/state/session/__tests__/session-test.ts +++ b/src/state/session/__tests__/session-test.ts @@ -1,7 +1,13 @@ import {type SessionData} from '@atproto/lex-password-session' import {describe, expect, it, jest} from '@jest/globals' -import {type Action, getInitialState, reducer, type State} from '../reducer' +import { + type Action, + getInitialState, + rebasePersistedSession, + reducer, + type State, +} from '../reducer' import {sessionDataToSessionAccount} from '../session-core' import {type SessionAccount} from '../types' @@ -1777,8 +1783,233 @@ describe('session', () => { expect(state.currentBundleState.did).toBe('alice-did') expect(state.needsPersist).toBe(true) }) + + describe('rebasePersistedSession', () => { + it('keeps a newer token generation when a stale metadata update is persisted', () => { + const stale = makeAccount('https://alice.com', { + active: true, + did: 'alice-did', + handle: 'alice.test', + accessJwt: 'alice-access-jwt-1', + refreshJwt: 'alice-refresh-jwt-1', + emailConfirmed: false, + emailAuthFactor: false, + }) + const fresh = makeAccount('https://alice.com', { + active: true, + did: 'alice-did', + handle: 'alice.test', + accessJwt: 'alice-access-jwt-2', + refreshJwt: 'alice-refresh-jwt-2', + emailConfirmed: false, + emailAuthFactor: false, + }) + const bob = makeAccount('https://bob.com', { + active: true, + did: 'bob-did', + handle: 'bob.test', + accessJwt: 'bob-access-jwt-1', + refreshJwt: 'bob-refresh-jwt-1', + }) + const desired = snapshot([{...stale, emailConfirmed: true}], 'alice-did') + const latest = snapshot([fresh, bob], 'alice-did') + + const rebased = rebasePersistedSession(latest, desired, { + type: 'partial-refresh-session', + accountDid: 'alice-did', + patch: {emailConfirmed: true, emailAuthFactor: false}, + }) + + expect(rebased.accounts).toEqual([{...fresh, emailConfirmed: true}, bob]) + expect(rebased.currentAccount?.did).toBe('alice-did') + }) + + it('accepts a fresh token pair delivered by the active session', () => { + const stale = makeAccount('https://alice.com', { + active: true, + did: 'alice-did', + handle: 'alice.test', + accessJwt: 'alice-access-jwt-1', + refreshJwt: 'alice-refresh-jwt-1', + }) + const fresh = makeAccount('https://alice.com', { + active: true, + did: 'alice-did', + handle: 'alice.test', + accessJwt: 'alice-access-jwt-2', + refreshJwt: 'alice-refresh-jwt-2', + }) + + const rebased = rebasePersistedSession( + snapshot([stale], 'alice-did'), + snapshot([fresh], 'alice-did'), + { + type: 'received-session-event', + bundle: makeBundle('https://alice.com'), + accountDid: 'alice-did', + refreshedAccount: fresh, + sessionEvent: 'update', + }, + ) + + expect(rebased.accounts).toEqual([fresh]) + }) + + it('preserves a concurrent newer generation instead of clearing it on expiry', () => { + const dying = makeAccount('https://alice.com', { + active: true, + did: 'alice-did', + handle: 'alice.test', + accessJwt: 'alice-access-jwt-1', + refreshJwt: 'alice-refresh-jwt-1', + }) + const fresh = makeAccount('https://alice.com', { + active: true, + did: 'alice-did', + handle: 'alice.test', + accessJwt: 'alice-access-jwt-2', + refreshJwt: 'alice-refresh-jwt-2', + }) + const expired = {...dying, accessJwt: undefined, refreshJwt: undefined} + + const rebased = rebasePersistedSession( + snapshot([fresh], 'alice-did'), + snapshot([expired]), + { + type: 'received-session-event', + bundle: makeBundle('https://alice.com'), + accountDid: 'alice-did', + refreshedAccount: undefined, + sessionEvent: 'expired', + expiredRefreshJwt: dying.refreshJwt, + }, + ) + + expect(rebased.accounts).toEqual([fresh]) + expect(rebased.currentAccount?.did).toBe('alice-did') + }) + + it('keeps explicit logout authoritative over a newer stored token', () => { + const fresh = makeAccount('https://alice.com', { + active: true, + did: 'alice-did', + handle: 'alice.test', + accessJwt: 'alice-access-jwt-2', + refreshJwt: 'alice-refresh-jwt-2', + }) + const loggedOut = {...fresh, accessJwt: undefined, refreshJwt: undefined} + + const rebased = rebasePersistedSession( + snapshot([fresh], 'alice-did'), + snapshot([loggedOut]), + {type: 'logged-out-current-account'}, + ) + + expect(rebased.accounts).toEqual([loggedOut]) + expect(rebased.currentAccount).toBeUndefined() + }) + + it('clears only the explicitly logged-out account', () => { + const alice = makeAccount('https://alice.com', { + active: true, + did: 'alice-did', + handle: 'alice.test', + accessJwt: 'alice-access-jwt-2', + refreshJwt: 'alice-refresh-jwt-2', + }) + const bobStale = makeAccount('https://bob.com', { + active: true, + did: 'bob-did', + handle: 'bob.test', + accessJwt: 'bob-access-jwt-1', + refreshJwt: 'bob-refresh-jwt-1', + }) + const bobFresh = { + ...bobStale, + accessJwt: 'bob-access-jwt-2', + refreshJwt: 'bob-refresh-jwt-2', + } + const loggedOutAlice = { + ...alice, + accessJwt: undefined, + refreshJwt: undefined, + } + + const rebased = rebasePersistedSession( + snapshot([alice, bobFresh], 'alice-did'), + snapshot([loggedOutAlice, bobStale]), + {type: 'logged-out-current-account', accountDid: 'alice-did'}, + ) + + expect(rebased.accounts).toEqual([loggedOutAlice, bobFresh]) + }) + + it('does not restore an account another tab removed', () => { + const stale = makeAccount('https://alice.com', { + active: true, + did: 'alice-did', + handle: 'alice.test', + accessJwt: 'alice-access-jwt-1', + refreshJwt: 'alice-refresh-jwt-1', + }) + + const rebased = rebasePersistedSession( + snapshot([]), + snapshot([{...stale, emailConfirmed: true}], 'alice-did'), + { + type: 'partial-refresh-session', + accountDid: 'alice-did', + patch: {emailConfirmed: true, emailAuthFactor: false}, + }, + ) + + expect(rebased.accounts).toEqual([]) + expect(rebased.currentAccount).toBeUndefined() + }) + + it('clears credentials for accounts added before an all-account logout', () => { + const alice = makeAccount('https://alice.com', { + active: true, + did: 'alice-did', + handle: 'alice.test', + accessJwt: 'alice-access-jwt-2', + refreshJwt: 'alice-refresh-jwt-2', + }) + const bob = makeAccount('https://bob.com', { + active: true, + did: 'bob-did', + handle: 'bob.test', + accessJwt: 'bob-access-jwt-2', + refreshJwt: 'bob-refresh-jwt-2', + }) + const loggedOutAlice = { + ...alice, + accessJwt: undefined, + refreshJwt: undefined, + } + + const rebased = rebasePersistedSession( + snapshot([alice, bob], 'alice-did'), + snapshot([loggedOutAlice]), + {type: 'logged-out-every-account'}, + ) + + expect(rebased.accounts).toEqual([ + loggedOutAlice, + {...bob, accessJwt: undefined, refreshJwt: undefined}, + ]) + expect(rebased.currentAccount).toBeUndefined() + }) + }) }) +function snapshot(accounts: SessionAccount[], currentDid?: string) { + return { + accounts, + currentAccount: accounts.find(account => account.did === currentDid), + } +} + function run(initialState: State, actions: Action[]): State { let state = initialState for (let action of actions) { diff --git a/src/state/session/index.tsx b/src/state/session/index.tsx index e8e3613ddc..69664de1d5 100644 --- a/src/state/session/index.tsx +++ b/src/state/session/index.tsx @@ -22,7 +22,13 @@ import {emitSessionDropped} from '../events' import {getPublicAppviewClient} from './clients' import {createSessionBundleAndCreateAccount} from './create-account' import {pickExpiryRescueCandidate} from './expiry-rescue' -import {type Action, getInitialState, reducer, type State} from './reducer' +import { + type Action, + getInitialState, + rebasePersistedSession, + reducer, + type State, +} from './reducer' import { type AtpSessionEvent, createSessionBundleAndLogin, @@ -108,17 +114,29 @@ class SessionStore { // Persist synchronously without waiting for the React render cycle. if (nextState.needsPersist) { nextState.needsPersist = false - const persistedData = { + const desiredPersistedData = { accounts: nextState.accounts, currentAccount: nextState.accounts.find( a => a.did === nextState.currentBundleState.did, ), } - addSessionDebugLog({ - type: 'persisted:broadcast', - data: redactPersistedSession(persistedData), + /* + * A background tab can have missed a token-rotation broadcast. Derive the + * outgoing value inside `write`, after its final web storage re-read, so an + * unrelated metadata update cannot restore an older token generation. + */ + void persisted.write('session', latest => { + const persistedData = rebasePersistedSession( + latest, + desiredPersistedData, + action, + ) + addSessionDebugLog({ + type: 'persisted:broadcast', + data: redactPersistedSession(persistedData), + }) + return persistedData }) - void persisted.write('session', persistedData) } this.listeners.forEach(listener => listener()) } @@ -257,6 +275,8 @@ export function Provider({children}: React.PropsWithChildren<{}>) { refreshedAccount, accountDid, sessionEvent, + expiredRefreshJwt: + sessionEvent === 'expired' ? sessionData?.refreshJwt : undefined, }) }, [store], @@ -290,6 +310,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) { type: 'switched-to-account', newBundle: bundle, newAccount: account, + tokenUpdate: 'login', }) ax.metric('account:create:success', metrics, { session: utils.accountToSessionMetadata(account), @@ -321,6 +342,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) { type: 'switched-to-account', newBundle: bundle, newAccount: account, + tokenUpdate: 'login', }) ax.metric( 'account:loggedIn', @@ -345,6 +367,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) { const prevState = store.getState() store.dispatch({ type: 'logged-out-current-account', + accountDid: prevState.currentBundleState.did, }) ax.metric( 'account:loggedOut', @@ -435,6 +458,10 @@ export function Provider({children}: React.PropsWithChildren<{}>) { type: 'switched-to-account', newBundle: bundle, newAccount: account, + tokenUpdate: + account.refreshJwt !== storedAccount.refreshJwt + ? 'refresh' + : undefined, }) addSessionDebugLog({ type: 'method:end', diff --git a/src/state/session/reducer.ts b/src/state/session/reducer.ts index 06ceb8ece6..32fc7bf834 100644 --- a/src/state/session/reducer.ts +++ b/src/state/session/reducer.ts @@ -1,5 +1,6 @@ import {unregisterPushToken} from '#/lib/notifications/notifications' import {logger} from '#/lib/notifications/util' +import {type Schema} from '#/state/persisted/schema' import {wrapSessionReducerForLogging} from './logging' import {createPublicSessionBundle} from './session-core' import {type AtpSessionEvent, type SessionAccount} from './types' @@ -28,11 +29,15 @@ export type Action = accountDid: string refreshedAccount: SessionAccount | undefined sessionEvent: AtpSessionEvent + /** The token that failed, used to preserve a newer concurrent generation. */ + expiredRefreshJwt?: string } | { type: 'switched-to-account' newBundle: OpaqueSessionBundle newAccount: SessionAccount + /** Whether this action obtained a fresh credential generation. */ + tokenUpdate?: 'login' | 'refresh' } | { // Replace an immutable session from synced or rescued tokens without rebroadcasting. @@ -46,6 +51,8 @@ export type Action = } | { type: 'logged-out-current-account' + /** Omitted only by legacy reducer tests; production always supplies it. */ + accountDid?: string } | { type: 'logged-out-every-account' @@ -61,6 +68,164 @@ export type Action = patch: Pick } +/** + * Rebase a pending session write over the newest persisted snapshot. + * + * Browser tabs hold independent in-memory session state. A tab which missed a + * token-rotation broadcast may still need to persist unrelated metadata. The + * persisted layer can protect other root keys, but it cannot tell whether the + * complete `session` value it receives contains an older credential generation. + * Keep the newest stored credential pair unless this action is known to have + * obtained a fresh pair or intentionally cleared it. + */ +export function rebasePersistedSession( + latest: Schema['session'], + desired: Schema['session'], + action: Action, +): Schema['session'] { + const latestByDid = new Map( + latest.accounts.map(account => [account.did, account]), + ) + const desiredDids = new Set(desired.accounts.map(account => account.did)) + const removedDid = + action.type === 'removed-account' ? action.accountDid : undefined + const preservedNewerTokenDids = new Set() + + const accounts = desired.accounts.flatMap(desiredAccount => { + const latestAccount = latestByDid.get(desiredAccount.did) + if (!latestAccount) { + /* Only an explicit fresh login may introduce an account absent from storage. */ + return shouldAddMissingAccount(action, desiredAccount.did) + ? [desiredAccount] + : [] + } + + const tokenMode = getTokenWriteMode( + action, + desiredAccount.did, + latestAccount, + ) + if (tokenMode === 'keep-latest') { + if (latestAccount.refreshJwt !== desiredAccount.refreshJwt) { + preservedNewerTokenDids.add(desiredAccount.did) + } + return [ + { + ...desiredAccount, + accessJwt: latestAccount.accessJwt, + refreshJwt: latestAccount.refreshJwt, + }, + ] + } + return [desiredAccount] + }) + + /* A stale tab must not discard an account another tab added meanwhile. */ + for (const latestAccount of latest.accounts) { + if ( + !desiredDids.has(latestAccount.did) && + latestAccount.did !== removedDid + ) { + accounts.push( + action.type === 'logged-out-every-account' + ? {...latestAccount, accessJwt: undefined, refreshJwt: undefined} + : latestAccount, + ) + } + } + + const desiredCurrentDid = desired.currentAccount?.did + const latestCurrentDid = latest.currentAccount?.did + const currentDid = selectCurrentDid({ + action, + desiredCurrentDid, + latestCurrentDid, + preservedNewerTokenDids, + }) + + return { + accounts, + currentAccount: accounts.find(account => account.did === currentDid), + } +} + +function shouldAddMissingAccount(action: Action, did: string): boolean { + return ( + action.type === 'switched-to-account' && + action.newAccount.did === did && + action.tokenUpdate === 'login' + ) +} + +function getTokenWriteMode( + action: Action, + did: string, + latestAccount: SessionAccount, +): 'replace' | 'clear' | 'keep-latest' { + switch (action.type) { + case 'logged-out-current-account': + return action.accountDid === undefined || action.accountDid === did + ? 'clear' + : 'keep-latest' + case 'logged-out-every-account': + return 'clear' + case 'received-session-event': + if (action.sessionEvent === 'update' && action.refreshedAccount) { + /* A remote logout wins over an in-flight local refresh. */ + return latestAccount.refreshJwt ? 'replace' : 'keep-latest' + } + if (action.sessionEvent === 'expired') { + return latestAccount.refreshJwt && + latestAccount.refreshJwt !== action.expiredRefreshJwt + ? 'keep-latest' + : 'clear' + } + return 'keep-latest' + case 'switched-to-account': + if (action.newAccount.did !== did) return 'keep-latest' + if (action.tokenUpdate === 'login') return 'replace' + if (action.tokenUpdate === 'refresh' && latestAccount.refreshJwt) { + return 'replace' + } + return 'keep-latest' + default: + return 'keep-latest' + } +} + +function selectCurrentDid({ + action, + desiredCurrentDid, + latestCurrentDid, + preservedNewerTokenDids, +}: { + action: Action + desiredCurrentDid: string | undefined + latestCurrentDid: string | undefined + preservedNewerTokenDids: Set +}): string | undefined { + switch (action.type) { + case 'switched-to-account': + case 'removed-account': + case 'logged-out-current-account': + case 'logged-out-every-account': + return desiredCurrentDid + case 'received-session-event': + if ( + action.sessionEvent === 'expired' && + action.accountDid !== undefined && + preservedNewerTokenDids.has(action.accountDid) + ) { + return latestCurrentDid + } + return action.sessionEvent === 'expired' + ? desiredCurrentDid + : (latestCurrentDid ?? desiredCurrentDid) + default: + return latestCurrentDid ?? desiredCurrentDid + } +} + function createPublicBundleState(): BundleState { return { bundle: createPublicSessionBundle(),