Skip to content

Commit 16e6460

Browse files
sam-gcsam-gc
committed
Add onAbort and refactor middleware
1 parent 3785b07 commit 16e6460

File tree

6 files changed

+236
-44
lines changed

6 files changed

+236
-44
lines changed

‎common/api-review/auth.api.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ export function applyActionCode(auth: Auth, oobCode: string): Promise<void>;
8181
// @public
8282
export interface Auth {
8383
readonly app: FirebaseApp;
84-
beforeAuthStateChanged(callback: (user: User | null) => void | Promise<void>): Unsubscribe;
84+
beforeAuthStateChanged(callback: (user: User | null) => void | Promise<void>, onAbort?: () => void): Unsubscribe;
8585
readonly config: Config;
8686
readonly currentUser: User | null;
8787
readonly emulatorConfig: EmulatorConfig | null;

‎packages/auth/src/core/auth/auth_impl.ts

+8-40
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ import { _getInstance } from '../util/instantiator';
6060
import { _getUserLanguage } from '../util/navigator';
6161
import { _getClientVersion } from '../util/version';
6262
import { HttpHeader } from '../../api';
63+
import { AuthMiddlewareQueue } from './middleware';
6364

6465
interface AsyncAction {
6566
(): Promise<void>;
@@ -79,7 +80,7 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
7980
private redirectPersistenceManager?: PersistenceUserManager;
8081
private authStateSubscription = new Subscription<User>(this);
8182
private idTokenSubscription = new Subscription<User>(this);
82-
private beforeStateQueue: Array<(user: User | null) => Promise<void>> = [];
83+
private readonly beforeStateQueue = new AuthMiddlewareQueue(this);
8384
private redirectUser: UserInternal | null = null;
8485
private isProactiveRefreshEnabled = false;
8586

@@ -225,7 +226,7 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
225226
// First though, ensure that we check the middleware is happy.
226227
if (needsTocheckMiddleware) {
227228
try {
228-
await this._runBeforeStateCallbacks(futureCurrentUser);
229+
await this.beforeStateQueue.runMiddleware(futureCurrentUser);
229230
} catch(e) {
230231
futureCurrentUser = previouslyStoredUser;
231232
// We know this is available since the bit is only set when the
@@ -347,7 +348,7 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
347348
}
348349

349350
if (!skipBeforeStateCallbacks) {
350-
await this._runBeforeStateCallbacks(user);
351+
await this.beforeStateQueue.runMiddleware(user);
351352
}
352353

353354
return this.queue(async () => {
@@ -356,23 +357,9 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
356357
});
357358
}
358359

359-
async _runBeforeStateCallbacks(user: User | null): Promise<void> {
360-
if (this.currentUser === user) {
361-
return;
362-
}
363-
try {
364-
for (const beforeStateCallback of this.beforeStateQueue) {
365-
await beforeStateCallback(user);
366-
}
367-
} catch (e) {
368-
throw this._errorFactory.create(
369-
AuthErrorCode.LOGIN_BLOCKED, { originalMessage: e.message });
370-
}
371-
}
372-
373360
async signOut(): Promise<void> {
374361
// Run first, to block _setRedirectUser() if any callbacks fail.
375-
await this._runBeforeStateCallbacks(null);
362+
await this.beforeStateQueue.runMiddleware(null);
376363
// Clear the redirect user when signOut is called
377364
if (this.redirectPersistenceManager || this._popupRedirectResolver) {
378365
await this._setRedirectUser(null);
@@ -415,29 +402,10 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
415402
}
416403

417404
beforeAuthStateChanged(
418-
callback: (user: User | null) => void | Promise<void>
405+
callback: (user: User | null) => void | Promise<void>,
406+
onAbort?: () => void,
419407
): Unsubscribe {
420-
// The callback could be sync or async. Wrap it into a
421-
// function that is always async.
422-
const wrappedCallback =
423-
(user: User | null): Promise<void> => new Promise((resolve, reject) => {
424-
try {
425-
const result = callback(user);
426-
// Either resolve with existing promise or wrap a non-promise
427-
// return value into a promise.
428-
resolve(result);
429-
} catch (e) {
430-
// Sync callback throws.
431-
reject(e);
432-
}
433-
});
434-
this.beforeStateQueue.push(wrappedCallback);
435-
const index = this.beforeStateQueue.length - 1;
436-
return () => {
437-
// Unsubscribe. Replace with no-op. Do not remove from array, or it will disturb
438-
// indexing of other elements.
439-
this.beforeStateQueue[index] = () => Promise.resolve();
440-
};
408+
return this.beforeStateQueue.pushCallback(callback, onAbort);
441409
}
442410

443411
onIdTokenChanged(

‎packages/auth/src/core/auth/middleware.test.ts

+132
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
import { expect, use } from 'chai';
2+
import chaiAsPromised from 'chai-as-promised';
3+
import * as sinon from 'sinon';
4+
import sinonChai from 'sinon-chai';
5+
import { testAuth, testUser } from '../../../test/helpers/mock_auth';
6+
import { AuthInternal } from '../../model/auth';
7+
import { User } from '../../model/public_types';
8+
import { AuthMiddlewareQueue } from './middleware';
9+
10+
use(chaiAsPromised);
11+
use(sinonChai);
12+
13+
describe('Auth middleware', () => {
14+
let middlewareQueue: AuthMiddlewareQueue;
15+
let user: User;
16+
let auth: AuthInternal;
17+
18+
beforeEach(async () => {
19+
auth = await testAuth();
20+
user = testUser(auth, 'uid');
21+
middlewareQueue = new AuthMiddlewareQueue(auth);
22+
});
23+
24+
afterEach(() => {
25+
sinon.restore();
26+
});
27+
28+
it('calls middleware in order', async () => {
29+
const calls: number[] = [];
30+
31+
middlewareQueue.pushCallback(() => {calls.push(1);});
32+
middlewareQueue.pushCallback(() => {calls.push(2);});
33+
middlewareQueue.pushCallback(() => {calls.push(3);});
34+
35+
await middlewareQueue.runMiddleware(user);
36+
37+
expect(calls).to.eql([1, 2, 3]);
38+
});
39+
40+
it('rejects on error', async () => {
41+
middlewareQueue.pushCallback(() => {
42+
throw new Error('no');
43+
});
44+
await expect(middlewareQueue.runMiddleware(user)).to.be.rejectedWith('auth/login-blocked');
45+
});
46+
47+
it('rejects on promise rejection', async () => {
48+
middlewareQueue.pushCallback(() => Promise.reject('no'));
49+
await expect(middlewareQueue.runMiddleware(user)).to.be.rejectedWith('auth/login-blocked');
50+
});
51+
52+
it('awaits middleware completion before calling next', async () => {
53+
const firstCb = sinon.spy();
54+
const secondCb = sinon.spy();
55+
56+
middlewareQueue.pushCallback(() => {
57+
// Force the first one to run one tick later
58+
return new Promise(resolve => {
59+
setTimeout(() => {
60+
firstCb();
61+
resolve();
62+
}, 1);
63+
});
64+
});
65+
middlewareQueue.pushCallback(secondCb);
66+
67+
await middlewareQueue.runMiddleware(user);
68+
expect(secondCb).to.have.been.calledAfter(firstCb);
69+
});
70+
71+
it('subsequent middleware not run after rejection', async () => {
72+
const spy = sinon.spy();
73+
74+
middlewareQueue.pushCallback(() => {
75+
throw new Error('no');
76+
});
77+
middlewareQueue.pushCallback(spy);
78+
79+
await expect(middlewareQueue.runMiddleware(user)).to.be.rejectedWith('auth/login-blocked');
80+
expect(spy).not.to.have.been.called;
81+
});
82+
83+
it('calls onAbort if provided but only for earlier runs', async () => {
84+
const firstOnAbort = sinon.spy();
85+
const secondOnAbort = sinon.spy();
86+
87+
middlewareQueue.pushCallback(() => {}, firstOnAbort);
88+
middlewareQueue.pushCallback(() => {
89+
throw new Error('no');
90+
}, secondOnAbort);
91+
92+
await expect(middlewareQueue.runMiddleware(user)).to.be.rejectedWith('auth/login-blocked');
93+
expect(firstOnAbort).to.have.been.called;
94+
expect(secondOnAbort).not.to.have.been.called;
95+
});
96+
97+
it('calls onAbort in order', async () => {
98+
const calls: number[] = [];
99+
100+
middlewareQueue.pushCallback(() => {}, () => {calls.push(1);});
101+
middlewareQueue.pushCallback(() => {}, () => {calls.push(2);});
102+
middlewareQueue.pushCallback(() => {}, () => {calls.push(3);});
103+
middlewareQueue.pushCallback(() => {
104+
throw new Error('no');
105+
});
106+
107+
await expect(middlewareQueue.runMiddleware(user)).to.be.rejectedWith('auth/login-blocked');
108+
expect(calls).to.eql([3, 2, 1]);
109+
});
110+
111+
it('does not call any middleware if user matches null', async () => {
112+
const spy = sinon.spy();
113+
114+
middlewareQueue.pushCallback(spy);
115+
await middlewareQueue.runMiddleware(null);
116+
117+
expect(spy).not.to.have.been.called;
118+
});
119+
120+
it('does not call any middleware if user matches object', async () => {
121+
const spy = sinon.spy();
122+
123+
// Directly set it manually since the public function creates a
124+
// copy of the user.
125+
auth.currentUser = user;
126+
127+
middlewareQueue.pushCallback(spy);
128+
await middlewareQueue.runMiddleware(user);
129+
130+
expect(spy).not.to.have.been.called;
131+
});
132+
});

‎packages/auth/src/core/auth/middleware.ts

+76
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
import { AuthInternal } from '../../model/auth';
2+
import { Unsubscribe, User } from '../../model/public_types';
3+
import { AuthErrorCode } from '../errors';
4+
5+
interface MiddlewareEntry {
6+
(user: User | null): Promise<void>;
7+
onAbort?: () => void;
8+
}
9+
10+
export class AuthMiddlewareQueue {
11+
private readonly queue: MiddlewareEntry[] = [];
12+
13+
constructor(private readonly auth: AuthInternal) {}
14+
15+
pushCallback(
16+
callback: (user: User | null) => void | Promise<void>,
17+
onAbort?: () => void): Unsubscribe {
18+
// The callback could be sync or async. Wrap it into a
19+
// function that is always async.
20+
const wrappedCallback: MiddlewareEntry =
21+
(user: User | null): Promise<void> => new Promise((resolve, reject) => {
22+
try {
23+
const result = callback(user);
24+
// Either resolve with existing promise or wrap a non-promise
25+
// return value into a promise.
26+
resolve(result);
27+
} catch (e) {
28+
// Sync callback throws.
29+
reject(e);
30+
}
31+
});
32+
// Attach the onAbort if present
33+
wrappedCallback.onAbort = onAbort;
34+
this.queue.push(wrappedCallback);
35+
36+
const index = this.queue.length - 1;
37+
return () => {
38+
// Unsubscribe. Replace with no-op. Do not remove from array, or it will disturb
39+
// indexing of other elements.
40+
this.queue[index] = () => Promise.resolve();
41+
};
42+
}
43+
44+
async runMiddleware(nextUser: User | null): Promise<void> {
45+
if (this.auth.currentUser === nextUser) {
46+
return;
47+
}
48+
49+
// While running the middleware, build a temporary stack of onAbort
50+
// callbacks to call if one middleware callback rejects.
51+
52+
const onAbortStack: Array<() => void> = [];
53+
try {
54+
for (const beforeStateCallback of this.queue) {
55+
await beforeStateCallback(nextUser);
56+
57+
// Only push the onAbort if the callback succeeds
58+
if (beforeStateCallback.onAbort) {
59+
onAbortStack.push(beforeStateCallback.onAbort);
60+
}
61+
}
62+
} catch (e) {
63+
// Run all onAbort, with separate try/catch to ignore any errors and
64+
// continue
65+
onAbortStack.reverse();
66+
for (const onAbort of onAbortStack) {
67+
try {
68+
onAbort();
69+
} catch (_) { /* swallow error */}
70+
}
71+
72+
throw this.auth._errorFactory.create(
73+
AuthErrorCode.LOGIN_BLOCKED, { originalMessage: e.message });
74+
}
75+
}
76+
}

‎packages/auth/src/model/public_types.ts

+4-1
Original file line numberDiff line numberDiff line change
@@ -260,9 +260,12 @@ export interface Auth {
260260
*
261261
* @param callback - callback triggered before new user value is set.
262262
* If this throws, it will block the user from being set.
263+
* @param onAbort - callback triggered if a later before state changed
264+
* callback throws, allowing you to undo any side effects.
263265
*/
264266
beforeAuthStateChanged(
265-
callback: (user: User | null) => void | Promise<void>
267+
callback: (user: User | null) => void | Promise<void>,
268+
onAbort?: () => void,
266269
): Unsubscribe;
267270
/**
268271
* Adds an observer for changes to the signed-in user's ID token.

‎packages/auth/test/integration/flows/middleware_test_generator.ts

+15-2
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,8 @@ export function generateMiddlewareTests(authGetter: () => Auth, signIn: () => Pr
4848
* automatically unsubscribe after every test (since some tests may
4949
* perform cleanup after that would be affected by the middleware)
5050
*/
51-
function beforeAuthStateChanged(callback: (user: User | null) => void | Promise<void>): void {
52-
unsubscribes.push(auth.beforeAuthStateChanged(callback));
51+
function beforeAuthStateChanged(callback: (user: User | null) => void | Promise<void>, onAbort?: () => void): void {
52+
unsubscribes.push(auth.beforeAuthStateChanged(callback, onAbort));
5353
}
5454

5555
it('can prevent user sign in', async () => {
@@ -192,5 +192,18 @@ export function generateMiddlewareTests(authGetter: () => Auth, signIn: () => Pr
192192
await expect(auth.signOut()).to.be.rejectedWith('auth/login-blocked');
193193
expect(auth.currentUser).to.eq(user);
194194
});
195+
196+
it('calls onAbort after rejection', async () => {
197+
const onAbort = sinon.spy();
198+
beforeAuthStateChanged(() => {
199+
// Pass
200+
}, onAbort);
201+
beforeAuthStateChanged(() => {
202+
throw new Error('block sign out');
203+
});
204+
205+
await expect(signIn()).to.be.rejectedWith('auth/login-blocked');
206+
expect(onAbort).to.have.been.called;
207+
});
195208
});
196209
}

0 commit comments

Comments
 (0)