diff --git a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/HandleHttpResponse.java b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/HandleHttpResponse.java index 2a5c1a6abd..41c6ecdec5 100644 --- a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/HandleHttpResponse.java +++ b/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 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 attributes = flowFile.getAttributes(); + for (final Map.Entry 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(); diff --git a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestHandleHttpResponse.java b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestHandleHttpResponse.java index 1925f7a603..02b1d1c8df 100644 --- a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestHandleHttpResponse.java +++ b/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 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);