Skip to content

Commit 72fbd32

Browse files
author
Timothy Mothra
authored
AAD: with InMemoryChannel (#2290)
* aad with InMemoryChannel
1 parent 9ecce67 commit 72fbd32

File tree

11 files changed

+277
-26
lines changed

11 files changed

+277
-26
lines changed

‎BASE/Test/Microsoft.ApplicationInsights.Test/Microsoft.ApplicationInsights.Tests/Extensibility/Implementation/Authentication/TelemetryConfigurationCredentialEnvelopeTests.cs‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ namespace Microsoft.ApplicationInsights.TestFramework.Extensibility.Implementati
44
using System;
55
using System.Threading.Tasks;
66

7+
using Microsoft.ApplicationInsights.Channel;
78
using Microsoft.ApplicationInsights.Extensibility;
89
using Microsoft.ApplicationInsights.Extensibility.Implementation.Authentication;
910
using Microsoft.VisualStudio.TestTools.UnitTesting;
@@ -47,6 +48,50 @@ public void VerifyCannotSetInvalidObjectOnTelemetryConfiguration()
4748
var telemetryConfiguration = new TelemetryConfiguration();
4849
telemetryConfiguration.SetAzureTokenCredential(Guid.Empty);
4950
}
51+
52+
[TestMethod]
53+
public void VerifySetCredential_CorrectlySetsTelemetryChannel_CredentialFirst()
54+
{
55+
// SETUP
56+
var tc = TelemetryConfiguration.CreateDefault();
57+
Assert.IsInstanceOfType(tc.TelemetryChannel, typeof(InMemoryChannel));
58+
Assert.IsTrue(tc.TelemetryChannel.EndpointAddress.Contains("v2")); // defaults to old api
59+
60+
// ACT
61+
// set credential first
62+
tc.SetAzureTokenCredential(new MockCredential());
63+
Assert.IsTrue(tc.TelemetryChannel.EndpointAddress.Contains("v2.1")); // api switch
64+
65+
// test new channel
66+
var channel = new InMemoryChannel();
67+
Assert.IsNull(channel.EndpointAddress); // new channel defaults null
68+
69+
// change config channel
70+
tc.TelemetryChannel = channel;
71+
Assert.IsTrue(channel.EndpointAddress.Contains("v2.1")); // configuration sets new api
72+
}
73+
74+
[TestMethod]
75+
public void VerifySetCredential_CorrectlySetsTelemetryChannel_TelemetryChannelFirst()
76+
{
77+
// SETUP
78+
var tc = TelemetryConfiguration.CreateDefault();
79+
Assert.IsInstanceOfType(tc.TelemetryChannel, typeof(InMemoryChannel));
80+
Assert.IsTrue(tc.TelemetryChannel.EndpointAddress.Contains("v2")); // defaults to old api
81+
82+
// ACT
83+
// set new channel first
84+
var channel = new InMemoryChannel();
85+
Assert.IsNull(channel.EndpointAddress); // new channel defaults null
86+
87+
// change config channel
88+
tc.TelemetryChannel = channel;
89+
Assert.IsTrue(channel.EndpointAddress.Contains("v2")); // configuration sets new api
90+
91+
// set credential second
92+
tc.SetAzureTokenCredential(new MockCredential());
93+
Assert.IsTrue(tc.TelemetryChannel.EndpointAddress.Contains("v2.1")); // api switch
94+
}
5095
}
5196
}
5297
#endif

