From f33698e7b4aa53c1a40f667de8408c6c44bdd97c Mon Sep 17 00:00:00 2001 From: Halit Ozkaya Date: Fri, 9 Oct 2026 01:01:44 +0000 Subject: [PATCH] fix(kafka): apply sts_header_overrides to MSK IAM SASL auth Route role-based AWS_MSK_IAM SASL (data-plane) authentication through a dedicated callback handler that sources credentials exclusively from Data Prepper's header-aware STS provider, so aws.sts_header_overrides (x-amz-source-arn / x-amz-source-account) are applied on the data-plane sts:AssumeRole used for broker auth. The handler implements AuthenticateCallbackHandler directly (never extends IAMClientCallbackHandler), so the vendored library never assembles its own header-less MSKCredentialProvider. The ROLE path is now unconditional: credentials always flow through our provider regardless of whether overrides are configured. Resolves #6923 Signed-off-by: Halit Ozkaya --- .../kafka/util/KafkaSecurityConfigurer.java | 18 +-- .../MskIamAuthCredentialsCallbackHandler.java | 87 ++++++++++++ .../util/KafkaSecurityConfigurerTest.java | 24 +++- ...IamAuthCredentialsCallbackHandlerTest.java | 130 ++++++++++++++++++ 4 files changed, 242 insertions(+), 17 deletions(-) create mode 100644 data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/MskIamAuthCredentialsCallbackHandler.java create mode 100644 data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/MskIamAuthCredentialsCallbackHandlerTest.java diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaSecurityConfigurer.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaSecurityConfigurer.java index 5428c92dc8..2c3b7059e8 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaSecurityConfigurer.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaSecurityConfigurer.java @@ -115,6 +115,10 @@ public class KafkaSecurityConfigurer { private static AwsCredentialsProvider awsGlueCredentialsProvider; private static GlueSchemaRegistryKafkaDeserializer glueDeserializer; + static AwsCredentialsProvider getMskCredentialsProvider() { + return mskCredentialsProvider; + } + /*public static void setSaslPlainTextProperties(final KafkaSourceConfig kafkaSourConfig, final Properties properties) { @@ -245,18 +249,8 @@ public static void setAwsIamAuthProperties(Properties properties, final AwsIamAu if (Objects.isNull(awsConfig)) { throw new RuntimeException("AWS Config needs to be specified when sasl/aws_msk_iam is set to \"role\""); } - String baseIamAuthConfig = "software.amazon.msk.auth.iam.IAMLoginModule required " + - "awsRoleArn=\"%s\" " + - "awsStsRegion=\"%s\""; - - baseIamAuthConfig = String.format(baseIamAuthConfig, awsConfig.getStsRoleArn(), awsConfig.getRegion()); - - if (Objects.nonNull(awsConfig.getStsRoleSessionName())) { - baseIamAuthConfig += String.format(" awsRoleSessionName=\"%s\"", awsConfig.getStsRoleSessionName()); - } - - baseIamAuthConfig += ";"; - properties.put(SASL_JAAS_CONFIG, baseIamAuthConfig); + properties.put(SASL_CLIENT_CALLBACK_HANDLER_CLASS, MskIamAuthCredentialsCallbackHandler.class.getName()); + properties.put(SASL_JAAS_CONFIG, "software.amazon.msk.auth.iam.IAMLoginModule required;"); } else if (awsIamAuthConfig == AwsIamAuthConfig.DEFAULT) { properties.put(SASL_JAAS_CONFIG, "software.amazon.msk.auth.iam.IAMLoginModule required;"); diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/MskIamAuthCredentialsCallbackHandler.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/MskIamAuthCredentialsCallbackHandler.java new file mode 100644 index 0000000000..29612db9d5 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/MskIamAuthCredentialsCallbackHandler.java @@ -0,0 +1,87 @@ +/* + * 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.util; + +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.auth.BasicAWSCredentials; +import com.amazonaws.auth.BasicSessionCredentials; +import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.awssdk.auth.credentials.AwsCredentials; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.auth.credentials.AwsSessionCredentials; +import software.amazon.msk.auth.iam.IAMLoginModule; +import software.amazon.msk.auth.iam.internals.AWSCredentialsCallback; + +import javax.security.auth.callback.Callback; +import javax.security.auth.callback.UnsupportedCallbackException; +import javax.security.auth.login.AppConfigurationEntry; +import java.util.List; +import java.util.Map; + +/** + * SASL client callback handler for the {@code AWS_MSK_IAM} mechanism that sources credentials + * exclusively from Data Prepper's own header-aware STS provider, so the configured + * {@code aws.sts_header_overrides} are applied on the data-plane {@code sts:AssumeRole} used for + * broker authentication. + */ +public class MskIamAuthCredentialsCallbackHandler implements AuthenticateCallbackHandler { + private static final Logger LOG = LoggerFactory.getLogger(MskIamAuthCredentialsCallbackHandler.class); + + @Override + public void configure(final Map configs, + final String saslMechanism, + final List jaasConfigEntries) { + if (!IAMLoginModule.MECHANISM.equals(saslMechanism)) { + throw new IllegalArgumentException("Unexpected SASL mechanism: " + saslMechanism); + } + } + + @Override + public void handle(final Callback[] callbacks) throws UnsupportedCallbackException { + for (final Callback callback : callbacks) { + if (callback instanceof AWSCredentialsCallback) { + handleCallback((AWSCredentialsCallback) callback); + } else { + throw new UnsupportedCallbackException(callback, + "Unsupported callback type: " + callback.getClass().getName()); + } + } + } + + private void handleCallback(final AWSCredentialsCallback callback) { + final AwsCredentialsProvider provider = KafkaSecurityConfigurer.getMskCredentialsProvider(); + try { + LOG.debug("Resolving MSK IAM credentials from header-aware provider {}", + provider == null ? "null" : provider.getClass().getSimpleName()); + callback.setAwsCredentials(toV1Credentials(provider.resolveCredentials())); + } catch (final Exception e) { + LOG.debug("Failed to resolve MSK IAM credentials with STS header overrides", e); + callback.setLoadingException(e); + } + } + + @Override + public void close() { + } + + private static AWSCredentials toV1Credentials(final AwsCredentials credentials) { + if (credentials instanceof AwsSessionCredentials) { + final AwsSessionCredentials sessionCredentials = (AwsSessionCredentials) credentials; + return new BasicSessionCredentials( + sessionCredentials.accessKeyId(), + sessionCredentials.secretAccessKey(), + sessionCredentials.sessionToken()); + } + return new BasicAWSCredentials(credentials.accessKeyId(), credentials.secretAccessKey()); + } +} diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaSecurityConfigurerTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaSecurityConfigurerTest.java index 4d95953bb1..eaa86feece 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaSecurityConfigurerTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaSecurityConfigurerTest.java @@ -182,15 +182,29 @@ public void testSetAuthPropertiesBootstrapServersWithSaslIAMRole() throws IOExce assertThat(props.getProperty("bootstrap.servers"), is("localhost:9092")); assertThat(props.getProperty("sasl.mechanism"), is("AWS_MSK_IAM")); assertThat(props.getProperty("sasl.jaas.config"), - is("software.amazon.msk.auth.iam.IAMLoginModule required " + - "awsRoleArn=\"test_sasl_iam_sts_role\" awsStsRegion=\"us-east-2\";")); + is("software.amazon.msk.auth.iam.IAMLoginModule required;")); assertThat(props.getProperty("security.protocol"), is("SASL_SSL")); assertThat(props.getProperty("certificateContent"), is(nullValue())); assertThat(props.getProperty("ssl.truststore.location"), is(nullValue())); assertThat(props.getProperty("ssl.truststore.password"), is(nullValue())); assertThat(props.get("ssl.engine.factory.class"), is(nullValue())); assertThat(props.get("sasl.client.callback.handler.class"), - is("software.amazon.msk.auth.iam.IAMClientCallbackHandler")); + is(MskIamAuthCredentialsCallbackHandler.class.getName())); + } + + @Test + public void testSetAuthPropertiesBootstrapServersWithSaslIAMRoleAndStsHeaderOverrides() throws IOException { + final Properties props = new Properties(); + final KafkaSourceConfig kafkaSourceConfig = + createKafkaSinkConfig("kafka-pipeline-bootstrap-servers-sasl-iam-role-with-headers.yaml"); + KafkaSecurityConfigurer.setAuthProperties(props, kafkaSourceConfig, null, LOG); + assertThat(props.getProperty("bootstrap.servers"), is("localhost:9092")); + assertThat(props.getProperty("sasl.mechanism"), is("AWS_MSK_IAM")); + assertThat(props.getProperty("security.protocol"), is("SASL_SSL")); + assertThat(props.getProperty("sasl.jaas.config"), + is("software.amazon.msk.auth.iam.IAMLoginModule required;")); + assertThat(props.get("sasl.client.callback.handler.class"), + is(MskIamAuthCredentialsCallbackHandler.class.getName())); } @Test @@ -230,14 +244,14 @@ public void testSetAuthPropertiesBootstrapServersOverrideByMSK() throws IOExcept assertThat(props.getProperty("bootstrap.servers"), is(testMSKEndpoint)); assertThat(props.getProperty("sasl.mechanism"), is("AWS_MSK_IAM")); assertThat(props.getProperty("sasl.jaas.config"), - is("software.amazon.msk.auth.iam.IAMLoginModule required awsRoleArn=\"sts_role_arn\" awsStsRegion=\"us-east-2\";")); + is("software.amazon.msk.auth.iam.IAMLoginModule required;")); assertThat(props.getProperty("security.protocol"), is("SASL_SSL")); assertThat(props.getProperty("certificateContent"), is(nullValue())); assertThat(props.getProperty("ssl.truststore.location"), is(nullValue())); assertThat(props.getProperty("ssl.truststore.password"), is(nullValue())); assertThat(props.get("ssl.engine.factory.class"), is(nullValue())); assertThat(props.get("sasl.client.callback.handler.class"), - is("software.amazon.msk.auth.iam.IAMClientCallbackHandler")); + is(MskIamAuthCredentialsCallbackHandler.class.getName())); } @Test diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/MskIamAuthCredentialsCallbackHandlerTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/MskIamAuthCredentialsCallbackHandlerTest.java new file mode 100644 index 0000000000..3233d08304 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/MskIamAuthCredentialsCallbackHandlerTest.java @@ -0,0 +1,130 @@ +/* + * 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.util; + +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.auth.BasicAWSCredentials; +import com.amazonaws.auth.BasicSessionCredentials; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.MockedStatic; +import org.mockito.junit.jupiter.MockitoExtension; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.auth.credentials.AwsSessionCredentials; +import software.amazon.msk.auth.iam.internals.AWSCredentialsCallback; + +import javax.security.auth.callback.Callback; +import javax.security.auth.callback.UnsupportedCallbackException; +import java.util.Collections; + +import static org.hamcrest.CoreMatchers.instanceOf; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.sameInstance; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +class MskIamAuthCredentialsCallbackHandlerTest { + + @Mock + private AwsCredentialsProvider mskCredentialsProvider; + + private final MskIamAuthCredentialsCallbackHandler objectUnderTest = new MskIamAuthCredentialsCallbackHandler(); + + @Test + void handleResolvesSessionCredentialsAndConvertsToV1() throws Exception { + final AwsSessionCredentials sessionCredentials = AwsSessionCredentials.create( + "accessKeyId", "secretAccessKey", "sessionToken"); + when(mskCredentialsProvider.resolveCredentials()).thenReturn(sessionCredentials); + final AWSCredentialsCallback callback = new AWSCredentialsCallback(); + + try (final MockedStatic mockedConfigurer = mockStatic(KafkaSecurityConfigurer.class)) { + mockedConfigurer.when(KafkaSecurityConfigurer::getMskCredentialsProvider).thenReturn(mskCredentialsProvider); + objectUnderTest.handle(new Callback[]{callback}); + } + + assertThat(callback.isSuccessful(), is(true)); + final AWSCredentials credentials = callback.getAwsCredentials(); + assertThat(credentials, instanceOf(BasicSessionCredentials.class)); + assertThat(credentials.getAWSAccessKeyId(), is("accessKeyId")); + assertThat(credentials.getAWSSecretKey(), is("secretAccessKey")); + assertThat(((BasicSessionCredentials) credentials).getSessionToken(), is("sessionToken")); + } + + @Test + void handleResolvesNonSessionCredentialsAndConvertsToV1() throws Exception { + final AwsBasicCredentials basicCredentials = AwsBasicCredentials.create("accessKeyId", "secretAccessKey"); + when(mskCredentialsProvider.resolveCredentials()).thenReturn(basicCredentials); + final AWSCredentialsCallback callback = new AWSCredentialsCallback(); + + try (final MockedStatic mockedConfigurer = mockStatic(KafkaSecurityConfigurer.class)) { + mockedConfigurer.when(KafkaSecurityConfigurer::getMskCredentialsProvider).thenReturn(mskCredentialsProvider); + objectUnderTest.handle(new Callback[]{callback}); + } + + assertThat(callback.isSuccessful(), is(true)); + final AWSCredentials credentials = callback.getAwsCredentials(); + assertThat(credentials, instanceOf(BasicAWSCredentials.class)); + assertThat(credentials.getAWSAccessKeyId(), is("accessKeyId")); + assertThat(credentials.getAWSSecretKey(), is("secretAccessKey")); + } + + @Test + void handleSetsLoadingExceptionWhenProviderThrows() throws Exception { + final RuntimeException resolveException = new RuntimeException("Access denied"); + when(mskCredentialsProvider.resolveCredentials()).thenThrow(resolveException); + final AWSCredentialsCallback callback = new AWSCredentialsCallback(); + + try (final MockedStatic mockedConfigurer = mockStatic(KafkaSecurityConfigurer.class)) { + mockedConfigurer.when(KafkaSecurityConfigurer::getMskCredentialsProvider).thenReturn(mskCredentialsProvider); + objectUnderTest.handle(new Callback[]{callback}); + } + + assertThat(callback.isSuccessful(), is(false)); + assertThat(callback.getLoadingException(), is(sameInstance(resolveException))); + } + + @Test + void handleSetsLoadingExceptionWhenProviderIsNull() throws Exception { + final AWSCredentialsCallback callback = new AWSCredentialsCallback(); + + try (final MockedStatic mockedConfigurer = mockStatic(KafkaSecurityConfigurer.class)) { + mockedConfigurer.when(KafkaSecurityConfigurer::getMskCredentialsProvider).thenReturn(null); + objectUnderTest.handle(new Callback[]{callback}); + } + + assertThat(callback.isSuccessful(), is(false)); + assertThat(callback.getLoadingException(), instanceOf(NullPointerException.class)); + } + + @Test + void configureAcceptsTheMskIamMechanismAndReadsNoJaasOptions() { + assertDoesNotThrow(() -> objectUnderTest.configure( + Collections.emptyMap(), "AWS_MSK_IAM", Collections.emptyList())); + } + + @Test + void configureRejectsAnUnexpectedMechanism() { + assertThrows(IllegalArgumentException.class, () -> objectUnderTest.configure( + Collections.emptyMap(), "SCRAM-SHA-512", Collections.emptyList())); + } + + @Test + void handleRejectsAnUnsupportedCallbackType() { + final Callback unsupported = new Callback() { }; + assertThrows(UnsupportedCallbackException.class, + () -> objectUnderTest.handle(new Callback[]{unsupported})); + } +}