Adding AAD handling (#841)

* Statsbeat update function resource provider

* WIP

* Adding AAD handling

* Build

* Change statbeat tests

* Test

* Test

.vscode/tasks.json

@@ -1,19 +1,32 @@
 {
-    "version": "0.1.0",
+    "version": "2.0.0",
     "command": "npm",
-    "isShellCommand": true,
-    "showOutput": "silent",
-    "suppressTaskName": true,
     "tasks": [
         {
-            "taskName": "test",
-            "args": ["run", "test"],
-            "isTestCommand": true
+            "label": "test",
+            "type": "shell",
+            "args": [
+                "run",
+                "test"
+            ],
+            "problemMatcher": [],
+            "group": {
+                "_id": "test",
+                "isDefault": false
+            }
         },
         {
-            "taskName": "build",
-            "args": ["run", "build"],
-            "isTestCommand": false
+            "label": "build",
+            "type": "shell",
+            "args": [
+                "run",
+                "build"
+            ],
+            "problemMatcher": [],
+            "group": {
+                "_id": "build",
+                "isDefault": false
+            }
         }
     ]
 }

Library/AuthorizationHandler.ts

@@ -0,0 +1,35 @@
+import http = require("http");
+import https = require("https");
+import azureCore = require("@azure/core-http");
+
+const applicationInsightsResource = "https://monitor.azure.com//.default";
+
+
+class AuthorizationHandler {
+
+    private _azureTokenPolicy: azureCore.RequestPolicy;
+
+    constructor(credential: azureCore.TokenCredential) {
+        let scopes: string[] = [applicationInsightsResource];
+        let emptyPolicy: azureCore.RequestPolicy = {
+            sendRequest(httpRequest: azureCore.WebResourceLike): Promise<azureCore.HttpOperationResponse> {
+                return null;
+            }
+        };
+        this._azureTokenPolicy = azureCore.bearerTokenAuthenticationPolicy(credential, scopes).create(emptyPolicy, new azureCore.RequestPolicyOptions());
+    }
+
+    /**
+    * Applies the Bearer token to the request through the Authorization header.
+    */
+    public async addAuthorizationHeader(requestOptions: http.RequestOptions | https.RequestOptions): Promise<void> {
+        let authHeaderName = azureCore.Constants.HeaderConstants.AUTHORIZATION;
+        let webResource = new azureCore.WebResource("https://");
+        this
+        await this._azureTokenPolicy.sendRequest(webResource);
+        requestOptions.headers[authHeaderName] = webResource.headers.get(authHeaderName);
+    }
+
+}
+
+export = AuthorizationHandler;

Library/Config.ts

@@ -54,9 +54,10 @@ class Config {
     public httpAgent: http.Agent;
     /** An https.Agent to use for SDK HTTPS traffic (Optional, Default undefined) */
     public httpsAgent: https.Agent;
-
     /** Disable including legacy headers in outgoing requests, x-ms-request-id */
     public ignoreLegacyHeaders?: boolean;
+    /** AAD TokenCredential to use to authenticate the app */
+    public aadTokenCredential?: azureCore.TokenCredential;
 
     private endpointBase: string = Constants.DEFAULT_BREEZE_ENDPOINT;
     private setCorrelationId: (v: string) => void;

Library/QuickPulseSender.ts

@@ -1,5 +1,6 @@
 import https = require("https");
 
+import AuthorizationHandler = require("./AuthorizationHandler");
 import Config = require("./Config");
 import AutoCollectHttpDependencies = require("../AutoCollection/HttpDependencies");
 import Logging = require("./Logging");
@@ -30,10 +31,12 @@ class QuickPulseSender {
 
     private _config: Config;
     private _consecutiveErrors: number;
+    private _getAuthorizationHandler: (config: Config) => AuthorizationHandler;
 
-    constructor(config: Config) {
+    constructor(config: Config, getAuthorizationHandler?: (config: Config) => AuthorizationHandler) {
         this._config = config;
         this._consecutiveErrors = 0;
+        this._getAuthorizationHandler = getAuthorizationHandler;
     }
 
     public ping(envelope: Contracts.EnvelopeQuickPulse,
@@ -85,6 +88,22 @@ class QuickPulseSender {
             additionalHeaders.forEach(header => options.headers[header.name] = header.value);
         }
 
+        if (postOrPing === "post") {
+            let authHandler = this._getAuthorizationHandler ? this._getAuthorizationHandler(this._config) : null;
+            if (authHandler) {
+                try {
+                    // Add bearer token
+                    await authHandler.addAuthorizationHeader(options);
+                }
+                catch (authError) {
+                    let notice = `Failed to get AAD bearer token for the Application. Error:`;
+                    Logging.info(QuickPulseSender.TAG, notice, authError);
+                    // Do not send request to Quickpulse if auth fails, data will be dropped
+                    return;
+                }
+            }
+        }
+
         // HTTPS only
         if (this._config.httpsAgent) {
             (<any>options).agent = this._config.httpsAgent;
@@ -131,4 +150,4 @@ class QuickPulseSender {
     }
 }
 
-export = QuickPulseSender;
+export = QuickPulseSender;

Library/QuickPulseStateManager.ts

@@ -1,3 +1,4 @@
+import AuthorizationHandler = require("./AuthorizationHandler");
 import Logging = require("./Logging");
 import Config = require("./Config");
 import QuickPulseEnvelopeFactory = require("./QuickPulseEnvelopeFactory");
@@ -13,6 +14,7 @@ import * as Contracts from "../Declarations/Contracts";
 class QuickPulseStateManager {
     public config: Config;
     public context: Context;
+    public authorizationHandler: AuthorizationHandler;
 
     private static MAX_POST_WAIT_TIME = 20000;
     private static MAX_PING_WAIT_TIME = 60000;
@@ -32,10 +34,10 @@ class QuickPulseStateManager {
     private _redirectedHost: string = null;
     private _pollingIntervalHint: number = -1;
 
-    constructor(config: Config, context?: Context) {
+    constructor(config: Config, context?: Context, getAuthorizationHandler?: (config: Config) => AuthorizationHandler) {
         this.config = config;
         this.context = context || new Context();
-        this._sender = new QuickPulseSender(this.config);
+        this._sender = new QuickPulseSender(this.config, getAuthorizationHandler);
         this._isEnabled = false;
     }
 
@@ -189,4 +191,4 @@ class QuickPulseStateManager {
     }
 }
 
-export = QuickPulseStateManager;
+export = QuickPulseStateManager;

Library/Sender.ts

@@ -5,6 +5,7 @@ import path = require("path");
 import zlib = require("zlib");
 import child_process = require("child_process");
 
+import AuthorizationHandler = require("./AuthorizationHandler");
 import Logging = require("./Logging");
 import Config = require("./Config")
 import Contracts = require("../Declarations/Contracts");
@@ -36,6 +37,7 @@ class Sender {
     private _statsbeat: Statsbeat;
     private _onSuccess: (response: string) => void;
     private _onError: (error: Error) => void;
+    private _getAuthorizationHandler: (config: Config) => AuthorizationHandler;
     private _enableDiskRetryMode: boolean;
     private _numConsecutiveFailures: number;
     private _numConsecutiveRedirects: number;
@@ -46,7 +48,7 @@ class Sender {
     protected _resendInterval: number;
     protected _maxBytesOnDisk: number;
 
-    constructor(config: Config, onSuccess?: (response: string) => void, onError?: (error: Error) => void, statsbeat?: Statsbeat) {
+    constructor(config: Config, getAuthorizationHandler?: (config: Config) => AuthorizationHandler, onSuccess?: (response: string) => void, onError?: (error: Error) => void, statsbeat?: Statsbeat) {
         this._config = config;
         this._onSuccess = onSuccess;
         this._onError = onError;
@@ -57,6 +59,7 @@ class Sender {
         this._numConsecutiveFailures = 0;
         this._numConsecutiveRedirects = 0;
         this._resendTimer = null;
+        this._getAuthorizationHandler = getAuthorizationHandler;
         this._fileCleanupTimer = null;
         // tmpdir is /tmp for *nix and USERDIR/AppData/Local/Temp for Windows
         this._tempDir = path.join(os.tmpdir(), Sender.TEMPDIR_PREFIX + this._config.instrumentationKey);
@@ -132,6 +135,27 @@ class Sender {
                 }
             };
 
+            let authHandler = this._getAuthorizationHandler ? this._getAuthorizationHandler(this._config) : null;
+            if (authHandler) {
+                if (this._statsbeat) {
+                    this._statsbeat.addFeature(Constants.StatsbeatFeature.AAD_HANDLING);
+                }
+                try {
+                    // Add bearer token
+                    await authHandler.addAuthorizationHeader(options);
+                }
+                catch (authError) {
+                    let errorMsg = "Failed to get AAD bearer token for the Application. Error:" + authError.toString();
+                    // If AAD auth fails do not send to Breeze
+                    if (typeof callback === "function") {
+                        callback(errorMsg);
+                    }
+                    this._storeToDisk(envelopes);
+                    Logging.warn(Sender.TAG, errorMsg);
+                    return;
+                }
+            }
+
             let batch: string = "";
             envelopes.forEach(envelope => {
                 var payload: string = this._stringify(envelope);
@@ -238,7 +262,7 @@ class Sender {
                         }
                         else {
                             if (this._statsbeat) {
-                                this._statsbeat.countRequest(Constants.StatsbeatNetworkCategory.Breeze, endpointHost,duration, res.statusCode === 200);
+                                this._statsbeat.countRequest(Constants.StatsbeatNetworkCategory.Breeze, endpointHost, duration, res.statusCode === 200);
                             }
                             this._numConsecutiveRedirects = 0;
                             if (typeof callback === "function") {
@@ -305,6 +329,8 @@ class Sender {
     private _isRetriable(statusCode: number) {
         return (
             statusCode === 206 || // Retriable
+            statusCode === 401 || // Unauthorized
+            statusCode === 403 || // Forbidden
             statusCode === 408 || // Timeout
             statusCode === 429 || // Throttle
             statusCode === 439 || // Quota
@@ -663,4 +689,4 @@ class Sender {
     }
 }
 
-export = Sender;
+export = Sender;

Library/TelemetryClient.ts

@@ -3,6 +3,7 @@ import os = require("os");
 import azureCore = require("@azure/core-http");
 
 import Config = require("./Config");
+import AuthorizationHandler = require("./AuthorizationHandler");
 import Context = require("./Context");
 import Contracts = require("../Declarations/Contracts");
 import Channel = require("./Channel");
@@ -31,6 +32,7 @@ class TelemetryClient {
     public commonProperties: { [key: string]: string; };
     public channel: Channel;
     public quickPulseClient: QuickPulseStateManager;
+    public authorizationHandler: AuthorizationHandler;
 
     /**
      * Constructs a new client of the client
@@ -41,11 +43,12 @@ class TelemetryClient {
         this.config = config;
         this.context = new Context();
         this.commonProperties = {};
+        this.authorizationHandler = null;
         if (!process.env["APPLICATION_INSIGHTS_NO_STATSBEAT"]) {
             this._statsbeat = new Statsbeat(this.config);
             this._statsbeat.enable(true);
         }
-        var sender = new Sender(this.config, null, null, this._statsbeat);
+        var sender = new Sender(this.config, this.getAuthorizationHandler, null, null, this._statsbeat);
         this.channel = new Channel(() => config.disableAppInsights, () => config.maxBatchSize, () => config.maxBatchIntervalMs, sender);
     }
 
@@ -187,6 +190,16 @@ class TelemetryClient {
         this._enableAzureProperties = value;
     }
 
+    /**
+     * Get Authorization handler
+     */
+    public getAuthorizationHandler(config: Config): AuthorizationHandler {
+        if (config && config.aadTokenCredential) {
+            return this.authorizationHandler || new AuthorizationHandler(config.aadTokenCredential);
+        }
+        return null;
+    }
+
     /**
      * Adds telemetry processor to the collection. Telemetry processors will be called one by one
      * before telemetry item is pushed for sending and in the order they were added.
@@ -252,4 +265,4 @@ class TelemetryClient {
     }
 }
 
-export = TelemetryClient;
+export = TelemetryClient;