‎BASE/Test/Microsoft.ApplicationInsights.Test/Microsoft.ApplicationInsights.Tests/Extensibility/Implementation/Authentication/TransmissionCredentialEnvelopeTests.cs‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
#if !NET452 && !NET46
2+
namespace Microsoft.ApplicationInsights.TestFramework.Extensibility.Implementation.Authentication
3+
{
4+
using System;
5+
using System.Collections.Generic;
6+
using System.Linq;
7+
using System.Net.Http;
8+
using System.Text;
9+
using System.Threading.Tasks;
10+
11+
using Microsoft.ApplicationInsights.Channel;
12+
using Microsoft.ApplicationInsights.DataContracts;
13+
using Microsoft.ApplicationInsights.Extensibility.Implementation.Authentication;
14+
using Microsoft.VisualStudio.TestTools.UnitTesting;
15+
16+
/// <summary>
17+
/// These tests verify that <see cref="Transmission"/> can receive and store an instance of <see cref="Azure.Core.TokenCredential"/>.
18+
/// </summary>
19+
/// <remarks>
20+
/// These tests do not run in NET452 OR NET46.
21+
/// In these cases, the test runner is NET452 or NET46 and Azure.Core.TokenCredential is NOT SUPPORTED in these frameworks.
22+
/// This does not affect the end user because we REQUIRE the end user to create their own instance of TokenCredential.
23+
/// This ensures that the end user is consuming the AI SDK in one of the newer frameworks.
24+
/// </remarks>
25+
[TestClass]
26+
[TestCategory("AAD")]
27+
public class TransmissionCredentialEnvelopeTests
28+
{
29+
private readonly Uri testUri = new Uri("https://127.0.0.1/");
30+
31+
[TestMethod]
32+
public async Task VerifyTransmissionSendAsync_Default()
33+
{
34+
var handler = new HandlerForFakeHttpClient
35+
{
36+
InnerHandler = new HttpClientHandler(),
37+
OnSendAsync = (req, cancellationToken) =>
38+
{
39+
// VALIDATE
40+
Assert.IsNull(req.Headers.Authorization);
41+
42+
return Task.FromResult<HttpResponseMessage>(new HttpResponseMessage());
43+
}
44+
};
45+
46+
using (var fakeHttpClient = new HttpClient(handler))
47+
{
48+
var expectedContentType = "content/type";
49+
var expectedContentEncoding = "contentEncoding";
50+
var items = new List<ITelemetry> { new EventTelemetry() };
51+
52+
// Instantiate Transmission with the mock HttpClient
53+
var transmission = new Transmission(testUri, new byte[] { 1, 2, 3, 4, 5 }, fakeHttpClient, expectedContentType, expectedContentEncoding);
54+
55+
var result = await transmission.SendAsync();
56+
}
57+
}
58+
59+
[TestMethod]
60+
public async Task VerifyTransmissionSendAsync_WithCredential_SetsAuthHeader()
61+
{
62+
var credendialEnvelope = new ReflectionCredentialEnvelope(new MockCredential());
63+
var token = credendialEnvelope.GetToken();
64+
65+
66+
var handler = new HandlerForFakeHttpClient
67+
{
68+
InnerHandler = new HttpClientHandler(),
69+
OnSendAsync = (req, cancellationToken) =>
70+
{
71+
// VALIDATE
72+
Assert.AreEqual(AuthConstants.AuthorizationTokenPrefix.Trim(), req.Headers.Authorization.Scheme);
73+
Assert.AreEqual(token, req.Headers.Authorization.Parameter);
74+
75+
return Task.FromResult<HttpResponseMessage>(new HttpResponseMessage());
76+
}
77+
};
78+
79+
using (var fakeHttpClient = new HttpClient(handler))
80+
{
81+
var expectedContentType = "content/type";
82+
var expectedContentEncoding = "contentEncoding";
83+
var items = new List<ITelemetry> { new EventTelemetry() };
84+
85+
// Instantiate Transmission with the mock HttpClient
86+
var transmission = new Transmission(testUri, new byte[] { 1, 2, 3, 4, 5 }, fakeHttpClient, expectedContentType, expectedContentEncoding);
87+
transmission.CredentialEnvelope = credendialEnvelope;
88+
89+
var result = await transmission.SendAsync();
90+
}
91+
}
92+
}
93+
}
94+
#endif

