From ee24df2830b1d880428f487081c0ea92b0ca0ca1 Mon Sep 17 00:00:00 2001
From: Paul Grey <greyp9@yahoo.com>
Date: Wed, 12 Apr 2023 23:44:41 -0400
Subject: [PATCH] NIFI-11439 Added Storage API URL property to GCS Processors

- Included Host Header override with Storage API URL based on Google Private Service Connect documentation

This closes #7172

Signed-off-by: David Handermann <exceptionfactory@apache.org>
---
 .../gcp/storage/AbstractGCSProcessor.java     | 25 ++++++++++++++++++-
 .../gcp/storage/AbstractGCSTest.java          | 19 ++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/storage/AbstractGCSProcessor.java b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/storage/AbstractGCSProcessor.java
index 439a2be8ba..bbfaa7e966 100644
--- a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/storage/AbstractGCSProcessor.java
+++ b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/storage/AbstractGCSProcessor.java
@@ -17,18 +17,22 @@
 package org.apache.nifi.processors.gcp.storage;
 
 import com.google.api.gax.retrying.RetrySettings;
+import com.google.api.gax.rpc.FixedHeaderProvider;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.BaseServiceException;
 import com.google.cloud.storage.Storage;
 import com.google.cloud.storage.StorageOptions;
+import com.google.common.collect.ImmutableMap;
 import org.apache.nifi.components.ConfigVerificationResult;
 import org.apache.nifi.components.PropertyDescriptor;
 import org.apache.nifi.components.ValidationContext;
 import org.apache.nifi.components.ValidationResult;
+import org.apache.nifi.expression.ExpressionLanguageScope;
 import org.apache.nifi.logging.ComponentLog;
 import org.apache.nifi.processor.ProcessContext;
 import org.apache.nifi.processor.Relationship;
 import org.apache.nifi.processor.VerifiableProcessor;
+import org.apache.nifi.processor.util.StandardValidators;
 import org.apache.nifi.processors.gcp.AbstractGCPProcessor;
 import org.apache.nifi.processors.gcp.ProxyAwareTransportFactory;
 import org.apache.nifi.proxy.ProxyConfiguration;
@@ -65,9 +69,21 @@ public abstract class AbstractGCSProcessor extends AbstractGCPProcessor<Storage,
         return relationships;
     }
 
+    // https://cloud.google.com/storage/docs/request-endpoints#storage-set-client-endpoint-java
+    public static final PropertyDescriptor STORAGE_API_URL = new PropertyDescriptor
+            .Builder().name("storage-api-url")
+            .displayName("Storage API URL")
+            .description("Overrides the default storage URL. Configuring an alternative Storage API URL also overrides the HTTP Host header on requests as described in the Google documentation for Private Service Connections.")
+            .addValidator(StandardValidators.URL_VALIDATOR)
+            .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY)
+            .required(false)
+            .build();
+
     @Override
     public List<PropertyDescriptor> getSupportedPropertyDescriptors() {
-        return Collections.unmodifiableList(super.getSupportedPropertyDescriptors());
+        final List<PropertyDescriptor> propertyDescriptors = new ArrayList<>(super.getSupportedPropertyDescriptors());
+        propertyDescriptors.add(STORAGE_API_URL);
+        return Collections.unmodifiableList(propertyDescriptors);
     }
 
     @Override
@@ -129,6 +145,7 @@ public abstract class AbstractGCSProcessor extends AbstractGCPProcessor<Storage,
     @Override
     protected StorageOptions getServiceOptions(ProcessContext context, GoogleCredentials credentials) {
         final String projectId = context.getProperty(PROJECT_ID).evaluateAttributeExpressions().getValue();
+        final String storageApiUrl = context.getProperty(STORAGE_API_URL).evaluateAttributeExpressions().getValue();
         final Integer retryCount = context.getProperty(RETRY_COUNT).asInteger();
 
         StorageOptions.Builder storageOptionsBuilder = StorageOptions.newBuilder()
@@ -141,6 +158,12 @@ public abstract class AbstractGCSProcessor extends AbstractGCPProcessor<Storage,
             storageOptionsBuilder.setProjectId(projectId);
         }
 
+        if (storageApiUrl != null && !storageApiUrl.isEmpty()) {
+            storageOptionsBuilder.setHost(storageApiUrl);
+            // https://codelabs.developers.google.com/cloudnet-psc#12
+            storageOptionsBuilder.setHeaderProvider(FixedHeaderProvider.create(ImmutableMap.of("Host", "www.googleapis.com")));
+        }
+
         return  storageOptionsBuilder.setTransportOptions(getTransportOptions(context)).build();
     }
 }
diff --git a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/storage/AbstractGCSTest.java b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/storage/AbstractGCSTest.java
index a35db57898..cdaaada1e5 100644
--- a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/storage/AbstractGCSTest.java
+++ b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/storage/AbstractGCSTest.java
@@ -45,6 +45,7 @@ import static org.mockito.Mockito.reset;
 @ExtendWith(MockitoExtension.class)
 public abstract class AbstractGCSTest {
     private static final String PROJECT_ID = System.getProperty("test.gcp.project.id", "nifi-test-gcp-project");
+    private static final String DEFAULT_STORAGE_URL = "https://storage.googleapis.com";
     private static final Integer RETRIES = 9;
 
     static final String BUCKET = RemoteStorageHelper.generateBucketName();
@@ -89,9 +90,27 @@ public abstract class AbstractGCSTest {
                 mockCredentials);
 
         assertEquals(PROJECT_ID, options.getProjectId(), "Project IDs should match");
+        assertEquals(DEFAULT_STORAGE_URL, options.getHost(), "Host URLs should match");
 
         assertEquals(RETRIES.intValue(), options.getRetrySettings().getMaxAttempts(), "Retry counts should match");
 
         assertSame(mockCredentials, options.getCredentials(), "Credentials should be configured correctly");
     }
+
+    @Test
+    public void testStorageOptionsConfigurationHostOverride() throws Exception {
+        reset(storage);
+        final TestRunner runner = buildNewRunner(getProcessor());
+
+        final String overrideStorageApiUrl = "https://localhost";
+        runner.setProperty(AbstractGCSProcessor.STORAGE_API_URL, overrideStorageApiUrl);
+
+        final AbstractGCSProcessor processor = getProcessor();
+        final GoogleCredentials mockCredentials = mock(GoogleCredentials.class);
+
+        final StorageOptions options = processor.getServiceOptions(runner.getProcessContext(),
+                mockCredentials);
+
+        assertEquals(overrideStorageApiUrl, options.getHost(), "Host URLs should match");
+    }
 }