Skip to content

Commit adb0623

Browse files
wunianze666-netizenlstein
wunianze666-netizen
and
lstein
authored
fix(auth): accept epoch-changing token refreshes (#9560)
* fix(auth): accept epoch-changing token refreshes * test(auth): preserve routine refresh throttling * test(auth): cover legacy epochless token refresh * test(auth): cover refresh after throttle window --------- Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com>
1 parent 910aec8 commit adb0623

4 files changed

Lines changed: 201 additions & 10 deletions

File tree

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,16 @@ import {
44
beginAuthTransition,
55
captureAuthGeneration,
66
createMediaAuthLock,
7+
markTokenRefreshAccepted,
78
runWithMediaAuthLock,
89
shouldAcceptRefreshedToken,
910
shouldEndSessionForUnauthorized,
11+
shouldThrottleRefreshedToken,
1012
} from './authTokenRefresh';
1113

14+
const tokenFor = (userId: string, nonce: number, epoch: number) =>
15+
`header.${btoa(JSON.stringify({ user_id: userId, nonce, token_epoch: epoch }))}.signature`;
16+
1217
describe('refreshed token acceptance', () => {
1318
beforeAll(() => {
1419
const values = new Map<string, string>();
@@ -60,6 +65,36 @@ describe('refreshed token acceptance', () => {
6065
expect(beginAuthTransition()).toBe(1);
6166
});
6267

68+
it('does not throttle the replacement token that advances the current user revocation epoch', () => {
69+
const now = vi.spyOn(Date, 'now').mockReturnValue(100_000);
70+
markTokenRefreshAccepted();
71+
72+
expect(shouldThrottleRefreshedToken(tokenFor('user', 1, 0), tokenFor('user', 2, 1))).toBe(false);
73+
74+
now.mockRestore();
75+
});
76+
77+
it('keeps routine, cross-user, and unreadable replacements throttled', () => {
78+
const now = vi.spyOn(Date, 'now').mockReturnValue(200_000);
79+
markTokenRefreshAccepted();
80+
81+
expect(shouldThrottleRefreshedToken(tokenFor('user', 1, 1), tokenFor('user', 2, 1))).toBe(true);
82+
expect(shouldThrottleRefreshedToken(tokenFor('user-a', 1, 0), tokenFor('user-b', 2, 1))).toBe(true);
83+
expect(shouldThrottleRefreshedToken('opaque-old', 'opaque-new')).toBe(true);
84+
85+
now.mockRestore();
86+
});
87+
88+
it('accepts a routine same-epoch replacement after the throttle window', () => {
89+
const now = vi.spyOn(Date, 'now').mockReturnValue(300_000);
90+
markTokenRefreshAccepted();
91+
now.mockReturnValue(360_001);
92+
93+
expect(shouldThrottleRefreshedToken(tokenFor('user', 1, 1), tokenFor('user', 2, 1))).toBe(false);
94+
95+
now.mockRestore();
96+
});
97+
6398
it('serializes media-cookie writes', async () => {
6499
const calls: string[] = [];
65100
let releaseFirst: (() => void) | undefined;

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

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { tokensBelongToSameUser } from 'features/auth/store/authSlice';
1+
import { getTokenSessionKey, tokensBelongToSameUser } from 'features/auth/store/authSlice';
22

33
const AUTH_GENERATION_KEY = 'auth_generation';
44
const MEDIA_AUTH_LOCK = 'invokeai-media-auth';
@@ -22,6 +22,18 @@ export const markTokenRefreshAccepted = () => {
2222
lastTokenRefreshAcceptedAt = Date.now();
2323
};
2424

25+
export const shouldThrottleRefreshedToken = (requestToken: string, refreshedToken: string): boolean => {
26+
if (!isTokenRefreshThrottled()) {
27+
return false;
28+
}
29+
// An epoch-changing replacement is the only credential that remains valid after revocation.
30+
// Keep every other replacement on the normal sliding-refresh throttle.
31+
return !(
32+
tokensBelongToSameUser(requestToken, refreshedToken) &&
33+
getTokenSessionKey(requestToken) !== getTokenSessionKey(refreshedToken)
34+
);
35+
};
36+
2537
type FallbackLockTicket = {
2638
choosing: boolean;
2739
expiresAt: number;

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

Lines changed: 144 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,20 @@
11
import { configureStore } from '@reduxjs/toolkit';
2+
import type { BaseQueryApi } from '@reduxjs/toolkit/query';
3+
import { tokenRefreshed } from 'features/auth/store/authSlice';
4+
import { markTokenRefreshAccepted } from 'features/auth/store/authTokenRefresh';
25
import { authApi } from 'services/api/endpoints/auth';
3-
import { beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
6+
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
47

5-
import { api } from '..';
8+
import { api, buildV1Url, dynamicBaseQuery } from '..';
69

710
/**
811
* `dynamicBaseQuery` reads the bearer token out of localStorage, and `getDeploymentBaseUrl`
912
* reads `window.location.origin`. Neither exists in the default (node) test environment.
1013
*/
11-
beforeAll(() => {
12-
const values = new Map<string, string>();
14+
const values = new Map<string, string>();
15+
16+
beforeEach(() => {
17+
values.clear();
1318
vi.stubGlobal('localStorage', {
1419
clear: () => values.clear(),
1520
getItem: (key: string) => values.get(key) ?? null,
@@ -23,8 +28,9 @@ beforeAll(() => {
2328
vi.stubGlobal('window', { location: { origin: 'http://localhost' } });
2429
});
2530

26-
beforeEach(() => {
27-
localStorage.clear();
31+
afterEach(() => {
32+
vi.restoreAllMocks();
33+
vi.unstubAllGlobals();
2834
});
2935

3036
const buildStore = () =>
@@ -33,6 +39,138 @@ const buildStore = () =>
3339
middleware: (getDefaultMiddleware) => getDefaultMiddleware().concat(api.middleware),
3440
});
3541

42+
const tokenFor = (nonce: number, epoch?: number) =>
43+
`header.${btoa(
44+
JSON.stringify({ user_id: 'user-1', nonce, ...(epoch === undefined ? {} : { token_epoch: epoch }) })
45+
)}.signature`;
46+
47+
describe('refreshed token acceptance', () => {
48+
it.each([
49+
['an explicit epoch-zero token', tokenFor(1, 0)],
50+
['a legacy token without an epoch claim', tokenFor(1)],
51+
])(
52+
'accepts an epoch-changing replacement for %s inside the routine refresh throttle window',
53+
async (_, requestToken) => {
54+
const refreshedToken = tokenFor(2, 1);
55+
localStorage.setItem('auth_token', requestToken);
56+
markTokenRefreshAccepted();
57+
58+
const events: string[] = [];
59+
const dispatch = vi.fn(() => events.push('dispatch'));
60+
const fetchMock = vi.fn((input: string | URL | Request, init?: RequestInit) => {
61+
const url = input instanceof Request ? input.url : input.toString();
62+
if (url.endsWith('/api/v1/auth/media-cookie')) {
63+
events.push('media-cookie');
64+
expect(new Headers(init?.headers).get('Authorization')).toBe(`Bearer ${refreshedToken}`);
65+
return Promise.resolve(new Response(null, { status: 204 }));
66+
}
67+
return Promise.resolve(
68+
new Response('{}', {
69+
headers: { 'content-type': 'application/json', 'X-Refreshed-Token': refreshedToken },
70+
})
71+
);
72+
});
73+
vi.stubGlobal('fetch', fetchMock);
74+
75+
await dynamicBaseQuery(
76+
buildV1Url('images/i/example.png'),
77+
{
78+
dispatch,
79+
getState: () => ({}),
80+
signal: new AbortController().signal,
81+
abort: () => {},
82+
endpoint: 'getImageDTO',
83+
type: 'query',
84+
forced: false,
85+
extra: undefined,
86+
} as unknown as BaseQueryApi,
87+
{}
88+
);
89+
90+
expect(fetchMock).toHaveBeenCalledTimes(2);
91+
expect(dispatch).toHaveBeenCalledWith(tokenRefreshed(refreshedToken));
92+
expect(events).toEqual(['media-cookie', 'dispatch']);
93+
}
94+
);
95+
96+
it('keeps a same-epoch replacement inside the routine refresh throttle window', async () => {
97+
const requestToken = tokenFor(1, 1);
98+
const refreshedToken = tokenFor(2, 1);
99+
localStorage.setItem('auth_token', requestToken);
100+
markTokenRefreshAccepted();
101+
102+
const dispatch = vi.fn();
103+
const fetchMock = vi.fn(() =>
104+
Promise.resolve(
105+
new Response('{}', {
106+
headers: { 'content-type': 'application/json', 'X-Refreshed-Token': refreshedToken },
107+
})
108+
)
109+
);
110+
vi.stubGlobal('fetch', fetchMock);
111+
112+
await dynamicBaseQuery(
113+
buildV1Url('images/i/example.png'),
114+
{
115+
dispatch,
116+
getState: () => ({}),
117+
signal: new AbortController().signal,
118+
abort: () => {},
119+
endpoint: 'getImageDTO',
120+
type: 'query',
121+
forced: false,
122+
extra: undefined,
123+
} as unknown as BaseQueryApi,
124+
{}
125+
);
126+
127+
expect(fetchMock).toHaveBeenCalledTimes(1);
128+
expect(dispatch).not.toHaveBeenCalled();
129+
});
130+
131+
it('commits a same-epoch replacement after the routine refresh throttle window', async () => {
132+
const now = vi.spyOn(Date, 'now').mockReturnValue(300_000);
133+
const requestToken = tokenFor(1, 1);
134+
const refreshedToken = tokenFor(2, 1);
135+
localStorage.setItem('auth_token', requestToken);
136+
markTokenRefreshAccepted();
137+
now.mockReturnValue(360_001);
138+
139+
const dispatch = vi.fn();
140+
const fetchMock = vi.fn((input: string | URL | Request, init?: RequestInit) => {
141+
const url = input instanceof Request ? input.url : input.toString();
142+
if (url.endsWith('/api/v1/auth/media-cookie')) {
143+
expect(new Headers(init?.headers).get('Authorization')).toBe(`Bearer ${refreshedToken}`);
144+
return Promise.resolve(new Response(null, { status: 204 }));
145+
}
146+
return Promise.resolve(
147+
new Response('{}', {
148+
headers: { 'content-type': 'application/json', 'X-Refreshed-Token': refreshedToken },
149+
})
150+
);
151+
});
152+
vi.stubGlobal('fetch', fetchMock);
153+
154+
await dynamicBaseQuery(
155+
buildV1Url('images/i/example.png'),
156+
{
157+
dispatch,
158+
getState: () => ({}),
159+
signal: new AbortController().signal,
160+
abort: () => {},
161+
endpoint: 'getImageDTO',
162+
type: 'query',
163+
forced: false,
164+
extra: undefined,
165+
} as unknown as BaseQueryApi,
166+
{}
167+
);
168+
169+
expect(fetchMock).toHaveBeenCalledTimes(2);
170+
expect(dispatch).toHaveBeenCalledWith(tokenRefreshed(refreshedToken));
171+
});
172+
});
173+
36174
describe('getCurrentUser', () => {
37175
it('does not let a replacement session read the 401 of the token it replaced', async () => {
38176
// 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‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,12 @@ import { sessionExpiredLogout, tokenRefreshed } from 'features/auth/store/authSl
1212
import {
1313
beginAuthTransition,
1414
captureAuthGeneration,
15-
isTokenRefreshThrottled,
1615
markTokenRefreshAccepted,
1716
MEDIA_COOKIE_SYNC_TIMEOUT_MS,
1817
runWithMediaAuthLock,
1918
shouldAcceptRefreshedToken,
2019
shouldEndSessionForUnauthorized,
20+
shouldThrottleRefreshedToken,
2121
} from 'features/auth/store/authTokenRefresh';
2222
import queryString from 'query-string';
2323
import stableHash from 'stable-hash';
@@ -174,11 +174,17 @@ export const acceptRefreshedToken = async (
174174
requestGeneration: number,
175175
dispatch: (action: ReturnType<typeof tokenRefreshed>) => unknown
176176
): Promise<void> => {
177-
if (isTokenRefreshThrottled() || !shouldAcceptRefreshedToken(requestToken, requestGeneration)) {
177+
if (
178+
shouldThrottleRefreshedToken(requestToken, refreshedToken) ||
179+
!shouldAcceptRefreshedToken(requestToken, requestGeneration)
180+
) {
178181
return;
179182
}
180183
await runWithMediaAuthLock(async () => {
181-
if (isTokenRefreshThrottled() || !shouldAcceptRefreshedToken(requestToken, requestGeneration)) {
184+
if (
185+
shouldThrottleRefreshedToken(requestToken, refreshedToken) ||
186+
!shouldAcceptRefreshedToken(requestToken, requestGeneration)
187+
) {
182188
return;
183189
}
184190
try {

0 commit comments

Comments
 (0)