refactor: streamline public key retrieval by consolidating procedures into ServerService

packages/core/src/index.ts

@@ -74,10 +74,7 @@ export * from './models/event.model';
 // Procedures
 export { makeJoinEventBuilder } from './procedures/makeJoin';
 export { makeGetMissingEventsProcedure } from './procedures/getMissingEvents';
-export {
-	makeGetPublicKeyFromServerProcedure,
-	getPublicKeyFromRemoteServer,
-} from './procedures/getPublicKeyFromServer';
+export { getPublicKeyFromRemoteServer } from './procedures/getPublicKeyFromServer';
 
 export { createLogger, logger } from './utils/logger';
 

packages/core/src/procedures/getPublicKeyFromServer.ts

@@ -7,38 +7,6 @@ import {
 	verifyJsonSignature,
 } from '../utils/signJson';
 
-export const makeGetPublicKeyFromServerProcedure = (
-	getFromLocal: (origin: string, key: string) => Promise<string | undefined>,
-	getFromOrigin: (
-		origin: string,
-		key: string,
-	) => Promise<{ key: string; validUntil: number }>,
-	store: (
-		origin: string,
-		key: string,
-		value: string,
-		validUntil: number,
-	) => Promise<void>,
-) => {
-	return async (origin: string, key: string) => {
-		const localPublicKey = await getFromLocal(origin, key);
-		if (localPublicKey) {
-			return localPublicKey;
-		}
-
-		const { key: remotePublicKey, validUntil } = await getFromOrigin(
-			origin,
-			key,
-		);
-		if (remotePublicKey) {
-			await store(origin, key, remotePublicKey, validUntil);
-			return remotePublicKey;
-		}
-
-		throw new Error('Public key not found');
-	};
-};
-
 export const getPublicKeyFromRemoteServer = async (
 	domain: string,
 	origin: string,

packages/federation-sdk/src/services/config.service.ts

@@ -1,14 +1,12 @@
 import {
+	SigningKey,
 	createLogger,
 	generateKeyPairsFromString,
-	getKeyPair,
 	toUnpaddedBase64,
 } from '@rocket.chat/federation-core';
 
 import { z } from 'zod';
 
-const CONFIG_FOLDER = process.env.CONFIG_FOLDER || '.';
-
 export interface AppConfig {
 	serverName: string;
 	instanceId: string;
@@ -76,6 +74,7 @@ export const AppConfigSchema = z.object({
 export class ConfigService {
 	private config: AppConfig;
 	private logger = createLogger('ConfigService');
+	private serverKeys: SigningKey[] = [];
 
 	constructor(values: AppConfig) {
 		try {
@@ -116,14 +115,18 @@ export class ConfigService {
 
 	async getSigningKey() {
 		// If config contains a signing key, use it
-		if (this.config.signingKey) {
+		if (!this.config.signingKey) {
+			throw new Error('Signing key is not configured');
+		}
+
+		if (!this.serverKeys.length) {
 			const signingKey = await generateKeyPairsFromString(
 				this.config.signingKey,
 			);
-			return [signingKey];
+			this.serverKeys = [signingKey];
 		}
-		// Otherwise load from file
-		return this.loadSigningKey();
+
+		return this.serverKeys;
 	}
 
 	async getSigningKeyId(): Promise<string> {
@@ -137,20 +140,8 @@ export class ConfigService {
 		return toUnpaddedBase64(signingKeys[0].privateKey);
 	}
 
-	async loadSigningKey() {
-		try {
-			const signingKeyPath = `${CONFIG_FOLDER}/${this.config.serverName}.signing.key`;
-			this.logger.info(`Loading signing key from ${signingKeyPath}`);
-			const keys = await getKeyPair({ signingKeyPath });
-			this.logger.info(
-				`Successfully loaded signing key for server ${this.config.serverName}`,
-			);
-			return keys;
-		} catch (error: unknown) {
-			const errorMessage =
-				error instanceof Error ? error.message : 'Unknown error';
-			this.logger.error(`Failed to load signing key: ${errorMessage}`);
-			throw error;
-		}
+	async getPublicSigningKeyBase64(): Promise<string> {
+		const signingKeys = await this.getSigningKey();
+		return toUnpaddedBase64(signingKeys[0].publicKey);
 	}
 }

packages/federation-sdk/src/services/event-authorization.service.ts

@@ -2,8 +2,6 @@ import {
 	createLogger,
 	extractSignaturesFromHeader,
 	generateId,
-	getPublicKeyFromRemoteServer,
-	makeGetPublicKeyFromServerProcedure,
 	validateAuthorizationHeader,
 } from '@rocket.chat/federation-core';
 import type {
@@ -12,10 +10,10 @@ import type {
 	PersistentEventBase,
 } from '@rocket.chat/federation-room';
 import { singleton } from 'tsyringe';
-import { KeyRepository } from '../repositories/key.repository';
 import { UploadRepository } from '../repositories/upload.repository';
 import { ConfigService } from './config.service';
 import { EventService } from './event.service';
+import { ServerService } from './server.service';
 import { StateService } from './state.service';
 
 @singleton()
@@ -27,7 +25,7 @@ export class EventAuthorizationService {
 		private readonly eventService: EventService,
 		private readonly configService: ConfigService,
 		private readonly uploadRepository: UploadRepository,
-		private readonly keyRepository: KeyRepository,
+		private readonly serverService: ServerService,
 	) {}
 
 	async authorizeEvent(event: Pdu, authEvents: Pdu[]): Promise<boolean> {
@@ -146,20 +144,7 @@ export class EventAuthorizationService {
 				return;
 			}
 
-			// TODO: move makeGetPublicKeyFromServerProcedure procedure to a proper service
-			const getPublicKeyFromServer = makeGetPublicKeyFromServerProcedure(
-				(origin, keyId) =>
-					this.keyRepository.getValidPublicKeyFromLocal(origin, keyId),
-				(origin, key) =>
-					getPublicKeyFromRemoteServer(
-						origin,
-						this.configService.serverName,
-						key,
-					),
-				(origin, keyId, publicKey) =>
-					this.keyRepository.storePublicKey(origin, keyId, publicKey),
-			);
-			const publicKey = await getPublicKeyFromServer(origin, key);
+			const publicKey = await this.serverService.getPublicKey(origin, key);
 			if (!publicKey) {
 				this.logger.warn(`Could not fetch public key for ${origin}:${key}`);
 				return;

packages/federation-sdk/src/services/event.service.ts

@@ -9,10 +9,6 @@ import { isPresenceEDU, isTypingEDU } from '@rocket.chat/federation-core';
 import type { RedactionEvent } from '@rocket.chat/federation-core';
 import { generateId } from '@rocket.chat/federation-core';
 import type { EventStore } from '@rocket.chat/federation-core';
-import {
-	getPublicKeyFromRemoteServer,
-	makeGetPublicKeyFromServerProcedure,
-} from '@rocket.chat/federation-core';
 import { pruneEventDict } from '@rocket.chat/federation-core';
 
 import { checkSignAndHashes } from '@rocket.chat/federation-core';
@@ -35,6 +31,7 @@ import { LockRepository } from '../repositories/lock.repository';
 import { eventSchemas } from '../utils/event-schemas';
 import { ConfigService } from './config.service';
 import { EventEmitterService } from './event-emitter.service';
+import { ServerService } from './server.service';
 import { StateService } from './state.service';
 
 export interface AuthEventParams {
@@ -57,6 +54,7 @@ export class EventService {
 
 		private readonly stagingAreaQueue: StagingAreaQueue,
 		private readonly stateService: StateService,
+		private readonly serverService: ServerService,
 
 		private readonly eventEmitterService: EventEmitterService,
 	) {
@@ -181,8 +179,9 @@ export class EventService {
 					} catch (err) {
 						this.logger.error({
 							msg: 'Event validation failed',
+							origin,
 							event,
-							error: err,
+							err,
 						});
 						continue;
 					}
@@ -251,41 +250,13 @@ export class EventService {
 			throw new Error('M_INVALID_EVENT');
 		}
 
-		const getPublicKeyFromServer = makeGetPublicKeyFromServerProcedure(
-			(origin, keyId) =>
-				this.keyRepository.getValidPublicKeyFromLocal(origin, keyId),
-			(origin, key) =>
-				getPublicKeyFromRemoteServer(
-					origin,
-					this.configService.serverName,
-					key,
-				),
-			(origin, keyId, publicKey) =>
-				this.keyRepository.storePublicKey(origin, keyId, publicKey),
-		);
-
 		if (!event.hashes && !event.signatures) {
 			throw new Error('M_MISSING_SIGNATURES_OR_HASHES');
 		}
 
-		let originToValidateSignatures = origin;
-
-		// If the event does not have a signature for the origin server,
-		// but has a signature for our server, we validate using our server name.
-		// This happens on sendJoin process, where the join event is signed by our server,
-		// but the origin is the remote server since it just returns the event as we sent it.
-		if (
-			!event.signatures[origin] &&
-			event.signatures[this.configService.serverName]
-		) {
-			originToValidateSignatures = this.configService.serverName;
-		}
-
-		await checkSignAndHashes(
-			event,
-			originToValidateSignatures,
-			getPublicKeyFromServer,
-		);
+		await checkSignAndHashes(event, origin, (origin, key) => {
+			return this.serverService.getPublicKey(origin, key);
+		});
 	}
 
 	private async processIncomingEDUs(edus: BaseEDU[]): Promise<void> {

packages/federation-sdk/src/services/server.service.ts

@@ -1,5 +1,6 @@
 import {
 	type SigningKey,
+	getPublicKeyFromRemoteServer,
 	signJson,
 	toUnpaddedBase64,
 } from '@rocket.chat/federation-core';
@@ -30,6 +31,32 @@ export class ServerService {
 		await this.serverRepository.storePublicKey(origin, key, value, validUntil);
 	}
 
+	async getPublicKey(origin: string, key: string): Promise<string> {
+		if (origin === this.configService.serverName) {
+			return this.configService.getPublicSigningKeyBase64();
+		}
+
+		const localPublicKey =
+			await this.serverRepository.getValidPublicKeyFromLocal(origin, key);
+		if (localPublicKey) {
+			return localPublicKey;
+		}
+
+		const { key: remotePublicKey, validUntil } =
+			await getPublicKeyFromRemoteServer(
+				origin,
+				this.configService.serverName,
+				key,
+			);
+
+		if (!remotePublicKey) {
+			throw new Error('Could not get public key from remote server');
+		}
+
+		await this.storePublicKey(origin, key, remotePublicKey, validUntil);
+		return remotePublicKey;
+	}
+
 	async getSignedServerKey() {
 		const signingKeys = await this.configService.getSigningKey();