Perf: cache Telegram MTProto Key across build

Build/build-telegram.ts

@@ -8,7 +8,7 @@ import { fastIpVersion } from 'foxts/fast-ip-version';
 import { appendArrayInPlace } from 'foxts/append-array-in-place';
 import { fetchTelegramBackupEndpoints } from './lib/fetch-telegram-backup-endpoints';
 
-const buildTelegramCIDR = task(require.main === module, __filename)(async (span) => {
+const buildTelegramCIDR = task(require.main === module, 'build-telegram-cidr')(async (span) => {
   const { timestamp, ipcidr, ipcidr6 } = await span.traceChildAsync('get telegram cidr', async (childSpan) => {
     const ipcidr: string[] = [
       // Unused secret Telegram backup CIDR, announced by AS62041
@@ -73,9 +73,13 @@ import path from 'node:path';
 import process from 'node:process';
 
 import { Api as TgApi, TelegramClient as TgClient } from 'telegram';
+import { AuthKey as TgAuthKey } from 'telegram/crypto/AuthKey';
 import { Logger as TgLogger, LogLevel as TgLogLevel } from 'telegram/extensions/Logger';
 import { ConnectionTCPAbridged as TgConnectionTCPAbridged } from 'telegram/network/connection';
 import { MemorySession as TgMemorySession } from 'telegram/sessions';
+import type { Buffer } from 'node:buffer';
+import type { Span } from './trace';
+import { mtprotoAuthKeyStore } from './lib/mtproto-auth-key-store';
 
 import { OUTPUT_INTERNAL_DIR } from './constants/dir';
 import { compareAndWriteFile } from './lib/create-file';
@@ -84,14 +88,40 @@ import {
   normalizeTelegramConfig,
   TELEGRAM_BOOTSTRAP_ENDPOINTS
 } from './lib/mtproto-dc-config';
-import type { MTProtoDCConfig } from './lib/mtproto-dc-config';
+import type { MTProtoDCConfig, MTProtoEndpoint } from './lib/mtproto-dc-config';
 
 const TELEGRAM_API_ID = 2040;
 const OUTPUT_PATH = path.join(OUTPUT_INTERNAL_DIR, 'mtproto-dc-config.json');
 
-async function fetchConfig(host: string, port: number, dcId: number) {
+/**
+ * How long a connect + help.getConfig with a persisted auth key may take before
+ * we give the key up. Normally 0.5-1.5s. It has to be a hard deadline: when a DC
+ * has forgotten the key it answers with transport error -404, and GramJS merely
+ * logs "Broken authorization key" and fires a connection-state event -- the
+ * pending request promise never settles.
+ */
+const PERSISTED_AUTH_KEY_TIMEOUT = 5000;
+
+class PersistedAuthKeyTimeoutError extends Error {
+  constructor(dcId: number, options?: ErrorOptions) {
+    super(`No help.getConfig response within ${PERSISTED_AUTH_KEY_TIMEOUT}ms using the persisted auth key for dc${dcId}`, options);
+    this.name = 'PersistedAuthKeyTimeoutError';
+  }
+}
+
+/**
+ * One connect + help.getConfig against a DC. With a persisted auth key the
+ * connect skips the DH handshake (three round trips); without one, GramJS
+ * negotiates a key and we persist it for the next build.
+ */
+async function fetchConfig(host: string, port: number, dcId: number, persistedAuthKey: Buffer | null) {
   const session = new TgMemorySession();
   session.setDC(dcId, host, port);
+  if (persistedAuthKey) {
+    const authKey = new TgAuthKey();
+    await authKey.setKey(persistedAuthKey);
+    session.setAuthKey(authKey, dcId);
+  }
 
   const client = new TgClient(session, TELEGRAM_API_ID, 'not-used-for-unauthenticated-rpc', {
     appVersion: '1.0',
@@ -111,19 +141,73 @@ async function fetchConfig(host: string, port: number, dcId: number) {
     useWSS: true
   });
 
-  try {
+  const work = async () => {
     const connected = await client.connect();
     if (!connected && !client.connected) {
       throw new Error('MTProto client did not connect');
     }
 
-    return normalizeTelegramConfig(await client.invoke(new TgApi.help.GetConfig()));
+    const config = normalizeTelegramConfig(await client.invoke(new TgApi.help.GetConfig()));
+
+    if (!persistedAuthKey) {
+      const negotiated = client.session.authKey?.getKey();
+      if (negotiated) {
+        await mtprotoAuthKeyStore.save(dcId, negotiated);
+      }
+    }
+
+    return config;
+  };
+
+  let timer: NodeJS.Timeout | null = null;
+  try {
+    if (!persistedAuthKey) {
+      return await work();
+    }
+    return await Promise.race([
+      work(),
+      new Promise<never>((_resolve, reject) => {
+        timer = setTimeout(reject, PERSISTED_AUTH_KEY_TIMEOUT, new PersistedAuthKeyTimeoutError(dcId));
+      })
+    ]);
   } finally {
-    await client.disconnect();
+    if (timer) {
+      clearTimeout(timer);
+    }
+    // destroy, not disconnect: a sender stuck on a rejected key must not keep the
+    // process alive after we have moved on to a fresh handshake
+    await client.destroy();
   }
 }
 
-async function fetchConfigFromBootstrapEndpoints() {
+async function fetchConfigFromEndpoint(span: Span, endpoint: MTProtoEndpoint) {
+  const label = `${endpoint.ip}:${endpoint.port}`;
+
+  const persistedAuthKey = await mtprotoAuthKeyStore.load(endpoint.dcId);
+  if (persistedAuthKey) {
+    try {
+      return await span.traceChildAsync(
+        `help.getConfig via ${label} (persisted auth key)`,
+        () => fetchConfig(endpoint.ip, endpoint.port, endpoint.dcId, persistedAuthKey),
+        SpanCategory.Network
+      );
+    } catch (error) {
+      // Most likely the DC no longer knows this key, which surfaces as the
+      // timeout above (see PERSISTED_AUTH_KEY_TIMEOUT). Whatever it was, a key
+      // that fails once is not worth a second try: fall through to a fresh handshake.
+      console.warn(`[telegram mtproto config] ${label} failed with the persisted auth key for dc${endpoint.dcId}, renegotiating`, error);
+      await mtprotoAuthKeyStore.drop(endpoint.dcId);
+    }
+  }
+
+  return span.traceChildAsync(
+    `help.getConfig via ${label} (fresh handshake)`,
+    () => fetchConfig(endpoint.ip, endpoint.port, endpoint.dcId, null),
+    SpanCategory.Network
+  );
+}
+
+async function fetchConfigFromBootstrapEndpoints(span: Span) {
   let lastError: unknown;
 
   for (let i = 0, len = TELEGRAM_BOOTSTRAP_ENDPOINTS.length; i < len; i++) {
@@ -132,7 +216,7 @@ async function fetchConfigFromBootstrapEndpoints() {
     try {
       // Bootstrap order is significant, and one successful response ends the loop.
       // eslint-disable-next-line no-await-in-loop -- Bootstrap endpoints must be attempted in order.
-      return await fetchConfig(endpoint.ip, endpoint.port, endpoint.dcId);
+      return await fetchConfigFromEndpoint(span, endpoint);
     } catch (error) {
       lastError = error;
       console.error(`[telegram mtproto config] ${endpoint.ip}:${endpoint.port} failed`, error);
@@ -145,7 +229,7 @@ async function fetchConfigFromBootstrapEndpoints() {
   );
 }
 
-export const buildMTProtoDCConfig = task(require.main === module, __filename)(async (span) => {
+export const buildMTProtoDCConfig = task(require.main === module, 'build-mtproto-dc-config')(async (span) => {
   const config = await span.traceChildAsync(
     'fetch help.getConfig',
     fetchConfigFromBootstrapEndpoints,