Support for AAD Audience in connection string (for sovereign cloud) (#4121)

Co-authored-by: Trask Stalnaker <[email protected]>

Author: harsimar

.github/workflows/build-common.yml

@@ -8,7 +8,7 @@ on:
         required: false
 
 env:
-  EXPORTER_VERSION: 1.0.0-beta.1 # to be updated with the latest version
+  EXPORTER_VERSION: 1.1.0 # to be updated with the latest version
 
 jobs:
   spotless:

CHANGELOG.md

@@ -1,4 +1,6 @@
 # CHANGELOG
+## Version 3.7.2 GA (Unreleased)
+* Support for using the AAD Audience from the connection string ([#4121](https://github.com/microsoft/ApplicationInsights-Java/pull/4121))
 
 ## Version 3.7.1 GA (02/26/2025)
 

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/httpclient/LazyHttpClient.java

@@ -8,6 +8,10 @@
 import com.azure.core.http.HttpClient;
 import com.azure.core.http.HttpPipeline;
 import com.azure.core.http.HttpPipelineBuilder;
+import com.azure.core.http.HttpPipelineCallContext;
+import com.azure.core.http.HttpPipelineNextPolicy;
+import com.azure.core.http.HttpPipelineNextSyncPolicy;
+import com.azure.core.http.HttpPipelinePosition;
 import com.azure.core.http.HttpRequest;
 import com.azure.core.http.HttpResponse;
 import com.azure.core.http.ProxyOptions;
@@ -31,15 +35,13 @@
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 import reactor.core.publisher.Mono;
 import reactor.netty.resources.LoopResources;
 
 public class LazyHttpClient implements HttpClient {
 
-  private static final String APPLICATIONINSIGHTS_AUTHENTICATION_SCOPE =
-      "https://monitor.azure.com//.default";
-
   private static final HttpClient INSTANCE = new LazyHttpClient();
 
   public static final CountDownLatch safeToInitLatch = new CountDownLatch(1);
@@ -113,16 +115,21 @@ private static HttpClient init() {
   }
 
   public static HttpPipeline newHttpPipeLineWithDefaultRedirect(
-      @Nullable Configuration.AadAuthentication aadConfiguration) {
-    return newHttpPipeLine(aadConfiguration, new RedirectPolicy(new DefaultRedirectStrategy()));
+      @Nullable Configuration.AadAuthentication aadConfiguration,
+      Supplier<String> aadAudienceWithScope) {
+    return newHttpPipeLine(
+        aadConfiguration, aadAudienceWithScope, new RedirectPolicy(new DefaultRedirectStrategy()));
   }
 
   public static HttpPipeline newHttpPipeLine(
       @Nullable Configuration.AadAuthentication aadConfiguration,
+      Supplier<String> aadAudienceWithScope,
       HttpPipelinePolicy... additionalPolicies) {
     List<HttpPipelinePolicy> policies = new ArrayList<>();
     if (aadConfiguration != null && aadConfiguration.enabled) {
-      policies.add(getAuthenticationPolicy(aadConfiguration));
+      policies.add(
+          new LazyHttpPipelinePolicy(
+              () -> getAuthenticationPolicy(aadConfiguration, aadAudienceWithScope.get())));
     }
     policies.addAll(asList(additionalPolicies));
     // Add Logging Policy. Can be enabled using AZURE_LOG_LEVEL.
@@ -144,31 +151,31 @@ public Mono<HttpResponse> send(HttpRequest request, Context context) {
   }
 
   private static HttpPipelinePolicy getAuthenticationPolicy(
-      Configuration.AadAuthentication configuration) {
+      Configuration.AadAuthentication configuration, String aadAudienceWithScope) {
     switch (configuration.type) {
       case UAMI:
-        return getAuthenticationPolicyWithUami(configuration);
+        return getAuthenticationPolicyWithUami(configuration, aadAudienceWithScope);
       case SAMI:
-        return getAuthenticationPolicyWithSami();
+        return getAuthenticationPolicyWithSami(aadAudienceWithScope);
       case VSCODE:
-        return getAuthenticationPolicyWithVsCode();
+        return getAuthenticationPolicyWithVsCode(aadAudienceWithScope);
       case CLIENTSECRET:
-        return getAuthenticationPolicyWithClientSecret(configuration);
+        return getAuthenticationPolicyWithClientSecret(configuration, aadAudienceWithScope);
     }
     throw new IllegalStateException(
         "Invalid Authentication Type used in AAD Authentication: " + configuration.type);
   }
 
   private static HttpPipelinePolicy getAuthenticationPolicyWithUami(
-      Configuration.AadAuthentication configuration) {
+      Configuration.AadAuthentication configuration, String aadAudienceWithScope) {
     ManagedIdentityCredentialBuilder managedIdentityCredential =
         new ManagedIdentityCredentialBuilder().clientId(configuration.clientId);
     return new BearerTokenAuthenticationPolicy(
-        managedIdentityCredential.build(), APPLICATIONINSIGHTS_AUTHENTICATION_SCOPE);
+        managedIdentityCredential.build(), aadAudienceWithScope);
   }
 
   private static HttpPipelinePolicy getAuthenticationPolicyWithClientSecret(
-      Configuration.AadAuthentication configuration) {
+      Configuration.AadAuthentication configuration, String aadAudienceWithScope) {
     ClientSecretCredentialBuilder credential =
         new ClientSecretCredentialBuilder()
             .tenantId(configuration.tenantId)
@@ -177,21 +184,54 @@ private static HttpPipelinePolicy getAuthenticationPolicyWithClientSecret(
     if (configuration.authorityHost != null) {
       credential.authorityHost(configuration.authorityHost);
     }
-    return new BearerTokenAuthenticationPolicy(
-        credential.build(), APPLICATIONINSIGHTS_AUTHENTICATION_SCOPE);
+    return new BearerTokenAuthenticationPolicy(credential.build(), aadAudienceWithScope);
   }
 
-  private static HttpPipelinePolicy getAuthenticationPolicyWithVsCode() {
+  private static HttpPipelinePolicy getAuthenticationPolicyWithVsCode(String aadAudienceWithScope) {
     VisualStudioCodeCredential visualStudioCodeCredential =
         new VisualStudioCodeCredentialBuilder().build();
-    return new BearerTokenAuthenticationPolicy(
-        visualStudioCodeCredential, APPLICATIONINSIGHTS_AUTHENTICATION_SCOPE);
+    return new BearerTokenAuthenticationPolicy(visualStudioCodeCredential, aadAudienceWithScope);
   }
 
-  private static HttpPipelinePolicy getAuthenticationPolicyWithSami() {
+  private static HttpPipelinePolicy getAuthenticationPolicyWithSami(String aadAudienceWithScope) {
     ManagedIdentityCredential managedIdentityCredential =
         new ManagedIdentityCredentialBuilder().build();
-    return new BearerTokenAuthenticationPolicy(
-        managedIdentityCredential, APPLICATIONINSIGHTS_AUTHENTICATION_SCOPE);
+    return new BearerTokenAuthenticationPolicy(managedIdentityCredential, aadAudienceWithScope);
+  }
+
+  private static class LazyHttpPipelinePolicy implements HttpPipelinePolicy {
+
+    private final Supplier<HttpPipelinePolicy> supplier;
+    private volatile HttpPipelinePolicy delegate;
+
+    LazyHttpPipelinePolicy(Supplier<HttpPipelinePolicy> supplier) {
+      this.supplier = supplier;
+    }
+
+    @Override
+    public Mono<HttpResponse> process(
+        HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
+      createDelegateFirstTime();
+      return delegate.process(context, next);
+    }
+
+    @Override
+    public HttpResponse processSync(
+        HttpPipelineCallContext context, HttpPipelineNextSyncPolicy next) {
+      createDelegateFirstTime();
+      return delegate.processSync(context, next);
+    }
+
+    @Override
+    public HttpPipelinePosition getPipelinePosition() {
+      createDelegateFirstTime();
+      return delegate.getPipelinePosition();
+    }
+
+    private void createDelegateFirstTime() {
+      if (delegate == null) {
+        delegate = supplier.get();
+      }
+    }
   }
 }

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/init/SecondEntryPoint.java

@@ -201,7 +201,9 @@ public void customize(AutoConfigurationCustomizer autoConfiguration) {
     if (telemetryClient.getConnectionString() != null) {
       statsbeatModule.start(
           AzureMonitorHelper.createStatsbeatTelemetryItemExporter(
-              LazyHttpClient.newHttpPipeLine(null), statsbeatModule, tempDir),
+              LazyHttpClient.newHttpPipeLine(null, telemetryClient::getAadAudienceWithScope),
+              statsbeatModule,
+              tempDir),
           telemetryClient::getStatsbeatConnectionString,
           telemetryClient::getInstrumentationKey,
           configuration.internal.statsbeat.disabledAll,
@@ -224,7 +226,8 @@ public void customize(AutoConfigurationCustomizer autoConfiguration) {
     if (configuration.preview.liveMetrics.enabled) {
       quickPulse =
           QuickPulse.create(
-              LazyHttpClient.newHttpPipeLineWithDefaultRedirect(configuration.authentication),
+              LazyHttpClient.newHttpPipeLineWithDefaultRedirect(
+                  configuration.authentication, telemetryClient::getAadAudienceWithScope),
               () -> {
                 ConnectionString connectionString = telemetryClient.getConnectionString();
                 return connectionString == null ? null : connectionString.getLiveEndpoint();

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/profiler/ProfilingInitializer.java

@@ -121,6 +121,7 @@ private synchronized void performInit() {
     httpPipeline =
         LazyHttpClient.newHttpPipeLine(
             telemetryClient.getAadAuthentication(),
+            telemetryClient::getAadAudienceWithScope,
             new RedirectPolicy(
                 new DefaultRedirectStrategy(
                     3,

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/telemetry/TelemetryClient.java

@@ -223,10 +223,10 @@ public BatchItemProcessor getMetricsBatchItemProcessor() {
 
   private BatchItemProcessor initBatchItemProcessor(
       int exportQueueCapacity, int maxExportBatchSize, String queueName) {
-
     HttpPipeline httpPipeline =
         LazyHttpClient.newHttpPipeLine(
             aadAuthentication,
+            this::getAadAudienceWithScope,
             new NetworkStatsbeatHttpPipelinePolicy(statsbeatModule.getNetworkStatsbeat()));
     // TODO (heya) refactor the following by using AzureMonitorHelper.createTelemetryItemExporter by
     // passing in getNonessentialStatsbeat
@@ -353,6 +353,11 @@ public ConnectionString getConnectionString() {
     return connectionString;
   }
 
+  @Nullable
+  public String getAadAudienceWithScope() {
+    return connectionString.getAadAudienceWithScope();
+  }
+
   @Nullable
   public String getRoleName() {
     return roleName;