This is an automated email from the ASF dual-hosted git repository.
exceptionfactory pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/nifi.git
The following commit(s) were added to refs/heads/main by this push:
new df1f9cd2573 NIFI-16367 Add Azure Entra Database Password Provider
(#11693)
df1f9cd2573 is described below
commit df1f9cd25737d02ebe0e667d0aca127afc742257
Author: Pierre Villard <[email protected]>
AuthorDate: Fri Sep 18 22:04:25 2026 +0200
NIFI-16367 Add Azure Entra Database Password Provider (#11693)
Signed-off-by: David Handermann <[email protected]>
---
.../nifi-azure-processors/pom.xml | 6 +
.../azure/AzureEntraDatabasePasswordProvider.java | 215 +++++++
.../org.apache.nifi.controller.ControllerService | 1 +
.../additionalDetails.md | 92 +++
.../AzureEntraDatabasePasswordProviderTest.java | 616 +++++++++++++++++++++
5 files changed, 930 insertions(+)
diff --git
a/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/pom.xml
b/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/pom.xml
index 724dda1c360..c4145b66d79 100644
--- a/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/pom.xml
+++ b/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/pom.xml
@@ -32,6 +32,12 @@
<artifactId>nifi-oauth2-provider-api</artifactId>
</dependency>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-dbcp-service-api</artifactId>
+ <scope>provided</scope>
+ </dependency>
+
<dependency>
<groupId>org.apache.nifi</groupId>
<artifactId>nifi-service-utils</artifactId>
diff --git
a/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/src/main/java/org/apache/nifi/services/azure/AzureEntraDatabasePasswordProvider.java
b/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/src/main/java/org/apache/nifi/services/azure/AzureEntraDatabasePasswordProvider.java
new file mode 100644
index 00000000000..53c4acf3e1e
--- /dev/null
+++
b/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/src/main/java/org/apache/nifi/services/azure/AzureEntraDatabasePasswordProvider.java
@@ -0,0 +1,215 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.nifi.services.azure;
+
+import com.azure.core.credential.AccessToken;
+import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRequestContext;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.nifi.annotation.documentation.CapabilityDescription;
+import org.apache.nifi.annotation.documentation.Tags;
+import org.apache.nifi.annotation.lifecycle.OnDisabled;
+import org.apache.nifi.annotation.lifecycle.OnEnabled;
+import org.apache.nifi.components.ConfigVerificationResult;
+import org.apache.nifi.components.ConfigVerificationResult.Outcome;
+import org.apache.nifi.components.PropertyDescriptor;
+import org.apache.nifi.controller.AbstractControllerService;
+import org.apache.nifi.controller.ConfigurationContext;
+import org.apache.nifi.controller.VerifiableControllerService;
+import org.apache.nifi.dbcp.api.DatabasePasswordProvider;
+import org.apache.nifi.dbcp.api.DatabasePasswordRequestContext;
+import org.apache.nifi.logging.ComponentLog;
+import org.apache.nifi.processor.exception.ProcessException;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+@Tags({"azure", "microsoft entra", "jdbc", "database password", "postgresql",
"mysql", "authentication"})
+@CapabilityDescription("""
+ Acquires short-lived Microsoft Entra database passwords for JDBC
authentication to Azure Database for PostgreSQL Flexible Server and
+ Azure Database for MySQL Flexible Server. Only Azure Database for
PostgreSQL Flexible Server and Azure Database for MySQL Flexible
+ Server are supported.
+ """)
+public class AzureEntraDatabasePasswordProvider extends
AbstractControllerService
+ implements DatabasePasswordProvider, VerifiableControllerService {
+
+ static final String OSS_RDBMS_SCOPE =
"https://ossrdbms-aad.database.windows.net/.default";;
+ static final String FAILED_CREDENTIAL_RESOLUTION_MESSAGE = "Failed to
resolve Azure credentials for Microsoft Entra database password.";
+ static final String FAILED_TOKEN_ACQUISITION_MESSAGE = "Failed to acquire
a valid Microsoft Entra database access token.";
+ static final String VERIFY_CREDENTIALS_STEP = "Resolve Azure credentials";
+ static final String VERIFY_TOKEN_STEP = "Acquire Microsoft Entra database
access token";
+ static final String VERIFY_CREDENTIALS_UNAVAILABLE = "Configured Azure
Credentials Service did not return Azure credentials.";
+ static final String VERIFY_TOKEN_ACQUISITION_FAILED = "Failed to acquire a
valid Microsoft Entra database access token.";
+ static final Duration DEFAULT_TOKEN_ACQUISITION_TIMEOUT =
Duration.ofMinutes(2);
+
+ static final PropertyDescriptor AZURE_CREDENTIALS_SERVICE = new
PropertyDescriptor.Builder()
+ .name("Azure Credentials Service")
+ .description("Controller Service that provides the Azure
credentials used to request Microsoft Entra database access tokens.")
+ .identifiesControllerService(AzureCredentialsService.class)
+ .required(true)
+ .build();
+
+ private static final List<PropertyDescriptor> PROPERTY_DESCRIPTORS =
List.of(
+ AZURE_CREDENTIALS_SERVICE
+ );
+
+ private final Duration tokenAcquisitionTimeout;
+ private volatile AzureCredentialsService azureCredentialsService;
+
+ public AzureEntraDatabasePasswordProvider() {
+ this(DEFAULT_TOKEN_ACQUISITION_TIMEOUT);
+ }
+
+ AzureEntraDatabasePasswordProvider(final Duration tokenAcquisitionTimeout)
{
+ this.tokenAcquisitionTimeout =
requirePositiveDuration(tokenAcquisitionTimeout);
+ }
+
+ @Override
+ protected List<PropertyDescriptor> getSupportedPropertyDescriptors() {
+ return PROPERTY_DESCRIPTORS;
+ }
+
+ @OnEnabled
+ public void onEnabled(final ConfigurationContext context) {
+ azureCredentialsService = resolveAzureCredentialsService(context);
+ }
+
+ @OnDisabled
+ public void onDisabled() {
+ azureCredentialsService = null;
+ }
+
+ @Override
+ public char[] getPassword(final DatabasePasswordRequestContext
requestContext) {
+ Objects.requireNonNull(requestContext, "Database Password Request
Context required");
+
+ final AzureCredentialsService configuredCredentialsService =
azureCredentialsService;
+ if (configuredCredentialsService == null) {
+ throw new ProcessException(FAILED_CREDENTIAL_RESOLUTION_MESSAGE);
+ }
+
+ final TokenCredential credential;
+ try {
+ credential = configuredCredentialsService.getCredentials();
+ } catch (final RuntimeException e) {
+ throw new ProcessException(FAILED_CREDENTIAL_RESOLUTION_MESSAGE);
+ }
+
+ if (credential == null) {
+ throw new ProcessException(FAILED_CREDENTIAL_RESOLUTION_MESSAGE);
+ }
+
+ final AccessToken accessToken;
+ try {
+ accessToken =
credential.getToken(createTokenRequestContext()).block(tokenAcquisitionTimeout);
+ } catch (final RuntimeException e) {
+ throw new ProcessException(FAILED_TOKEN_ACQUISITION_MESSAGE);
+ }
+
+ if (!isValidAccessToken(accessToken)) {
+ throw new ProcessException(FAILED_TOKEN_ACQUISITION_MESSAGE);
+ }
+
+ return accessToken.getToken().toCharArray();
+ }
+
+ @Override
+ public List<ConfigVerificationResult> verify(final ConfigurationContext
context, final ComponentLog verificationLogger,
+ final Map<String, String>
attributes) {
+ final List<ConfigVerificationResult> results = new ArrayList<>(2);
+
+ final TokenCredential verificationCredential =
resolveVerificationCredential(context, verificationLogger);
+ if (verificationCredential == null) {
+ results.add(buildVerificationResult(VERIFY_CREDENTIALS_STEP,
Outcome.FAILED, VERIFY_CREDENTIALS_UNAVAILABLE));
+ return results;
+ }
+
+ results.add(buildVerificationResult(
+ VERIFY_CREDENTIALS_STEP,
+ Outcome.SUCCESSFUL,
+ "Resolved Azure credentials service and current
TokenCredential."
+ ));
+ results.add(verifyAccessToken(verificationCredential,
verificationLogger));
+ return results;
+ }
+
+ private TokenCredential resolveVerificationCredential(final
ConfigurationContext context, final ComponentLog verificationLogger) {
+ final AzureCredentialsService verificationCredentialsService =
resolveAzureCredentialsService(context);
+ if (verificationCredentialsService == null) {
+ return null;
+ }
+
+ try {
+ return verificationCredentialsService.getCredentials();
+ } catch (final RuntimeException e) {
+ verificationLogger.error(VERIFY_CREDENTIALS_UNAVAILABLE);
+ return null;
+ }
+ }
+
+ private ConfigVerificationResult verifyAccessToken(final TokenCredential
credential, final ComponentLog verificationLogger) {
+ final AccessToken accessToken;
+ try {
+ accessToken =
credential.getToken(createTokenRequestContext()).block(tokenAcquisitionTimeout);
+ } catch (final RuntimeException e) {
+ verificationLogger.error(VERIFY_TOKEN_ACQUISITION_FAILED);
+ return buildVerificationResult(VERIFY_TOKEN_STEP, Outcome.FAILED,
VERIFY_TOKEN_ACQUISITION_FAILED);
+ }
+
+ if (!isValidAccessToken(accessToken)) {
+ verificationLogger.error(VERIFY_TOKEN_ACQUISITION_FAILED);
+ return buildVerificationResult(VERIFY_TOKEN_STEP, Outcome.FAILED,
VERIFY_TOKEN_ACQUISITION_FAILED);
+ }
+
+ return buildVerificationResult(
+ VERIFY_TOKEN_STEP,
+ Outcome.SUCCESSFUL,
+ "Acquired a Microsoft Entra database access token. Use DBCP
Verify to validate database connectivity."
+ );
+ }
+
+ private AzureCredentialsService resolveAzureCredentialsService(final
ConfigurationContext context) {
+ return
context.getProperty(AZURE_CREDENTIALS_SERVICE).asControllerService(AzureCredentialsService.class);
+ }
+
+ private TokenRequestContext createTokenRequestContext() {
+ return new TokenRequestContext().addScopes(OSS_RDBMS_SCOPE);
+ }
+
+ private Duration requirePositiveDuration(final Duration
tokenAcquisitionTimeout) {
+ final Duration configuredTimeout =
Objects.requireNonNull(tokenAcquisitionTimeout, "Token acquisition timeout
required");
+ if (configuredTimeout.isZero() || configuredTimeout.isNegative()) {
+ throw new IllegalArgumentException("Token acquisition timeout must
be positive");
+ }
+ return configuredTimeout;
+ }
+
+ private boolean isValidAccessToken(final AccessToken accessToken) {
+ return accessToken != null &&
StringUtils.isNotBlank(accessToken.getToken());
+ }
+
+ private ConfigVerificationResult buildVerificationResult(final String
stepName, final Outcome outcome, final String explanation) {
+ return new ConfigVerificationResult.Builder()
+ .verificationStepName(stepName)
+ .outcome(outcome)
+ .explanation(explanation)
+ .build();
+ }
+}
diff --git
a/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/src/main/resources/META-INF/services/org.apache.nifi.controller.ControllerService
b/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/src/main/resources/META-INF/services/org.apache.nifi.controller.ControllerService
index e76b973a781..f1483760db6 100644
---
a/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/src/main/resources/META-INF/services/org.apache.nifi.controller.ControllerService
+++
b/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/src/main/resources/META-INF/services/org.apache.nifi.controller.ControllerService
@@ -22,5 +22,6 @@
org.apache.nifi.services.azure.storage.AzureStorageCredentialsControllerService_
org.apache.nifi.services.azure.storage.AzureStorageCredentialsControllerServiceLookup_v12
org.apache.nifi.services.azure.StandardAzureCredentialsControllerService
org.apache.nifi.services.azure.StandardAzureIdentityFederationTokenProvider
+org.apache.nifi.services.azure.AzureEntraDatabasePasswordProvider
org.apache.nifi.services.azure.storage.AzureBlobStorageFileResourceService
org.apache.nifi.services.azure.storage.AzureDataLakeStorageFileResourceService
diff --git
a/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/src/main/resources/docs/org.apache.nifi.services.azure.AzureEntraDatabasePasswordProvider/additionalDetails.md
b/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/src/main/resources/docs/org.apache.nifi.services.azure.AzureEntraDatabasePasswordProvider/additionalDetails.md
new file mode 100644
index 00000000000..84e3988f5e2
--- /dev/null
+++
b/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/src/main/resources/docs/org.apache.nifi.services.azure.AzureEntraDatabasePasswordProvider/additionalDetails.md
@@ -0,0 +1,92 @@
+<!--
+ Licensed to the Apache Software Foundation (ASF) under one or more
+ contributor license agreements. See the NOTICE file distributed with
+ this work for additional information regarding copyright ownership.
+ The ASF licenses this file to You under the Apache License, Version 2.0
+ (the "License"); you may not use this file except in compliance with
+ the License. You may obtain a copy of the License at
+ http://www.apache.org/licenses/LICENSE-2.0
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+-->
+
+## Summary
+
+`AzureEntraDatabasePasswordProvider` acquires a short-lived Microsoft Entra
access token and supplies it as the
+database password for a DBCP service. Use it to connect NiFi to Azure Database
for PostgreSQL Flexible Server or Azure
+Database for MySQL Flexible Server without storing a long-lived database
password in NiFi. Other database products and
+sovereign Azure clouds are not supported.
+
+The provider supplies a password when DBCP creates a new physical connection.
Existing pooled connections are not
+reauthenticated when the token expires.
+
+## Usage
+
+1. Configure an Azure Credentials Service that can obtain credentials for the
public Azure cloud.
+2. Create and enable `AzureEntraDatabasePasswordProvider`.
+3. Set **Azure Credentials Service** to the configured credentials service.
+4. Configure the DBCP service with the JDBC URL, driver, database user, and
**Database Password Provider** set to
+ `AzureEntraDatabasePasswordProvider`.
+5. Run **Verify** on the provider, then run **Verify** on the DBCP service.
+
+Create the Microsoft Entra principal in the database separately and grant the
database privileges required by NiFi.
+
+## Workload Identity Federation
+
+For Workload Identity Federation, configure
`StandardAzureCredentialsControllerService` with **Credentials Strategy**
+set to **Identity Federation** and select a
`StandardAzureIdentityFederationTokenProvider`. The federation token
+provider accepts an OAuth2 access token provider that supplies the external
client assertion. The federated identity
+credential in Microsoft Entra must match the assertion issuer, subject, and
audience.
+
+The password provider requests the public-cloud Azure OSS RDBMS resource. It
does not require the Azure JDBC
+authentication plugins and does not depend on a particular external assertion
issuer.
+
+## PostgreSQL Configuration
+
+Configure Microsoft Entra authentication on Azure Database for PostgreSQL
Flexible Server and create the database
+role that corresponds to the Entra principal. Set the DBCP **Database User**
to that mapped role name.
+
+The token is sent using PostgreSQL cleartext-password authentication and must
be protected by TLS. Use direct port
+5432 as the baseline and set `sslmode=verify-full`, with the client configured
to trust the server certificate chain.
+
+| Setting | Value |
+|---|---|
+| Driver Class Name | `org.postgresql.Driver` |
+| JDBC URL |
`jdbc:postgresql://<SERVER>.postgres.database.azure.com:5432/<DATABASE>?sslmode=verify-full`
|
+| Database User | Microsoft Entra principal's mapped PostgreSQL role name |
+
+Azure's built-in PgBouncer on port 6432 must be validated separately. Large
Entra tokens containing many group claims
+can exceed intermediary protocol buffers even when the same token works
against direct port 5432. This provider does
+not claim unrestricted PgBouncer compatibility.
+
+## MySQL Configuration
+
+Configure Microsoft Entra authentication on Azure Database for MySQL Flexible
Server and create an Entra user or alias
+for the principal. Set the DBCP **Database User** to the mapped alias. MySQL
user names are limited to 32 characters, so
+an alias can be required for a longer Entra principal name.
+
+MySQL Connector/J must use TLS, verify the server certificate, and permit the
`mysql_clear_password` authentication
+mechanism used to transmit the token. Configure Connector/J with a trust store
containing the issuing CA certificate,
+use `sslMode=VERIFY_IDENTITY`, and set `allowCleartextPasswords=true`. The
Azure MySQL JDBC authentication plugin is not
+required when this provider supplies the token.
+
+| Setting | Value |
+|---|---|
+| Driver Class Name | `com.mysql.cj.jdbc.Driver` |
+| Driver Location(s) | compatible MySQL Connector/J driver jar provided to the
DBCP service |
+| JDBC URL |
`jdbc:mysql://<SERVER>.mysql.database.azure.com:3306/<DATABASE>?sslMode=VERIFY_IDENTITY&allowCleartextPasswords=true`
|
+| Database User | Microsoft Entra principal's mapped MySQL user or alias |
+
+## Verify and Troubleshooting
+
+`AzureEntraDatabasePasswordProvider` **Verify** checks that the configured
Azure Credentials Service can acquire a
+nonblank token for the Azure OSS RDBMS resource. It does not validate the JDBC
URL, database server, network
+path, TLS configuration, driver, principal mapping, or database grants. DBCP
**Verify** checks the actual database
+connection using those settings.
+
+If provider **Verify** fails, confirm that the Azure Credentials Service and
any upstream identity-federation services
+are enabled and configured for the intended Entra application and tenant. If
provider **Verify** succeeds but DBCP
+**Verify** fails, investigate the JDBC connection, network, TLS, driver,
database principal mapping, and grants.
\ No newline at end of file
diff --git
a/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/src/test/java/org/apache/nifi/services/azure/AzureEntraDatabasePasswordProviderTest.java
b/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/src/test/java/org/apache/nifi/services/azure/AzureEntraDatabasePasswordProviderTest.java
new file mode 100644
index 00000000000..aa7861f75aa
--- /dev/null
+++
b/nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/src/test/java/org/apache/nifi/services/azure/AzureEntraDatabasePasswordProviderTest.java
@@ -0,0 +1,616 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.nifi.services.azure;
+
+import com.azure.core.credential.AccessToken;
+import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRequestContext;
+import org.apache.nifi.components.ConfigVerificationResult;
+import org.apache.nifi.components.PropertyDescriptor;
+import org.apache.nifi.controller.AbstractControllerService;
+import org.apache.nifi.dbcp.api.DatabasePasswordProvider;
+import org.apache.nifi.dbcp.api.DatabasePasswordRequestContext;
+import org.apache.nifi.processor.exception.ProcessException;
+import org.apache.nifi.util.LogMessage;
+import org.apache.nifi.util.MockComponentLog;
+import org.apache.nifi.util.NoOpProcessor;
+import org.apache.nifi.util.TestRunner;
+import org.apache.nifi.util.TestRunners;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.slf4j.helpers.MessageFormatter;
+import reactor.core.publisher.Mono;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.time.OffsetDateTime;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Supplier;
+import java.util.stream.Stream;
+
+import static
org.apache.nifi.components.ConfigVerificationResult.Outcome.FAILED;
+import static
org.apache.nifi.components.ConfigVerificationResult.Outcome.SUCCESSFUL;
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class AzureEntraDatabasePasswordProviderTest {
+
+ private static final String CREDENTIALS_SERVICE_ID = "azureCredentials";
+ private static final String PASSWORD_PROVIDER_ID =
"azureEntraPasswordProvider";
+ private static final String DRIVER_CLASS = "org.postgresql.Driver";
+ private static final String DATABASE_USER = "[email protected]";
+ private static final String JDBC_URL =
"jdbc:postgresql://example.postgres.database.azure.com:5432/database";
+ private static final String TOKEN_VALUE = "entra-database-token";
+ private static final String REFRESHED_TOKEN_VALUE =
"refreshed-entra-database-token";
+ private static final String LEAK_SENTINEL = "sentinel-entra-secret-value";
+ private static final Duration SHORT_TOKEN_ACQUISITION_TIMEOUT =
Duration.ofMillis(50);
+
+ private ExecutorService executorService;
+
+ @AfterEach
+ void tearDown() {
+ if (executorService != null) {
+ executorService.shutdownNow();
+ }
+ }
+
+ @Test
+ void
testSupportedPropertyDescriptorsContainOnlyCredentialsServiceAndCredentialsPropertyIsRequired()
throws Exception {
+ final AzureEntraDatabasePasswordProvider provider = new
AzureEntraDatabasePasswordProvider();
+ final List<PropertyDescriptor> descriptors =
provider.getSupportedPropertyDescriptors();
+
+ assertEquals(1, descriptors.size());
+
assertEquals(AzureEntraDatabasePasswordProvider.AZURE_CREDENTIALS_SERVICE,
descriptors.get(0));
+ assertTrue(descriptors.get(0).isRequired());
+ final TestRunner runner =
TestRunners.newTestRunner(NoOpProcessor.class);
+
+ runner.addControllerService(PASSWORD_PROVIDER_ID, provider);
+ runner.assertNotValid(provider);
+ }
+
+ @Test
+ void testGetPasswordUsesCurrentCredentialForEveryCall() throws Exception {
+ final TestAzureCredentialsService credentialsService = new
TestAzureCredentialsService(
+ new StaticTokenCredential(validToken(TOKEN_VALUE)),
+ new StaticTokenCredential(validToken(REFRESHED_TOKEN_VALUE))
+ );
+ final DatabasePasswordProvider provider =
getProvider(configureRunner(credentialsService));
+
+ final char[] firstPassword = provider.getPassword(requestContext());
+ final char[] firstPasswordCopy = firstPassword.clone();
+ firstPassword[0] = 'X';
+ final char[] secondPassword = provider.getPassword(requestContext());
+
+ assertEquals(2, credentialsService.getGetCredentialsCount());
+ assertNotSame(firstPassword, secondPassword);
+ assertArrayEquals(TOKEN_VALUE.toCharArray(), firstPasswordCopy);
+ assertArrayEquals(REFRESHED_TOKEN_VALUE.toCharArray(), secondPassword);
+ }
+
+ @Test
+ void testGetPasswordRequestsExactOssRdbmsScope() throws Exception {
+ final RecordingTokenCredential credential = new
RecordingTokenCredential(Mono.just(validToken(TOKEN_VALUE)));
+ final DatabasePasswordProvider provider =
getProvider(configureRunner(new TestAzureCredentialsService(credential)));
+
+ provider.getPassword(requestContext());
+
+
assertEquals(List.of(AzureEntraDatabasePasswordProvider.OSS_RDBMS_SCOPE),
credential.getLastRequestedScopes());
+ }
+
+ @Test
+ void
testVerifyBeforeEnableUsesCurrentCredentialAndLeavesEnabledStateUntouched()
throws Exception {
+ final RecordingTokenCredential credential = new
RecordingTokenCredential(Mono.just(validToken(TOKEN_VALUE)));
+ final TestRunner runner = configureRunner(new
TestAzureCredentialsService(credential), false);
+ final AzureEntraDatabasePasswordProvider provider =
getProviderImplementation(runner);
+
+ final List<ConfigVerificationResult> results = runner.verify(provider,
Map.of());
+
+ assertEquals(1, credential.getGetTokenCount());
+ assertEquals(2, results.size());
+ assertVerificationResult(results.get(0),
AzureEntraDatabasePasswordProvider.VERIFY_CREDENTIALS_STEP, SUCCESSFUL,
+ "Resolved Azure credentials service and current
TokenCredential.");
+ assertVerificationResult(results.get(1),
AzureEntraDatabasePasswordProvider.VERIFY_TOKEN_STEP, SUCCESSFUL,
+ "DBCP Verify");
+
+ final ProcessException exception =
assertThrows(ProcessException.class, () ->
provider.getPassword(requestContext()));
+
assertEquals(AzureEntraDatabasePasswordProvider.FAILED_CREDENTIAL_RESOLUTION_MESSAGE,
exception.getMessage());
+ assertNull(exception.getCause());
+ }
+
+ @Test
+ void testVerifyUsesFreshCredentialWithoutMutatingEnabledState() throws
Exception {
+ final RecordingTokenCredential verifyCredential = new
RecordingTokenCredential(Mono.just(validToken(TOKEN_VALUE)));
+ final RecordingTokenCredential passwordCredential = new
RecordingTokenCredential(Mono.just(validToken(REFRESHED_TOKEN_VALUE)));
+ final TestAzureCredentialsService credentialsService = new
TestAzureCredentialsService(verifyCredential, passwordCredential);
+ final TestRunner runner = configureRunner(credentialsService);
+ final AzureEntraDatabasePasswordProvider provider =
getProviderImplementation(runner);
+
+ final List<ConfigVerificationResult> results = runner.verify(provider,
Map.of());
+ final char[] password = provider.getPassword(requestContext());
+
+ assertEquals(2, credentialsService.getGetCredentialsCount());
+ assertEquals(1, verifyCredential.getGetTokenCount());
+ assertEquals(1, passwordCredential.getGetTokenCount());
+ assertVerificationResult(results.get(1),
AzureEntraDatabasePasswordProvider.VERIFY_TOKEN_STEP, SUCCESSFUL, "DBCP
Verify");
+ assertArrayEquals(REFRESHED_TOKEN_VALUE.toCharArray(), password);
+ }
+
+ @Test
+ void testVerifyCredentialsResolutionFailureIsSanitized() throws Exception {
+ final TestRunner runner = configureRunner(new
TestAzureCredentialsService(new IllegalStateException(LEAK_SENTINEL)), false);
+ final AzureEntraDatabasePasswordProvider provider =
getProviderImplementation(runner);
+
+ final List<ConfigVerificationResult> results = runner.verify(provider,
Map.of());
+
+ assertEquals(1, results.size());
+ assertVerificationResult(results.get(0),
AzureEntraDatabasePasswordProvider.VERIFY_CREDENTIALS_STEP, FAILED,
+
AzureEntraDatabasePasswordProvider.VERIFY_CREDENTIALS_UNAVAILABLE);
+ assertFalse(results.get(0).getExplanation().contains(LEAK_SENTINEL));
+
assertNoSensitiveLogging(runner.getControllerServiceLogger(PASSWORD_PROVIDER_ID),
LEAK_SENTINEL);
+ }
+
+ @Test
+ void testVerifyNullCredentialFailsCredentialsStep() throws Exception {
+ final TestRunner runner = configureRunner(new
TestAzureCredentialsService((TokenCredential) null), false);
+
+ final List<ConfigVerificationResult> results =
runner.verify(getProviderImplementation(runner), Map.of());
+
+ assertEquals(1, results.size());
+ assertVerificationResult(results.get(0),
AzureEntraDatabasePasswordProvider.VERIFY_CREDENTIALS_STEP, FAILED,
+
AzureEntraDatabasePasswordProvider.VERIFY_CREDENTIALS_UNAVAILABLE);
+ }
+
+ @ParameterizedTest(name = "{0}")
+ @MethodSource("invalidAccessTokens")
+ void testVerifyInvalidAccessTokenFailsTokenStep(final String testName,
final AccessToken accessToken) throws Exception {
+ final TestRunner runner = configureRunner(new
TestAzureCredentialsService(new StaticTokenCredential(accessToken)), false);
+
+ final List<ConfigVerificationResult> results =
runner.verify(getProviderImplementation(runner), Map.of());
+
+ assertEquals(2, results.size());
+ assertVerificationResult(results.get(1),
AzureEntraDatabasePasswordProvider.VERIFY_TOKEN_STEP, FAILED,
+
AzureEntraDatabasePasswordProvider.VERIFY_TOKEN_ACQUISITION_FAILED);
+ }
+
+ @ParameterizedTest(name = "{0}")
+ @MethodSource("acceptedAccessTokens")
+ void testVerifyAcceptsAccessTokensWithoutFutureExpiration(final String
testName, final AccessToken accessToken,
+ final String
expectedPassword) throws Exception {
+ final TestRunner runner = configureRunner(new
TestAzureCredentialsService(new StaticTokenCredential(accessToken)), false);
+
+ final List<ConfigVerificationResult> results =
runner.verify(getProviderImplementation(runner), Map.of());
+
+ assertEquals(2, results.size());
+ assertVerificationResult(results.get(1),
AzureEntraDatabasePasswordProvider.VERIFY_TOKEN_STEP, SUCCESSFUL,
+ "DBCP Verify");
+ }
+
+ @Test
+ void testVerifyTokenAcquisitionTimeoutFailsTokenStep() throws Exception {
+ final AzureEntraDatabasePasswordProvider provider = new
AzureEntraDatabasePasswordProvider(SHORT_TOKEN_ACQUISITION_TIMEOUT);
+ final TestRunner runner = configureRunner(provider,
+ new TestAzureCredentialsService(new
RecordingTokenCredential(Mono.never())), false);
+
+ assertTimeoutPreemptively(Duration.ofSeconds(5), () -> {
+ final List<ConfigVerificationResult> results =
runner.verify(provider, Map.of());
+
+ assertEquals(2, results.size());
+ assertVerificationResult(results.get(1),
AzureEntraDatabasePasswordProvider.VERIFY_TOKEN_STEP, FAILED,
+
AzureEntraDatabasePasswordProvider.VERIFY_TOKEN_ACQUISITION_FAILED);
+ });
+ }
+
+ @Test
+ void testVerifyTokenAcquisitionFailureIsSanitized() throws Exception {
+ final TestRunner runner = configureRunner(
+ new TestAzureCredentialsService(new
RecordingTokenCredential(Mono.error(new IllegalStateException(LEAK_SENTINEL)))),
+ false
+ );
+
+ final List<ConfigVerificationResult> results =
runner.verify(getProviderImplementation(runner), Map.of());
+
+ assertEquals(2, results.size());
+ assertVerificationResult(results.get(1),
AzureEntraDatabasePasswordProvider.VERIFY_TOKEN_STEP, FAILED,
+
AzureEntraDatabasePasswordProvider.VERIFY_TOKEN_ACQUISITION_FAILED);
+ assertFalse(results.get(1).getExplanation().contains(LEAK_SENTINEL));
+
assertNoSensitiveLogging(runner.getControllerServiceLogger(PASSWORD_PROVIDER_ID),
LEAK_SENTINEL);
+ }
+
+ @Test
+ void testVerifyUsesExactOssRdbmsScope() throws Exception {
+ final RecordingTokenCredential credential = new
RecordingTokenCredential(Mono.just(validToken(TOKEN_VALUE)));
+ final TestRunner runner = configureRunner(new
TestAzureCredentialsService(credential), false);
+
+ runner.verify(getProviderImplementation(runner), Map.of());
+
+
assertEquals(List.of(AzureEntraDatabasePasswordProvider.OSS_RDBMS_SCOPE),
credential.getLastRequestedScopes());
+ }
+
+ @Test
+ void testDisabledCallsFailBeforeEnable() throws Exception {
+ final DatabasePasswordProvider provider = getProvider(configureRunner(
+ new TestAzureCredentialsService(new
StaticTokenCredential(validToken(TOKEN_VALUE))), false));
+
+ final ProcessException exception =
assertThrows(ProcessException.class, () ->
provider.getPassword(requestContext()));
+
+
assertEquals(AzureEntraDatabasePasswordProvider.FAILED_CREDENTIAL_RESOLUTION_MESSAGE,
exception.getMessage());
+ assertNull(exception.getCause());
+ }
+
+ @Test
+ void testOnDisabledClearsRetainedCredentialsServiceAndFailsClosed() throws
Exception {
+ final TestRunner runner = configureRunner(new
TestAzureCredentialsService(new
StaticTokenCredential(validToken(TOKEN_VALUE))));
+ final AzureEntraDatabasePasswordProvider provider =
getProviderImplementation(runner);
+
+ runner.disableControllerService(provider);
+
+ final ProcessException exception =
assertThrows(ProcessException.class, () ->
provider.getPassword(requestContext()));
+
assertEquals(AzureEntraDatabasePasswordProvider.FAILED_CREDENTIAL_RESOLUTION_MESSAGE,
exception.getMessage());
+ assertNull(exception.getCause());
+ }
+
+ @ParameterizedTest(name = "{0}")
+ @MethodSource("invalidAccessTokens")
+ void testGetPasswordRejectsInvalidAccessToken(final String testName, final
AccessToken accessToken) throws Exception {
+ final DatabasePasswordProvider provider = getProvider(configureRunner(
+ new TestAzureCredentialsService(new
StaticTokenCredential(accessToken))));
+
+ final ProcessException exception =
assertThrows(ProcessException.class, () ->
provider.getPassword(requestContext()));
+
+
assertEquals(AzureEntraDatabasePasswordProvider.FAILED_TOKEN_ACQUISITION_MESSAGE,
exception.getMessage());
+ assertNull(exception.getCause());
+ }
+
+ @ParameterizedTest(name = "{0}")
+ @MethodSource("acceptedAccessTokens")
+ void testGetPasswordPassesThroughAccessTokensWithoutFutureExpiration(final
String testName, final AccessToken accessToken,
+ final
String expectedPassword) throws Exception {
+ final DatabasePasswordProvider provider = getProvider(configureRunner(
+ new TestAzureCredentialsService(new
StaticTokenCredential(accessToken))));
+
+ final char[] password = provider.getPassword(requestContext());
+
+ assertArrayEquals(expectedPassword.toCharArray(), password);
+ }
+
+ @Test
+ void testGetPasswordTokenAcquisitionTimeoutFailsClosed() throws Exception {
+ final AzureEntraDatabasePasswordProvider provider = new
AzureEntraDatabasePasswordProvider(SHORT_TOKEN_ACQUISITION_TIMEOUT);
+ final DatabasePasswordProvider configuredProvider =
getProvider(configureRunner(provider,
+ new TestAzureCredentialsService(new
RecordingTokenCredential(Mono.never()))));
+
+ assertTimeoutPreemptively(Duration.ofSeconds(5), () -> {
+ final ProcessException exception =
assertThrows(ProcessException.class, () ->
configuredProvider.getPassword(requestContext()));
+
+
assertEquals(AzureEntraDatabasePasswordProvider.FAILED_TOKEN_ACQUISITION_MESSAGE,
exception.getMessage());
+ assertNull(exception.getCause());
+ });
+ }
+
+ @ParameterizedTest(name = "{0}")
+ @MethodSource("sanitizedPasswordFailures")
+ void testPasswordGenerationFailuresAreSanitized(final String testName,
+ final
Supplier<TestAzureCredentialsService> credentialsServiceSupplier,
+ final String
expectedMessage) throws Exception {
+ final TestRunner runner =
configureRunner(credentialsServiceSupplier.get());
+ final DatabasePasswordProvider provider = getProvider(runner);
+
+ final ProcessException exception =
assertThrows(ProcessException.class, () ->
provider.getPassword(requestContext()));
+
+ assertEquals(expectedMessage, exception.getMessage());
+ assertNull(exception.getCause());
+ assertFalse(exception.getMessage().contains(LEAK_SENTINEL));
+
assertNoSensitiveLogging(runner.getControllerServiceLogger(PASSWORD_PROVIDER_ID),
LEAK_SENTINEL);
+ }
+
+ @Test
+ void testNullCredentialRejectedForPasswordGeneration() throws Exception {
+ final DatabasePasswordProvider provider =
getProvider(configureRunner(new TestAzureCredentialsService((TokenCredential)
null)));
+
+ final ProcessException exception =
assertThrows(ProcessException.class, () ->
provider.getPassword(requestContext()));
+
+
assertEquals(AzureEntraDatabasePasswordProvider.FAILED_CREDENTIAL_RESOLUTION_MESSAGE,
exception.getMessage());
+ assertNull(exception.getCause());
+ }
+
+ @Test
+ void testConcurrentGetPasswordUsesSeparateCredentialResolutionPerCall()
throws Exception {
+ final BlockingTokenCredential firstCredential = new
BlockingTokenCredential(validToken(TOKEN_VALUE));
+ final BlockingTokenCredential secondCredential = new
BlockingTokenCredential(validToken(REFRESHED_TOKEN_VALUE));
+ final TestAzureCredentialsService credentialsService = new
TestAzureCredentialsService(firstCredential, secondCredential);
+ final DatabasePasswordProvider provider =
getProvider(configureRunner(credentialsService));
+
+ executorService = Executors.newFixedThreadPool(2);
+ final CountDownLatch startLatch = new CountDownLatch(1);
+ final Future<char[]> first = executorService.submit(() ->
getPasswordAfterStart(provider, startLatch));
+ final Future<char[]> second = executorService.submit(() ->
getPasswordAfterStart(provider, startLatch));
+
+ startLatch.countDown();
+ assertTrue(firstCredential.awaitGetTokenEntry());
+ assertTrue(secondCredential.awaitGetTokenEntry());
+ firstCredential.releaseGetToken();
+ secondCredential.releaseGetToken();
+
+ final List<String> passwords = List.of(
+ new String(first.get(5, TimeUnit.SECONDS)),
+ new String(second.get(5, TimeUnit.SECONDS))
+ );
+
+ assertTrue(passwords.contains(TOKEN_VALUE));
+ assertTrue(passwords.contains(REFRESHED_TOKEN_VALUE));
+ assertEquals(2, credentialsService.getGetCredentialsCount());
+ assertEquals(1, firstCredential.getGetTokenCount());
+ assertEquals(1, secondCredential.getGetTokenCount());
+ }
+
+ @Test
+ void testControllerServiceRegistrationContainsProvider() throws
IOException {
+ final String resourcePath =
"META-INF/services/org.apache.nifi.controller.ControllerService";
+ try (InputStream inputStream =
AzureEntraDatabasePasswordProvider.class.getClassLoader().getResourceAsStream(resourcePath))
{
+ assertNotNull(inputStream);
+ final String registeredServices = new
String(inputStream.readAllBytes(), StandardCharsets.UTF_8);
+
assertTrue(registeredServices.contains(AzureEntraDatabasePasswordProvider.class.getName()));
+ }
+ }
+
+ private TestRunner configureRunner(final TestAzureCredentialsService
credentialsService) throws Exception {
+ return configureRunner(new AzureEntraDatabasePasswordProvider(),
credentialsService, true);
+ }
+
+ private TestRunner configureRunner(final TestAzureCredentialsService
credentialsService, final boolean enableProvider) throws Exception {
+ return configureRunner(new AzureEntraDatabasePasswordProvider(),
credentialsService, enableProvider);
+ }
+
+ private TestRunner configureRunner(final
AzureEntraDatabasePasswordProvider provider,
+ final TestAzureCredentialsService
credentialsService) throws Exception {
+ return configureRunner(provider, credentialsService, true);
+ }
+
+ private TestRunner configureRunner(final
AzureEntraDatabasePasswordProvider provider,
+ final TestAzureCredentialsService
credentialsService,
+ final boolean enableProvider) throws
Exception {
+ final TestRunner runner =
TestRunners.newTestRunner(NoOpProcessor.class);
+
+ runner.addControllerService(CREDENTIALS_SERVICE_ID,
credentialsService);
+ runner.enableControllerService(credentialsService);
+
+ runner.addControllerService(PASSWORD_PROVIDER_ID, provider);
+ runner.setProperty(provider,
AzureEntraDatabasePasswordProvider.AZURE_CREDENTIALS_SERVICE,
CREDENTIALS_SERVICE_ID);
+ if (enableProvider) {
+ runner.enableControllerService(provider);
+ runner.assertValid(provider);
+ }
+
+ return runner;
+ }
+
+ private DatabasePasswordProvider getProvider(final TestRunner runner) {
+ return (DatabasePasswordProvider) runner.getProcessContext()
+ .getControllerServiceLookup()
+ .getControllerService(PASSWORD_PROVIDER_ID);
+ }
+
+ private AzureEntraDatabasePasswordProvider getProviderImplementation(final
TestRunner runner) {
+ return (AzureEntraDatabasePasswordProvider) getProvider(runner);
+ }
+
+ private DatabasePasswordRequestContext requestContext() {
+ return DatabasePasswordRequestContext.builder()
+ .jdbcUrl(JDBC_URL)
+ .databaseUser(DATABASE_USER)
+ .driverClassName(DRIVER_CLASS)
+ .connectionProperties(Map.of())
+ .build();
+ }
+
+ private char[] getPasswordAfterStart(final DatabasePasswordProvider
provider, final CountDownLatch startLatch) throws InterruptedException {
+ startLatch.await(5, TimeUnit.SECONDS);
+ return provider.getPassword(requestContext());
+ }
+
+ private static Stream<Arguments> invalidAccessTokens() {
+ return Stream.of(
+ Arguments.of("null access token", null),
+ Arguments.of("null token string", token(null,
OffsetDateTime.now().plusMinutes(15))),
+ Arguments.of("blank token string", validToken(" "))
+ );
+ }
+
+ private static Stream<Arguments> acceptedAccessTokens() {
+ return Stream.of(
+ Arguments.of("expired token", expiredToken(TOKEN_VALUE),
TOKEN_VALUE),
+ Arguments.of("null expiration", token(REFRESHED_TOKEN_VALUE,
null), REFRESHED_TOKEN_VALUE)
+ );
+ }
+
+ private static Stream<Arguments> sanitizedPasswordFailures() {
+ return Stream.of(
+ Arguments.of("credentials resolution failure",
+ (Supplier<TestAzureCredentialsService>) () -> new
TestAzureCredentialsService(new IllegalStateException(LEAK_SENTINEL)),
+
AzureEntraDatabasePasswordProvider.FAILED_CREDENTIAL_RESOLUTION_MESSAGE),
+ Arguments.of("token acquisition failure",
+ (Supplier<TestAzureCredentialsService>) () -> new
TestAzureCredentialsService(
+ new RecordingTokenCredential(Mono.error(new
IllegalStateException(LEAK_SENTINEL)))),
+
AzureEntraDatabasePasswordProvider.FAILED_TOKEN_ACQUISITION_MESSAGE)
+ );
+ }
+
+ private static AccessToken validToken(final String tokenValue) {
+ return token(tokenValue, OffsetDateTime.now().plusMinutes(15));
+ }
+
+ private static AccessToken expiredToken(final String tokenValue) {
+ return token(tokenValue, OffsetDateTime.now().minusMinutes(15));
+ }
+
+ private static AccessToken token(final String tokenValue, final
OffsetDateTime expiresAt) {
+ return new AccessToken(tokenValue, expiresAt);
+ }
+
+ private static void assertVerificationResult(final
ConfigVerificationResult result, final String stepName,
+ final
ConfigVerificationResult.Outcome outcome, final String explanationFragment) {
+ assertEquals(stepName, result.getVerificationStepName());
+ assertEquals(outcome, result.getOutcome());
+ assertTrue(result.getExplanation().contains(explanationFragment),
result::getExplanation);
+ }
+
+ private static void assertNoSensitiveLogging(final MockComponentLog
logger, final String value) {
+ final List<LogMessage> logMessages = new ArrayList<>();
+ logMessages.addAll(logger.getDebugMessages());
+ logMessages.addAll(logger.getInfoMessages());
+ logMessages.addAll(logger.getWarnMessages());
+ logMessages.addAll(logger.getErrorMessages());
+
+ for (final LogMessage logMessage : logMessages) {
+ final String rawMessage = logMessage.getMsg();
+ assertFalse(rawMessage != null && rawMessage.contains(value));
+ final Object[] args = logMessage.getArgs();
+ final String formattedMessage =
MessageFormatter.arrayFormat(rawMessage, args == null ? new Object[0] :
args).getMessage();
+ assertFalse(formattedMessage != null &&
formattedMessage.contains(value));
+ if (args != null) {
+ for (final Object arg : args) {
+ final String argValue = arg == null ? null :
arg.toString();
+ assertFalse(argValue != null && argValue.contains(value));
+ assertFalse(arg instanceof Throwable);
+ }
+ }
+ assertNull(logMessage.getThrowable());
+ }
+ }
+
+ private static final class TestAzureCredentialsService extends
AbstractControllerService implements AzureCredentialsService {
+ private final List<TokenCredential> credentials;
+ private final RuntimeException credentialsException;
+ private final AtomicInteger getCredentialsCount = new AtomicInteger();
+
+ private TestAzureCredentialsService(final TokenCredential...
credentials) {
+ this.credentials = new ArrayList<>(Arrays.asList(credentials));
+ this.credentialsException = null;
+ }
+
+ private TestAzureCredentialsService(final RuntimeException
credentialsException) {
+ this.credentials = List.of();
+ this.credentialsException = credentialsException;
+ }
+
+ @Override
+ public TokenCredential getCredentials() {
+ final int callIndex = getCredentialsCount.getAndIncrement();
+ if (credentialsException != null) {
+ throw credentialsException;
+ }
+
+ if (credentials.isEmpty()) {
+ return null;
+ }
+
+ return credentials.get(Math.min(callIndex, credentials.size() -
1));
+ }
+
+ private int getGetCredentialsCount() {
+ return getCredentialsCount.get();
+ }
+ }
+
+ private static class RecordingTokenCredential implements TokenCredential {
+ private final Mono<AccessToken> tokenMono;
+ private final AtomicInteger getTokenCount = new AtomicInteger();
+ private final List<List<String>> requestedScopes = new
CopyOnWriteArrayList<>();
+
+ private RecordingTokenCredential(final Mono<AccessToken> tokenMono) {
+ this.tokenMono = tokenMono;
+ }
+
+ @Override
+ public Mono<AccessToken> getToken(final TokenRequestContext request) {
+ getTokenCount.incrementAndGet();
+ requestedScopes.add(List.copyOf(request.getScopes()));
+ return tokenMono;
+ }
+
+ private int getGetTokenCount() {
+ return getTokenCount.get();
+ }
+
+ private List<String> getLastRequestedScopes() {
+ return requestedScopes.isEmpty() ? List.of() :
requestedScopes.get(requestedScopes.size() - 1);
+ }
+ }
+
+ private static final class StaticTokenCredential extends
RecordingTokenCredential {
+ private StaticTokenCredential(final AccessToken accessToken) {
+ super(Mono.justOrEmpty(accessToken));
+ }
+ }
+
+ private static final class BlockingTokenCredential implements
TokenCredential {
+ private final AccessToken accessToken;
+ private final CountDownLatch getTokenEnteredLatch = new
CountDownLatch(1);
+ private final CountDownLatch releaseGetTokenLatch = new
CountDownLatch(1);
+ private final AtomicInteger getTokenCount = new AtomicInteger();
+
+ private BlockingTokenCredential(final AccessToken accessToken) {
+ this.accessToken = accessToken;
+ }
+
+ @Override
+ public Mono<AccessToken> getToken(final TokenRequestContext request) {
+ getTokenCount.incrementAndGet();
+ return Mono.fromCallable(() -> {
+ getTokenEnteredLatch.countDown();
+ if (!releaseGetTokenLatch.await(5, TimeUnit.SECONDS)) {
+ throw new IllegalStateException("Timed out waiting for
token release");
+ }
+ return accessToken;
+ });
+ }
+
+ private boolean awaitGetTokenEntry() throws InterruptedException {
+ return getTokenEnteredLatch.await(5, TimeUnit.SECONDS);
+ }
+
+ private void releaseGetToken() {
+ releaseGetTokenLatch.countDown();
+ }
+
+ private int getGetTokenCount() {
+ return getTokenCount.get();
+ }
+ }
+}
