Support shared OAuth token

Author: acogoluegnes

pom.xml

@@ -45,15 +45,16 @@
         <netty4.version>4.1.115.Final</netty4.version>
         <netty4.iouring.version>0.0.25.Final</netty4.iouring.version>
         <micrometer.version>1.14.1</micrometer.version>
+        <gson.version>2.11.0</gson.version>
         <logback.version>1.2.13</logback.version>
         <junit.jupiter.version>5.11.3</junit.jupiter.version>
         <assertj.version>3.26.3</assertj.version>
         <mockito.version>5.14.2</mockito.version>
         <jqwik.version>1.9.2</jqwik.version>
-        <amqp-client.version>5.22.0</amqp-client.version>
         <micrometer-tracing-test.version>1.4.0</micrometer-tracing-test.version>
         <micrometer-docs-generator.version>1.0.4</micrometer-docs-generator.version>
         <jose4j.version>0.9.6</jose4j.version>
+        <commons-lang3.version>3.17.0</commons-lang3.version>
         <maven.compiler.plugin.version>3.13.0</maven.compiler.plugin.version>
         <maven.dependency.plugin.version>3.8.1</maven.dependency.plugin.version>
         <maven-surefire-plugin.version>3.5.2</maven-surefire-plugin.version>
@@ -76,7 +77,6 @@
         <gpg.skip>true</gpg.skip>
         <gpg.keyname>6026DFCA</gpg.keyname>
         <spotbugs.skip>false</spotbugs.skip>
-        <gson.version>2.11.0</gson.version>
     </properties>
 
     <dependencies>
@@ -251,6 +251,11 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-lang3</artifactId>
+            <version>${commons-lang3.version}</version>
+        </dependency>
 
     </dependencies>
 

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

@@ -29,5 +29,7 @@ public interface OAuthSettings<T> {
 
   OAuthSettings<T> parameter(String name, String value);
 
+  OAuthSettings<T> shared(boolean shared);
+
   T connection();
 }

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

