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,

Build/lib/mtproto-auth-key-store.test.ts

@@ -0,0 +1,52 @@
+import { describe, it } from 'mocha';
+import { expect } from 'earl';
+import fs from 'node:fs';
+import os from 'node:os';
+import path from 'node:path';
+import { randomBytes } from 'node:crypto';
+
+import { createMTProtoAuthKeyStore } from './mtproto-auth-key-store';
+
+function tmpStorePath() {
+  return path.join(fs.mkdtempSync(path.join(os.tmpdir(), 'mtproto-auth-key-')), 'keys');
+}
+
+describe('mtproto auth key store', () => {
+  it('round-trips a key per DC and survives reopening the directory', async () => {
+    const dir = tmpStorePath();
+    const dc1 = randomBytes(256);
+    const dc2 = randomBytes(256);
+
+    const store = createMTProtoAuthKeyStore(dir);
+    expect(await store.load(1)).toEqual(null);
+    await store.save(1, dc1);
+    await store.save(2, dc2);
+    expect((await store.load(1))?.equals(dc1)).toEqual(true);
+    expect((await store.load(2))?.equals(dc2)).toEqual(true);
+
+    // a fresh handle on the same directory is what the next build sees
+    const reopened = createMTProtoAuthKeyStore(dir);
+    expect((await reopened.load(1))?.equals(dc1)).toEqual(true);
+  });
+
+  it('replaces on save and forgets on drop', async () => {
+    const store = createMTProtoAuthKeyStore(tmpStorePath());
+    const first = randomBytes(256);
+    const second = randomBytes(256);
+
+    await store.save(4, first);
+    await store.save(4, second);
+    expect((await store.load(4))?.equals(second)).toEqual(true);
+
+    await store.drop(4);
+    expect(await store.load(4)).toEqual(null);
+    // dropping an absent key is a no-op
+    await store.drop(4);
+  });
+
+  it('treats a key of the wrong size as absent', async () => {
+    const store = createMTProtoAuthKeyStore(tmpStorePath());
+    await store.save(5, randomBytes(16));
+    expect(await store.load(5)).toEqual(null);
+  });
+});

Build/lib/mtproto-auth-key-store.ts

@@ -0,0 +1,48 @@
+import path from 'node:path';
+import cacache from 'cacache';
+import type { Buffer } from 'node:buffer';
+import { CACHE_DIR } from '../constants/dir';
+
+export interface MTProtoAuthKeyStore {
+  load(dcId: number): Promise<Buffer | null>,
+  save(dcId: number, key: Buffer): Promise<void>,
+  drop(dcId: number): Promise<void>
+}
+
+/** a permanent auth key is exactly 2048 bits */
+const AUTH_KEY_BYTES = 256;
+
+const VERIFY_INTERVAL = 30 * 24 * 60 * 60 * 1000;
+async function verifyIfStale(cachePath: string) {
+  const lastRun = await cacache.verify.lastRun(cachePath).catch(() => null);
+  if (lastRun === null || Date.now() - lastRun.getTime() > VERIFY_INTERVAL) {
+    await cacache.verify(cachePath);
+  }
+}
+
+const entryKey = (dcId: number) => `dc${dcId}`;
+
+export function createMTProtoAuthKeyStore(cachePath: string): MTProtoAuthKeyStore {
+  return {
+    async load(dcId) {
+      let data: Buffer;
+      try {
+        ({ data } = await cacache.get(cachePath, entryKey(dcId)));
+      } catch {
+        // ENOENT (never negotiated) or EINTEGRITY (corrupt), both mean "no key"
+        return null;
+      }
+      return data.length === AUTH_KEY_BYTES ? data : null;
+    },
+    async save(dcId, key) {
+      await cacache.put(cachePath, entryKey(dcId), key);
+      await verifyIfStale(cachePath);
+    },
+    async drop(dcId) {
+      await cacache.rm.entry(cachePath, entryKey(dcId));
+      await verifyIfStale(cachePath);
+    }
+  };
+}
+
+export const mtprotoAuthKeyStore = createMTProtoAuthKeyStore(path.join(CACHE_DIR, 'telegram-mtproto-auth-keys'));

Build/trace/index.ts

@@ -65,22 +65,30 @@ export function makeSpan(rawSpan: RawSpan): Span {
       return span;
     },
     traceChild,
+    // Spans are stopped in `finally` so a throwing function is still measured
+    // (and the tree does not end up with an "(unfinished)" hole where it failed).
     traceSyncFn<T>(fn: (span: Span) => T) {
       traceResult.sync = true;
-      const res = fn(span);
-      span.stop();
-      return res;
+      try {
+        return fn(span);
+      } finally {
+        span.stop();
+      }
     },
     async traceAsyncFn<T>(fn: (span: Span) => T | Promise<T>): Promise<T> {
-      const res = await fn(span);
-      span.stop();
-      return res;
+      try {
+        return await fn(span);
+      } finally {
+        span.stop();
+      }
     },
     traceResult,
     async tracePromise<T>(promise: Promise<T>): Promise<T> {
-      const res = await promise;
-      span.stop();
-      return res;
+      try {
+        return await promise;
+      } finally {
+        span.stop();
+      }
     },
     traceChildSync: <T>(name: string, fn: (span: Span) => T, category?: SpanCategory): T => traceChild(name, category).traceSyncFn(fn),
     traceChildAsync: <T>(name: string, fn: (span: Span) => T | Promise<T>, category?: SpanCategory): Promise<T> => traceChild(name, category).traceAsyncFn(fn),