‎BASE/src/Microsoft.ApplicationInsights/Channel/InMemoryChannel.cs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
11
namespace Microsoft.ApplicationInsights.Channel
22
{
33
using System;
4+
using System.ComponentModel;
45
using System.Diagnostics;
56
using System.Threading;
67
using System.Threading.Tasks;
8+
79
using Microsoft.ApplicationInsights.Common;
10+
using Microsoft.ApplicationInsights.Extensibility;
11+
using Microsoft.ApplicationInsights.Extensibility.Implementation.Authentication;
812
using Microsoft.ApplicationInsights.Extensibility.Implementation.Tracing;
913

1014
/// <summary>
@@ -122,6 +126,20 @@ public int BacklogSize
122126
set { this.buffer.BacklogSize = value; }
123127
}
124128

129+
/// <summary>
130+
/// Gets or sets the <see cref="CredentialEnvelope"/> which is used for AAD.
131+
/// FOR INTERNAL USE. Customers should use <see cref="TelemetryConfiguration.SetAzureTokenCredential"/> instead.
132+
/// </summary>
133+
/// <remarks>
134+
/// <see cref="InMemoryChannel.CredentialEnvelope"/> sets <see cref="InMemoryTransmitter.CredentialEnvelope"/>
135+
/// which is used to set <see cref="Transmission.CredentialEnvelope"/> just before calling <see cref="Transmission.SendAsync"/>.
136+
/// </remarks>
137+
internal CredentialEnvelope CredentialEnvelope
138+
{
139+
get => this.transmitter.CredentialEnvelope;
140+
set => this.transmitter.CredentialEnvelope = value;
141+
}
142+
125143
internal bool IsDisposed => this.isDisposed;
126144

127145
/// <summary>

‎BASE/src/Microsoft.ApplicationInsights/Channel/InMemoryTransmitter.cs‎

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,11 @@ namespace Microsoft.ApplicationInsights.Channel
1111
using System.Net.Http;
1212
using System.Threading;
1313
using System.Threading.Tasks;
14+
1415
using Microsoft.ApplicationInsights.Common.Extensions;
1516
using Microsoft.ApplicationInsights.Extensibility;
1617
using Microsoft.ApplicationInsights.Extensibility.Implementation;
18+
using Microsoft.ApplicationInsights.Extensibility.Implementation.Authentication;
1719
using Microsoft.ApplicationInsights.Extensibility.Implementation.Tracing;
1820

1921
/// <summary>
@@ -32,13 +34,13 @@ internal class InMemoryTransmitter : IDisposable
3234
[SuppressMessage("Microsoft.Usage", "CA2213:DisposableFieldsShouldBeDisposed", Justification = "Object is disposed within the using statement of the " + nameof(Runner) + " method.")]
3335
private AutoResetEvent startRunnerEvent;
3436
private bool enabled = true;
35-
37+
3638
/// <summary>
3739
/// The number of times this object was disposed.
3840
/// </summary>
3941
private int disposeCount = 0;
4042
private TimeSpan sendingInterval = TimeSpan.FromSeconds(30);
41-
43+
4244
internal InMemoryTransmitter(TelemetryBuffer buffer)
4345
{
4446
this.buffer = buffer;
@@ -47,11 +49,11 @@ internal InMemoryTransmitter(TelemetryBuffer buffer)
4749
// Starting the Runner
4850
Task.Factory.StartNew(this.Runner, CancellationToken.None, TaskCreationOptions.LongRunning, TaskScheduler.Default)
4951
.ContinueWith(
50-
task =>
52+
task =>
5153
{
5254
string msg = string.Format(CultureInfo.InvariantCulture, "InMemoryTransmitter: Unhandled exception in Runner: {0}", task.Exception);
5355
CoreEventSource.Log.LogVerbose(msg);
54-
},
56+
},
5557
TaskContinuationOptions.OnlyOnFaulted);
5658
}
5759

@@ -63,6 +65,15 @@ internal TimeSpan SendingInterval
6365
set { this.sendingInterval = value; }
6466
}
6567

