Add OAuth configuration

WIP

Author: acogoluegnes

src/main/java/com/rabbitmq/client/amqp/ConnectionSettings.java

@@ -154,6 +154,8 @@ public interface ConnectionSettings<T> {
    */
   Affinity<? extends T> affinity();
 
+  OAuthSettings<? extends T> oauth();
+
   /**
    * TLS settings.
    *

src/main/java/com/rabbitmq/client/amqp/OAuthSettings.java

@@ -0,0 +1,33 @@
+// Copyright (c) 2024 Broadcom. All Rights Reserved.
+// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
+//
+// Licensed 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.
+//
+// If you have any questions regarding licensing, please contact us at
+// [email protected].
+package com.rabbitmq.client.amqp;
+
+public interface OAuthSettings<T> {
+
+  OAuthSettings<T> tokenEndpointUri(String uri);
+
+  OAuthSettings<T> clientId(String clientId);
+
+  OAuthSettings<T> clientSecret(String clientSecret);
+
+  OAuthSettings<T> grantType(String grantType);
+
+  OAuthSettings<T> parameter(String name, String value);
+
+  T connection();
+}

src/main/java/com/rabbitmq/client/amqp/impl/AmqpConnection.java

@@ -20,6 +20,8 @@
 import static com.rabbitmq.client.amqp.Resource.State.*;
 import static com.rabbitmq.client.amqp.impl.Utils.supportFilterExpressions;
 import static com.rabbitmq.client.amqp.impl.Utils.supportSetToken;
+import static java.lang.System.nanoTime;
+import static java.time.Duration.ofNanos;
 
 import com.rabbitmq.client.amqp.*;
 import com.rabbitmq.client.amqp.ObservationCollector;
@@ -81,8 +83,9 @@ final class AmqpConnection extends ResourceBase implements Connection {
   AmqpConnection(AmqpConnectionBuilder builder) {
     super(builder.listeners());
     this.id = ID_SEQUENCE.getAndIncrement();
-    this.name = builder.name();
     this.environment = builder.environment();
+    this.name =
+        builder.name() == null ? this.environment.toString() + "-" + this.id : builder.name();
     this.connectionSettings = builder.connectionSettings().consolidate();
     this.sessionHandlerSupplier =
         builder.isolateResources()
@@ -115,19 +118,18 @@ final class AmqpConnection extends ResourceBase implements Connection {
       this.affinityStrategy = null;
     }
     this.management = createManagement();
-    if (this.connectionSettings.credentialsProvider()
-        instanceof UsernamePasswordCredentialsProvider) {
-      UsernamePasswordCredentialsProvider credentialsProvider =
-          (UsernamePasswordCredentialsProvider) connectionSettings.credentialsProvider();
-      Credentials credentials = new UsernamePasswordCredentials(credentialsProvider);
-      this.credentialsRegistration =
-          credentials.register(
-              (u, p) -> {
-                ((AmqpManagement) management()).setToken(p);
-              });
-    } else {
-      this.credentialsRegistration = Credentials.NO_OP.register((u, p) -> {});
-    }
+    Credentials credentials = builder.credentials();
+    this.credentialsRegistration =
+        credentials.register(
+            (username, password) -> {
+              LOGGER.debug("Setting new token for connection {}", this.name);
+              long start = nanoTime();
+              ((AmqpManagement) management()).setToken(password);
+              LOGGER.debug(
+                  "Set new token for connection {} in {} ms",
+                  this.name,
+                  ofNanos(nanoTime() - start).toMillis());
+            });
     LOGGER.debug("Opening native connection for connection '{}'...", this.name());
     NativeConnectionWrapper ncw =
         ConnectionUtils.enforceAffinity(
@@ -611,7 +613,7 @@ ExecutorService dispatchingExecutorService() {
       if (this.dispatchingExecutorService == null) {
         this.dispatchingExecutorService =
             Executors.newSingleThreadExecutor(
-                Utils.threadFactory("dispatching-" + this.name + "-"));
+                Utils.threadFactory("dispatching-" + this.name() + "-"));
       }
       return this.dispatchingExecutorService;
     } finally {
@@ -754,7 +756,9 @@ private void close(Throwable cause) {
         consumer.close(cause);
       }
       try {
-        this.dispatchingExecutorService.shutdownNow();
+        if (this.dispatchingExecutorService != null) {
+          this.dispatchingExecutorService.shutdownNow();
+        }
       } catch (Exception e) {
         LOGGER.info(
             "Error while shutting down dispatching executor service for connection '{}': {}",
@@ -776,7 +780,7 @@ private void close(Throwable cause) {
 
   @Override
   public String toString() {
-    return this.environment.toString() + "-" + this.id;
+    return this.name();
   }
 
   static class NativeConnectionWrapper {

src/main/java/com/rabbitmq/client/amqp/impl/AmqpConnectionBuilder.java

@@ -104,6 +104,11 @@ public DefaultConnectionSettings.DefaultAffinity<? extends ConnectionBuilder> af
     return this.connectionSettings.affinity();
   }
 
+  @Override
+  public OAuthSettings<? extends ConnectionBuilder> oauth() {
+    return this.connectionSettings.oauth();
+  }
+
   @Override
   public ConnectionBuilder listeners(Resource.StateListener... listeners) {
     if (listeners == null || listeners.length == 0) {
@@ -157,6 +162,10 @@ AmqpEnvironment environment() {
     return environment;
   }
 
+  Credentials credentials() {
+    return environment().credentialsFactory().credentials(this.connectionSettings);
+  }
+
   AmqpRecoveryConfiguration recoveryConfiguration() {
     return recoveryConfiguration;
   }

src/main/java/com/rabbitmq/client/amqp/impl/AmqpEnvironment.java

@@ -51,6 +51,7 @@ class AmqpEnvironment implements Environment {
   private final ConnectionUtils.AffinityCache affinityCache = new ConnectionUtils.AffinityCache();
   private final EventLoop recoveryEventLoop;
   private final ExecutorService recoveryEventLoopExecutorService;
+  private final CredentialsFactory credentialsFactory = new CredentialsFactory(this);
 
   AmqpEnvironment(
       ExecutorService executorService,
@@ -120,6 +121,10 @@ Clock clock() {
     return this.clock;
   }
 
+  CredentialsFactory credentialsFactory() {
+    return this.credentialsFactory;
+  }
+
   @Override
   public void close() {
     if (this.closed.compareAndSet(false, true)) {

src/main/java/com/rabbitmq/client/amqp/impl/CredentialsFactory.java

@@ -0,0 +1,87 @@
+// Copyright (c) 2024 Broadcom. All Rights Reserved.
+// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
+//
+// Licensed 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.
+//
+// If you have any questions regarding licensing, please contact us at
+// [email protected].
+package com.rabbitmq.client.amqp.impl;
+
+import com.rabbitmq.client.amqp.CredentialsProvider;
+import com.rabbitmq.client.amqp.UsernamePasswordCredentialsProvider;
+import com.rabbitmq.client.amqp.oauth.GsonTokenParser;
+import com.rabbitmq.client.amqp.oauth.HttpTokenRequester;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+
+final class CredentialsFactory {
+
+  private volatile Credentials oauthCredentials;
+  private final Lock oauthCredentialsLock = new ReentrantLock();
+  private final AmqpEnvironment environment;
+
+  CredentialsFactory(AmqpEnvironment environment) {
+    this.environment = environment;
+  }
+
+  Credentials credentials(DefaultConnectionSettings<?> settings) {
+    CredentialsProvider provider = settings.credentialsProvider();
+    Credentials credentials;
+    if (settings.oauth().enabled()) {
+      // TODO consider OAuth credentials are not shared
+      credentials = oauthCredentials(settings);
+    } else {
+      if (provider instanceof UsernamePasswordCredentialsProvider) {
+        UsernamePasswordCredentialsProvider credentialsProvider =
+            (UsernamePasswordCredentialsProvider) provider;
+        credentials = new UsernamePasswordCredentials(credentialsProvider);
+      } else {
+        credentials = Credentials.NO_OP;
+      }
+    }
+    return credentials;
+  }
+
+  private Credentials oauthCredentials(DefaultConnectionSettings<?> connectionSettings) {
+    Credentials result = this.oauthCredentials;
+    if (result != null) {
+      return result;
+    }
+
+    this.oauthCredentialsLock.lock();
+    try {
+      if (this.oauthCredentials == null) {
+        DefaultConnectionSettings.DefaultOAuthSettings<?> settings = connectionSettings.oauth();
+        // TODO set TLS configuration on TLS requester
+        // TODO use pre-configured token requester if any
+        HttpTokenRequester tokenRequester =
+            new HttpTokenRequester(
+                settings.tokenEndpointUri(),
+                settings.clientId(),
+                settings.clientSecret(),
+                settings.grantType(),
+                settings.parameters(),
+                null,
+                null,
+                null,
+                null,
+                new GsonTokenParser());
+        this.oauthCredentials =
+            new TokenCredentials(tokenRequester, environment.scheduledExecutorService());
+      }
+      return this.oauthCredentials;
+    } finally {
+      this.oauthCredentialsLock.unlock();
+    }
+  }
+}

src/main/java/com/rabbitmq/client/amqp/impl/DefaultConnectionSettings.java

@@ -28,9 +28,7 @@
 import java.security.KeyManagementException;
 import java.security.NoSuchAlgorithmException;
 import java.time.Duration;
-import java.util.Collections;
-import java.util.List;
-import java.util.Random;
+import java.util.*;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.stream.Collectors;
 import javax.net.ssl.SSLContext;
@@ -74,6 +72,7 @@ abstract class DefaultConnectionSettings<T> implements ConnectionSettings<T> {
   private String saslMechanism = ConnectionSettings.SASL_MECHANISM_ANONYMOUS;
   private final DefaultTlsSettings<T> tlsSettings = new DefaultTlsSettings<>(this);
   private final DefaultAffinity<T> affinity = new DefaultAffinity<>(this);
+  private final DefaultOAuthSettings<T> oAuthSettings = new DefaultOAuthSettings<>(this);
 
   @Override
   public T uri(String uriString) {
@@ -223,6 +222,10 @@ void copyTo(DefaultConnectionSettings<?> copy) {
     }
 
     this.affinity.copyTo(copy.affinity);
+
+    if (this.oAuthSettings.enabled()) {
+      this.oAuthSettings.copyTo((DefaultOAuthSettings<?>) copy.oauth());
+    }
   }
 
   DefaultConnectionSettings<?> consolidate() {
@@ -295,6 +298,11 @@ public DefaultAffinity<? extends T> affinity() {
     return this.affinity;
   }
 
+  @Override
+  public DefaultOAuthSettings<? extends T> oauth() {
+    return this.oAuthSettings;
+  }
+
   static DefaultConnectionSettings<?> instance() {
     return new DefaultConnectionSettings<>() {
       @Override
@@ -490,4 +498,89 @@ void validate() {
       }
     }
   }
+
+  static class DefaultOAuthSettings<T> implements OAuthSettings<T> {
+
+    private final DefaultConnectionSettings<T> connectionSettings;
+    private final Map<String, String> parameters = new HashMap<>();
+    private String tokenEndpointUri;
+    private String clientId;
+    private String clientSecret;
+    private String grantType = "client_credentials";
+
+    DefaultOAuthSettings(DefaultConnectionSettings<T> connectionSettings) {
+      this.connectionSettings = connectionSettings;
+    }
+
+    @Override
+    public OAuthSettings<T> tokenEndpointUri(String uri) {
+      this.tokenEndpointUri = uri;
+      return this;
+    }
+
+    @Override
+    public OAuthSettings<T> clientId(String clientId) {
+      this.clientId = clientId;
+      return this;
+    }
+
+    @Override
+    public OAuthSettings<T> clientSecret(String clientSecret) {
+      this.clientSecret = clientSecret;
+      return this;
+    }
+
+    @Override
+    public OAuthSettings<T> grantType(String grantType) {
+      this.grantType = grantType;
+      return this;
+    }
+
+    @Override
+    public OAuthSettings<T> parameter(String name, String value) {
+      if (value == null) {
+        this.parameters.remove(name);
+      } else {
+        this.parameters.put(name, value);
+      }
+      return this;
+    }
+
+    @Override
+    public T connection() {
+      return this.connectionSettings.toReturn();
+    }
+
+    void copyTo(DefaultOAuthSettings<?> copy) {
+      copy.tokenEndpointUri(this.tokenEndpointUri);
+      copy.clientId(this.clientId);
+      copy.clientSecret(this.clientSecret);
+      copy.grantType(this.grantType);
+      this.parameters.forEach(copy::parameter);
+    }
+
+    String tokenEndpointUri() {
+      return this.tokenEndpointUri;
+    }
+
+    String clientId() {
+      return this.clientId;
+    }
+
+    String clientSecret() {
+      return this.clientSecret;
+    }
+
+    String grantType() {
+      return this.grantType;
+    }
+
+    Map<String, String> parameters() {
+      return Map.copyOf(this.parameters);
+    }
+
+    boolean enabled() {
+      return this.tokenEndpointUri != null;
+    }
+  }
 }