Add azure_federated authentication to Kafka source

data-prepper-plugins/kafka-plugins/README.md

@@ -6,6 +6,8 @@ This source allows Data Prepper to use Kafka as source. This source reads record
 
 For usage and configuration, please refer to the documentation [here] (https://opensearch.org/docs/2.9/data-prepper/pipelines/configuration/sources/sources/kafka-source).
 
+The `azure_federated` SASL/OAUTHBEARER mechanism additionally lets the Kafka source consume from an Azure
+Event Hubs Kafka endpoint using AWS to Azure Entra workload-identity federation, with no stored secret.
 
 ## Developer guide
 

data-prepper-plugins/kafka-plugins/build.gradle

@@ -62,6 +62,7 @@ dependencies {
     testImplementation 'org.yaml:snakeyaml:2.2'
     testImplementation testLibs.spring.test
     testImplementation 'com.fasterxml.jackson.datatype:jackson-datatype-jsr310'
+    testImplementation 'org.hibernate.validator:hibernate-validator:8.0.2.Final'
     testImplementation project(':data-prepper-test:test-common')
     testImplementation project(':data-prepper-plugins:blocking-buffer')
     testImplementation project(':data-prepper-core')
@@ -143,6 +144,14 @@ task integrationTest(type: Test) {
     systemProperty 'tests.kafka.authconfig.password', System.getProperty('tests.kafka.authconfig.password')
     systemProperty 'tests.kafka.authconfig.mechanism', System.getProperty('tests.kafka.authconfig.mechanism')
     systemProperty 'tests.kafka.kms_key', System.getProperty('tests.kafka.kms_key')
+    systemProperty 'tests.azure.eventhubs.bootstrap_servers', System.getProperty('tests.azure.eventhubs.bootstrap_servers')
+    systemProperty 'tests.azure.eventhubs.topic', System.getProperty('tests.azure.eventhubs.topic')
+    systemProperty 'tests.azure.tenant_id', System.getProperty('tests.azure.tenant_id')
+    systemProperty 'tests.azure.client_id', System.getProperty('tests.azure.client_id')
+    systemProperty 'tests.azure.token_endpoint', System.getProperty('tests.azure.token_endpoint')
+    systemProperty 'tests.azure.scope', System.getProperty('tests.azure.scope')
+    systemProperty 'tests.aws.sts_role_arn', System.getProperty('tests.aws.sts_role_arn')
+    systemProperty 'tests.aws.region', System.getProperty('tests.aws.region')
 
     filter {
         includeTestsMatching '*IT'

data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/admin/KafkaAdminAccessor.java

@@ -31,7 +31,8 @@ public class KafkaAdminAccessor {
 
     public KafkaAdminAccessor(final KafkaClusterAuthConfig kafkaClusterAuthConfig, final List<String> consumerGroupIds) {
         Properties authProperties = new Properties();
-        KafkaSecurityConfigurer.setAuthProperties(authProperties, kafkaClusterAuthConfig, LOG);
+        // no AwsCredentialsSupplier on the admin path
+        KafkaSecurityConfigurer.setAuthProperties(authProperties, kafkaClusterAuthConfig, null, LOG);
         this.kafkaAdminClient = KafkaAdminClient.create(authProperties);
         this.topicEmptinessMetadata = new TopicEmptinessMetadata();
         this.consumerGroupIds = consumerGroupIds;

data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AwsCredentialsSupplierProvider.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ */
+
+package org.opensearch.dataprepper.plugins.kafka.authenticator;
+
+import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier;
+
+public class AwsCredentialsSupplierProvider {
+    private static final AwsCredentialsSupplierProvider singleton = new AwsCredentialsSupplierProvider();
+
+    public static AwsCredentialsSupplierProvider getInstance() {
+        return singleton;
+    }
+
+    // set once before Kafka client construction; volatile suffices for cross-thread visibility
+    private volatile AwsCredentialsSupplier awsCredentialsSupplier;
+
+    // Used for testing only
+    protected AwsCredentialsSupplierProvider() {}
+
+    public AwsCredentialsSupplier getAwsCredentialsSupplier() {
+        return awsCredentialsSupplier;
+    }
+
+    public void set(final AwsCredentialsSupplier awsCredentialsSupplier) {
+        this.awsCredentialsSupplier = awsCredentialsSupplier;
+    }
+}

data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedCallbackHandler.java

@@ -0,0 +1,97 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ */
+
+package org.opensearch.dataprepper.plugins.kafka.authenticator;
+
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerTokenCallback;
+import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier;
+
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.AppConfigurationEntry;
+import java.io.IOException;
+import java.util.Base64;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+public class AzureFederatedCallbackHandler implements AuthenticateCallbackHandler {
+    public static final String OPT_REGION = "azureFederatedRegion";
+    public static final String OPT_STS_ROLE_ARN = "azureFederatedStsRoleArn";
+    public static final String OPT_TOKEN_ENDPOINT = "azureFederatedTokenEndpoint";
+    public static final String OPT_CLIENT_ID = "azureFederatedClientId";
+    public static final String OPT_SCOPE = "azureFederatedScope";
+    public static final String OPT_STS_HEADER_OVERRIDES = "azureFederatedStsHeaderOverrides";
+
+    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+    private AzureFederatedTokenProvider tokenProvider;
+
+    public AzureFederatedCallbackHandler() {
+    }
+
+    AzureFederatedCallbackHandler(final AzureFederatedTokenProvider tokenProvider) {
+        this.tokenProvider = tokenProvider;
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public void configure(final Map<String, ?> configs, final String saslMechanism,
+                          final List<AppConfigurationEntry> jaasConfigEntries) {
+        if (!OAuthBearerLoginModule.OAUTHBEARER_MECHANISM.equals(saslMechanism)) {
+            throw new IllegalArgumentException("Unexpected SASL mechanism: " + saslMechanism);
+        }
+        final Map<String, String> options = (Map<String, String>) jaasConfigEntries.get(0).getOptions();
+        final Map<String, String> stsHeaderOverrides = decodeStsHeaderOverrides(options.get(OPT_STS_HEADER_OVERRIDES));
+        final AwsCredentialsSupplier awsCredentialsSupplier =
+                AwsCredentialsSupplierProvider.getInstance().getAwsCredentialsSupplier();
+        this.tokenProvider = new AzureFederatedTokenProvider(
+                options.get(OPT_REGION), options.get(OPT_STS_ROLE_ARN), options.get(OPT_TOKEN_ENDPOINT),
+                options.get(OPT_CLIENT_ID), options.get(OPT_SCOPE), stsHeaderOverrides, awsCredentialsSupplier);
+    }
+
+    private static Map<String, String> decodeStsHeaderOverrides(final String encoded) {
+        if (encoded == null || encoded.isEmpty()) {
+            return Collections.emptyMap();
+        }
+        try {
+            final byte[] json = Base64.getDecoder().decode(encoded);
+            return OBJECT_MAPPER.readValue(json, new TypeReference<Map<String, String>>() {});
+        } catch (final IOException e) {
+            throw new RuntimeException("Failed to decode azure_federated STS header overrides", e);
+        }
+    }
+
+    @Override
+    public void handle(final Callback[] callbacks) throws UnsupportedCallbackException {
+        for (final Callback callback : callbacks) {
+            if (callback instanceof OAuthBearerTokenCallback) {
+                ((OAuthBearerTokenCallback) callback).token(tokenProvider.getToken());
+            } else {
+                throw new UnsupportedCallbackException(callback);
+            }
+        }
+    }
+
+    @Override
+    public void close() {
+        if (tokenProvider != null) {
+            tokenProvider.close();
+        }
+    }
+
+    AzureFederatedTokenProvider getTokenProvider() {
+        return tokenProvider;
+    }
+}

data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedOAuthBearerToken.java

@@ -0,0 +1,57 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ */
+
+package org.opensearch.dataprepper.plugins.kafka.authenticator;
+
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerToken;
+
+import java.util.Set;
+
+public class AzureFederatedOAuthBearerToken implements OAuthBearerToken {
+    private final String value;
+    private final long lifetimeMs;
+    private final long startTimeMs;
+    private final String scope;
+    private final String principalName;
+
+    public AzureFederatedOAuthBearerToken(final String value, final long expiresInSeconds,
+                                          final String scope, final String principalName) {
+        this.value = value;
+        this.startTimeMs = System.currentTimeMillis();
+        this.lifetimeMs = this.startTimeMs + (expiresInSeconds * 1000L);
+        this.scope = scope;
+        this.principalName = principalName;
+    }
+
+    @Override
+    public String value() {
+        return value;
+    }
+
+    @Override
+    public long lifetimeMs() {
+        return lifetimeMs;
+    }
+
+    @Override
+    public Long startTimeMs() {
+        return startTimeMs;
+    }
+
+    @Override
+    public String principalName() {
+        return principalName;
+    }
+
+    @Override
+    public Set<String> scope() {
+        return scope == null ? Set.of() : Set.of(scope);
+    }
+}