NIFI-8124 Added Cookie Strategy property to InvokeHTTP

This closes #5887

Signed-off-by: David Handermann <exceptionfactory@apache.org>

nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/pom.xml

@@ -256,6 +256,10 @@
             <groupId>com.squareup.okhttp3</groupId>
             <artifactId>okhttp</artifactId>
         </dependency>
+        <dependency>
+            <groupId>com.squareup.okhttp3</groupId>
+            <artifactId>okhttp-urlconnection</artifactId>
+        </dependency>
         <dependency>
             <groupId>io.github.rburgst</groupId>
             <artifactId>okhttp-digest</artifactId>

nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/InvokeHTTP.java

@@ -23,6 +23,8 @@ import com.burgstaller.okhttp.digest.DigestAuthenticator;
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
+import java.net.CookieManager;
+import java.net.CookiePolicy;
 import java.net.Proxy;
 import java.net.Proxy.Type;
 import java.net.URL;
@@ -60,6 +62,7 @@ import okhttp3.ConnectionPool;
 import okhttp3.Credentials;
 import okhttp3.Handshake;
 import okhttp3.Headers;
+import okhttp3.JavaNetCookieJar;
 import okhttp3.MediaType;
 import okhttp3.MultipartBody;
 import okhttp3.MultipartBody.Builder;
@@ -100,6 +103,7 @@ import org.apache.nifi.processor.Relationship;
 import org.apache.nifi.processor.exception.ProcessException;
 import org.apache.nifi.processor.util.StandardValidators;
 import org.apache.nifi.processors.standard.http.FlowFileNamingStrategy;
+import org.apache.nifi.processors.standard.http.CookieStrategy;
 import org.apache.nifi.processors.standard.util.ProxyAuthenticator;
 import org.apache.nifi.processors.standard.util.SoftLimitBoundedByteArrayOutputStream;
 import org.apache.nifi.proxy.ProxyConfiguration;
@@ -516,6 +520,15 @@ public class InvokeHTTP extends AbstractProcessor {
             )
             .build();
 
+    public static final PropertyDescriptor PROP_COOKIE_STRATEGY = new PropertyDescriptor.Builder()
+            .name("cookie-strategy")
+            .description("Strategy for accepting and persisting HTTP cookies. Accepting cookies enables persistence across multiple requests.")
+            .displayName("Cookie Strategy")
+            .required(true)
+            .defaultValue(CookieStrategy.DISABLED.name())
+            .allowableValues(CookieStrategy.values())
+            .build();
+
     private static final ProxySpec[] PROXY_SPECS = {ProxySpec.HTTP_AUTH, ProxySpec.SOCKS};
     public static final PropertyDescriptor PROXY_CONFIGURATION_SERVICE
             = ProxyConfiguration.createProxyConfigPropertyDescriptor(true, PROXY_SPECS);
