Skip to content

Commit 9547faf

Browse files
pvillard31pvillard31
committed
NIFI-15161 - Add support for Azure Federated Identity Credentials
1 parent 5a87d32 commit 9547faf

File tree

32 files changed

+1328
-83
lines changed

32 files changed

+1328
-83
lines changed

‎nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/pom.xml‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,20 @@
2626
<artifactId>nifi-utils</artifactId>
2727
</dependency>
2828

29+
<!-- OAuth2 Access Token Provider API for Web Identity (OIDC) support -->
30+
<dependency>
31+
<groupId>org.apache.nifi</groupId>
32+
<artifactId>nifi-oauth2-provider-api</artifactId>
33+
</dependency>
34+
<dependency>
35+
<groupId>org.apache.nifi</groupId>
36+
<artifactId>nifi-web-client-api</artifactId>
37+
</dependency>
38+
<dependency>
39+
<groupId>org.apache.nifi</groupId>
40+
<artifactId>nifi-web-client-provider-api</artifactId>
41+
</dependency>
42+
2943
<dependency>
3044
<groupId>org.apache.nifi</groupId>
3145
<artifactId>nifi-service-utils</artifactId>
@@ -52,10 +66,6 @@
5266
<groupId>org.apache.nifi</groupId>
5367
<artifactId>nifi-record</artifactId>
5468
</dependency>
55-
<dependency>
56-
<groupId>org.apache.nifi</groupId>
57-
<artifactId>nifi-oauth2-provider-api</artifactId>
58-
</dependency>
5969
<dependency>
6070
<groupId>org.apache.nifi</groupId>
6171
<artifactId>nifi-proxy-configuration-api</artifactId>

‎nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/src/main/java/org/apache/nifi/processors/azure/eventhub/ConsumeAzureEventHub.java‎

Lines changed: 62 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@
8484
import org.apache.nifi.shared.azure.eventhubs.AzureEventHubComponent;
8585
import org.apache.nifi.shared.azure.eventhubs.AzureEventHubTransportType;
8686
import org.apache.nifi.shared.azure.eventhubs.BlobStorageAuthenticationStrategy;
87+
import org.apache.nifi.services.azure.AzureIdentityFederationTokenProvider;
8788
import org.apache.nifi.util.StopWatch;
8889
import org.apache.nifi.util.StringUtils;
8990

@@ -150,6 +151,7 @@ public class ConsumeAzureEventHub extends AbstractSessionFactoryProcessor implem
150151
static final PropertyDescriptor SERVICE_BUS_ENDPOINT = AzureEventHubUtils.SERVICE_BUS_ENDPOINT;
151152
static final PropertyDescriptor AUTHENTICATION_STRATEGY = AzureEventHubComponent.AUTHENTICATION_STRATEGY;
152153
static final PropertyDescriptor EVENT_HUB_OAUTH2_ACCESS_TOKEN_PROVIDER = AzureEventHubComponent.OAUTH2_ACCESS_TOKEN_PROVIDER;
154+
static final PropertyDescriptor EVENT_HUB_IDENTITY_FEDERATION_TOKEN_PROVIDER = AzureEventHubComponent.IDENTITY_FEDERATION_TOKEN_PROVIDER;
153155
static final PropertyDescriptor ACCESS_POLICY_NAME = new PropertyDescriptor.Builder()
154156
.name("Shared Access Policy Name")
155157
.description("The name of the shared access policy. This policy must have Listen claims.")
@@ -267,6 +269,13 @@ public class ConsumeAzureEventHub extends AbstractSessionFactoryProcessor implem
267269
.required(true)
268270
.dependsOn(BLOB_STORAGE_AUTHENTICATION_STRATEGY, BlobStorageAuthenticationStrategy.OAUTH2)
269271
.build();
272+
static final PropertyDescriptor BLOB_STORAGE_IDENTITY_FEDERATION_TOKEN_PROVIDER = new PropertyDescriptor.Builder()
273+
.name("Storage Identity Federation Token Provider")
274+
.description("Controller Service exchanging workload identity tokens for Azure AD access tokens when using Identity Federation with Azure Blob Storage.")
275+
.identifiesControllerService(AzureIdentityFederationTokenProvider.class)
276+
.required(true)
277+
.dependsOn(BLOB_STORAGE_AUTHENTICATION_STRATEGY, BlobStorageAuthenticationStrategy.IDENTITY_FEDERATION)
278+
.build();
270279
static final PropertyDescriptor STORAGE_ACCOUNT_KEY = new PropertyDescriptor.Builder()
271280
.name("Storage Account Key")
272281
.description("The Azure Storage account key to store event hub consumer group state.")
@@ -326,6 +335,7 @@ public class ConsumeAzureEventHub extends AbstractSessionFactoryProcessor implem
326335
POLICY_PRIMARY_KEY,
327336
AUTHENTICATION_STRATEGY,
328337
EVENT_HUB_OAUTH2_ACCESS_TOKEN_PROVIDER,
338+
EVENT_HUB_IDENTITY_FEDERATION_TOKEN_PROVIDER,
329339
CONSUMER_GROUP,
330340
RECORD_READER,
331341
RECORD_WRITER,
@@ -340,6 +350,7 @@ public class ConsumeAzureEventHub extends AbstractSessionFactoryProcessor implem
340350
STORAGE_ACCOUNT_KEY,
341351
STORAGE_SAS_TOKEN,
342352
BLOB_STORAGE_OAUTH2_ACCESS_TOKEN_PROVIDER,
353+
BLOB_STORAGE_IDENTITY_FEDERATION_TOKEN_PROVIDER,
343354
PROXY_CONFIGURATION_SERVICE
344355
);
345356

