NIFI-16078 - Fix StandardHashiCorpVaultClientService for non-token authentication

nifi-commons/nifi-hashicorp-vault/src/main/java/org/apache/nifi/vault/hashicorp/StandardHashiCorpVaultCommunicationService.java

@@ -107,6 +107,8 @@ public StandardHashiCorpVaultCommunicationService(final PropertySource<?>... pro
         final VaultEndpoint vaultEndpoint = vaultConfiguration.vaultEndpoint();
         final RestTemplate restTemplate = VaultClients.createRestTemplate(vaultEndpoint, clientHttpRequestFactory);
 
+        vaultConfiguration.setClientHttpRequestFactory(clientHttpRequestFactory);
+
         final VaultClient.Builder vaultClientBuilder = VaultClient.builder(restTemplate).endpoint(vaultEndpoint);
         final String namespace = vaultConfiguration.getNamespace();
         if (namespace != null && !namespace.isEmpty()) {

nifi-commons/nifi-hashicorp-vault/src/main/java/org/apache/nifi/vault/hashicorp/config/HashiCorpVaultConfiguration.java

@@ -23,6 +23,9 @@
 import org.springframework.core.env.StandardEnvironment;
 import org.springframework.core.io.FileSystemResource;
 import org.springframework.core.io.support.ResourcePropertySource;
+import org.springframework.http.client.ClientHttpRequestFactory;
+import org.springframework.vault.client.RestTemplateFactory;
+import org.springframework.vault.config.AbstractVaultConfiguration;
 import org.springframework.vault.config.EnvironmentVaultConfiguration;
 import org.springframework.vault.core.VaultKeyValueOperationsSupport.KeyValueBackend;
 import org.springframework.vault.support.ClientOptions;
@@ -65,9 +68,14 @@ public String getKey() {
 
     private static final String HTTPS = "https";
 
+    private static final String CLIENT_HTTP_REQUEST_FACTORY_WRAPPER_BEAN = "clientHttpRequestFactoryWrapper";
+
     private final SslConfiguration sslConfiguration;
     private final ClientOptions clientOptions;
     private final KeyValueBackend keyValueBackend;
+    private final HashiCorpVaultApplicationContext applicationContext;
+
+    private AbstractVaultConfiguration.ClientFactoryWrapper clientFactoryWrapper;
 
     /**
      * Creates a HashiCorpVaultConfiguration from property sources, in increasing precedence.
@@ -118,7 +126,8 @@ public HashiCorpVaultConfiguration(final ConfigurableEnvironment env, final Prop
         }
         this.keyValueBackend = keyValueBackend;
 
-        this.setApplicationContext(new HashiCorpVaultApplicationContext(env));
+        this.applicationContext = new HashiCorpVaultApplicationContext(env);
+        this.setApplicationContext(applicationContext);
 
         sslConfiguration = env.getProperty(VaultConfigurationKey.URI.key).contains(HTTPS)
                 ? super.sslConfiguration() : SslConfiguration.unconfigured();
@@ -130,6 +139,23 @@ public KeyValueBackend getKeyValueBackend() {
         return keyValueBackend;
     }
 
+    /**
+     * Registers the given request factory as the {@code clientHttpRequestFactoryWrapper} bean and uses it
+     * to build the {@link RestTemplateFactory}. Spring Vault authentication methods other than token
+     * authentication build their own Vault client from these, so this must be called before
+     * {@link #clientAuthentication()} to supply NiFi's configured request factory.
+     * @param clientHttpRequestFactory The request factory used for Vault communication
+     */
+    public void setClientHttpRequestFactory(final ClientHttpRequestFactory clientHttpRequestFactory) {
+        this.clientFactoryWrapper = new AbstractVaultConfiguration.ClientFactoryWrapper(clientHttpRequestFactory);
+        applicationContext.getBeanFactory().registerSingleton(CLIENT_HTTP_REQUEST_FACTORY_WRAPPER_BEAN, clientFactoryWrapper);
+    }
+
+    @Override
+    protected RestTemplateFactory getRestTemplateFactory() {
+        return clientFactoryWrapper == null ? super.getRestTemplateFactory() : restTemplateFactory(clientFactoryWrapper);
+    }
+
     /**
      * A convenience method to create a PropertySource from a file on disk.
      * @param filename The properties filename.

nifi-commons/nifi-hashicorp-vault/src/test/java/org/apache/nifi/vault/hashicorp/TestHashiCorpVaultConfiguration.java

@@ -25,6 +25,7 @@
 import org.junit.jupiter.api.Test;
 import org.springframework.core.env.ConfigurableEnvironment;
 import org.springframework.core.env.StandardEnvironment;
+import org.springframework.http.client.SimpleClientHttpRequestFactory;
 import org.springframework.vault.authentication.ClientAuthentication;
 import org.springframework.vault.client.VaultEndpoint;
 import org.springframework.vault.core.VaultKeyValueOperationsSupport;
@@ -148,6 +149,60 @@ public void testBasicProperties() {
         this.runTest("http");
     }
 
+    @Test
+    public void testKubernetesClientAuthentication() throws IOException {
+        File k8sAuthProps = null;
+        Path tokenFile = null;
+        try {
+            tokenFile = Files.createTempFile("vault-k8s-token", ".jwt");
+            Files.writeString(tokenFile, "test-jwt-token");
+
+            final Map<String, String> props = new HashMap<>();
+            props.put(VAULT_AUTHENTICATION, "KUBERNETES");
+            props.put("vault.kubernetes.role", "test-role");
+            props.put("vault.kubernetes.service-account-token-file", tokenFile.toFile().getAbsolutePath());
+            k8sAuthProps = writeVaultAuthProperties(props);
+            propertiesBuilder.setAuthPropertiesFilename(k8sAuthProps.getAbsolutePath());
+
+            config = new HashiCorpVaultConfiguration(
+                    createIsolatedEnvironment(), new HashiCorpVaultPropertySource(propertiesBuilder.build()));
+            config.setClientHttpRequestFactory(new SimpleClientHttpRequestFactory());
+
+            final ClientAuthentication clientAuthentication = config.clientAuthentication();
+            assertNotNull(clientAuthentication);
+        } finally {
+            if (k8sAuthProps != null) {
+                Files.deleteIfExists(k8sAuthProps.toPath());
+            }
+            if (tokenFile != null) {
+                Files.deleteIfExists(tokenFile);
+            }
+        }
+    }
+
+    @Test
+    public void testAwsEc2ClientAuthentication() throws IOException {
+        File awsAuthProps = null;
+        try {
+            final Map<String, String> props = new HashMap<>();
+            props.put(VAULT_AUTHENTICATION, "AWS_EC2");
+            props.put("vault.aws-ec2.role", "test-role");
+            awsAuthProps = writeVaultAuthProperties(props);
+            propertiesBuilder.setAuthPropertiesFilename(awsAuthProps.getAbsolutePath());
+
+            config = new HashiCorpVaultConfiguration(
+                    createIsolatedEnvironment(), new HashiCorpVaultPropertySource(propertiesBuilder.build()));
+            config.setClientHttpRequestFactory(new SimpleClientHttpRequestFactory());
+
+            final ClientAuthentication clientAuthentication = config.clientAuthentication();
+            assertNotNull(clientAuthentication);
+        } finally {
+            if (awsAuthProps != null) {
+                Files.deleteIfExists(awsAuthProps.toPath());
+            }
+        }
+    }
+
     @Test
     public void testKvVersion() {
         config = new HashiCorpVaultConfiguration(new HashiCorpVaultPropertySource(propertiesBuilder.build()));