[core-rest-pipeline] Add support for custom agents. (Azure#14450)

On NodeJS it's often useful to be able to control the agent used to make a request.

This change adds support for setting a custom agent on a per-request basis and also refactors ProxyPolicy to be in charge of proxy agents.

Author: xirzec

sdk/core/core-rest-pipeline/review/core-rest-pipeline.api.md

@@ -17,6 +17,15 @@ export interface AddPipelineOptions {
     phase?: PipelinePhase;
 }
 
+// @public
+export interface Agent {
+    destroy(): void;
+    maxFreeSockets: number;
+    maxSockets: number;
+    requests: unknown;
+    sockets: unknown;
+}
+
 // @public
 export function bearerTokenAuthenticationPolicy(options: BearerTokenAuthenticationPolicyOptions): PipelinePolicy;
 
@@ -153,6 +162,7 @@ export interface PipelinePolicy {
 // @public
 export interface PipelineRequest {
     abortSignal?: AbortSignalLike;
+    agent?: Agent;
     allowInsecureConnection?: boolean;
     body?: RequestBodyType;
     disableKeepAlive?: boolean;

sdk/core/core-rest-pipeline/src/index.ts

@@ -2,6 +2,7 @@
 // Licensed under the MIT license.
 
 export {
+  Agent,
   HttpClient,
   PipelineRequest,
   PipelineResponse,

sdk/core/core-rest-pipeline/src/interfaces.ts

@@ -56,6 +56,34 @@ export type RequestBodyType =
   | string
   | null;
 
+/**
+ * An interface compatible with NodeJS's `http.Agent`.
+ * We want to avoid publicly re-exporting the actual interface,
+ * since it might vary across runtime versions.
+ */
+export interface Agent {
+  /**
+   * Destroy any sockets that are currently in use by the agent.
+   */
+  destroy(): void;
+  /**
+   * For agents with keepAlive enabled, this sets the maximum number of sockets that will be left open in the free state.
+   */
+  maxFreeSockets: number;
+  /**
+   * Determines how many concurrent sockets the agent can have open per origin.
+   */
+  maxSockets: number;
+  /**
+   * An object which contains queues of requests that have not yet been assigned to sockets.
+   */
+  requests: unknown;
+  /**
+   * An object which contains arrays of sockets currently in use by the agent.
+   */
+  sockets: unknown;
+}
+
 /**
  * Metadata about a request being made by the pipeline.
  */
@@ -138,6 +166,14 @@ export interface PipelineRequest {
 
   /** Set to true if the request is sent over HTTP instead of HTTPS */
   allowInsecureConnection?: boolean;
+
+  /**
+   * NODEJS ONLY
+   *
+   * A Node-only option to provide a custom `http.Agent`/`https.Agent`.
+   * Does nothing when running in the browser.
+   */
+  agent?: Agent;
 }
 
 /**

sdk/core/core-rest-pipeline/src/nodeHttpClient.ts

@@ -5,17 +5,14 @@ import * as http from "http";
 import * as https from "https";
 import * as zlib from "zlib";
 import { Transform } from "stream";
-import { HttpsProxyAgent, HttpsProxyAgentOptions } from "https-proxy-agent";
-import { HttpProxyAgent, HttpProxyAgentOptions } from "http-proxy-agent";
 import { AbortController, AbortError } from "@azure/abort-controller";
 import {
   HttpClient,
   PipelineRequest,
   PipelineResponse,
   TransferProgressEvent,
   HttpHeaders,
-  RequestBodyType,
-  ProxySettings
+  RequestBodyType
 } from "./interfaces";
 import { createHttpHeaders } from "./httpHeaders";
 import { RestError } from "./restError";
@@ -67,9 +64,7 @@ class ReportTransform extends Transform {
  */
 class NodeHttpClient implements HttpClient {
   private httpsKeepAliveAgent?: https.Agent;
-  private httpsProxyAgent?: https.Agent;
   private httpKeepAliveAgent?: http.Agent;
-  private httpProxyAgent?: http.Agent;
 
   /**
    * Makes a request over an underlying transport layer and returns the response.
@@ -205,7 +200,7 @@ class NodeHttpClient implements HttpClient {
       throw new Error(`Cannot connect to ${request.url} while allowInsecureConnection is false.`);
     }
 
-    const agent = this.getOrCreateAgent(request, isInsecure);
+    const agent = request.agent ?? this.getOrCreateAgent(request, isInsecure);
     const options: http.RequestOptions = {
       agent,
       hostname: url.hostname,
@@ -221,66 +216,25 @@ class NodeHttpClient implements HttpClient {
     }
   }
 
-  private getProxyAgentOptions(
-    proxySettings: ProxySettings,
-    requestHeaders: HttpHeaders
-  ): HttpProxyAgentOptions {
-    let parsedProxyUrl: URL;
-    try {
-      parsedProxyUrl = new URL(proxySettings.host);
-    } catch (_error) {
-      throw new Error(
-        `Expecting a valid host string in proxy settings, but found "${proxySettings.host}".`
-      );
-    }
-
-    const proxyAgentOptions: HttpsProxyAgentOptions = {
-      hostname: parsedProxyUrl.hostname,
-      port: proxySettings.port,
-      protocol: parsedProxyUrl.protocol,
-      headers: requestHeaders.toJSON()
-    };
-    if (proxySettings.username && proxySettings.password) {
-      proxyAgentOptions.auth = `${proxySettings.username}:${proxySettings.password}`;
-    }
-    return proxyAgentOptions;
-  }
-
   private getOrCreateAgent(request: PipelineRequest, isInsecure: boolean): http.Agent {
-    // At the moment, proxy settings and keepAlive are mutually
-    // exclusive because the proxy library currently lacks the
-    // ability to create a proxy with keepAlive turned on.
-    const proxySettings = request.proxySettings;
-    if (proxySettings) {
+    if (!request.disableKeepAlive) {
       if (isInsecure) {
-        if (!this.httpProxyAgent) {
-          const proxyAgentOptions = this.getProxyAgentOptions(proxySettings, request.headers);
-          this.httpProxyAgent = new HttpProxyAgent(proxyAgentOptions);
+        if (!this.httpKeepAliveAgent) {
+          this.httpKeepAliveAgent = new http.Agent({
+            keepAlive: true
+          });
         }
-        return this.httpProxyAgent;
+
+        return this.httpKeepAliveAgent;
       } else {
-        if (!this.httpsProxyAgent) {
-          const proxyAgentOptions = this.getProxyAgentOptions(proxySettings, request.headers);
-          this.httpsProxyAgent = new HttpsProxyAgent(proxyAgentOptions);
+        if (!this.httpsKeepAliveAgent) {
+          this.httpsKeepAliveAgent = new https.Agent({
+            keepAlive: true
+          });
         }
-        return this.httpsProxyAgent;
-      }
-    } else if (!request.disableKeepAlive && !isInsecure) {
-      if (!this.httpsKeepAliveAgent) {
-        this.httpsKeepAliveAgent = new https.Agent({
-          keepAlive: true
-        });
-      }
 
-      return this.httpsKeepAliveAgent;
-    } else if (!request.disableKeepAlive && isInsecure) {
-      if (!this.httpKeepAliveAgent) {
-        this.httpKeepAliveAgent = new http.Agent({
-          keepAlive: true
-        });
+        return this.httpsKeepAliveAgent;
       }
-
-      return this.httpKeepAliveAgent;
     } else if (isInsecure) {
       return http.globalAgent;
     } else {

sdk/core/core-rest-pipeline/src/policies/proxyPolicy.ts

@@ -1,7 +1,17 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-import { PipelineResponse, PipelineRequest, SendRequest, ProxySettings } from "../interfaces";
+import * as http from "http";
+import * as https from "https";
+import { HttpsProxyAgent, HttpsProxyAgentOptions } from "https-proxy-agent";
+import { HttpProxyAgent, HttpProxyAgentOptions } from "http-proxy-agent";
+import {
+  PipelineResponse,
+  PipelineRequest,
+  SendRequest,
+  ProxySettings,
+  HttpHeaders
+} from "../interfaces";
 import { PipelinePolicy } from "../pipeline";
 import { URL } from "../util/url";
 
@@ -18,6 +28,9 @@ export const proxyPolicyName = "proxyPolicy";
 export const noProxyList: string[] = loadNoProxy();
 const byPassedList: Map<string, boolean> = new Map();
 
+let httpsProxyAgent: https.Agent | undefined;
+let httpProxyAgent: http.Agent | undefined;
+
 function getEnvironmentValue(name: string): string | undefined {
   if (process.env[name]) {
     return process.env[name];
@@ -107,6 +120,54 @@ export function getDefaultProxySettings(proxyUrl?: string): ProxySettings | unde
   };
 }
 
+function getProxyAgentOptions(
+  proxySettings: ProxySettings,
+  requestHeaders: HttpHeaders
+): HttpProxyAgentOptions {
+  let parsedProxyUrl: URL;
+  try {
+    parsedProxyUrl = new URL(proxySettings.host);
+  } catch (_error) {
+    throw new Error(
+      `Expecting a valid host string in proxy settings, but found "${proxySettings.host}".`
+    );
+  }
+
+  const proxyAgentOptions: HttpsProxyAgentOptions = {
+    hostname: parsedProxyUrl.hostname,
+    port: proxySettings.port,
+    protocol: parsedProxyUrl.protocol,
+    headers: requestHeaders.toJSON()
+  };
+  if (proxySettings.username && proxySettings.password) {
+    proxyAgentOptions.auth = `${proxySettings.username}:${proxySettings.password}`;
+  }
+  return proxyAgentOptions;
+}
+
+function setProxyAgentOnRequest(request: PipelineRequest): void {
+  const url = new URL(request.url);
+
+  const isInsecure = url.protocol !== "https:";
+
+  const proxySettings = request.proxySettings;
+  if (proxySettings) {
+    if (isInsecure) {
+      if (!httpProxyAgent) {
+        const proxyAgentOptions = getProxyAgentOptions(proxySettings, request.headers);
+        httpProxyAgent = new HttpProxyAgent(proxyAgentOptions);
+      }
+      request.agent = httpProxyAgent;
+    } else {
+      if (!httpsProxyAgent) {
+        const proxyAgentOptions = getProxyAgentOptions(proxySettings, request.headers);
+        httpsProxyAgent = new HttpsProxyAgent(proxyAgentOptions);
+      }
+      request.agent = httpsProxyAgent;
+    }
+  }
+}
+
 /**
  * A policy that allows one to apply proxy settings to all requests.
  * If not passed static settings, they will be retrieved from the HTTPS_PROXY
@@ -120,6 +181,10 @@ export function proxyPolicy(proxySettings = getDefaultProxySettings()): Pipeline
       if (!request.proxySettings && !isBypassed(request.url)) {
         request.proxySettings = proxySettings;
       }
+
+      if (request.proxySettings) {
+        setProxyAgentOnRequest(request);
+      }
       return next(request);
     }
   };