@@ -530,6 +543,7 @@ public class InvokeHTTP extends AbstractProcessor {
             PROP_MAX_IDLE_CONNECTIONS,
             PROP_DATE_HEADER,
             PROP_FOLLOW_REDIRECTS,
+            PROP_COOKIE_STRATEGY,
             DISABLE_HTTP2_PROTOCOL,
             FLOW_FILE_NAMING_STRATEGY,
             PROP_ATTRIBUTES_TO_SEND,
@@ -824,6 +838,17 @@ public class InvokeHTTP extends AbstractProcessor {
             okHttpClientBuilder.sslSocketFactory(socketFactory, trustManager);
         }
 
+        // Configure cookie strategy
+        switch(CookieStrategy.valueOf(context.getProperty(PROP_COOKIE_STRATEGY).getValue())) {
+            case DISABLED:
+                break;
+            case ACCEPT_ALL:
+                final CookieManager cookieManager = new CookieManager();
+                cookieManager.setCookiePolicy(CookiePolicy.ACCEPT_ALL);
+                okHttpClientBuilder.cookieJar(new JavaNetCookieJar(cookieManager));
+                break;
+        }
+
         setAuthenticator(okHttpClientBuilder, context);
 
         useChunked = context.getProperty(PROP_USE_CHUNKED_ENCODING).asBoolean();

nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/http/CookieStrategy.java

@@ -0,0 +1,32 @@
+/*
+ * 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.processors.standard.http;
+
+public enum CookieStrategy {
+    DISABLED("A cookie strategy that ignores cookies."),
+    ACCEPT_ALL("A cookie strategy that stores all cookies from incoming HTTP responses.");
+
+    private final String description;
+
+    CookieStrategy(final String description) {
+        this.description = description;
+    }
+
+    public String getDescription() {
+        return description;
+    }
+}

nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/InvokeHTTPTest.java

@@ -24,6 +24,7 @@ import org.apache.nifi.flowfile.attributes.CoreAttributes;
 import org.apache.nifi.oauth2.OAuth2AccessTokenProvider;
 import org.apache.nifi.processor.Relationship;
 import org.apache.nifi.processors.standard.http.FlowFileNamingStrategy;
+import org.apache.nifi.processors.standard.http.CookieStrategy;
 import org.apache.nifi.reporting.InitializationException;
 import org.apache.nifi.security.util.StandardTlsConfiguration;
 import org.apache.nifi.security.util.TemporaryKeyStoreBuilder;
@@ -104,6 +105,14 @@ public class InvokeHTTPTest {
 
     private static final String LOCATION_HEADER = "Location";
 
+    private static final String SET_COOKIE_HEADER = "Set-Cookie";
+
+    private static final String COOKIE_HEADER = "Cookie";
+
+    private static final String COOKIE_1 = "a=apple";
+
+    private static final String COOKIE_2 = "b=banana";
+
     private static final String TRANSFER_ENCODING_HEADER = "Transfer-Encoding";
 
     private static final String USER_AGENT_HEADER = "User-Agent";
@@ -603,6 +612,46 @@ public class InvokeHTTPTest {
         assertRelationshipStatusCodeEquals(InvokeHTTP.REL_NO_RETRY, HTTP_MOVED_TEMP);
     }
 
+    @Test
+    public void testRunGetHttp302CookieStrategyAcceptAll() throws InterruptedException {
+        runner.setProperty(InvokeHTTP.PROP_COOKIE_STRATEGY, CookieStrategy.ACCEPT_ALL.name());
+        mockWebServer.enqueue(new MockResponse().setResponseCode(HTTP_MOVED_TEMP)
+            .addHeader(SET_COOKIE_HEADER, COOKIE_1)
+            .addHeader(SET_COOKIE_HEADER, COOKIE_2)
+            .addHeader(LOCATION_HEADER, getMockWebServerUrl()));
+        enqueueResponseCodeAndRun(HTTP_OK);
+
+        RecordedRequest request1 = mockWebServer.takeRequest();
+        assertNull(request1.getHeader(COOKIE_HEADER));
+
+        RecordedRequest request2 = mockWebServer.takeRequest();
+        final String expectedHeader = String.format("%s; %s", COOKIE_1, COOKIE_2);
+        assertEquals(expectedHeader, request2.getHeader(COOKIE_HEADER));
+
+        runner.assertTransferCount(InvokeHTTP.REL_FAILURE, 0);
+        runner.assertTransferCount(InvokeHTTP.REL_NO_RETRY, 0);
+        assertRelationshipStatusCodeEquals(InvokeHTTP.REL_RESPONSE, HTTP_OK);
+    }
+
+    @Test
+    public void testRunGetHttp302CookieStrategyDefaultDisabled() throws InterruptedException {
+        mockWebServer.enqueue(new MockResponse().setResponseCode(HTTP_MOVED_TEMP)
+            .addHeader(SET_COOKIE_HEADER, COOKIE_1)
+            .addHeader(SET_COOKIE_HEADER, COOKIE_2)
+            .addHeader(LOCATION_HEADER, getMockWebServerUrl()));
+        enqueueResponseCodeAndRun(HTTP_OK);
+
+        RecordedRequest request1 = mockWebServer.takeRequest();
+        assertNull(request1.getHeader(COOKIE_HEADER));
+
+        RecordedRequest request2 = mockWebServer.takeRequest();
+        assertNull(request2.getHeader(COOKIE_HEADER));
+
+        runner.assertTransferCount(InvokeHTTP.REL_FAILURE, 0);
+        runner.assertTransferCount(InvokeHTTP.REL_NO_RETRY, 0);
+        assertRelationshipStatusCodeEquals(InvokeHTTP.REL_RESPONSE, HTTP_OK);
+    }
+
     @Test
     public void testRunGetHttp400NoRetryMinimumProperties() {
         enqueueResponseCodeAndRun(HTTP_BAD_REQUEST);