AAD: with ServerTelemetryChannel (#2292)

* AAD: with ServerTelemetryChannel
microsoft · Jun 3, 2021 · 9d7e28a · 9d7e28a
1 parent d0ce00d
commit 9d7e28a
Show file tree

Hide file tree

Showing 13 changed files with 211 additions and 47 deletions.
diff --git a/BASE/Test/ServerTelemetryChannel.Test/TelemetryChannel.Tests/Helpers/MockCredential.cs b/BASE/Test/ServerTelemetryChannel.Test/TelemetryChannel.Tests/Helpers/MockCredential.cs
@@ -0,0 +1,27 @@
+#if NET461 || NETCOREAPP2_1 || NETCOREAPP3_1 || NET5_0
+namespace Microsoft.ApplicationInsights.TestFramework.Helpers
+{
+    using System;
+    using System.Threading;
+    using System.Threading.Tasks;
+
+    using Azure.Core;
+
+
+    /// <remarks>
+    /// Copied from (https://github.com/Azure/azure-sdk-for-net/blob/master/sdk/core/Azure.Core.TestFramework/src/MockCredential.cs).
+    /// </remarks>
+    public class MockCredential : TokenCredential
+    {
+        public override ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken)
+        {
+            return new ValueTask<AccessToken>(GetToken(requestContext, cancellationToken));
+        }
+
+        public override AccessToken GetToken(TokenRequestContext requestContext, CancellationToken cancellationToken)
+        {
+            return new AccessToken("TEST TOKEN " + string.Join(" ", requestContext.Scopes), DateTimeOffset.MaxValue);
+        }
+    }
+}
+#endif
diff --git a/...st/TelemetryChannel.Tests/Implementation/ServerTelemetryChannelCredentialEnvelopeTests.cs b/...st/TelemetryChannel.Tests/Implementation/ServerTelemetryChannelCredentialEnvelopeTests.cs
@@ -0,0 +1,68 @@
+#if !NET452 && !NET46
+namespace Microsoft.ApplicationInsights.TestFramework.Implementation
+{
+    using Microsoft.ApplicationInsights.Channel;
+    using Microsoft.ApplicationInsights.Extensibility;
+    using Microsoft.ApplicationInsights.TestFramework.Helpers;
+    using Microsoft.ApplicationInsights.WindowsServer.TelemetryChannel;
+    using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+    /// <summary>
+    /// These tests verify that <see cref="TelemetryConfiguration"/> can receive and store an instance of <see cref="Azure.Core.TokenCredential"/>.
+    /// </summary>
+    /// <remarks>
+    /// These tests do not run in NET452 OR NET46. 
+    /// In these cases, the test runner is NET452 or NET46 and Azure.Core.TokenCredential is NOT SUPPORTED in these frameworks.
+    /// This does not affect the end user because we REQUIRE the end user to create their own instance of TokenCredential.
+    /// This ensures that the end user is consuming the AI SDK in one of the newer frameworks.
+    /// </remarks>
+    [TestClass]
+    [TestCategory("AAD")]
+    public class ServerTelemetryChannelCredentialEnvelopeTests
+    {
+        [TestMethod]
+        public void VerifySetCredential_CorrectlySetsTelemetryChannel_CredentialFirst()
+        {
+            // SETUP
+            var tc = TelemetryConfiguration.CreateDefault();
+            Assert.IsInstanceOfType(tc.TelemetryChannel, typeof(InMemoryChannel));
+            Assert.IsTrue(tc.TelemetryChannel.EndpointAddress.Contains("v2")); // defaults to old api
+
+            // ACT
+            // set credential first
+            tc.SetAzureTokenCredential(new MockCredential());
+            Assert.IsTrue(tc.TelemetryChannel.EndpointAddress.Contains("v2.1")); // api switch
+
+            // test new channel
+            var channel = new ServerTelemetryChannel();
+            Assert.IsNull(channel.EndpointAddress); // new channel defaults null
+
+            // change config channel
+            tc.TelemetryChannel = channel;
+            Assert.IsTrue(channel.EndpointAddress.Contains("v2.1")); // configuration sets new api
+        }
+
+        [TestMethod]
+        public void VerifySetCredential_CorrectlySetsTelemetryChannel_TelemetryChannelFirst()
+        {
+            // SETUP
+            var tc = TelemetryConfiguration.CreateDefault();
+            Assert.IsInstanceOfType(tc.TelemetryChannel, typeof(InMemoryChannel));
+            Assert.IsTrue(tc.TelemetryChannel.EndpointAddress.Contains("v2")); // defaults to old api
+
+            // ACT
+            // set new channel first
+            var channel = new ServerTelemetryChannel();
+            Assert.IsNull(channel.EndpointAddress); // new channel defaults null
+
+            // change config channel
+            tc.TelemetryChannel = channel;
+            Assert.IsTrue(channel.EndpointAddress.Contains("v2")); // configuration sets new api
+
+            // set credential second
+            tc.SetAzureTokenCredential(new MockCredential());
+            Assert.IsTrue(tc.TelemetryChannel.EndpointAddress.Contains("v2.1")); // api switch
+        }
+    }
+}
+#endif
diff --git a/BASE/Test/ServerTelemetryChannel.Test/TelemetryChannel.Tests/TelemetryChannel.Tests.csproj b/BASE/Test/ServerTelemetryChannel.Test/TelemetryChannel.Tests/TelemetryChannel.Tests.csproj
@@ -22,6 +22,9 @@
   </ItemGroup>
 
   <ItemGroup Condition="'$(TargetFramework)' == 'netcoreapp2.1'">
