Repository navigation
fix(kafka): apply sts_header_overrides to MSK IAM SASL auth #7187
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
dlvenable
merged 1 commit into
opensearch-project:main
from
hozkaya2000:kafka-sts-header-overrides
Oct 13, 2026
+242
−17
Merged
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
87 changes: 87 additions & 0 deletions
87
...a/org/opensearch/dataprepper/plugins/kafka/util/MskIamAuthCredentialsCallbackHandler.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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<String, ?> configs, | ||
| final String saslMechanism, | ||
| final List<AppConfigurationEntry> 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()); | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
130 changes: 130 additions & 0 deletions
130
...g/opensearch/dataprepper/plugins/kafka/util/MskIamAuthCredentialsCallbackHandlerTest.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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<KafkaSecurityConfigurer> 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<KafkaSecurityConfigurer> 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<KafkaSecurityConfigurer> 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<KafkaSecurityConfigurer> 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})); | ||
| } | ||
| } |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I was curious how this even worked since I don't see this here.
We already have code to use the STS headers. They are just not used on communicating with the Kafka client directly.
data-prepper/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaSecurityConfigurer.java
Lines 324 to 334 in e48a6ff
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yep, that STS-header code is only used for the initial GetBootstrapBrokers call. The Kafka client's own SASL handshake goes through the vendored IAMClientCallbackHandler, which builds its own provider from the JAAS role with no header overrides — so the data-plane AssumeRole was missing the headers. This callback handler is what routes that handshake through the header-aware provider too.