@@ -435,6 +446,7 @@ protected Collection<ValidationResult> customValidate(ValidationContext validati
435446
final String storageSasToken = validationContext.getProperty(STORAGE_SAS_TOKEN).evaluateAttributeExpressions().getValue();
436447
final CheckpointStrategy checkpointStrategy = CheckpointStrategy.valueOf(validationContext.getProperty(CHECKPOINT_STRATEGY).getValue());
437448
final boolean blobOauthProviderSet = validationContext.getProperty(BLOB_STORAGE_OAUTH2_ACCESS_TOKEN_PROVIDER).isSet();
449+
final boolean blobIdentityFederationProviderSet = validationContext.getProperty(BLOB_STORAGE_IDENTITY_FEDERATION_TOKEN_PROVIDER).isSet();
438450

439451
if ((recordReader != null && recordWriter == null) || (recordReader == null && recordWriter != null)) {
440452
results.add(new ValidationResult.Builder()
@@ -527,9 +539,42 @@ protected Collection<ValidationResult> customValidate(ValidationContext validati
527539
.valid(false)
528540
.build());
529541
}
542+
if (blobIdentityFederationProviderSet) {
543+
results.add(new ValidationResult.Builder()
544+
.subject(BLOB_STORAGE_IDENTITY_FEDERATION_TOKEN_PROVIDER.getDisplayName())
545+
.explanation("%s must not be set when %s is %s."
546+
.formatted(BLOB_STORAGE_IDENTITY_FEDERATION_TOKEN_PROVIDER.getDisplayName(),
547+
BLOB_STORAGE_AUTHENTICATION_STRATEGY.getDisplayName(),
548+
BlobStorageAuthenticationStrategy.OAUTH2.getDisplayName()))
549+
.valid(false)
550+
.build());
551+
}
552+
} else if (blobStorageAuthenticationStrategy == BlobStorageAuthenticationStrategy.IDENTITY_FEDERATION) {
553+
if (StringUtils.isNotBlank(storageAccountKey)) {
554+
results.add(new ValidationResult.Builder()
555+
.subject(STORAGE_ACCOUNT_KEY.getDisplayName())
556+
.explanation("%s must not be set when %s is %s."
557+
.formatted(STORAGE_ACCOUNT_KEY.getDisplayName(),
558+
BLOB_STORAGE_AUTHENTICATION_STRATEGY.getDisplayName(),
559+
BlobStorageAuthenticationStrategy.IDENTITY_FEDERATION.getDisplayName()))
560+
.valid(false)
561+
.build());
562+
}
563+
564+
if (StringUtils.isNotBlank(storageSasToken)) {
565+
results.add(new ValidationResult.Builder()
566+
.subject(STORAGE_SAS_TOKEN.getDisplayName())
567+
.explanation("%s must not be set when %s is %s."
568+
.formatted(STORAGE_SAS_TOKEN.getDisplayName(),
569+
BLOB_STORAGE_AUTHENTICATION_STRATEGY.getDisplayName(),
570+
BlobStorageAuthenticationStrategy.IDENTITY_FEDERATION.getDisplayName()))
571+
.valid(false)
572+
.build());
573+
}
530574
}
531575
}
532-
results.addAll(AzureEventHubUtils.customValidate(ACCESS_POLICY_NAME, POLICY_PRIMARY_KEY, EVENT_HUB_OAUTH2_ACCESS_TOKEN_PROVIDER, validationContext));
576+
results.addAll(AzureEventHubUtils.customValidate(ACCESS_POLICY_NAME, POLICY_PRIMARY_KEY,
577+
EVENT_HUB_OAUTH2_ACCESS_TOKEN_PROVIDER, EVENT_HUB_IDENTITY_FEDERATION_TOKEN_PROVIDER, validationContext));
533578
return results;
534579
}
535580

