NIFI-8639: Add incoming flowfile to ConnectWebSocket processor to configure custom headers and dynamic URL through flowfile attributes in JettyWebSocketClient service.

This resolves #5130.

Signed-off-by: Tamas Palfy <tamas.bertalan.palfy@gmail.com>

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

@@ -52,5 +52,11 @@
             <version>1.15.0-SNAPSHOT</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.apache.nifi</groupId>
+            <artifactId>nifi-websocket-services-jetty</artifactId>
+            <version>1.14.0-SNAPSHOT</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 </project>

nifi-nar-bundles/nifi-websocket-bundle/nifi-websocket-processors/src/main/java/org/apache/nifi/processors/websocket/AbstractWebSocketGatewayProcessor.java

@@ -43,6 +43,7 @@ import java.io.IOException;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 
 import static org.apache.nifi.processors.websocket.WebSocketProcessorAttributes.ATTR_WS_CS_ID;
@@ -89,6 +90,7 @@ public abstract class AbstractWebSocketGatewayProcessor extends AbstractSessionF
         logger = getLogger();
     }
 
+    @FunctionalInterface
     public interface WebSocketFunction {
         void execute(final WebSocketService webSocketService) throws IOException, WebSocketConfigurationException;
     }
@@ -118,12 +120,24 @@ public abstract class AbstractWebSocketGatewayProcessor extends AbstractSessionF
 
     // @OnScheduled can not report error messages well on bulletin since it's an async method.
     // So, let's do it in onTrigger().
-    public void onWebSocketServiceReady(final WebSocketService webSocketService) throws IOException {
+    public void onWebSocketServiceReady(final WebSocketService webSocketService, final ProcessContext context) throws IOException {
-
         if (webSocketService instanceof WebSocketClientService) {
             // If it's a ws client, then connect to the remote here.
             // Otherwise, ws server is already started at WebSocketServerService
-            ((WebSocketClientService) webSocketService).connect(endpointId);
+            WebSocketClientService webSocketClientService = (WebSocketClientService) webSocketService;
+            if (context.hasIncomingConnection()) {
+                final ProcessSession session = processSessionFactory.createSession();
+                final FlowFile flowFile = session.get();
+                final Map<String, String> attributes = flowFile.getAttributes();
+                try {
+                    webSocketClientService.connect(endpointId, attributes);
+                } finally {
+                    session.remove(flowFile);
+                    session.commitAsync();
+                }
+            } else {
+                webSocketClientService.connect(endpointId);
+            }
         }
 
     }
@@ -137,6 +151,7 @@ public abstract class AbstractWebSocketGatewayProcessor extends AbstractSessionF
     }
 
     protected abstract WebSocketService getWebSocketService(final ProcessContext context);
+
     protected abstract String getEndpointId(final ProcessContext context);
 
     protected boolean isProcessorRegisteredToService() {
@@ -146,7 +161,7 @@ public abstract class AbstractWebSocketGatewayProcessor extends AbstractSessionF
     }
 
     @OnStopped
-    public void onStopped(final ProcessContext context) throws IOException {
+    public void onStopped(final ProcessContext context) {
         deregister();
     }
 
@@ -165,27 +180,36 @@ public abstract class AbstractWebSocketGatewayProcessor extends AbstractSessionF
     }
 
     @Override