68+
/// <summary>
69+
/// Gets or sets the <see cref="CredentialEnvelope"/> which is used for AAD.
70+
/// </summary>
71+
/// <remarks>
72+
/// <see cref="InMemoryChannel.CredentialEnvelope"/> sets <see cref="InMemoryTransmitter.CredentialEnvelope"/>
73+
/// which is used to set <see cref="Transmission.CredentialEnvelope"/> just before calling <see cref="Transmission.SendAsync"/>.
74+
/// </remarks>
75+
internal CredentialEnvelope CredentialEnvelope { get; set; }
76+
6677
/// <summary>
6778
/// Performs application-defined tasks associated with freeing, releasing, or resetting unmanaged resources.
6879
/// </summary>
@@ -163,7 +174,7 @@ private Task Send(IEnumerable<ITelemetry> telemetryItems, TimeSpan timeout)
163174
}
164175

165176
var transmission = new Transmission(this.EndpointAddress, data, JsonSerializer.ContentType, JsonSerializer.CompressionType, timeout);
166-
177+
transmission.CredentialEnvelope = this.CredentialEnvelope;
167178
return transmission.SendAsync();
168179
}
169180

‎BASE/src/Microsoft.ApplicationInsights/Channel/Transmission.cs‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,9 @@
1010
using System.Net.Http.Headers;
1111
using System.Threading;
1212
using System.Threading.Tasks;
13+
1314
using Microsoft.ApplicationInsights.Extensibility.Implementation;
15+
using Microsoft.ApplicationInsights.Extensibility.Implementation.Authentication;
1416
using Microsoft.ApplicationInsights.Extensibility.Implementation.Tracing;
1517

1618
/// <summary>
@@ -141,6 +143,12 @@ public ICollection<ITelemetry> TelemetryItems
141143
get; private set;
142144
}
143145

146+
/// <summary>
147+
/// Gets or sets the <see cref="CredentialEnvelope"/> which is used for AAD.
148+
/// This is used include an AAD token on HTTP Requests sent to ingestion.
149+
/// </summary>
150+
internal CredentialEnvelope CredentialEnvelope { get; set; }
151+
144152
/// <summary>
145153
/// Gets the flush async id for the transmission.
146154
/// </summary>
@@ -404,6 +412,22 @@ protected virtual HttpRequestMessage CreateRequestMessage(Uri address, Stream co
404412
request.Content.Headers.Add(ContentEncodingHeader, this.ContentEncoding);
405413
}
406414

415+
if (this.CredentialEnvelope != null)
416+
{
417+
// TODO: NEED TO USE CACHING HERE
418+
var authToken = this.CredentialEnvelope.GetToken();
419+
420+
if (authToken == null)
421+
{
422+
// TODO: DO NOT SEND. RETURN FAILURE AND LET CHANNEL DECIDE WHEN TO RETRY.
423+
// This could be either a configuration error or the AAD service is unavailable.
424+
}
425+
else
426+
{
427+
request.Headers.TryAddWithoutValidation(AuthConstants.AuthorizationHeaderName, AuthConstants.AuthorizationTokenPrefix + authToken);
428+
}
429+
}
430+
407431
return request;
408432
}
409433