@@ -627,6 +672,14 @@ protected EventProcessorClient createClient(final ProcessContext context) {
627672
blobContainerClientBuilder.endpoint(endpoint);
628673
blobContainerClientBuilder.credential(tokenCredential);
629674
}
675+
case IDENTITY_FEDERATION -> {
676+
final AzureIdentityFederationTokenProvider tokenProvider =
677+
context.getProperty(BLOB_STORAGE_IDENTITY_FEDERATION_TOKEN_PROVIDER).asControllerService(AzureIdentityFederationTokenProvider.class);
678+
final TokenCredential tokenCredential = AzureEventHubUtils.createTokenCredential(tokenProvider);
679+
final String endpoint = createBlobEndpoint(storageAccountName, domainName);
680+
blobContainerClientBuilder.endpoint(endpoint);
681+
blobContainerClientBuilder.credential(tokenCredential);
682+
}
630683
}
631684
blobContainerClientBuilder.containerName(containerName);
632685

@@ -682,6 +735,13 @@ protected EventProcessorClient createClient(final ProcessContext context) {
682735
final TokenCredential tokenCredential = AzureEventHubUtils.createTokenCredential(tokenProvider);
683736
eventProcessorClientBuilder.credential(fullyQualifiedNamespace, eventHubName, tokenCredential);
684737
}
738+
case IDENTITY_FEDERATION -> {
739+
final AzureIdentityFederationTokenProvider tokenProvider =
740+
context.getProperty(EVENT_HUB_IDENTITY_FEDERATION_TOKEN_PROVIDER)
741+
.asControllerService(AzureIdentityFederationTokenProvider.class);
742+
final TokenCredential tokenCredential = AzureEventHubUtils.createTokenCredential(tokenProvider);
743+
eventProcessorClientBuilder.credential(fullyQualifiedNamespace, eventHubName, tokenCredential);
744+
}
685745
}
686746