+    <PackageReference Include="Azure.Core">
+      <Version>1.14.0</Version>
+    </PackageReference>
     <PackageReference Include="Microsoft.AspNetCore.Hosting" Version="2.2.0" />
     <PackageReference Include="Microsoft.AspNetCore.Hosting.Abstractions" Version="2.2.0" />
     <PackageReference Include="Microsoft.AspNetCore.Server.Kestrel" Version="2.2.0" />
@@ -40,5 +43,11 @@
     <Service Include="{82a7f48d-3b50-4b1e-b82e-3ada8210c358}" />
   </ItemGroup>
 
+  <ItemGroup Condition="'$(TargetFramework)' == 'netcoreapp3.1'">
+    <PackageReference Include="Azure.Core">
+      <Version>1.14.0</Version>
+    </PackageReference>
+  </ItemGroup>
+
   <Import Project="..\..\..\Test\TestFramework\Shared\TestFramework.Shared.projitems" Label="TestFramework.Shared" />
 </Project>
diff --git a/BASE/src/Microsoft.ApplicationInsights/Channel/InMemoryChannel.cs b/BASE/src/Microsoft.ApplicationInsights/Channel/InMemoryChannel.cs
@@ -15,7 +15,7 @@
     /// Represents a communication channel for sending telemetry to Application Insights via HTTPS. There will be a buffer that will not be persisted, to enforce the 
     /// queued telemetry items to be sent, <see cref="ITelemetryChannel.Flush"/> should be called.    
     /// </summary>
-    public class InMemoryChannel : ITelemetryChannel, IAsyncFlushable
+    public class InMemoryChannel : ITelemetryChannel, IAsyncFlushable, ISupportCredentialEnvelope
     {
         private readonly TelemetryBuffer buffer;
         private readonly InMemoryTransmitter transmitter;
@@ -96,6 +96,19 @@ public TimeSpan SendingInterval
             }
         }
 
+        /// <summary>
+        /// Gets or sets the <see cref="CredentialEnvelope"/> which is used for AAD.
+        /// </summary>
+        /// <remarks>
+        /// <see cref="ISupportCredentialEnvelope.CredentialEnvelope"/> on <see cref="InMemoryChannel"/> sets <see cref="InMemoryTransmitter.CredentialEnvelope"/> 
+        /// which is used to set <see cref="Transmission.CredentialEnvelope"/> just before calling <see cref="Transmission.SendAsync"/>.
+        /// </remarks>
+        CredentialEnvelope ISupportCredentialEnvelope.CredentialEnvelope
+        {
+            get => this.transmitter.CredentialEnvelope;
+            set => this.transmitter.CredentialEnvelope = value;
+        }
+
         /// <summary>
         /// Gets or sets the HTTP address where the telemetry is sent.
         /// </summary>
@@ -126,20 +139,6 @@ public int BacklogSize
             set { this.buffer.BacklogSize = value; }
         }
 
-        /// <summary>
-        /// Gets or sets the <see cref="CredentialEnvelope"/> which is used for AAD.
-        /// FOR INTERNAL USE. Customers should use <see cref="TelemetryConfiguration.SetAzureTokenCredential"/> instead.
-        /// </summary>
-        /// <remarks>
-        /// <see cref="InMemoryChannel.CredentialEnvelope"/> sets <see cref="InMemoryTransmitter.CredentialEnvelope"/> 
-        /// which is used to set <see cref="Transmission.CredentialEnvelope"/> just before calling <see cref="Transmission.SendAsync"/>.
-        /// </remarks>
-        internal CredentialEnvelope CredentialEnvelope
-        {
-            get => this.transmitter.CredentialEnvelope;
-            set => this.transmitter.CredentialEnvelope = value;
-        }
-
         internal bool IsDisposed => this.isDisposed;
 
         /// <summary>

