Use DoH for query fallback

ts/test-node/util/dns_test.ts

@@ -0,0 +1,67 @@
+// Copyright 2023 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+import { assert } from 'chai';
+import * as sinon from 'sinon';
+
+import { DNSCache } from '../../util/dns';
+import { SECOND } from '../../util/durations';
+
+const NOW = 1680726906000;
+
+describe('dns/DNSCache', () => {
+  let sandbox: sinon.SinonSandbox;
+  let cache: DNSCache;
+  beforeEach(() => {
+    sandbox = sinon.createSandbox();
+    cache = new DNSCache();
+  });
+
+  afterEach(() => {
+    sandbox.restore();
+  });
+
+  it('should cache records and pick a random one', () => {
+    sandbox.useFakeTimers({
+      now: NOW,
+    });
+
+    const result = cache.setAndPick('signal.org', 4, [
+      {
+        data: '10.0.0.1',
+        expiresAt: NOW + SECOND,
+      },
+      {
+        data: '10.0.0.2',
+        expiresAt: NOW + SECOND,
+      },
+    ]);
+
+    assert.oneOf(result, ['10.0.0.1', '10.0.0.2']);
+  });
+
+  it('should invalidate cache after expiration', () => {
+    const clock = sandbox.useFakeTimers({
+      now: NOW,
+    });
+
+    cache.setAndPick('signal.org', 4, [
+      {
+        data: '10.0.0.1',
+        expiresAt: NOW + SECOND,
+      },
+      {
+        data: '10.0.0.2',
+        expiresAt: NOW + 2 * SECOND,
+      },
+    ]);
+
+    assert.oneOf(cache.get('signal.org', 4), ['10.0.0.1', '10.0.0.2']);
+
+    clock.tick(SECOND);
+    assert.strictEqual(cache.get('signal.org', 4), '10.0.0.2');
+
+    clock.tick(SECOND);
+    assert.strictEqual(cache.get('signal.org', 4), undefined);
+  });
+});

ts/util/dns.ts

@@ -1,21 +1,163 @@
 // Copyright 2023 Signal Messenger, LLC
 // SPDX-License-Identifier: AGPL-3.0-only
 
-import {
-  lookup as nativeLookup,
-  resolve4,
-  resolve6,
-  getServers,
-  setServers,
-} from 'dns';
+import { lookup as nativeLookup } from 'dns';
 import type { LookupOneOptions } from 'dns';
+import fetch from 'node-fetch';
+import { z } from 'zod';
 
 import * as log from '../logging/log';
 import * as Errors from '../types/errors';
 import { strictAssert } from './assert';
+import { SECOND } from './durations';
 
