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

fluxer_api/pkgs/cache/src/ICacheService.ts

@@ -10,6 +10,8 @@ interface CacheMSetEntry<T> {
 
 export type CacheLookupResult<T> = {hit: true; value: T} | {hit: false};
 
+type CacheTtlSeconds<T> = number | ((value: T) => number);
+
 export abstract class ICacheService {
 	private readonly inflightValues = new Map<string, Promise<unknown>>();
 
@@ -56,7 +58,7 @@ export abstract class ICacheService {
 		return entry.hit ? entry.value : null;
 	}
 
-	async getOrSet<T>(key: string, valueFactory: () => Promise<T>, ttlSeconds?: number): Promise<T> {
+	async getOrSet<T>(key: string, valueFactory: () => Promise<T>, ttlSeconds?: CacheTtlSeconds<T>): Promise<T> {
 		const existing = await this.getEntry<T>(key);
 		if (existing.hit) {
 			return existing.value;
@@ -75,9 +77,13 @@ export abstract class ICacheService {
 		return await pending;
 	}
 
-	private async produceAndStore<T>(key: string, valueFactory: () => Promise<T>, ttlSeconds?: number): Promise<T> {
+	private async produceAndStore<T>(
+		key: string,
+		valueFactory: () => Promise<T>,
+		ttlSeconds?: CacheTtlSeconds<T>,
+	): Promise<T> {
 		const value = await valueFactory();
-		await this.set(key, value, ttlSeconds);
+		await this.set(key, value, typeof ttlSeconds === 'function' ? ttlSeconds(value) : ttlSeconds);
 		return value;
 	}
 }

fluxer_api/pkgs/cache/src/__tests__/CacheGetOrSet.test.ts

@@ -5,19 +5,25 @@ import {KVCacheProvider} from '@pkgs/cache/src/providers/KVCacheProvider';
 import type {IKVProvider} from '@pkgs/kv_client/src/IKVProvider';
 import {describe, expect, it, vi} from 'vitest';
 
-function createKVCacheProvider(): {provider: KVCacheProvider; store: Map<string, string>} {
+function createKVCacheProvider(): {
+	provider: KVCacheProvider;
+	store: Map<string, string>;
+	ttls: Array<[string, number]>;
+} {
 	const store = new Map<string, string>();
+	const ttls: Array<[string, number]> = [];
 	const client = {
 		get: async (key: string) => store.get(key) ?? null,
 		set: async (key: string, value: string) => {
 			store.set(key, value);
 			return 'OK';
 		},
-		setex: async (key: string, _ttlSeconds: number, value: string) => {
+		setex: async (key: string, ttlSeconds: number, value: string) => {
+			ttls.push([key, ttlSeconds]);
 			store.set(key, value);
 		},
 	} as unknown as IKVProvider;
-	return {provider: new KVCacheProvider({client}), store};
+	return {provider: new KVCacheProvider({client}), store, ttls};
 }
 
 function delay(ms: number): Promise<void> {
@@ -76,6 +82,18 @@ describe('ICacheService.getOrSet', () => {
 		expect(factory).toHaveBeenCalledTimes(1);
 	});
 
+	it('resolves the ttl from the produced value', async () => {
+		const {provider, store, ttls} = createKVCacheProvider();
+		const resolver = (value: number | null) => (value === null ? 5 : 30);
+		await expect(provider.getOrSet<number | null>('present', async () => 1, resolver)).resolves.toBe(1);
+		await expect(provider.getOrSet<number | null>('absent', async () => null, resolver)).resolves.toBeNull();
+		expect(ttls).toEqual([
+			['present', 30],
+			['absent', 5],
+		]);
+		expect(store.get('absent')).toBe('null');
+	});
+
 	it('rejects every waiter and retries on the next call when the factory fails', async () => {
 		const cache = new InMemoryProvider();
 		const failing = vi.fn(async () => {

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,19 +264,20 @@ 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,
 				}),
 			);
 		}
 		for (const original of originals) {
 			batch.addPrepared(AuthSessionTombstones.insert(toTombstoneRow(original, deletedAt)));
 		}
 		await batch.execute();
+		await invalidateAuthSessionCache(sessionIdHashes);
 	}
 }