‎BASE/src/Microsoft.ApplicationInsights/Extensibility/Implementation/Authentication/AuthConstants.cs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,10 @@
22
{
33
internal static class AuthConstants
44
{
5+
public const string AuthorizationHeaderName = "Authorization";
6+
7+
public const string AuthorizationTokenPrefix = "Bearer ";
8+
59
/// <summary>
610
/// Source:
711
/// (https://docs.microsoft.com/azure/active-directory/develop/msal-acquire-cache-tokens#scopes-when-acquiring-tokens).

‎BASE/src/Microsoft.ApplicationInsights/Extensibility/Implementation/Authentication/CredentialEnvelope.cs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,19 @@ public abstract class CredentialEnvelope
1616
/// <summary>
1717
/// Gets an Azure.Core.AccessToken.
1818
/// </summary>
19+
/// <remarks>
20+
/// Whomever uses this MUST verify that it's called within <see cref="SdkInternalOperationsMonitor.Enter"/> otherwise dependency calls will be tracked.
21+
/// </remarks>
1922
/// <param name="cancellationToken">The System.Threading.CancellationToken to use.</param>
2023
/// <returns>A valid Azure.Core.AccessToken.</returns>
2124
public abstract string GetToken(CancellationToken cancellationToken = default);
2225

2326
/// <summary>
2427
/// Gets an Azure.Core.AccessToken.
2528
/// </summary>
29+
/// <remarks>
30+
/// Whomever uses this MUST verify that it's called within <see cref="SdkInternalOperationsMonitor.Enter"/> otherwise dependency calls will be tracked.
31+
/// </remarks>
2632
/// <param name="cancellationToken">The System.Threading.CancellationToken to use.</param>
2733
/// <returns>A valid Azure.Core.AccessToken.</returns>
2834
public abstract Task<string> GetTokenAsync(CancellationToken cancellationToken = default);

‎BASE/src/Microsoft.ApplicationInsights/Extensibility/Implementation/Authentication/ReflectionCredentialEnvelope.cs‎

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -46,11 +46,14 @@ public ReflectionCredentialEnvelope(object tokenCredential)
4646
/// <summary>
4747
/// Gets an Azure.Core.AccessToken.
4848
/// </summary>
49+
/// <remarks>
50+
/// Whomever uses this MUST verify that it's called within <see cref="SdkInternalOperationsMonitor.Enter"/> otherwise dependency calls will be tracked.
51+
/// </remarks>
4952
/// <param name="cancellationToken">The System.Threading.CancellationToken to use.</param>
5053
/// <returns>A valid Azure.Core.AccessToken.</returns>
5154
public override string GetToken(CancellationToken cancellationToken = default)
5255
{
53-
SdkInternalOperationsMonitor.Enter();
56+
// TODO: NEED TO FULLY TEST IF WE NEED TO CALL SdkInternalOperationsMonitor.Enter
5457
try
5558
{
5659
return AzureCore.InvokeGetToken(this.tokenCredential, this.tokenRequestContext, cancellationToken);
@@ -60,20 +63,19 @@ public override string GetToken(CancellationToken cancellationToken = default)
6063
CoreEventSource.Log.FailedToGetToken(ex.ToInvariantString());
6164
return null;
6265
}
63-
finally
64-
{
65-
SdkInternalOperationsMonitor.Exit();
66-
}
6766
}
6867

6968
/// <summary>
7069
/// Gets an Azure.Core.AccessToken.
7170
/// </summary>
71+
/// <remarks>
72+
/// Whomever uses this MUST verify that it's called within <see cref="SdkInternalOperationsMonitor.Enter"/> otherwise dependency calls will be tracked.
73+
/// </remarks>
7274
/// <param name="cancellationToken">The System.Threading.CancellationToken to use.</param>
7375
/// <returns>A valid Azure.Core.AccessToken.</returns>
7476
public override async Task<string> GetTokenAsync(CancellationToken cancellationToken = default)
7577
{
76-
SdkInternalOperationsMonitor.Enter();
78+
// TODO: NEED TO FULLY TEST IF WE NEED TO CALL SdkInternalOperationsMonitor.Enter
7779
try
7880
{
7981
return await AzureCore.InvokeGetTokenAsync(this.tokenCredential, this.tokenRequestContext, cancellationToken).ConfigureAwait(false);
@@ -83,10 +85,6 @@ public override async Task<string> GetTokenAsync(CancellationToken cancellationT
8385
CoreEventSource.Log.FailedToGetToken(ex.ToInvariantString());
8486
return null;
8587
}
86-
finally
87-
{
88-
SdkInternalOperationsMonitor.Exit();
89-
}
9088
}
9189

9290
/// <summary>

0 commit comments

Comments
 (0)