687747
final Integer prefetchCount = context.getProperty(PREFETCH_COUNT).evaluateAttributeExpressions().asInteger();
@@ -921,7 +981,7 @@ private String createStorageConnectionString(final ProcessContext context,
921981
String.format(FORMAT_STORAGE_CONNECTION_STRING_FOR_ACCOUNT_KEY, storageAccountName, storageAccountKey, domainName);
922982
case SHARED_ACCESS_SIGNATURE ->
923983
String.format(FORMAT_STORAGE_CONNECTION_STRING_FOR_SAS_TOKEN, storageAccountName, domainName, storageSasToken);
924-
case OAUTH2 -> throw new IllegalArgumentException(String.format(
984+
case OAUTH2, IDENTITY_FEDERATION -> throw new IllegalArgumentException(String.format(
925985
"Blob Storage Authentication Strategy %s does not support connection string authentication", blobStorageAuthenticationStrategy));
926986
};
927987
}

‎nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/src/main/java/org/apache/nifi/processors/azure/eventhub/GetAzureEventHub.java‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@
6060
import org.apache.nifi.shared.azure.eventhubs.AzureEventHubComponent;
6161
import org.apache.nifi.shared.azure.eventhubs.AzureEventHubTransportType;
6262
import org.apache.nifi.oauth2.OAuth2AccessTokenProvider;
63+
import org.apache.nifi.services.azure.AzureIdentityFederationTokenProvider;
6364
import org.apache.nifi.util.StopWatch;
6465

6566
import java.time.Duration;
@@ -116,6 +117,7 @@ public class GetAzureEventHub extends AbstractProcessor implements AzureEventHub
116117
static final PropertyDescriptor SERVICE_BUS_ENDPOINT = AzureEventHubUtils.SERVICE_BUS_ENDPOINT;
117118
static final PropertyDescriptor AUTHENTICATION_STRATEGY = AzureEventHubComponent.AUTHENTICATION_STRATEGY;
118119
static final PropertyDescriptor EVENT_HUB_OAUTH2_ACCESS_TOKEN_PROVIDER = AzureEventHubComponent.OAUTH2_ACCESS_TOKEN_PROVIDER;
120+
static final PropertyDescriptor EVENT_HUB_IDENTITY_FEDERATION_TOKEN_PROVIDER = AzureEventHubComponent.IDENTITY_FEDERATION_TOKEN_PROVIDER;
119121
static final PropertyDescriptor ACCESS_POLICY = new PropertyDescriptor.Builder()
120122
.name("Shared Access Policy Name")
121123
.description("The name of the shared access policy. This policy must have Listen claims.")
@@ -172,6 +174,7 @@ public class GetAzureEventHub extends AbstractProcessor implements AzureEventHub
172174
POLICY_PRIMARY_KEY,
173175
AUTHENTICATION_STRATEGY,
174176
EVENT_HUB_OAUTH2_ACCESS_TOKEN_PROVIDER,
177+
EVENT_HUB_IDENTITY_FEDERATION_TOKEN_PROVIDER,
175178
CONSUMER_GROUP,
176179
ENQUEUE_TIME,
177180
RECEIVER_FETCH_SIZE,
@@ -209,7 +212,8 @@ public final List<PropertyDescriptor> getSupportedPropertyDescriptors() {
209212

210213
@Override
211214
protected Collection<ValidationResult> customValidate(ValidationContext context) {
212-
return AzureEventHubUtils.customValidate(ACCESS_POLICY, POLICY_PRIMARY_KEY, EVENT_HUB_OAUTH2_ACCESS_TOKEN_PROVIDER, context);
215+
return AzureEventHubUtils.customValidate(ACCESS_POLICY, POLICY_PRIMARY_KEY,
216+
EVENT_HUB_OAUTH2_ACCESS_TOKEN_PROVIDER, EVENT_HUB_IDENTITY_FEDERATION_TOKEN_PROVIDER, context);
213217
}
214218

215219
@OnPrimaryNodeStateChange
@@ -435,6 +439,13 @@ private EventHubClientBuilder createEventHubClientBuilder(final ProcessContext c
435439
final TokenCredential tokenCredential = AzureEventHubUtils.createTokenCredential(tokenProvider);
436440
eventHubClientBuilder.credential(fullyQualifiedNamespace, eventHubName, tokenCredential);
437441
}
442+
case IDENTITY_FEDERATION -> {
443+
final AzureIdentityFederationTokenProvider tokenProvider =
444+
context.getProperty(EVENT_HUB_IDENTITY_FEDERATION_TOKEN_PROVIDER)
445+
.asControllerService(AzureIdentityFederationTokenProvider.class);
446+
final TokenCredential tokenCredential = AzureEventHubUtils.createTokenCredential(tokenProvider);
447+
eventHubClientBuilder.credential(fullyQualifiedNamespace, eventHubName, tokenCredential);
448+
}
438449
}
439450

440451
// Set Azure Event Hub Client Identifier using Processor Identifier instead of default random UUID

‎nifi-extension-bundles/nifi-azure-bundle/nifi-azure-processors/src/main/java/org/apache/nifi/processors/azure/eventhub/PutAzureEventHub.java‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353
import org.apache.nifi.shared.azure.eventhubs.AzureEventHubAuthenticationStrategy;
5454
import org.apache.nifi.shared.azure.eventhubs.AzureEventHubComponent;
5555
import org.apache.nifi.shared.azure.eventhubs.AzureEventHubTransportType;
56+
import org.apache.nifi.services.azure.AzureIdentityFederationTokenProvider;
5657
import org.apache.nifi.stream.io.StreamUtils;
5758
import org.apache.nifi.util.StopWatch;
5859

@@ -89,6 +90,7 @@ public class PutAzureEventHub extends AbstractProcessor implements AzureEventHub
8990
static final PropertyDescriptor SERVICE_BUS_ENDPOINT = AzureEventHubUtils.SERVICE_BUS_ENDPOINT;
9091
static final PropertyDescriptor AUTHENTICATION_STRATEGY = AzureEventHubComponent.AUTHENTICATION_STRATEGY;
9192
static final PropertyDescriptor EVENT_HUB_OAUTH2_ACCESS_TOKEN_PROVIDER = AzureEventHubComponent.OAUTH2_ACCESS_TOKEN_PROVIDER;
93+
static final PropertyDescriptor EVENT_HUB_IDENTITY_FEDERATION_TOKEN_PROVIDER = AzureEventHubComponent.IDENTITY_FEDERATION_TOKEN_PROVIDER;
9294
static final PropertyDescriptor ACCESS_POLICY = new PropertyDescriptor.Builder()
9395
.name("Shared Access Policy Name")
9496
.description("The name of the shared access policy. This policy must have Send claims.")
@@ -133,6 +135,7 @@ public class PutAzureEventHub extends AbstractProcessor implements AzureEventHub
133135
POLICY_PRIMARY_KEY,
134136
AUTHENTICATION_STRATEGY,
135137
EVENT_HUB_OAUTH2_ACCESS_TOKEN_PROVIDER,
138+
EVENT_HUB_IDENTITY_FEDERATION_TOKEN_PROVIDER,
136139
PARTITIONING_KEY_ATTRIBUTE_NAME,
137140
MAX_BATCH_SIZE,
138141
PROXY_CONFIGURATION_SERVICE
@@ -171,7 +174,8 @@ public void closeClient() {
171174

172175
@Override
173176
protected Collection<ValidationResult> customValidate(ValidationContext context) {
174-
return AzureEventHubUtils.customValidate(ACCESS_POLICY, POLICY_PRIMARY_KEY, EVENT_HUB_OAUTH2_ACCESS_TOKEN_PROVIDER, context);
177+
return AzureEventHubUtils.customValidate(ACCESS_POLICY, POLICY_PRIMARY_KEY,
178+
EVENT_HUB_OAUTH2_ACCESS_TOKEN_PROVIDER, EVENT_HUB_IDENTITY_FEDERATION_TOKEN_PROVIDER, context);
175179
}
176180

177181
@Override
@@ -259,6 +263,13 @@ protected EventHubProducerClient createEventHubProducerClient(final ProcessConte
259263
final TokenCredential tokenCredential = AzureEventHubUtils.createTokenCredential(tokenProvider);
260264
eventHubClientBuilder.credential(fullyQualifiedNamespace, eventHubName, tokenCredential);
261265
}
266+
case IDENTITY_FEDERATION -> {
267+
final AzureIdentityFederationTokenProvider tokenProvider =
268+
context.getProperty(EVENT_HUB_IDENTITY_FEDERATION_TOKEN_PROVIDER)
269+
.asControllerService(AzureIdentityFederationTokenProvider.class);
270+
final TokenCredential tokenCredential = AzureEventHubUtils.createTokenCredential(tokenProvider);
271+
eventHubClientBuilder.credential(fullyQualifiedNamespace, eventHubName, tokenCredential);
272+
}
262273
}
263274
AzureEventHubUtils.getProxyOptions(context).ifPresent(eventHubClientBuilder::proxyOptions);
264275
return eventHubClientBuilder.buildProducerClient();

0 commit comments

Comments
 (0)