perf(auth): cache auth session lookups by token hash (#2225)

fluxer_api/src/api/auth/tests/AuthSessionCacheInvalidation.test.ts

@@ -0,0 +1,128 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+import {createHash} from 'node:crypto';
+import type {AuthSessionResponse} from '@fluxer/schema/src/domains/auth/AuthSchemas';
+import {afterAll, beforeAll, beforeEach, describe, expect, it} from 'vitest';
+import type {ApiTestHarness} from '../../test/ApiTestHarness';
+import {createBuilder} from '../../test/TestRequestBuilder';
+import {
+	createAuthHarness,
+	createFakeAuthToken,
+	createTestAccount,
+	loginAccount,
+	setUserACLs,
+	type TestAccount,
+} from './AuthTestUtils';
+
+function authSessionCacheKey(token: string): string {
+	return `auth:session:${createHash('sha256').update(token).digest('base64url')}`;
+}
+
+async function readCachedSession(harness: ApiTestHarness, token: string): Promise<string | null> {
+	return harness.kvProvider.get(authSessionCacheKey(token));
+}
+
+describe('Auth session cache invalidation', () => {
+	let harness: ApiTestHarness;
+	beforeAll(async () => {
+		harness = await createAuthHarness();
+	});
+	beforeEach(async () => {
+		await harness.reset();
+	});
+	afterAll(async () => {
+		await harness?.shutdown();
+	});
+	it('caches the session on the token hash and never on the raw token', async () => {
+		const account = await createTestAccount(harness);
+		await createBuilder(harness, account.token).get('/users/@me').expect(200).execute();
+		const cached = await readCachedSession(harness, account.token);
+		expect(cached).not.toBeNull();
+		expect(cached).not.toContain(account.token);
+		await expect(harness.kvProvider.get(`auth:session:${account.token}`)).resolves.toBeNull();
+	});
+	it('stops serving a session revoked by logout', async () => {
+		const account = await createTestAccount(harness);
+		await createBuilder(harness, account.token).get('/users/@me').expect(200).execute();
+		expect(await readCachedSession(harness, account.token)).not.toBeNull();
+		await createBuilder(harness, account.token).post('/auth/logout').expect(204).execute();
+		expect(await readCachedSession(harness, account.token)).toBeNull();
+		await createBuilder(harness, account.token).get('/users/@me').expect(401).execute();
+	});
+	it('stops serving a session revoked from another session list', async () => {
+		const account = await createTestAccount(harness);
+		const second = await loginAccount(harness, account);
+		await createBuilder(harness, second.token).get('/users/@me').expect(200).execute();
+		expect(await readCachedSession(harness, second.token)).not.toBeNull();
+		const sessions = await createBuilder<Array<AuthSessionResponse>>(harness, account.token)
+			.get('/auth/sessions')
+			.execute();
+		const secondSessionIdHash = createHash('sha256').update(second.token).digest('base64url');
+		expect(sessions.some((session) => session.id_hash === secondSessionIdHash)).toBe(true);
+		await createBuilder(harness, account.token)
+			.post('/auth/sessions/logout')
+			.body({session_id_hashes: [secondSessionIdHash], password: account.password})
+			.expect(204)
+			.execute();
+		expect(await readCachedSession(harness, second.token)).toBeNull();
+		await createBuilder(harness, second.token).get('/users/@me').expect(401).execute();
+		await createBuilder(harness, account.token).get('/users/@me').expect(200).execute();
+	});
+	it('stops serving every session of a user after a password change', async () => {
+		const account = await createTestAccount(harness);
+		const second = await loginAccount(harness, account);
+		const third = await loginAccount(harness, account);
+		for (const token of [account.token, second.token, third.token]) {
+			await createBuilder(harness, token).get('/users/@me').expect(200).execute();
+			expect(await readCachedSession(harness, token)).not.toBeNull();
+		}
+		await createBuilder(harness, account.token)
+			.patch('/users/@me')
+			.body({password: account.password, new_password: `cache-rotation-${Date.now()}`})
+			.execute();
+		for (const token of [account.token, second.token, third.token]) {
+			expect(await readCachedSession(harness, token)).toBeNull();
+			await createBuilder(harness, token).get('/users/@me').expect(401).execute();
+		}
+	});
+	it('stops serving every session of a user after an admin temp ban', async () => {
+		const target = await createTestAccount(harness);
+		const targetSecond = await loginAccount(harness, target);
+		let admin: TestAccount = await createTestAccount(harness);
+		admin = await setUserACLs(harness, admin, ['admin:authenticate', 'user:temp_ban']);
+		for (const token of [target.token, targetSecond.token]) {
+			await createBuilder(harness, token).get('/users/@me').expect(200).execute();
+			expect(await readCachedSession(harness, token)).not.toBeNull();
+		}
+		await createBuilder(harness, admin.token)
+			.post('/admin/users/temp-ban')
+			.body({user_id: target.userId, duration_hours: 24, reason: 'cache invalidation coverage'})
+			.execute();
+		for (const token of [target.token, targetSecond.token]) {
+			expect(await readCachedSession(harness, token)).toBeNull();
+			await createBuilder(harness, token).get('/users/@me').expect(401).execute();
+		}
+	});
+	it('stops serving sessions after the account disables itself', async () => {
+		const account = await createTestAccount(harness);
+		await createBuilder(harness, account.token).get('/users/@me').expect(200).execute();
+		expect(await readCachedSession(harness, account.token)).not.toBeNull();
+		await createBuilder(harness, account.token)
+			.post('/users/@me/disable')
+			.body({password: account.password})
+			.expect(204)
+			.execute();
+		expect(await readCachedSession(harness, account.token)).toBeNull();
+		await createBuilder(harness, account.token).get('/users/@me').expect(401).execute();
+	});
+	it('caches an unknown token hash without stranding sessions created afterwards', async () => {
+		const unknownToken = createFakeAuthToken();
+		await createBuilder(harness, unknownToken).get('/users/@me').expect(401).execute();
+		await expect(readCachedSession(harness, unknownToken)).resolves.toBe('null');
+		await createBuilder(harness, unknownToken).get('/users/@me').expect(401).execute();
+		const account = await createTestAccount(harness);
+		await createBuilder(harness, account.token).get('/users/@me').expect(200).execute();
+		const rotated = await loginAccount(harness, account);
+		await createBuilder(harness, rotated.token).get('/users/@me').expect(200).execute();
+	});
+});

fluxer_api/src/api/user/repositories/auth/AuthSessionRepository.ts

@@ -1,15 +1,70 @@
 // SPDX-License-Identifier: AGPL-3.0-or-later
 
-import type {UserID} from '../../../BrandedTypes';
+import {createUserID, type UserID} from '../../../BrandedTypes';
 import {BatchBuilder, fetchMany, fetchOne, upsertOne} from '../../../database/CassandraQueryExecution';
 import {Db} from '../../../database/CassandraTypes';
 import type {AuthSessionRow, AuthSessionTombstoneRow, UserCountryHistoryRow} from '../../../database/types/AuthTypes';
 import {Logger} from '../../../Logger';
-import {getPhoneFraudGraphService} from '../../../middleware/ServiceSingletons';
+import {getCacheService, getPhoneFraudGraphService} from '../../../middleware/ServiceSingletons';
 import {AuthSession, AuthSessionTombstone} from '../../../models/AuthSession';
 import {AuthSessions, AuthSessionsByUserId, AuthSessionTombstones, UserCountryHistory} from '../../../Tables';
 
-function invalidateAuthSessionCache(_sessionIdHash: Buffer): void {}
+const AUTH_SESSION_CACHE_TTL_SECONDS = 30;
+const AUTH_SESSION_MISS_CACHE_TTL_SECONDS = 5;
+
+interface CachedAuthSession {
+	user_id: string;
+	session_id_hash: string;
+	created_at: number;
+	approx_last_used_at: number;
+	client_ip: string;
+	client_user_agent: string | null;
+	client_os: string | null;
+	client_country: string | null;
+	version: number;
+}
+
+function authSessionCacheKey(sessionIdHash: Buffer): string {
+	return `auth:session:${sessionIdHash.toString('base64url')}`;
+}
+
+function encodeCachedAuthSession(row: AuthSessionRow): CachedAuthSession {
+	return {
+		user_id: row.user_id.toString(),
+		session_id_hash: row.session_id_hash.toString('base64url'),
+		created_at: row.created_at.getTime(),
+		approx_last_used_at: row.approx_last_used_at.getTime(),
+		client_ip: row.client_ip,
+		client_user_agent: row.client_user_agent,
+		client_os: row.client_os,
+		client_country: row.client_country,
+		version: row.version,
+	};
+}
+
+function decodeCachedAuthSession(cached: CachedAuthSession): AuthSessionRow {
+	return {
+		user_id: createUserID(BigInt(cached.user_id)),
+		session_id_hash: Buffer.from(cached.session_id_hash, 'base64url'),
+		created_at: new Date(cached.created_at),
+		approx_last_used_at: new Date(cached.approx_last_used_at),
+		client_ip: cached.client_ip,
+		client_user_agent: cached.client_user_agent,
+		client_os: cached.client_os,
+		client_country: cached.client_country,
+		version: cached.version,
+	};
+}
+
+async function invalidateAuthSessionCache(sessionIdHashes: ReadonlyArray<Buffer>): Promise<void> {
+	if (sessionIdHashes.length === 0) return;
+	try {
+		const cache = getCacheService();
+		await Promise.all(sessionIdHashes.map((sessionIdHash) => cache.delete(authSessionCacheKey(sessionIdHash))));
+	} catch (error) {
+		Logger.error({error}, 'Failed to invalidate cached auth sessions; they expire with the cache ttl');
+	}
+}
 
 const FETCH_AUTH_SESSIONS_CQL = AuthSessions.selectCql({
 	where: AuthSessions.where.in('session_id_hash', 'session_id_hashes'),
@@ -46,6 +101,7 @@ export class AuthSessionRepository {
 			}),
 		);
 		await batch.execute();
+		await invalidateAuthSessionCache([sessionData.session_id_hash]);
 		try {
 			await getPhoneFraudGraphService().recordSessionForCohortGraph(
 				sessionData.user_id,
@@ -107,10 +163,27 @@ export class AuthSessionRepository {
 	}
 
 	async getAuthSessionByToken(sessionIdHash: Buffer): Promise<AuthSession | null> {
-		const session = await fetchOne<AuthSessionRow>(FETCH_AUTH_SESSION_BY_TOKEN_CQL, {
+		try {
+			const cached = await getCacheService().getOrSet<CachedAuthSession | null>(
+				authSessionCacheKey(sessionIdHash),
+				async () => {
+					const session = await this.fetchAuthSessionByToken(sessionIdHash);
+					return session ? encodeCachedAuthSession(session) : null;
+				},
+				(value) => (value === null ? AUTH_SESSION_MISS_CACHE_TTL_SECONDS : AUTH_SESSION_CACHE_TTL_SECONDS),
+			);
+			return cached ? new AuthSession(decodeCachedAuthSession(cached)) : null;
+		} catch (error) {
+			Logger.warn({error}, 'Auth session cache lookup failed; falling back to the datastore');
+			const session = await this.fetchAuthSessionByToken(sessionIdHash);
+			return session ? new AuthSession(session) : null;
+		}
+	}
+
+	private async fetchAuthSessionByToken(sessionIdHash: Buffer): Promise<AuthSessionRow | null> {
+		return fetchOne<AuthSessionRow>(FETCH_AUTH_SESSION_BY_TOKEN_CQL, {
 			session_id_hash: sessionIdHash,
 		});
-		return session ? new AuthSession(session) : null;
 	}
 
 	async listAuthSessions(userId: UserID): Promise<Array<AuthSession>> {
@@ -138,11 +211,12 @@ export class AuthSessionRepository {
 		await upsertOne(
 			AuthSessions.patchByPk({session_id_hash: sessionIdHash}, {approx_last_used_at: Db.set(approximateLastUsedAt)}),
 		);
-		invalidateAuthSessionCache(sessionIdHash);
+		await invalidateAuthSessionCache([sessionIdHash]);
 	}
 
 	async deleteAuthSessions(userId: UserID, sessionIdHashes: Array<Buffer>): Promise<void> {
 		if (sessionIdHashes.length === 0) return;
+		await invalidateAuthSessionCache(sessionIdHashes);
 		let originals: Array<AuthSessionRow> = [];
 		try {
 			originals = await fetchMany<AuthSessionRow>(FETCH_AUTH_SESSIONS_CQL, {
@@ -165,6 +239,7 @@ export class AuthSessionRepository {
 			batch.addPrepared(AuthSessionTombstones.insert(toTombstoneRow(original, deletedAt)));
 		}
 		await batch.execute();
+		await invalidateAuthSessionCache(sessionIdHashes);
 	}
 
 	async deleteAllAuthSessions(userId: UserID): Promise<void> {
@@ -174,10 +249,12 @@ export class AuthSessionRepository {
 			user_id: userId,
 		});
 		if (sessionRefs.length === 0) return;
+		const sessionIdHashes = sessionRefs.map((session) => session.session_id_hash);
+		await invalidateAuthSessionCache(sessionIdHashes);
 		let originals: Array<AuthSessionRow> = [];
 		try {
 			originals = await fetchMany<AuthSessionRow>(FETCH_AUTH_SESSIONS_CQL, {
-				session_id_hashes: sessionRefs.map((s) => s.session_id_hash),
+				session_id_hashes: sessionIdHashes,
 			});
 		} catch (error) {
 			Logger.warn(
@@ -187,12 +264,12 @@ export class AuthSessionRepository {
 		}
 		const deletedAt = new Date();
 		const batch = new BatchBuilder();
-		for (const session of sessionRefs) {
-			batch.addPrepared(AuthSessions.deleteByPk({session_id_hash: session.session_id_hash}));
+		for (const sessionIdHash of sessionIdHashes) {
+			batch.addPrepared(AuthSessions.deleteByPk({session_id_hash: sessionIdHash}));
 			batch.addPrepared(
 				AuthSessionsByUserId.deleteByPk({
 					user_id: userId,
-					session_id_hash: session.session_id_hash,
+					session_id_hash: sessionIdHash,
 				}),
 			);
 		}
@@ -200,6 +277,7 @@ export class AuthSessionRepository {
 			batch.addPrepared(AuthSessionTombstones.insert(toTombstoneRow(original, deletedAt)));
 		}
 		await batch.execute();
+		await invalidateAuthSessionCache(sessionIdHashes);
 	}
 }