fix(auth): accept epoch-changing token refreshes

invokeai/frontend/web/src/features/auth/store/authTokenRefresh.test.ts

@@ -4,11 +4,16 @@ import {
   beginAuthTransition,
   captureAuthGeneration,
   createMediaAuthLock,
+  markTokenRefreshAccepted,
   runWithMediaAuthLock,
   shouldAcceptRefreshedToken,
   shouldEndSessionForUnauthorized,
+  shouldThrottleRefreshedToken,
 } from './authTokenRefresh';
 
+const tokenFor = (userId: string, nonce: number, epoch: number) =>
+  `header.${btoa(JSON.stringify({ user_id: userId, nonce, token_epoch: epoch }))}.signature`;
+
 describe('refreshed token acceptance', () => {
   beforeAll(() => {
     const values = new Map<string, string>();
@@ -60,6 +65,26 @@ describe('refreshed token acceptance', () => {
     expect(beginAuthTransition()).toBe(1);
   });
 
+  it('does not throttle the replacement token that advances the current user revocation epoch', () => {
+    const now = vi.spyOn(Date, 'now').mockReturnValue(100_000);
+    markTokenRefreshAccepted();
+
+    expect(shouldThrottleRefreshedToken(tokenFor('user', 1, 0), tokenFor('user', 2, 1))).toBe(false);
+
+    now.mockRestore();
+  });
+
+  it('keeps routine, cross-user, and unreadable replacements throttled', () => {
+    const now = vi.spyOn(Date, 'now').mockReturnValue(200_000);
+    markTokenRefreshAccepted();
+
+    expect(shouldThrottleRefreshedToken(tokenFor('user', 1, 1), tokenFor('user', 2, 1))).toBe(true);
+    expect(shouldThrottleRefreshedToken(tokenFor('user-a', 1, 0), tokenFor('user-b', 2, 1))).toBe(true);
+    expect(shouldThrottleRefreshedToken('opaque-old', 'opaque-new')).toBe(true);
+
+    now.mockRestore();
+  });
+
   it('serializes media-cookie writes', async () => {
     const calls: string[] = [];
     let releaseFirst: (() => void) | undefined;

invokeai/frontend/web/src/features/auth/store/authTokenRefresh.ts

@@ -1,4 +1,4 @@
-import { tokensBelongToSameUser } from 'features/auth/store/authSlice';
+import { getTokenSessionKey, tokensBelongToSameUser } from 'features/auth/store/authSlice';
 
 const AUTH_GENERATION_KEY = 'auth_generation';
 const MEDIA_AUTH_LOCK = 'invokeai-media-auth';
@@ -22,6 +22,18 @@ export const markTokenRefreshAccepted = () => {
   lastTokenRefreshAcceptedAt = Date.now();
 };
 
+export const shouldThrottleRefreshedToken = (requestToken: string, refreshedToken: string): boolean => {
+  if (!isTokenRefreshThrottled()) {
+    return false;
+  }
+  // An epoch-changing replacement is the only credential that remains valid after revocation.
+  // Keep every other replacement on the normal sliding-refresh throttle.
+  return !(
+    tokensBelongToSameUser(requestToken, refreshedToken) &&
+    getTokenSessionKey(requestToken) !== getTokenSessionKey(refreshedToken)
+  );
+};
+
 type FallbackLockTicket = {
   choosing: boolean;
   expiresAt: number;

invokeai/frontend/web/src/services/api/endpoints/auth.test.ts

@@ -1,8 +1,11 @@
 import { configureStore } from '@reduxjs/toolkit';
+import type { BaseQueryApi } from '@reduxjs/toolkit/query';
+import { tokenRefreshed } from 'features/auth/store/authSlice';
+import { markTokenRefreshAccepted } from 'features/auth/store/authTokenRefresh';
 import { authApi } from 'services/api/endpoints/auth';
 import { beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
 
-import { api } from '..';
+import { api, buildV1Url, dynamicBaseQuery } from '..';
 
 /**
  * `dynamicBaseQuery` reads the bearer token out of localStorage, and `getDeploymentBaseUrl`
@@ -33,6 +36,93 @@ const buildStore = () =>
     middleware: (getDefaultMiddleware) => getDefaultMiddleware().concat(api.middleware),
   });
 
+const tokenFor = (nonce: number, epoch?: number) =>
+  `header.${btoa(
+    JSON.stringify({ user_id: 'user-1', nonce, ...(epoch === undefined ? {} : { token_epoch: epoch }) })
+  )}.signature`;
+
+describe('refreshed token acceptance', () => {
+  it.each([
+    ['an explicit epoch-zero token', tokenFor(1, 0)],
+    ['a legacy token without an epoch claim', tokenFor(1)],
+  ])(
+    'accepts an epoch-changing replacement for %s inside the routine refresh throttle window',
+    async (_, requestToken) => {
+      const refreshedToken = tokenFor(2, 1);
+      localStorage.setItem('auth_token', requestToken);
+      markTokenRefreshAccepted();
+
+      const dispatch = vi.fn();
+      const fetchMock = vi.fn((input: string | URL | Request, init?: RequestInit) => {
+        const url = input instanceof Request ? input.url : input.toString();
+        if (url.endsWith('/api/v1/auth/media-cookie')) {
+          expect(new Headers(init?.headers).get('Authorization')).toBe(`Bearer ${refreshedToken}`);
+          return Promise.resolve(new Response(null, { status: 204 }));
+        }
+        return Promise.resolve(
+          new Response('{}', {
+            headers: { 'content-type': 'application/json', 'X-Refreshed-Token': refreshedToken },
+          })
+        );
+      });
+      vi.stubGlobal('fetch', fetchMock);
+
+      await dynamicBaseQuery(
+        buildV1Url('images/i/example.png'),
+        {
+          dispatch,
+          getState: () => ({}),
+          signal: new AbortController().signal,
+          abort: () => {},
+          endpoint: 'getImageDTO',
+          type: 'query',
+          forced: false,
+          extra: undefined,
+        } as unknown as BaseQueryApi,
+        {}
+      );
+
+      expect(fetchMock).toHaveBeenCalledTimes(2);
+      expect(dispatch).toHaveBeenCalledWith(tokenRefreshed(refreshedToken));
+    }
+  );
+
+  it('keeps a same-epoch replacement inside the routine refresh throttle window', async () => {
+    const requestToken = tokenFor(1, 1);
+    const refreshedToken = tokenFor(2, 1);
+    localStorage.setItem('auth_token', requestToken);
+    markTokenRefreshAccepted();
+
+    const dispatch = vi.fn();
+    const fetchMock = vi.fn(() =>
+      Promise.resolve(
+        new Response('{}', {
+          headers: { 'content-type': 'application/json', 'X-Refreshed-Token': refreshedToken },
+        })
+      )
+    );
+    vi.stubGlobal('fetch', fetchMock);
+
+    await dynamicBaseQuery(
+      buildV1Url('images/i/example.png'),
+      {
+        dispatch,
+        getState: () => ({}),
+        signal: new AbortController().signal,
+        abort: () => {},
+        endpoint: 'getImageDTO',
+        type: 'query',
+        forced: false,
+        extra: undefined,
+      } as unknown as BaseQueryApi,
+      {}
+    );
+
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+    expect(dispatch).not.toHaveBeenCalled();
+  });
+});
+
 describe('getCurrentUser', () => {
   it('does not let a replacement session read the 401 of the token it replaced', async () => {
     // The sequence this exists for: a tab page-loads with an expired token and asks who it is;

invokeai/frontend/web/src/services/api/index.ts

@@ -12,12 +12,12 @@ import { sessionExpiredLogout, tokenRefreshed } from 'features/auth/store/authSl
 import {
   beginAuthTransition,
   captureAuthGeneration,
-  isTokenRefreshThrottled,
   markTokenRefreshAccepted,
   MEDIA_COOKIE_SYNC_TIMEOUT_MS,
   runWithMediaAuthLock,
   shouldAcceptRefreshedToken,
   shouldEndSessionForUnauthorized,
+  shouldThrottleRefreshedToken,
 } from 'features/auth/store/authTokenRefresh';
 import queryString from 'query-string';
 import stableHash from 'stable-hash';
@@ -174,11 +174,17 @@ export const acceptRefreshedToken = async (
   requestGeneration: number,
   dispatch: (action: ReturnType<typeof tokenRefreshed>) => unknown
 ): Promise<void> => {
-  if (isTokenRefreshThrottled() || !shouldAcceptRefreshedToken(requestToken, requestGeneration)) {
+  if (
+    shouldThrottleRefreshedToken(requestToken, refreshedToken) ||
+    !shouldAcceptRefreshedToken(requestToken, requestGeneration)
+  ) {
     return;
   }
   await runWithMediaAuthLock(async () => {
-    if (isTokenRefreshThrottled() || !shouldAcceptRefreshedToken(requestToken, requestGeneration)) {
+    if (
+      shouldThrottleRefreshedToken(requestToken, refreshedToken) ||
+      !shouldAcceptRefreshedToken(requestToken, requestGeneration)
+    ) {
       return;
     }
     try {