don't cache stale actor keys

packages/backend/src/core/activitypub/ApDbResolverService.ts

@@ -14,10 +14,10 @@ import { UtilityService } from '@/core/UtilityService.js';
 import type { MiNote } from '@/models/Note.js';
 import { bindThis } from '@/decorators.js';
 import type { MiLocalUser, MiRemoteUser } from '@/models/User.js';
+import { ApLoggerService } from '@/core/activitypub/ApLoggerService.js';
 import { getApId } from './type.js';
 import { ApPersonService } from './models/ApPersonService.js';
 import type { IObject } from './type.js';
-import { ApLoggerService } from '@/core/activitypub/ApLoggerService.js';
 
 export type UriParseResult = {
 	/** wether the URI was generated by us */
@@ -37,9 +37,6 @@ export type UriParseResult = {
 
 @Injectable()
 export class ApDbResolverService implements OnApplicationShutdown {
-	private publicKeyCache: MemoryKVCache<MiUserPublickey | null>;
-	private publicKeyByUserIdCache: MemoryKVCache<MiUserPublickey | null>;
-
 	constructor(
 		@Inject(DI.config)
 		private config: Config,
@@ -58,8 +55,7 @@ export class ApDbResolverService implements OnApplicationShutdown {
 		private apLoggerService: ApLoggerService,
 		private utilityService: UtilityService,
 	) {
-		this.publicKeyCache = new MemoryKVCache<MiUserPublickey | null>(1000 * 60 * 60 * 12); // 12h
-		this.publicKeyByUserIdCache = new MemoryKVCache<MiUserPublickey | null>(1000 * 60 * 60 * 12); // 12h
+		// Caches moved to ApPersonService to avoid circular dependency
 	}
 
 	@bindThis
@@ -131,15 +127,7 @@ export class ApDbResolverService implements OnApplicationShutdown {
 		user: MiRemoteUser;
 		key: MiUserPublickey;
 	} | null> {
-		const key = await this.publicKeyCache.fetch(keyId, async () => {
-			const key = await this.userPublickeysRepository.findOneBy({
-				keyId,
-			});
-
-			if (key == null) return null;
-
-			return key;
-		}, key => key != null);
+		const key = await this.apPersonService.findPublicKeyByKeyId(keyId);
 
 		if (key == null) return null;
 
@@ -164,11 +152,7 @@ export class ApDbResolverService implements OnApplicationShutdown {
 		const user = await this.apPersonService.resolvePerson(uri) as MiRemoteUser;
 		if (user.isDeleted) return null;
 
-		const key = await this.publicKeyByUserIdCache.fetch(
-			user.id,
-			() => this.userPublickeysRepository.findOneBy({ userId: user.id }),
-			v => v != null,
-		);
+		const key = await this.apPersonService.findPublicKeyByUserId(user.id);
 
 		return {
 			user,
@@ -184,21 +168,19 @@ export class ApDbResolverService implements OnApplicationShutdown {
 		this.apLoggerService.logger.debug('Re-fetching public key for user', { userId: user.id, uri: user.uri });
 		await this.apPersonService.updatePerson(user.uri);
 
-		const key = await this.userPublickeysRepository.findOneBy({ userId: user.id });
-		this.publicKeyByUserIdCache.set(user.id, key);
+		const key = await this.apPersonService.findPublicKeyByUserId(user.id);
 
 		if (key) {
 			this.apLoggerService.logger.info('Re-fetched public key for user', { userId: user.id, uri: user.uri });
 		} else {
 			this.apLoggerService.logger.warn('Failed to re-fetch key for user', { userId: user.id, uri: user.uri });
 		}
+
 		return key;
 	}
 
 	@bindThis
 	public dispose(): void {
-		this.publicKeyCache.dispose();
-		this.publicKeyByUserIdCache.dispose();
 	}
 
 	@bindThis

packages/backend/src/core/activitypub/models/ApPersonService.ts

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: AGPL-3.0-only
  */
 
-import { Inject, Injectable } from '@nestjs/common';
+import { Inject, Injectable, OnApplicationShutdown } from '@nestjs/common';
 import promiseLimit from 'promise-limit';
 import { DataSource } from 'typeorm';
 import { ModuleRef } from '@nestjs/core';
@@ -40,6 +40,7 @@ import { RoleService } from '@/core/RoleService.js';
 import { DriveFileEntityService } from '@/core/entities/DriveFileEntityService.js';
 import type { AccountMoveService } from '@/core/AccountMoveService.js';
 import { ApUtilityService } from '@/core/activitypub/ApUtilityService.js';
+import { MemoryKVCache } from '@/misc/cache.js';
 import { getApId, getApType, isActor, isCollection, isCollectionOrOrderedCollection, isPropertyValue } from '../type.js';
 import { extractApHashtags } from './tag.js';
 import type { OnModuleInit } from '@nestjs/common';
@@ -57,7 +58,11 @@ const summaryLength = 2048;
 type Field = Record<'name' | 'value', string>;
 
 @Injectable()
-export class ApPersonService implements OnModuleInit {
+export class ApPersonService implements OnModuleInit, OnApplicationShutdown {
+	// Moved from ApDbResolverService
+	private readonly publicKeyByKeyIdCache = new MemoryKVCache<MiUserPublickey | null>(1000 * 60 * 60 * 12); // 12h
+	private readonly publicKeyByUserIdCache = new MemoryKVCache<MiUserPublickey | null>(1000 * 60 * 60 * 12); // 12h
+
 	private utilityService: UtilityService;
 	private userEntityService: UserEntityService;
 	private driveFileEntityService: DriveFileEntityService;
@@ -132,6 +137,10 @@ export class ApPersonService implements OnModuleInit {
 		this.logger = this.apLoggerService.logger;
 	}
 
+	onApplicationShutdown(): void {
+		this.dispose();
+	}
+
 	/**
 	 * Validate and convert to actor object
 	 * @param x Fetched object
@@ -355,6 +364,7 @@ export class ApPersonService implements OnModuleInit {
 
 		// Create user
 		let user: MiRemoteUser | null = null;
+		let publicKey: MiUserPublickey | null = null;
 
 		//#region カスタム絵文字取得
 		const emojis = await this.apNoteService.extractEmojis(person.tag ?? [], host)
@@ -435,7 +445,7 @@ export class ApPersonService implements OnModuleInit {
 				}));
 
 				if (person.publicKey) {
-					await transactionalEntityManager.save(new MiUserPublickey({
+					publicKey = await transactionalEntityManager.save(new MiUserPublickey({
 						userId: user.id,
 						keyId: person.publicKey.id,
 						keyPem: person.publicKey.publicKeyPem.trim(),
@@ -450,6 +460,7 @@ export class ApPersonService implements OnModuleInit {
 				if (u == null) throw new UnrecoverableError(`already registered a user with conflicting data: ${uri}`);
 
 				user = u as MiRemoteUser;
+				publicKey = await this.userPublickeysRepository.findOneBy({ userId: user.id });
 			} else {
 				this.logger.error(e instanceof Error ? e : new Error(e as string));
 				throw e;
@@ -461,6 +472,11 @@ export class ApPersonService implements OnModuleInit {
 		// Register to the cache
 		this.cacheService.uriPersonCache.set(user.uri, user);
 
+		// Register public key to the cache.
+		// Value may be null, which indicates that the user has no defined key. (optimization)
+		this.publicKeyByUserIdCache.set(user.id, publicKey);
+		if (publicKey) this.publicKeyByKeyIdCache.set(publicKey.keyId, publicKey);
+
 		// Register host
 		if (this.meta.enableStatsForFederatedInstances) {
 			this.federatedInstanceService.fetchOrRegister(host).then(i => {
@@ -611,10 +627,27 @@ export class ApPersonService implements OnModuleInit {
 		await this.usersRepository.update(exist.id, updates);
 
 		if (person.publicKey) {
-			await this.userPublickeysRepository.update({ userId: exist.id }, {
+			const publicKey = new MiUserPublickey({
+				userId: exist.id,
 				keyId: person.publicKey.id,
 				keyPem: person.publicKey.publicKeyPem,
 			});
+
+			// Create or update key
+			await this.userPublickeysRepository.save(publicKey);
+
+			this.publicKeyByKeyIdCache.set(person.publicKey.id, publicKey);
+			this.publicKeyByUserIdCache.set(exist.id, publicKey);
+		} else {
+			const existingPublicKey = await this.userPublickeysRepository.findOneBy({ userId: exist.id });
+			if (existingPublicKey) {
+				// Delete key
+				await this.userPublickeysRepository.delete({ userId: exist.id });
+				this.publicKeyByKeyIdCache.delete(existingPublicKey.keyId);
+			}
+
+			// Null indicates that the user has no key. (optimization)
+			this.publicKeyByUserIdCache.set(exist.id, null);
 		}
 
 		let _description: string | null = null;
@@ -825,4 +858,38 @@ export class ApPersonService implements OnModuleInit {
 
 		return false;
 	}
+
+	@bindThis
+	public async findPublicKeyByUserId(userId: string): Promise<MiUserPublickey | null> {
+		const publicKey = this.publicKeyByUserIdCache.get(userId) ?? await this.userPublickeysRepository.findOneBy({ userId });
+
+		// This can technically keep a key cached "forever" if it's used enough, but that's ok.
+		// We can never have stale data because the publicKey caches are coherent. (cache updates whenever data changes)
+		if (publicKey) {
+			this.publicKeyByUserIdCache.set(publicKey.userId, publicKey);
+			this.publicKeyByKeyIdCache.set(publicKey.keyId, publicKey);
+		}
+
+		return publicKey;
+	}
+
+	@bindThis
+	public async findPublicKeyByKeyId(keyId: string): Promise<MiUserPublickey | null> {
+		const publicKey = this.publicKeyByKeyIdCache.get(keyId) ?? await this.userPublickeysRepository.findOneBy({ keyId });
+
+		// This can technically keep a key cached "forever" if it's used enough, but that's ok.
+		// We can never have stale data because the publicKey caches are coherent. (cache updates whenever data changes)
+		if (publicKey) {
+			this.publicKeyByUserIdCache.set(publicKey.userId, publicKey);
+			this.publicKeyByKeyIdCache.set(publicKey.keyId, publicKey);
+		}
+
+		return publicKey;
+	}
+
+	@bindThis
+	public dispose(): void {
+		this.publicKeyByUserIdCache.dispose();
+		this.publicKeyByKeyIdCache.dispose();
+	}
 }