NIFI-9397 Added Custom Authorization property to JettyWebSocketClient

This closes apache#5574 Signed-off-by: David Handermann <exceptionfactory@apache.org>
pvillard31 · Dec 7, 2021 · 12015a1 · 12015a1
1 parent 70822b9
commit 12015a1
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 2 deletions.
diff --git a/...et-services-jetty/src/main/java/org/apache/nifi/websocket/jetty/JettyWebSocketClient.java b/...et-services-jetty/src/main/java/org/apache/nifi/websocket/jetty/JettyWebSocketClient.java
@@ -145,6 +145,19 @@ public class JettyWebSocketClient extends AbstractJettyWebSocketService implemen
             .defaultValue("US-ASCII")
             .build();
 
+    public static final PropertyDescriptor CUSTOM_AUTH = new PropertyDescriptor.Builder()
+            .name("custom-authorization")
+            .displayName("Custom Authorization")
+            .description(
+                    "Configures a custom HTTP Authorization Header as described in RFC 7235 Section 4.2." +
+                    " Setting a custom Authorization Header excludes configuring the User Name and User Password properties for Basic Authentication.")
+            .required(false)
+            .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY)
+            .addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
+            .sensitive(true)
+            .build();
+
+
     public static final PropertyDescriptor PROXY_HOST = new PropertyDescriptor.Builder()
             .name("proxy-host")
             .displayName("HTTP Proxy Host")
@@ -174,6 +187,7 @@ public class JettyWebSocketClient extends AbstractJettyWebSocketService implemen
         props.add(USER_NAME);
         props.add(USER_PASSWORD);
         props.add(AUTH_CHARSET);
+        props.add(CUSTOM_AUTH);
         props.add(PROXY_HOST);
         props.add(PROXY_PORT);
 
@@ -184,10 +198,10 @@ public class JettyWebSocketClient extends AbstractJettyWebSocketService implemen
     private final ReentrantLock connectionLock = new ReentrantLock();
     private WebSocketClient client;
     private URI webSocketUri;
-    private String authorizationHeader;
     private long connectionTimeoutMillis;
     private volatile ScheduledExecutorService sessionMaintenanceScheduler;
     private ConfigurationContext configurationContext;
+    protected String authorizationHeader;
 
     @Override
     protected List<PropertyDescriptor> getSupportedPropertyDescriptors() {
@@ -219,7 +233,11 @@ public void startClient(final ConfigurationContext context) throws Exception {
         configurePolicy(context, client.getPolicy());
         final String userName = context.getProperty(USER_NAME).evaluateAttributeExpressions().getValue();
         final String userPassword = context.getProperty(USER_PASSWORD).evaluateAttributeExpressions().getValue();
-        if (!StringUtils.isEmpty(userName) && !StringUtils.isEmpty(userPassword)) {
+        final String customAuth = context.getProperty(CUSTOM_AUTH).evaluateAttributeExpressions().getValue();
+
+        if (!StringUtils.isEmpty(customAuth)) {
+            authorizationHeader = customAuth;
+        } else if (!StringUtils.isEmpty(userName) && !StringUtils.isEmpty(userPassword)) {
             final String charsetName = context.getProperty(AUTH_CHARSET).evaluateAttributeExpressions().getValue();
             if (StringUtils.isEmpty(charsetName)) {
                 throw new IllegalArgumentException(AUTH_CHARSET.getDisplayName() + " was not specified.");
@@ -258,6 +276,14 @@ protected Collection<ValidationResult> customValidate(ValidationContext validati
             results.add(new ValidationResult.Builder().subject("HTTP Proxy Host and Port").valid(false).explanation(
                     "If HTTP Proxy Host or HTTP Proxy Port is set, both must be set").build());
         }
+
+        final boolean isBaseAuthUsed = validationContext.getProperty(USER_NAME).isSet() || validationContext.getProperty(USER_PASSWORD).isSet();
+
+        if (isBaseAuthUsed && validationContext.getProperty(CUSTOM_AUTH).isSet()) {
+            results.add((new ValidationResult.Builder().subject("Authentication").valid(false).explanation(
+                    "Properties related to Basic Authentication (\"User Name\" and \"User Password\") cannot be used together with \"Custom Authorization\"")).build());
+        }
+
         return results;
     }
 

diff --git a/...ervices-jetty/src/test/java/org/apache/nifi/websocket/jetty/TestJettyWebSocketClient.java b/...ervices-jetty/src/test/java/org/apache/nifi/websocket/jetty/TestJettyWebSocketClient.java
@@ -17,7 +17,11 @@
 package org.apache.nifi.websocket.jetty;
 
 import org.apache.nifi.components.ValidationResult;
+import org.apache.nifi.processor.Processor;
+import org.apache.nifi.util.TestRunner;
+import org.apache.nifi.util.TestRunners;
 import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
 
 import java.util.Collection;
 
@@ -27,6 +31,8 @@
 
 public class TestJettyWebSocketClient {
 
+    public static final String CUSTOM_AUTH = "Apikey 8743b52063cd84097a65d1633f5c74f5";
+
     @Test
     public void testValidationRequiredProperties() throws Exception {
         final JettyWebSocketClient service = new JettyWebSocketClient();
@@ -97,4 +103,22 @@ public void testValidationSuccessWithProxy() throws Exception {
         final Collection<ValidationResult> results = service.validate(context.getValidationContext());
         assertEquals(0, results.size());
     }
+
+    @Test
+    public void testCustomAuthHeader() throws Exception {
+        final TestRunner runner = TestRunners.newTestRunner(Mockito.mock(Processor.class));
+        final TestableJettyWebSocketClient testSubject = new TestableJettyWebSocketClient();
+        runner.addControllerService("client", testSubject);
+        runner.setProperty(testSubject, JettyWebSocketClient.WS_URI, "wss://localhost:9001/test");
+        runner.setProperty(testSubject, JettyWebSocketClient.CUSTOM_AUTH, CUSTOM_AUTH);
+        runner.assertValid(testSubject);
+        runner.enableControllerService(testSubject);
+        assertEquals(CUSTOM_AUTH, testSubject.getAuthHeaderValue());
+    }
+
+    private static class TestableJettyWebSocketClient extends JettyWebSocketClient {
+        public String getAuthHeaderValue() {
+            return authorizationHeader;
+        }
+    }
 }