Added regex support for attribute header selection on HandleHTTPResponse

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

@@ -80,6 +80,12 @@ public class HandleHttpResponse extends AbstractProcessor {
             .required(true)
             .identifiesControllerService(HttpContextMap.class)
             .build();
+    public static final PropertyDescriptor ATTRIBUTES_AS_HEADERS_REGEX = new PropertyDescriptor.Builder()
+            .name("Attributes to add to the HTTP Response (Regex)")
+            .description("Specifies the Regular Expression that determines the names of FlowFile attributes that should be added to the HTTP response")
+            .addValidator(StandardValidators.REGULAR_EXPRESSION_VALIDATOR)
+            .required(false)
+            .build();
 
     public static final Relationship REL_SUCCESS = new Relationship.Builder()
             .name("success")
@@ -96,6 +102,7 @@ public class HandleHttpResponse extends AbstractProcessor {
         final List<PropertyDescriptor> properties = new ArrayList<>();
         properties.add(STATUS_CODE);
         properties.add(HTTP_CONTEXT_MAP);
+        properties.add(ATTRIBUTES_AS_HEADERS_REGEX);
         return properties;
     }
 
@@ -166,6 +173,21 @@ public class HandleHttpResponse extends AbstractProcessor {
             }
         }
 
+        final String attributeHeaderRegex = context.getProperty(ATTRIBUTES_AS_HEADERS_REGEX).getValue();
+        if (attributeHeaderRegex != null) {
+            final Pattern pattern = Pattern.compile(attributeHeaderRegex);
+
+            final Map<String, String> attributes = flowFile.getAttributes();
+            for (final Map.Entry<String, String> entry : attributes.entrySet()) {
+                final String key = entry.getKey();
+                if (pattern.matcher(key).matches()) {
+                    if (!entry.getValue().trim().isEmpty()){
+                        response.setHeader(entry.getKey(), entry.getValue());
+                    }
+                }
+            }
+        }
+
         try {
             session.exportTo(flowFile, response.getOutputStream());
             response.flushBuffer();

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

@@ -93,6 +93,45 @@ public class TestHandleHttpResponse {
         assertTrue(contextMap.headersWithNoValue.isEmpty());
     }
 
+    @Test
+    public void testRegexHeaders() throws InitializationException {
+        final TestRunner runner = TestRunners.newTestRunner(HandleHttpResponse.class);
+
+        final MockHttpContextMap contextMap = new MockHttpContextMap("my-id", "");
+        runner.addControllerService("http-context-map", contextMap);
+        runner.enableControllerService(contextMap);
+        runner.setProperty(HandleHttpResponse.HTTP_CONTEXT_MAP, "http-context-map");
+        runner.setProperty(HandleHttpResponse.STATUS_CODE, "${status.code}");
+        runner.setProperty(HandleHttpResponse.ATTRIBUTES_AS_HEADERS_REGEX, "^(my.*)$");
+
+        final Map<String, String> attributes = new HashMap<>();
+        attributes.put(HTTPUtils.HTTP_CONTEXT_ID, "my-id");
+        attributes.put(HTTPUtils.HTTP_REQUEST_URI, "/test");
+        attributes.put(HTTPUtils.HTTP_LOCAL_NAME, "server");
+        attributes.put(HTTPUtils.HTTP_PORT, "8443");
+        attributes.put(HTTPUtils.HTTP_REMOTE_HOST, "client");
+        attributes.put(HTTPUtils.HTTP_SSL_CERT, "sslDN");
+        attributes.put("my-attr", "hello");
+        attributes.put("my-blank-attr", "");
+        attributes.put("status.code", "201");
+
+        runner.enqueue("hello".getBytes(), attributes);
+
+        runner.run();
+
+        runner.assertAllFlowFilesTransferred(HandleHttpResponse.REL_SUCCESS, 1);
+        assertTrue(runner.getProvenanceEvents().size() == 1);
+        assertEquals(ProvenanceEventType.SEND, runner.getProvenanceEvents().get(0).getEventType());
+        assertEquals("https://client@server:8443/test", runner.getProvenanceEvents().get(0).getTransitUri());
+
+        assertEquals("hello", contextMap.baos.toString());
+        assertEquals("hello", contextMap.headersSent.get("my-attr"));
+        assertNull(contextMap.headersSent.get("my-blank-attr"));
+        assertEquals(201, contextMap.statusCode);
+        assertEquals(1, contextMap.getCompletionCount());
+        assertTrue(contextMap.headersWithNoValue.isEmpty());
+    }
+
     @Test
     public void testWithExceptionThrown() throws InitializationException {
         final TestRunner runner = TestRunners.newTestRunner(HandleHttpResponse.class);