From 6278da6aec2fde629459e3f256dcb21575da108c Mon Sep 17 00:00:00 2001
From: George Kampitakis <george.kampitakis@datadoghq.com>
Date: Wed, 19 Jun 2024 09:43:35 +0100
Subject: [PATCH] support selecting resolver via env

---
 packages/grpc-js/src/resolver-dns.ts | 58 ++++++++++++++++++----------
 1 file changed, 38 insertions(+), 20 deletions(-)

diff --git a/packages/grpc-js/src/resolver-dns.ts b/packages/grpc-js/src/resolver-dns.ts
index 5932396c..d049c959 100644
--- a/packages/grpc-js/src/resolver-dns.ts
+++ b/packages/grpc-js/src/resolver-dns.ts
@@ -20,7 +20,6 @@ import {
   registerResolver,
   registerDefaultScheme,
 } from './resolver';
-import { AnyRecord } from 'dns';
 import * as dns from 'dns/promises';
 import { extractAndSelectServiceConfig, ServiceConfig } from './service-config';
 import { Status } from './constants';
@@ -45,7 +44,8 @@ function trace(text: string): void {
  */
 export const DEFAULT_PORT = 443;
 
-const DEFAULT_MIN_TIME_BETWEEN_RESOLUTIONS_MS = 30_000;
+const DEFAULT_MIN_TIME_BETWEEN_RESOLUTIONS_MS = 30_000,
+  DNS_RESOLUTION_ENV = 'GRPC_DNS_RESOLVER_TODO';
 
 /**
  * Resolver implementation that handles DNS names and IP addresses.
@@ -60,7 +60,7 @@ class DnsResolver implements Resolver {
    * Failures are handled by the backoff timer.
    */
   private readonly minTimeBetweenResolutionsMs: number;
-  private pendingLookupPromise: Promise<AnyRecord[]> | null = null;
+  private pendingLookupPromise: Promise<TcpSubchannelAddress[]> | null = null;
   private pendingTxtPromise: Promise<string[][]> | null = null;
   private latestLookupResult: Endpoint[] | null = null;
   private latestServiceConfig: ServiceConfig | null = null;
@@ -73,7 +73,7 @@ class DnsResolver implements Resolver {
   private isNextResolutionTimerRunning = false;
   private isServiceConfigEnabled = true;
   private returnedIpResult = false;
-  private resolver = new dns.Resolver();
+  private independentResolver = new dns.Resolver();
 
   constructor(
     private target: GrpcUri,
@@ -82,7 +82,7 @@ class DnsResolver implements Resolver {
   ) {
     trace('Resolver constructed for target ' + uriToString(target));
     if (target.authority) {
-      this.resolver.setServers([target.authority]);
+      this.independentResolver.setServers([target.authority]);
     }
     const hostPort = splitHostPort(target.path);
     if (hostPort === null) {
@@ -187,11 +187,7 @@ class DnsResolver implements Resolver {
        * revert to an effectively blank one. */
       this.latestLookupResult = null;
       const hostname: string = this.dnsHostname;
-      /* We lookup both address families here and then split them up later
-       * because when looking up a single family, dns.lookup outputs an error
-       * if the name exists but there are no records for that family, and that
-       * error is indistinguishable from other kinds of errors */
-      this.pendingLookupPromise = this.resolver.resolveAny(hostname);
+      this.pendingLookupPromise = this.lookup(hostname);
       this.pendingLookupPromise.then(
         addressList => {
           if (this.pendingLookupPromise === null) {
@@ -200,19 +196,12 @@ class DnsResolver implements Resolver {
           this.pendingLookupPromise = null;
           this.backoff.reset();
           this.backoff.stop();
-          const subchannelAddresses: TcpSubchannelAddress[] = addressList
-            .filter(addr => {
-              addr.type === 'A' || addr.type === 'AAAA';
-            })
-            .map(addr => ({ host: addr.address, port: +this.port! }));
-          this.latestLookupResult = subchannelAddresses.map(address => ({
+          this.latestLookupResult = addressList.map(address => ({
             addresses: [address],
           }));
           const allAddressesString: string =
             '[' +
-            subchannelAddresses
-              .map(addr => addr.host + ':' + addr.port)
-              .join(',') +
+            addressList.map(addr => addr.host + ':' + addr.port).join(',') +
             ']';
           trace(
             'Resolved addresses for target ' +
@@ -257,7 +246,7 @@ class DnsResolver implements Resolver {
         /* We handle the TXT query promise differently than the others because
          * the name resolution attempt as a whole is a success even if the TXT
          * lookup fails */
-        this.pendingTxtPromise = this.resolver.resolveTxt(hostname);
+        this.pendingTxtPromise = this.resolveTxt(hostname);
         this.pendingTxtPromise.then(
           txtRecord => {
             if (this.pendingTxtPromise === null) {
@@ -306,6 +295,35 @@ class DnsResolver implements Resolver {
     }
   }
 
+  private async lookup(hostname: string): Promise<TcpSubchannelAddress[]> {
+    if (process.env[DNS_RESOLUTION_ENV] === 'true') {
+      const records = await this.independentResolver.resolveAny(hostname);
+      const addressList = records.filter(addr => {
+        addr.type === 'A' || addr.type === 'AAAA';
+      }) as unknown as { address: string }[];
+
+      return addressList.map(addr => ({
+        host: addr.address,
+        port: +this.port!,
+      }));
+    }
+
+    /* We lookup both address families here and then split them up later
+     * because when looking up a single family, dns.lookup outputs an error
+     * if the name exists but there are no records for that family, and that
+     * error is indistinguishable from other kinds of errors */
+    const addressList = await dns.lookup(hostname, { all: true });
+    return addressList.map(addr => ({ host: addr.address, port: +this.port! }));
+  }
+
+  private async resolveTxt(hostname: string): Promise<string[][]> {
+    if (process.env[DNS_RESOLUTION_ENV] === 'true') {
+      return this.independentResolver.resolveTxt(hostname);
+    }
+
+    return dns.resolveTxt(hostname);
+  }
+
   private startNextResolutionTimer() {
     clearTimeout(this.nextResolutionTimer);
     this.nextResolutionTimer = setTimeout(() => {