-const ORIGINAL_SERVERS = getServers();
-const FALLBACK_SERVERS = ['1.1.1.1'];
+const HOST_ALLOWLIST = new Set([
+  // Production
+  'chat.signal.org',
+  'storage.signal.org',
+  'cdsi.signal.org',
+  'cdn.signal.org',
+  'cdn2.signal.org',
+  'create.signal.art',
+
+  // Staging
+  'chat.staging.signal.org',
+  'storage-staging.signal.org',
+  'cdsi.staging.signal.org',
+  'cdn-staging.signal.org',
+  'cdn2-staging.signal.org',
+  'create.staging.signal.art',
+
+  // Common
+  'updates2.signal.org',
+  'sfu.voip.signal.org',
+]);
+
+const dohResponseSchema = z.object({
+  Status: z.number(),
+  Answer: z.array(
+    z.object({
+      data: z.string(),
+      TTL: z.number(),
+    })
+  ),
+  Comment: z.string().optional(),
+});
+
+type CacheEntry = Readonly<{
+  data: string;
+  expiresAt: number;
+}>;
+
+export class DNSCache {
+  private readonly ipv4 = new Map<string, Array<CacheEntry>>();
+  private readonly ipv6 = new Map<string, Array<CacheEntry>>();
+
+  public get(hostname: string, family: 4 | 6): string | undefined {
+    const map = this.getMap(family);
+
+    const entries = map.get(hostname);
+    if (!entries) {
+      return undefined;
+    }
+
+    // Cleanup old records
+    this.cleanup(entries);
+    if (entries.length === 0) {
+      map.delete(hostname);
+      return undefined;
+    }
+
+    // Pick a random record
+    return this.pick(entries);
+  }
+
+  public setAndPick(
+    hostname: string,
+    family: 4 | 6,
+    entries: Array<CacheEntry>
+  ): string {
+    strictAssert(entries.length !== 0, 'should have at least on entry');
+
+    const map = this.getMap(family);
+
+    // Just overwrite the entries - we shouldn't get here unless it was a cache
+    // miss.
+    map.set(hostname, entries);
+
+    return this.pick(entries);
+  }
+
+  // Private
+
+  private getMap(family: 4 | 6): Map<string, Array<CacheEntry>> {
+    return family === 4 ? this.ipv4 : this.ipv6;
+  }
+
+  private pick(entries: Array<CacheEntry>): string {
+    const index = Math.floor(Math.random() * entries.length);
+    return entries[index].data;
+  }
+
+  private cleanup(entries: Array<CacheEntry>): void {
+    const now = Date.now();
+    for (let i = entries.length - 1; i >= 0; i -= 1) {
+      const { expiresAt } = entries[i];
+      if (expiresAt <= now) {
+        entries.splice(i, 1);
+      }
+    }
+  }
+}
+
+const cache = new DNSCache();
+
+export async function doh(hostname: string, family: 4 | 6): Promise<string> {
+  const cached = cache.get(hostname, family);
+  if (cached !== undefined) {
+    log.info(`dns/doh: using cached value for ${hostname}/IPv${family}`);
+    return cached;
+  }
+
+  const url = new URL('https://1.1.1.1/dns-query');
+  url.searchParams.append('name', hostname);
+  url.searchParams.append('type', family === 4 ? 'A' : 'AAAA');
+  const res = await fetch(url.toString(), {
+    headers: {
+      accept: 'application/dns-json',
+      'user-agent': 'Electron',
+    },
+  });
+
+  if (!res.ok) {
+    throw new Error(
+      `DoH request for ${hostname} failed with http status: ${res.status}`
+    );
+  }
+
+  const {
+    Status: status,
+    Answer: answer,
+    Comment: comment,
+  } = dohResponseSchema.parse(await res.json());
+
+  if (status !== 0) {
+    throw new Error(`DoH request for ${hostname} failed: ${status}/${comment}`);
+  }
+
+  if (answer.length === 0) {
+    throw new Error(`DoH request for ${hostname} failed: empty answer`);
+  }
+
+  const now = Date.now();
+  return cache.setAndPick(
+    hostname,
+    family,
+    answer.map(({ data, TTL }) => {
+      return { data, expiresAt: now + TTL * SECOND };
+    })
+  );
+}
 
 export function lookupWithFallback(
   hostname: string,
@@ -31,43 +173,31 @@ export function lookupWithFallback(
   strictAssert(Boolean(opts.all) !== true, 'options.all is not supported');
   strictAssert(typeof callback === 'function', 'missing callback');
 
-  nativeLookup(hostname, opts, (err, ...nativeArgs) => {
+  nativeLookup(hostname, opts, async (err, ...nativeArgs) => {
     if (!err) {
       return callback(err, ...nativeArgs);
     }
 
+    if (!HOST_ALLOWLIST.has(hostname)) {
+      log.error(
+        `dns/lookup: failed for ${hostname}, ` +
+          `err: ${Errors.toLogFormat(err)}. not retrying`
+      );
+      return callback(err, ...nativeArgs);
+    }
+
     const family = opts.family === 6 ? 6 : 4;
 
     log.error(
-      `lookup: failed for ${hostname}, error: ${Errors.toLogFormat(err)}. ` +
-        `Retrying with c-ares (IPv${family})`
+      `dns/lookup: failed for ${hostname}, err: ${Errors.toLogFormat(err)}. ` +
+        `Retrying with DoH (IPv${family})`
     );
-    const onRecords = (
-      fallbackErr: NodeJS.ErrnoException | null,
-      records: Array<string>
-    ): void => {
-      setServers(ORIGINAL_SERVERS);
-      if (fallbackErr) {
-        return callback(fallbackErr, '', 0);
-      }
-
-      if (!Array.isArray(records) || records.length === 0) {
-        return callback(
-          new Error(`No DNS records returned for: ${hostname}`),
-          '',
-          0
-        );
-      }
-
-      const index = Math.floor(Math.random() * records.length);
-      callback(null, records[index], family);
-    };
 
-    setServers(FALLBACK_SERVERS);
-    if (family === 4) {
-      resolve4(hostname, onRecords);
-    } else {
-      resolve6(hostname, onRecords);
+    try {
+      const answer = await doh(hostname, family);
+      callback(null, answer, family);
+    } catch (fallbackErr) {
+      callback(fallbackErr, '', 0);
     }
   });
 }