diff --git a/invokeai/frontend/web/src/features/auth/store/authSlice.ts b/invokeai/frontend/web/src/features/auth/store/authSlice.ts index 353b7ab9730..f69f289f923 100644 --- a/invokeai/frontend/web/src/features/auth/store/authSlice.ts +++ b/invokeai/frontend/web/src/features/auth/store/authSlice.ts @@ -22,7 +22,8 @@ const zAuthState = z.object({ type User = z.infer; type AuthState = z.infer; -const getTokenUserId = (token: string | null): string | null => { +/** The token's claims, or null when it is absent or not a readable JWT. */ +const decodeTokenPayload = (token: string | null): Record | null => { if (!token) { return null; } @@ -32,13 +33,51 @@ const getTokenUserId = (token: string | null): string | null => { return null; } const normalizedPayload = encodedPayload.replace(/-/g, '+').replace(/_/g, '/'); - const payload = JSON.parse(atob(normalizedPayload.padEnd(Math.ceil(normalizedPayload.length / 4) * 4, '='))); - return typeof payload.user_id === 'string' ? payload.user_id : null; + return JSON.parse(atob(normalizedPayload.padEnd(Math.ceil(normalizedPayload.length / 4) * 4, '='))); } catch { return null; } }; +const getTokenUserId = (token: string | null): string | null => { + const payload = decodeTokenPayload(token); + return typeof payload?.user_id === 'string' ? payload.user_id : null; +}; + +/** + * The session a token belongs to, as opposed to the bytes it happens to be made of. + * + * The sliding-window middleware mints a replacement token on every mutating request, so a live + * session's token changes bytes on a fixed cadence (throttled to once a minute by + * `acceptRefreshedToken`) while the login behind it never changes. Anything that must be rebuilt + * when the *session* changes — the socket, most notably — keys on this rather than on the token, + * so a routine refresh does not tear down work that belongs to the same user. + * + * The revocation epoch is part of the identity, not incidental to it. A password change bumps + * `token_epoch` on the user record, and the server force-disconnects every socket that + * authenticated under the superseded epoch (`sockets.py`, `_handle_user_access_changed`). A + * server-initiated disconnect is terminal for socket.io — the client sets `skipReconnect` and + * never retries (`socket.io-client`, `Socket.ondisconnect` -> `Manager._close`) — so the only way + * back is to build a new socket, and the replacement token the server hands out carries the new + * epoch. Keying on `user_id` alone would leave that socket dead until a full page reload, with + * `$isConnected` stuck false and Invoke disabled with it. + * + * A token carrying no user id falls back to its own bytes: with no identity to compare, byte + * equality is the only safe answer, and consumers keep their pre-existing behaviour. + */ +export const getTokenSessionKey = (token: string | null): string | null => { + if (!token) { + return null; + } + const payload = decodeTokenPayload(token); + if (typeof payload?.user_id !== 'string') { + return token; + } + // Absent on tokens minted before the claim existed; the server reads those as epoch 0 too. + const epoch = typeof payload.token_epoch === 'number' ? payload.token_epoch : 0; + return `${payload.user_id}:${epoch}`; +}; + export const tokensBelongToSameUser = (first: string | null, second: string | null): boolean => { const firstUserId = getTokenUserId(first); return firstUserId !== null && firstUserId === getTokenUserId(second); @@ -164,5 +203,6 @@ export const authSliceConfig: SliceConfig = { export const selectIsAuthenticated = (state: { auth: AuthState }) => state.auth.isAuthenticated; export const selectCurrentUser = (state: { auth: AuthState }) => state.auth.user; export const selectAuthToken = (state: { auth: AuthState }) => state.auth.token; +export const selectAuthSessionKey = (state: { auth: AuthState }) => getTokenSessionKey(state.auth.token); export const selectIsAuthLoading = (state: { auth: AuthState }) => state.auth.isLoading; export const selectSessionExpired = (state: { auth: AuthState }) => state.auth.sessionExpired; diff --git a/invokeai/frontend/web/src/features/auth/store/refreshedTokenConsumers.test.ts b/invokeai/frontend/web/src/features/auth/store/refreshedTokenConsumers.test.ts index 39974d9a4d4..daa53dae211 100644 --- a/invokeai/frontend/web/src/features/auth/store/refreshedTokenConsumers.test.ts +++ b/invokeai/frontend/web/src/features/auth/store/refreshedTokenConsumers.test.ts @@ -4,7 +4,14 @@ import { fileURLToPath } from 'node:url'; import { describe, expect, it } from 'vitest'; -import { authSliceConfig, currentUserUpdated, externalTokenAdopted, setCredentials, tokenRefreshed } from './authSlice'; +import { + authSliceConfig, + currentUserUpdated, + externalTokenAdopted, + getTokenSessionKey, + setCredentials, + tokenRefreshed, +} from './authSlice'; const user = { user_id: 'user', @@ -14,17 +21,37 @@ const user = { is_active: true, }; -const tokenFor = (userId: string) => - `header.${Buffer.from(JSON.stringify({ user_id: userId })).toString('base64url')}.signature`; +const tokenFor = (userId: string, nonce = 0, epoch = 0) => + `header.${Buffer.from(JSON.stringify({ user_id: userId, nonce, token_epoch: epoch })).toString('base64url')}.signature`; describe('refreshed token consumers', () => { - it('updates the Redux token used by socket reconnects', () => { + it('updates the Redux token the next request will carry', () => { let state = authSliceConfig.slice.reducer(undefined, setCredentials({ token: 'old', user })); state = authSliceConfig.slice.reducer(state, tokenRefreshed('new')); expect(state.token).toBe('new'); }); + it('keeps one session key across a refresh, so the socket is not rebuilt for it', () => { + // The middleware mints new bytes for the same login on every mutating request. Consumers that + // key work to a session — `useSocketIO` — must not see that as a new session. + expect(getTokenSessionKey(tokenFor('user', 1))).toBe(getTokenSessionKey(tokenFor('user', 2))); + expect(getTokenSessionKey(tokenFor('user', 1))).not.toBe(getTokenSessionKey(tokenFor('other', 1))); + expect(getTokenSessionKey(null)).toBeNull(); + }); + + it('changes the session key when a revoked epoch supersedes the token', () => { + // A password change bumps the epoch and the server drops every socket authenticated under the + // old one; socket.io does not retry that, so the replacement token must count as a new session. + expect(getTokenSessionKey(tokenFor('user', 1, 0))).not.toBe(getTokenSessionKey(tokenFor('user', 2, 1))); + }); + + it('falls back to the token itself when it carries no user id', () => { + // Nothing to compare identities with, so byte equality is the only safe answer. + expect(getTokenSessionKey('opaque-token')).toBe('opaque-token'); + expect(getTokenSessionKey('')).toBeNull(); + }); + it('clears the previous user when adopting a token from another tab', () => { let state = authSliceConfig.slice.reducer(undefined, setCredentials({ token: tokenFor(user.user_id), user })); state = authSliceConfig.slice.reducer(state, externalTokenAdopted(tokenFor('other-user'))); diff --git a/invokeai/frontend/web/src/services/events/useSocketIO.test.tsx b/invokeai/frontend/web/src/services/events/useSocketIO.test.tsx new file mode 100644 index 00000000000..8aeeb3b72c2 --- /dev/null +++ b/invokeai/frontend/web/src/services/events/useSocketIO.test.tsx @@ -0,0 +1,317 @@ +// @vitest-environment happy-dom +/** + * Mounted-DOM coverage for the socket's lifecycle against the auth session. + * + * The rule this pins down is which auth changes may replace the live socket. A sliding-window + * token refresh must not: it lands about once a minute during any activity, and rebuilding the + * socket for it disposed the event listeners and glitched the preview mid-generation. A change of + * user must, or the socket keeps the previous account's rooms and private events. + */ +import { Buffer } from 'node:buffer'; + +import { createStore } from 'app/store/store'; +import { + currentUserUpdated, + externalTokenAdopted, + logout, + setCredentials, + tokenRefreshed, +} from 'features/auth/store/authSlice'; +import { act } from 'react'; +import type { Root } from 'react-dom/client'; +import { createRoot } from 'react-dom/client'; +import { Provider } from 'react-redux'; +import type { ManagerOptions, SocketOptions } from 'socket.io-client'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import { useSocketIO } from './useSocketIO'; + +type SocketIOOptions = Partial; + +declare global { + var IS_REACT_ACT_ENVIRONMENT: boolean; +} +globalThis.IS_REACT_ACT_ENVIRONMENT = true; + +const createFakeSocket = () => { + const handlers = new Map void>>(); + return { + connect: vi.fn(), + disconnect: vi.fn(), + on: vi.fn((event: string, handler: (...args: unknown[]) => void) => { + handlers.set(event, (handlers.get(event) ?? new Set()).add(handler)); + }), + off: vi.fn((event: string, handler: (...args: unknown[]) => void) => { + handlers.get(event)?.delete(handler); + }), + emit: vi.fn(), + /** Drive the socket's own listeners, as socket.io does when the connection ends. */ + fire: (event: string, ...args: unknown[]) => { + handlers.get(event)?.forEach((handler) => handler(...args)); + }, + }; +}; + +const sockets: ReturnType[] = []; +const io = vi.fn((_url: string, _options: SocketIOOptions) => { + const socket = createFakeSocket(); + sockets.push(socket); + return socket; +}); + +vi.mock('socket.io-client', () => ({ io: (url: string, options: SocketIOOptions) => io(url, options) })); +vi.mock('services/events/setEventListeners', () => ({ setEventListeners: () => () => {} })); + +const user = { + user_id: 'user-1', + email: 'user@example.com', + display_name: null, + is_admin: false, + is_active: true, +}; + +/** + * A token for `userId`. `nonce` varies the bytes without varying the identity, as a sliding-window + * refresh does; `epoch` is the revocation epoch the server bumps when a password change kills the + * account's earlier sessions. + */ +const tokenFor = (userId: string, nonce = 0, epoch = 0) => + `header.${Buffer.from(JSON.stringify({ user_id: userId, nonce, token_epoch: epoch })).toString('base64url')}.signature`; + +const Probe = () => { + useSocketIO(); + return null; +}; + +/** The `auth` option as socket.io will invoke it: a callback, resolved at each connection attempt. */ +const resolveAuth = (options: SocketIOOptions | undefined): object | undefined => { + const auth = options?.auth; + if (typeof auth !== 'function') { + return auth; + } + let resolved: object | undefined; + auth((data) => { + resolved = data; + }); + return resolved; +}; + +describe('useSocketIO (mounted)', () => { + let container: HTMLDivElement; + let root: Root; + let store: ReturnType; + + const mount = () => { + act(() => { + root.render( + + + + ); + }); + }; + + beforeEach(() => { + vi.useFakeTimers(); + sockets.length = 0; + io.mockClear(); + localStorage.clear(); + store = createStore(); + container = document.createElement('div'); + document.body.appendChild(container); + root = createRoot(container); + }); + + afterEach(() => { + act(() => { + root.unmount(); + }); + container.remove(); + vi.useRealTimers(); + }); + + it('keeps the live socket across a sliding-window token refresh', () => { + act(() => { + store.dispatch(setCredentials({ token: tokenFor(user.user_id, 1), user })); + }); + mount(); + expect(io).toHaveBeenCalledTimes(1); + + act(() => { + store.dispatch(tokenRefreshed(tokenFor(user.user_id, 2))); + }); + + expect(io).toHaveBeenCalledTimes(1); + expect(sockets[0]?.disconnect).not.toHaveBeenCalled(); + }); + + it('presents the current token on every connection attempt, not the one it was built with', () => { + act(() => { + store.dispatch(setCredentials({ token: tokenFor(user.user_id, 1), user })); + }); + mount(); + + const options = io.mock.calls[0]?.[1]; + expect(resolveAuth(options)).toEqual({ token: tokenFor(user.user_id, 1) }); + + act(() => { + store.dispatch(tokenRefreshed(tokenFor(user.user_id, 2))); + }); + + // Same socket, same options object — but a reconnect now sends the refreshed token. + expect(resolveAuth(options)).toEqual({ token: tokenFor(user.user_id, 2) }); + }); + + // Note this one is the hydration gate's, not the session key's: `externalTokenAdopted` nulls + // `auth.user`, which closes the gate on its own. The session key's account-switch case is the + // same-tab `setCredentials` test below, where the gate never closes. + it('waits for the new account to hydrate before reconnecting when another tab takes over', () => { + act(() => { + store.dispatch(setCredentials({ token: tokenFor(user.user_id, 1), user })); + }); + mount(); + + act(() => { + store.dispatch(externalTokenAdopted(tokenFor('user-2', 1))); + }); + + // The adopted token nulls the current user, so no socket may connect until /me rehydrates: + // events arriving before that would be attributed to the wrong owner. + expect(sockets[0]?.disconnect).toHaveBeenCalledTimes(1); + expect(io).toHaveBeenCalledTimes(1); + + act(() => { + store.dispatch(currentUserUpdated({ ...user, user_id: 'user-2' })); + }); + + expect(io).toHaveBeenCalledTimes(2); + expect(resolveAuth(io.mock.calls[1]?.[1])).toEqual({ token: tokenFor('user-2', 1) }); + }); + + it('rebuilds the socket when a password change revokes the earlier session', () => { + act(() => { + store.dispatch(setCredentials({ token: tokenFor(user.user_id, 1), user })); + }); + mount(); + + // A password change bumps the revocation epoch and the server force-disconnects every socket + // that authenticated under the old one. socket.io never retries a server-initiated + // disconnect, so the replacement token — which carries the new epoch — has to bring a new + // socket with it, or this session has no events until the page is reloaded. + act(() => { + store.dispatch(tokenRefreshed(tokenFor(user.user_id, 2, 1))); + }); + + expect(sockets[0]?.disconnect).toHaveBeenCalledTimes(1); + expect(io).toHaveBeenCalledTimes(2); + expect(resolveAuth(io.mock.calls[1]?.[1])).toEqual({ token: tokenFor(user.user_id, 2, 1) }); + }); + + it('rebuilds the socket when a different account signs in without an intervening logout', () => { + act(() => { + store.dispatch(setCredentials({ token: tokenFor(user.user_id, 1), user })); + }); + mount(); + + // One action swaps token and user together, so the hydration gate never closes — only the + // session key separates this from a refresh. + act(() => { + store.dispatch(setCredentials({ token: tokenFor('user-2', 1), user: { ...user, user_id: 'user-2' } })); + }); + + expect(sockets[0]?.disconnect).toHaveBeenCalledTimes(1); + expect(io).toHaveBeenCalledTimes(2); + expect(resolveAuth(io.mock.calls[1]?.[1])).toEqual({ token: tokenFor('user-2', 1) }); + }); + + it('reconnects after a server-initiated disconnect, which socket.io never retries', () => { + act(() => { + store.dispatch(setCredentials({ token: tokenFor(user.user_id, 1), user })); + }); + mount(); + expect(sockets[0]?.connect).toHaveBeenCalledTimes(1); + + act(() => { + sockets[0]?.fire('disconnect', 'io server disconnect'); + }); + act(() => { + vi.advanceTimersByTime(1000); + }); + + // The same socket, reconnected — not a rebuild, so listeners and rooms survive. + expect(sockets[0]?.connect).toHaveBeenCalledTimes(2); + expect(io).toHaveBeenCalledTimes(1); + }); + + it('leaves transport-level drops to socket.io', () => { + act(() => { + store.dispatch(setCredentials({ token: tokenFor(user.user_id, 1), user })); + }); + mount(); + + act(() => { + sockets[0]?.fire('disconnect', 'transport close'); + }); + act(() => { + vi.advanceTimersByTime(60_000); + }); + + // socket.io's own backoff owns this one; a second driver would race it. + expect(sockets[0]?.connect).toHaveBeenCalledTimes(1); + }); + + it('gives up after a bounded number of server disconnects', () => { + act(() => { + store.dispatch(setCredentials({ token: tokenFor(user.user_id, 1), user })); + }); + mount(); + + for (let attempt = 0; attempt < 8; attempt++) { + act(() => { + sockets[0]?.fire('disconnect', 'io server disconnect'); + }); + act(() => { + vi.advanceTimersByTime(60_000); + }); + } + + // 1 mount + 5 retries: a server that keeps dropping this socket must not be retried forever. + expect(sockets[0]?.connect).toHaveBeenCalledTimes(6); + }); + + it('does not reconnect a socket the session has already discarded', () => { + act(() => { + store.dispatch(setCredentials({ token: tokenFor(user.user_id, 1), user })); + }); + mount(); + + act(() => { + sockets[0]?.fire('disconnect', 'io server disconnect'); + }); + // Logout lands inside the retry delay: the pending timer must not revive the old socket. + act(() => { + store.dispatch(logout()); + }); + act(() => { + vi.advanceTimersByTime(60_000); + }); + + expect(sockets[0]?.connect).toHaveBeenCalledTimes(1); + }); + + it('tears the socket down on logout', () => { + act(() => { + store.dispatch(setCredentials({ token: tokenFor(user.user_id, 1), user })); + }); + mount(); + + act(() => { + store.dispatch(logout()); + }); + + expect(sockets[0]?.disconnect).toHaveBeenCalledTimes(1); + // No token is single-user mode, where an unauthenticated socket is the correct one. + expect(io).toHaveBeenCalledTimes(2); + expect(resolveAuth(io.mock.calls[1]?.[1])).toBeUndefined(); + }); +}); diff --git a/invokeai/frontend/web/src/services/events/useSocketIO.ts b/invokeai/frontend/web/src/services/events/useSocketIO.ts index 62dc9c1e249..41153a4679e 100644 --- a/invokeai/frontend/web/src/services/events/useSocketIO.ts +++ b/invokeai/frontend/web/src/services/events/useSocketIO.ts @@ -1,7 +1,7 @@ import { useAppSelector, useAppStore } from 'app/store/storeHooks'; import { useAssertSingleton } from 'common/hooks/useAssertSingleton'; import { getBasePath, getDeploymentBaseUrl } from 'common/util/baseUrl'; -import { selectAuthToken, selectCurrentUser } from 'features/auth/store/authSlice'; +import { selectAuthSessionKey, selectAuthToken, selectCurrentUser } from 'features/auth/store/authSlice'; import type { MapStore } from 'nanostores'; import { useEffect, useMemo } from 'react'; import { selectQueueStatus } from 'services/api/endpoints/queue'; @@ -19,6 +19,22 @@ declare global { } } +/** + * A server-initiated disconnect is terminal for socket.io: `Socket.ondisconnect` destroys the + * manager subscriptions and sets `skipReconnect`, so the client never retries on its own. The + * server uses that disconnect to act on authorization changes — a deactivated account, a + * superseded token epoch (`_handle_user_access_changed` in `sockets.py`) — and expects a client + * that still holds usable credentials to come back. Without a retry here this tab has no events + * and no Invoke button until the page is reloaded, even after the change that caused it is + * reverted. + * + * Bounded, because retrying is only ever right when the client is still welcome: a client that + * is not gets a connect error, which socket.io does not retry, and the attempts stop there. The + * count lives with the socket, so a session change (which builds a new one) starts it over. + */ +const MAX_SERVER_DISCONNECT_RECONNECTS = 5; +const SERVER_DISCONNECT_RECONNECT_DELAY_MS = 1000; + /** * Initializes the socket.io connection and sets up event listeners. */ @@ -34,11 +50,19 @@ export const useSocketIO = () => { // stale token that ProtectedRoute has cleared), where every event is the client's own and the // socket can connect immediately. // - // The token also feeds socketOptions, making it a dependency of the connect effect: an in-tab - // logout or session expiry (which nulls the token) tears the authenticated socket down instead - // of letting it keep the old user's room membership — and private events — until the next full - // page reload. + // The session identity also feeds socketOptions, making it a dependency of the connect effect: + // an in-tab logout, a session expiry (which nulls the token) or another account taking over the + // tab tears the authenticated socket down instead of letting it keep the old user's room + // membership — and private events — until the next full page reload. + // + // Identity, deliberately, and not the token itself. The sliding-window middleware mints a + // replacement token on every mutating request, and the client commits one about once a minute + // during any activity (TOKEN_REFRESH_THROTTLE_MS). Keying on the token's bytes meant every one + // of those routine refreshes disconnected and rebuilt the live socket: listeners disposed, the + // progress store cleared, and a visible glitch in the preview mid-generation, once a minute, + // for a credential change that does not change who is connected or which rooms they belong to. const token = useAppSelector(selectAuthToken); + const sessionKey = useAppSelector(selectAuthSessionKey); const currentUser = useAppSelector(selectCurrentUser); const isAuthHydrated = !token || currentUser !== null; @@ -49,25 +73,38 @@ export const useSocketIO = () => { return `${wsProtocol}://${base.host}`; }, []); - // Derived from the redux token (hydrated synchronously from localStorage) rather than a - // one-time localStorage read, so the socket always authenticates with the current session's - // token and reconnects when it changes. + // Derived from the redux session (hydrated synchronously from localStorage) rather than a + // one-time localStorage read, so the socket always authenticates as the current session's user + // and is rebuilt when that user changes. const socketOptions = useMemo(() => { const options: Partial = { timeout: 60000, path: `${getBasePath()}/ws/socket.io`, autoConnect: false, // achtung! removing this breaks the dynamic middleware forceNew: true, - auth: token ? { token } : undefined, - extraHeaders: token - ? { - Authorization: `Bearer ${token}`, + // A callback, so the token is read at each connection attempt instead of being baked into + // the socket: this socket now outlives the token it was built with, and a reconnect — the + // client's own retry, or one after the network drops — must present the token that is live + // then, not the one that was live when the session started. + // + // This payload is the socket's only credential. It used to be backed by an `Authorization` + // extra header, which `_handle_connect` falls back to when the payload carries no token — + // but a header is fixed for the life of the manager, so it could only ever hold the token + // that was live when the socket was built. That made it a way to authenticate a credential + // the client had already discarded: log out in another tab while a reconnect is in flight + // and the callback correctly sends `{}`, while the stale header would still have been + // accepted. The payload is read live and cannot go stale, so the fallback only ever + // weakened it. + auth: sessionKey + ? (cb: (data: object) => void) => { + const currentToken = selectAuthToken(store.getState()); + cb(currentToken ? { token: currentToken } : {}); } : undefined, }; return options; - }, [token]); + }, [sessionKey, store]); useEffect(() => { if (!isAuthHydrated) { @@ -78,6 +115,23 @@ export const useSocketIO = () => { const disposeEventListeners = setEventListeners({ socket, store, setIsConnected: $isConnected.set }); + // See MAX_SERVER_DISCONNECT_RECONNECTS. Only `io server disconnect` is ours to retry — + // socket.io reconnects transport-level drops itself, and `io client disconnect` is this + // effect's own teardown. + let serverDisconnects = 0; + let reconnectTimeout: ReturnType | undefined; + const reconnectAfterServerDisconnect = (reason: string) => { + if (reason !== 'io server disconnect' || serverDisconnects >= MAX_SERVER_DISCONNECT_RECONNECTS) { + return; + } + const delay = SERVER_DISCONNECT_RECONNECT_DELAY_MS * 2 ** serverDisconnects; + serverDisconnects++; + reconnectTimeout = setTimeout(() => { + socket.connect(); + }, delay); + }; + socket.on('disconnect', reconnectAfterServerDisconnect); + socket.connect(); if (import.meta.env.MODE === 'development') { @@ -102,6 +156,10 @@ export const useSocketIO = () => { } unsubscribeQueueStatusListener(); // Before the socket goes: anything this session scheduled must not land in the next one. + // A pending retry is exactly that — it would reconnect a socket this session has discarded, + // with credentials that may no longer be the live ones. + clearTimeout(reconnectTimeout); + socket.off('disconnect', reconnectAfterServerDisconnect); disposeEventListeners(); socket.disconnect(); $socket.set(null);