diff --git a/BASE/src/Microsoft.ApplicationInsights/Channel/InMemoryTransmitter.cs b/BASE/src/Microsoft.ApplicationInsights/Channel/InMemoryTransmitter.cs
@@ -69,7 +69,7 @@ internal TimeSpan SendingInterval
         /// Gets or sets the <see cref="CredentialEnvelope"/> which is used for AAD.
         /// </summary>
         /// <remarks>
-        /// <see cref="InMemoryChannel.CredentialEnvelope"/> sets <see cref="InMemoryTransmitter.CredentialEnvelope"/> 
+        /// <see cref="ISupportCredentialEnvelope.CredentialEnvelope"/> on <see cref="InMemoryChannel"/> sets <see cref="InMemoryTransmitter.CredentialEnvelope"/> 
         /// which is used to set <see cref="Transmission.CredentialEnvelope"/> just before calling <see cref="Transmission.SendAsync"/>.
         /// </remarks>
         internal CredentialEnvelope CredentialEnvelope { get; set; }

diff --git a/...icrosoft.ApplicationInsights/Extensibility/Implementation/Authentication/AuthConstants.cs b/...icrosoft.ApplicationInsights/Extensibility/Implementation/Authentication/AuthConstants.cs
@@ -11,7 +11,7 @@ internal static class AuthConstants
         /// (https://docs.microsoft.com/azure/active-directory/develop/msal-acquire-cache-tokens#scopes-when-acquiring-tokens).
         /// (https://docs.microsoft.com/azure/active-directory/develop/v2-permissions-and-consent#the-default-scope).
         /// </summary>
-        private const string AzureMonitorScope = "https://monitor.azure.com//.default"; // TODO: THIS SCOPE IS UNVERIFIED. WAITING FOR SERVICES TEAM TO PROVIDE AN INT ENVIRONMENT FOR E2E TESTING.
+        private const string AzureMonitorScope = "https://monitor.azure.com//.default";
 
         /// <summary>
         /// Get scopes for Azure Monitor as an array.

