AMQP websocket custom proxy authenticator

anuchandy · anuchandy · commit 3e730d238346 · 2024-04-12T13:20:45.000-07:00
diff --git a/src/main/java/com/microsoft/azure/proton/transport/proxy/ProxyAuthenticator.java b/src/main/java/com/microsoft/azure/proton/transport/proxy/ProxyAuthenticator.java
@@ -0,0 +1,76 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.microsoft.azure.proton.transport.proxy;
+
+import com.microsoft.azure.proton.transport.proxy.impl.ChallengeResponseAccessHelper;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * A contract to authenticate a proxy server to tunnel a websocket connection to an AMQP broker.
+ */
+public interface ProxyAuthenticator {
+    /**
+     * Authenticate a proxy server to tunnel a websocket connection to an AMQP broker.
+     * <p>
+     * This method is called when the proxy server replies to the CONNECT with 407 (Proxy Authentication Required)
+     * challenge. The proxy server's challenge response includes a 'Proxy-Authenticate' header indicating
+     * the authentication scheme(s) that the proxy supports. The implementation of this method should
+     * <ul>
+     *     <li>enumerate the schemes using {@link ChallengeResponse#getAuthenticationSchemes()}) and choose the most
+     *     secure scheme the client supports,</li>
+     *     <li>identify the credential for the chosen scheme, </li>
+     *     <li>compute and return authorization value.The RFC7325 defines authorization format as a value that starts
+     *     with the selected scheme, followed by a space and the base64 encoded credentials for the scheme.</li>
+     * </ul>
+     * The returned authorization value will be sent to the proxy server in 'Proxy-Authorization' header to complete
+     * the authentication.
+     * @see <a href="https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/407">407 Proxy Authentication Required</a>
+     * @see <a href="https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Proxy-Authenticate">Proxy-Authenticate</a>
+     * @see <a href="https://datatracker.ietf.org/doc/html/rfc7235#section-4.4">RFC7235</a>
+     *
+     * @param response the challenge response from the proxy server.
+     * @return the authorization value to send to the proxy server using 'Proxy-Authorization' header.
+     */
+    String authenticate(ChallengeResponse response);
+
+    /**
+     * Represents the 407 challenge response from the proxy server.
+     */
+    final class ChallengeResponse {
+        static {
+            ChallengeResponseAccessHelper.setAccessor(ChallengeResponse::new);
+        }
+        private static final String PROXY_AUTHENTICATE = "Proxy-Authenticate";
+        private final Map<String, List<String>> headers;
+
+        /**
+         * Creates the ChallengeResponse.
+         *
+         * @param headers the response headers
+         */
+        ChallengeResponse(Map<String, List<String>> headers) {
+            this.headers = headers;
+        }
+
+        /**
+         * Gets the headers.
+         *
+         * @return the headers.
+         */
+        public Map<String, List<String>> getHeaders() {
+            return headers;
+        }
+
+        /**
+         * Gets the authentication schemes supported by the proxy server.
+         *
+         * @return the authentication schemes supported by the proxy server.
+         */
+        public List<String> getAuthenticationSchemes() {
+            return headers.get(PROXY_AUTHENTICATE);
+        }
+    }
+}
diff --git a/src/main/java/com/microsoft/azure/proton/transport/proxy/ProxyConfiguration.java b/src/main/java/com/microsoft/azure/proton/transport/proxy/ProxyConfiguration.java
@@ -20,6 +20,7 @@ public class ProxyConfiguration implements AutoCloseable {
 
     private final java.net.Proxy proxyAddress;
     private final ProxyAuthenticationType authentication;
+    private final ProxyAuthenticator authenticator;
     private final PasswordAuthentication credentials;
 
     /**
@@ -32,6 +33,7 @@ public class ProxyConfiguration implements AutoCloseable {
      */
     private ProxyConfiguration() {
         this.authentication = null;
+        this.authenticator = null;
         this.credentials = null;
         this.proxyAddress = null;
     }
@@ -65,6 +67,21 @@ public ProxyConfiguration(ProxyAuthenticationType authentication, java.net.Proxy
 
             this.credentials = null;
         }
+        this.authenticator = null;
+    }
+
+    /**
+     * Creates a proxy configuration that uses the {@code proxyAddress} and authenticates with provided {@code authenticator}.
+     *
+     * @param authenticator the proxy authenticator to use.
+     * @param proxyAddress Proxy to use.
+     * @throws NullPointerException if {@code proxyAddress} or {@code proxyAuthenticator} is {@code null}.
+     */
+    public ProxyConfiguration(ProxyAuthenticator authenticator, java.net.Proxy proxyAddress) {
+        this.authenticator = Objects.requireNonNull(authenticator, "'authenticator' cannot be null.");
+        this.proxyAddress = Objects.requireNonNull(proxyAddress, "'proxyAddress' cannot be null.");
+        this.authentication = null;
+        this.credentials = null;
     }
 
     /**
@@ -97,6 +114,19 @@ public ProxyAuthenticationType authentication() {
         return authentication;
     }
 
+    /**
+     * Gets the proxy authenticator to set up the web socket connection to the AMQP broker via a proxy.
+     * <p>
+     * The authenticator is responsible for selecting one of the authorization schemes that the proxy presents, identify
+     * the credentials for the scheme it selects then compute and return the authorization value to be sent through
+     * the 'Proxy-Authorization' Header.
+     * </p>
+     * @return the proxy authenticator.
+     */
+    public ProxyAuthenticator getAuthenticator() {
+        return this.authenticator;
+    }
+
     /**
      * Gets whether the user has defined credentials.
      *
diff --git a/src/main/java/com/microsoft/azure/proton/transport/proxy/impl/ChallengeResponseAccessHelper.java b/src/main/java/com/microsoft/azure/proton/transport/proxy/impl/ChallengeResponseAccessHelper.java
@@ -0,0 +1,62 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.microsoft.azure.proton.transport.proxy.impl;
+
+import java.util.List;
+import java.util.Map;
+import com.microsoft.azure.proton.transport.proxy.ProxyAuthenticator;
+
+/**
+ * The accessor helper for {@link ProxyAuthenticator.ChallengeResponse}.
+ */
+public final class ChallengeResponseAccessHelper {
+    private static ChallengeResponseAccessor accessor;
+
+    /**
+     * The accessor interface for {@link ProxyAuthenticator.ChallengeResponse} construction.
+     */
+    public interface ChallengeResponseAccessor {
+        /**
+         * Create an instance of {@link ProxyAuthenticator.ChallengeResponse} with the provided headers.
+         *
+         * @param headers the headers.
+         * @return the created instance of {@link ProxyAuthenticator.ChallengeResponse}.
+         */
+        ProxyAuthenticator.ChallengeResponse internalCreate(Map<String, List<String>> headers);
+    }
+
+    /**
+     * Sets the accessor.
+     *
+     * @param accessor the accessor.
+     */
+    public static void setAccessor(ChallengeResponseAccessor accessor) {
+        ChallengeResponseAccessHelper.accessor = accessor;
+    }
+
+    /**
+     * Creates an instance of {@link ProxyAuthenticator.ChallengeResponse} with the provided headers.
+     *
+     * @param headers the headers.
+     * @return the created instance of {@link ProxyAuthenticator.ChallengeResponse}.
+     */
+    public static ProxyAuthenticator.ChallengeResponse internalCreate(Map<String, List<String>> headers) {
+        if (accessor == null) {
+            try {
+                Class.forName(ProxyAuthenticator.ChallengeResponse.class.getName(), true,
+                    ChallengeResponseAccessHelper.class.getClassLoader());
+            } catch (ClassNotFoundException e) {
+                throw new RuntimeException(e);
+            }
+        }
+        assert accessor != null;
+        return accessor.internalCreate(headers);
+    }
+
+    /**
+     * Private constructor to prevent instantiation.
+     */
+    private ChallengeResponseAccessHelper() {
+    }
+}
diff --git a/src/main/java/com/microsoft/azure/proton/transport/proxy/impl/ProxyImpl.java b/src/main/java/com/microsoft/azure/proton/transport/proxy/impl/ProxyImpl.java
@@ -22,10 +22,12 @@
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicReference;
@@ -253,12 +255,13 @@ public void process() throws TransportException {
                     proxyResponse.set(null);
 
                     final boolean isSuccess = proxyHandler.validateProxyResponse(connectResponse);
-                    // When connecting to proxy, it does not challenge us for authentication. If the user has specified
-                    // a configuration, and it is not NONE, then we fail due to misconfiguration.
                     if (isSuccess) {
-                        if (proxyConfiguration == null || proxyConfiguration.authentication() == ProxyAuthenticationType.NONE) {
+                        if (proxyConfiguration == null || proxyConfiguration.authentication() == ProxyAuthenticationType.NONE
+                            || proxyConfiguration.getAuthenticator() == null) {
                             proxyState = ProxyState.PN_PROXY_CONNECTED;
                         } else {
+                            // In response to CONNECT, the proxy didn't challenge us for authentication. Given the user has specified an
+                            // authentication configuration, we fail due to misconfiguration.
                             if (LOGGER.isErrorEnabled()) {
                                 LOGGER.error("ProxyConfiguration mismatch. User configured: '{}', but authentication is not required",
                                     proxyConfiguration.authentication());
@@ -269,29 +272,47 @@ public void process() throws TransportException {
                     }
 
                     final Map<String, List<String>> headers = connectResponse.getHeaders();
-                    final Set<ProxyAuthenticationType> supportedTypes = getAuthenticationTypes(headers);
-
-                    // The proxy did not successfully connect, user has specified that they want a particular
-                    // authentication method, but it is not in list of supported authentication methods.
-                    if (proxyConfiguration != null && !supportedTypes.contains(proxyConfiguration.authentication())) {
-                        if (LOGGER.isErrorEnabled()) {
-                            LOGGER.error("Proxy authentication required. User configured: '{}', but supported proxy authentication methods are: {}",
-                                proxyConfiguration.authentication(),
-                                supportedTypes.stream().map(type -> type.toString()).collect(Collectors.joining(",")));
+                    final List<String> challenges = headers.getOrDefault(PROXY_AUTHENTICATE, new ArrayList<>());
+                    final ProxyChallengeProcessor processor;
+                    if (proxyConfiguration != null && proxyConfiguration.getAuthenticator() != null) {
+                        final boolean is407 = connectResponse.getStatus().getStatusCode() == 407;
+                        if (!is407) {
+                            closeTailProxyError(PROXY_CONNECT_FAILED + connectResponse);
+                            break;
+                        }
+                        if (challenges.isEmpty()) {
+                            closeTailProxyError("'407 Proxy Authentication Required' received without " + PROXY_AUTHENTICATE
+                                + "header or authentication schemes." + connectResponse);
+                            break;
                         }
-                        closeTailProxyError(PROXY_CONNECT_USER_ERROR + PROXY_CONNECT_FAILED
+                        processor = new DelegatedProxyChallengeProcessor(headers, proxyConfiguration.getAuthenticator());
+                    } else {
+                        final Set<ProxyAuthenticationType> supportedTypes = getAuthenticationTypes(headers);
+                        // The proxy did not successfully connect, user has specified that they want a particular
+                        // authentication method, but it is not in list of supported authentication methods.
+                        if (proxyConfiguration != null && !supportedTypes.contains(proxyConfiguration.authentication())) {
+                            if (LOGGER.isErrorEnabled()) {
+                                LOGGER.error(
+                                    "Proxy authentication required. User configured: '{}', but supported proxy authentication methods are: {}",
+                                    proxyConfiguration.authentication(),
+                                    supportedTypes.stream().map(type -> type.toString()).collect(Collectors.joining(",")));
+                            }
+                            closeTailProxyError(PROXY_CONNECT_USER_ERROR + PROXY_CONNECT_FAILED
                                 + connectResponse);
-                        break;
-                    }
-
-                    final List<String> challenges = headers.getOrDefault(PROXY_AUTHENTICATE, new ArrayList<>());
-                    final ProxyChallengeProcessor processor = proxyConfiguration != null
+                            break;
+                        }
+                        processor = proxyConfiguration != null
                             ? getChallengeProcessor(host, challenges, proxyConfiguration.authentication())
                             : getChallengeProcessor(host, challenges, supportedTypes);
+                    }
 
                     if (processor != null) {
                         proxyState = ProxyState.PN_PROXY_CHALLENGE;
                         ProxyImpl.this.headers = processor.getHeader();
+                        if (proxyConfiguration != null
+                            && proxyConfiguration.getAuthenticator() != null && !testAuthorizeHeader(challenges, ProxyImpl.this.headers)) {
+                            closeTailProxyError("User error: ProxyAuthenticator did not provide a valid authorization header.");
+                        }
                     } else {
                         LOGGER.warn("Could not get ProxyChallengeProcessor for challenges.");
                         closeTailProxyError(PROXY_CONNECT_FAILED + String.join(";", challenges));
@@ -540,4 +561,39 @@ private ProxyResponse readProxyResponse(ByteBuffer buffer) {
             return proxyResponse.get();
         }
     }
+
+    private static final class DelegatedProxyChallengeProcessor implements ProxyChallengeProcessor {
+        private final Map<String, List<String>> headers;
+        private final com.microsoft.azure.proton.transport.proxy.ProxyAuthenticator authenticator;
+
+        DelegatedProxyChallengeProcessor(Map<String, List<String>> headers,
+            com.microsoft.azure.proton.transport.proxy.ProxyAuthenticator authenticator) {
+            this.headers = Objects.requireNonNull(headers, "'headers' cannot be null.");
+            this.authenticator = Objects.requireNonNull(authenticator, "'authenticator' cannot be null.");
+        }
+
+        @Override
+        public Map<String, String> getHeader() {
+            // the call site ensured that the 'headers' contain 'Proxy-Authenticate' header with the authentication schemes.
+            final String authorizedHeader = authenticator.authenticate(ChallengeResponseAccessHelper.internalCreate(headers));
+            final Map<String, String> headers = new HashMap<>(1);
+            headers.put(Constants.PROXY_AUTHORIZATION, authorizedHeader);
+            return headers;
+        }
+    }
+
+    private boolean testAuthorizeHeader(List<String> challenges, Map<String, String> headers) {
+        assert !challenges.isEmpty();
+        final String authorizeHeader = headers.get(Constants.PROXY_AUTHORIZATION);
+        if (authorizeHeader == null || authorizeHeader.trim().isEmpty()) {
+            return false;
+        }
+        final String value = authorizeHeader.toLowerCase(Locale.ROOT);
+        for (String scheme : challenges) {
+            if (value.startsWith(scheme.toLowerCase(Locale.ROOT) + " ")) {
+                return true;
+            }
+        }
+        return false;
+    }
 }