-    public final void onTrigger(final ProcessContext context, final ProcessSessionFactory sessionFactory) throws ProcessException {
+    public final void onTrigger(final ProcessContext context, final ProcessSessionFactory sessionFactory) {
         if (processSessionFactory == null) {
             processSessionFactory = sessionFactory;
         }
 
         if (!isProcessorRegisteredToService()) {
             try {
-                registerProcessorToService(context, webSocketService -> onWebSocketServiceReady(webSocketService));
+                registerProcessorToService(context, webSocketService -> onWebSocketServiceReady(webSocketService, context));
-            } catch (IOException|WebSocketConfigurationException e) {
+            } catch (IOException | WebSocketConfigurationException e) {
                 // Deregister processor if it failed so that it can retry next onTrigger.
                 deregister();
                 context.yield();
                 throw new ProcessException("Failed to register processor to WebSocket service due to: " + e, e);
             }
+
+        } else {
+            try {
+                onWebSocketServiceReady(webSocketService, context);
+            } catch (IOException e) {
+                deregister();
+                context.yield();
+                throw new ProcessException("Failed to renew session and connect to WebSocket service due to: " + e, e);
+            }
         }
 
         context.yield();//nothing really to do here since handling WebSocket messages is done at ControllerService.
     }
 
 
-    private void enqueueMessage(final WebSocketMessage incomingMessage){
+    private void enqueueMessage(final WebSocketMessage incomingMessage) {
         final ProcessSession session = processSessionFactory.createSession();
         try {
             FlowFile messageFlowFile = session.create();
@@ -206,9 +230,9 @@ public abstract class AbstractWebSocketGatewayProcessor extends AbstractSessionF
 
             final byte[] payload = incomingMessage.getPayload();
             if (payload != null) {
-                messageFlowFile = session.write(messageFlowFile, out -> {
+                messageFlowFile = session.write(messageFlowFile, out ->
-                    out.write(payload, incomingMessage.getOffset(), incomingMessage.getLength());
+                        out.write(payload, incomingMessage.getOffset(), incomingMessage.getLength())
-                });
+                );
             }
 
             session.getProvenanceReporter().receive(messageFlowFile, getTransitUri(sessionInfo));
@@ -216,7 +240,7 @@ public abstract class AbstractWebSocketGatewayProcessor extends AbstractSessionF
             if (incomingMessage instanceof WebSocketConnectedMessage) {
                 session.transfer(messageFlowFile, REL_CONNECTED);
             } else {
-                switch (messageType) {
+                switch (Objects.requireNonNull(messageType)) {
                     case TEXT:
                         session.transfer(messageFlowFile, REL_MESSAGE_TEXT);
                         break;
@@ -233,6 +257,6 @@ public abstract class AbstractWebSocketGatewayProcessor extends AbstractSessionF
         }
     }
 
-    abstract protected String getTransitUri(final WebSocketSessionInfo sessionInfo);
+    protected abstract String getTransitUri(final WebSocketSessionInfo sessionInfo);
 
 }

nifi-nar-bundles/nifi-websocket-bundle/nifi-websocket-processors/src/main/java/org/apache/nifi/processors/websocket/ConnectWebSocket.java

@@ -44,7 +44,7 @@ import static org.apache.nifi.processors.websocket.WebSocketProcessorAttributes.
 import static org.apache.nifi.processors.websocket.WebSocketProcessorAttributes.ATTR_WS_SESSION_ID;
 
 @Tags({"subscribe", "WebSocket", "consume", "listen"})
-@InputRequirement(InputRequirement.Requirement.INPUT_FORBIDDEN)
+@InputRequirement(InputRequirement.Requirement.INPUT_ALLOWED)
 @TriggerSerially
 @CapabilityDescription("Acts as a WebSocket client endpoint to interact with a remote WebSocket server." +
         " FlowFiles are transferred to downstream relationships according to received message types" +
@@ -80,7 +80,7 @@ public class ConnectWebSocket extends AbstractWebSocketGatewayProcessor {
     private static final List<PropertyDescriptor> descriptors;
     private static final Set<Relationship> relationships;
 
-    static{
+    static {
         final List<PropertyDescriptor> innerDescriptorsList = new ArrayList<>();
         innerDescriptorsList.add(PROP_WEBSOCKET_CLIENT_SERVICE);
         innerDescriptorsList.add(PROP_WEBSOCKET_CLIENT_ID);
@@ -113,6 +113,6 @@ public class ConnectWebSocket extends AbstractWebSocketGatewayProcessor {
 
     @Override
     protected String getTransitUri(final WebSocketSessionInfo sessionInfo) {
-        return ((WebSocketClientService)webSocketService).getTargetUri();
+        return ((WebSocketClientService) webSocketService).getTargetUri();
     }
 }

nifi-nar-bundles/nifi-websocket-bundle/nifi-websocket-processors/src/main/resources/docs/org.apache.nifi.processors.websocket.ConnectWebSocket/additionalDetails.html

@@ -0,0 +1,47 @@
+<!DOCTYPE html>
+<html lang="en">
+<!--
+  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.
+-->
+<head>
+    <meta charset="utf-8"/>
+    <title>ConnectWebSocket</title>
+    <link rel="stylesheet" href="../../../../../css/component-usage.css" type="text/css"/>
+</head>
+
+<body>
+<h2>Summary</h2>
+<p>
+    This processor acts as a WebSocket client endpoint to interact with a remote WebSocket server.
+    It is capable of receiving messages from a websocket server and it transfers them to downstream relationships
+    according to the received message types.
+</p>
+<p>
+    The processor may have an incoming relationship, in which case flowfile attributes are passed down to its WebSocket Client Service.
+    This can be used to fine-tune the connection configuration (url and headers for example). For example "dynamic.url =
+    currentValue" flowfile attribute can be referenced in the WebSocket Client Service with the ${dynamic.url} expression.
+</p>
+<p>
+    You can define custom websocket headers in the incoming flowfile as additional attributes. The attribute key
+    shall start with "header." and continue with they header key. For example: "header.Authorization". The attribute
+    value will be the corresponding header value.
+<ol>
+    <li>header.Autorization | Basic base64UserNamePassWord</li>
+    <li>header.Content-Type | application, audio, example</li>
+</ol>
+<p>
+    For multiple header values provide a comma separated list.
+</p>
+</body>
+</html>

nifi-nar-bundles/nifi-websocket-bundle/nifi-websocket-processors/src/test/java/org/apache/nifi/processors/websocket/TestConnectWebSocket.java

@@ -20,6 +20,8 @@ import org.apache.nifi.processor.ProcessSessionFactory;
 import org.apache.nifi.processor.Relationship;
 import org.apache.nifi.provenance.ProvenanceEventRecord;
 import org.apache.nifi.provenance.ProvenanceEventType;
+import org.apache.nifi.remote.io.socket.NetworkUtils;
+import org.apache.nifi.reporting.InitializationException;
 import org.apache.nifi.util.MockFlowFile;
 import org.apache.nifi.util.MockProcessSession;
 import org.apache.nifi.util.SharedSessionState;
@@ -29,9 +31,13 @@ import org.apache.nifi.websocket.AbstractWebSocketSession;
 import org.apache.nifi.websocket.WebSocketClientService;
 import org.apache.nifi.websocket.WebSocketMessage;
 import org.apache.nifi.websocket.WebSocketSession;
+import org.apache.nifi.websocket.jetty.JettyWebSocketClient;
+import org.apache.nifi.websocket.jetty.JettyWebSocketServer;
+import org.junit.Assert;
 import org.junit.Test;
 
 import java.net.InetSocketAddress;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -51,7 +57,9 @@ public class TestConnectWebSocket extends TestListenWebSocket {
     @Test
     public void testSuccess() throws Exception {
         final TestRunner runner = TestRunners.newTestRunner(ConnectWebSocket.class);
-        final ConnectWebSocket processor = (ConnectWebSocket)runner.getProcessor();
+        runner.setIncomingConnection(false);
+
+        final ConnectWebSocket processor = (ConnectWebSocket) runner.getProcessor();
 
         final SharedSessionState sharedSessionState = new SharedSessionState(processor, new AtomicLong(0));
         // Use this custom session factory implementation so that createdSessions can be read from test case,
@@ -121,4 +129,63 @@ public class TestConnectWebSocket extends TestListenWebSocket {
         assertTrue(provenanceEvents.stream().allMatch(event -> ProvenanceEventType.RECEIVE.equals(event.getEventType())));
     }
 
+    @Test
+    public void testDynamicUrlsParsedFromFlowFileAndAbleToConnect() throws InitializationException {
+        // Start websocket server
+        final int port = NetworkUtils.availablePort();
+        TestRunner webSocketListener = getListenWebSocket(port);
+        webSocketListener.run(1, false);
+
+        final TestRunner runner = TestRunners.newTestRunner(ConnectWebSocket.class);
+
+        final String serviceId = "ws-service";
+        final String endpointId = "client-1";
+
+        Map<String, String> attributes = new HashMap<>();
+        attributes.put("dynamicUrlPart", "test");
+        MockFlowFile flowFile = new MockFlowFile(1L);
+        flowFile.putAttributes(attributes);
+        runner.enqueue(flowFile);
+
+        attributes.put("dynamicUrlPart", "test2");
+        MockFlowFile flowFileWithWrongUrl = new MockFlowFile(2L);
+        flowFileWithWrongUrl.putAttributes(attributes);
+        runner.enqueue(flowFileWithWrongUrl);
+
+        JettyWebSocketClient service = new JettyWebSocketClient();
+
+
+        runner.addControllerService(serviceId, service);
+        runner.setProperty(service, JettyWebSocketClient.WS_URI, String.format("ws://localhost:%s/${dynamicUrlPart}", port));
+        runner.enableControllerService(service);
+
+        runner.setProperty(ConnectWebSocket.PROP_WEBSOCKET_CLIENT_SERVICE, serviceId);
+        runner.setProperty(ConnectWebSocket.PROP_WEBSOCKET_CLIENT_ID, endpointId);
+
+        runner.run(1, false);
+
+        final List<MockFlowFile> flowFilesForRelationship = runner.getFlowFilesForRelationship(ConnectWebSocket.REL_CONNECTED);
+        assertEquals(1, flowFilesForRelationship.size());
+
+        final AssertionError assertionError = Assert.assertThrows(AssertionError.class, () -> runner.run(1));
+        assertTrue(assertionError.getCause().getLocalizedMessage().contains("Failed to renew session and connect to WebSocket service"));
+
+        runner.stop();
+        webSocketListener.stop();
+    }
+
+    private TestRunner getListenWebSocket(final int port) throws InitializationException {
+        final TestRunner runner = TestRunners.newTestRunner(ListenWebSocket.class);
+
+        final String serviceId = "ws-server-service";
+        JettyWebSocketServer service = new JettyWebSocketServer();
+        runner.addControllerService(serviceId, service);
+        runner.setProperty(service, JettyWebSocketServer.LISTEN_PORT, String.valueOf(port));
+        runner.enableControllerService(service);
+
+        runner.setProperty(ListenWebSocket.PROP_WEBSOCKET_SERVER_SERVICE, serviceId);
+        runner.setProperty(ListenWebSocket.PROP_SERVER_URL_PATH, "/test");
+
+        return runner;
+    }
 }

nifi-nar-bundles/nifi-websocket-bundle/nifi-websocket-services-api/src/main/java/org/apache/nifi/websocket/WebSocketClientService.java

@@ -20,6 +20,7 @@ import org.apache.nifi.annotation.documentation.CapabilityDescription;
 import org.apache.nifi.controller.ConfigurationContext;
 
 import java.io.IOException;
+import java.util.Map;
 
 /**
  * Control a WebSocket client instance.
@@ -34,6 +35,10 @@ public interface WebSocketClientService extends WebSocketService {
 
     void connect(final String clientId) throws IOException;
 
+    default void connect(final String clientId, final Map<String, String> flowFileAttributes) throws IOException {
+        connect(clientId);
+    }
+
     String getTargetUri();
 
 }

nifi-nar-bundles/nifi-websocket-bundle/nifi-websocket-services-jetty/src/main/java/dto/SessionInfo.java

@@ -0,0 +1,38 @@
+/*
+ * 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 dto;
+
+import java.util.Map;
+
+public class SessionInfo {
+
+    private final String sessionId;
+    private final Map<String, String> flowFileAttributes;
+
+    public SessionInfo(final String sessionId, final Map<String, String> flowFileAttributes) {
+        this.sessionId = sessionId;
+        this.flowFileAttributes = flowFileAttributes;
+    }
+
+    public String getSessionId() {
+        return sessionId;
+    }
+
+    public Map<String, String> getFlowFileAttributes() {
+        return flowFileAttributes;
+    }
+}

nifi-nar-bundles/nifi-websocket-bundle/nifi-websocket-services-jetty/src/main/java/org/apache/nifi/websocket/jetty/JettyWebSocketClient.java

@@ -16,6 +16,7 @@
  */
 package org.apache.nifi.websocket.jetty;
 
+import dto.SessionInfo;
 import org.apache.nifi.annotation.documentation.CapabilityDescription;
 import org.apache.nifi.annotation.documentation.Tags;
 import org.apache.nifi.annotation.lifecycle.OnDisabled;
@@ -27,6 +28,7 @@ import org.apache.nifi.components.ValidationResult;
 import org.apache.nifi.controller.ConfigurationContext;
 import org.apache.nifi.expression.ExpressionLanguageScope;
 import org.apache.nifi.logging.ComponentLog;
+import org.apache.nifi.processor.exception.ProcessException;
 import org.apache.nifi.processor.util.StandardValidators;
 import org.apache.nifi.ssl.SSLContextService;
 import org.apache.nifi.util.StringUtils;
@@ -40,14 +42,17 @@ import org.eclipse.jetty.util.ssl.SslContextFactory;
 import org.eclipse.jetty.websocket.api.Session;
 import org.eclipse.jetty.websocket.client.ClientUpgradeRequest;
 import org.eclipse.jetty.websocket.client.WebSocketClient;
+import util.HeaderMapExtractor;
 
 import java.io.IOException;
 import java.net.URI;
+import java.net.URISyntaxException;
 import java.nio.charset.Charset;
 import java.util.ArrayList;
 import java.util.Base64;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
@@ -68,7 +73,7 @@ public class JettyWebSocketClient extends AbstractJettyWebSocketService implemen
             .displayName("WebSocket URI")
             .description("The WebSocket URI this client connects to.")
             .required(true)
-            .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY)
+            .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES)
             .addValidator(StandardValidators.URI_VALIDATOR)
             .addValidator((subject, input, context) -> {
                 final ValidationResult.Builder result = new ValidationResult.Builder()
@@ -79,7 +84,7 @@ public class JettyWebSocketClient extends AbstractJettyWebSocketService implemen
                     result.explanation("Expression Language Present").valid(true);
                 } else {
                     result.explanation("Protocol should be either 'ws' or 'wss'.")
-                        .valid(input.startsWith("ws://") || input.startsWith("wss://"));
+                            .valid(input.startsWith("ws://") || input.startsWith("wss://"));
                 }
 
                 return result.build();
@@ -161,8 +166,7 @@ public class JettyWebSocketClient extends AbstractJettyWebSocketService implemen
     private static final List<PropertyDescriptor> properties;
 
     static {
-        final List<PropertyDescriptor> props = new ArrayList<>();
+        final List<PropertyDescriptor> props = new ArrayList<>(getAbstractPropertyDescriptors());
-        props.addAll(getAbstractPropertyDescriptors());
         props.add(WS_URI);
         props.add(SSL_CONTEXT);
         props.add(CONNECTION_TIMEOUT);
@@ -176,12 +180,14 @@ public class JettyWebSocketClient extends AbstractJettyWebSocketService implemen
         properties = Collections.unmodifiableList(props);
     }
 
+    private final Map<String, SessionInfo> activeSessions = new ConcurrentHashMap<>();
+    private final ReentrantLock connectionLock = new ReentrantLock();
     private WebSocketClient client;
     private URI webSocketUri;
     private String authorizationHeader;
     private long connectionTimeoutMillis;
     private volatile ScheduledExecutorService sessionMaintenanceScheduler;
-    private final ReentrantLock connectionLock = new ReentrantLock();
+    private ConfigurationContext configurationContext;
 
     @Override
     protected List<PropertyDescriptor> getSupportedPropertyDescriptors() {
@@ -190,8 +196,8 @@ public class JettyWebSocketClient extends AbstractJettyWebSocketService implemen
 
     @OnEnabled
     @Override
-    public void startClient(final ConfigurationContext context) throws Exception{
+    public void startClient(final ConfigurationContext context) throws Exception {
-
+        configurationContext = context;
         final SSLContextService sslService = context.getProperty(SSL_CONTEXT).asControllerService(SSLContextService.class);
         SslContextFactory sslContextFactory = null;
         if (sslService != null) {
@@ -227,8 +233,7 @@ public class JettyWebSocketClient extends AbstractJettyWebSocketService implemen
 
         client.start();
         activeSessions.clear();
-
+        webSocketUri = new URI(context.getProperty(WS_URI).evaluateAttributeExpressions(new HashMap<>()).getValue());
-        webSocketUri = new URI(context.getProperty(WS_URI).evaluateAttributeExpressions().getValue());
         connectionTimeoutMillis = context.getProperty(CONNECTION_TIMEOUT).evaluateAttributeExpressions().asTimePeriod(TimeUnit.MILLISECONDS);
 
         final Long sessionMaintenanceInterval = context.getProperty(SESSION_MAINTENANCE_INTERVAL).evaluateAttributeExpressions().asTimePeriod(TimeUnit.MILLISECONDS);
@@ -281,10 +286,20 @@ public class JettyWebSocketClient extends AbstractJettyWebSocketService implemen
 
     @Override
     public void connect(final String clientId) throws IOException {
-        connect(clientId, null);
+        connect(clientId, null, Collections.emptyMap());
     }
 
-    private void connect(final String clientId, String sessionId) throws IOException {
+    @Override
+    public void connect(final String clientId, final Map<String, String> flowFileAttributes) throws IOException {
+        connect(clientId, null, flowFileAttributes);
+    }
+
+    private void connect(final String clientId, final String sessionId, final Map<String, String> flowFileAttributes) throws IOException {
+        try {
+            webSocketUri = new URI(configurationContext.getProperty(WS_URI).evaluateAttributeExpressions(flowFileAttributes).getValue());
+        } catch (URISyntaxException e) {
+            throw new ProcessException("Could not create websocket URI", e);
+        }
 
         connectionLock.lock();
 
@@ -293,17 +308,21 @@ public class JettyWebSocketClient extends AbstractJettyWebSocketService implemen
             try {
                 router = routers.getRouterOrFail(clientId);
             } catch (WebSocketConfigurationException e) {
-                throw new IllegalStateException("Failed to get router due to: "  + e, e);
+                throw new IllegalStateException("Failed to get router due to: " + e, e);
             }
             final RoutingWebSocketListener listener = new RoutingWebSocketListener(router);
             listener.setSessionId(sessionId);
 
             final ClientUpgradeRequest request = new ClientUpgradeRequest();
+
+            if (!flowFileAttributes.isEmpty()) {
+                request.setHeaders(HeaderMapExtractor.getHeaderMap(flowFileAttributes));
+            }
             if (!StringUtils.isEmpty(authorizationHeader)) {
                 request.setHeader(HttpHeader.AUTHORIZATION.asString(), authorizationHeader);
             }
             final Future<Session> connect = client.connect(listener, webSocketUri, request);
-            getLogger().info("Connecting to : {}", new Object[]{webSocketUri});
+            getLogger().info("Connecting to : {}", webSocketUri);
 
             final Session session;
             try {
@@ -311,8 +330,8 @@ public class JettyWebSocketClient extends AbstractJettyWebSocketService implemen
             } catch (Exception e) {
                 throw new IOException("Failed to connect " + webSocketUri + " due to: " + e, e);
             }
-            getLogger().info("Connected, session={}", new Object[]{session});
+            getLogger().info("Connected, session={}", session);
-            activeSessions.put(clientId, listener.getSessionId());
+            activeSessions.put(clientId, new SessionInfo(listener.getSessionId(), flowFileAttributes));
 
         } finally {
             connectionLock.unlock();
@@ -320,8 +339,6 @@ public class JettyWebSocketClient extends AbstractJettyWebSocketService implemen
 
     }
 
-    private Map<String, String> activeSessions = new ConcurrentHashMap<>();
-
     void maintainSessions() throws Exception {
         if (client == null) {
             return;
@@ -338,19 +355,19 @@ public class JettyWebSocketClient extends AbstractJettyWebSocketService implemen
                     router = routers.getRouterOrFail(clientId);
                 } catch (final WebSocketConfigurationException e) {
                     if (logger.isDebugEnabled()) {
-                        logger.debug("The clientId {} is no longer active. Discarding the clientId.", new Object[]{clientId});
+                        logger.debug("The clientId {} is no longer active. Discarding the clientId.", clientId);
                     }
                     activeSessions.remove(clientId);
                     continue;
                 }
 
-                final String sessionId = activeSessions.get(clientId);
+                final SessionInfo sessionInfo = activeSessions.get(clientId);
                 // If this session is still alive, do nothing.
-                if (!router.containsSession(sessionId)) {
+                if (!router.containsSession(sessionInfo.getSessionId())) {
                     // This session is no longer active, reconnect it.
                     // If it fails, the sessionId will remain in activeSessions, and retries later.
                     // This reconnect attempt is continued until user explicitly stops a processor or this controller service.
-                    connect(clientId, sessionId);
+                    connect(clientId, sessionInfo.getSessionId(), sessionInfo.getFlowFileAttributes());
                 }
             }
         } finally {
@@ -358,7 +375,7 @@ public class JettyWebSocketClient extends AbstractJettyWebSocketService implemen
         }
 
         if (logger.isDebugEnabled()) {
-            logger.debug("Session maintenance completed. activeSessions={}", new Object[]{activeSessions});
+            logger.debug("Session maintenance completed. activeSessions={}", activeSessions);
         }
     }
 

nifi-nar-bundles/nifi-websocket-bundle/nifi-websocket-services-jetty/src/main/java/util/HeaderMapExtractor.java

@@ -0,0 +1,47 @@
+/*
+ * 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 util;
+
+import org.apache.nifi.util.StringUtils;
+
+import java.util.AbstractMap;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public final class HeaderMapExtractor {
+
+    private HeaderMapExtractor(){
+        // Utility class, not meant to be instantiated.
+    }
+
+    public static final String HEADER_PREFIX = "header.";
+
+    public static Map<String, List<String>> getHeaderMap(final Map<String, String> flowFileAttributes) {
+        return flowFileAttributes.entrySet().stream()
+                .filter(entry -> entry.getKey().startsWith(HEADER_PREFIX))
+                .filter(entry -> StringUtils.isNotBlank(entry.getValue()))
+                .map(entry -> new AbstractMap.SimpleImmutableEntry<>(StringUtils.substringAfter(entry.getKey(), HEADER_PREFIX), entry.getValue()))
+                .collect(Collectors.toMap(Map.Entry::getKey, HeaderMapExtractor::headerValueMapper));
+    }
+
+    private static List<String> headerValueMapper(Map.Entry<String, String> entry) {
+        return Arrays.stream(entry.getValue().split(",")).map(String::trim).collect(Collectors.toList());
+    }
+
+}

nifi-nar-bundles/nifi-websocket-bundle/nifi-websocket-services-jetty/src/test/java/org/apache/nifi/websocket/jetty/ITJettyWebSocketCommunication.java

@@ -30,6 +30,7 @@ import org.mockito.invocation.InvocationOnMock;
 
 import java.net.ServerSocket;
 import java.nio.ByteBuffer;
+import java.util.Collections;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
@@ -234,7 +235,7 @@ public class ITJettyWebSocketCommunication {
 
         clientService.registerProcessor(clientId, clientProcessor);
 
-        clientService.connect(clientId);
+        clientService.connect(clientId, Collections.emptyMap());
 
         assertTrue("WebSocket client should be able to fire connected event.", clientConnectedServer.await(5, TimeUnit.SECONDS));
         assertTrue("WebSocket server should be able to fire connected event.", serverIsConnectedByClient.await(5, TimeUnit.SECONDS));

nifi-nar-bundles/nifi-websocket-bundle/nifi-websocket-services-jetty/src/test/java/org/apache/nifi/websocket/util/HeaderMapExtractorTest.java

@@ -0,0 +1,67 @@
+/*
+ * 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.websocket.util;
+
+import org.junit.Test;
+import util.HeaderMapExtractor;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class HeaderMapExtractorTest {
+
+    @Test
+    public void testMapExtractor() {
+
+        // GIVEN
+        final Map<String, String> attributes = new HashMap<>();
+        attributes.put("header.AUTHORIZATION", "AUTH_VALUE");
+        attributes.put("header.MULTI_HEADER_KEY", "FIRST, SECOND   ,THIRD");
+        attributes.put("header.dots.dots.dots.DOTS_KEY", "DOTS_VALUE");
+        attributes.put("something.else.SOMETHING_ELSE_KEY", "SOMETHING_ELSE_VALUE");
+        attributes.put("header.EMPTY_VALUE_KEY", "");
+        attributes.put("headerButNotReally.UNRECOGNIZED", "NOT_A_HEADER_VALUE");
+
+        final Map<String, List<String>> expected = new HashMap<>();
+        expected.put("AUTHORIZATION", Collections.singletonList("AUTH_VALUE"));
+        expected.put("MULTI_HEADER_KEY", Arrays.asList("FIRST", "SECOND", "THIRD"));
+        expected.put("dots.dots.dots.DOTS_KEY", Collections.singletonList("DOTS_VALUE"));
+
+        // WHEN
+        final Map<String, List<String>> actual = HeaderMapExtractor.getHeaderMap(attributes);
+
+        // THEN
+        assertEquals(expected, actual);
+
+        assertEquals(expected.size(), actual.size());
+        for (Map.Entry<String, List<String>> entry : actual.entrySet()) {
+            assertTrue(expected.containsKey(entry.getKey()));
+            final List<String> actualHeaderValues = entry.getValue();
+            final List<String> expectedHeaderValues = expected.get(entry.getKey());
+            for (int i = 0; i < actualHeaderValues.size(); i++) {
+                assertEquals(expectedHeaderValues.get(i), actualHeaderValues.get(i));
+            }
+        }
+    }
+
+}