diff --git a/...icationInsights/Extensibility/Implementation/Authentication/ISupportCredentialEnvelope.cs b/...icationInsights/Extensibility/Implementation/Authentication/ISupportCredentialEnvelope.cs
@@ -0,0 +1,13 @@
+namespace Microsoft.ApplicationInsights.Extensibility.Implementation.Authentication
+{
+    /// <summary>
+    /// This interface defines a class that accepts the <see cref="Authentication.CredentialEnvelope"/> as a property.
+    /// </summary>
+    internal interface ISupportCredentialEnvelope
+    {
+        /// <summary>
+        /// Gets or sets the <see cref="Authentication.CredentialEnvelope"/>.
+        /// </summary>
+        CredentialEnvelope CredentialEnvelope { get; set; }
+    }
+}
diff --git a/...ationInsights/Extensibility/Implementation/Authentication/ReflectionCredentialEnvelope.cs b/...ationInsights/Extensibility/Implementation/Authentication/ReflectionCredentialEnvelope.cs
@@ -53,7 +53,6 @@ public ReflectionCredentialEnvelope(object tokenCredential)
         /// <returns>A valid Azure.Core.AccessToken.</returns>
         public override string GetToken(CancellationToken cancellationToken = default)
         {
-            // TODO: NEED TO FULLY TEST IF WE NEED TO CALL SdkInternalOperationsMonitor.Enter
             try
             {
                 return AzureCore.InvokeGetToken(this.tokenCredential, this.tokenRequestContext, cancellationToken);
@@ -75,7 +74,6 @@ public override string GetToken(CancellationToken cancellationToken = default)
         /// <returns>A valid Azure.Core.AccessToken.</returns>
         public override async Task<string> GetTokenAsync(CancellationToken cancellationToken = default)
         {
-            // TODO: NEED TO FULLY TEST IF WE NEED TO CALL SdkInternalOperationsMonitor.Enter
             try
             {
                 return await AzureCore.InvokeGetTokenAsync(this.tokenCredential, this.tokenRequestContext, cancellationToken).ConfigureAwait(false);

diff --git a/BASE/src/Microsoft.ApplicationInsights/Extensibility/TelemetryConfiguration.cs b/BASE/src/Microsoft.ApplicationInsights/Extensibility/TelemetryConfiguration.cs
@@ -339,7 +339,9 @@ public string ConnectionString
 
         /// <summary>
         /// Gets an envelope for Azure.Core.TokenCredential which provides an AAD Authenticated token.
+        /// FOR INTERNAL USE ONLY. To set the Credential use <see cref="SetAzureTokenCredential"/>.
         /// </summary>
+        [EditorBrowsable(EditorBrowsableState.Never)]
         public CredentialEnvelope CredentialEnvelope { get; private set; }
 
         /// <summary>
@@ -504,9 +506,9 @@ private static void SetTelemetryChannelEndpoint(ITelemetryChannel channel, strin
 
         private static void SetTelemetryChannelCredentialEnvelope(ITelemetryChannel telemetryChannel, CredentialEnvelope credentialEnvelope)
         {
-            if (telemetryChannel is InMemoryChannel inMemoryChannel)
+            if (telemetryChannel is ISupportCredentialEnvelope tc)
             {
-                inMemoryChannel.CredentialEnvelope = credentialEnvelope;
+                tc.CredentialEnvelope = credentialEnvelope;
             }
         }
 

diff --git a/BASE/src/ServerTelemetryChannel/Implementation/ResponseStatusCodes.cs b/BASE/src/ServerTelemetryChannel/Implementation/ResponseStatusCodes.cs
@@ -11,6 +11,8 @@ internal class ResponseStatusCodes
         public const int InternalServerError = 500;
         public const int BadGateway = 502;
         public const int ServiceUnavailable = 503;
-        public const int GatewayTimeout = 504;
+        public const int GatewayTimeout = 504; 
+        public const int Unauthorized = 401; // AAD
+        public const int Forbidden = 403; // AAD
     }
 }
diff --git a/BASE/src/ServerTelemetryChannel/Implementation/TransmissionSender.cs b/BASE/src/ServerTelemetryChannel/Implementation/TransmissionSender.cs
@@ -9,13 +9,15 @@
     using System.Net.Http;
     using System.Threading;
     using System.Threading.Tasks;
+
     using Microsoft.ApplicationInsights.Channel;
     using Microsoft.ApplicationInsights.Common.Extensions;
     using Microsoft.ApplicationInsights.Extensibility;
     using Microsoft.ApplicationInsights.Extensibility.Implementation;
+    using Microsoft.ApplicationInsights.Extensibility.Implementation.Authentication;
 
     internal class TransmissionSender
-    {      
+    {
         private static readonly HttpWebResponseWrapper DefaultHttpWebResponseWrapper = default(HttpWebResponseWrapper);
         // Stores all inflight requests using this dictionary, before SendAsync.
         // Removes entry from dictionary after response.
@@ -107,6 +109,15 @@ public virtual int ThrottleWindow
             }
         }
 
+        /// <summary>
+        /// Gets or sets the <see cref="CredentialEnvelope"/> which is used for AAD.
+        /// </summary>
+        /// <remarks>
+        /// <see cref="ISupportCredentialEnvelope.CredentialEnvelope"/> on <see cref="ServerTelemetryChannel"/> sets <see cref="Transmitter.CredentialEnvelope"/> and then sets <see cref="TransmissionSender.CredentialEnvelope"/> 
+        /// which is used to set <see cref="Transmission.CredentialEnvelope"/> just before calling <see cref="Transmission.SendAsync"/>.
+        /// </remarks>
+        internal CredentialEnvelope CredentialEnvelope { get; set; }
+
         public virtual bool Enqueue(Func<Transmission> transmissionGetter)
         {
             bool enqueueSucceded = false;
@@ -197,9 +208,11 @@ private async Task StartSending(Transmission transmission)
                 try
                 {
                     TelemetryChannelEventSource.Log.TransmissionSendStarted(acceptedTransmission.Id);
+                    acceptedTransmission.CredentialEnvelope = this.CredentialEnvelope;
+
                     transmissionTask = acceptedTransmission.SendAsync();
                     this.inFlightTransmissions.TryAdd(transmission.FlushAsyncId, transmissionTask);
-                    responseContent = await transmissionTask.ConfigureAwait(false); 
+                    responseContent = await transmissionTask.ConfigureAwait(false);
                 }
                 catch (Exception e)
                 {
@@ -302,7 +315,7 @@ private Transmission Throttle(Transmission transmission)
 
             int attemptedItemsCount = -1;
             int acceptedItemsCount = -1;
-            Tuple<Transmission, Transmission> transmissions = transmission.Split((transmissionLength) => 
+            Tuple<Transmission, Transmission> transmissions = transmission.Split((transmissionLength) =>
             {
                 attemptedItemsCount = transmissionLength;
                 acceptedItemsCount = this.IsTransmissionSendable(transmissionLength);