@@ -98,7 +98,7 @@ final class AmqpConnection extends ResourceBase implements Connection {
     this.topologyListener = createTopologyListener(builder);
 
     if (recoveryConfiguration.activated()) {
-      disconnectHandler = recoveryDisconnectHandler(recoveryConfiguration, builder.name());
+      disconnectHandler = recoveryDisconnectHandler(recoveryConfiguration, this.name());
     } else {
       disconnectHandler =
           (c, e) -> {
@@ -121,21 +121,30 @@ final class AmqpConnection extends ResourceBase implements Connection {
     Credentials credentials = builder.credentials();
     this.credentialsRegistration =
         credentials.register(
+            this.name(),
             (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());
+              State state = this.state();
+              if (state == OPEN) {
+                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());
+              } else {
+                LOGGER.debug(
+                    "Could not set new token for connection {} because its state is {}",
+                    this.name(),
+                    state);
+              }
             });
     LOGGER.debug("Opening native connection for connection '{}'...", this.name());
     NativeConnectionWrapper ncw =
         ConnectionUtils.enforceAffinity(
             addrs -> {
               NativeConnectionWrapper wrapper =
-                  connect(this.connectionSettings, builder.name(), disconnectHandler, addrs);
+                  connect(this.connectionSettings, this.name(), disconnectHandler, addrs);
               this.nativeConnection = wrapper.connection();
               return wrapper;
             },
@@ -419,7 +428,10 @@ private void recoverAfterConnectionFailure(
                           () -> {
                             if (!this.recoveringConnection.get()) {
                               recoverAfterConnectionFailure(
-                                  recoveryConfiguration, name, ex, disconnectedHandlerReference);
+                                  recoveryConfiguration,
+                                  this.name(),
+                                  ex,
+                                  disconnectedHandlerReference);
                             }
                           });
                 }
@@ -822,7 +834,7 @@ public int hashCode() {
     return Objects.hashCode(id);
   }
 
-  private static class TokenConnectionCallback implements Credentials.ConnectionCallback {
+  private static class TokenConnectionCallback implements Credentials.AuthenticationCallback {
 
     private final ConnectionOptions options;
 
@@ -831,15 +843,8 @@ private TokenConnectionCallback(ConnectionOptions options) {
     }
 
     @Override
-    public Credentials.ConnectionCallback username(String username) {
-      options.user(username);
-      return this;
-    }
-
-    @Override
-    public Credentials.ConnectionCallback password(String password) {
-      options.password(password);
-      return this;
+    public void authenticate(String username, String password) {
+      options.user(username).password(password);
     }
   }
 }

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

@@ -21,39 +21,32 @@ interface Credentials {
 
   Credentials NO_OP = new NoOpCredentials();
 
-  Registration register(RefreshCallback refreshCallback);
+  Registration register(String name, AuthenticationCallback refreshCallback);
 
   interface Registration {
 
-    void connect(ConnectionCallback callback);
+    void connect(AuthenticationCallback callback);
 
     void unregister();
   }
 
-  interface ConnectionCallback {
+  interface AuthenticationCallback {
 
-    ConnectionCallback username(String username);
-
-    ConnectionCallback password(String password);
-  }
-
-  interface RefreshCallback {
-
-    void refresh(String username, String password);
+    void authenticate(String username, String password);
   }
 
   class NoOpCredentials implements Credentials {
 
     @Override
-    public Registration register(RefreshCallback refreshCallback) {
+    public Registration register(String name, AuthenticationCallback refreshCallback) {
       return new NoOpRegistration();
     }
   }
 
   class NoOpRegistration implements Registration {
 
     @Override
-    public void connect(ConnectionCallback callback) {}
+    public void connect(AuthenticationCallback callback) {}
 
     @Override
     public void unregister() {}

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

@@ -26,7 +26,7 @@
 
 final class CredentialsFactory {
 
-  private volatile Credentials oauthCredentials;
+  private volatile Credentials globalOAuthCredentials;
   private final Lock oauthCredentialsLock = new ReentrantLock();
   private final AmqpEnvironment environment;
 
@@ -38,8 +38,11 @@ Credentials credentials(DefaultConnectionSettings<?> settings) {
     CredentialsProvider provider = settings.credentialsProvider();
     Credentials credentials;
     if (settings.oauth().enabled()) {
-      // TODO consider OAuth credentials are not shared
-      credentials = oauthCredentials(settings);
+      if (settings.oauth().shared()) {
+        credentials = globalOAuthCredentials(settings);
+      } else {
+        credentials = createOAuthCredentials(settings);
+      }
     } else {
       if (provider instanceof UsernamePasswordCredentialsProvider) {
         UsernamePasswordCredentialsProvider credentialsProvider =
@@ -52,36 +55,39 @@ Credentials credentials(DefaultConnectionSettings<?> settings) {
     return credentials;
   }
 
-  private Credentials oauthCredentials(DefaultConnectionSettings<?> connectionSettings) {
-    Credentials result = this.oauthCredentials;
+  private Credentials globalOAuthCredentials(DefaultConnectionSettings<?> connectionSettings) {
+    Credentials result = this.globalOAuthCredentials;
     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());
+      if (this.globalOAuthCredentials == null) {
+        this.globalOAuthCredentials = createOAuthCredentials(connectionSettings);
       }
-      return this.oauthCredentials;
+      return this.globalOAuthCredentials;
     } finally {
       this.oauthCredentialsLock.unlock();
     }
   }
+
+  private Credentials createOAuthCredentials(DefaultConnectionSettings<?> connectionSettings) {
+    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());
+    return new TokenCredentials(tokenRequester, environment.scheduledExecutorService());
+  }
 }

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

@@ -507,13 +507,15 @@ static class DefaultOAuthSettings<T> implements OAuthSettings<T> {
     private String clientId;
     private String clientSecret;
     private String grantType = "client_credentials";
+    private boolean shared = true;
 
     DefaultOAuthSettings(DefaultConnectionSettings<T> connectionSettings) {
       this.connectionSettings = connectionSettings;
     }
 
     @Override
     public OAuthSettings<T> tokenEndpointUri(String uri) {
+      this.connectionSettings.saslMechanism(SASL_MECHANISM_PLAIN);
       this.tokenEndpointUri = uri;
       return this;
     }
@@ -546,6 +548,12 @@ public OAuthSettings<T> parameter(String name, String value) {
       return this;
     }
 
+    @Override
+    public OAuthSettings<T> shared(boolean shared) {
+      this.shared = shared;
+      return this;
+    }
+
     @Override
     public T connection() {
       return this.connectionSettings.toReturn();
@@ -556,6 +564,7 @@ void copyTo(DefaultOAuthSettings<?> copy) {
       copy.clientId(this.clientId);
       copy.clientSecret(this.clientSecret);
       copy.grantType(this.grantType);
+      copy.shared(this.shared);
       this.parameters.forEach(copy::parameter);
     }
 
@@ -579,6 +588,10 @@ Map<String, String> parameters() {
       return Map.copyOf(this.parameters);
     }
 
+    boolean shared() {
+      return this.shared;
+    }
+
     boolean enabled() {
       return this.tokenEndpointUri != null;
     }